yandex tank psycopg gun example

import json
import logging
#import sys
#sys.path.append('/usr/lib/python2.7/dist-packages/')
#import psycopg2
import os
from psycopg2.pool import ThreadedConnectionPool

gun_log = logging.getLogger(__name__)
gun_log.setLevel(logging.INFO)
fh = logging.FileHandler('%s.log' % __name__)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
fh.setFormatter(formatter)
gun_log.addHandler(fh)
gun_log.propagate=False

class ProcessSafePoolManager:

    def __init__(self, *args, **kwargs):
        self.last_seen_process_id = os.getpid()
        self.args = args
        self.kwargs = kwargs
        self._init()

    def _init(self):
        self._pool = ThreadedConnectionPool(*self.args, **self.kwargs)

    def getconn(self):
        current_pid = os.getpid()
        if not (current_pid == self.last_seen_process_id):
            self._init()
            self.last_seen_process_id = current_pid
        return self._pool.getconn()

    def putconn(self, conn, cl):
        return self._pool.putconn(conn, close=cl)

    def closeall(self):
        return self._pool.closeall()

def select(pool, query, params):
    try:
        pctx = pool.getconn()
        c = pctx.cursor()
        c.execute(query, params)
        results_count = len(c.fetchall())
        pctx.commit()
        c.close()
        if results_count>0:
            return 0, 200, results_count
        else:
            return 0, 404, results_count
    except Exception as e:
        gun_log.error("Query failed: %s\n%s", query, e)
        return 1, 0, None

    finally:
        pool.putconn(pctx, False)


class LoadTest(object):
    def __init__(self, gun):
        self.gun = gun
        self.sql_pool = None
        self.log_empty = int(self.gun.get_option('log_empty', '0'))
        self.host = self.gun.get_option('host')
        self.port = self.gun.get_option('port', '6432')
        self.db = self.gun.get_option('db')
        return self._pool.closeall()

def select(pool, query, params):
    try:
        pctx = pool.getconn()
        c = pctx.cursor()
        c.execute(query, params)
        results_count = len(c.fetchall())
        pctx.commit()
        c.close()
        if results_count>0:
            return 0, 200, results_count
        else:
            return 0, 404, results_count
    except Exception as e:
        gun_log.error("Query failed: %s\n%s", query, e)
        return 1, 0, None

    finally:
        pool.putconn(pctx, False)


class LoadTest(object):
    def __init__(self, gun):
        self.gun = gun
        self.sql_pool = None
        self.log_empty = int(self.gun.get_option('log_empty', '0'))
        self.host = self.gun.get_option('host')
        self.port = self.gun.get_option('port', '6432')
        self.db = self.gun.get_option('db')
        self.user = self.gun.get_option('user')
        self.password = self.gun.get_option('password')
        self.tpls = {'select1' : "SELECT %(num)s",
                     'select2' : "SELECT %(str)s"}

    def select(self, missile):
        try:
            tpl, params = missile.split('\t')
            query = self.tpls[tpl]
            params = json.loads(params)
            with self.gun.measure(tpl) as sample:
                sample["net_code"], sample["proto_code"], results = select(self.sql_pool, query, params)
                gun_log.info("AMMO %s, %s, %s\n", query, params, results)
        except Exception as e:
            gun_log.error('Scenario failed with %s', e, exc_info=True)

    def setup(self, settings):
        self.sql_pool = ProcessSafePoolManager(1, 1, host=self.host, port=self.port, database=self.db, user=self.user, password=self.password)
        gun_log.info('Pool created')

    def teardown(self):
        self.sql_pool.closeall()
        gun_log.info('Pool closed')
        os._exit(0)