[AP-741] send new schema message when wal payload has new columns - LOG_BASED (#59)

* rebuild schema and send schema singer message when wal message s has new columns

* add step to check dependency conflict

Author: Samira-El

.circleci/config.yml

@@ -32,9 +32,15 @@ jobs:
           command: |
             python3 -m venv ./virtualenvs/tap-postgres
             . ./virtualenvs/tap-postgres/bin/activate
-            pip install --upgrade pip
+            pip install --upgrade pip setuptools
             pip install -e .[test]
 
+      - run:
+          name: Check dependencies for conflict
+          command: |
+            . ./virtualenvs/tap-postgres/bin/activate
+            pip check && echo "No conflicts" || exit 1
+
       - run:
           name: 'Pylinting'
           command: |

tap_postgres/__init__.py

@@ -201,7 +201,8 @@ def sync_logical_streams(conn_config, logical_streams, state, end_lsn, state_fil
         LOGGER.info("Pure Logical Replication upto lsn %s for (%s)", end_lsn,
                     [s['tap_stream_id'] for s in logical_streams])
 
-        logical_streams = [logical_replication.add_automatic_properties(s, conn_config) for s in logical_streams]
+        logical_streams = [logical_replication.add_automatic_properties(
+            s, conn_config.get('debug_lsn', False)) for s in logical_streams]
 
         # Remove LOG_BASED stream bookmarks from state if it has been de-selected
         # This is to avoid sending very old starting and flushing positions to source

tap_postgres/sync_strategies/logical_replication.py

@@ -2,17 +2,19 @@
 import pytz
 import decimal
 import psycopg2
-from psycopg2 import sql
 import copy
 import json
 import singer
 import singer.metadata as metadata
+
+from psycopg2 import sql
 from singer import utils, get_bookmark
 from dateutil.parser import parse
 from functools import reduce
 
 import tap_postgres.db as post_db
 import tap_postgres.sync_strategies.common as sync_common
+from tap_postgres.stream_utils import refresh_streams_schema
 
 LOGGER = singer.get_logger('tap_postgres')
 
@@ -22,6 +24,9 @@
 class ReplicationSlotNotFoundError(Exception):
     """Custom exception when replication slot not found"""
 
+class UnsupportedPayloadKindError(Exception):
+    """Custom exception when waljson payload is not insert, update nor delete"""
+
 
 # pylint: disable=invalid-name,missing-function-docstring,too-many-branches,too-many-statements,too-many-arguments
 def get_pg_version(conn_info):
@@ -98,9 +103,10 @@ def fetch_current_lsn(conn_config):
             return lsn_to_int(current_lsn)
 
 
-def add_automatic_properties(stream, conn_config):
+def add_automatic_properties(stream, debug_lsn: bool = False):
     stream['schema']['properties']['_sdc_deleted_at'] = {'type': ['null', 'string'], 'format': 'date-time'}
-    if conn_config.get('debug_lsn'):
+
+    if debug_lsn:
         LOGGER.debug('debug_lsn is ON')
         stream['schema']['properties']['_sdc_lsn'] = {'type': ['null', 'string']}
     else:
@@ -295,7 +301,7 @@ def row_to_singer_message(stream, row, version, columns, time_extracted, md_map,
 
 
 # pylint: disable=unused-argument,too-many-locals
-def consume_message(streams, state, msg, time_extracted, conn_info, end_lsn):
+def consume_message(streams, state, msg, time_extracted, conn_info):
     # Strip leading comma generated by write-in-chunks and parse valid JSON
     try:
         payload = json.loads(msg.payload.lstrip(','))
@@ -304,56 +310,58 @@ def consume_message(streams, state, msg, time_extracted, conn_info, end_lsn):
 
     lsn = msg.data_start
 
-    streams_lookup = {}
-    for s in streams:
-        streams_lookup[s['tap_stream_id']] = s
+    streams_lookup = {s['tap_stream_id']: s for s in streams}
 
     tap_stream_id = post_db.compute_tap_stream_id(payload['schema'], payload['table'])
     if streams_lookup.get(tap_stream_id) is None:
         return state
 
     target_stream = streams_lookup[tap_stream_id]
+
+    if payload['kind'] not in {'insert', 'update', 'delete'}:
+        raise UnsupportedPayloadKindError("unrecognized replication operation: {}".format(payload['kind']))
+
+    # Get the additional fields in payload that are not in schema properties:
+    # only inserts and updates have the list of columns that can be used to detect any different in columns
+    diff = set()
+    if payload['kind'] in {'insert', 'update'}:
+        diff = set(payload['columnnames']).difference(target_stream['schema']['properties'].keys())
+
+    # if there is new columns in the payload that are not in the schema properties then refresh the stream schema
+    if diff:
+        LOGGER.info('Detected new columns "%s", refreshing schema of stream %s', diff, target_stream['stream'])
+        # encountered a column that is not in the schema
+        # refresh the stream schema and metadata by running discovery
+        refresh_streams_schema(conn_info, [target_stream])
+
+        # add the automatic properties back to the stream
+        add_automatic_properties(target_stream, conn_info.get('debug_lsn', False))
+
+        # publish new schema
+        sync_common.send_schema_message(target_stream, ['lsn'])
+
     stream_version = get_stream_version(target_stream['tap_stream_id'], state)
     stream_md_map = metadata.to_map(target_stream['metadata'])
 
-    desired_columns = [c for c in target_stream['schema']['properties'].keys() if sync_common.should_sync_column(
-        stream_md_map, c)]
+    desired_columns = {c for c in target_stream['schema']['properties'].keys() if sync_common.should_sync_column(
+        stream_md_map, c)}
 
-    if payload['kind'] == 'insert':
+    if payload['kind'] in {'insert', 'update'}:
         col_names = []
         col_vals = []
-        for idx, col in enumerate(payload['columnnames']):
-            if col in set(desired_columns):
-                col_names.append(col)
-                col_vals.append(payload['columnvalues'][idx])
 
-        col_names = col_names + ['_sdc_deleted_at']
-        col_vals = col_vals + [None]
-        if conn_info.get('debug_lsn'):
-            col_names = col_names + ['_sdc_lsn']
-            col_vals = col_vals + [str(lsn)]
-        record_message = row_to_singer_message(target_stream,
-                                               col_vals,
-                                               stream_version,
-                                               col_names,
-                                               time_extracted,
-                                               stream_md_map,
-                                               conn_info)
-
-    elif payload['kind'] == 'update':
-        col_names = []
-        col_vals = []
         for idx, col in enumerate(payload['columnnames']):
-            if col in set(desired_columns):
+            if col in desired_columns:
                 col_names.append(col)
                 col_vals.append(payload['columnvalues'][idx])
 
-        col_names = col_names + ['_sdc_deleted_at']
-        col_vals = col_vals + [None]
+        col_names.append('_sdc_deleted_at')
+        col_vals.append(None)
 
         if conn_info.get('debug_lsn'):
-            col_vals = col_vals + [str(lsn)]
-            col_names = col_names + ['_sdc_lsn']
+            col_names.append('_sdc_lsn')
+            col_vals.append(str(lsn))
+
         record_message = row_to_singer_message(target_stream,
                                                col_vals,
                                                stream_version,
@@ -366,15 +374,17 @@ def consume_message(streams, state, msg, time_extracted, conn_info, end_lsn):
         col_names = []
         col_vals = []
         for idx, col in enumerate(payload['oldkeys']['keynames']):
-            if col in set(desired_columns):
+            if col in desired_columns:
                 col_names.append(col)
                 col_vals.append(payload['oldkeys']['keyvalues'][idx])
 
-        col_names = col_names + ['_sdc_deleted_at']
-        col_vals = col_vals + [singer.utils.strftime(time_extracted)]
+        col_names.append('_sdc_deleted_at')
+        col_vals.append(singer.utils.strftime(time_extracted))
+
         if conn_info.get('debug_lsn'):
-            col_vals = col_vals + [str(lsn)]
-            col_names = col_names + ['_sdc_lsn']
+            col_names.append('_sdc_lsn')
+            col_vals.append(str(lsn))
+
         record_message = row_to_singer_message(target_stream,
                                                col_vals,
                                                stream_version,
@@ -383,9 +393,6 @@ def consume_message(streams, state, msg, time_extracted, conn_info, end_lsn):
                                                stream_md_map,
                                                conn_info)
 
-    else:
-        raise Exception("unrecognized replication operation: {}".format(payload['kind']))
-
     singer.write_message(record_message)
     state = singer.write_bookmark(state, target_stream['tap_stream_id'], 'lsn', lsn)
 
@@ -544,7 +551,7 @@ def sync_tables(conn_info, logical_streams, state, end_lsn, state_file):
                 LOGGER.info('Breaking - reached max_run_seconds of %i', max_run_seconds)
                 break
 
-            state = consume_message(logical_streams, state, msg, time_extracted, conn_info, end_lsn)
+            state = consume_message(logical_streams, state, msg, time_extracted, conn_info)
 
             # When using wal2json with write-in-chunks, multiple messages can have the same lsn
             # This is to ensure we only flush to lsn that has completed entirely

tests/test_logical_replication.py

@@ -1,6 +1,11 @@
+import json
 import unittest
 
+from collections import namedtuple
+from unittest.mock import patch
+
 from tap_postgres.sync_strategies import logical_replication
+from tap_postgres.sync_strategies.logical_replication import UnsupportedPayloadKindError
 
 
 class PostgresCurReplicationSlotMock:
@@ -31,7 +36,7 @@ class TestLogicalReplication(unittest.TestCase):
     maxDiff = None
 
     def setUp(self):
-        pass
+        self.WalMessage = namedtuple('WalMessage', ['payload', 'data_start'])
 
     def test_streams_to_wal2json_tables(self):
         """Validate if table names are escaped to wal2json format"""
@@ -115,3 +120,126 @@ def test_locate_replication_slot_by_cur(self):
                                                                                  'some_db',
                                                                                  'some_tap'),
                               'pipelinewise_some_db_some_tap')
+
+    def test_consume_with_message_payload_is_not_json_expect_same_state(self):
+
+        output = logical_replication.consume_message([],
+                                            {},
+                                            self.WalMessage(payload='this is an invalid json message', data_start=None),
+                                            None,
+                                            {}
+                                            )
+        self.assertDictEqual({}, output)
+
+    def test_consume_with_message_stream_in_payload_is_not_selected_expect_same_state(self):
+        output = logical_replication.consume_message(
+            [{'tap_stream_id': 'myschema-mytable'}],
+            {},
+            self.WalMessage(payload='{"schema": "myschema", "table": "notmytable"}',
+                            data_start='some lsn'),
+            None,
+            {}
+        )
+
+        self.assertDictEqual({}, output)
+
+    def test_consume_with_payload_kind_is_not_supported_expect_exception(self):
+
+        with self.assertRaises(UnsupportedPayloadKindError):
+            logical_replication.consume_message(
+                [{'tap_stream_id': 'myschema-mytable'}],
+                {},
+                self.WalMessage(payload='{"kind":"truncate", "schema": "myschema", "table": "mytable"}',
+                                data_start='some lsn'),
+                None,
+                {}
+            )
+    @patch('tap_postgres.logical_replication.singer.write_message')
+    @patch('tap_postgres.logical_replication.sync_common.send_schema_message')
+    @patch('tap_postgres.logical_replication.refresh_streams_schema')
+    def test_consume_message_with_new_column_in_payload_will_refresh_schema(self,
+                                                                            refresh_schema_mock,
+                                                                            send_schema_mock,
+                                                                            write_message_mock):
+        streams = [
+                {
+                    'tap_stream_id': 'myschema-mytable',
+                    'stream': 'mytable',
+                    'schema': {
+                        'properties': {
+                            'id': {},
+                            'date_created': {}
+                        }
+                    },
+                    'metadata': [
+                        {
+                            'breadcrumb': [],
+                            'metadata': {
+                                'is-view': False,
+                                'table-key-properties': ['id'],
+                                'schema-name': 'myschema'
+                            }
+                        },
+                        {
+                            "breadcrumb": [
+                                "properties",
+                                "id"
+                            ],
+                            "metadata": {
+                                "sql-datatype": "integer",
+                                "inclusion": "automatic",
+                            }
+                        },
+                        {
+                            "breadcrumb": [
+                                "properties",
+                                "date_created"
+                            ],
+                            "metadata": {
+                                "sql-datatype": "datetime",
+                                "inclusion": "available",
+                                "selected": True
+                            }
+                        }
+                    ],
+                }
+            ]
+
+        return_v = logical_replication.consume_message(
+            streams,
+            {
+                'bookmarks': {
+                    "myschema-mytable": {
+                        "last_replication_method": "LOG_BASED",
+                        "lsn": None,
+                        "version": 1000,
+                        "xmin": None
+                    }
+                }
+            },
+            self.WalMessage(payload='{"kind": "insert", '
+                                    '"schema": "myschema", '
+                                    '"table": "mytable",'
+                                    '"columnnames": ["id", "date_created", "new_col"],'
+                                    '"columnnames": [1, null, "some random text"]'
+                                    '}',
+                            data_start='some lsn'),
+            None,
+            {}
+        )
+
+        self.assertDictEqual(return_v,
+                          {
+                              'bookmarks': {
+                                  "myschema-mytable": {
+                                      "last_replication_method": "LOG_BASED",
+                                      "lsn": "some lsn",
+                                      "version": 1000,
+                                      "xmin": None
+                                  }
+                              }
+                          })
+
+        refresh_schema_mock.assert_called_once_with({}, [streams[0]])
+        send_schema_mock.assert_called_once()
+        write_message_mock.assert_called_once()