[DOP-22132] Implement increment for transfers with DB source

Author: Ilyas Gasanov

syncmaster/worker/controller.py

@@ -42,6 +42,7 @@
 )
 from syncmaster.dto.transfers_strategy import Strategy
 from syncmaster.exceptions.connection import ConnectionTypeNotRecognizedError
+from syncmaster.schemas.v1.connection_types import FILE_CONNECTION_TYPES
 from syncmaster.worker.handlers.base import Handler
 from syncmaster.worker.handlers.db.clickhouse import ClickhouseHandler
 from syncmaster.worker.handlers.db.hive import HiveHandler
@@ -242,11 +243,15 @@ def _perform_incremental_transfer(self) -> None:
         ).force_create_namespace() as hwm_store:
 
             with IncrementalStrategy():
+                if self.source_handler.connection_dto.type in FILE_CONNECTION_TYPES:
+                    hwm_name_suffix = self.source_handler.transfer_dto.directory_path
+                else:
+                    hwm_name_suffix = self.source_handler.transfer_dto.table_name
                 hwm_name = "_".join(
                     [
                         str(self.source_handler.transfer_dto.id),
                         self.source_handler.connection_dto.type,
-                        self.source_handler.transfer_dto.directory_path,
+                        hwm_name_suffix,
                     ],
                 )
                 hwm = hwm_store.get_hwm(hwm_name)

syncmaster/worker/handlers/db/base.py

@@ -36,11 +36,20 @@ class DBHandler(Handler):
     }
 
     def read(self) -> DataFrame:
+        reader_params = {}
+        if self.transfer_dto.strategy.type == "incremental":
+            hwm_name = f"{self.transfer_dto.id}_{self.connection_dto.type}_{self.transfer_dto.table_name}"
+            reader_params["hwm"] = DBReader.AutoDetectHWM(
+                name=hwm_name,
+                expression=self.transfer_dto.strategy.increment_by,
+            )
+
         reader = DBReader(
             connection=self.connection,
             table=self.transfer_dto.table_name,
             where=self._get_rows_filter_expression(),
             columns=self._get_columns_filter_expressions(),
+            **reader_params,
         )
         return reader.run()
 

syncmaster/worker/handlers/db/postgres.py

@@ -9,6 +9,7 @@
 
 from syncmaster.dto.connections import PostgresConnectionDTO
 from syncmaster.dto.transfers import PostgresTransferDTO
+from syncmaster.dto.transfers_strategy import IncrementalStrategy
 from syncmaster.worker.handlers.db.base import DBHandler
 
 if TYPE_CHECKING:
@@ -36,6 +37,11 @@ def connect(self, spark: SparkSession):
             spark=spark,
         ).check()
 
+    def read(self) -> DataFrame:
+        if isinstance(self.transfer_dto.strategy, IncrementalStrategy):
+            self.transfer_dto.strategy.increment_by = DBHandler._quote_field(self.transfer_dto.strategy.increment_by)
+        return super().read()
+
     def _normalize_column_names(self, df: DataFrame) -> DataFrame:
         for column_name in df.columns:
             df = df.withColumnRenamed(column_name, column_name.lower())

tests/test_integration/test_run_transfer/connection_fixtures/__init__.py

@@ -123,6 +123,7 @@
     full_strategy,
     incremental_strategy_by_file_modified_since,
     incremental_strategy_by_file_name,
+    incremental_strategy_by_number_column,
 )
 from tests.test_integration.test_run_transfer.connection_fixtures.webdav_fixtures import (
     prepare_webdav,

tests/test_integration/test_run_transfer/connection_fixtures/filters_fixtures.py

@@ -74,7 +74,7 @@ def dataframe_columns_filter_transformations(source_type: str):
                 },
                 {
                     "type": "include",
-                    "field": "PHONE_NUMBER",
+                    "field": "NUMBER",
                 },
                 {
                     "type": "rename",
@@ -83,7 +83,7 @@ def dataframe_columns_filter_transformations(source_type: str):
                 },
                 {
                     "type": "cast",
-                    "field": "NUMBER",
+                    "field": "BIRTH_DATE",
                     "as_type": string_types_per_source_type[source_type],
                 },
                 {
@@ -100,9 +100,9 @@ def expected_dataframe_columns_filter():
     return lambda df, source_type: df.selectExpr(
         "ID",
         "REGION",
-        "PHONE_NUMBER",
+        "NUMBER",
         "REGION AS REGION2",
-        "CAST(NUMBER AS STRING) AS NUMBER",
+        "CAST(BIRTH_DATE AS STRING) AS BIRTH_DATE",
         "REGISTERED_AT",
     )
 

tests/test_integration/test_run_transfer/connection_fixtures/strategy_fixtures.py

@@ -22,3 +22,11 @@ def incremental_strategy_by_file_name():
         "type": "incremental",
         "increment_by": "file_name",
     }
+
+
+@pytest.fixture
+def incremental_strategy_by_number_column():
+    return {
+        "type": "incremental",
+        "increment_by": "NUMBER",
+    }