Some minor refactoring of logical replication

zilder · zilder · commit b1cba73cf385 · 2018-03-22T15:25:53.000+03:00
diff --git a/testgres/node.py b/testgres/node.py
@@ -847,8 +847,7 @@ def poll_query_until(self,
                          expected=True,
                          commit=True,
                          raise_programming_error=True,
-                         raise_internal_error=True,
-                         zero_rows_is_ok=False):
+                         raise_internal_error=True):
         """
         Run a query once per second until it returns 'expected'.
         Query should return a single value (1 row, 1 column).
@@ -884,18 +883,14 @@ def poll_query_until(self,
                 if res is None:
                     raise QueryException('Query returned None', query)
 
-                if len(res) == 0:
-                    if zero_rows_is_ok:
-                        time.sleep(sleep_time)
-                        attempts += 1
-                        continue
-                    else:
-                        raise QueryException('Query returned 0 rows', query)
-
-                if len(res[0]) == 0:
-                    raise QueryException('Query returned 0 columns', query)
-
-                if res[0][0] == expected:
+                # result set is not empty
+                if len(res):
+                    if len(res[0]) == 0:
+                        raise QueryException('Query returned 0 columns', query)
+                    if res[0][0] == expected:
+                        return    # done
+                # empty result set is considered as None
+                elif expected is None:
                     return    # done
 
             except ProgrammingError as e:
diff --git a/testgres/pubsub.py b/testgres/pubsub.py
@@ -138,24 +138,17 @@ def catchup(self, username=None):
         Args:
             username: remote node's user name
         """
-        if pg_version_ge('10'):
-            query = (
-                "select pg_current_wal_lsn() - replay_lsn = 0 "
-                "from pg_stat_replication where application_name = '{}'"
-            ).format(self.name)
-        else:
-            query = (
-                "select pg_current_xlog_location() - replay_location = 0 "
-                "from pg_stat_replication where application_name = '{}'"
-            ).format(self.name)
+        query = (
+            "select pg_current_wal_lsn() - replay_lsn = 0 "
+            "from pg_stat_replication where application_name = '{}'"
+        ).format(self.name)
 
         try:
             # wait until this LSN reaches subscriber
             self.pub.node.poll_query_until(
                 query=query,
                 dbname=self.pub.dbname,
                 username=username or self.pub.username,
-                max_attempts=60,
-                zero_rows_is_ok=True)  # statistics may have not updated yet
+                max_attempts=60)
         except Exception as e:
             raise_from(CatchUpException("Failed to catch up", query), e)
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -495,11 +495,6 @@ def test_poll_query_until(self):
 
             self.assertTrue(end_time - start_time >= 5)
 
-            # check 0 rows
-            with self.assertRaises(QueryException):
-                node.poll_query_until(
-                    query='select * from pg_class where true = false')
-
             # check 0 columns
             with self.assertRaises(QueryException):
                 node.poll_query_until(query='select from pg_class limit 1')
@@ -512,6 +507,10 @@ def test_poll_query_until(self):
             node.poll_query_until(
                 query='create table def()', expected=None)    # returns nothing
 
+            # check 0 rows equivalent to expected=None
+            node.poll_query_until(
+                query='select * from pg_class where true = false', expected=None)
+
             # check arbitrary expected value, fail
             with self.assertRaises(TimeoutException):
                 node.poll_query_until(
