[DOP-22146] Set the names of saved files (#201)

Author: IlyasDevelopment

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/connections/hdfs.py

@@ -25,7 +25,7 @@ class ReadHDFSConnectionDataSchema(BaseModel):
 
 
 class UpdateHDFSConnectionDataSchema(BaseModel):
-    cluster: str
+    cluster: str | None = None
 
 
 class CreateHDFSConnectionSchema(CreateConnectionBaseSchema):

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

@@ -2,8 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import re
 from pathlib import PurePosixPath
-from typing import Any
+from typing import Any, ClassVar
 
 from pydantic import BaseModel, Field, field_validator
 
@@ -63,11 +64,13 @@ 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)
 
+    FILE_NAME_PATTERN: ClassVar[re.Pattern] = re.compile(r"^[a-zA-Z0-9_.{}-]+$")
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -80,7 +83,10 @@ def _directory_path_is_valid_path(cls, value):
 
     @field_validator("file_name_template")
     @classmethod
-    def validate_file_name_template(cls, value):
+    def validate_file_name_template(cls, value: str) -> str:
+        if not cls.FILE_NAME_PATTERN.match(value):
+            raise ValueError("Template contains invalid characters. Allowed: letters, numbers, '.', '_', '-', '{', '}'")
+
         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
 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 onetl.file.filter import Glob
 
 from syncmaster.dto.connections import ConnectionDTO
 from syncmaster.dto.transfers import FileTransferDTO
@@ -15,9 +17,18 @@
 if TYPE_CHECKING:
     from pyspark.sql.dataframe import DataFrame
 
+COLUMN_FORMATS = ("parquet", "orc")
+
 
 class FileHandler(Handler):
-    connection: BaseFileDFConnection
+    """
+    TODO: FileHandler is actually handler for FileDFWriter with remote FS (direct write).
+    FileProtocolHandler is handler for FileDFWriter with local FS (write via upload).
+    Maybe we should keep here only common methods,
+    like file name generator and split other ones to classes where the method is really used.
+    """
+
+    df_connection: BaseFileDFConnection
     connection_dto: ConnectionDTO
     transfer_dto: FileTransferDTO
     _operators = {
@@ -35,12 +46,29 @@ class FileHandler(Handler):
         "not_ilike": "NOT ILIKE",
         "regexp": "RLIKE",
     }
+    _compression_to_file_suffix = {
+        "gzip": "gz",
+        "snappy": "snappy",
+        "zlib": "zlib",
+        "lz4": "lz4",
+        "bzip2": "bz2",
+        "deflate": "deflate",
+    }
+    _file_format_to_file_suffix = {
+        "json": "json",
+        "jsonline": "jsonl",
+        "csv": "csv",
+        "xml": "xml",
+        "excel": "xlsx",
+        "parquet": "parquet",
+        "orc": "orc",
+    }
 
     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 +87,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,
+                # ignore .crc and other metadata files
+                filters=[Glob(f"*.{self._get_file_extension()}")],
+            )
+            mover.run()
+        finally:
+            self.file_connection.remove_dir(tmp_path, recursive=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()
+            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) -> str:
+        file_format = self.transfer_dto.file_format.__class__.__name__.lower()
+        extension_suffix = self._file_format_to_file_suffix[file_format]
+
+        compression = getattr(self.transfer_dto.file_format, "compression", "none")
+        if compression == "none":
+            return extension_suffix
+
+        compression_suffix = self._compression_to_file_suffix[compression]
+
+        # https://github.com/apache/parquet-java/blob/fb6f0be0323f5f52715b54b8c6602763d8d0128d/parquet-common/src/main/java/org/apache/parquet/hadoop/metadata/CompressionCodecName.java#L26-L33
+        if extension_suffix == "parquet" and compression_suffix == "lz4":
+            return "lz4hadoop.parquet"
+
+        if extension_suffix in COLUMN_FORMATS:
+            return f"{compression_suffix}.{extension_suffix}"
+
+        return f"{extension_suffix}.{compression_suffix}"
 
     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()