[DOP-28167] Slots on worker handlers support inheritance

dolfinus · dolfinus · commit 384b018b8efe · 2025-10-29T17:27:05.000+03:00
diff --git a/syncmaster/worker/handlers/db/base.py b/syncmaster/worker/handlers/db/base.py
@@ -8,6 +8,7 @@
 
 from onetl.base import BaseDBConnection
 from onetl.db import DBReader, DBWriter
+from onetl.hooks import slot, support_hooks
 
 from syncmaster.dto.transfers import DBTransferDTO
 from syncmaster.worker.handlers.base import Handler
@@ -16,6 +17,7 @@
     from pyspark.sql.dataframe import DataFrame
 
 
+@support_hooks
 class DBHandler(Handler):
     connection: BaseDBConnection
     transfer_dto: DBTransferDTO
@@ -35,6 +37,7 @@ class DBHandler(Handler):
         "not_ilike": "NOT ILIKE",
     }
 
+    @slot
     def read(self) -> DataFrame:
         reader_params = {}
         if self.transfer_dto.strategy.type == "incremental":
@@ -58,6 +61,7 @@ def read(self) -> DataFrame:
         )
         return reader.run()
 
+    @slot
     def write(self, df: DataFrame) -> None:
         if self.transfer_dto.strategy.type == "incremental" and self.hwm and self.hwm.value:
             self.transfer_dto.options["if_exists"] = "append"
diff --git a/syncmaster/worker/handlers/db/clickhouse.py b/syncmaster/worker/handlers/db/clickhouse.py
@@ -42,10 +42,6 @@ def connect(self, spark: SparkSession):
             spark=spark,
         ).check()
 
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
     @slot
     def write(self, df: DataFrame) -> None:
         normalized_df = self._normalize_column_names(df)
diff --git a/syncmaster/worker/handlers/db/hive.py b/syncmaster/worker/handlers/db/hive.py
@@ -38,10 +38,6 @@ def read(self) -> DataFrame:
         self.connection.spark.catalog.refreshTable(self.transfer_dto.table_name)
         return super().read()
 
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
-
     def _normalize_column_names(self, df: DataFrame) -> DataFrame:
         for column_name in df.columns:
             df = df.withColumnRenamed(column_name, column_name.lower())
diff --git a/syncmaster/worker/handlers/db/iceberg.py b/syncmaster/worker/handlers/db/iceberg.py
@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 from onetl.connection import Iceberg
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import IcebergRESTCatalogS3ConnectionDTO
 from syncmaster.dto.transfers import IcebergRESTCatalogS3TransferDTO
@@ -51,14 +51,6 @@ def connect(self, spark: SparkSession):
             ),
         ).check()
 
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
-
     def _normalize_column_names(self, df: DataFrame) -> DataFrame:
         for column_name in df.columns:
             df = df.withColumnRenamed(column_name, column_name.lower())
diff --git a/syncmaster/worker/handlers/db/mssql.py b/syncmaster/worker/handlers/db/mssql.py
@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 from onetl.connection import MSSQL
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import MSSQLConnectionDTO
 from syncmaster.dto.transfers import MSSQLTransferDTO
@@ -40,14 +40,6 @@ def connect(self, spark: SparkSession):
             spark=spark,
         ).check()
 
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
-
     def _normalize_column_names(self, df: DataFrame) -> DataFrame:
         for column_name in df.columns:
             df = df.withColumnRenamed(column_name, column_name.lower())
diff --git a/syncmaster/worker/handlers/db/mysql.py b/syncmaster/worker/handlers/db/mysql.py
@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 from onetl.connection import MySQL
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import MySQLConnectionDTO
 from syncmaster.dto.transfers import MySQLTransferDTO
@@ -37,14 +37,6 @@ def connect(self, spark: SparkSession):
             spark=spark,
         ).check()
 
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
-
     def _normalize_column_names(self, df: DataFrame) -> DataFrame:
         for column_name in df.columns:
             df = df.withColumnRenamed(column_name, column_name.lower())
diff --git a/syncmaster/worker/handlers/db/oracle.py b/syncmaster/worker/handlers/db/oracle.py
@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 from onetl.connection import Oracle
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import OracleConnectionDTO
 from syncmaster.dto.transfers import OracleTransferDTO
@@ -39,14 +39,6 @@ def connect(self, spark: SparkSession):
             spark=spark,
         ).check()
 
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
-
     def _normalize_column_names(self, df: DataFrame) -> DataFrame:
         for column_name in df.columns:
             df = df.withColumnRenamed(column_name, column_name.upper())
diff --git a/syncmaster/worker/handlers/db/postgres.py b/syncmaster/worker/handlers/db/postgres.py
@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 from onetl.connection import Postgres
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import PostgresConnectionDTO
 from syncmaster.dto.transfers import PostgresTransferDTO
@@ -38,14 +38,6 @@ def connect(self, spark: SparkSession):
             spark=spark,
         ).check()
 
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
-
     def _normalize_column_names(self, df: DataFrame) -> DataFrame:
         for column_name in df.columns:
             df = df.withColumnRenamed(column_name, column_name.lower())
diff --git a/syncmaster/worker/handlers/file/ftp.py b/syncmaster/worker/handlers/file/ftp.py
@@ -5,13 +5,13 @@
 
 from typing import TYPE_CHECKING
 
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import FTPConnectionDTO
 from syncmaster.worker.handlers.file.local_df import LocalDFFileHandler
 
 if TYPE_CHECKING:
-    from pyspark.sql import DataFrame, SparkSession
+    from pyspark.sql import SparkSession
 
 
 @support_hooks
