[DOP-22146] Set the names of saved files

Author: Ilyas Gasanov

docs/changelog/next_release/201.feature.rst

@@ -0,0 +1 @@
+Use the `file_name_template` field to specify the names of saved files

syncmaster/dto/runs.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from datetime import datetime
+
+
+@dataclass
+class RunDTO:
+    id: str
+    created_at: datetime

syncmaster/dto/transfers.py

@@ -23,6 +23,7 @@ class FileTransferDTO(TransferDTO):
     directory_path: str
     file_format: CSV | JSONLine | JSON | Excel | XML | ORC | Parquet
     options: dict
+    file_name_template: str | None = None
     df_schema: dict | None = None
     transformations: list[dict] | None = None
 

syncmaster/schemas/v1/transfers/file/base.py

@@ -63,7 +63,7 @@ class CreateFileTransferTarget(BaseModel):
         discriminator="type",
     )
     file_name_template: str = Field(
-        default="{run_created_at}_{index}.{extension}",
+        default="{run_created_at}-{index}.{extension}",
         description="Template for file naming with required placeholders 'index' and 'extension'",
     )
     options: dict[str, Any] = Field(default_factory=dict)

syncmaster/worker/controller.py

@@ -20,6 +20,7 @@
     SFTPConnectionDTO,
     WebDAVConnectionDTO,
 )
+from syncmaster.dto.runs import RunDTO
 from syncmaster.dto.transfers import (
     ClickhouseTransferDTO,
     FTPSTransferDTO,
@@ -60,66 +61,79 @@
         HiveHandler,
         HiveConnectionDTO,
         HiveTransferDTO,
+        RunDTO,
     ),
     "oracle": (
         OracleHandler,
         OracleConnectionDTO,
         OracleTransferDTO,
+        RunDTO,
     ),
     "clickhouse": (
         ClickhouseHandler,
         ClickhouseConnectionDTO,
         ClickhouseTransferDTO,
+        RunDTO,
     ),
     "mssql": (
         MSSQLHandler,
         MSSQLConnectionDTO,
         MSSQLTransferDTO,
+        RunDTO,
     ),
     "mysql": (
         MySQLHandler,
         MySQLConnectionDTO,
         MySQLTransferDTO,
+        RunDTO,
     ),
     "postgres": (
         PostgresHandler,
         PostgresConnectionDTO,
         PostgresTransferDTO,
+        RunDTO,
     ),
     "s3": (
         S3Handler,
         S3ConnectionDTO,
         S3TransferDTO,
+        RunDTO,
     ),
     "hdfs": (
         HDFSHandler,
         HDFSConnectionDTO,
         HDFSTransferDTO,
+        RunDTO,
     ),
     "sftp": (
         SFTPHandler,
         SFTPConnectionDTO,
         SFTPTransferDTO,
+        RunDTO,
     ),
     "ftp": (
         FTPHandler,
         FTPConnectionDTO,
         FTPTransferDTO,
+        RunDTO,
     ),
     "ftps": (
         FTPSHandler,
         FTPSConnectionDTO,
         FTPSTransferDTO,
+        RunDTO,
     ),
     "samba": (
         SambaHandler,
         SambaConnectionDTO,
         SambaTransferDTO,
+        RunDTO,
     ),
     "webdav": (
         WebDAVHandler,
         WebDAVConnectionDTO,
         WebDAVTransferDTO,
+        RunDTO,
     ),
 }
 
