# This file is part of Tryton. The COPYRIGHT file at the top level of # this repository contains the full copyright notices and license terms. import time import logging import os import urllib import json from decimal import Decimal from threading import RLock try: from psycopg2cffi import compat compat.register() except ImportError: pass from psycopg2 import connect, Binary from psycopg2.pool import ThreadedConnectionPool, PoolError from psycopg2.extensions import cursor from psycopg2.extensions import ISOLATION_LEVEL_REPEATABLE_READ from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT from psycopg2.extensions import register_type, register_adapter from psycopg2.extensions import UNICODE, AsIs try: from psycopg2.extensions import PYDATE, PYDATETIME, PYTIME, PYINTERVAL except ImportError: PYDATE, PYDATETIME, PYTIME, PYINTERVAL = None, None, None, None from psycopg2 import IntegrityError as DatabaseIntegrityError from psycopg2 import OperationalError as DatabaseOperationalError from psycopg2.extras import register_default_json, register_default_jsonb from sql import Flavor from sql.functions import Function from trytond.backend.database import DatabaseInterface, SQLType from trytond.config import config, parse_uri from trytond.protocols.jsonrpc import JSONDecoder __all__ = ['Database', 'DatabaseIntegrityError', 'DatabaseOperationalError'] logger = logging.getLogger(__name__) os.environ['PGTZ'] = os.environ.get('TZ', '') def unescape_quote(s): if s.startswith('"') and s.endswith('"'): return s.strip('"').replace('""', '"') return s def replace_special_values(s, **mapping): for name, value in mapping.iteritems(): s = s.replace('$' + name, value) return s class LoggingCursor(cursor): def execute(self, sql, args=None): if logger.isEnabledFor(logging.DEBUG): logger.debug(self.mogrify(sql, args)) cursor.execute(self, sql, args) class Unaccent(Function): __slots__ = () _function = 'unaccent' class Database(DatabaseInterface): _lock = RLock() _databases = {} _connpool = None _list_cache = None _list_cache_timestamp = None _version_cache = {} _search_path = None _current_user = None _has_returning = None _has_unaccent = {} flavor = Flavor(ilike=True) TYPES_MAPPING = { 'INTEGER': SQLType('INT4', 'INT4'), 'BIGINT': SQLType('INT8', 'INT8'), 'FLOAT': SQLType('FLOAT8', 'FLOAT8'), 'BLOB': SQLType('BYTEA', 'BYTEA'), 'DATETIME': SQLType('TIMESTAMP', 'TIMESTAMP(0)'), 'TIMESTAMP': SQLType('TIMESTAMP', 'TIMESTAMP(6)'), } def __new__(cls, name='template1'): with cls._lock: if name in cls._databases: return cls._databases[name] inst = DatabaseInterface.__new__(cls, name=name) logger.info('connect to "%s"', name) minconn = config.getint('database', 'minconn', default=1) maxconn = config.getint('database', 'maxconn', default=64) inst._connpool = ThreadedConnectionPool( minconn, maxconn, cls.dsn(name), cursor_factory=LoggingCursor) cls._databases[name] = inst return inst @classmethod def dsn(cls, name): uri = parse_uri(config.get('database', 'uri')) 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): return self def get_connection(self, autocommit=False, readonly=False): for count in range(config.getint('database', 'retry'), -1, -1): try: conn = self._connpool.getconn() break except PoolError: if count and not self._connpool.closed: logger.info('waiting a connection') time.sleep(1) continue raise 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') return conn def put_connection(self, connection, close=False): self._connpool.putconn(connection, close=close) def close(self): with self._lock: self._connpool.closeall() self._databases.pop(self.name) @classmethod def create(cls, connection, database_name, template='template0'): cursor = connection.cursor() cursor.execute('CREATE DATABASE "' + database_name + '" ' 'TEMPLATE "' + template + '" ENCODING \'unicode\'') connection.commit() cls._list_cache = None def drop(self, connection, database_name): cursor = connection.cursor() cursor.execute('DROP DATABASE "' + database_name + '"') self.__class__._list_cache = None def get_version(self, connection): if self.name not in self._version_cache: cursor = connection.cursor() cursor.execute('SHOW server_version_num') version, = cursor.fetchone() major, rest = divmod(int(version), 10000) minor, patch = divmod(rest, 100) self._version_cache[self.name] = (major, minor, patch) return self._version_cache[self.name] def list(self): now = time.time() timeout = config.getint('session', 'timeout') res = self.__class__._list_cache if res and abs(self.__class__._list_cache_timestamp - now) < timeout: return res connection = self.get_connection() try: 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 finally: self.put_connection(connection) self.__class__._list_cache = res self.__class__._list_cache_timestamp = now return res def init(self): from trytond.modules import get_module_info connection = self.get_connection() try: 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 = 'not activated' if module in ('ir', 'res'): state = 'to activate' 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() finally: 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, constraint): 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 def has_returning(self): if self._has_returning is None: connection = self.get_connection() try: # RETURNING clause is available since PostgreSQL 8.2 self._has_returning = self.get_version(connection) >= (8, 2) finally: self.put_connection(connection) return self._has_returning def has_select_for(self): return True def has_window_functions(self): return True @classmethod def has_sequence(cls): return True def has_unaccent(self): if self.name in self._has_unaccent: return self._has_unaccent[self.name] connection = self.get_connection() unaccent = False try: cursor = connection.cursor() cursor.execute( "SELECT 1 FROM pg_proc WHERE proname=%s", (Unaccent._function,)) unaccent = bool(cursor.rowcount) finally: self.put_connection(connection) self._has_unaccent[self.name] = unaccent return unaccent def sql_type(self, type_): if type_ in self.TYPES_MAPPING: return self.TYPES_MAPPING[type_] if type_.startswith('VARCHAR'): return SQLType('VARCHAR', type_) return SQLType(type_, type_) def sql_format(self, type_, value): if type_ == 'BLOB': if value is not None: return Binary(value) return value def unaccent(self, value): if self.has_unaccent(): return Unaccent(value) return value def sequence_exist(self, connection, name): cursor = connection.cursor() for schema in self.search_path: cursor.execute('SELECT 1 ' 'FROM information_schema.sequences ' 'WHERE sequence_name = %s AND sequence_schema = %s', (name, schema)) if cursor.rowcount: return True return False def sequence_create( self, connection, name, number_increment=1, start_value=1): cursor = connection.cursor() param = self.flavor.param cursor.execute( 'CREATE SEQUENCE "%s" ' 'INCREMENT BY %s ' 'START WITH %s' % (name, param, param), (number_increment, start_value)) def sequence_update( self, connection, name, number_increment=1, start_value=1): cursor = connection.cursor() param = self.flavor.param cursor.execute( 'ALTER SEQUENCE "%s" ' 'INCREMENT BY %s ' 'RESTART WITH %s' % (name, param, param), (number_increment, start_value)) def sequence_rename(self, connection, old_name, new_name): cursor = connection.cursor() if (self.sequence_exist(connection, old_name) and not self.sequence_exist(connection, new_name)): cursor.execute('ALTER TABLE "%s" RENAME TO "%s"' % (old_name, new_name)) def sequence_delete(self, connection, name): cursor = connection.cursor() cursor.execute('DROP SEQUENCE "%s"' % name) def sequence_next_number(self, connection, name): cursor = connection.cursor() version = self.get_version(connection) if version >= (10, 0): cursor.execute( 'SELECT increment_by ' 'FROM pg_sequences ' 'WHERE sequencename=%s ' % self.flavor.param, (name,)) increment, = cursor.fetchone() cursor.execute( 'SELECT CASE WHEN NOT is_called THEN last_value ' 'ELSE last_value + %s ' 'END ' 'FROM "%s"' % (self.flavor.param, name), (increment,)) else: cursor.execute( 'SELECT CASE WHEN NOT is_called THEN last_value ' 'ELSE last_value + increment_by ' 'END ' 'FROM "%s"' % name) return cursor.fetchone()[0] register_type(UNICODE) if PYDATE: register_type(PYDATE) if PYDATETIME: register_type(PYDATETIME) if PYTIME: register_type(PYTIME) if PYINTERVAL: register_type(PYINTERVAL) register_adapter(float, lambda value: AsIs(repr(value))) register_adapter(Decimal, lambda value: AsIs(str(value))) def convert_json(value): return json.loads(value, object_hook=JSONDecoder()) register_default_json(loads=convert_json) register_default_jsonb(loads=convert_json)