[DOP-21444] Add Excel integration tests

Ilyas Gasanov · Ilyas Gasanov · commit e2e1bd5e5969 · 2024-12-02T09:47:06.000+03:00
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -59,6 +59,7 @@ python-jose = { version = "^3.3.0", extras = ["cryptography"], optional = true }
 jinja2 = { version = "^3.1.4", optional = true }
 python-multipart = { version = ">=0.0.9,<0.0.18", optional = true }
 celery = { version = "^5.4.0", optional = true }
+pyspark = { version = ">=3.5.1,<3.5.2", optional = true }
 onetl = { version = "^0.12.0", extras = ["spark"], optional = true }
 pyyaml = {version = "*", optional = true}
 # due to not supporting MacOS 14.x https://www.psycopg.org/psycopg3/docs/news.html#psycopg-3-1-20
diff --git a/syncmaster/dto/transfers.py b/syncmaster/dto/transfers.py
@@ -4,7 +4,7 @@
 from dataclasses import dataclass
 from typing import ClassVar
 
-from onetl.file.format import CSV, JSON, JSONLine
+from onetl.file.format import CSV, JSON, Excel, JSONLine
 
 
 @dataclass
@@ -20,10 +20,17 @@ class DBTransferDTO(TransferDTO):
 @dataclass
 class FileTransferDTO(TransferDTO):
     directory_path: str
-    file_format: CSV | JSONLine | JSON
+    file_format: CSV | JSONLine | JSON | Excel
     options: dict
     df_schema: dict | None = None
 
+    _format_parsers = {
+        "csv": CSV,
+        "jsonline": JSONLine,
+        "json": JSON,
+        "excel": Excel,
+    }
+
     def __post_init__(self):
         if isinstance(self.file_format, dict):
             self.file_format = self._get_format(self.file_format.copy())
@@ -32,13 +39,10 @@ def __post_init__(self):
 
     def _get_format(self, file_format: dict):
         file_type = file_format.pop("type", None)
-        if file_type == "csv":
-            return CSV.parse_obj(file_format)
-        if file_type == "jsonline":
-            return JSONLine.parse_obj(file_format)
-        if file_type == "json":
-            return JSON.parse_obj(file_format)
-        raise ValueError("Unknown file type")
+        parser_class = self._format_parsers.get(file_type)
+        if parser_class is not None:
+            return parser_class.parse_obj(file_format)
+        raise ValueError(f"Unknown file type: {file_type}")
 
 
 @dataclass
diff --git a/syncmaster/worker/handlers/file/s3.py b/syncmaster/worker/handlers/file/s3.py
@@ -6,12 +6,13 @@
 from typing import TYPE_CHECKING
 
 from onetl.connection import SparkS3
+from onetl.file import FileDFReader
 
 from syncmaster.dto.connections import S3ConnectionDTO
 from syncmaster.worker.handlers.file.base import FileHandler
 
 if TYPE_CHECKING:
-    from pyspark.sql import SparkSession
+    from pyspark.sql import DataFrame, SparkSession
 
 
 class S3Handler(FileHandler):
@@ -29,3 +30,20 @@ def connect(self, spark: SparkSession):
             extra=self.connection_dto.additional_params,
             spark=spark,
         ).check()
