root / logilab.pylintinstaller / logilab / common / db.py

Revision 202:d67e86292521, 23.0 kB (checked in by tziade@…, 9 months ago)

added logilab.pylintinstaller

Line 
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
17modules and database connection whatever the database and client lib used.
18
19Currently support:
20
21- postgresql (pgdb, psycopg, psycopg2, pyPgSQL)
22- mysql (MySQLdb)
23- sqlite (pysqlite2, sqlite, sqlite3)
24
25just use the `get_connection` function from this module to get a
26wrapped connection.  If multiple drivers for a database are available,
27you can control which one you want to use using the
28`set_prefered_driver` function.
29
30Additional helpers are also provided for advanced functionalities such
31as listing existing users or databases, creating database... Get the
32helper for your database using the `get_adv_func_helper` function.
33"""
34
35import sys
36import re
37from warnings import warn
38
39from logilab.common.deprecation import obsolete
40try:
41    from mx.DateTime import DateTimeType, DateTimeDeltaType, strptime
42    HAS_MX_DATETIME = True
43except:
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
52class UnknownDriver(Exception):
53    """raised when a unknown driver is given to get connexion"""
54
55class 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
67def _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       
97class 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
124class 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   
159class 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
211class _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
230class _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   
247class _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
281class _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
304class _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   
409class _SqliteAdapter(DBAPIAdapter):