1 """Miscellaneous goodies for psycopg2
3 This module is a generic place used to hold little helper functions
4 and classes untill a better place in the distribution is found.
6 # psycopg/extras.py - miscellaneous extra goodies for psycopg
8 # Copyright (C) 2003-2004 Federico Di Gregorio <fog@debian.org>
10 # This program is free software; you can redistribute it and/or modify
11 # it under the terms of the GNU General Public License as published by the
12 # Free Software Foundation; either version 2, or (at your option) any later
15 # This program is distributed in the hope that it will be useful, but
16 # WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTIBILITY
17 # or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
28 from psycopg2.extensions import cursor as _cursor
29 from psycopg2.extensions import connection as _connection
30 from psycopg2.extensions import register_adapter as _RA
31 from psycopg2.extensions import adapt as _A
34 class DictConnection(_connection):
35 """A connection that uses DictCursor automatically."""
37 return _connection.cursor(self, cursor_factory=DictCursor)
39 class DictCursor(_cursor):
40 """A cursor that keeps a list of column name -> index mappings."""
44 def execute(self, query, vars=None, async=0):
45 self.row_factory = DictRow
47 self.__query_executed = 1
48 return _cursor.execute(self, query, vars, async)
50 def callproc(self, procname, vars=None):
51 self.row_factory = DictRow
53 self.__query_executed = 1
54 return _cursor.callproc(self, procname, vars)
56 def _build_index(self):
57 if self.__query_executed == 1 and self.description:
58 for i in range(len(self.description)):
59 self.index[self.description[i][0]] = i
60 self.__query_executed = 0
63 res = _cursor.fetchone(self)
64 if self.__query_executed:
68 def fetchmany(self, size=None):
69 res = _cursor.fetchmany(self, size)
70 if self.__query_executed:
75 res = _cursor.fetchall(self)
76 if self.__query_executed:
81 res = _cursor.fetchone(self)
84 if self.__query_executed:
89 """A row object that allow by-colun-name access to data."""
91 def __init__(self, cursor):
92 self._index = cursor.index
93 self[:] = [None] * len(cursor.description)
95 def __getitem__(self, x):
98 return list.__getitem__(self, x)
102 for n, v in self._index.items():
103 res.append((n, list.__getitem__(self, v)))
107 return self._index.keys()
110 return tuple(self[:])
112 def has_key(self, x):
113 return self._index.has_key(x)
115 def get(self, x, default=None):
122 class SQL_IN(object):
123 """Adapt any iterable to an SQL quotable object."""
125 def __init__(self, seq):
128 def prepare(self, conn):
132 # this is the important line: note how every object in the
133 # list is adapted and then how getquoted() is called on it
134 pobjs = [_A(o) for o in self._seq]
136 if hasattr(obj, 'prepare'):
137 obj.prepare(self._conn)
138 qobjs = [str(o.getquoted()) for o in pobjs]
139 return '(' + ', '.join(qobjs) + ')'
146 class LoggingConnection(_connection):
147 """A connection that logs all queries to a file or logger object."""
149 def initialize(self, logobj):
150 """Initialize the connection to log to `logobj`.
152 The `logobj` parameter can be an open file object or a Logger instance
153 from the standard logging module.
155 self._logobj = logobj
156 if logging and isinstance(logobj, logging.Logger):
157 self.log = self._logtologger
159 self.log = self._logtofile
161 def filter(self, msg, curs):
162 """Filter the query before logging it.
164 This is the method to overwrite to filter unwanted queries out of the
165 log or to add some extra data to the output. The default implementation
170 def _logtofile(self, msg, curs):
171 msg = self.filter(msg, curs)
172 if msg: self._logobj.write(msg + os.linesep)
174 def _logtologger(self, msg, curs):
175 msg = self.filter(msg, curs)
176 if msg: self._logobj.debug(msg)
179 if not hasattr(self, '_logobj'):
180 raise self.ProgrammingError(
181 "LoggingConnection object has not been initialize()d")
185 return _connection.cursor(self, cursor_factory=LoggingCursor)
187 class LoggingCursor(_cursor):
188 """A cursor that logs queries using its connection logging facilities."""
190 def execute(self, query, vars=None, async=0):
192 return _cursor.execute(self, query, vars, async)
194 self.connection.log(self.query, self)
196 def callproc(self, procname, vars=None):
198 return _cursor.callproc(self, procname, vars)
200 self.connection.log(self.query, self)
203 class MinTimeLoggingConnection(LoggingConnection):
204 """A connection that logs queries based on execution time.
206 This is just an example of how to sub-class LoggingConnection to provide
207 some extra filtering for the logged queries. Both the `.inizialize()` and
208 `.filter()` methods are overwritten to make sure that only queries
209 executing for more than `mintime` ms are logged.
211 Note that this connection uses the specialized cursor MinTimeLoggingCursor.
213 def initialize(self, logobj, mintime=0):
214 LoggingConnection.initialize(self, logobj)
215 self._mintime = mintime
217 def filter(self, msg, curs):
218 t = (time.time() - curs.timestamp) * 1000
219 if t > self._mintime:
220 return msg + os.linesep + " (execution time: %d ms)" % t
224 return _connection.cursor(self, cursor_factory=MinTimeLoggingCursor)
226 class MinTimeLoggingCursor(LoggingCursor):
227 """The cursor sub-class companion to MinTimeLoggingConnection."""
229 def execute(self, query, vars=None, async=0):
230 self.timestamp = time.time()
231 return LoggingCursor.execute(self, query, vars, async)
233 def callproc(self, procname, vars=None):
234 self.timestamp = time.time()
235 return LoggingCursor.execute(self, procname, var)