Perform logical replication after initial sync

tap_postgres/__init__.py

@@ -131,6 +131,8 @@ def sync_method_for_streams(streams, state, default_replication_method):
             # finishing previously interrupted full-table (first stage of logical replication)
             lookup[stream['tap_stream_id']] = 'logical_initial_interrupted'
             traditional_steams.append(stream)
+            # do any required logical replication after inital sync is complete
+            logical_streams.append(stream)
 
         # inconsistent state
         elif get_bookmark(state, stream['tap_stream_id'], 'xmin') and \
@@ -142,6 +144,8 @@ def sync_method_for_streams(streams, state, default_replication_method):
             # initial full-table phase of logical replication
             lookup[stream['tap_stream_id']] = 'logical_initial'
             traditional_steams.append(stream)
+            # do any required logical replication after inital sync is complete
+            logical_streams.append(stream)
 
         else:  # no xmin but we have an lsn
             # initial stage of logical replication(full-table) has been completed. moving onto pure logical replication

tests/test_full_table_interruption.py

@@ -1,5 +1,7 @@
 import unittest
+import unittest.mock
 import tap_postgres
+from tap_postgres.sync_strategies import logical_replication
 import tap_postgres.sync_strategies.full_table as full_table
 import tap_postgres.sync_strategies.common as pg_common
 import singer
@@ -45,9 +47,31 @@ def do_not_dump_catalog(catalog):
 tap_postgres.dump_catalog = do_not_dump_catalog
 full_table.UPDATE_BOOKMARK_PERIOD = 1
 
+
+@unittest.mock.patch('tap_postgres.sync_logical_streams', wraps=tap_postgres.sync_logical_streams)
 class LogicalInterruption(unittest.TestCase):
     maxDiff = None
 
+    @classmethod
+    def setUpClass(cls):
+        conn_config = get_test_connection_config()
+        slot_name = logical_replication.generate_replication_slot_name(
+            dbname=conn_config['dbname'], tap_id=conn_config['tap_id']
+        )
+        with get_test_connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(f"SELECT * FROM pg_create_logical_replication_slot('{slot_name}', 'wal2json')")
+
+    @classmethod
+    def tearDownClass(cls):
+        conn_config = get_test_connection_config()
+        slot_name = logical_replication.generate_replication_slot_name(
+            dbname=conn_config['dbname'], tap_id=conn_config['tap_id']
+        )
+        with get_test_connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(f"SELECT * FROM pg_drop_replication_slot('{slot_name}')")
+
     def setUp(self):
         table_spec_1 = {"columns": [{"name": "id", "type" : "serial",       "primary_key" : True},
                                     {"name" : 'name', "type": "character varying"},
@@ -62,12 +86,13 @@ def setUp(self):
         global CAUGHT_MESSAGES
         CAUGHT_MESSAGES.clear()
 
-    def test_catalog(self):
+    def test_catalog(self, mock_sync_logical_streams):
         singer.write_message = singer_write_message_no_cow
         pg_common.write_schema_message = singer_write_message_ok
 
         conn_config = get_test_connection_config()
         streams = tap_postgres.do_discovery(conn_config)
+
         cow_stream = [s for s in streams if s['table_name'] == 'COW'][0]
         self.assertIsNotNone(cow_stream)
         cow_stream = select_all_of_stream(cow_stream)
@@ -90,7 +115,7 @@ def test_catalog(self):
             insert_record(cur, 'COW', cow_rec)
 
         conn.close()
-           
+
         blew_up_on_cow = False
         state = {}
         #the initial phase of cows logical replication will be a full table.
@@ -102,6 +127,8 @@ def test_catalog(self):
 
         self.assertTrue(blew_up_on_cow)
 
+        mock_sync_logical_streams.assert_not_called()
+
         self.assertEqual(7, len(CAUGHT_MESSAGES))
 
         self.assertEqual(CAUGHT_MESSAGES[0]['type'], 'SCHEMA')
@@ -151,7 +178,9 @@ def test_catalog(self):
         CAUGHT_MESSAGES.clear()
         tap_postgres.do_sync(get_test_connection_config(), {'streams' : streams}, None, old_state)
 
-        self.assertEqual(8, len(CAUGHT_MESSAGES))
+        mock_sync_logical_streams.assert_called_once()
+
+        self.assertEqual(10, len(CAUGHT_MESSAGES))
 
         self.assertEqual(CAUGHT_MESSAGES[0]['type'], 'SCHEMA')
 
@@ -200,8 +229,16 @@ def test_catalog(self):
         self.assertEqual(CAUGHT_MESSAGES[7].value['bookmarks']['public-COW'].get('lsn'), end_lsn)
         self.assertEqual(CAUGHT_MESSAGES[7].value['bookmarks']['public-COW'].get('version'), new_version)
 
+        assert CAUGHT_MESSAGES[8]['type'] == 'SCHEMA'
+
+        assert isinstance(CAUGHT_MESSAGES[9], singer.messages.StateMessage)
+        assert CAUGHT_MESSAGES[9].value['bookmarks']['public-COW'].get('xmin') is None
+        assert CAUGHT_MESSAGES[9].value['bookmarks']['public-COW'].get('lsn') == end_lsn
+        assert CAUGHT_MESSAGES[9].value['bookmarks']['public-COW'].get('version') == new_version
+
 class FullTableInterruption(unittest.TestCase):
     maxDiff = None
+
     def setUp(self):
         table_spec_1 = {"columns": [{"name": "id", "type" : "serial",       "primary_key" : True},
                                     {"name" : 'name', "type": "character varying"},
@@ -238,7 +275,7 @@ def test_catalog(self):
 
         conn = get_test_connection()
         conn.autocommit = True
-        
+
         with conn.cursor() as cur:
             cow_rec = {'name': 'betty', 'colour': 'blue'}
             insert_record(cur, 'COW', {'name': 'betty', 'colour': 'blue'})
@@ -256,7 +293,7 @@ def test_catalog(self):
 
         state = {}
         blew_up_on_cow = False
-        
+
         #this will sync the CHICKEN but then blow up on the COW
         try:
             tap_postgres.do_sync(get_test_connection_config(), {'streams' : streams}, None, state)

tests/utils.py

@@ -18,7 +18,11 @@ def get_test_connection_config(target_db='postgres'):
     if len(missing_envs) != 0:
         raise Exception("set TAP_POSTGRES_HOST, TAP_POSTGRES_USER, TAP_POSTGRES_PASSWORD, TAP_POSTGRES_PORT")
 
-    conn_config = {'host': os.environ.get('TAP_POSTGRES_HOST'),
+    conn_config = {'tap_id': 'test-postgres',
+                   'max_run_seconds': 5,
+                   'break_at_end_lsn': True,
+                   'logical_poll_total_seconds': 2,
+                   'host': os.environ.get('TAP_POSTGRES_HOST'),
                    'user': os.environ.get('TAP_POSTGRES_USER'),
                    'password': os.environ.get('TAP_POSTGRES_PASSWORD'),
                    'port': os.environ.get('TAP_POSTGRES_PORT'),