@@ -141,13 +155,15 @@ def __init__(
         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_params=run.transfer.source_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_params=run.transfer.target_params,
             transformations=run.transfer.transformations,
             connection_auth_data=target_auth_data,
@@ -175,6 +191,7 @@ def get_handler(
         self,
         connection_data: dict[str, Any],
         connection_auth_data: dict,
+        run_data: dict[str, Any],
         transfer_params: dict[str, Any],
         transformations: list[dict],
         temp_dir: TemporaryDirectory,
@@ -186,10 +203,11 @@ def get_handler(
         if connection_handler_proxy.get(handler_type, None) is None:
             raise ConnectionTypeNotRecognizedError
 
-        handler, connection_dto, transfer_dto = connection_handler_proxy[handler_type]
+        handler, connection_dto, transfer_dto, run_dto = connection_handler_proxy[handler_type]
 
         return handler(
             connection_dto=connection_dto(**connection_data),
             transfer_dto=transfer_dto(**transfer_params, transformations=transformations),
+            run_dto=run_dto(**run_data),
             temp_dir=temp_dir,
         )

syncmaster/worker/handlers/base.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING
 
 from syncmaster.dto.connections import ConnectionDTO
+from syncmaster.dto.runs import RunDTO
 from syncmaster.dto.transfers import TransferDTO
 
 if TYPE_CHECKING:
@@ -20,10 +21,12 @@ def __init__(
         self,
         connection_dto: ConnectionDTO,
         transfer_dto: TransferDTO,
+        run_dto: RunDTO,
         temp_dir: TemporaryDirectory,
     ):
         self.connection_dto = connection_dto
         self.transfer_dto = transfer_dto
+        self.run_dto = run_dto
         self.temp_dir = temp_dir
 
     @abstractmethod

syncmaster/worker/handlers/file/base.py

@@ -3,10 +3,11 @@
 
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
 from onetl.base.base_file_df_connection import BaseFileDFConnection
-from onetl.file import FileDFReader, FileDFWriter
+from onetl.file import FileDFReader, FileDFWriter, FileMover
 
 from syncmaster.dto.connections import ConnectionDTO
 from syncmaster.dto.transfers import FileTransferDTO
@@ -17,7 +18,7 @@
 
 
 class FileHandler(Handler):
-    connection: BaseFileDFConnection
+    df_connection: BaseFileDFConnection
     connection_dto: ConnectionDTO
     transfer_dto: FileTransferDTO
     _operators = {
@@ -40,7 +41,7 @@ def read(self) -> DataFrame:
         from pyspark.sql.types import StructType
 
         reader = FileDFReader(
-            connection=self.connection,
+            connection=self.df_connection,
             format=self.transfer_dto.file_format,
             source_path=self.transfer_dto.directory_path,
             df_schema=StructType.fromJson(self.transfer_dto.df_schema) if self.transfer_dto.df_schema else None,
@@ -59,14 +60,74 @@ def read(self) -> DataFrame:
         return df
 
     def write(self, df: DataFrame) -> None:
+        tmp_path = os.path.join(self.transfer_dto.directory_path, ".tmp", str(self.run_dto.id))
         writer = FileDFWriter(
-            connection=self.connection,
+            connection=self.df_connection,
             format=self.transfer_dto.file_format,
-            target_path=self.transfer_dto.directory_path,
+            target_path=tmp_path,
             options=self.transfer_dto.options,
         )
+        writer.run(df=df)
+
+        self._rename_files(tmp_path)
+
+        mover = FileMover(
+            connection=self.connection,
+            source_path=tmp_path,
+            target_path=self.transfer_dto.directory_path,
+        )
+        mover.run()
+
+    def _rename_files(self, tmp_path: str) -> None:
+        files = self.connection.list_dir(tmp_path)
+
+        for index, file_name in enumerate(files):
+            extension = self._get_file_extension(str(file_name))
+            new_name = self._get_file_name(str(index), extension)
+            old_path = os.path.join(tmp_path, file_name)
+            new_path = os.path.join(tmp_path, new_name)
+            self.connection.rename_file(old_path, new_path)
+
+    def _get_file_name(self, index: str, extension: str) -> str:
+        return (
+            self.transfer_dto.file_name_template.replace(
+                "{index}",
+                index,
+            )
+            .replace(
+                "{extension}",
+                extension,
+            )
+            .replace(
+                "{run_id}",
+                str(self.run_dto.id),
+            )
+            .replace(
+                "{run_created_at}",
+                self.run_dto.created_at.strftime("%Y_%m_%d_%H_%M_%S"),
+            )
+        )
+
+    def _get_file_extension(self, file_name: str) -> str:
+        extension = self.transfer_dto.file_format.name
+        parts = file_name.split(".")
+
+        if extension == "xml":  # spark-xml does not write any extension to files
+            if len(parts) <= 1:
+                return extension
+
+            compression = parts[-1]
+
+        else:
+            if len(parts) <= 2:
+                return extension
+
+            compression = parts[-1] if parts[-1] != extension else parts[-2]
+
+        if extension in ("parquet", "orc"):
+            return f"{compression}.{extension}"
 
-        return writer.run(df=df)
+        return f"{extension}.{compression}"
 
     def _make_rows_filter_expression(self, filters: list[dict]) -> str | None:
         expressions = []

syncmaster/worker/handlers/file/hdfs.py

@@ -5,7 +5,7 @@
 
 from typing import TYPE_CHECKING
 
-from onetl.connection import SparkHDFS
+from onetl.connection import HDFS, SparkHDFS
 
 from syncmaster.dto.connections import HDFSConnectionDTO
 from syncmaster.worker.handlers.file.base import FileHandler
@@ -18,7 +18,10 @@ class HDFSHandler(FileHandler):
     connection_dto: HDFSConnectionDTO
 
     def connect(self, spark: SparkSession):
-        self.connection = SparkHDFS(
+        self.df_connection = SparkHDFS(
             cluster=self.connection_dto.cluster,
             spark=spark,
         ).check()
+        self.connection = HDFS(
+            cluster=self.connection_dto.cluster,
+        ).check()

syncmaster/worker/handlers/file/protocol.py

@@ -60,6 +60,8 @@ def write(self, df: DataFrame) -> None:
         for file in crc_files:
             os.remove(os.path.join(self.temp_dir.name, file))
 
+        self._rename_files()
+
         uploader = FileUploader(
             connection=self.connection,
             local_path=self.temp_dir.name,
@@ -68,6 +70,16 @@ def write(self, df: DataFrame) -> None:
         )
         uploader.run()
 
+    def _rename_files(self):
+        files = os.listdir(self.temp_dir.name)
+
+        for index, file_name in enumerate(files):
+            extension = self._get_file_extension(file_name)
+            new_name = self._get_file_name(str(index), extension)
+            old_path = os.path.join(self.temp_dir.name, file_name)
+            new_path = os.path.join(self.temp_dir.name, new_name)
+            os.rename(old_path, new_path)
+
     def _make_file_metadata_filters(self, filters: list[dict]) -> list[Glob | Regexp | FileSizeRange]:
         processed_filters = []
         for filter in filters:

syncmaster/worker/handlers/file/s3.py

@@ -5,7 +5,7 @@
 
 from typing import TYPE_CHECKING
 
-from onetl.connection import SparkS3
+from onetl.connection import S3, SparkS3
 from onetl.file import FileDFReader
 
 from syncmaster.dto.connections import S3ConnectionDTO
@@ -19,7 +19,7 @@ class S3Handler(FileHandler):
     connection_dto: S3ConnectionDTO
 
     def connect(self, spark: SparkSession):
-        self.connection = SparkS3(
+        self.df_connection = SparkS3(
             host=self.connection_dto.host,
             port=self.connection_dto.port,
             access_key=self.connection_dto.access_key,
@@ -30,6 +30,15 @@ def connect(self, spark: SparkSession):
             extra=self.connection_dto.additional_params,
             spark=spark,
         ).check()
+        self.connection = S3(
+            host=self.connection_dto.host,
+            port=self.connection_dto.port,
+            access_key=self.connection_dto.access_key,
+            secret_key=self.connection_dto.secret_key,
+            bucket=self.connection_dto.bucket,
+            protocol=self.connection_dto.protocol,
+            region=self.connection_dto.region,
+        ).check()
 
     def read(self) -> DataFrame:
         from pyspark.sql.types import StructType
@@ -39,7 +48,7 @@ def read(self) -> DataFrame:
             options = {"inferSchema": True}
 
         reader = FileDFReader(
-            connection=self.connection,
+            connection=self.df_connection,
             format=self.transfer_dto.file_format,
             source_path=self.transfer_dto.directory_path,
             df_schema=StructType.fromJson(self.transfer_dto.df_schema) if self.transfer_dto.df_schema else None,