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)
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 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
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 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)
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()
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)
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
]
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: # noqa: E722
# We're fine with the bare except cos we raise in any case.
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)
# Override configs with env variables
if os.environ.get('WIKILABELS_DATABASE'):
params['dbname'] = os.environ.get('WIKILABELS_DATABASE')
if os.environ.get('WIKILABELS_DATABASE_HOST'):
params['host'] = os.environ.get('WIKILABELS_DATABASE_HOST')
if os.environ.get('WIKILABELS_DATABASE_USER'):
params['user'] = os.environ.get('WIKILABELS_DATABASE_USER')
if os.environ.get('WIKILABELS_DATABASE_PASSWORD'):
params['password'] = os.environ.get('WIKILABELS_DATABASE_PASSWORD')
return cls(**params)
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 remove_user(email):
try:
params = config()
tcp = ThreadedConnectionPool(1, 10, **params)
conn = tcp.getconn()
cur = conn.cursor()
cur.execute(remove_user_by_email_sql, (email, ))
conn.commit()
except Exception as err:
pass
class BuildingsApi(Flask):
def __init__(self, config_path):
super().__init__(__name__)
# Register the routes and their handlers
self.route('/')(self.index)
self.route('/buildings')(self.buildings)
self.route('/buildings/<building_id>')(self.building_get)
# Read the configuration file
self.app_config = load_config(config_path)
# Connect to the database and create our connection pool
try:
self.db_pool = ThreadedConnectionPool(1, 20,
**self.app_config['db'])
except:
# This a bad way of handling connection errors, but it's good enough for this project
sys.exit("Failed to connect to database.")
def run(self):
super().run(**self.app_config['web'])
@contextmanager
def _get_db_connection(self):
try:
connection = self.db_pool.getconn()
yield connection
finally:
self.db_pool.putconn(connection)
@contextmanager
def _get_db_cursor(self, commit=False):
with self._get_db_connection() as connection:
cursor = connection.cursor(cursor_factory=RealDictCursor)
try:
yield cursor
if commit:
connection.commit()
finally:
cursor.close()
def index(self):
content = "{'message':'Hello world!'}"
return Response(content, mimetype='text/json')
def buildings(self):
with self._get_db_cursor() as cursor:
buildings = db_buildings_all(cursor)
return Response(str(buildings), mimetype='text/json')
return Response("Database error", status_code=500)
def building_get(self, building_id):
with self._get_db_cursor() as cursor:
building = db_buildings_get(cursor, building_id)
return Response(str(building), mimetype='text/json')
return Response("Database error", status_code=500)
def create(name: str, description: str, pool: ThreadedConnectionPool):
"""
Create a new tag and store the information in db.
"""
conn = pool.getconn()
with conn.cursor() as curs:
curs.execute("INSERT INTO tags(name, description) VALUES(%s, %s)",
(name, description))
conn.commit()
pool.putconn(conn)
return Tag(name=name, description=description, pool=pool)
def from_name(name: str, pool: ThreadedConnectionPool):
"""
Retrieve information of tag ``name`` and return a ``Tag`` object.
"""
connection = pool.getconn()
with connection.cursor() as curs:
curs.execute("SELECT description FROM tags WHERE name=%s",
(name, ))
(description, ) = curs.fetchone()
pool.putconn(connection)
return Tag(name=name, description=description, pool=pool)
def connect():
"Connect to a database. This context manager yields the connection, and closes it when exited."
global _pool
if _pool is None:
_log.info('connecting to %s', db_url())
_pool = ThreadedConnectionPool(1, 5, db_url())
conn = _pool.getconn()
try:
yield conn
finally:
_pool.putconn(conn)
class ConnectionPool: # no test coverage
"""https://gist.github.com/jeorgen/4eea9b9211bafeb18ada"""
is_setup = False
def setup(self):
self.last_seen_process_id = os.getpid()
self._init()
self.is_setup = True
def _init(self):
self._pool = ThreadedConnectionPool(1,
10,
database=config.db_name,
**credentials)
def _getconn(self) -> connection:
current_pid = os.getpid()
if not (current_pid == self.last_seen_process_id):
self._init()
log.debug(
f"New id is {current_pid}, old id was {self.last_seen_process_id}"
)
self.last_seen_process_id = current_pid
conn = self._pool.getconn()
return conn
def _putconn(self, conn: connection):
return self._pool.putconn(conn)
def closeall(self):
self._pool.closeall()
@contextmanager
def get_connection(self) -> t.Generator[connection, None, None]:
try:
conn = self._getconn()
yield conn
finally:
self._putconn(conn)
@contextmanager
def get_cursor(self,
commit=False) -> t.Generator[LoggingCursor, None, None]:
with self.get_connection() as conn:
cursor = conn.cursor(cursor_factory=LoggingCursor)
try:
yield cursor
if commit:
conn.commit()
finally:
cursor.close()
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)
class PostgresConnectionPool(AbstractDatabaseConnectionPool):
def __init__(self, **kwargs):
try:
self.pconnect = ThreadedConnectionPool(1, poolsize, dsn=pdsn)
except:
global _pgpool
_pgpool = None
raise ConnectorError('Database Connection Failed')
maxsize = kwargs.pop('maxsize', None)
self.kwargs = kwargs
AbstractDatabaseConnectionPool.__init__(self, maxsize)
def create_connection(self):
self.conn = self.pconnect.getconn()
self.conn.autocommit = True
return self.conn
def from_qid(qid: int,
pool: ThreadedConnectionPool):
"""
Construct ``Question`` from ``qid``.
"""
conn = pool.getconn()
with conn.cursor() as curs:
curs.execute("SELECT author, title, description, upvotes, downvotes "
"FROM questions "
"WHERE qid=%s", (qid, ))
author, title, body, upvotes, downvotes = curs.fetchone()
curs.execute("SELECT tag FROM question_tags WHERE qid=%s", (qid, ))
tags = [res[0] for res in curs]
pool.putconn(conn)
return Question(author=author, title=title, body=body, upvotes=upvotes,
downvotes=downvotes, qid=qid, pool=pool, tags=tags)
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)
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)
class PostgresDataContext(object):
def __init__(self):
self._pool = ThreadedConnectionPool(1, 12, database="forums", user="******",
password='******', host='localhost')
def _get_connection(self):
return self._pool.getconn()
def put_connection(self, conn):
self._pool.putconn(conn)
def create_connection(self):
conn = self._get_connection()
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
cursor = conn.cursor(cursor_factory=RealDictCursor)
return conn, cursor
class psqlServer():
def __init__(self):
self.user = user
self.dbname = dbname
self.port = port
self.pswd = pswd
self.pool = ThreadedConnectionPool(1,
20,
host=host,
port=port,
database=dbname,
user=user,
password=pswd)
@contextmanager
def get_conn(self):
try:
conn = self.pool.getconn()
yield conn
finally:
self.pool.putconn(conn)
@contextmanager
def get_curs(self, commit=False):
with self.get_conn() as conn:
curs = conn.cursor(cursor_factory=RealDictCursor)
try:
yield curs
finally:
curs.close()
def execute(self, statement, *args):
with self.get_curs() as curs:
curs.execute(statement, args)
return curs.fetchall()
def execute_insert(self, statement, *args):
with self.get_conn() as conn:
curs = conn.cursor()
curs.execute(statement, args)
conn.commit()
curs.close()
if curs.description is not None:
return curs.fetchall()
def save_partition(iter):
pool = ThreadedConnectionPool(1,
5,
database='postgres',
user='******',
password='******',
host='10.0.0.6',
port='5432')
conn = pool.getconn()
cur = conn.cursor()
for record in iter:
qry = "INSERT INTO sentiment (timestamp, sentiment, count) VALUES ('%s','%s',%s);" % (
time, record[0], record[1])
cur.execute(qry)
conn.commit()
cur.close()
pool.putconn(conn)
class DB(object):
def __init__(self, queries, minconn, maxconn, *args, **kwargs):
self.queries = queries
for method in self.queries.keys():
setattr(DB, method, api_method(method,
lambda self, **kwargs: None))
logger.debug("Добавлен метод API базы данных: {}".format(method))
self.pool = ThreadedConnectionPool(minconn, maxconn, *args, **kwargs)
logger.debug('Создан пул соеднинений с базой данных размером {}'\
'и минимальным количеством соединений {}'.format(maxconn, minconn))
@contextmanager
def _connect(self):
logger.debug('Резервируется соединение с базой данных')
con = self.pool.getconn()
try:
yield con.cursor()
con.commit()
finally:
logger.debug('Освобождается соединение с базой данных')
self.pool.putconn(con)
def execute_api_method(self, method_name, **kwargs):
if method_name not in self.queries:
logger.error("DBAPI: Нет такого метода '{}'".format(method_name))
return
with self._connect() as cursor:
logger.info("DBAPI: Вызов '{}'".format(method_name))
params = self.queries[method_name].format(**kwargs)
if 'get_data' not in kwargs:
query = "select {0}({1});".format(method_name, params)
else:
query = "select * from {0}({1});".format(method_name, params)
logger.debug('ЗАПРОС: {}'.format(query))
cursor.execute(query)
if 'get_data' not in kwargs:
result = Result(cursor.fetchone())
logger.log(LEVELS[result.type], 'DBAPI: {}'.format(result))
return result.value
return cursor.fetchall()
class ParallelTerrainInterface(BaseInterface):
def __init__(self, DSN, lidar_table, processes=2):
super(ParallelTerrainInterface, self).__init__(DSN, lidar_table)
self.processes = processes
self.querier = []
# Connection to PSQL
self.tcp = ThreadedConnectionPool(1, 100, DSN)
# MT Queryier
self.workers_query_order_q = mp.Queue()
self.workers_query_result_q = mp.Queue()
self.conns = [self.tcp.getconn() for i in range(processes)]
for i in range(self.processes):
t = mp.Process(target=self._query_worker, args=[self.conns[i]])
self.querier.append(t)
t.daemon = True
t.start()
def _query_worker(self, conn):
while (True):
order = self.workers_query_order_q.get(block=True)
link = self._profile_osm(order, conn)
self.workers_query_result_q.put(link)
def get_link_parallel(self, src, dst_list):
"""Calculate the path loss between two lists of building
"""
links = []
params = [{
'src': src,
'dst': dst_list[i],
'srid': self.srid,
'lidar_table': self.lidar_table,
'buff': self.buff,
'latlong': False
} for i in range(len(dst_list))]
# add orders in the queue
for order in params:
self.workers_query_order_q.put(order)
# wait for all the orders to come back
while len(links) < len(dst_list):
links.append(self.workers_query_result_q.get(block=True))
return links
class PoolWrapper:
"""Exists to provide an acquire method for easy usage.
pool = PoolWrapper(...)
with pool.acquire() as conneciton:
connection.execute(...)
"""
def __init__(self, max_pool_size: int, *, dsn):
self._pool = ThreadedConnectionPool(
1, max_pool_size, dsn=dsn, cursor_factory=RealDictCursor,
)
@contextmanager
def acquire(self):
try:
connection = self._pool.getconn()
yield connection
finally:
self._pool.putconn(connection)
def make_endpoints(dbtype, from_list):
global PGPOOL, CPU_COUNT
endpoints = []
for server, db, auth in from_list:
user, password = "", ""
if auth:
if len(auth) == 2:
user, password = auth
else:
user = auth[0]
if isinstance(server, str):
try:
host, port = server.split(":", 1)
except ValueError:
host, port = server, None
else:
host, port = server
kargs = {}
if port: kargs["port"] = port
if password: kargs["password"] = password
if dbtype == DB_SQLITE3:
conn = sqlite3.connect(host)
elif dbtype == DB_PGSQL:
if user: kargs["user"] = user
if PGPOOL is None:
PGPOOL = ThreadedConnectionPool(0,
CPU_COUNT * 3,
host=host,
database=db,
**kargs)
conn = PGPOOL.getconn()
elif dbtype == DB_REDIS:
conn = redis.Redis(host=host, port=port, db=db)
elif dbtype == DB_MONGODB:
if user: kargs["username"] = user
conn = pymongo.MongoClient(host=host, **kargs)
endpoints.append(conn)
return endpoints
class SAEService(TornadoService):
name = "sae"
api = [SAEAPI]
# Some activator pools may be unavailable
# while SAE remains healthy
require_dcs_health = False
use_mongo = True
def __init__(self):
super().__init__()
self.pool_cache = {}
self.activators = {}
self.pg_pool = None
self.pg_pool_ready = threading.Event()
def load_pools(self):
self.logger.info("Loading pools")
for p in Pool.objects.all():
self.pool_cache[str(p.id)] = p.name
async def on_activate(self):
self.load_pools()
self.pg_pool = ThreadedConnectionPool(1, config.sae.db_threads, **config.pg_connection_args)
self.pg_pool_ready.set()
def get_pool_name(self, pool_id):
"""
Returns pool name by id
"""
return self.pool_cache.get(str(pool_id))
@contextlib.contextmanager
def get_pg_connect(self):
self.pg_pool_ready.wait()
connect = self.pg_pool.getconn()
if not connect.autocommit:
connect.autocommit = True
try:
yield connect
finally:
self.pg_pool.putconn(connect)
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()
print("New id is %s, old id was %s" %
(current_pid, self.last_seen_process_id))
self.last_seen_process_id = current_pid
return self._pool.getconn()
def putconn(self, conn):
return self._pool.putconn(conn)
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)
def cur(sql_stmt, col_names=True):
"""
Database postgres cursor
:param sql_stmt: - sql statement string to execute
:returns data: - Dictionary response from the database.
"""
# minimum and maximum threads to allow.
min = 5
max = 50
HOST = '0.0.0.0'
SCHEMA = 'testdb'
PORT = '5444'
connection_pool = ThreadedConnectionPool(
min, max,
host=HOST,
database=SCHEMA,
port=PORT,
user='******')
# RealDictCursor plugin makes the cursor return a python dictionary.
conn = connection_pool.getconn()
cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
if 'DROP' in sql_stmt.upper():
return "Permission denied. DROP commands are not supported."
cursor.execute(sql_stmt)
colnames = [desc[0] for desc in cursor.description]
if 'INSERT' in sql_stmt.upper() or "UPDATE" in sql_stmt.upper():
conn.commit()
cursor.close()
return
data = cursor.fetchall()
cursor.close()
if col_names:
return data, colnames
return data
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)
def from_qid_author(qid: int,
author: str,
pool: ThreadedConnectionPool = None):
"""
Retrieve data from database and construct ``Answer`` using it.
"""
conn = pool.getconn()
with conn.cursor() as curs:
curs.execute(
"SELECT body, upvotes, downvotes "
"FROM answers "
"WHERE author=%(username)s AND qid=%(qid)s", {
'username': author,
'qid': qid
})
body, upvotes, downvotes = curs.fetchone()
pool.putconn(conn)
return Answer(author=author,
body=body,
upvotes=upvotes,
downvotes=downvotes,
qid=qid,
pool=pool)
# --------------------------
# 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
'''
if len(value.shape) == 1: # 1-d array
class PgPool:
def __init__(self):
logger.debug('initializing postgres threaded pool')
self.host, self.port = None, None
self.database, self.pool = None, None
self.user, self.passwd = None, None
self.pool = None
logger.debug('Server Addr: {host}:{port}; Database: {db}; User: {user}; Password: {passwd}'.format(
host=self.host, port=self.port,
db=self.database, user=self.user, passwd=self.passwd
))
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)
def get_conn(self):
"""
Get a connection from pool and return connection and cursor
:return: conn, cursor
"""
logger.debug('getting connection from pool')
try:
conn = self.pool.getconn()
cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
return conn, cursor
except Exception as e:
logger.exception(e.message)
return None, None
@staticmethod
def execute_query(cursor, query, params):
"""
Execute a query on database
:param cursor: cursor object
:param query: database query
:type query: str
:param params: query parameters
:type params: tuple
:return: query results or bool
"""
logger.info('executing query')
logger.debug('Cursor: {cursor}, Query: {query}'.format(
cursor=cursor, query=query))
try:
if query.split()[0].lower() == 'select':
cursor.execute(query, params)
return cursor.fetchall()
else:
return cursor.execute(query, params)
except Exception as e:
logger.exception(e.message)
return False
# commit changes to db permanently
@staticmethod
def commit_changes(conn):
"""
Commit changes to the databse permanently
:param conn: connection object
:return: bool
"""
logger.debug('commiting changes to database')
try:
return conn.commit()
except Exception as e:
logger.exception(e.message)
return False
def put_conn(self, conn):
"""
Put connection back to the pool
:param conn: connection object
:return: bool
"""
logger.debug('putting connection {conn} back to pool'.format(conn=conn))
try:
return self.pool.putconn(conn)
except Exception as e:
logger.exception(e.message)
return False
def close_pool(self):
"""
Closes connection pool
:return: bool
"""
logger.debug('closing connections pool')
try:
return self.pool.closeall()
except Exception as e:
logger.exception(e.message)
return False
class Database:
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 initialize(self):
# Initialize Database, Recreate Tables
try:
self.curs.execute("""CREATE TABLE Users (
user_id text, balance bigint)""")
except:
self.conn.rollback()
self.curs.execute("""DROP TABLE Users""")
self.curs.execute("""CREATE TABLE Users (
user_id text, balance bigint)""")
self.conn.commit()
try:
self.curs.execute("""CREATE TABLE Stock (
stock_id text, user_id text, amount bigint)""")
except:
self.conn.rollback()
self.curs.execute("""DROP TABLE Stock""")
self.curs.execute("""CREATE TABLE Stock (
stock_id text, user_id text, amount bigint)""")
self.conn.commit()
try:
self.curs.execute("""CREATE TABLE PendingTrans (
type text, user_id text, stock_id text, amount bigint, timestamp bigint)""")
except:
self.conn.rollback()
self.curs.execute("""DROP TABLE PendingTrans""")
self.curs.execute("""CREATE TABLE PendingTrans (
type text, user_id text, stock_id text, amount bigint, timestamp bigint)""")
self.conn.commit()
try:
self.curs.execute("""CREATE TABLE Trigger (
type text, user_id text, stock_id text, amount bigint, trigger bigint)""")
except:
self.conn.rollback()
self.curs.execute("""DROP TABLE Trigger""")
self.curs.execute("""CREATE TABLE Trigger (
type text, user_id text, stock_id text, amount bigint, trigger bigint)""")
self.conn.commit()
print "DB Initialized"
# Return a Database Connection from the pool
def get_connection(self):
connection = self.pool.getconn()
cursor = connection.cursor()
return connection, cursor
def close_connection(self, connection):
self.pool.putconn(connection)
# call like: select_record("Users", "id,balance", "id='jim' AND balance=200")
def select_record(self, values, table, constraints):
connection, cursor = self.get_connection()
try:
command = """SELECT %s FROM %s WHERE %s""" % (values, table, constraints)
cursor.execute(command)
connection.commit()
except Exception as e:
print 'PG Select error - ' + str(e)
result = cursor.fetchall()
self.close_connection(connection)
# Format to always return a tuple of the single record, with each value.
if len(result) > 1:
print 'PG Select returned more than one value.'
return (None,None)
elif len(result) == 0:
return (None,None)
else:
return result[0]
def filter_records(self, values, table, constraints):
connection, cursor = self.get_connection()
try:
command = """SELECT %s FROM %s WHERE %s""" % (values, table, constraints)
cursor.execute(command)
connection.commit()
except Exception as e:
print 'PG Select error - ' + str(e)
result = cursor.fetchall()
self.close_connection(connection)
# Return array of tuples
return result
def insert_record(self, table, columns, values):
connection, cursor = self.get_connection()
try:
command = """INSERT INTO %s (%s) VALUES (%s)""" % (table, columns, values)
cursor.execute(command)
connection.commit()
except Exception as e:
print 'PG Insert error - ' + str(e)
self.close_connection(connection)
def update_record(self, table, values, constraints):
connection, cursor = self.get_connection()
try:
command = """UPDATE %s SET %s WHERE %s""" % (table, values, constraints)
cursor.execute(command)
connection.commit()
except Exception as e:
print 'PG Update error %s \n table=%s values=%s constraints=%s command=%s' % (str(e), table, values, constraints, command)
self.close_connection(connection)
def delete_record(self, table, constraints):
connection, cursor = self.get_connection()
try:
command = """DELETE FROM %s WHERE %s""" % (table, constraints)
cursor.execute(command)
connection.commit()
except Exception as e:
print 'PG Delete error - ' + str(e)
self.close_connection(connection)
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 Database(DatabaseInterface):
_databases = {}
_connpool = None
_list_cache = None
_list_cache_timestamp = None
_version_cache = {}
def __new__(cls, database_name="template1"):
if database_name in cls._databases:
return cls._databases[database_name]
return DatabaseInterface.__new__(cls, database_name=database_name)
def __init__(self, database_name="template1"):
super(Database, self).__init__(database_name=database_name)
self._databases.setdefault(database_name, 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)
self._connpool = ThreadedConnectionPool(minconn, maxconn, dsn)
return self
def cursor(self, autocommit=False, readonly=False):
if self._connpool is None:
self.connect()
conn = self._connpool.getconn()
if autocommit:
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
else:
conn.set_isolation_level(ISOLATION_LEVEL_REPEATABLE_READ)
cursor = Cursor(self._connpool, conn, self)
if readonly:
cursor.execute("SET TRANSACTION READ ONLY")
return cursor
def close(self):
if self._connpool is None:
return
self._connpool.closeall()
self._connpool = None
def create(self, cursor, database_name):
cursor.execute('CREATE DATABASE "' + database_name + '" ' "TEMPLATE template0 ENCODING 'unicode'")
Database._list_cache = None
def drop(self, cursor, database_name):
cursor.execute('DROP DATABASE "' + database_name + '"')
Database._list_cache = None
def get_version(self, cursor):
if self.database_name not in self._version_cache:
cursor.execute("SELECT version()")
version, = cursor.fetchone()
self._version_cache[self.database_name] = tuple(map(int, RE_VERSION.search(version).groups()))
return self._version_cache[self.database_name]
@staticmethod
def dump(database_name):
from trytond.tools import exec_pg_command_pipe
cmd = ["pg_dump", "--format=c", "--no-owner"]
if CONFIG["db_user"]:
cmd.append("--username="******"db_user"])
if CONFIG["db_host"]:
cmd.append("--host=" + CONFIG["db_host"])
if CONFIG["db_port"]:
cmd.append("--port=" + CONFIG["db_port"])
cmd.append(database_name)
pipe = exec_pg_command_pipe(*tuple(cmd))
pipe.stdin.close()
data = pipe.stdout.read()
res = pipe.wait()
if res:
raise Exception("Couldn't dump database!")
return data
@staticmethod
def restore(database_name, data):
from trytond.tools import exec_pg_command_pipe
database = Database().connect()
cursor = database.cursor(autocommit=True)
database.create(cursor, database_name)
cursor.commit()
cursor.close()
cmd = ["pg_restore", "--no-owner"]
if CONFIG["db_user"]:
cmd.append("--username="******"db_user"])
if CONFIG["db_host"]:
cmd.append("--host=" + CONFIG["db_host"])
if CONFIG["db_port"]:
cmd.append("--port=" + CONFIG["db_port"])
cmd.append("--dbname=" + database_name)
args2 = tuple(cmd)
if os.name == "nt":
tmpfile = (os.environ["TMP"] or "C:\\") + os.tmpnam()
with open(tmpfile, "wb") as fp:
fp.write(data)
args2 = list(args2)
args2.append(" " + tmpfile)
args2 = tuple(args2)
pipe = exec_pg_command_pipe(*args2)
if not os.name == "nt":
pipe.stdin.write(data)
pipe.stdin.close()
res = pipe.wait()
if res:
raise Exception("Couldn't restore database")
database = Database(database_name).connect()
cursor = database.cursor()
if not cursor.test():
cursor.close()
database.close()
raise Exception("Couldn't restore database!")
cursor.close()
database.close()
Database._list_cache = None
return True
@staticmethod
def list(cursor):
now = time.time()
timeout = int(CONFIG["session_timeout"])
res = Database._list_cache
if res and abs(Database._list_cache_timestamp - now) < timeout:
return res
db_user = CONFIG["db_user"]
if not db_user and os.name == "posix":
db_user = pwd.getpwuid(os.getuid())[0]
if not db_user:
cursor.execute(
"SELECT usename "
"FROM pg_user "
"WHERE usesysid = ("
"SELECT datdba "
"FROM pg_database "
"WHERE datname = %s)",
(CONFIG["db_name"],),
)
res = cursor.fetchone()
db_user = res and res[0]
if db_user:
cursor.execute(
"SELECT datname "
"FROM pg_database "
"WHERE datdba = ("
"SELECT usesysid "
"FROM pg_user "
"WHERE usename=%s) "
"AND datname not in "
"('template0', 'template1', 'postgres') "
"ORDER BY datname",
(db_user,),
)
else:
cursor.execute(
"SELECT datname "
"FROM pg_database "
"WHERE datname not in "
"('template0', 'template1','postgres') "
"ORDER BY datname"
)
res = []
for (db_name,) in cursor.fetchall():
db_name = db_name.encode("utf-8")
try:
database = Database(db_name).connect()
except Exception:
continue
cursor2 = database.cursor()
if cursor2.test():
res.append(db_name)
cursor2.close(close=True)
else:
cursor2.close(close=True)
database.close()
Database._list_cache = res
Database._list_cache_timestamp = now
return res
@staticmethod
def init(cursor):
from trytond.modules import get_module_info
sql_file = os.path.join(os.path.dirname(__file__), "init.sql")
with open(sql_file) as fp:
for line in fp.read().split(";"):
if (len(line) > 0) and (not line.isspace()):
cursor.execute(line)
for module in ("ir", "res", "webdav"):
state = "uninstalled"
if module in ("ir", "res"):
state = "to install"
info = get_module_info(module)
cursor.execute("SELECT NEXTVAL('ir_module_module_id_seq')")
module_id = cursor.fetchone()[0]
cursor.execute(
"INSERT INTO ir_module_module "
"(id, create_uid, create_date, name, state) "
"VALUES (%s, %s, now(), %s, %s)",
(module_id, 0, module, state),
)
for dependency in info.get("depends", []):
cursor.execute(
"INSERT INTO ir_module_module_dependency "
"(create_uid, create_date, module, name) "
"VALUES (%s, now(), %s, %s)",
(0, module_id, dependency),
)
class Database:
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']))
@contextmanager
def getcursor(self):
conn = self.pool.getconn()
try:
yield conn.cursor()
conn.commit()
except Exception as e:
conn.rollback()
raise e.with_traceback(sys.exc_info()[2])
finally:
self.pool.putconn(conn)
def insert(self, table, data_list, id_col='id'):
"""
TODO: rename `id_col` -> `return_col`
Create a bulk insert statement which is much faster (~2x in tests with 10k & 100k rows and 4 cols)
for inserting data then executemany()
TODO: Is there a limit of length the query can be? If so handle it.
"""
# Make sure that `data_list` is a list
if not isinstance(data_list, list):
data_list = [data_list]
if len(data_list) == 0:
# No need to continue
return []
# Data in the list must be dicts (just check the first one)
if not isinstance(data_list[0], dict):
logger.critical("Data must be a list of dicts")
# Do not return here, let the exception handle the error that will be thrown when the query runs
try:
with self.getcursor() as cur:
query = "INSERT INTO {table} ({fields}) VALUES {values} RETURNING {id_col}"\
.format(table=table,
fields='{0}{1}{0}'.format('"', '", "'.join(data_list[0].keys())),
values=','.join(['%s'] * len(data_list)),
id_col=id_col
)
query = cur.mogrify(query, [tuple(v.values()) for v in data_list])
cur.execute(query)
return [t[0] for t in cur.fetchall()]
except Exception as e:
logger.exception("Error inserting data")
logger.debug("Error inserting data: {data}".format(data=data_list))
raise e.with_traceback(sys.exc_info()[2])
def update(self, table, data_list, matched_field=None, return_col='id'):
"""
Create a bulk insert statement which is much faster (~2x in tests with 10k & 100k rows and 4 cols)
for inserting data then executemany()
TODO: Is there a limit of length the query can be? If so handle it.
"""
if matched_field is None:
# Assume the id field
logger.info("Matched field not defined, assuming the `id` field")
matched_field = 'id'
# Make sure that `data_list` is a list
if not isinstance(data_list, list):
data_list = [data_list]
if len(data_list) == 0:
# No need to continue
return []
# Data in the list must be dicts (just check the first one)
if not isinstance(data_list[0], dict):
logger.critical("Data must be a list of dicts")
# Do not return here, let the exception handle the error that will be thrown when the query runs
try:
with self.getcursor() as cur:
query_list = []
# TODO: change to return data from the database, not just what you passed in
return_list = []
for row in data_list:
if row.get(matched_field) is None:
logger.debug("Cannot update row. Missing field {field} in data {data}"
.format(field=matched_field, data=row))
logger.error("Cannot update row. Missing field {field} in data".format(field=matched_field))
continue
# Pull matched_value from data to be updated and remove that key
matched_value = row.get(matched_field)
del row[matched_field]
query = "UPDATE {table} SET {data} WHERE {matched_field}=%s RETURNING {return_col}"\
.format(table=table,
data=','.join("%s=%%s" % u for u in row.keys()),
matched_field=matched_field,
return_col=return_col
)
values = list(row.values())
values.append(matched_value)
query = cur.mogrify(query, values)
query_list.append(query)
return_list.append(matched_value)
finial_query = b';'.join(query_list)
cur.execute(finial_query)
return return_list
except Exception as e:
logger.exception("Error updating data")
logger.debug("Error updating data: {data}".format(data=data_list))
raise e.with_traceback(sys.exc_info()[2])
class Database(rigor.database.Database):
""" Container for a database connection pool """
def __init__(self, database):
super(Database, self).__init__(database)
register_type(psycopg2.extensions.UNICODE)
register_uuid()
dsn = Database.build_dsn(database)
self._pool = ThreadedConnectionPool(config.get('database', 'min_database_connections'), config.get('database', 'max_database_connections'), dsn)
@staticmethod
def build_dsn(database):
""" Builds the database connection string from config values """
dsn = "dbname='{0}' host='{1}'".format(database, config.get('database', 'host'))
try:
ssl = config.getboolean('database', 'ssl')
if ssl:
dsn += " sslmode='require'"
except ConfigParser.Error:
pass
try:
username = config.get('database', 'username')
dsn += " user='******'".format(username)
except ConfigParser.Error:
pass
try:
password = config.get('database', 'password')
dsn += " password='******'".format(password)
except ConfigParser.Error:
pass
return dsn
@staticmethod
@template
def create(name):
""" Creates a new database with the given name """
return "CREATE DATABASE {0};".format(name)
@staticmethod
@template
def drop(name):
""" Drops the database with the given name """
return "DROP DATABASE {0};".format(name)
@staticmethod
@template
def clone(source, destination):
"""
Copies the source database to a new destination database. This may fail if
the source database is in active use.
"""
return "CREATE DATABASE {0} WITH TEMPLATE {1};".format(destination, source)
@contextmanager
def get_cursor(self, commit=True):
""" Gets a cursor from a connection in the pool """
connection = self._pool.getconn()
cursor = connection.cursor(cursor_factory=RigorCursor)
try:
yield cursor
except psycopg2.IntegrityError as error:
exc_info = sys.exc_info()
self.rollback(cursor)
raise rigor.database.IntegrityError, exc_info[1], exc_info[2]
except psycopg2.DatabaseError as error:
exc_info = sys.exc_info()
self.rollback(cursor)
raise rigor.database.DatabaseError, exc_info[1], exc_info[2]
except:
exc_info = sys.exc_info()
self.rollback(cursor)
raise exc_info[0], exc_info[1], exc_info[2]
else:
if commit:
self.commit(cursor)
else:
self.rollback(cursor)
def _close_cursor(self, cursor):
""" Closes a cursor and releases the connection to the pool """
cursor.close()
self._pool.putconn(cursor.connection)
def commit(self, cursor):
""" Commits the transaction, then closes the cursor """
cursor.connection.commit()
self._close_cursor(cursor)
def rollback(self, cursor):
""" Rolls back the transaction, then closes the cursor """
cursor.connection.rollback()
self._close_cursor(cursor)
def __del__(self):
self._pool.closeall()
class FlaskWithPool(Flask):
""" Add some PGPool operations so we can reset from a request context when the DB is restarted on us """
def __init__(self, name):
Flask.__init__(self, name)
self.resetlock = threading.Lock()
self.pool = None
def db_prepare(self):
""" Check if we have the series in the URL, set the schema path if available, return an error message otherwise """
if hasattr(g, 'db'):
raise EnvironmentError('Database has already been prepared. Preparing again is an error.')
g.db = self._get_from_pool()
if hasattr(g, 'series') and g.series:
# Set up the schema path if we have a series
g.seriestype = Series.type(g.series)
if g.seriestype == Series.INVALID:
abort(404, "%s is not a valid series" % g.series)
with g.db.cursor() as cur:
cur.execute("SET search_path=%s,'public'; commit; begin", (g.series,))
else:
g.seriestype = Series.UNKNOWN
def db_return(self):
""" Return a connection to the pool and clear the attribute """
if hasattr(g, 'db'):
self.pool.putconn(g.db)
del g.db # Removes 'db' from g dictionary
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 _get_from_pool(self):
""" Get a database connection from the pool and make sure its connected, attempt reset once if needed """
try:
ret = self.pool.getconn()
with ret.cursor() as cur:
cur.execute("select 1")
except (DatabaseError, OperationalError) as e:
log.warning("Possible database restart. Reseting pool and trying again!")
try:
self._reset_pool()
ret = self.pool.getconn()
with ret.cursor() as cur:
cur.execute("select 1")
except (DatabaseError, OperationalError) as e:
raise EnvironmentError("Errors with postgresql connection pool. Bailing")
#log.debug("{} setup: {} connections used".format(threading.current_thread(), len(self.pool._used)))
return ret
class Database(DatabaseInterface):
_databases = {}
_connpool = None
_list_cache = None
_list_cache_timestamp = None
_version_cache = {}
flavor = Flavor(ilike=True)
def __new__(cls, database_name='template1'):
if database_name in cls._databases:
return cls._databases[database_name]
return DatabaseInterface.__new__(cls, database_name=database_name)
def __init__(self, database_name='template1'):
super(Database, self).__init__(database_name=database_name)
self._databases.setdefault(database_name, 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 cursor(self, autocommit=False, readonly=False):
if self._connpool is None:
self.connect()
conn = self._connpool.getconn()
if autocommit:
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
else:
conn.set_isolation_level(ISOLATION_LEVEL_REPEATABLE_READ)
cursor = Cursor(self._connpool, conn, self)
if readonly:
cursor.execute('SET TRANSACTION READ ONLY')
return cursor
def close(self):
if self._connpool is None:
return
self._connpool.closeall()
self._connpool = None
@classmethod
def create(cls, cursor, database_name):
cursor.execute('CREATE DATABASE "' + database_name + '" '
'TEMPLATE template0 ENCODING \'unicode\'')
cls._list_cache = None
@classmethod
def drop(cls, cursor, database_name):
cursor.execute('DROP DATABASE "' + database_name + '"')
cls._list_cache = None
def get_version(self, cursor):
if self.database_name not in self._version_cache:
cursor.execute('SELECT version()')
version, = cursor.fetchone()
self._version_cache[self.database_name] = tuple(map(int,
RE_VERSION.search(version).groups()))
return self._version_cache[self.database_name]
@staticmethod
def dump(database_name):
from trytond.tools import exec_command_pipe
cmd = ['pg_dump', '--format=c', '--no-owner']
env = {}
uri = parse_uri(config.get('database', 'uri'))
if uri.username:
cmd.append('--username='******'--host=' + uri.hostname)
if uri.port:
cmd.append('--port=' + str(uri.port))
if uri.password:
# if db_password is set in configuration we should pass
# an environment variable PGPASSWORD to our subprocess
# see libpg documentation
env['PGPASSWORD'] = uri.password
cmd.append(database_name)
pipe = exec_command_pipe(*tuple(cmd), env=env)
pipe.stdin.close()
data = pipe.stdout.read()
res = pipe.wait()
if res:
raise Exception('Couldn\'t dump database!')
return data
@staticmethod
def restore(database_name, data):
from trytond.tools import exec_command_pipe
database = Database().connect()
cursor = database.cursor(autocommit=True)
database.create(cursor, database_name)
cursor.commit()
cursor.close()
database.close()
cmd = ['pg_restore', '--no-owner']
env = {}
uri = parse_uri(config.get('database', 'uri'))
if uri.username:
cmd.append('--username='******'--host=' + uri.hostname)
if uri.port:
cmd.append('--port=' + str(uri.port))
if uri.password:
env['PGPASSWORD'] = uri.password
cmd.append('--dbname=' + database_name)
args2 = tuple(cmd)
if os.name == "nt":
tmpfile = (os.environ['TMP'] or 'C:\\') + os.tmpnam()
with open(tmpfile, 'wb') as fp:
fp.write(data)
args2 = list(args2)
args2.append(' ' + tmpfile)
args2 = tuple(args2)
pipe = exec_command_pipe(*args2, env=env)
if not os.name == "nt":
pipe.stdin.write(data)
pipe.stdin.close()
res = pipe.wait()
if res:
raise Exception('Couldn\'t restore database')
database = Database(database_name).connect()
cursor = database.cursor()
if not cursor.test():
cursor.close()
database.close()
raise Exception('Couldn\'t restore database!')
cursor.close()
database.close()
Database._list_cache = None
return True
@staticmethod
def list(cursor):
now = time.time()
timeout = config.getint('session', 'timeout')
res = Database._list_cache
if res and abs(Database._list_cache_timestamp - now) < timeout:
return res
cursor.execute('SELECT datname FROM pg_database '
'WHERE datistemplate = false ORDER BY datname')
res = []
for db_name, in cursor.fetchall():
try:
database = Database(db_name).connect()
except Exception:
continue
cursor2 = database.cursor()
if cursor2.test():
res.append(db_name)
cursor2.close(close=True)
else:
cursor2.close(close=True)
database.close()
Database._list_cache = res
Database._list_cache_timestamp = now
return res
@staticmethod
def init(cursor):
from trytond.modules import get_module_info
sql_file = os.path.join(os.path.dirname(__file__), 'init.sql')
with open(sql_file) as fp:
for line in fp.read().split(';'):
if (len(line) > 0) and (not line.isspace()):
cursor.execute(line)
for module in ('ir', 'res'):
state = 'uninstalled'
if module in ('ir', 'res'):
state = 'to install'
info = get_module_info(module)
cursor.execute('SELECT NEXTVAL(\'ir_module_id_seq\')')
module_id = cursor.fetchone()[0]
cursor.execute('INSERT INTO ir_module '
'(id, create_uid, create_date, name, state) '
'VALUES (%s, %s, now(), %s, %s)',
(module_id, 0, module, state))
for dependency in info.get('depends', []):
cursor.execute('INSERT INTO ir_module_dependency '
'(create_uid, create_date, module, name) '
'VALUES (%s, now(), %s, %s)',
(0, module_id, dependency))
class DbPool(object):
"""DB class that makes connection transparently. Thread-safe - every
thread get its own database connection. Not meant to be used directly,
there is no reason to have more than one instance - global variable Db
- in this module."""
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 cursor(self, **kwargs):
"""Creates and returns cursor for current thread's connection.
Cursor is a "dict" cursor, so you can access the columns by
names (not just indices), e.g.:
cursor.execute("SELECT id, name FROM ... WHERE ...", sql_args)
row = cursor.fetchone()
id = row['id']
Server-side cursors (named cursors) should be closed explicitly.
@param kwargs: currently string parameter 'name' is supported.
Named cursors are for server-side cursors, which
are useful when fetching result of a large query via fetchmany()
method. See http://initd.org/psycopg/docs/usage.html#server-side-cursors
"""
return self.connection().cursor(cursor_factory=DictCursor, **kwargs)
def connection(self):
"""Return connection for this thread"""
return self.pool.getconn(id(threading.current_thread()))
def commit(self):
"""Commit all the commands in this transaction in this thread's
connection. If errors (e.g. duplicate key) arose, this will
cause transaction rollback.
"""
self.connection().commit()
def rollback(self):
"""Rollback last transaction on this thread's connection"""
self.connection().rollback()
def putconn(self):
"""Put back connection used by this thread. Necessary upon finishing of
spawned threads, otherwise new threads won't get connection if the pool
is depleted."""
conn = self.connection()
self.pool.putconn(conn, id(threading.current_thread()))
def close(self):
"""Close connection."""
self.connection().close()
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)
class Database(DatabaseInterface):
_databases = {}
_connpool = None
_list_cache = None
_list_cache_timestamp = None
_version_cache = {}
flavor = Flavor(ilike=True)
def __new__(cls, name='template1'):
if name in cls._databases:
return cls._databases[name]
return DatabaseInterface.__new__(cls, name=name)
def __init__(self, name='template1'):
super(Database, self).__init__(name=name)
self._databases.setdefault(name, self)
self._search_path = None
self._current_user = None
@classmethod
def dsn(cls, 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" % name
user = uri.username and "user=%s" % uri.username or ''
password = ("password=%s" % urllib.unquote_plus(uri.password)
if uri.password else '')
return '%s %s %s %s %s' % (host, port, name, user, password)
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 get_connection(self, autocommit=False, readonly=False):
if self._connpool is None:
self.connect()
conn = self._connpool.getconn()
if autocommit:
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
else:
conn.set_isolation_level(ISOLATION_LEVEL_REPEATABLE_READ)
if readonly:
cursor = conn.cursor()
cursor.execute('SET TRANSACTION READ ONLY')
conn.cursor_factory = PerfCursor
return conn
def put_connection(self, connection, close=False):
self._connpool.putconn(connection, close=close)
def close(self):
if self._connpool is None:
return
self._connpool.closeall()
self._connpool = None
@classmethod
def create(cls, connection, database_name):
cursor = connection.cursor()
cursor.execute('CREATE DATABASE "' + database_name + '" '
'TEMPLATE template0 ENCODING \'unicode\'')
connection.commit()
cls._list_cache = None
def drop(self, connection, database_name):
cursor = connection.cursor()
cursor.execute('DROP DATABASE "' + database_name + '"')
Database._list_cache = None
def get_version(self, connection):
if self.name not in self._version_cache:
cursor = connection.cursor()
cursor.execute('SELECT version()')
version, = cursor.fetchone()
self._version_cache[self.name] = tuple(map(int,
RE_VERSION.search(version).groups()))
return self._version_cache[self.name]
@staticmethod
def dump(database_name):
from trytond.tools import exec_command_pipe
cmd = ['pg_dump', '--format=c', '--no-owner']
env = {}
uri = parse_uri(config.get('database', 'uri'))
if uri.username:
cmd.append('--username='******'--host=' + uri.hostname)
if uri.port:
cmd.append('--port=' + str(uri.port))
if uri.password:
# if db_password is set in configuration we should pass
# an environment variable PGPASSWORD to our subprocess
# see libpg documentation
env['PGPASSWORD'] = uri.password
cmd.append(database_name)
pipe = exec_command_pipe(*tuple(cmd), env=env)
pipe.stdin.close()
data = pipe.stdout.read()
res = pipe.wait()
if res:
raise Exception('Couldn\'t dump database!')
return data
@staticmethod
def restore(database_name, data):
from trytond.tools import exec_command_pipe
database = Database().connect()
connection = database.get_connection(autocommit=True)
database.create(connection, database_name)
database.close()
cmd = ['pg_restore', '--no-owner']
env = {}
uri = parse_uri(config.get('database', 'uri'))
if uri.username:
cmd.append('--username='******'--host=' + uri.hostname)
if uri.port:
cmd.append('--port=' + str(uri.port))
if uri.password:
env['PGPASSWORD'] = uri.password
cmd.append('--dbname=' + database_name)
args2 = tuple(cmd)
if os.name == "nt":
tmpfile = (os.environ['TMP'] or 'C:\\') + os.tmpnam()
with open(tmpfile, 'wb') as fp:
fp.write(data)
args2 = list(args2)
args2.append(' ' + tmpfile)
args2 = tuple(args2)
pipe = exec_command_pipe(*args2, env=env)
if not os.name == "nt":
pipe.stdin.write(data)
pipe.stdin.close()
res = pipe.wait()
if res:
raise Exception('Couldn\'t restore database')
database = Database(database_name).connect()
cursor = database.get_connection().cursor()
if not database.test():
cursor.close()
database.close()
raise Exception('Couldn\'t restore database!')
cursor.close()
database.close()
Database._list_cache = None
return True
def list(self):
now = time.time()
timeout = config.getint('session', 'timeout')
res = Database._list_cache
if res and abs(Database._list_cache_timestamp - now) < timeout:
return res
connection = self.get_connection()
cursor = connection.cursor()
cursor.execute('SELECT datname FROM pg_database '
'WHERE datistemplate = false ORDER BY datname')
res = []
for db_name, in cursor:
try:
with connect(self.dsn(db_name)) as conn:
if self._test(conn):
res.append(db_name)
except Exception:
continue
self.put_connection(connection)
Database._list_cache = res
Database._list_cache_timestamp = now
return res
def init(self):
from trytond.modules import get_module_info
connection = self.get_connection()
cursor = connection.cursor()
sql_file = os.path.join(os.path.dirname(__file__), 'init.sql')
with open(sql_file) as fp:
for line in fp.read().split(';'):
if (len(line) > 0) and (not line.isspace()):
cursor.execute(line)
for module in ('ir', 'res'):
state = 'uninstalled'
if module in ('ir', 'res'):
state = 'to install'
info = get_module_info(module)
cursor.execute('SELECT NEXTVAL(\'ir_module_id_seq\')')
module_id = cursor.fetchone()[0]
cursor.execute('INSERT INTO ir_module '
'(id, create_uid, create_date, name, state) '
'VALUES (%s, %s, now(), %s, %s)',
(module_id, 0, module, state))
for dependency in info.get('depends', []):
cursor.execute('INSERT INTO ir_module_dependency '
'(create_uid, create_date, module, name) '
'VALUES (%s, now(), %s, %s)',
(0, module_id, dependency))
connection.commit()
self.put_connection(connection)
def test(self):
connection = self.get_connection()
is_tryton_database = self._test(connection)
self.put_connection(connection)
return is_tryton_database
@classmethod
def _test(cls, connection):
cursor = connection.cursor()
cursor.execute('SELECT 1 FROM information_schema.tables '
'WHERE table_name IN %s',
(('ir_model', 'ir_model_field', 'ir_ui_view', 'ir_ui_menu',
'res_user', 'res_group', 'ir_module',
'ir_module_dependency', 'ir_translation',
'ir_lang'),))
return len(cursor.fetchall()) != 0
def nextid(self, connection, table):
cursor = connection.cursor()
cursor.execute("SELECT NEXTVAL('" + table + "_id_seq')")
return cursor.fetchone()[0]
def setnextid(self, connection, table, value):
cursor = connection.cursor()
cursor.execute("SELECT SETVAL('" + table + "_id_seq', %d)" % value)
def currid(self, connection, table):
cursor = connection.cursor()
cursor.execute('SELECT last_value FROM "' + table + '_id_seq"')
return cursor.fetchone()[0]
def lock(self, connection, table):
cursor = connection.cursor()
cursor.execute('LOCK "%s" IN EXCLUSIVE MODE NOWAIT' % table)
def has_constraint(self):
return True
def has_multirow_insert(self):
return True
def get_table_schema(self, connection, table_name):
cursor = connection.cursor()
for schema in self.search_path:
cursor.execute('SELECT 1 '
'FROM information_schema.tables '
'WHERE table_name = %s AND table_schema = %s',
(table_name, schema))
if cursor.rowcount:
return schema
@property
def current_user(self):
if self._current_user is None:
connection = self.get_connection()
try:
cursor = connection.cursor()
cursor.execute('SELECT current_user')
self._current_user = cursor.fetchone()[0]
finally:
self.put_connection(connection)
return self._current_user
@property
def search_path(self):
if self._search_path is None:
connection = self.get_connection()
try:
cursor = connection.cursor()
cursor.execute('SHOW search_path')
path, = cursor.fetchone()
special_values = {
'user': self.current_user,
}
self._search_path = [
unescape_quote(replace_special_values(
p.strip(), **special_values))
for p in path.split(',')]
finally:
self.put_connection(connection)
return self._search_path
class Database():
def __init__(self, config, verbose_start=False):
self.config = config
if verbose_start:
print('Connecting to database...'.ljust(
LogDecorator.get_ljust_offset()),
end='')
# get DB parameters
try:
self.database = config.getProperty('Database', 'name').lower()
self.host = config.getProperty('Database', 'host')
self.port = config.getProperty('Database', 'port')
self.user = config.getProperty('Database', 'user', fallback=None)
self.password = config.getProperty('Database',
'password',
fallback=None)
if self.user is None or self.password is None:
# load from credentials file instead
credentials = config.getProperty('Database', 'credentials')
with open(credentials, 'r') as c:
lines = c.readlines()
for line in lines:
line = line.lstrip().rstrip('\r').rstrip('\n')
if line.startswith('#') or line.startswith(';'):
continue
tokens = line.split('=')
if len(tokens) >= 2:
idx = line.find('=') + 1
field = tokens[0].strip().lower()
if field == 'username':
self.user = line[idx:]
elif field == 'password':
self.password = line[idx:]
self.user = self.user.lower()
except Exception as e:
if verbose_start:
LogDecorator.print_status('fail')
raise Exception(
f'Incomplete database credentials provided in configuration file (message: "{str(e)}").'
)
try:
self._createConnectionPool()
except Exception as e:
if verbose_start:
LogDecorator.print_status('fail')
raise Exception(
f'Could not connect to database (message: "{str(e)}").')
if verbose_start:
LogDecorator.print_status('ok')
def _createConnectionPool(self):
self.connectionPool = ThreadedConnectionPool(
0,
max(
2,
self.config.getProperty('Database',
'max_num_connections',
type=int,
fallback=20)
), # 2 connections are needed as minimum for retrying of execution
host=self.host,
database=self.database,
port=self.port,
user=self.user,
password=self.password,
connect_timeout=10)
def canConnect(self):
with self._get_connection() as conn:
return conn is not None and not conn.closed
@contextmanager
def _get_connection(self):
conn = self.connectionPool.getconn()
conn.autocommit = True
try:
yield conn
finally:
self.connectionPool.putconn(conn, close=False)
def execute(self, query, arguments, numReturn=None):
with self._get_connection() as conn:
cursor = conn.cursor(cursor_factory=RealDictCursor)
# execute statement
try:
cursor.execute(query, arguments)
conn.commit()
except Exception as e:
if not conn.closed:
conn.rollback()
# self.connectionPool.putconn(conn, close=False) #TODO: this still causes connection to close
conn = self.connectionPool.getconn()
# retry execution
try:
cursor = conn.cursor(cursor_factory=RealDictCursor)
cursor.execute(query, arguments)
conn.commit()
except:
if not conn.closed:
conn.rollback()
print(e)
# get results
try:
returnValues = []
if numReturn is None:
# cursor.close()
return
elif numReturn == 'all':
returnValues = cursor.fetchall()
# cursor.close()
return returnValues
else:
for _ in range(numReturn):
rv = cursor.fetchone()
if rv is None:
return returnValues
returnValues.append(rv)
# cursor.close()
return returnValues
except Exception as e:
print(e)
def insert(self, query, values, numReturn=None):
with self._get_connection() as conn:
cursor = conn.cursor()
try:
execute_values(cursor, query, values)
conn.commit()
except Exception as e:
if not conn.closed:
conn.rollback()
# cursor.close()
# retry execution
conn = self.connectionPool.getconn()
try:
cursor = conn.cursor(cursor_factory=RealDictCursor)
execute_values(cursor, query, values)
conn.commit()
except Exception as e:
if not conn.closed:
conn.rollback()
print(e)
# get results
try:
returnValues = []
if numReturn is None:
# cursor.close()
return
elif numReturn == 'all':
returnValues = cursor.fetchall()
# cursor.close()
return returnValues
else:
for _ in range(numReturn):
rv = cursor.fetchone()
if rv is None:
return returnValues
returnValues.append(rv)
# cursor.close()
return returnValues
except Exception as e:
print(e)
class Database(DatabaseInterface):
_databases = {}
_connpool = None
_list_cache = None
_list_cache_timestamp = None
_version_cache = {}
def __new__(cls, database_name='template1'):
if database_name in cls._databases:
return cls._databases[database_name]
return DatabaseInterface.__new__(cls, database_name=database_name)
def __init__(self, database_name='template1'):
super(Database, self).__init__(database_name=database_name)
self._databases.setdefault(database_name, 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)
self._connpool = ThreadedConnectionPool(minconn, maxconn, dsn)
return self
def cursor(self, autocommit=False, readonly=False):
if self._connpool is None:
self.connect()
conn = self._connpool.getconn()
if autocommit:
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
else:
conn.set_isolation_level(ISOLATION_LEVEL_REPEATABLE_READ)
cursor = Cursor(self._connpool, conn, self)
# TODO change for set_session
if readonly:
cursor.execute('SET TRANSACTION READ ONLY')
return cursor
def close(self):
if self._connpool is None:
return
self._connpool.closeall()
self._connpool = None
def create(self, cursor, database_name):
cursor.execute('CREATE DATABASE "' + database_name + '" ' \
'TEMPLATE template0 ENCODING \'unicode\'')
Database._list_cache = None
def drop(self, cursor, database_name):
cursor.execute('DROP DATABASE "' + database_name + '"')
Database._list_cache = None
def get_version(self, cursor):
if self.database_name not in self._version_cache:
cursor.execute('SELECT version()')
version, = cursor.fetchone()
self._version_cache[self.database_name] = tuple(map(int,
RE_VERSION.search(version).groups()))
return self._version_cache[self.database_name]
@staticmethod
def dump(database_name):
from trytond.tools import exec_pg_command_pipe
cmd = ['pg_dump', '--format=c', '--no-owner']
if CONFIG['db_user']:
cmd.append('--username='******'db_user'])
if CONFIG['db_host']:
cmd.append('--host=' + CONFIG['db_host'])
if CONFIG['db_port']:
cmd.append('--port=' + CONFIG['db_port'])
cmd.append(database_name)
pipe = exec_pg_command_pipe(*tuple(cmd))
pipe.stdin.close()
data = pipe.stdout.read()
res = pipe.wait()
if res:
raise Exception('Couldn\'t dump database!')
return data
@staticmethod
def restore(database_name, data):
from trytond.tools import exec_pg_command_pipe
database = Database().connect()
cursor = database.cursor(autocommit=True)
database.create(cursor, database_name)
cursor.commit()
cursor.close()
cmd = ['pg_restore', '--no-owner']
if CONFIG['db_user']:
cmd.append('--username='******'db_user'])
if CONFIG['db_host']:
cmd.append('--host=' + CONFIG['db_host'])
if CONFIG['db_port']:
cmd.append('--port=' + CONFIG['db_port'])
cmd.append('--dbname=' + database_name)
args2 = tuple(cmd)
if os.name == "nt":
tmpfile = (os.environ['TMP'] or 'C:\\') + os.tmpnam()
with open(tmpfile, 'wb') as fp:
fp.write(data)
args2 = list(args2)
args2.append(' ' + tmpfile)
args2 = tuple(args2)
pipe = exec_pg_command_pipe(*args2)
if not os.name == "nt":
pipe.stdin.write(data)
pipe.stdin.close()
res = pipe.wait()
if res:
raise Exception('Couldn\'t restore database')
database = Database(database_name).connect()
cursor = database.cursor()
if not cursor.test():
cursor.close()
database.close()
raise Exception('Couldn\'t restore database!')
cursor.close()
database.close()
Database._list_cache = None
return True
@staticmethod
def list(cursor):
now = time.time()
timeout = int(CONFIG['session_timeout'])
res = Database._list_cache
if res and abs(Database._list_cache_timestamp - now) < timeout:
return res
db_user = CONFIG['db_user']
if not db_user and os.name == 'posix':
db_user = pwd.getpwuid(os.getuid())[0]
if not db_user:
cursor.execute("SELECT usename " \
"FROM pg_user " \
"WHERE usesysid = (" \
"SELECT datdba " \
"FROM pg_database " \
"WHERE datname = %s)",
(CONFIG["db_name"],))
res = cursor.fetchone()
db_user = res and res[0]
if db_user:
cursor.execute("SELECT datname " \
"FROM pg_database " \
"WHERE datdba = (" \
"SELECT usesysid " \
"FROM pg_user " \
"WHERE usename=%s) " \
"AND datname not in " \
"('template0', 'template1', 'postgres') " \
"ORDER BY datname",
(db_user,))
else:
cursor.execute("SELECT datname " \
"FROM pg_database " \
"WHERE datname not in " \
"('template0', 'template1','postgres') " \
"ORDER BY datname")
res = []
for db_name, in cursor.fetchall():
db_name = db_name.encode('utf-8')
try:
database = Database(db_name).connect()
except Exception:
continue
cursor2 = database.cursor()
if cursor2.test():
res.append(db_name)
cursor2.close(close=True)
else:
cursor2.close(close=True)
database.close()
Database._list_cache = res
Database._list_cache_timestamp = now
return res
@staticmethod
def init(cursor):
from trytond.tools import safe_eval
sql_file = os.path.join(os.path.dirname(__file__), 'init.sql')
with open(sql_file) as fp:
for line in fp.read().split(';'):
if (len(line) > 0) and (not line.isspace()):
cursor.execute(line)
for i in ('ir', 'res', 'webdav'):
root_path = os.path.join(os.path.dirname(__file__), '..', '..')
tryton_file = os.path.join(root_path, i, '__tryton__.py')
with open(tryton_file) as fp:
info = safe_eval(fp.read())
active = info.get('active', False)
if active:
state = 'to install'
else:
state = 'uninstalled'
cursor.execute('SELECT NEXTVAL(\'ir_module_module_id_seq\')')
module_id = cursor.fetchone()[0]
cursor.execute('INSERT INTO ir_module_module ' \
'(id, create_uid, create_date, author, website, name, ' \
'shortdesc, description, state) ' \
'VALUES (%s, %s, now(), %s, %s, %s, %s, %s, %s)',
(module_id, 0, info.get('author', ''),
info.get('website', ''), i, info.get('name', False),
info.get('description', ''), state))
dependencies = info.get('depends', [])
for dependency in dependencies:
cursor.execute('INSERT INTO ir_module_module_dependency ' \
'(create_uid, create_date, module, name) ' \
'VALUES (%s, now(), %s, %s)',
(0, module_id, dependency))
