[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)
@@ -81,6 +81,9 @@ def _directory_path_is_valid_path(cls, value):
     @field_validator("file_name_template")
     @classmethod
     def validate_file_name_template(cls, value):
+        if ":" in value:
+            raise ValueError("The ':' symbol is not allowed by Spark in filenames")
+
         required_keys = {"index", "extension"}
         placeholders = {key for key in required_keys if f"{{{key}}}" in value}
 

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,12 @@
 
 from __future__ import annotations
 
+import os
+import shutil
 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 +19,7 @@
 
 
 class FileHandler(Handler):
-    connection: BaseFileDFConnection
+    df_connection: BaseFileDFConnection
     connection_dto: ConnectionDTO
     transfer_dto: FileTransferDTO
     _operators = {
@@ -40,7 +42,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 +61,65 @@ def read(self) -> DataFrame:
         return df
 
     def write(self, df: DataFrame) -> None:
-        writer = FileDFWriter(
-            connection=self.connection,
-            format=self.transfer_dto.file_format,
-            target_path=self.transfer_dto.directory_path,
-            options=self.transfer_dto.options,
+        tmp_path = os.path.join(self.transfer_dto.directory_path, ".tmp", str(self.run_dto.id))
+        try:
+            writer = FileDFWriter(
+                connection=self.df_connection,
+                format=self.transfer_dto.file_format,
+                target_path=tmp_path,
+                options=self.transfer_dto.options,
+            )
+            writer.run(df=df)
+
+            self._rename_files(tmp_path)
+
+            mover = FileMover(
+                connection=self.file_connection,
+                source_path=tmp_path,
+                target_path=self.transfer_dto.directory_path,
+            )
+            mover.run()
+        finally:
+            shutil.rmtree(tmp_path, ignore_errors=True)
+
+    def _rename_files(self, tmp_path: str) -> None:
+        files = self.file_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.file_connection.rename_file(old_path, new_path)
+
+    def _get_file_name(self, index: str, extension: str) -> str:
+        return self.transfer_dto.file_name_template.format(
+            index=index,
+            extension=extension,
+            run_id=self.run_dto.id,
+            run_created_at=self.run_dto.created_at.strftime("%Y_%m_%d_%H_%M_%S"),
         )
 
-        return writer.run(df=df)
+    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 f"{extension}.{compression}"
 
     def _make_rows_filter_expression(self, filters: list[dict]) -> str | None:
         expressions = []

syncmaster/worker/handlers/file/ftp.py

@@ -18,12 +18,12 @@ class FTPHandler(FileProtocolHandler):
     connection_dto: FTPConnectionDTO
 
     def connect(self, spark: SparkSession) -> None:
-        self.connection = FTP(
+        self.file_connection = FTP(
             host=self.connection_dto.host,
             port=self.connection_dto.port,
             user=self.connection_dto.user,
             password=self.connection_dto.password,
         ).check()
-        self.local_connection = SparkLocalFS(
+        self.local_df_connection = SparkLocalFS(
             spark=spark,
         ).check()

syncmaster/worker/handlers/file/ftps.py

@@ -18,12 +18,12 @@ class FTPSHandler(FileProtocolHandler):
     connection_dto: FTPSConnectionDTO
 
     def connect(self, spark: SparkSession) -> None:
-        self.connection = FTPS(
+        self.file_connection = FTPS(
             host=self.connection_dto.host,
             port=self.connection_dto.port,
             user=self.connection_dto.user,
             password=self.connection_dto.password,
         ).check()
-        self.local_connection = SparkLocalFS(
+        self.local_df_connection = SparkLocalFS(
             spark=spark,
         ).check()

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,12 @@ 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.file_connection = HDFS(
+            cluster=self.connection_dto.cluster,
+            user=self.connection_dto.user,
+            password=self.connection_dto.password,
+        ).check()