def __init__(self, min_con, max_con, *args, **kwargs):
self._pool = ThreadedConnectionPool(
min_con,
max_con,
*args,
connection_factory=NamedTupleConnection,
**kwargs)
class Postgres(object):
def __init__(self, app=None, pool_size=10):
self.app = app
self.pool = None
self.pool_size = pool_size
if app is not None:
self.init_app(app)
def init_app(self, app):
self.pool = ThreadedConnectionPool(1, self.pool_size,
app.config['POSTGRES_URL'])
app.teardown_appcontext(self.teardown)
def teardown(self, exception):
ctx = _app_ctx_stack.top
conn = getattr(ctx, 'postgres_conn', None)
if conn is not None:
self.pool.putconn(conn)
del ctx.postgres_conn
def _connect(self):
return self.pool.getconn()
@property
def connection(self):
ctx = _app_ctx_stack.top
if ctx is not None:
if not hasattr(ctx, 'postgres_conn'):
ctx.postgres_conn = self._connect()
return ctx.postgres_conn
def cursor(self):
return self.connection.cursor()
def setUp(self):
self.createdb('quast_test')
# Run setup.sql
file_path = os.path.abspath(
os.path.join(os.path.dirname(os.path.abspath(__file__)), '..',
'sql/setup.sql'))
self.conn = psycopg2.connect("user=postgres dbname=quast_test")
self.conn.autocommit = True
with open(file_path) as f:
with self.conn.cursor() as curs:
sql = f.read()
curs.execute(sql)
# Run populate.sql
file_path = os.path.abspath(
os.path.join(os.path.dirname(os.path.abspath(__file__)), '..',
'sql/populate.sql'))
with open(file_path) as f:
with self.conn.cursor() as curs:
sql = f.read()
curs.execute(sql)
self.pool = ThreadedConnectionPool(1, 10,
"user=postgres dbname=quast_test")
def main(n=3000):
dsn = os.environ['POSTGRES_DSN']
pool = ThreadedConnectionPool(1, 2, dsn=dsn)
conn = pool.getconn()
conn.autocommit = True
write_centroids_to_db(conn, n)
pool.putconn(conn)
def _initialize_pool(self):
if self.pool is None:
logger.info("Initializing connection pool.")
args, kwargs = self.pool_params
self.pool = ThreadedConnectionPool(*args,
cursor_factory=RealDictCursor,
**kwargs)
class PgConnectionPool:
def __init__(self,
*args,
min_conns=1,
keep_conns=10,
max_conns=10,
**kwargs):
self._pool = ThreadedConnectionPool(min_conns, max_conns, *args,
**kwargs)
self._keep_conns = keep_conns
def acquire(self):
pool = self._pool
conn = pool.getconn()
pool.minconn = min(self._keep_conns, len(pool._used))
return conn
def release(self, conn):
self._pool.putconn(conn)
def close(self):
if hasattr(self, "_pool"):
self._pool.closeall()
__del__ = close
class _pool(SingletonMixin):
""" Pg pool class to provied connections,
it should be singleton and thread safe
"""
def __init__(self):
log.trace('Starting init pg pool...')
try:
log.trace('Database type: %s' % conf.db_type)
log.trace('Generated "dsn": "%s"' % conf.dsn)
self.pool = ThreadedConnectionPool(conf.pool_min,
conf.pool_max, conf.dsn)
except Exception as err:
log.error('Failed to init pg pool. %s' % str(err))
def get_con(self):
try:
# Except when pool is poor
__con = self.pool.getconn()
except Exception as err:
log.error('Error get connection from pool: %s' % str(err))
__con.autocommit = conf.autocommit
return __con
def put_con(self, conn):
if conn is not None and not conn.closed:
self.pool.putconn(conn=conn)
def __init__(self,
db_info=None,
hstore=False,
log=None,
logf=None,
min=1,
max=100,
default_cursor=DictCursor):
if db_info is None:
raise ValueError("Invalid connection Params")
self.pool = ThreadedConnectionPool(
min,
max,
database=db_info['database'],
user=db_info['user'],
password=db_info['password'],
host=db_info['host'],
port=db_info['port'],
)
self.hstore = hstore
self.log = log
self.logf = logf or (lambda cursor: cursor.query.decode())
self.default_cursor = default_cursor
self.prepared_statement_id = 0
def open_pool(self, minconns: int = 1,
maxconns: Optional[int] = None) -> None:
"""Initialises and opens a connection pool to the PostgreSQL database
(psycopg2.pool.ThreadedConnectionPool).
'minconn' new connections are created immediately. The connection pool
will support a maximum of about 'maxconn' connections.
Args:
minconns: minimum amount of available connections in the pool,
created on startup.
maxconns: maximum amount of available connections supported by
the pool. Defaults to 'minconns'.
"""
# If the pool hasn't been opened yet
if self.pool.closed:
# initialize max number of supported connections
maxconns = maxconns if maxconns is not None else minconns
# create a connection pool based on jumbo's configuration settings
self.pool = ThreadedConnectionPool(
minconns, maxconns,
host=self.config.DATABASE_HOST,
user=self.config.DATABASE_USERNAME,
password=self.config.DATABASE_PASSWORD,
port=self.config.DATABASE_PORT,
dbname=self.config.DATABASE_NAME,
sslmode='disable')
logger.info(f"Connection pool created to PostgreSQL database: "
f"{maxconns} connections available.")
class PgUtil():
__metaclass__ = Singleton
def __init__(self):
self.conn_pool = ThreadedConnectionPool(
minconn=4,
maxconn=100,
database='coda',
user='******',
password='******',
host='127.0.0.1',
port=5432
)
def get_conn(self):
conn = self.conn_pool.getconn()
return conn
def put_conn(self, conn):
self.conn_pool.putconn(conn)
def execute_insert_sql(self, sql, values):
conn = self.get_conn()
cur = conn.cursor(cursor_factory=DictCursor)
cur.execute(sql,values)
cur.close()
conn.commit()
self.put_conn(conn)
def __init__(self, database_config):
self.database_config = database_config
pool_config = self.database_config['pool']
conn_config = self.database_config['conn']
self.conn_pool = ThreadedConnectionPool(minconn=pool_config['min'],
maxconn=pool_config['max'],
**conn_config)
def __init__(
self,
table="footprints",
database_url=os.getenv("DATABASE_URL"),
geometry_column="geom",
):
if database_url is None:
raise Exception("Database URL must be provided.")
urlparse.uses_netloc.append("postgis")
urlparse.uses_netloc.append("postgres")
url = urlparse.urlparse(database_url)
self._pool = ThreadedConnectionPool(
0,
16,
database=url.path[1:],
user=url.username,
password=url.password,
host=url.hostname,
port=url.port,
)
self._log = logging.getLogger(__name__)
self.table = table
self.geometry_column = geometry_column
def __init__(self, max_pool_size: int, *, dsn):
self._pool = ThreadedConnectionPool(
1,
max_pool_size,
dsn=dsn,
cursor_factory=RealDictCursor,
)
class Pgsql(PgsqlBase):
def __init__(self, min_con, max_con, *args, **kwargs):
self._pool = ThreadedConnectionPool(
min_con,
max_con,
*args,
connection_factory=NamedTupleConnection,
**kwargs)
def exec(self, sql, args=None):
def t(conn):
with conn.cursor() as c:
c.execute(sql, args)
try:
res = c.fetchall()
return res
except ProgrammingError:
pass
r = self.transaction(t)
return r
def transaction(self, operation):
conn = self._pool.getconn()
try:
with conn:
return operation(conn)
finally:
self._pool.putconn(conn)
def connect(self, ctx_list=None, mark=None):
"""Connect to database."""
self.log.debug("Connecting to database: %(dbname)s as %(user)s" %
self.__params__)
try:
self.lock.acquire_lock()
try:
if not self.__conn_pool__:
connector = (
"host=%(host)s port=%(port)s user=%(user)s password=%(password)s dbname=%(dbname)s"
% self.__params__)
self.__conn_pool__ = ThreadedConnectionPool(
1, 20, connector)
finally:
self.lock.release_lock()
if not ctx_list:
ctx_list = []
ctx_list.append(self._new_ctx(mark))
self.log.debug(
"ctx:(%(mark)s, %(ctxid)s) Connected to database: %(dbname)s as %(user)s"
% {
"mark": ctx_list[-1]["mark"],
"ctxid": str(id(ctx_list[-1])),
"dbname": self.__params__["dbname"],
"user": self.__params__["user"],
})
return ctx_list
except psycopg2.Error as _error:
# We do not want our users to have to 'import psycopg2' to
# handle the module's underlying database errors
_, value, traceback = sys.exc_info()
raise self.DatabaseError(value).with_traceback(traceback)
def test_context_manager(self):
"""
test using the context manager to access the pool
"""
min_connections = 1
max_connections = 5
test_number = 42
connection_pool = ThreadedConnectionPool(min_connections,
max_connections,
**_database_credentials)
test_greenlet = ContextWriteGreenlet(connection_pool, test_number, 3.0)
rollback_greenlet = ContextRollbackGreenlet(connection_pool, 3.0)
test_greenlet.start()
rollback_greenlet.start()
test_greenlet.join()
self.assertTrue(test_greenlet.successful())
rollback_greenlet.join()
self.assertTrue(rollback_greenlet.successful())
result = test_greenlet.value
self.assertEqual(result, [(test_number, )])
connection_pool.closeall()
def test_decorator(self):
"""
test using the decorator to access the pool
"""
global _connection_pool
min_connections = 1
max_connections = 5
test_number = 42
_connection_pool = ThreadedConnectionPool(min_connections,
max_connections,
**_database_credentials)
test_greenlet = DecoratorWriteGreenlet(test_number, 3.0)
rollback_greenlet = DecoratorRollbackGreenlet(3.0)
test_greenlet.start()
rollback_greenlet.start()
test_greenlet.join()
self.assertTrue(test_greenlet.successful())
rollback_greenlet.join()
self.assertTrue(rollback_greenlet.successful())
result = test_greenlet.value
self.assertEqual(result, [(test_number, )])
_connection_pool.closeall()
def __init__(self, env_path: Optional[str] = None) -> None:
"""Initializes manager to handle connections to a given PostgreSQL
database.
Args:
env_path: path where to look for the jumbo.env configuration
file. If not provided, looks for configuration file
in the working directory of the script invoking this
constructor.
Attributes:
config (jumbo.config.Config): jumbo's database configuration
settings.
pool (psycopg2.pool): connection pool to the PostgreSQL
database. Initially closed.
"""
# Database configuration settings
self.config = Config(env_path)
# initialize a placeholder connection pool
self.pool = ThreadedConnectionPool(0, 0)
self.pool.closed = True # keep it closed on construction
# Log configuration settings
logger.info(f"Jumbo Connection Manager created:\n{self.config}")
def __init__(self,
url=None,
hstore=False,
log=None,
logf=None,
min=1,
max=5,
default_cursor=DictCursor):
params = urlparse(url or os.environ.get('DATABASE_URL')
or 'postgres://localhost/')
if params.scheme != 'postgres':
raise ValueError(
"Invalid connection string (postgres://user@pass:host/db?param=value)"
)
self.pool = ThreadedConnectionPool(
min,
max,
database=params.path[1:] or parse_qs(params.query).get('dbname'),
user=params.username or parse_qs(params.query).get('user'),
password=params.password or parse_qs(params.query).get('password'),
host=params.hostname or parse_qs(params.query).get('host'),
port=params.port or parse_qs(params.query).get('port'),
)
self.hstore = hstore
self.log = log
self.logf = logf or (lambda cursor: cursor.query.decode())
self.default_cursor = default_cursor
self.prepared_statement_id = 0
class PostgresqlConnection:
def __init__(self, **kwargs):
self.pool = ThreadedConnectionPool(minconn=kwargs['minconn'],
maxconn=kwargs['maxconn'],
database=kwargs['database'],
user=kwargs['username'],
password=kwargs['password'],
host=kwargs['host'],
port=kwargs['port'])
@contextmanager
def cursor(self, auto_commit=True):
conn = self.pool.getconn()
cursor = conn.cursor()
try:
yield cursor
finally:
if auto_commit:
conn.commit()
if cursor is not None and not cursor.closed:
cursor.close()
self.pool.putconn(conn)
def close(self):
self.pool.closeall()
def test_context_manager(self):
"""
test using the context manager to access the pool
"""
min_connections = 1
max_connections = 5
test_number = 42
connection_pool = ThreadedConnectionPool(min_connections,
max_connections,
**_database_credentials)
test_greenlet = ContextWriteGreenlet(connection_pool, test_number, 3.0)
rollback_greenlet = ContextRollbackGreenlet(connection_pool, 3.0)
test_greenlet.start()
rollback_greenlet.start()
test_greenlet.join()
self.assertTrue(test_greenlet.successful())
rollback_greenlet.join()
self.assertTrue(rollback_greenlet.successful())
result = test_greenlet.value
self.assertEqual(result, [(test_number, )])
connection_pool.closeall()
def test_decorator(self):
"""
test using the decorator to access the pool
"""
global _connection_pool
min_connections = 1
max_connections = 5
test_number = 42
_connection_pool = ThreadedConnectionPool(min_connections,
max_connections,
**_database_credentials)
test_greenlet = DecoratorWriteGreenlet(test_number, 3.0)
rollback_greenlet = DecoratorRollbackGreenlet(3.0)
test_greenlet.start()
rollback_greenlet.start()
test_greenlet.join()
self.assertTrue(test_greenlet.successful())
rollback_greenlet.join()
self.assertTrue(rollback_greenlet.successful())
result = test_greenlet.value
self.assertEqual(result, [(test_number, )])
_connection_pool.closeall()
def create_pool(self, conn_dict):
"""
Create a connection pool
:param conn_dict: connection params dictionary
:type conn_dict: dict
"""
if conn_dict["Host"] is None:
self.host = 'localhost'
else:
self.host = conn_dict["Host"]
if conn_dict["Port"] is None:
self.port = '5432'
else:
self.port = conn_dict["Port"]
self.database = conn_dict["Database"]
self.user = conn_dict["User"]
self.passwd = conn_dict["Password"]
conn_params = "host='{host}' dbname='{db}' user='******' password='******' port='{port}'".format(
host=self.host,
db=self.database,
user=self.user,
passwd=self.passwd,
port=self.port)
try:
self.pool = ThreadedConnectionPool(1, 50, conn_params)
except Exception as e:
print(e.message)
def __init__(self, dbconn):
if Manager.SINGLETON is not None:
raise RuntimeError(
"Attempt to create multiple MonitorManager instances")
else:
Manager.SINGLETON = self
self.log = logging.getLogger("Manager")
# set up database pool
self.dbconn = dbconn
self.pool = ThreadedConnectionPool(20, 100, dbconn)
# prepare database to recover jobs
self.getDao().ready_active_jobs()
# create executors
self.execs = {} # the executor registry
self.default_executor = None # the default executor for jobs
self._create_executors()
# monitor registry
self.monitors = {} # the monitor registry
self._create_monitors()
self.log.info("Manager initialized and running")
class DBConnection:
def __init__(self, db, *args, **kwargs):
self.pool = ThreadedConnectionPool(1, 20, db)
self.obj = None
@contextmanager
def connect(self):
try:
conn = self.pool.getconn()
yield conn
finally:
self.pool.putconn(conn)
@contextmanager
def dict_cursor(self, commit=False):
with self.connect() as connection:
cursor = connection.cursor(cursor_factory=RealDictCursor)
try:
yield cursor
finally:
if commit:
connection.commit()
cursor.close()
@contextmanager
def cursor(self, commit=False):
with self.connect() as connection:
cursor = connection.cursor()
try:
yield cursor
finally:
if commit:
connection.commit()
cursor.close()
def execute(self, query, named=True, commit=False):
cur = self.dict_cursor if named else self.cursor
with cur(commit) as cursor:
try:
cursor.execute(query)
data = list(cursor)
except (
psycopg2.OperationalError,
psycopg2.errors.UniqueViolation,
) as error:
raise ValidationError(error.args[0].split("DETAIL:")[1],
status_code=400)
return data
def drop_all(self):
query = """SELECT table_schema,table_name FROM information_schema.tables\
WHERE table_schema = 'public' ORDER BY table_schema,table_name"""
rows = self.execute(query, named=False)
with self.cursor(commit=True) as cursor:
[
cursor.execute("drop table " + row[1] + " cascade")
for row in rows
]
def __init__(self):
self.conn_pool = ThreadedConnectionPool(minconn=4,
maxconn=100,
database='coda',
user='******',
password='******',
host='127.0.0.1',
port=5432)
def __init__(self):
self._pool = ThreadedConnectionPool(1,
12,
database="forum_server",
user="******",
password="******",
host="localhost",
port="5454")
def __init__(self, host, port, user, password):
min_conns, max_conns = 1, 20
self.pool = ThreadedConnectionPool(
min_conns, max_conns,
host=host, port=port,
user=user, password=password,
database="postgres",
)
def __init__(self, **kwargs):
self.pool = ThreadedConnectionPool(minconn=kwargs['minconn'],
maxconn=kwargs['maxconn'],
database=kwargs['database'],
user=kwargs['username'],
password=kwargs['password'],
host=kwargs['host'],
port=kwargs['port'])
def __init__(self):
self.pool = ThreadedConnectionPool(minconn=1,
maxconn=5,
user="******",
password="******",
host="localhost",
port="5432",
database="irdb")
def create_connection_pool(self):
""" Create the thread safe threaded postgres connection pool"""
max_conn = self.pid_max
min_conn = max_conn / 2
# creating separate connection for read and write purpose
self._select_conn_pool = ThreadedConnectionPool(
min_conn, max_conn, DB_URL)
def connect(self):
if self._connpool is not None:
return self
logger.info('connect to "%s"', self.name)
minconn = config.getint('database', 'minconn', default=1)
maxconn = config.getint('database', 'maxconn', default=64)
self._connpool = ThreadedConnectionPool(minconn, maxconn,
self.dsn(self.name))
return self
def _connect(self):
global _CONNECTION_POOL
if _CONNECTION_POOL is None:
_CONNECTION_POOL = ThreadedConnectionPool(
config.DB_MIN_CONNECTIONS, config.DB_MAX_CONNECTIONS,
**config.DB_PARAMS)
if self._connection is not None:
raise RuntimeError("Connection still exists.")
self._connection = _CONNECTION_POOL.getconn()
self._connection.set_session(autocommit=True)
def putconn(self, conn):
"""
Returns connection back to pool.
"""
#calledBy = traceback.extract_stack()[-2]
#logging.info("PUTCONN - FILE: " + calledBy[0] + ", LINE: " + str(calledBy[1]) + ", METHOD: " + calledBy[2])
ThreadedConnectionPool.putconn(self, conn)
class Database:
def __init__(self, connect_param):
self.__connect_param = connect_param
self.__pool = ThreadedConnectionPool(0, 10, self.__connect_param)
# get cursor and test it
# cur = self.cursor()
# cur.execute('SHOW transaction_read_only')
# standby = cur.fetchone()
# cur.close()
def get_connection(self):
return self.__pool.getconn()
def put_connection(self, connection):
self.__pool.putconn(connection)
def create_pool(self, conn_dict, limits):
"""
Create a connection pool
:param conn_dict: connection params dictionary
:type conn_dict: dict
"""
if conn_dict["Host"] is None:
self.host = 'localhost'
else:
self.host = conn_dict["Host"]
if conn_dict["Port"] is None:
self.port = '5432'
else:
self.port = conn_dict["Port"]
self.database = conn_dict["Database"]
self.user = conn_dict["User"]
self.passwd = conn_dict["Password"]
conn_params = "host='{host}' dbname='{db}' user='******' password='******' port='{port}'".format(
host=self.host, db=self.database, user=self.user, passwd=self.passwd, port=self.port
)
try:
logger.debug('creating pool')
self.pool = ThreadedConnectionPool(int(limits["Min"]), int(limits["Max"]), conn_params)
except Exception as e:
logger.exception(e.message)
class Database():
def __init__(self, config):
logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s', level=logging.INFO)
self._pool = ThreadedConnectionPool(1, 10,
database=config['DB_DATABASE'],
user=config['DB_USER'],
password=config['DB_PASSWORD'],
host=config['DB_HOST'],
async=False)
def get_connection(self):
return self._pool.getconn()
def put_away_connection(self, con):
self._pool.putconn(con)
def getconn(self):
"""
Gets connection from parent class, enables AUTOCOMMIT and returns requested connection.
@rtype: object
@return: connection with isolation level set to autocommit
"""
#calledBy = traceback.extract_stack()[-2]
#logging.info("GETCONN - FILE: " + calledBy[0] + ", LINE: " + str(calledBy[1]) + ", METHOD: " + calledBy[2])
conn = ThreadedConnectionPool.getconn(self)
try:
#conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
conn.cursor().execute("SELECT 1")
except (psycopg2.OperationalError, psycopg2.InterfaceError, psycopg2.InternalError):
key = self._rused[id(conn)]
del self._rused[id(conn)]
conn = psycopg2.connect(self.dsn)
self._rused[id(conn)] = key
if Config.hstoreEnabled == True:
try:
psycopg2.extras.register_hstore(conn)
except Exception, e:
Config.hstoreEnabled = False
def __init__(self, host, port, database, username, password, session_ttl=24*60*60*365, anon_session_ttl=24*60*60, session_renew=0, anon_session_renew=0, min_connections=1, max_connections=10): self.db = ThreadedConnectionPool(min_connections, max_connections, database=database, user=username, password=password, host=host, port=port) self.create_tables() self.session_ttl = session_ttl self.anon_session_ttl = anon_session_ttl or self.session_ttl self.session_renew = session_renew or self.session_ttl self.anon_session_renew = anon_session_renew or self.anon_session_ttl
class DB:
def __init__(self, *args, **kwargs):
self.pool_params = (args, kwargs)
self.pool = None
self.campaigns = Campaigns(self)
self.worksets = Worksets(self)
self.tasks = Tasks(self)
self.labels = Labels(self)
self.logger = logging.getLogger(__name__)
def _initialize_pool(self):
if self.pool is None:
logger.info("Initializing connection pool.")
args, kwargs = self.pool_params
self.pool = ThreadedConnectionPool(
*args, cursor_factory=RealDictCursor, **kwargs)
def execute(self, sql):
with self.transaction() as transactor:
cursor = transactor.cursor()
cursor.execute(sql)
return cursor
@contextmanager
def transaction(self):
"""Provides a transactional scope around a series of operations."""
self._initialize_pool()
conn = self.pool.getconn()
try:
yield conn
conn.commit()
except:
conn.rollback()
raise
finally:
self.pool.putconn(conn)
@classmethod
def from_config(cls, config):
# Copy config as kwargs
params = {k: v for k, v in config['database'].items()}
params['minconn'] = params.get('minconn', 1)
params['maxconn'] = params.get('maxconn', 5)
return cls(**params)
def __init__(self, config):
logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s', level=logging.INFO)
self._pool = ThreadedConnectionPool(1, 10,
database=config['DB_DATABASE'],
user=config['DB_USER'],
password=config['DB_PASSWORD'],
host=config['DB_HOST'],
async=False)
def connect(self):
if self._connpool is not None:
return self
logger.info('connect to "%s"', self.name)
minconn = config.getint('database', 'minconn', default=1)
maxconn = config.getint('database', 'maxconn', default=64)
self._connpool = ThreadedConnectionPool(
minconn, maxconn, self.dsn(self.name))
return self
def connect(self):
if self._connpool is not None:
return self
logger = logging.getLogger('database')
logger.info('connect to "%s"' % self.database_name)
host = CONFIG['db_host'] and "host=%s" % CONFIG['db_host'] or ''
port = CONFIG['db_port'] and "port=%s" % CONFIG['db_port'] or ''
name = "dbname=%s" % self.database_name
user = CONFIG['db_user'] and "user=%s" % CONFIG['db_user'] or ''
password = (CONFIG['db_password']
and "password=%s" % CONFIG['db_password'] or '')
minconn = int(CONFIG['db_minconn']) or 1
maxconn = int(CONFIG['db_maxconn']) or 64
dsn = '%s %s %s %s %s' % (host, port, name, user, password)
if dbgis:
self._connpool = GisThreadedConnectionPool(minconn, maxconn, dsn)
else:
self._connpool = ThreadedConnectionPool(minconn, maxconn, dsn)
return self
def __init__(self, settings):
from psycopg2.pool import ThreadedConnectionPool
dbsettings = settings['database']
self.pool = ThreadedConnectionPool(
minconn=1,
maxconn=settings['database']['conn_pool_size'],
database=dbsettings['name'],
user=dbsettings['username'],
password=dbsettings['password'],
host=dbsettings['host'],
port=dbsettings.get('port')
)
class PostgresThreadPool:
provides = ['db_connection_pool', 'postgres']
requires_configured = ['json_settings']
def __init__(self, settings):
from psycopg2.pool import ThreadedConnectionPool
dbsettings = settings['database']
self.pool = ThreadedConnectionPool(
minconn=1,
maxconn=settings['database']['conn_pool_size'],
database=dbsettings['name'],
user=dbsettings['username'],
password=dbsettings['password'],
host=dbsettings['host'],
port=dbsettings.get('port')
)
def getconn(self):
return self.pool.getconn()
def putconn(self, connection):
return self.pool.putconn(connection)
class ConnectionPool(object):
def __init__(self, conn_params, minconn=5, maxconn=5):
self._conn_params = conn_params.copy()
self._conn_params['minconn'] = minconn
self._conn_params['maxconn'] = maxconn
self._conn_pool = None
def initialize(self):
self._conn_pool = ThreadedConnectionPool(**self._conn_params)
@contextmanager
def cursor(self):
conn = self._conn_pool.getconn()
cursor = conn.cursor()
try:
yield cursor
conn.commit()
except Exception:
conn.rollback()
raise
finally:
self._conn_pool.putconn(conn)
class PgConnectionPool:
def __init__(self, *args, min_conns=1, keep_conns=10, max_conns=10,
**kwargs):
self._pool = ThreadedConnectionPool(
min_conns, max_conns, *args, **kwargs)
self._keep_conns = keep_conns
def acquire(self):
pool = self._pool
conn = pool.getconn()
pool.minconn = min(self._keep_conns, len(pool._used))
return conn
def release(self, conn):
self._pool.putconn(conn)
def close(self):
if hasattr(self, '_pool'):
self._pool.closeall()
__del__ = close
def __init__(self, db_config, table_raw=None, max_connections=10):
from psycopg2.pool import ThreadedConnectionPool
self.table_raw = table_raw
try:
self.pool = ThreadedConnectionPool(minconn=1,
maxconn=max_connections,
dsn="dbname={db_name} user={db_user} host={db_host} password={db_pass}"
.format(**db_config))
except Exception:
logger.exception("Error in db connection")
sys.exit(1)
logger.debug("Connected to database: {host}".format(host=db_config['db_host']))
def connect(self):
if self._connpool is not None:
return self
logger = logging.getLogger("database")
logger.info('connect to "%s"' % self.database_name)
host = CONFIG["db_host"] and "host=%s" % CONFIG["db_host"] or ""
port = CONFIG["db_port"] and "port=%s" % CONFIG["db_port"] or ""
name = "dbname=%s" % self.database_name
user = CONFIG["db_user"] and "user=%s" % CONFIG["db_user"] or ""
password = CONFIG["db_password"] and "password=%s" % CONFIG["db_password"] or ""
minconn = int(CONFIG["db_minconn"]) or 1
maxconn = int(CONFIG["db_maxconn"]) or 64
dsn = "%s %s %s %s %s" % (host, port, name, user, password)
self._connpool = ThreadedConnectionPool(minconn, maxconn, dsn)
return self
def connect(self):
if self._connpool is not None:
return self
logger.info('connect to "%s"', self.database_name)
uri = parse_uri(config.get('database', 'uri'))
assert uri.scheme == 'postgresql'
host = uri.hostname and "host=%s" % uri.hostname or ''
port = uri.port and "port=%s" % uri.port or ''
name = "dbname=%s" % self.database_name
user = uri.username and "user=%s" % uri.username or ''
password = ("password=%s" % urllib.unquote_plus(uri.password)
if uri.password else '')
minconn = config.getint('database', 'minconn', default=1)
maxconn = config.getint('database', 'maxconn', default=64)
dsn = '%s %s %s %s %s' % (host, port, name, user, password)
self._connpool = ThreadedConnectionPool(minconn, maxconn, dsn)
return self
def __init__(self,url=None,hstore=False,log=None,logf=None,min=1,max=5,
default_cursor=DictCursor):
params = urlparse.urlparse(url or
os.environ.get('DATABASE_URL') or
'postgres://localhost/')
self.pool = ThreadedConnectionPool(min,max,
database=params.path[1:],
user=params.username,
password=params.password,
host=params.hostname,
port=params.port,
)
self.hstore = hstore
self.log = log
self.logf = logf or (lambda cursor : cursor.query)
self.default_cursor = default_cursor
self.prepared_statement_id = 0
def _reset_pool(self):
""" First person here gets to reset it, others can continue on and try again """
if self.resetlock.acquire(False):
try:
if self.pool and not self.pool.closed:
self.pool.closeall()
# Make sure the basic database is present and create a PG connection pool
connkeys = { 'host':self.config['DBHOST'], 'port':self.config['DBPORT'], 'user':'******' }
ensure_database_created(connkeys)
ensure_public_schema(connkeys)
# Create a new pool of connections. Server should support 100, leave 10 for applications
self.pool = ThreadedConnectionPool(5, 80, cursor_factory=DictCursor, application_name="webserver", dbname="scorekeeper",
host=self.config['DBHOST'], port=self.config['DBPORT'], user=self.config['DBUSER'])
except Exception as e:
log.error("Error in pool create/reset: %s", str(e))
finally:
self.resetlock.release()
def __init__(self, config): """Configures the Db, connection is not created yet. @param config: instance of config.NotaryServerConfig.""" self.host = config.db_host self.port = config.db_port self.user = config.db_user self.password = config.db_password self.db_name = config.db_name self.min_connections = config.db_min_conn self.max_connections = config.db_max_conn self.pool = ThreadedConnectionPool( minconn = self.min_connections, maxconn = self.max_connections, host = self.host, port = self.port, user = self.user, password = self.password, database = self.db_name)
def __init__(self, host, port, dbname, dbuser, dbpass, minconn=1, maxconn=1):
# Thread pool
self.pool = ThreadedConnectionPool(
minconn=minconn,
maxconn=maxconn,
host=host,
database=dbname,
user=dbuser,
password=dbpass,
port=port
)
# Base connection for initialization
self.conn = psycopg2.connect(
host=host,
database=dbname,
user=dbuser,
password=dbpass,
port=port
)
self.curs = self.conn.cursor()
def __init__(self,url=None,hstore=False,log=None,logf=None,min=1,max=5,
default_cursor=DictCursor):
params = urlparse(url or
os.environ.get('DATABASE_URL') or
'postgres://localhost/')
if params.scheme != 'postgres':
raise ValueError("Invalid connection string (postgres://user@pass:host/db?param=value)")
self.pool = ThreadedConnectionPool(min,max,
database=params.path[1:] or parse_qs(params.query).get('dbname'),
user=params.username or parse_qs(params.query).get('user'),
password=params.password or parse_qs(params.query).get('password'),
host=params.hostname or parse_qs(params.query).get('host'),
port=params.port or parse_qs(params.query).get('port'),
)
self.hstore = hstore
self.log = log
self.logf = logf or (lambda cursor : cursor.query.decode())
self.default_cursor = default_cursor
self.prepared_statement_id = 0
def _create_connection_pool(self):
try:
self.connection_pool = ThreadedConnectionPool(
1, 16, **self.connection_params)
except Error as ex:
raise BackendError(str(ex)) from ex
class Psycopg2Backend(Backend):
"""Backend for accessing data stored in a Postgres database
"""
display_name = "PostgreSQL"
connection_pool = None
auto_create_extensions = True
def __init__(self, connection_params):
super().__init__(connection_params)
if self.connection_pool is None:
self._create_connection_pool()
if self.auto_create_extensions:
self._create_extensions()
def _create_connection_pool(self):
try:
self.connection_pool = ThreadedConnectionPool(
1, 16, **self.connection_params)
except Error as ex:
raise BackendError(str(ex)) from ex
def _create_extensions(self):
for ext in EXTENSIONS:
try:
query = "CREATE EXTENSION IF NOT EXISTS {}".format(ext)
with self.execute_sql_query(query):
pass
except OperationalError:
warnings.warn("Database is missing extension {}".format(ext))
def create_sql_query(self, table_name, fields, filters=(),
group_by=None, order_by=None,
offset=None, limit=None,
use_time_sample=None):
sql = ["SELECT", ', '.join(fields),
"FROM", table_name]
if use_time_sample is not None:
sql.append("TABLESAMPLE system_time(%i)" % use_time_sample)
if filters:
sql.extend(["WHERE", " AND ".join(filters)])
if group_by is not None:
sql.extend(["GROUP BY", ", ".join(group_by)])
if order_by is not None:
sql.extend(["ORDER BY", ",".join(order_by)])
if offset is not None:
sql.extend(["OFFSET", str(offset)])
if limit is not None:
sql.extend(["LIMIT", str(limit)])
return " ".join(sql)
@contextmanager
def execute_sql_query(self, query, params=None):
connection = self.connection_pool.getconn()
cur = connection.cursor()
try:
utfquery = cur.mogrify(query, params).decode('utf-8')
log.debug("Executing: %s", utfquery)
t = time()
cur.execute(query, params)
yield cur
log.info("%.2f ms: %s", 1000 * (time() - t), utfquery)
finally:
connection.commit()
self.connection_pool.putconn(connection)
def quote_identifier(self, name):
return '"%s"' % name
def unquote_identifier(self, quoted_name):
if quoted_name.startswith('"'):
return quoted_name[1:len(quoted_name) - 1]
else:
return quoted_name
def list_tables_query(self, schema=None):
if schema:
schema_clause = "AND n.nspname = '{}'".format(schema)
else:
schema_clause = "AND pg_catalog.pg_table_is_visible(c.oid)"
return """SELECT n.nspname as "Schema",
c.relname AS "Name"
FROM pg_catalog.pg_class c
LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
WHERE c.relkind IN ('r','v','m','S','f','')
AND n.nspname <> 'pg_catalog'
AND n.nspname <> 'information_schema'
AND n.nspname !~ '^pg_toast'
{}
AND NOT c.relname LIKE '\\_\\_%'
ORDER BY 1;""".format(schema_clause)
def create_variable(self, field_name, field_metadata,
type_hints, inspect_table=None):
if field_name in type_hints:
var = type_hints[field_name]
else:
var = self._guess_variable(field_name, field_metadata,
inspect_table)
field_name_q = self.quote_identifier(field_name)
if var.is_continuous:
if isinstance(var, TimeVariable):
var.to_sql = ToSql("extract(epoch from {})"
.format(field_name_q))
else:
var.to_sql = ToSql("({})::double precision"
.format(field_name_q))
else: # discrete or string
var.to_sql = ToSql("({})::text"
.format(field_name_q))
return var
def _guess_variable(self, field_name, field_metadata, inspect_table):
type_code = field_metadata[0]
FLOATISH_TYPES = (700, 701, 1700) # real, float8, numeric
INT_TYPES = (20, 21, 23) # bigint, int, smallint
CHAR_TYPES = (25, 1042, 1043,) # text, char, varchar
BOOLEAN_TYPES = (16,) # bool
DATE_TYPES = (1082, 1114, 1184, ) # date, timestamp, timestamptz
# time, timestamp, timestamptz, timetz
TIME_TYPES = (1083, 1114, 1184, 1266,)
if type_code in FLOATISH_TYPES:
return ContinuousVariable(field_name)
if type_code in TIME_TYPES + DATE_TYPES:
tv = TimeVariable(field_name)
tv.have_date |= type_code in DATE_TYPES
tv.have_time |= type_code in TIME_TYPES
return tv
if type_code in INT_TYPES: # bigint, int, smallint
if inspect_table:
values = self.get_distinct_values(field_name, inspect_table)
if values:
return DiscreteVariable(field_name, values)
return ContinuousVariable(field_name)
if type_code in BOOLEAN_TYPES:
return DiscreteVariable(field_name, ['false', 'true'])
if type_code in CHAR_TYPES:
if inspect_table:
values = self.get_distinct_values(field_name, inspect_table)
if values:
return DiscreteVariable(field_name, values)
return StringVariable(field_name)
def count_approx(self, query):
sql = "EXPLAIN " + query
with self.execute_sql_query(sql) as cur:
s = ''.join(row[0] for row in cur.fetchall())
return int(re.findall(r'rows=(\d*)', s)[0])
def __getstate__(self):
# Drop connection_pool from state as it cannot be pickled
state = dict(self.__dict__)
state.pop('connection_pool', None)
return state
def __setstate__(self, state):
# Create a new connection pool if none exists
self.__dict__.update(state)
if self.connection_pool is None:
self._create_connection_pool()
class DatabaseManager:
"""
This class provides abstraction over underlying database.
"""
def __init__(self, db_name="test_db", db_pass="", host="127.0.0.1" , port="5432"):
self.connection_pool = ThreadedConnectionPool(10, 50, database=db_name, user="******", \
password=db_pass, host=host, port=port)
self.logger = get_logger()
def __execute_query(self, query):
connection = self.connection_pool.getconn()
cursor = connection.cursor()
self.logger.debug("Going to execute query {}".format(query))
try:
cursor.execute(query)
except ProgrammingError:
self.logger.error("Error occurred while executing query {}".format(query))
except IntegrityError:
self.logger.error("Query failed. Duplicate row for query {}".format(query))
finally:
connection.commit()
self.connection_pool.putconn(connection)
"""
Inserts multiple rows in table_name. column_headers contain tuple of table headers.
rows contain the list of tuples where each tuple has values for each rows. The values in
tuple are ordered according to column_headers tuple.
"""
def insert_batch(self, table_name, column_headers, rows):
query = "INSERT INTO {} {} VALUES {}".format(table_name, '(' + ','.join(column_headers) + ')', str(rows)[1:-1])
self.__execute_query(query)
"""
Updates a row(uid) with new values from column_vs_value dict.
"""
def update(self, table_name, column_vs_value, uid):
update_str = ''.join('{}={},'.format(key, val) for key, val in column_vs_value.items())[:-1]
query = "UPDATE {} SET {} WHERE id = {} ".format(table_name, update_str, uid)
self.__execute_query(query)
"""
Deletes all rows from table_name with uids. uids is a tuple.
"""
def delete_batch(self, table_name , uids, uid_column_name='id'):
query = "DELETE from {} WHERE {} in {}".format(table_name, uid_column_name, str(uids))
self.__execute_query(query)
"""
Returns the dict a row by uid.
"""
def get_row(self, table_name, uid, uid_column_name='id'):
query = "Select * from {} where {} = {}".format(table_name, uid_column_name, uid)
connection = self.connection_pool.getconn()
cursor = connection.cursor()
cursor.execute(query)
column_names = [desc[0] for desc in cursor.description]
values = cursor.fetchall()
result = {}
if len(values) > 0:
for x, y in itertools.izip(column_names, values[0]):
result[x] = y
self.connection_pool.putconn(connection)
return result
"""
Returns all distinct values of column_name from table_name.
"""
def get_all_values_for_attr(self, table_name, column_name):
query = "Select distinct {} from {}".format(column_name, table_name)
connection = self.connection_pool.getconn()
cursor = connection.cursor()
cursor.execute(query)
rows = cursor.fetchall()
uids = [row[0] for row in rows]
self.connection_pool.putconn(connection)
return uids
"""
Returns all rows from table_name satisfying where_clause. The number of returned rows are limited to
limit.
"""
def get_all_rows(self, table_name, where_clause='1=1', limit=20, order_by=None):
query = "Select * from {} where {} ".format(table_name, where_clause)
if order_by:
query = '{} order by {} desc'.format(query, order_by)
query = '{} limit {}'.format(query, limit)
connection = self.connection_pool.getconn()
cursor = connection.cursor()
cursor.execute(query)
column_names = [desc[0] for desc in cursor.description]
rows = cursor.fetchall()
result = []
for row in rows:
result_row = {}
for x, y in itertools.izip(column_names, row):
result_row[x] = str(y)
result.append(result_row)
self.connection_pool.putconn(connection)
return result
"""
Gets a new connection from the pool and returns the connection object.
"""
def get_connection(self):
return self.connection_pool.getconn()
"""
Releases the connection back to pool.
"""
def release_connection(self, connection):
self.connection_pool.putconn(connection)
# ON ENTER OF WITH STATEMENT
# --------------------------
# perform any setup/initialization here
dbconnection = pool.getconn()
yield dbconnection.cursor()
# -------------------------
# ON EXIT OF WITH STATEMENT
# -------------------------
# this gets called on exit of the "with" block, no matter what
#app.logger.debug("Returning connection to the pool")
pool.putconn(dbconnection)
#db_connection = psycopg2.connect(args.database_connection_string)
pool = ThreadedConnectionPool(minconn=1, maxconn=10, dsn=db.dsn)
db_connection = pool.getconn()
def copy_into_db(copy_data=None, columns=None):
with database_cursor(pool) as cursor:
cursor.copy_from(file=os.fdopen(copy_data), table='sdssphoto.field', columns=columns, null="NULL")
cursor.close()
cursor.connection.commit()
def ndarray2pgcopy(a=None, convert_null=False):
''' Convert an ndarray object to a string appropriate for use with the COPY command.
Format is (including quotes!):
{1,2,3,4}
{{1,2,3,4}, {5,6,7,8}}
convert_null converts -9999 values to NULL
'''
def __init__(self, db_name="test_db", db_pass="", host="127.0.0.1" , port="5432"):
self.connection_pool = ThreadedConnectionPool(10, 50, database=db_name, user="******", \
password=db_pass, host=host, port=port)
self.logger = get_logger()