[DOP-22133] Implement increment for transfers with file sources

Author: Ilyas Gasanov

.env.docker

@@ -38,6 +38,13 @@ SYNCMASTER__AUTH__ACCESS_TOKEN__SECRET_KEY=generate_another_random_string
 # Scheduler options
 SYNCMASTER__SCHEDULER__TRANSFER_FETCHING_TIMEOUT_SECONDS=200
 
+# Horizon
+SYNCMASTER__HORIZON__URL=http://horizon:8000
+SYNCMASTER__HORIZON__NAMESPACE=syncmaster_namespace
+SYNCMASTER__HORIZON__USER=admin
+SYNCMASTER__HORIZON__PASSWORD=password
+
+# Tests-only
 TEST_S3_HOST_FOR_CONFTEST=test-s3
 TEST_S3_PORT_FOR_CONFTEST=9000
 TEST_S3_HOST_FOR_WORKER=test-s3

docker-compose.test.yml

@@ -131,6 +131,49 @@ services:
         condition: service_healthy
     profiles: [worker, scheduler, s3, oracle, hdfs, hive, clickhouse, mysql, mssql, sftp, ftp, ftps, samba, webdav, all]
 
+  horizon:
+    image: mtsrus/horizon-backend:develop
+    restart: unless-stopped
+    env_file: .env.docker
+    environment:
+      # list here usernames which should be assigned SUPERADMIN role on application start
+      HORIZON__ENTRYPOINT__ADMIN_USERS: admin
+      HORIZON__DATABASE__URL: postgresql+asyncpg://horizon:changeme@horizon-db:5432/horizon
+      HORIZON__AUTH__ACCESS_TOKEN__SECRET_KEY: generate_another_random_string
+      HORIZON__AUTH__PROVIDER: horizon.backend.providers.auth.dummy.DummyAuthProvider
+      HORIZON__SERVER__LOGGING__PRESET: colored
+      # PROMETHEUS_MULTIPROC_DIR is required for multiple workers, see:
+      # https://prometheus.github.io/client_python/multiprocess/
+      PROMETHEUS_MULTIPROC_DIR: /tmp/prometheus-metrics
+    # tmpfs dir is cleaned up each container restart
+    tmpfs:
+      - /tmp/prometheus-metrics
+    ports:
+      - 8020:8000
+    depends_on:
+      horizon-db:
+        condition: service_healthy
+    profiles: [horizon, s3, oracle, hdfs, hive, clickhouse, mysql, mssql, sftp, ftp, ftps, samba, webdav, all]
+
+  horizon-db:
+    image: postgres
+    restart: unless-stopped
+    environment:
+      TZ: UTC
+      POSTGRES_DB: horizon
+      POSTGRES_USER: horizon
+      POSTGRES_PASSWORD: changeme
+    ports:
+      - 5434:5432
+    volumes:
+      - horizon_test_data:/var/lib/postgresql/data
+    healthcheck:
+      test: pg_isready
+      start_period: 5s
+      interval: 30s
+      timeout: 5s
+      retries: 3
+
   test-postgres:
     image: postgres
     restart: unless-stopped
@@ -366,4 +409,5 @@ services:
 volumes:
   postgres_test_data:
   rabbitmq_test_data:
+  horizon_test_data:
   keycloak_data:

docs/changelog/next_release/209.feature.rst

@@ -0,0 +1 @@
+Implement increment for transfers with file sources

poetry.lock

@@ -71,6 +71,7 @@ apscheduler = { version = "^3.10.4", optional = true }
 starlette-exporter = {version = "^0.23.0", optional = true}
 itsdangerous = {version = "*", optional = true}
 python-keycloak = {version = ">=4.7,<6.0", optional = true}
+horizon-hwm-store = {version = ">=1.0.2", optional = true }
 
 [tool.poetry.extras]
 server = [
@@ -109,6 +110,7 @@ worker = [
     "coloredlogs",
     "python-json-logger",
     "pyyaml",
+    "horizon-hwm-store",
 ]
 
 scheduler = [

pyproject.toml

@@ -6,6 +6,8 @@
 
 from onetl.file.format import CSV, JSON, ORC, XML, Excel, JSONLine, Parquet
 
+from syncmaster.dto.transfers_strategy import FullStrategy, IncrementalStrategy
+
 
 @dataclass
 class TransferDTO:
@@ -14,14 +16,18 @@ class TransferDTO:
 
 @dataclass
 class DBTransferDTO(TransferDTO):
+    id: int
     table_name: str
+    strategy: FullStrategy | IncrementalStrategy
     transformations: list[dict] | None = None
 
 
 @dataclass
 class FileTransferDTO(TransferDTO):
+    id: int
     directory_path: str
     file_format: CSV | JSONLine | JSON | Excel | XML | ORC | Parquet
+    strategy: FullStrategy | IncrementalStrategy
     options: dict
     file_name_template: str | None = None
     df_schema: dict | None = None
@@ -43,7 +49,7 @@ def __post_init__(self):
         if isinstance(self.df_schema, str):
             self.df_schema = json.loads(self.df_schema)
 
-        self.options.setdefault("if_exists", "replace_entire_directory")  # TODO: use "append" for incremental strategy
+        self.options.setdefault("if_exists", "replace_entire_directory")
 
     def _get_file_format(self, file_format: dict) -> CSV | JSONLine | JSON | Excel | XML | ORC | Parquet:
         file_type = file_format.pop("type", None)

syncmaster/dto/transfers.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class Strategy:
+    type: str
+
+    @classmethod
+    def from_dict(cls, data: dict) -> Strategy:
+        strategy_classes = {
+            "full": FullStrategy,
+            "incremental": IncrementalStrategy,
+        }
+
+        strategy_type = data.get("type")
+        if strategy_type not in strategy_classes:
+            raise ValueError(f"Unknown strategy type: {strategy_type}")
+
+        return strategy_classes[strategy_type](**data)
+
+
+@dataclass
+class FullStrategy(Strategy):
+    pass
+
+
+@dataclass
+class IncrementalStrategy(Strategy):
+    increment_by: str

syncmaster/dto/transfers_strategy.py

@@ -255,8 +255,10 @@ def validate_increment_by(cls, values):
         source_type = values.source_params.type
         increment_by = values.strategy_params.increment_by
 
-        if source_type in FILE_CONNECTION_TYPES and increment_by != "modified_since":
-            raise ValueError("Field 'increment_by' must be equal to 'modified_since' for file source types")
+        if source_type in FILE_CONNECTION_TYPES and increment_by not in ("modified_since", "file_name"):
+            raise ValueError(
+                "Field 'increment_by' must be equal to 'modified_since' or 'file_name' for file source types",
+            )
 
         return values
 

syncmaster/schemas/v1/transfers/__init__.py

@@ -4,6 +4,15 @@
 from tempfile import TemporaryDirectory
 from typing import Any
 
+from horizon.client.auth import LoginPassword
+from horizon.client.sync import HorizonClientSync
+from horizon.commons.schemas.v1 import (
+    NamespaceCreateRequestV1,
+    NamespacePaginateQueryV1,
+)
+from horizon_hwm_store import HorizonHWMStore
+from onetl.strategy import IncrementalStrategy
+
 from syncmaster.db.models import Connection, Run
 from syncmaster.dto.connections import (
     ClickhouseConnectionDTO,
@@ -36,6 +45,7 @@
     SFTPTransferDTO,
     WebDAVTransferDTO,
 )
+from syncmaster.dto.transfers_strategy import Strategy
 from syncmaster.exceptions.connection import ConnectionTypeNotRecognizedError
 from syncmaster.worker.handlers.base import Handler
 from syncmaster.worker.handlers.db.clickhouse import ClickhouseHandler
@@ -139,11 +149,13 @@
 
 
 class TransferController:
+    settings: WorkerAppSettings
     source_handler: Handler
     target_handler: Handler
 
     def __init__(
         self,
+        settings: WorkerAppSettings,
         run: Run,
         source_connection: Connection,
         source_auth_data: dict,
@@ -152,27 +164,32 @@ def __init__(
     ):
         self.temp_dir = TemporaryDirectory(prefix=f"syncmaster_{run.id}_")
 
+        self.settings = settings
         self.run = run
         self.source_handler = self.get_handler(
             connection_data=source_connection.data,
             run_data={"id": run.id, "created_at": run.created_at},
+            transfer_id=run.transfer.id,
             transfer_params=run.transfer.source_params,
+            strategy_params=run.transfer.strategy_params,
             transformations=run.transfer.transformations,
             connection_auth_data=source_auth_data,
             temp_dir=TemporaryDirectory(dir=self.temp_dir.name, prefix="downloaded_"),
         )
         self.target_handler = self.get_handler(
             connection_data=target_connection.data,
             run_data={"id": run.id, "created_at": run.created_at},
+            transfer_id=run.transfer.id,
             transfer_params=run.transfer.target_params,
+            strategy_params=run.transfer.strategy_params,
             transformations=run.transfer.transformations,
             connection_auth_data=target_auth_data,
             temp_dir=TemporaryDirectory(dir=self.temp_dir.name, prefix="written_"),
         )
 
-    def perform_transfer(self, settings: WorkerAppSettings) -> None:
+    def perform_transfer(self) -> None:
         try:
-            spark = settings.worker.CREATE_SPARK_SESSION_FUNCTION(
+            spark = self.settings.worker.CREATE_SPARK_SESSION_FUNCTION(
                 run=self.run,
                 source=self.source_handler.connection_dto,
                 target=self.target_handler.connection_dto,
@@ -182,6 +199,9 @@ def perform_transfer(self, settings: WorkerAppSettings) -> None:
                 self.source_handler.connect(spark)
                 self.target_handler.connect(spark)
 
+                if self.source_handler.transfer_dto.strategy.type == "incremental":
+                    return self._perform_incremental_transfer()
+
                 df = self.source_handler.read()
                 self.target_handler.write(df)
         finally:
@@ -192,7 +212,9 @@ def get_handler(
         connection_data: dict[str, Any],
         connection_auth_data: dict,
         run_data: dict[str, Any],
+        transfer_id: int,
         transfer_params: dict[str, Any],
+        strategy_params: dict[str, Any],
         transformations: list[dict],
         temp_dir: TemporaryDirectory,
     ) -> Handler:
@@ -207,7 +229,48 @@ def get_handler(
 
         return handler(
             connection_dto=connection_dto(**connection_data),
-            transfer_dto=transfer_dto(**transfer_params, transformations=transformations),
+            transfer_dto=transfer_dto(
+                id=transfer_id,
+                strategy=Strategy.from_dict(strategy_params),
+                transformations=transformations,
+                **transfer_params,
+            ),
             run_dto=run_dto(**run_data),
             temp_dir=temp_dir,
         )
+
+    def _perform_incremental_transfer(self) -> None:
+        self._ensure_horizon_namespace_exists()
+
+        with HorizonHWMStore(
+            api_url=self.settings.horizon.url,
+            auth=LoginPassword(login=self.settings.horizon.user, password=self.settings.horizon.password),
+            namespace=self.settings.horizon.namespace,
+        ) as hwm_store:
+            with IncrementalStrategy():
+                hwm_name = "_".join(
+                    [
+                        str(self.source_handler.transfer_dto.id),
+                        self.source_handler.connection_dto.type,
+                        self.source_handler.transfer_dto.directory_path,
+                    ],
+                )
+                hwm = hwm_store.get_hwm(hwm_name)
+
+                # S3 & HDFS sources do not currently support incremental reading
+                if not isinstance(self.source_handler, (S3Handler, HDFSHandler)):
+                    self.source_handler.hwm = hwm
+                self.target_handler.hwm = hwm
+
+                df = self.source_handler.read()
+                self.target_handler.write(df)
+
+    def _ensure_horizon_namespace_exists(self) -> None:
+        client = HorizonClientSync(
+            base_url=self.settings.horizon.url,
+            auth=LoginPassword(login=self.settings.horizon.user, password=self.settings.horizon.password),
+        )
+        client.authorize()
+        namespace_query = NamespacePaginateQueryV1(name=self.settings.horizon.namespace)
+        if not client.paginate_namespaces(query=namespace_query).items:
+            client.create_namespace(NamespaceCreateRequestV1(name=self.settings.horizon.namespace))

syncmaster/worker/controller.py

@@ -7,6 +7,8 @@
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING
 
+from etl_entities.hwm import HWM
+
 from syncmaster.dto.connections import ConnectionDTO
 from syncmaster.dto.runs import RunDTO
 from syncmaster.dto.transfers import TransferDTO
@@ -17,6 +19,8 @@
 
 
 class Handler(ABC):
+    hwm: HWM | None = None
+
     def __init__(
         self,
         connection_dto: ConnectionDTO,