+
+    def read(self) -> DataFrame:
+        from pyspark.sql.types import StructType
+
+        options = {}
+        if self.transfer_dto.file_format.__class__.__name__ == "Excel":
+            options = {"inferSchema": True}
+
+        reader = FileDFReader(
+            connection=self.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,
+            options={**options, **self.transfer_dto.options},
+        )
+
+        return reader.run()
diff --git a/syncmaster/worker/spark.py b/syncmaster/worker/spark.py
@@ -37,6 +37,7 @@ def get_worker_spark_session(
 
 def get_packages(db_type: str) -> list[str]:
     from onetl.connection import MSSQL, Clickhouse, Oracle, Postgres, SparkS3
+    from onetl.file.format import Excel
 
     if db_type == "postgres":
         return Postgres.get_packages()
@@ -51,7 +52,12 @@ def get_packages(db_type: str) -> list[str]:
         import pyspark
 
         spark_version = pyspark.__version__
-        return SparkS3.get_packages(spark_version=spark_version)
+        return SparkS3.get_packages(spark_version=spark_version) + Excel.get_packages(spark_version=spark_version)
+    if db_type == "hdfs":
+        import pyspark
+
+        spark_version = pyspark.__version__
+        return Excel.get_packages(spark_version=spark_version)
 
     # If the database type does not require downloading .jar packages
     return []
diff --git a/tests/test_integration/test_run_transfer/conftest.py b/tests/test_integration/test_run_transfer/conftest.py
@@ -3,14 +3,15 @@
 import os
 import secrets
 from collections import namedtuple
-from pathlib import Path, PurePosixPath
+from pathlib import Path, PosixPath, PurePosixPath
 
 import pyspark
 import pytest
 import pytest_asyncio
 from onetl.connection import MSSQL, Clickhouse, Hive, Oracle, Postgres, SparkS3
+from onetl.connection.file_connection.s3 import S3
 from onetl.db import DBWriter
-from onetl.file.format import CSV, JSON, JSONLine
+from onetl.file.format import CSV, JSON, Excel, JSONLine
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql.types import (
     DateType,
@@ -107,6 +108,9 @@ def spark(settings: Settings, request: FixtureRequest) -> SparkSession:
             )
         )
 
+    if "hdfs" in markers or "s3" in markers:
+        maven_packages.extend(Excel.get_packages(spark_version=pyspark.__version__))
+
     if maven_packages:
         spark = spark.config("spark.jars.packages", ",".join(maven_packages))
 
@@ -427,12 +431,22 @@ def s3_file_df_connection(s3_file_connection, spark, s3_server):
 
 
 @pytest.fixture(scope="session")
-def prepare_s3(resource_path, s3_file_connection, s3_file_df_connection_with_path: tuple[SparkS3, PurePosixPath]):
-    logger.info("START PREPARE HDFS")
-    connection, upload_to = s3_file_df_connection_with_path
-    files = upload_files(resource_path, upload_to, s3_file_connection)
-    logger.info("END PREPARE HDFS")
-    return connection, upload_to, files
+def prepare_s3(
+    resource_path: PosixPath,
+    s3_file_connection: S3,
+    s3_file_df_connection_with_path: tuple[SparkS3, PurePosixPath],
+):
+    logger.info("START PREPARE S3")
+    connection, remote_path = s3_file_df_connection_with_path
+
+    s3_file_connection.remove_dir(remote_path, recursive=True)
+    files = upload_files(resource_path, remote_path, s3_file_connection)
+
+    yield connection, remote_path, files
+
+    logger.info("START POST-CLEANUP S3")
+    s3_file_connection.remove_dir(remote_path, recursive=True)
+    logger.info("END POST-CLEANUP S3")
 
 
 @pytest.fixture(scope="session")
@@ -600,14 +614,14 @@ def prepare_clickhouse(
         pass
 
     def fill_with_data(df: DataFrame):
-        logger.info("START PREPARE ORACLE")
+        logger.info("START PREPARE CLICKHOUSE")
         db_writer = DBWriter(
             connection=onetl_conn,
             target=f"{clickhouse.user}.source_table",
             options=Clickhouse.WriteOptions(createTableOptions="ENGINE = Memory"),
         )
         db_writer.run(df)
-        logger.info("END PREPARE ORACLE")
+        logger.info("END PREPARE CLICKHOUSE")
 
     yield onetl_conn, fill_with_data
 
@@ -666,7 +680,7 @@ def fill_with_data(df: DataFrame):
         pass
 
 
-@pytest.fixture(params=[("csv", {}), ("jsonline", {}), ("json", {})])
+@pytest.fixture(params=[("csv", {}), ("jsonline", {}), ("json", {}), ("excel", {})])
 def source_file_format(request: FixtureRequest):
     name, params = request.param
     if name == "csv":
@@ -690,10 +704,17 @@ def source_file_format(request: FixtureRequest):
             **params,
         )
 
+    if name == "excel":
+        return "excel", Excel(
+            header=True,
+            inferSchema=True,
+            **params,
+        )
+
     raise ValueError(f"Unsupported file format: {name}")
 
 
