[DOP-22297] Parallel reading from JDBC sources

Author: Ilyas Gasanov

syncmaster/dto/transfers.py

@@ -6,6 +6,7 @@
 
 from onetl.file.format import CSV, JSON, ORC, XML, Excel, JSONLine, Parquet
 
+from syncmaster.dto.transfers_resources import Resources
 from syncmaster.dto.transfers_strategy import FullStrategy, IncrementalStrategy
 
 
@@ -19,6 +20,7 @@ class DBTransferDTO(TransferDTO):
     id: int
     table_name: str
     strategy: FullStrategy | IncrementalStrategy
+    resources: Resources
     transformations: list[dict] | None = None
     options: dict | None = None
 

syncmaster/dto/transfers_resources.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+
+@dataclass
+class Resources:
+    max_parallel_tasks: int
+    cpu_cores_per_task: int
+    ram_bytes_per_task: int

syncmaster/worker/controller.py

@@ -41,6 +41,7 @@
     SFTPTransferDTO,
     WebDAVTransferDTO,
 )
+from syncmaster.dto.transfers_resources import Resources
 from syncmaster.dto.transfers_strategy import Strategy
 from syncmaster.exceptions.connection import ConnectionTypeNotRecognizedError
 from syncmaster.schemas.v1.connection_types import FILE_CONNECTION_TYPES
@@ -169,6 +170,7 @@ def __init__(
             transfer_id=run.transfer.id,
             transfer_params=run.transfer.source_params,
             strategy_params=run.transfer.strategy_params,
+            resources=run.transfer.resources,
             transformations=run.transfer.transformations,
             connection_auth_data=source_auth_data,
             temp_dir=TemporaryDirectory(dir=self.temp_dir.name, prefix="downloaded_"),
@@ -179,6 +181,7 @@ def __init__(
             transfer_id=run.transfer.id,
             transfer_params=run.transfer.target_params,
             strategy_params=run.transfer.strategy_params,
+            resources=run.transfer.resources,
             transformations=run.transfer.transformations,
             connection_auth_data=target_auth_data,
             temp_dir=TemporaryDirectory(dir=self.temp_dir.name, prefix="written_"),
@@ -215,6 +218,7 @@ def get_handler(
         transfer_id: int,
         transfer_params: dict[str, Any],
         strategy_params: dict[str, Any],
+        resources: dict[str, Any],
         transformations: list[dict],
         temp_dir: TemporaryDirectory,
     ) -> Handler:
@@ -232,6 +236,7 @@ def get_handler(
             transfer_dto=transfer_dto(
                 id=transfer_id,
                 strategy=Strategy.from_dict(strategy_params),
+                resources=Resources(**resources),
                 transformations=transformations,
                 **transfer_params,
             ),

syncmaster/worker/handlers/db/base.py

@@ -44,6 +44,27 @@ def read(self) -> DataFrame:
                 name=hwm_name,
                 expression=self.transfer_dto.strategy.increment_by,
             )
+            if self.transfer_dto.resources.max_parallel_tasks > 1:
+                reader_params["options"] = {
+                    "partition_column": self.transfer_dto.strategy.increment_by,
+                    "num_partitions": self.transfer_dto.resources.max_parallel_tasks,
+                    "partitioning_mode": "range" if self.hwm and self.hwm.value else "hash",
+                }
+
+        elif self.transfer_dto.strategy.type == "full" and self.transfer_dto.resources.max_parallel_tasks > 1:
+            schema_reader = DBReader(
+                connection=self.connection,
+                table=self.transfer_dto.table_name,
+                columns=self._get_columns_filter_expressions(),
+                options={"fetchsize": 0},
+            )
+            df = schema_reader.run()
+
+            reader_params["options"] = {
+                "partition_column": self._quote_field(df.schema[0].name),
+                "num_partitions": self.transfer_dto.resources.max_parallel_tasks,
+                "partitioning_mode": "hash",
+            }
 
         reader = DBReader(
             connection=self.connection,

syncmaster/worker/handlers/db/clickhouse.py

@@ -46,9 +46,9 @@ def write(self, df: DataFrame) -> None:
             (col for col in normalized_df.columns if col.lower().endswith("id")),
             normalized_df.columns[0],  # if there is no column with "id", take the first column
         )
-        quoted_sort_column = f'"{sort_column}"'
-
-        self.transfer_dto.options["createTableOptions"] = f"ENGINE = MergeTree() ORDER BY {quoted_sort_column}"
+        self.transfer_dto.options["createTableOptions"] = (
+            f"ENGINE = MergeTree() ORDER BY {self._quote_field(sort_column)}"
+        )
 
         if self.transfer_dto.strategy.type == "incremental" and self.hwm and self.hwm.value:
             self.transfer_dto.options["if_exists"] = "append"

tests/test_integration/test_run_transfer/test_clickhouse.py

@@ -47,6 +47,11 @@ async def postgres_to_clickhouse(
             "table_name": f"{clickhouse_for_conftest.user}.target_table",
         },
         strategy_params=strategy,
+        resources={
+            "max_parallel_tasks": 2,
+            "cpu_cores_per_task": 1,
+            "ram_bytes_per_task": 1024**3,
+        },
         transformations=transformations,
         queue_id=queue.id,
     )