#-*- coding: utf-8 -*- """Threadsafe database wrapper library over psycopg Very simple---- Oops, not so simple. File descriptors are copied on forks, so the thread-safe variant must take pid into account. >>> db = pgdb.get_cursor('host=localhost dbname=eyedb ' ... 'user=eyeinvoice password=Kthulhu4presidente', ... {'encoding': 'UNICODE'}) >>> >>> db.execute("SELECT COUNT(*) FROM accounting.invoices") >>> db.dictfetchall() [{'count': 138145L}] @author: Anders Eurenius @author: Ulf Renman @organization: Favoptic Glasögondirekt AB @requires: config : dbstring, dbparams @var con: The current database connection @type con: psycopg connection """ import sys, os, time, thread v = None try: import psycopg2 as pg import psycopg2.extras as ex v = 2 except: import psycopg as pg v = 1 try: import config # import global config if there is one config.dsn = getattr(config, 'dsn', '') config.bug = getattr(config, 'bug', {}) except: class old: pass # make our own, reachable from outside. config = old() # we don't insert it because that might del old # interfere with loading that module later. config.dsn = '' # We delete it, because it's junk. config.bug = {} if '--dsn' in sys.argv and len(sys.argv) < sys.index('--dsn')+1: config.dsn = sys.pop(index('--dsn')+1) sys.remove('--dsn') def init(): global con, thr, bug, thr_lvl, iso_lvl thr_lvl = ['Threads may not share the module.', 'Threads may share the module, but not connections.', 'Threads may share the module and connections.', 'Threads may share the module, connections and cursors.'] iso_lvl = ['Autocommit','Read committed','Read uncommitted', 'Repeatable read', 'Serializable'] con, thr, bug = None, {}, getattr(config, 'bug', {}) if bug: print 'psycopg: ', pg.__version__ print 'API level: ', pg.apilevel print 'Param style: ', pg.paramstyle print 'Thread safety:', thr_lvl[pg.threadsafety] init() def dsn_clean(dsn=''): l = dsn.split('=') # split kvskv to k,vsk,v d, k = {}, l[0] # k for vk in l[1:-1]: d[k], k = vk.rsplit(' ', 1) # =v,k d[k] = l[-1] # =v dsn = '' for k,v in d.items(): v = v.rstrip() if v: v = v[0] == v[-1] and v[0] in '\'"' and v[1:-1] or v d[k] = v for k in 'dbname','host','port','user','password','sslmode': if k in d: dsn = dsn + (k.lower() +'='+ d[k]) return dsn, d def connect(dsn): dsn, d = dsn_clean(dsn) return pg.connect(dsn), d def get_cursor(dsn=''): """Creates a new connection (if necessary) and a new cursor""" global con cxn = con if dsn: cxn, d = connect(dsn) elif not con or con.closed: con, d = connect(config.dsn) cxn = con else: d = {} if v == 1: cur = cxn.cursor() elif v == 2: cur = cxn.cursor(cursor_factory = ex.DictCursor) if d.get('autocommit'): if v == 1: cxn.autocommit(1) elif v == 2: cxn.set_isolation_level(0) if d.get('encoding'): if v == 1: cur.execute("SET CLIENT_ENCODING=%s", (d['encoding'],)) elif v == 2: cxn.set_client_encoding(d['encoding']) return cur def get_thread_cursor(dsn='', pid=[-1]): """Utility to create separate cursors for different (pid x thread)s.""" global thr, con if pid[0] != os.getpid(): con, pid[0], thr = None, os.getpid(), {} me = (dsn, thread.get_ident()) return thr.setdefault(me, get_cursor(dsn)) class Query( object ): """Utility for reusing a query in a safe and convenient way. The instance is created with 1. the query string, 2. an ordered list of pairs of parameters and their casting (or conversion, or..) functions. (optional) 3. a dictionary containing default values. When called the argument defaults to an empty dictionary. If a dict is given on the other hand, the defaults are copied, the copy is then updated with the argument dict. A list of query paramenters is then constructed by picking the dict items according to the keys list and mapping them with their corresponding functions. @note: The same dict key can be used more than once in the query. @note: If the query fails because of a db restart, it tries to reconnect. Although the class is crafted for that use, 1. The query need not be a select, 2. The keys need not be strings and 3. The functions need not be constructors; Notably, they can return None """ def __init__(self, sql, keys=(), defaults={}): self.sql, self.keys, self.defaults = sql, keys, defaults def __repr__(self): return ('Query("""%s""" x (%s) x {%s})' % (self.sql, ', '.join([ k for k, f in self.keys]), ', '.join( self.defaults.keys() ) )) def __call__(self, d={}, retry=1): e = dict(self.defaults) # (shallow) copy the defaults if isinstance( d, dict ): e.update(d) # args override defaults l = [ f(e[k]) for k, f in self.keys ] # map out the args else: # same, but assume object l = [ f( getattr(d, k, e.get(k)) ) for k, f in self.keys ] try: tc = get_thread_cursor() if bug.get('pgdb') > 2: print self.sql % l tc.execute(self.sql, l) except pg.OperationalError, x: if bug.get('pgdb') > 1: print x if retry: global thr if bug.get('pgdb') > 0: print '--- RECONNECTING ---' make_connection() thr = {} return self(d, retry-1) except Exception, x: if bug.get('pgdb') > 0: print x if tc.statusmessage == 'SELECT': return [ dict(x) for x in tc.fetchall() ] or [] else: return tc.rowcount class CachedQuery( list ): """Result caching list class that refreshes itself. The instance is given 1. a function or functor, (L{Query}, hint, hint.) 2. an optional cache time in seconds 3. a function the query results are mapped through The point of the exercise is to get stuff from the db while balancing 1. not doing a query every time 2. allowing change without restarting the application It can be refreshed manually with C{refresh}, so if you like, you can set the cache time to 2**64 and refresh it explicitly. >>> cq = CachedQuery(Query("SELECT * FROM pg_catalog.pg_class"), ... f=lambda x: (x['start'].strftime('%F'), x['uid'])) @warning: Failure semantics are not so well thought out. (If you have a better idea, tell me.) """ def __init__(self, q, to=300, f=lambda x: x): self.q, self.to, self.f, self.t = q, to, f, 0 self.refresh() def refresh(self): t = time.time() if t - self.t > self.to: if bug.get('pgdb') > 1: print 'REFRESH: ', self.q try: self.t, self[:] = t, map(self.f, self.q()) except: self[:] = [] def __getitem__(self, *x): self.refresh();return list.__getitem__(self,*x) def __getslice__(self,*x): self.refresh();return list.__getslice__(self,*x) def __repr__(self,*x): self.refresh();return list.__repr__(self, *x) def __str__(self,*x): self.refresh();return list.__str__(self, *x) def __len__(self,*x): self.refresh();return list.__len__(self, *x)