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))