diff -r a9e12b7cc207 -r e526ca146fa9 pgtoolkit/pgmanager.py --- a/pgtoolkit/pgmanager.py Fri Nov 25 18:17:27 2011 +0100 +++ b/pgtoolkit/pgmanager.py Sun Nov 27 17:02:42 2011 +0100 @@ -78,7 +78,7 @@ import psycopg2 import psycopg2.extensions -from psycopg2 import DatabaseError, IntegrityError +from psycopg2 import DatabaseError, IntegrityError, OperationalError class PgManagerError(Exception): @@ -88,9 +88,12 @@ class ConnectionInfo: - def __init__(self, dsn, isolation_level=None, init_statement=None, keep_open=1): + def __init__(self, dsn, isolation_level=None, keep_alive=True, + init_statement=None, keep_open=1): + self.dsn = dsn self.isolation_level = isolation_level + self.keep_alive = keep_alive self.init_statement = init_statement self.keep_open = keep_open @@ -161,12 +164,15 @@ '''Set socket to keepalive mode. Must be called before any query.''' sock = socket.fromfd(self.fileno(), socket.AF_INET, socket.SOCK_STREAM) sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) - # Maximum keep-alive probes before asuming the connection is lost - sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5) - # Interval (in seconds) between keep-alive probes - sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 2) - # Maximum idle time (in seconds) before start sending keep-alive probes - sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 10) + try: + # Maximum keep-alive probes before asuming the connection is lost + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5) + # Interval (in seconds) between keep-alive probes + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 2) + # Maximum idle time (in seconds) before start sending keep-alive probes + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 10) + except socket.error: + pass class PgManager: @@ -180,16 +186,25 @@ for conn in tuple(self.conn_known.keys()): self.destroy_conn(conn) - def create_conn(self, name='default', isolation_level=None, dsn=None, **kw): - '''Create named connection.''' + def create_conn(self, name='default', isolation_level=None, keep_alive=True, dsn=None, **kw): + '''Create named connection. + + name -- name for connection (default is "default") + 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()]) + dsn = ' '.join([x[0]+'='+str(x[1]) for x in kw.items() if x[1] is not None]) isolation_level = self._normalize_isolation_level(isolation_level) - ci = ConnectionInfo(dsn, isolation_level) + ci = ConnectionInfo(dsn, isolation_level, keep_alive) self.conn_known[name] = ci self.conn_pool[name] = [] @@ -315,7 +330,8 @@ def _connect(self, ci): conn = psycopg2.connect(ci.dsn, connection_factory=Connection) - conn.keep_alive() + 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: