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 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_DEBUG):
  64     logger = netsvc.Logger()
  65     logger.notifyChannel('sql', lvl, msg)
  66
  67 class Cursor(object):
  68     IN_MAX = 1000
  69
  70     def check(f):
  71         @wraps(f)
  72         def wrapper(self, *args, **kwargs):
  73             if self.__closed:
  74                 raise psycopg2.ProgrammingError('Unable to use the cursor after having closing it')
  75             return f(self, *args, **kwargs)
  76         return wrapper
  77
  78     def __init__(self, pool, dbname, serialized=False):
  79         self.sql_from_log = {}
  80         self.sql_into_log = {}
  81         self.sql_log = False
  82         self.sql_log_count = 0
  83
  84         self.__closed = True    # avoid the call of close() (by __del__) if an exception
  85                                 # is raised by any of the following initialisations
  86         self._pool = pool
  87         self.dbname = dbname
  88         self._serialized = serialized
  89         self._cnx = pool.borrow(dsn(dbname))
  90         self._obj = self._cnx.cursor(cursor_factory=psycopg1cursor)
  91         self.__closed = False   # real initialisation value
  92         self.autocommit(False)
  93
  94         if tools.config['log_level'] in (netsvc.LOG_DEBUG, netsvc.LOG_DEBUG_RPC):
  95             from inspect import stack
  96             self.__caller = tuple(stack()[2][1:3])
  97
  98     def __del__(self):
  99         if not self.__closed:
 100             # Oops. 'self' has not been closed explicitly.
 101             # The cursor will be deleted by the garbage collector,
 102             # but the database connection is not put back into the connection
 103             # pool, preventing some operation on the database like dropping it.
 104             # This can also lead to a server overload.
 105             if tools.config['log_level'] in (netsvc.LOG_DEBUG, netsvc.LOG_DEBUG_RPC):
 106                 msg = "Cursor not closed explicitly\n"  \
 107                       "Cursor was created at %s:%s" % self.__caller
 108                 log(msg, netsvc.LOG_WARNING)
 109             self.close()
 110
 111     @check
 112     def execute(self, query, params=None):
 113         if '%d' in query or '%f' in query:
 114             log(query, netsvc.LOG_WARNING)
 115             log("SQL queries mustn't contain %d or %f anymore. Use only %s", netsvc.LOG_WARNING)
 116             if params:
 117                 query = query.replace('%d', '%s').replace('%f', '%s')
 118
 119         if self.sql_log:
 120             now = mdt.now()
 121
 122         try:
 123             params = params or None
 124             res = self._obj.execute(query, params)
 125         except Exception, e:
 126             log("bad query: %s" % self._obj.query)
 127             log(e)
 128             raise
 129
 130         if self.sql_log:
 131             log("query: %s" % self._obj.query)
 132             self.sql_log_count+=1
 133             res_from = re_from.match(query.lower())
 134             if res_from:
 135                 self.sql_from_log.setdefault(res_from.group(1), [0, 0])
 136                 self.sql_from_log[res_from.group(1)][0] += 1
 137                 self.sql_from_log[res_from.group(1)][1] += mdt.now() - now
 138             res_into = re_into.match(query.lower())
 139             if res_into:
 140                 self.sql_into_log.setdefault(res_into.group(1), [0, 0])
 141                 self.sql_into_log[res_into.group(1)][0] += 1
 142                 self.sql_into_log[res_into.group(1)][1] += mdt.now() - now
 143         return res
 144
 145     def print_log(self):
 146         if not self.sql_log:
 147             return
 148
 149         def process(type):
 150             sqllogs = {'from':self.sql_from_log, 'into':self.sql_into_log}
 151             if not sqllogs[type]:
 152                 return
 153             sqllogitems = sqllogs[type].items()
 154             sqllogitems.sort(key=lambda k: k[1][1])
 155             sum = 0
 156             log("SQL LOG %s:" % (type,))
 157             for r in sqllogitems:
 158                 log("table: %s: %s/%s" %(r[0], str(r[1][1]), r[1][0]))
 159                 sum+= r[1][1]
 160             log("SUM:%s/%d" % (sum, self.sql_log_count))
 161             sqllogs[type].clear()
 162         process('from')
 163         process('into')
 164         self.sql_log_count = 0
 165         self.sql_log = False
 166
 167     @check
 168     def close(self):
 169         if not self._obj:
 170             return
 171
 172         self.print_log()
 173
 174         if not self._serialized:
 175             self.rollback() # Ensure we close the current transaction.
 176
 177         self._obj.close()
 178
 179         # This force the cursor to be freed, and thus, available again. It is
 180         # important because otherwise we can overload the server very easily
 181         # because of a cursor shortage (because cursors are not garbage
 182         # collected as fast as they should). The problem is probably due in
 183         # part because browse records keep a reference to the cursor.
 184         del self._obj
 185         self.__closed = True
 186         self._pool.give_back(self._cnx)
 187
 188     @check
 189     def autocommit(self, on):
 190         offlevel = [ISOLATION_LEVEL_READ_COMMITTED, ISOLATION_LEVEL_SERIALIZABLE][bool(self._serialized)]
 191         self._cnx.set_isolation_level([offlevel, ISOLATION_LEVEL_AUTOCOMMIT][bool(on)])
 192
 193     @check
 194     def commit(self):
 195         return self._cnx.commit()
 196
 197     @check
 198     def rollback(self):
 199         return self._cnx.rollback()
 200
 201     @check
 202     def __getattr__(self, name):
 203         return getattr(self._obj, name)
 204
 205
 206 class ConnectionPool(object):
 207
 208     def locked(fun):
 209         @wraps(fun)
 210         def _locked(self, *args, **kwargs):
 211             self._lock.acquire()
 212             try:
 213                 return fun(self, *args, **kwargs)
 214             finally:
 215                 self._lock.release()
 216         return _locked
 217
 218
 219     def __init__(self, maxconn=64):
 220         self._connections = []
 221         self._maxconn = max(maxconn, 1)
 222         self._lock = threading.Lock()
 223         self._logger = netsvc.Logger()
 224
 225     def _log(self, msg):
 226         #self._logger.notifyChannel('ConnectionPool', netsvc.LOG_INFO, msg)
 227         pass
 228     def _debug(self, msg):
 229         #self._logger.notifyChannel('ConnectionPool', netsvc.LOG_DEBUG, msg)
 230         pass
 231
 232     @locked
 233     def borrow(self, dsn):
 234         self._log('Borrow connection to %s' % (dsn,))
 235
 236         result = None
 237         for i, (cnx, used) in enumerate(self._connections):
 238             if not used and cnx.dsn == dsn:
 239                 self._debug('Existing connection found at index %d' % i)
 240
 241                 self._connections.pop(i)
 242                 self._connections.append((cnx, True))
 243
 244                 result = cnx
 245                 break
 246         if result:
 247             return result
 248
 249         if len(self._connections) >= self._maxconn:
 250             # try to remove the older connection not used
 251             for i, (cnx, used) in enumerate(self._connections):
 252                 if not used:
 253                     self._debug('Removing old connection at index %d: %s' % (i, cnx.dsn))
 254                     self._connections.pop(i)
 255                     break
 256             else:
 257                 # note: this code is called only if the for loop has completed (no break)
 258                 raise PoolError('Connection Pool Full')
 259
 260         self._debug('Create new connection')
 261         result = psycopg2.connect(dsn=dsn)
 262         self._connections.append((result, True))
 263         return result
 264
 265     @locked
 266     def give_back(self, connection):
 267         self._log('Give back connection to %s' % (connection.dsn,))
 268         for i, (cnx, used) in enumerate(self._connections):
 269             if cnx is connection:
 270                 self._connections.pop(i)
 271                 self._connections.append((cnx, False))
 272                 break
 273         else:
 274             raise PoolError('This connection does not below to the pool')
 275
 276     @locked
 277     def close_all(self, dsn):
 278         for i, (cnx, used) in tools.reverse_enumerate(self._connections):
 279             if cnx.dsn == dsn:
 280                 cnx.close()
 281                 self._connections.pop(i)
 282
 283
 284 class Connection(object):
 285     __LOCKS = {}
 286
 287     def __init__(self, pool, dbname, unique=False):
 288         self.dbname = dbname
 289         self._pool = pool
 290         self._unique = unique
 291         if unique:
 292             if dbname not in self.__LOCKS:
 293                 self.__LOCKS[dbname] = threading.Lock()
 294             self.__LOCKS[dbname].acquire()
 295
 296     def __del__(self):
 297         if self._unique:
 298             self.__LOCKS[self.dbname].release()
 299
 300     def cursor(self, serialized=False):
 301         return Cursor(self._pool, self.dbname, serialized=serialized)
 302
 303     def serialized_cursor(self):
 304         return self.cursor(True)
 305
 306
 307 _dsn = ''
 308 for p in ('host', 'port', 'user', 'password'):
 309     cfg = tools.config['db_' + p]
 310     if cfg:
 311         _dsn += '%s=%s ' % (p, cfg)
 312
 313 def dsn(db_name):
 314     return '%sdbname=%s' % (_dsn, db_name)
 315
 316
 317 _Pool = ConnectionPool(int(tools.config['db_maxconn']))
 318
 319 def db_connect(db_name):
 320     unique = db_name in ['template1', 'template0']
 321     return Connection(_Pool, db_name, unique)
 322
 323 def close_db(db_name):
 324     _Pool.close_all(dsn(db_name))
 325     tools.cache.clean_caches_for_db(db_name)
 326
 327
 328 # vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
 329