# HG changeset patch
# User Radek Brich <radek.brich@devl.cz>
# Date 1313500373 -7200
# Node ID 2911935c524d392b55bb4f5a728c74d65efa5fc2
# Parent  685b20d2d3ab51a612130a82a9922947c3220dfd
pgmanager: Add keep_alive support.

diff -r 685b20d2d3ab -r 2911935c524d tools/pgmanager.py
--- a/tools/pgmanager.py	Fri Aug 12 14:39:49 2011 +0200
+++ b/tools/pgmanager.py	Tue Aug 16 15:12:53 2011 +0200
@@ -70,6 +70,7 @@
 import logging
 import threading
 import select
+import socket
 
 import psycopg2
 import psycopg2.extensions
@@ -153,6 +154,17 @@
         else:
             return super(Connection, self).cursor(name, cursor_factory=Cursor)
 
+    def keep_alive(self):
+        '''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)
+
 
 class PgManager:
 
@@ -218,13 +230,7 @@
 
             if conn is None:
                 ci = self.conn_known[name]
-                conn = psycopg2.connect(ci.dsn, connection_factory=Connection)
-                if not ci.isolation_level is None:
-                    conn.set_isolation_level(ci.isolation_level)
-                if ci.init_statement:
-                    curs = conn.cursor()
-                    curs.execute(ci.init_statement)
-                    curs.close()
+                conn = self._connect(ci)
         finally:
             self.lock.release()
         return conn
@@ -304,6 +310,17 @@
                 conn.notifies.pop()
             self.put_conn(conn, name)
 
+    def _connect(self, ci):
+        conn = psycopg2.connect(ci.dsn, connection_factory=Connection)
+        conn.keep_alive()
+        if not ci.isolation_level is None:
+            conn.set_isolation_level(ci.isolation_level)
+        if ci.init_statement:
+            curs = conn.cursor()
+            curs.execute(ci.init_statement)
+            curs.close()
+        return conn
+
     def _normalize_isolation_level(self, level):
         if type(level) == str:
             if level.lower() == 'autocommit':