pydbkit: comparison pgtoolkit/pgmanager.py

equal deleted inserted replaced

-:e7f79c4a27ce
+:e67101c22e83
 # Requires: Python 3.2, psycopg2
 #
 # Part of pgtoolkit
 # http://hg.devl.cz/pgtoolkit
 #
-# Copyright (c) 2010, 2011  Radek Brich <radek.brich@devl.cz>
+# Copyright (c) 2010, 2011, 2012  Radek Brich <radek.brich@devl.cz>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
 # in the Software without restriction, including without limitation the rights
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 On second line we have created connection named 'default' (this name can be left out).
 The with statement obtains connection (actually connects to database when needed),
 then returns cursor for this connection. At the end of with statement,
 the connection is returned to the pool or closed (depending on number of connections
-in pool and on setting of keep_open parameter).
+in pool and on setting of pool_size parameter).
 The row returned by fetchone_dict() is special dict object, which can be accessed
 using item or attribute access, that is row['now'] or row.now.
 """
 pass
 class ConnectionInfo:
-def __init__(self, dsn, isolation_level=None, keep_alive=True,
+def __init__(self, name, dsn, isolation_level=None, keep_alive=True,
-init_statement=None, keep_open=1):
+init_statement=None, pool_size=1):
+self.name = name  # connection name is logged with SQL queries
 self.dsn = dsn
 self.isolation_level = isolation_level
 self.keep_alive = keep_alive
 self.init_statement = init_statement
-self.keep_open = keep_open
+self.pool_size = pool_size
 class RowDict(OrderedDict):
 def __getattr__(self, key):
 class Cursor(psycopg2.extensions.cursor):
 def execute(self, query, args=None):
+self._log_query(query, args)
 try:
 return super(Cursor, self).execute(query, args)
-finally:
+except DatabaseError:
-self._log_query()
+self._log_exception()
 def callproc(self, procname, args=None):
+self._log_query('CALL %s(%s)' % (procname, ','.join(args)))
 try:
 return super(Cursor, self).callproc(procname, args)
-finally:
+except DatabaseError:
-self._log_query()
+self._log_exception()
 def row_dict(self, row, lstrip=None):
 adjustname = lambda a: a
 if lstrip:
 adjustname = lambda a: a.lstrip(lstrip)
 def fetchall_adapted(self, lstrip=None):
 '''Like fetchall_dict() but values are quoted for direct inclusion in SQL query.'''
 rows = super(Cursor, self).fetchall()
 return [self.row_dict([self.mogrify('%s', [x]).decode('utf8') for x in row], lstrip) for row in rows]
-def _log_query(self):
+def _log_query(self, query, args=None):
-if self.query:
+name = self.connection.name if hasattr(self.connection, 'name') else '-'
-name = self.connection.name if hasattr(self.connection, 'name') else '-'
+log_sql.info('[%s] %s' % (name, self.mogrify(query, args).decode('utf8')))
-log_sql.info('[%s] %s' % (name, self.query.decode('utf8')))
+def _log_exception(self):
+name = self.connection.name if hasattr(self.connection, 'name') else '-'
+log_sql.exception('[%s] exception:' % (name,))
 class Connection(psycopg2.extensions.connection):
 def cursor(self, name=None):
 if name is None:
 def __del__(self):
 for conn in tuple(self.conn_known.keys()):
 self.destroy_conn(conn)
-def create_conn(self, name='default', keep_open=1, isolation_level=None, keep_alive=True, dsn=None, **kw):
+def create_conn(self, name='default', isolation_level=None, keep_alive=True, init_statement=None,
+pool_size=1, dsn=None, **kwargs):
 '''Create named connection.
 name -- name for connection (default is "default")
-keep_open -- how many connections will be kept open in pool (more connections will still be created,
+pool_size -- how many connections will be kept open in pool
-but they will be closed by put_conn)
+(more connections will still be created but they will be closed by put_conn)
+None - disable pool, always return same connection
 isolation_level -- "autocommit", "read_committed", "serializable" or None for driver default
 keep_alive -- set socket to keepalive mode
 dsn -- string with connection parameters (dsn means Data Source Name)
 Alternative for dsn is keyword args (same names as in dsn).
 '''
 if name in self.conn_known:
 raise PgManagerError('Connection name "%s" already registered.' % name)
 if dsn is None:
-dsn = ' '.join([x[0]+'='+str(x[1]) for x in kw.items() if x[1] is not None])
+dsn = ' '.join([x[0]+'='+str(x[1]) for x in kwargs.items() if x[1] is not None])
 isolation_level = self._normalize_isolation_level(isolation_level)
-ci = ConnectionInfo(dsn, isolation_level, keep_alive, keep_open=keep_open)
+ci = ConnectionInfo(name, dsn, isolation_level, keep_alive, init_statement, pool_size)
 self.conn_known[name] = ci
 self.conn_pool[name] = []
+def create_conn_listen(self, name, channel, dsn=None, copy_dsn=None, **kwargs):
+'''Create connection listening for notifies.
+Disables pool. If you want to use pool, create other connection for that.
+This connection can be used as usual: conn.cursor() etc.
+Don't use PgManager's cursor() and put_conn().
+name -- name for connection
+channel -- listen on this channel
+copy_dsn -- specify name of other connection and its dsn will be used
+Other parameters forwarded to create_conn().
+'''
+if dsn is None and copy_dsn:
+try:
+dsn = self.conn_known[copy_dsn].dsn
+except KeyError:
+raise PgManagerError("Connection name '%s' not registered." % copy_dsn)
+listen_query = "LISTEN " + channel
+self.create_conn(name=name, pool_size=None, isolation_level='autocommit', init_statement=listen_query,
+dsn=dsn, **kwargs)
 def close_conn(self, name='default'):
 '''Close all connections of given name.
 Connection credentials are still saved.
 def get_conn(self, name='default'):
 '''Get connection of name 'name' from pool.'''
 self._check_fork()
 self.lock.acquire()
 try:
-if not name in self.conn_known:
+try:
+ci = self.conn_known[name]
+except KeyError:
 raise PgManagerError("Connection name '%s' not registered." % name)
-conn = None
+# no pool, just one static connection
-while len(self.conn_pool[name]) and conn is None:
+if ci.pool_size is None:
-conn = self.conn_pool[name].pop()
+# check for existing connection
-if conn.closed:
+try:
+conn = self.conn_pool[name][0]
+if conn.closed:
+conn = None
+except IndexError:
 conn = None
+self.conn_pool[name].append(conn)
-if conn is None:
+# if no existing connection is valid, connect new one and save it
-ci = self.conn_known[name]
+if conn is None:
 conn = self._connect(ci)
-# add our name to connection instance (this is then logged with SQL queries)
+self.conn_pool[name][0] = conn
-conn.name = name
+# connection from pool
+else:
+conn = None
+while len(self.conn_pool[name]) and conn is None:
+conn = self.conn_pool[name].pop()
+if conn.closed:
+conn = None
+if conn is None:
+conn = self._connect(ci)
 finally:
 self.lock.release()
 return conn
 def put_conn(self, conn, name='default'):
 self.lock.acquire()
 try:
 if not name in self.conn_known:
 raise PgManagerError("Connection name '%s' not registered." % name)
-if len(self.conn_pool[name]) >= self.conn_known[name].keep_open:
+if len(self.conn_pool[name]) >= self.conn_known[name].pool_size:
 conn.close()
 return
 if conn.get_transaction_status() == psycopg2.extensions.TRANSACTION_STATUS_UNKNOWN:
 conn.close()
 def log_notices(self, conn):
 for notice in conn.notices:
 log_notices.info(notice.rstrip())
 conn.notices[:] = []
-def wait_for_notify(self, name='default', timeout=5):
+def wait_for_notify(self, name='default', timeout=5.0):
 '''Wait for asynchronous notifies, return the last one.
+name -- name of connection, must be created using create_conn_listen()
+timeout -- in seconds, floating point (None - wait forever)
 Returns None on timeout.
 '''
 conn = self.get_conn(name)
-try:
+# return any notifies on stack
-# any residual notify?
+if conn.notifies:
-# then return it, that should not break anything
+return conn.notifies.pop()
+if select.select([conn], [], [], timeout) == ([], [], []):
+# timeout
+return None
+else:
+conn.poll()
+# return just the last notify (we do not care for older ones)
 if conn.notifies:
 return conn.notifies.pop()
+return None
-if select.select([conn], [], [], timeout) == ([], [], []):
-# timeout
-return None
-else:
-conn.poll()
-# return just the last notify (we do not care for older ones)
-if conn.notifies:
-return conn.notifies.pop()
-return None
-finally:
-# clean notifies
-while conn.notifies:
-conn.notifies.pop()
-self.put_conn(conn, name)
 def _connect(self, ci):
 conn = psycopg2.connect(ci.dsn, connection_factory=Connection)
+conn.name = ci.name
 if ci.keep_alive:
 conn.keep_alive()
 if not ci.isolation_level is None:
 conn.set_isolation_level(ci.isolation_level)
 if ci.init_statement:

changeset 36	e67101c22e83
parent 33	bd0beda49bcb
child 37	5b0eb4b11940