bin/sql_db.py

   1 # -*- coding: utf-8 -*-
   2 ##############################################################################
   3 #
   4 #    OpenERP, Open Source Management Solution
   5 #    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>).
   6 #
   7 #    This program is free software: you can redistribute it and/or modify
   8 #    it under the terms of the GNU Affero General Public License as
   9 #    published by the Free Software Foundation, either version 3 of the
  10 #    License, or (at your option) any later version.
  11 #
  12 #    This program is distributed in the hope that it will be useful,
  13 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 #    GNU Affero General Public License for more details.
  16 #
  17 #    You should have received a copy of the GNU Affero General Public License
  18 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
  19 #
  20 ##############################################################################
  21
  22 __all__ = ['db_connect', 'close_db']
  23
  24 import netsvc
  25 from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT, ISOLATION_LEVEL_READ_COMMITTED, ISOLATION_LEVEL_SERIALIZABLE
  26 from psycopg2.psycopg1 import cursor as psycopg1cursor
  27 from psycopg2.pool import PoolError
  28
  29 import psycopg2.extensions
  30
  31 psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
  32
  33 types_mapping = {
  34     'date': (1082,),
  35     'time': (1083,),
  36     'datetime': (1114,),
  37 }
  38
  39 def unbuffer(symb, cr):
  40     if symb is None: return None
  41     return str(symb)
  42
  43 def undecimalize(symb, cr):
  44     if symb is None: return None
  45     return float(symb)
  46
  47 for name, typeoid in types_mapping.items():
  48     psycopg2.extensions.register_type(psycopg2.extensions.new_type(typeoid, name, lambda x, cr: x))
  49 psycopg2.extensions.register_type(psycopg2.extensions.new_type((700, 701, 1700,), 'float', undecimalize))
  50
  51
  52 import tools
  53 from tools.func import wraps
  54 from datetime import datetime as mdt
  55 from datetime import timedelta
  56 import threading
  57
  58 import re
  59 re_from = re.compile('.* from "?([a-zA-Z_0-9]+)"? .*$');
  60 re_into = re.compile('.* into "?([a-zA-Z_0-9]+)"? .*$');
  61
  62
  63 def log(msg, lvl=netsvc.LOG_DEBUG2):
  64     logger = netsvc.Logger()
  65     logger.notifyChannel('sql', lvl, msg)
  66
  67 sql_counter = 0
  68
  69 class Cursor(object):
  70     IN_MAX = 1000
  71
  72     def check(f):
  73         @wraps(f)
  74         def wrapper(self, *args, **kwargs):
  75             if self.__closed:
  76                 raise psycopg2.ProgrammingError('Unable to use the cursor after having closed it')
  77             return f(self, *args, **kwargs)
  78         return wrapper
  79
  80     def __init__(self, pool, dbname, serialized=False):
  81         self.sql_from_log = {}
  82         self.sql_into_log = {}
  83         self.sql_log = False
  84         self.sql_log_count = 0
  85         self.__closed = True    # avoid the call of close() (by __del__) if an exception
  86                                 # is raised by any of the following initialisations
  87         self._pool = pool
  88         self.dbname = dbname
  89         self._serialized = serialized
  90         self._cnx = pool.borrow(dsn(dbname))
  91         self._obj = self._cnx.cursor(cursor_factory=psycopg1cursor)
  92         self.__closed = False   # real initialisation value
  93         self.autocommit(False)
  94
  95         if tools.config['log_level'] in (netsvc.LOG_DEBUG, netsvc.LOG_DEBUG_RPC):
  96             from inspect import stack
  97             self.__caller = tuple(stack()[2][1:3])
  98
  99     def __del__(self):
 100         if not self.__closed:
 101             # Oops. 'self' has not been closed explicitly.
 102             # The cursor will be deleted by the garbage collector,
 103             # but the database connection is not put back into the connection
 104             # pool, preventing some operation on the database like dropping it.
 105             # This can also lead to a server overload.
 106             if tools.config['log_level'] in (netsvc.LOG_DEBUG, netsvc.LOG_DEBUG_RPC):
 107                 msg = "Cursor not closed explicitly\n"  \
 108                       "Cursor was created at %s:%s" % self.__caller
 109                 log(msg, netsvc.LOG_WARNING)
 110             self.close()
 111
 112     @check
 113     def execute(self, query, params=None):
 114         if '%d' in query or '%f' in query:
 115             log(query, netsvc.LOG_WARNING)
 116             log("SQL queries cannot contain %d or %f anymore. Use only %s", netsvc.LOG_WARNING)
 117             if params:
 118                 query = query.replace('%d', '%s').replace('%f', '%s')
 119
 120         if self.sql_log:
 121             now = mdt.now()
 122
 123         try:
 124             params = params or None
 125             res = self._obj.execute(query, params)
 126         except psycopg2.ProgrammingError, pe:
 127             logger= netsvc.Logger()
 128             logger.notifyChannel('sql_db', netsvc.LOG_ERROR, "Programming error: %s, in query %s" % (pe, query))
 129             raise
 130         except Exception, e:
 131             log("bad query: %s" % self._obj.query)
 132             log(e)
 133             raise
 134
 135         if self.sql_log:
 136             delay = mdt.now() - now
 137             delay = delay.seconds * 1E6 + delay.microseconds
 138
 139             log("query: %s" % self._obj.query)
 140             self.sql_log_count+=1
 141             res_from = re_from.match(query.lower())
 142             if res_from:
 143                 self.sql_from_log.setdefault(res_from.group(1), [0, 0])
 144                 self.sql_from_log[res_from.group(1)][0] += 1
 145                 self.sql_from_log[res_from.group(1)][1] += delay
 146             res_into = re_into.match(query.lower())
 147             if res_into:
 148                 self.sql_into_log.setdefault(res_into.group(1), [0, 0])
 149                 self.sql_into_log[res_into.group(1)][0] += 1
 150                 self.sql_into_log[res_into.group(1)][1] += delay
 151         return res
 152
 153     def print_log(self):
 154         global sql_counter
 155         sql_counter += self.sql_log_count
 156         if not self.sql_log:
 157             return
 158         def process(type):
 159             sqllogs = {'from':self.sql_from_log, 'into':self.sql_into_log}
 160             sum = 0
 161             if sqllogs[type]:
 162                 sqllogitems = sqllogs[type].items()
 163                 sqllogitems.sort(key=lambda k: k[1][1])
 164                 log("SQL LOG %s:" % (type,))
 165                 for r in sqllogitems:
 166                     delay = timedelta(microseconds=r[1][1])
 167                     log("table: %s: %s/%s" %(r[0], str(delay), r[1][0]))
 168                     sum+= r[1][1]
 169                 sqllogs[type].clear()
 170             sum = timedelta(microseconds=sum)
 171             log("SUM %s:%s/%d [%d]" % (type, str(sum), self.sql_log_count, sql_counter))
 172             sqllogs[type].clear()
 173         process('from')
 174         process('into')
 175         self.sql_log_count = 0
 176         self.sql_log = False
 177
 178     @check
 179     def close(self):
 180         if not self._obj:
 181             return
 182
 183         self.print_log()
 184
 185         if not self._serialized:
 186             self.rollback() # Ensure we close the current transaction.
 187
 188         self._obj.close()
 189
 190         # This force the cursor to be freed, and thus, available again. It is
 191         # important because otherwise we can overload the server very easily
 192         # because of a cursor shortage (because cursors are not garbage
 193         # collected as fast as they should). The problem is probably due in
 194         # part because browse records keep a reference to the cursor.
 195         del self._obj
 196         self.__closed = True
 197         self._pool.give_back(self._cnx)
 198
 199     @check
 200     def autocommit(self, on):
 201         offlevel = [ISOLATION_LEVEL_READ_COMMITTED, ISOLATION_LEVEL_SERIALIZABLE][bool(self._serialized)]
 202         self._cnx.set_isolation_level([offlevel, ISOLATION_LEVEL_AUTOCOMMIT][bool(on)])
 203
 204     @check
 205     def commit(self):
 206         return self._cnx.commit()
 207
 208     @check
 209     def rollback(self):
 210         return self._cnx.rollback()
 211
 212     @check
 213     def __getattr__(self, name):
 214         return getattr(self._obj, name)
 215
 216
 217 class ConnectionPool(object):
 218
 219     def locked(fun):
 220         @wraps(fun)
 221         def _locked(self, *args, **kwargs):
 222             self._lock.acquire()
 223             try:
 224                 return fun(self, *args, **kwargs)
 225             finally:
 226                 self._lock.release()
 227         return _locked
 228
 229
 230     def __init__(self, maxconn=64):
 231         self._connections = []
 232         self._maxconn = max(maxconn, 1)
 233         self._lock = threading.Lock()
 234         self._logger = netsvc.Logger()
 235
 236     def _log(self, msg):
 237         #self._logger.notifyChannel('ConnectionPool', netsvc.LOG_INFO, msg)
 238         pass
 239     def _debug(self, msg):
 240         #self._logger.notifyChannel('ConnectionPool', netsvc.LOG_DEBUG, msg)
 241         pass
 242
 243     @locked
 244     def borrow(self, dsn):
 245         self._log('Borrow connection to %s' % (dsn,))
 246
 247         result = None
 248         for i, (cnx, used) in enumerate(self._connections):
 249             if not used and dsn_are_equals(cnx.dsn, dsn):
 250                 self._debug('Existing connection found at index %d' % i)
 251
 252                 self._connections.pop(i)
 253                 self._connections.append((cnx, True))
 254
 255                 result = cnx
 256                 break
 257         if result:
 258             return result
 259
 260         if len(self._connections) >= self._maxconn:
 261             # try to remove the older connection not used
 262             for i, (cnx, used) in enumerate(self._connections):
 263                 if not used:
 264                     self._debug('Removing old connection at index %d: %s' % (i, cnx.dsn))
 265                     self._connections.pop(i)
 266                     break
 267             else:
 268                 # note: this code is called only if the for loop has completed (no break)
 269                 raise PoolError('Connection Pool Full')
 270
 271         self._debug('Create new connection')
 272         result = psycopg2.connect(dsn=dsn)
 273         self._connections.append((result, True))
 274         return result
 275
 276     @locked
 277     def give_back(self, connection):
 278         self._log('Give back connection to %s' % (connection.dsn,))
 279         for i, (cnx, used) in enumerate(self._connections):
 280             if cnx is connection:
 281                 self._connections.pop(i)
 282                 self._connections.append((cnx, False))
 283                 break
 284         else:
 285             raise PoolError('This connection does not below to the pool')
 286
 287     @locked
 288     def close_all(self, dsn):
 289         for i, (cnx, used) in tools.reverse_enumerate(self._connections):
 290             if dsn_are_equals(cnx.dsn, dsn):
 291                 cnx.close()
 292                 self._connections.pop(i)
 293
 294
 295 class Connection(object):
 296     __LOCKS = {}
 297
 298     def __init__(self, pool, dbname, unique=False):
 299         self.dbname = dbname
 300         self._pool = pool
 301         self._unique = unique
 302
 303     def __enter__(self):
 304         if self._unique:
 305             self.lock()
 306         return self
 307
 308     def __exit__(self, exc_type, exc_value, traceback):
 309         if self._unique:
 310             self.release()
 311
 312     def lock(self):
 313         if self.dbname not in self.__LOCKS:
 314             self.__LOCKS[self.dbname] = threading.Lock()
 315         self.__LOCKS[self.dbname].acquire()
 316
 317     def release(self):
 318         close_db(self.dbname)
 319         self.__LOCKS[self.dbname].release()
 320
 321     def cursor(self, serialized=False):
 322         if self._unique:
 323             lock = self.__LOCKS.get(self.dbname, None)
 324             if not (lock and lock.locked()):
 325                 netsvc.Logger().notifyChannel('Connection', netsvc.LOG_WARNING, 'Unprotected connection to %s' % (self.dbname,))
 326
 327         return Cursor(self._pool, self.dbname, serialized=serialized)
 328
 329     def serialized_cursor(self):
 330         return self.cursor(True)
 331
 332     def __nonzero__(self):
 333         """Check if connection is possible"""
 334         try:
 335             cr = self.cursor()
 336             cr.close()
 337             return True
 338         except:
 339             return False
 340
 341
 342 _dsn = ''
 343 for p in ('host', 'port', 'user', 'password'):
 344     cfg = tools.config['db_' + p]
 345     if cfg:
 346         _dsn += '%s=%s ' % (p, cfg)
 347
 348 def dsn(db_name):
 349     return '%sdbname=%s' % (_dsn, db_name)
 350
 351 def dsn_are_equals(first, second):
 352     def key(dsn):
 353         k = dict(x.split('=', 1) for x in dsn.strip().split())
 354         k.pop('password', None) # password is not relevant
 355         return k
 356     return key(first) == key(second)
 357
 358
 359 _Pool = ConnectionPool(int(tools.config['db_maxconn']))
 360
 361 def db_connect(db_name):
 362     unique = db_name in ['template1', 'template0']
 363     return Connection(_Pool, db_name, unique)
 364
 365 def close_db(db_name):
 366     _Pool.close_all(dsn(db_name))
 367     tools.cache.clean_caches_for_db(db_name)
 368
 369
 370 # vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
 371