bin/sql_db.py

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