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