[DOP-31721] Implement SQL transformations

Author: Ilyas Gasanov

docs/changelog/next_release/330.feature.rst

@@ -0,0 +1 @@
+Implement SQL transformations in worker

syncmaster/worker/handlers/db/base.py

@@ -58,7 +58,14 @@ def read(self) -> DataFrame:
             columns=self._get_columns_filter_expressions(),
             **reader_params,
         )
-        return reader.run()
+        df = reader.run()
+
+        sql_query = self._get_sql_query()
+        if sql_query:
+            df.createOrReplaceTempView("source")
+            df = self.connection.spark.sql(sql_query)
+
+        return df
 
     def write(self, df: DataFrame) -> None:
         if self.transfer_dto.strategy.type == "incremental" and self.hwm and self.hwm.value:
@@ -110,6 +117,12 @@ def _get_columns_filter_expressions(self) -> list[str] | None:
 
         return self._make_columns_filter_expressions(expressions)
 
+    def _get_sql_query(self) -> str | None:
+        for transformation in self.transfer_dto.transformations:
+            if transformation["type"] == "sql":
+                return transformation["query"]
+        return None
+
     def _get_reading_options(self) -> dict:
         options: dict[str, Any] = {}
 

syncmaster/worker/handlers/file/base.py

@@ -135,3 +135,9 @@ def _get_columns_filter_expressions(self) -> list[str] | None:
                 expressions.extend(transformation["filters"])
 
         return self._make_columns_filter_expressions(expressions)
+
+    def _get_sql_query(self) -> str | None:
+        for transformation in self.transfer_dto.transformations:
+            if transformation["type"] == "sql":
+                return transformation["query"]
+        return None

syncmaster/worker/handlers/file/local_df.py

@@ -58,6 +58,11 @@ def read(self) -> DataFrame:
         if columns_filter_expressions:
             df = df.selectExpr(*columns_filter_expressions)
 
+        sql_query = self._get_sql_query()
+        if sql_query:
+            df.createOrReplaceTempView("source")
+            df = self.df_connection.spark.sql(sql_query)
+
         return df
 
     def write(self, df: DataFrame) -> None:

syncmaster/worker/handlers/file/remote_df.py

@@ -35,6 +35,11 @@ def read(self) -> DataFrame:
         if columns_filter_expressions:
             df = df.selectExpr(*columns_filter_expressions)
 
+        sql_query = self._get_sql_query()
+        if sql_query:
+            df.createOrReplaceTempView("source")
+            df = self.df_connection.spark.sql(sql_query)
+
         return df
 
     def write(self, df: DataFrame) -> None:

syncmaster/worker/handlers/file/s3.py

@@ -71,6 +71,11 @@ def read(self) -> DataFrame:
         if columns_filter_expressions:
             df = df.selectExpr(*columns_filter_expressions)
 
+        sql_query = self._get_sql_query()
+        if sql_query:
+            df.createOrReplaceTempView("source")
+            df = self.df_connection.spark.sql(sql_query)
+
         return df
 
     @slot

…t_integration/test_scheduler/conftest.py tests/test_integration/conftest.pytests/test_integration/test_scheduler/conftest.py renamed to tests/test_integration/conftest.py tests/test_integration/test_scheduler/conftest.py renamed to tests/test_integration/conftest.py

@@ -1,4 +1,5 @@
 pytest_plugins = [
     "tests.test_unit.test_scheduler.scheduler_fixtures",
     "tests.test_integration.test_scheduler.scheduler_fixtures",
+    "tests.test_integration.test_run_transfer.connection_fixtures",
 ]

tests/test_integration/test_run_transfer/conftest.py

@@ -24,7 +24,9 @@
     dataframe_rows_filter_transformations,
     expected_dataframe_columns_filter,
     expected_dataframe_rows_filter,
+    expected_sql_transformation,
     file_metadata_filter_transformations,
+    sql_transformation,
 )
 from tests.test_integration.test_run_transfer.connection_fixtures.ftp_fixtures import (
     ftp_connection,
@@ -151,6 +153,7 @@
     "dataframe_rows_filter_transformations",
     "expected_dataframe_columns_filter",
     "expected_dataframe_rows_filter",
+    "expected_sql_transformation",
     "file_format_flavor",
     "file_metadata_filter_transformations",
     "ftp_connection",
@@ -239,6 +242,7 @@
     "sftp_for_worker",
     "source_file_format",
     "spark",
+    "sql_transformation",
     "target_file_format",
     "update_transfer_strategy",
     "webdav_connection",

tests/test_integration/test_run_transfer/connection_fixtures/__init__.py

@@ -133,3 +133,19 @@ def file_metadata_filter_transformations():
             ],
         },
     ]
+
+
+@pytest.fixture
+def sql_transformation():
+    return [
+        {
+            "type": "sql",
+            "query": "SELECT * FROM source WHERE NUMBER <= 20",
+            "dialect": "spark",
+        },
+    ]
+
+
+@pytest.fixture
+def expected_sql_transformation():
+    return lambda df, source_type: df.filter(df.NUMBER <= 20)