| 1 | # Copyright (c) 2002-2007 LOGILAB S.A. (Paris, FRANCE). |
|---|
| 2 | # http://www.logilab.fr/ -- mailto:contact@logilab.fr |
|---|
| 3 | # |
|---|
| 4 | # This program is free software; you can redistribute it and/or modify it under |
|---|
| 5 | # the terms of the GNU General Public License as published by the Free Software |
|---|
| 6 | # Foundation; either version 2 of the License, or (at your option) any later |
|---|
| 7 | # version. |
|---|
| 8 | # |
|---|
| 9 | # This program is distributed in the hope that it will be useful, but WITHOUT |
|---|
| 10 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS |
|---|
| 11 | # FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. |
|---|
| 12 | # |
|---|
| 13 | # You should have received a copy of the GNU General Public License along with |
|---|
| 14 | # this program; if not, write to the Free Software Foundation, Inc., |
|---|
| 15 | # 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA. |
|---|
| 16 | """This modules contains wrappers to get actually replaceable DBAPI2 compliant |
|---|
| 17 | modules and database connection whatever the database and client lib used. |
|---|
| 18 | |
|---|
| 19 | Currently support: |
|---|
| 20 | |
|---|
| 21 | - postgresql (pgdb, psycopg, psycopg2, pyPgSQL) |
|---|
| 22 | - mysql (MySQLdb) |
|---|
| 23 | - sqlite (pysqlite2, sqlite, sqlite3) |
|---|
| 24 | |
|---|
| 25 | just use the `get_connection` function from this module to get a |
|---|
| 26 | wrapped connection. If multiple drivers for a database are available, |
|---|
| 27 | you can control which one you want to use using the |
|---|
| 28 | `set_prefered_driver` function. |
|---|
| 29 | |
|---|
| 30 | Additional helpers are also provided for advanced functionalities such |
|---|
| 31 | as listing existing users or databases, creating database... Get the |
|---|
| 32 | helper for your database using the `get_adv_func_helper` function. |
|---|
| 33 | """ |
|---|
| 34 | |
|---|
| 35 | import sys |
|---|
| 36 | import re |
|---|
| 37 | from warnings import warn |
|---|
| 38 | |
|---|
| 39 | from logilab.common.deprecation import obsolete |
|---|
| 40 | try: |
|---|
| 41 | from mx.DateTime import DateTimeType, DateTimeDeltaType, strptime |
|---|
| 42 | HAS_MX_DATETIME = True |
|---|
| 43 | except: |
|---|
| 44 | HAS_MX_DATETIME = False |
|---|
| 45 | |
|---|
| 46 | __all__ = ['get_dbapi_compliant_module', |
|---|
| 47 | 'get_connection', 'set_prefered_driver', |
|---|
| 48 | 'PyConnection', 'PyCursor', |
|---|
| 49 | 'UnknownDriver', 'NoAdapterFound', |
|---|
| 50 | ] |
|---|
| 51 | |
|---|
| 52 | class UnknownDriver(Exception): |
|---|
| 53 | """raised when a unknown driver is given to get connexion""" |
|---|
| 54 | |
|---|
| 55 | class NoAdapterFound(Exception): |
|---|
| 56 | """Raised when no Adpater to DBAPI was found""" |
|---|
| 57 | def __init__(self, obj, objname=None, protocol='DBAPI'): |
|---|
| 58 | if objname is None: |
|---|
| 59 | objname = obj.__name__ |
|---|
| 60 | Exception.__init__(self, "Could not adapt %s to protocol %s" % |
|---|
| 61 | (objname, protocol)) |
|---|
| 62 | self.adapted_obj = obj |
|---|
| 63 | self.objname = objname |
|---|
| 64 | self._protocol = protocol |
|---|
| 65 | |
|---|
| 66 | |
|---|
| 67 | def _import_driver_module(driver, drivers, imported_elements=None, quiet=True): |
|---|
| 68 | """Imports the first module found in 'drivers' for 'driver' |
|---|
| 69 | |
|---|
| 70 | :rtype: tuple |
|---|
| 71 | :returns: the tuple module_object, module_name where module_object |
|---|
| 72 | is the dbapi module, and modname the module's name |
|---|
| 73 | """ |
|---|
| 74 | if not driver in drivers: |
|---|
| 75 | raise UnknownDriver(driver) |
|---|
| 76 | imported_elements = imported_elements or [] |
|---|
| 77 | for modname in drivers[driver]: |
|---|
| 78 | try: |
|---|
| 79 | if not quiet: |
|---|
| 80 | print >> sys.stderr, 'Trying %s' % modname |
|---|
| 81 | module = __import__(modname, globals(), locals(), imported_elements) |
|---|
| 82 | break |
|---|
| 83 | except ImportError: |
|---|
| 84 | if not quiet: |
|---|
| 85 | print >> sys.stderr, '%s is not available' % modname |
|---|
| 86 | continue |
|---|
| 87 | else: |
|---|
| 88 | raise ImportError('Unable to import a %s module' % driver) |
|---|
| 89 | if not imported_elements: |
|---|
| 90 | for part in modname.split('.')[1:]: |
|---|
| 91 | module = getattr(module, part) |
|---|
| 92 | return module, modname |
|---|
| 93 | |
|---|
| 94 | |
|---|
| 95 | ## Connection and cursor wrappers ############################################# |
|---|
| 96 | |
|---|
| 97 | class PyConnection: |
|---|
| 98 | """A simple connection wrapper in python (useful for profiling)""" |
|---|
| 99 | def __init__(self, cnx): |
|---|
| 100 | """Wraps the original connection object""" |
|---|
| 101 | self._cnx = cnx |
|---|
| 102 | |
|---|
| 103 | # XXX : Would it work if only __getattr__ was defined |
|---|
| 104 | def cursor(self): |
|---|
| 105 | """Wraps cursor()""" |
|---|
| 106 | return PyCursor(self._cnx.cursor()) |
|---|
| 107 | |
|---|
| 108 | def commit(self): |
|---|
| 109 | """Wraps commit()""" |
|---|
| 110 | return self._cnx.commit() |
|---|
| 111 | |
|---|
| 112 | def rollback(self): |
|---|
| 113 | """Wraps rollback()""" |
|---|
| 114 | return self._cnx.rollback() |
|---|
| 115 | |
|---|
| 116 | def close(self): |
|---|
| 117 | """Wraps close()""" |
|---|
| 118 | return self._cnx.close() |
|---|
| 119 | |
|---|
| 120 | def __getattr__(self, attrname): |
|---|
| 121 | return getattr(self._cnx, attrname) |
|---|
| 122 | |
|---|
| 123 | |
|---|
| 124 | class PyCursor: |
|---|
| 125 | """A simple cursor wrapper in python (useful for profiling)""" |
|---|
| 126 | def __init__(self, cursor): |
|---|
| 127 | self._cursor = cursor |
|---|
| 128 | |
|---|
| 129 | def close(self): |
|---|
| 130 | """Wraps close()""" |
|---|
| 131 | return self._cursor.close() |
|---|
| 132 | |
|---|
| 133 | def execute(self, *args, **kwargs): |
|---|
| 134 | """Wraps execute()""" |
|---|
| 135 | return self._cursor.execute(*args, **kwargs) |
|---|
| 136 | |
|---|
| 137 | def executemany(self, *args, **kwargs): |
|---|
| 138 | """Wraps executemany()""" |
|---|
| 139 | return self._cursor.executemany(*args, **kwargs) |
|---|
| 140 | |
|---|
| 141 | def fetchone(self, *args, **kwargs): |
|---|
| 142 | """Wraps fetchone()""" |
|---|
| 143 | return self._cursor.fetchone(*args, **kwargs) |
|---|
| 144 | |
|---|
| 145 | def fetchmany(self, *args, **kwargs): |
|---|
| 146 | """Wraps execute()""" |
|---|
| 147 | return self._cursor.fetchmany(*args, **kwargs) |
|---|
| 148 | |
|---|
| 149 | def fetchall(self, *args, **kwargs): |
|---|
| 150 | """Wraps fetchall()""" |
|---|
| 151 | return self._cursor.fetchall(*args, **kwargs) |
|---|
| 152 | |
|---|
| 153 | def __getattr__(self, attrname): |
|---|
| 154 | return getattr(self._cursor, attrname) |
|---|
| 155 | |
|---|
| 156 | |
|---|
| 157 | ## Adapters list ############################################################## |
|---|
| 158 | |
|---|
| 159 | class DBAPIAdapter: |
|---|
| 160 | """Base class for all DBAPI adpaters""" |
|---|
| 161 | |
|---|
| 162 | def __init__(self, native_module, pywrap=False): |
|---|
| 163 | """ |
|---|
| 164 | :type native_module: module |
|---|
| 165 | :param native_module: the database's driver adapted module |
|---|
| 166 | """ |
|---|
| 167 | self._native_module = native_module |
|---|
| 168 | self._pywrap = pywrap |
|---|
| 169 | |
|---|
| 170 | def connect(self, host='', database='', user='', password='', port=''): |
|---|
| 171 | """Wraps the native module connect method""" |
|---|
| 172 | kwargs = {'host' : host, 'port' : port, 'database' : database, |
|---|
| 173 | 'user' : user, 'password' : password} |
|---|
| 174 | cnx = self._native_module.connect(**kwargs) |
|---|
| 175 | return self._wrap_if_needed(cnx) |
|---|
| 176 | |
|---|
| 177 | def _wrap_if_needed(self, cnx): |
|---|
| 178 | """Wraps the connection object if self._pywrap is True, and returns it |
|---|
| 179 | If false, returns the original cnx object |
|---|
| 180 | """ |
|---|
| 181 | if self._pywrap: |
|---|
| 182 | return PyConnection(cnx) |
|---|
| 183 | else: |
|---|
| 184 | return cnx |
|---|
| 185 | |
|---|
| 186 | def __getattr__(self, attrname): |
|---|
| 187 | return getattr(self._native_module, attrname) |
|---|
| 188 | |
|---|
| 189 | def process_value(self, value, description, encoding='utf-8', binarywrap=None): |
|---|
| 190 | typecode = description[1] |
|---|
| 191 | assert typecode is not None, self # dbapi module isn't supporting type codes, override to return value directly |
|---|
| 192 | if typecode == self.STRING: |
|---|
| 193 | if isinstance(value, str): |
|---|
| 194 | return unicode(value, encoding, 'replace') |
|---|
| 195 | elif typecode == self.BOOLEAN: |
|---|
| 196 | return bool(value) |
|---|
| 197 | elif typecode == self.BINARY and not binarywrap is None: |
|---|
| 198 | return binarywrap(value) |
|---|
| 199 | ## elif typecode == dbapimod.DATETIME: |
|---|
| 200 | ## pass |
|---|
| 201 | ## elif typecode == dbapimod.NUMBER: |
|---|
| 202 | ## pass |
|---|
| 203 | ## else: |
|---|
| 204 | ## self.warning("type -%s- unknown for %r (%s) ", |
|---|
| 205 | ## typecode, value, type(value)) |
|---|
| 206 | return value |
|---|
| 207 | |
|---|
| 208 | |
|---|
| 209 | # Postgresql ######################################################### |
|---|
| 210 | |
|---|
| 211 | class _PgdbAdapter(DBAPIAdapter): |
|---|
| 212 | """Simple PGDB Adapter to DBAPI (pgdb modules lacks Binary() and NUMBER) |
|---|
| 213 | """ |
|---|
| 214 | def __init__(self, native_module, pywrap=False): |
|---|
| 215 | DBAPIAdapter.__init__(self, native_module, pywrap) |
|---|
| 216 | self.NUMBER = native_module.pgdbType('int2', 'int4', 'serial', |
|---|
| 217 | 'int8', 'float4', 'float8', |
|---|
| 218 | 'numeric', 'bool', 'money') |
|---|
| 219 | |
|---|
| 220 | def connect(self, host='', database='', user='', password='', port=''): |
|---|
| 221 | """Wraps the native module connect method""" |
|---|
| 222 | if port: |
|---|
| 223 | warn("pgdb doesn't support 'port' parameter in connect()", UserWarning) |
|---|
| 224 | kwargs = {'host' : host, 'database' : database, |
|---|
| 225 | 'user' : user, 'password' : password} |
|---|
| 226 | cnx = self._native_module.connect(**kwargs) |
|---|
| 227 | return self._wrap_if_needed(cnx) |
|---|
| 228 | |
|---|
| 229 | |
|---|
| 230 | class _PsycopgAdapter(DBAPIAdapter): |
|---|
| 231 | """Simple Psycopg Adapter to DBAPI (cnx_string differs from classical ones) |
|---|
| 232 | """ |
|---|
| 233 | def connect(self, host='', database='', user='', password='', port=''): |
|---|
| 234 | """Handles psycopg connexion format""" |
|---|
| 235 | if host: |
|---|
| 236 | cnx_string = 'host=%s dbname=%s user=%s' % (host, database, user) |
|---|
| 237 | else: |
|---|
| 238 | cnx_string = 'dbname=%s user=%s' % (database, user) |
|---|
| 239 | if port: |
|---|
| 240 | cnx_string += ' port=%s' % port |
|---|
| 241 | if password: |
|---|
| 242 | cnx_string = '%s password=%s' % (cnx_string, password) |
|---|
| 243 | cnx = self._native_module.connect(cnx_string) |
|---|
| 244 | cnx.set_isolation_level(1) |
|---|
| 245 | return self._wrap_if_needed(cnx) |
|---|
| 246 | |
|---|
| 247 | class _Psycopg2Adapter(_PsycopgAdapter): |
|---|
| 248 | """Simple Psycopg2 Adapter to DBAPI (cnx_string differs from classical ones) |
|---|
| 249 | """ |
|---|
| 250 | BOOLEAN = 16 # XXX see additional types in psycopg2.extensions |
|---|
| 251 | def __init__(self, native_module, pywrap=False): |
|---|
| 252 | DBAPIAdapter.__init__(self, native_module, pywrap) |
|---|
| 253 | self._init_psycopg2() |
|---|
| 254 | |
|---|
| 255 | def _init_psycopg2(self): |
|---|
| 256 | """initialize psycopg2 to use mx.DateTime for date and timestamps |
|---|
| 257 | instead for datetime.datetime""" |
|---|
| 258 | psycopg2 = self._native_module |
|---|
| 259 | if hasattr(psycopg2, '_lc_initialized'): |
|---|
| 260 | return |
|---|
| 261 | psycopg2._lc_initialized = 1 |
|---|
| 262 | # use mxDateTime instead of datetime if available |
|---|
| 263 | if HAS_MX_DATETIME: |
|---|
| 264 | from psycopg2 import extensions |
|---|
| 265 | extensions.register_type(psycopg2._psycopg.MXDATETIME) |
|---|
| 266 | extensions.register_type(psycopg2._psycopg.MXINTERVAL) |
|---|
| 267 | extensions.register_type(psycopg2._psycopg.MXDATE) |
|---|
| 268 | extensions.register_type(psycopg2._psycopg.MXTIME) |
|---|
| 269 | # StringIO/cStringIO adaptation |
|---|
| 270 | # XXX (syt) todo, see my december discussion on the psycopg2 list |
|---|
| 271 | # for a working solution |
|---|
| 272 | #def adapt_stringio(stringio): |
|---|
| 273 | # print 'ADAPTING', stringio |
|---|
| 274 | # return psycopg2.Binary(stringio.getvalue()) |
|---|
| 275 | #import StringIO |
|---|
| 276 | #extensions.register_adapter(StringIO.StringIO, adapt_stringio) |
|---|
| 277 | #import cStringIO |
|---|
| 278 | #extensions.register_adapter(cStringIO.StringIO, adapt_stringio) |
|---|
| 279 | |
|---|
| 280 | |
|---|
| 281 | class _PgsqlAdapter(DBAPIAdapter): |
|---|
| 282 | """Simple pyPgSQL Adapter to DBAPI |
|---|
| 283 | """ |
|---|
| 284 | def connect(self, host='', database='', user='', password='', port=''): |
|---|
| 285 | """Handles psycopg connexion format""" |
|---|
| 286 | kwargs = {'host' : host, 'port': port or None, |
|---|
| 287 | 'database' : database, |
|---|
| 288 | 'user' : user, 'password' : password or None} |
|---|
| 289 | cnx = self._native_module.connect(**kwargs) |
|---|
| 290 | return self._wrap_if_needed(cnx) |
|---|
| 291 | |
|---|
| 292 | |
|---|
| 293 | def Binary(self, string): |
|---|
| 294 | """Emulates the Binary (cf. DB-API) function""" |
|---|
| 295 | return str |
|---|
| 296 | |
|---|
| 297 | def __getattr__(self, attrname): |
|---|
| 298 | # __import__('pyPgSQL.PgSQL', ...) imports the toplevel package |
|---|
| 299 | return getattr(self._native_module, attrname) |
|---|
| 300 | |
|---|
| 301 | |
|---|
| 302 | # Sqlite ############################################################# |
|---|
| 303 | |
|---|
| 304 | class _PySqlite2Adapter(DBAPIAdapter): |
|---|
| 305 | """Simple pysqlite2 Adapter to DBAPI |
|---|
| 306 | """ |
|---|
| 307 | def __init__(self, native_module, pywrap=False): |
|---|
| 308 | DBAPIAdapter.__init__(self, native_module, pywrap) |
|---|
| 309 | self._init_pysqlite2() |
|---|
| 310 | # no type code in pysqlite2 |
|---|
| 311 | self.BINARY = 'XXX' |
|---|
| 312 | self.STRING = 'XXX' |
|---|
| 313 | self.DATETIME = 'XXX' |
|---|
| 314 | self.NUMBER = 'XXX' |
|---|
| 315 | |
|---|
| 316 | def _init_pysqlite2(self): |
|---|
| 317 | """initialize pysqlite2 to use mx.DateTime for date and timestamps""" |
|---|
| 318 | sqlite = self._native_module |
|---|
| 319 | if hasattr(sqlite, '_lc_initialized'): |
|---|
| 320 | return |
|---|
| 321 | sqlite._lc_initialized = 1 |
|---|
| 322 | |
|---|
| 323 | # bytea type handling |
|---|
| 324 | from StringIO import StringIO |
|---|
| 325 | def adapt_bytea(data): |
|---|
| 326 | return data.getvalue() |
|---|
| 327 | sqlite.register_adapter(StringIO, adapt_bytea) |
|---|
| 328 | def convert_bytea(data): |
|---|
| 329 | return StringIO(data) |
|---|
| 330 | sqlite.register_converter('bytea', convert_bytea) |
|---|
| 331 | |
|---|
| 332 | # boolean type handling |
|---|
| 333 | def convert_boolean(ustr): |
|---|
| 334 | if ustr.upper() in ('F', 'FALSE'): |
|---|
| 335 | return False |
|---|
| 336 | return True |
|---|
| 337 | sqlite.register_converter('boolean', convert_boolean) |
|---|
| 338 | def adapt_boolean(bval): |
|---|
| 339 | return str(bval).upper() |
|---|
| 340 | sqlite.register_adapter(bool, adapt_boolean) |
|---|
| 341 | |
|---|
| 342 | # date/time types handling |
|---|
| 343 | if HAS_MX_DATETIME: |
|---|
| 344 | def adapt_mxdatetime(mxd): |
|---|
| 345 | return mxd.strftime('%Y-%m-%d %H:%M:%S') |
|---|
| 346 | sqlite.register_adapter(DateTimeType, adapt_mxdatetime) |
|---|
| 347 | def adapt_mxdatetimedelta(mxd): |
|---|
| 348 | return mxd.strftime('%H:%M:%S') |
|---|
| 349 | sqlite.register_adapter(DateTimeDeltaType, adapt_mxdatetimedelta) |
|---|
| 350 | |
|---|
| 351 | def convert_mxdate(ustr): |
|---|
| 352 | return strptime(ustr, '%Y-%m-%d %H:%M:%S') |
|---|
| 353 | sqlite.register_converter('date', convert_mxdate) |
|---|
| 354 | def convert_mxdatetime(ustr): |
|---|
| 355 | return strptime(ustr, '%Y-%m-%d %H:%M:%S') |
|---|
| 356 | sqlite.register_converter('timestamp', convert_mxdatetime) |
|---|
| 357 | def convert_mxtime(ustr): |
|---|
| 358 | try: |
|---|
| 359 | return strptime(ustr, '%H:%M:%S') |
|---|
| 360 | except: |
|---|
| 361 | # DateTime used as Time? |
|---|
| 362 | return strptime(ustr, '%Y-%m-%d %H:%M:%S') |
|---|
| 363 | sqlite.register_converter('time', convert_mxtime) |
|---|
| 364 | # XXX else use datetime.datetime |
|---|
| 365 | |
|---|
| 366 | |
|---|
| 367 | def connect(self, host='', database='', user='', password='', port=None): |
|---|
| 368 | """Handles sqlite connexion format""" |
|---|
| 369 | sqlite = self._native_module |
|---|
| 370 | |
|---|
| 371 | class PySqlite2Cursor(sqlite.Cursor): |
|---|
| 372 | """cursor adapting usual dict format to pysqlite named format |
|---|
| 373 | in SQL queries |
|---|
| 374 | """ |
|---|
| 375 | def _replace_parameters(self, sql, kwargs): |
|---|
| 376 | if isinstance(kwargs, dict): |
|---|
| 377 | return re.sub(r'%\(([^\)]+)\)s', r':\1', sql) |
|---|
| 378 | # XXX dumb |
|---|
| 379 | return re.sub(r'%s', r'?', sql) |
|---|
| 380 | |
|---|
| 381 | def execute(self, sql, kwargs=None): |
|---|
| 382 | if kwargs is None: |
|---|
| 383 | self.__class__.__bases__[0].execute(self, sql) |
|---|
| 384 | else: |
|---|
| 385 | self.__class__.__bases__[0].execute(self, self._replace_parameters(sql, kwargs), kwargs) |
|---|
| 386 | |
|---|
| 387 | def executemany(self, sql, kwargss): |
|---|
| 388 | if not isinstance(kwargss, (list, tuple)): |
|---|
| 389 | kwargss = tuple(kwargss) |
|---|
| 390 | self.__class__.__bases__[0].executemany(self, self._replace_parameters(sql, kwargss[0]), kwargss) |
|---|
| 391 | |
|---|
| 392 | class PySqlite2CnxWrapper: |
|---|
| 393 | def __init__(self, cnx): |
|---|
| 394 | self._cnx = cnx |
|---|
| 395 | |
|---|
| 396 | def cursor(self): |
|---|
| 397 | return self._cnx.cursor(PySqlite2Cursor) |
|---|
| 398 | def __getattr__(self, attrname): |
|---|
| 399 | return getattr(self._cnx, attrname) |
|---|
| 400 | cnx = sqlite.connect(database, detect_types=sqlite.PARSE_DECLTYPES) |
|---|
| 401 | return self._wrap_if_needed(PySqlite2CnxWrapper(cnx)) |
|---|
| 402 | |
|---|
| 403 | def process_value(self, value, description, encoding='utf-8', binarywrap=None): |
|---|
| 404 | if not binarywrap is None and isinstance(value, self._native_module.Binary): |
|---|
| 405 | return binarywrap(value) |
|---|
| 406 | return value # no type code support, can't do anything |
|---|
| 407 | |
|---|
| 408 | |
|---|
| 409 | class _SqliteAdapter(DBAPIAdapter): |
|---|
|
|---|