@@ -31,11 +31,3 @@ def connect(self, spark: SparkSession) -> None:
         self.local_df_connection = SparkLocalFS(
             spark=spark,
         ).check()
-
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
diff --git a/syncmaster/worker/handlers/file/ftps.py b/syncmaster/worker/handlers/file/ftps.py
@@ -5,13 +5,13 @@
 
 from typing import TYPE_CHECKING
 
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import FTPSConnectionDTO
 from syncmaster.worker.handlers.file.local_df import LocalDFFileHandler
 
 if TYPE_CHECKING:
-    from pyspark.sql import DataFrame, SparkSession
+    from pyspark.sql import SparkSession
 
 
 @support_hooks
@@ -31,11 +31,3 @@ def connect(self, spark: SparkSession) -> None:
         self.local_df_connection = SparkLocalFS(
             spark=spark,
         ).check()
-
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
diff --git a/syncmaster/worker/handlers/file/hdfs.py b/syncmaster/worker/handlers/file/hdfs.py
@@ -5,13 +5,13 @@
 
 from typing import TYPE_CHECKING
 
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import HDFSConnectionDTO
 from syncmaster.worker.handlers.file.remote_df import RemoteDFFileHandler
 
 if TYPE_CHECKING:
-    from pyspark.sql import DataFrame, SparkSession
+    from pyspark.sql import SparkSession
 
 
 @support_hooks
@@ -31,11 +31,3 @@ def connect(self, spark: SparkSession):
             user=self.connection_dto.user,
             password=self.connection_dto.password,
         ).check()
-
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
diff --git a/syncmaster/worker/handlers/file/local_df.py b/syncmaster/worker/handlers/file/local_df.py
@@ -9,15 +9,18 @@
 from etl_entities.hwm import FileListHWM, FileModifiedTimeHWM
 from onetl.file import FileDFReader, FileDFWriter, FileDownloader, FileUploader
 from onetl.file.filter import FileSizeRange, Glob, Regexp
+from onetl.hooks import slot, support_hooks
 
 from syncmaster.worker.handlers.file.base import FileHandler
 
 if TYPE_CHECKING:
     from pyspark.sql import DataFrame
 
 
+@support_hooks
 class LocalDFFileHandler(FileHandler):
 
+    @slot
     def read(self) -> DataFrame:
         from pyspark.sql.types import StructType
 
@@ -59,8 +62,8 @@ def read(self) -> DataFrame:
 
         return df
 
+    @slot
     def write(self, df: DataFrame) -> None:
-
         writer = FileDFWriter(
             connection=self.local_df_connection,
             format=self.transfer_dto.file_format,
diff --git a/syncmaster/worker/handlers/file/remote_df.py b/syncmaster/worker/handlers/file/remote_df.py
@@ -7,15 +7,18 @@
 
 from onetl.file import FileDFReader, FileDFWriter, FileMover
 from onetl.file.filter import Glob
+from onetl.hooks import slot, support_hooks
 
 from syncmaster.worker.handlers.file.base import FileHandler
 
 if TYPE_CHECKING:
     from pyspark.sql import DataFrame
 
 
+@support_hooks
 class RemoteDFFileHandler(FileHandler):
 
+    @slot
     def read(self) -> DataFrame:
         from pyspark.sql.types import StructType
 
@@ -38,6 +41,7 @@ def read(self) -> DataFrame:
 
         return df
 
+    @slot
     def write(self, df: DataFrame) -> None:
         tmp_path = os.path.join(self.transfer_dto.directory_path, ".tmp", str(self.run_dto.id))
         try:
diff --git a/syncmaster/worker/handlers/file/s3.py b/syncmaster/worker/handlers/file/s3.py
@@ -70,7 +70,3 @@ def read(self) -> DataFrame:
             df = df.selectExpr(*columns_filter_expressions)
 
         return df
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
diff --git a/syncmaster/worker/handlers/file/samba.py b/syncmaster/worker/handlers/file/samba.py
@@ -5,13 +5,13 @@
 
 from typing import TYPE_CHECKING
 
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import SambaConnectionDTO
 from syncmaster.worker.handlers.file.local_df import LocalDFFileHandler
 
 if TYPE_CHECKING:
-    from pyspark.sql import DataFrame, SparkSession
+    from pyspark.sql import SparkSession
 
 
 @support_hooks
@@ -35,11 +35,3 @@ def connect(self, spark: SparkSession) -> None:
         self.local_df_connection = SparkLocalFS(
             spark=spark,
         ).check()
-
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
diff --git a/syncmaster/worker/handlers/file/sftp.py b/syncmaster/worker/handlers/file/sftp.py
@@ -5,13 +5,13 @@
 
 from typing import TYPE_CHECKING
 
-from onetl.hooks import slot, support_hooks
+from onetl.hooks import support_hooks
 
 from syncmaster.dto.connections import SFTPConnectionDTO
 from syncmaster.worker.handlers.file.local_df import LocalDFFileHandler
 
 if TYPE_CHECKING:
-    from pyspark.sql import DataFrame, SparkSession
+    from pyspark.sql import SparkSession
 
 
 @support_hooks
@@ -32,11 +32,3 @@ def connect(self, spark: SparkSession) -> None:
         self.local_df_connection = SparkLocalFS(
             spark=spark,
         ).check()
-
-    @slot
-    def read(self) -> DataFrame:
-        return super().read()
-
-    @slot
-    def write(self, df: DataFrame) -> None:
-        return super().write(df)
diff --git a/syncmaster/worker/handlers/file/webdav.py b/syncmaster/worker/handlers/file/webdav.py
