Feature/replication keys

elx/catalog.py

@@ -158,3 +158,45 @@ def select(self, streams: Optional[List[str]]) -> "Catalog":
             )
 
         return catalog
+
+    def set_replication_keys(self, replication_keys: Optional[dict]) -> "Catalog":
+        """
+        Set the replication key for streams and updates the catalog.
+
+        Args:
+            keys (Optional[dict]): Dictionary stream replication_key value-pairs.
+            E.g. {"stream_one": "updated_at", "stream_two": "modified_at"}
+
+        Returns:
+            Catalog: A new catalog with updated replication settings.
+        """
+        # Make a copy of the existing catalog.
+        catalog = self.copy(deep=True)
+
+        # Loop over the streams
+        for stream in catalog.streams:
+            # If the stream is specified in `replication_keys` dictionary
+            if stream.tap_stream_id in replication_keys:
+                # Set the replication method to INCREMENTAL
+                stream.replication_method = "INCREMENTAL"
+
+                # Update the replication key value
+                stream.replication_key = replication_keys[stream.tap_stream_id]
+
+                # Add the replication key to stream metadata
+                stream.upsert_metadata(
+                    breadcrumb=[],
+                    metadata={
+                        "valid-replication-keys": [
+                            replication_keys[stream.tap_stream_id]
+                        ]
+                    },
+                )
+
+                # Set inclusion of replication property metadata to `automatic`
+                stream.upsert_metadata(
+                    breadcrumb=["properties", replication_keys[stream.tap_stream_id]],
+                    metadata={"inclusion": "automatic"},
+                )
+
+        return catalog

elx/tap.py

@@ -17,9 +17,11 @@ def __init__(
         executable: str | None = None,
         config: dict = {},
         deselected: List[str] = None,
+        replication_keys: dict = {},
     ):
         super().__init__(spec, executable, config)
         self.deselected = deselected
+        self.replication_keys = replication_keys
 
     def discover(self, config_path: Path) -> dict:
         """
@@ -45,7 +47,11 @@ def catalog(self) -> Catalog:
         with json_temp_file(self.config) as config_path:
             catalog = self.discover(config_path)
             catalog = Catalog(**catalog)
-            return catalog.deselect(patterns=self.deselected)
+            catalog = catalog.deselect(patterns=self.deselected)
+            catalog = catalog.set_replication_keys(
+                replication_keys=self.replication_keys
+            )
+            return catalog
 
     @contextlib.asynccontextmanager
     @require_install

tests/fixtures/tap.py

@@ -12,4 +12,5 @@ def tap() -> Generator[Tap, None, None]:
         executable="tap-mock-fixture",
         spec="git+https://github.com/quantile-taps/tap-mock-fixture.git",
         config={},
+        replication_keys={"users": "updated_at"},
     )

tests/test_elx/test_catalog.py

@@ -178,3 +178,11 @@ def test_catalog_valid_replication_keys(tap: Tap):
 
     # Checks that value of `valid-replication-keys` equals to the replication-key
     assert replication_keys == [DEFAULT_CATALOG["streams"][1]["replication_key"]]
+
+
+def test_catalog_set_stream_replication_key(tap: Tap):
+    """If we define a replication key, the catalog should be updated."""
+    catalog = tap.catalog
+
+    assert catalog.streams[1].replication_method == "INCREMENTAL"
+    assert catalog.streams[1].replication_key == "updated_at"