From ce267d9f1e07c531ffa6fb6844e9a1ce9e2b465f Mon Sep 17 00:00:00 2001
From: Alexander Kukushkin <wmd@mail.ru>
Date: Sun, 24 May 2015 09:01:23 +0200
Subject: [PATCH] query method will retry only in case of communication error

---
 helpers/ha.py            |  4 ++--
 helpers/postgresql.py    | 15 +++++++++++----
 tests/test_postgresql.py | 12 ++++++++++--
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/helpers/ha.py b/helpers/ha.py
index 74c6a892..773c9b57 100644
--- a/helpers/ha.py
+++ b/helpers/ha.py
@@ -1,7 +1,7 @@
 import logging
 
 from helpers.errors import EtcdError, HealthiestMemberError
-from psycopg2 import OperationalError
+from psycopg2 import InterfaceError, OperationalError
 
 logger = logging.getLogger(__name__)
 
@@ -92,7 +92,7 @@ class Ha:
             if self.state_handler.is_leader():
                 self.state_handler.demote(None)
                 return 'demoted self because etcd is not accessible and i was a leader'
-        except OperationalError:
+        except (InterfaceError, OperationalError):
             logger.error('Error communicating with Postgresql.  Will try again')
         except HealthiestMemberError:
             logger.error('failed to determine healthiest member fromt etcd')
diff --git a/helpers/postgresql.py b/helpers/postgresql.py
index 11be38c0..720d0775 100644
--- a/helpers/postgresql.py
+++ b/helpers/postgresql.py
@@ -54,13 +54,13 @@ class Postgresql:
         return self.listen_addresses.split(',')[0].strip() + ':' + self.port
 
     def connection(self):
-        if not self._connection or self._connection.closed:
+        if not self._connection or self._connection.closed != 0:
             self._connection = psycopg2.connect('postgres://{}/postgres'.format(self.local_address))
             self._connection.autocommit = True
         return self._connection
 
     def _cursor(self):
-        if not self._cursor_holder or self._cursor_holder.closed != 0:
+        if not self._cursor_holder or self._cursor_holder.closed:
             self._cursor_holder = self.connection().cursor()
         return self._cursor_holder
 
@@ -71,15 +71,22 @@ class Postgresql:
     def query(self, sql, *params):
         max_attempts = 0
         while True:
+            ex = None
             try:
                 cursor = self._cursor()
                 cursor.execute(sql, params)
                 return cursor
-            except (psycopg2.OperationalError, psycopg2.InterfaceError) as e:
+            except psycopg2.InterfaceError as e:
+                ex = e
+            except psycopg2.OperationalError as e:
+                if self._connection and self._connection.closed == 0:
+                    raise e
+                ex = e
+            if ex:
                 self.disconnect()
                 max_attempts += 1
                 if max_attempts >= 3:
-                    raise e
+                    raise ex
                 time.sleep(5)
 
     def data_directory_empty(self):
diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
index 5442e58c..e0a21645 100644
--- a/tests/test_postgresql.py
+++ b/tests/test_postgresql.py
@@ -22,13 +22,15 @@ def xlog_position():
 class MockCursor:
 
     def __init__(self):
-        self.closed = 0
+        self.closed = False
         self.current = 0
         self.results = []
 
     def execute(self, sql, *params):
         if sql.startswith('blabla'):
             raise psycopg2.OperationalError()
+        elif sql.startswith('InterfaceError'):
+            raise psycopg2.InterfaceError()
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla',), ('foobar',)]
         elif sql.startswith('SELECT pg_current_xlog_location()'):
@@ -61,7 +63,7 @@ class MockConnect:
 
     def __init__(self):
         self.autocommit = False
-        self.closed = False
+        self.closed = 0
 
     def cursor(self):
         return MockCursor()
@@ -128,6 +130,12 @@ class TestPostgresql(unittest.TestCase):
 
     def test_query(self):
         self.p.query('select 1')
+        self.assertRaises(psycopg2.InterfaceError, self.p.query, 'InterfaceError')
+        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
+        self.p._connection.closed = 2
+        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
+        self.p._connection.closed = 2
+        self.p.disconnect = false
         self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
 
     def test_is_healthiest_node(self):