-@pytest.fixture(params=[("csv", {}), ("jsonline", {})])
+@pytest.fixture(params=[("csv", {}), ("jsonline", {}), ("excel", {})])
 def target_file_format(request: FixtureRequest):
     name, params = request.param
     if name == "csv":
@@ -712,6 +733,12 @@ def target_file_format(request: FixtureRequest):
             **params,
         )
 
+    if name == "excel":
+        return "excel", Excel(
+            header=False,
+            **params,
+        )
+
     raise ValueError(f"Unsupported file format: {name}")
 
 
diff --git a/tests/test_integration/test_run_transfer/test_hdfs.py b/tests/test_integration/test_run_transfer/test_hdfs.py
@@ -8,6 +8,7 @@
 from onetl.db import DBReader
 from onetl.file import FileDFReader
 from pyspark.sql import DataFrame
+from pyspark.sql.functions import col, date_format, date_trunc, to_timestamp
 from pytest import FixtureRequest
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -37,6 +38,7 @@ async def hdfs_to_postgres(
     file_format_flavor: str,
 ):
     format_name, file_format = source_file_format
+    format_name_in_path = "xlsx" if format_name == "excel" else format_name
     _, source_path, _ = prepare_hdfs
 
     result = await create_transfer(
@@ -47,7 +49,7 @@ async def hdfs_to_postgres(
         target_connection_id=postgres_connection.id,
         source_params={
             "type": "hdfs",
-            "directory_path": os.fspath(source_path / "file_df_connection" / format_name / file_format_flavor),
+            "directory_path": os.fspath(source_path / "file_df_connection" / format_name_in_path / file_format_flavor),
             "file_format": {
                 "type": format_name,
                 **file_format.dict(),
@@ -121,6 +123,11 @@ async def postgres_to_hdfs(
             "without_compression",
             id="jsonline",
         ),
+        pytest.param(
+            ("excel", {}),
+            "with_header",
+            id="excel",
+        ),
     ],
     indirect=["source_file_format", "file_format_flavor"],
 )
@@ -135,6 +142,7 @@ async def test_run_transfer_hdfs_to_postgres(
 ):
     # Arrange
     postgres, _ = prepare_postgres
+    file_format, _ = source_file_format
 
     # Act
     result = await client.post(
@@ -164,6 +172,12 @@ async def test_run_transfer_hdfs_to_postgres(
         table="public.target_table",
     )
     df = reader.run()
+
+    # as initial datetime values in excel are rounded to milliseconds
+    if file_format == "excel":
+        df = df.withColumn("REGISTERED_AT", date_trunc("second", col("REGISTERED_AT")))
+        init_df = init_df.withColumn("REGISTERED_AT", date_trunc("second", col("REGISTERED_AT")))
+
     for field in init_df.schema:
         df = df.withColumn(field.name, df[field.name].cast(field.dataType))
 
@@ -183,6 +197,11 @@ async def test_run_transfer_hdfs_to_postgres(
             "without_compression",
             id="jsonline",
         ),
+        pytest.param(
+            ("excel", {}),
+            "with_header",
+            id="excel",
+        ),
     ],
     indirect=["target_file_format", "file_format_flavor"],
 )
@@ -235,6 +254,13 @@ async def test_run_transfer_postgres_to_hdfs(
     )
     df = reader.run()
 
+    # when spark writes to excel, the datetime precision is truncated to milliseconds (without rounding)
+    if format_name == "excel":
+        init_df = init_df.withColumn(
+            "REGISTERED_AT",
+            to_timestamp(date_format(col("REGISTERED_AT"), "yyyy-MM-dd HH:mm:ss.SSS")),
+        )
+
     for field in init_df.schema:
         df = df.withColumn(field.name, df[field.name].cast(field.dataType))
 
diff --git a/tests/test_integration/test_run_transfer/test_mssql.py b/tests/test_integration/test_run_transfer/test_mssql.py
diff --git a/tests/test_integration/test_run_transfer/test_s3.py b/tests/test_integration/test_run_transfer/test_s3.py
