[DOP-21976] Add compression options for XML

MobileTeleSystems · Dec 5, 2024 · 3e63d9b · 3e63d9b
1 parent beecd7a
commit 3e63d9b
Show file tree

Hide file tree

Showing 8 changed files with 30 additions and 7 deletions.
diff --git a/docs/changelog/next_release/161.feature.rst b/docs/changelog/next_release/161.feature.rst
@@ -0,0 +1 @@
+Add compression options to XML file format
diff --git a/syncmaster/dto/transfers.py b/syncmaster/dto/transfers.py
@@ -36,17 +36,24 @@ class FileTransferDTO(TransferDTO):
 
     def __post_init__(self):
         if isinstance(self.file_format, dict):
-            self.file_format = self._get_format(self.file_format.copy())
+            self.file_format = self._get_file_format(self.file_format.copy())
         if isinstance(self.df_schema, str):
             self.df_schema = json.loads(self.df_schema)
 
-    def _get_format(self, file_format: dict):
-        file_type = file_format.pop("type", None)
+    def _get_file_format(self, file_format: dict) -> CSV | JSONLine | JSON | Excel | XML | ORC | Parquet:
+        file_type = self._prepare_file_format(file_format)
         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}")
 
+    @staticmethod
+    def _prepare_file_format(file_format: dict) -> str | None:
+        file_type = file_format.pop("type", None)
+        if file_type == "xml" and file_format.get("compression") == "none":
+            file_format.pop("compression")
+        return file_type
+
 
 @dataclass
 class PostgresTransferDTO(DBTransferDTO):

diff --git a/syncmaster/schemas/v1/transfers/file_format.py b/syncmaster/schemas/v1/transfers/file_format.py
@@ -49,6 +49,14 @@ class CSVCompression(str, Enum):
     DEFLATE = "deflate"
 
 
+class XMLCompression(str, Enum):
+    NONE = "none"
+    BZIP2 = "bzip2"
+    GZIP = "gzip"
+    LZ4 = "lz4"
+    SNAPPY = "snappy"
+
+
 class CSV(BaseModel):
     type: CSV_FORMAT
     delimiter: str = ","
@@ -84,6 +92,7 @@ class XML(BaseModel):
     type: XML_FORMAT
     root_tag: str
     row_tag: str
+    compression: XMLCompression = XMLCompression.GZIP
 
 
 class ORC(BaseModel):

diff --git a/tests/test_integration/test_run_transfer/test_hdfs.py b/tests/test_integration/test_run_transfer/test_hdfs.py
@@ -229,8 +229,8 @@ async def test_run_transfer_hdfs_to_postgres(
             id="parquet",
         ),
         pytest.param(
-            ("xml", {}),
-            "without_compression",
+            ("xml", {"compression": "snappy"}),
+            "with_compression",
             id="xml",
         ),
     ],

diff --git a/tests/test_integration/test_run_transfer/test_s3.py b/tests/test_integration/test_run_transfer/test_s3.py
@@ -230,8 +230,8 @@ async def test_run_transfer_s3_to_postgres(
             id="parquet",
         ),
         pytest.param(
-            ("xml", {}),
-            "without_compression",
+            ("xml", {"compression": "none"}),
+            "with_compression",
             id="xml",
         ),
     ],

diff --git a/tests/test_unit/test_transfers/test_file_transfers/test_create_transfer.py b/tests/test_unit/test_transfers/test_file_transfers/test_create_transfer.py
@@ -59,6 +59,7 @@
                 "type": "xml",
                 "root_tag": "data",
                 "row_tag": "record",
+                "compression": "lz4",
             },
             "options": {
                 "some": "option",
@@ -166,6 +167,7 @@ async def test_developer_plus_can_create_s3_transfer(
             "type": "xml",
             "root_tag": "data",
             "row_tag": "record",
+            "compression": "lz4",
         },
         "orc": {
             "type": "orc",
@@ -221,6 +223,7 @@ async def test_developer_plus_can_create_s3_transfer(
                 "type": "xml",
                 "root_tag": "data",
                 "row_tag": "record",
+                "compression": "bzip2",
             },
         },
         {
@@ -320,6 +323,7 @@ async def test_developer_plus_can_create_hdfs_transfer(
             "type": "xml",
             "root_tag": "data",
             "row_tag": "record",
+            "compression": "bzip2",
         },
         "orc": {
             "type": "orc",

diff --git a/tests/test_unit/test_transfers/test_file_transfers/test_read_transfer.py b/tests/test_unit/test_transfers/test_file_transfers/test_read_transfer.py
@@ -41,6 +41,7 @@
                 "type": "xml",
                 "root_tag": "data",
                 "row_tag": "record",
+                "compression": "bzip2",
             },
             "options": {},
         },

diff --git a/tests/test_unit/test_transfers/test_file_transfers/test_update_transfer.py b/tests/test_unit/test_transfers/test_file_transfers/test_update_transfer.py
@@ -41,6 +41,7 @@
                 "type": "xml",
                 "root_tag": "data",
                 "row_tag": "record",
+                "compression": "bzip2",
             },
             "options": {},
         },