Move uncompress_file function from airflow.utils to Hive provider

providers/src/airflow/providers/apache/hive/transfers/s3_to_hive.py

@@ -22,6 +22,7 @@
 import bz2
 import gzip
 import os
+import shutil
 import tempfile
 from collections.abc import Sequence
 from tempfile import NamedTemporaryFile, TemporaryDirectory
@@ -31,7 +32,6 @@
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.apache.hive.hooks.hive import HiveCliHook
-from airflow.utils.compression import uncompress_file
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -277,3 +277,20 @@ def _delete_top_row_and_compress(input_file_name, output_file_ext, dest_dir):
             for line in f_in:
                 f_out.write(line)
         return fn_output
+
+
+def uncompress_file(input_file_name, file_extension, dest_dir):
+    """Uncompress gz and bz2 files."""
+    if file_extension.lower() not in (".gz", ".bz2"):
+        raise NotImplementedError(
+            f"Received {file_extension} format. Only gz and bz2 files can currently be uncompressed."
+        )
+    if file_extension.lower() == ".gz":
+        fmodule = gzip.GzipFile
+    elif file_extension.lower() == ".bz2":
+        fmodule = bz2.BZ2File
+    with fmodule(input_file_name, mode="rb") as f_compressed, NamedTemporaryFile(
+        dir=dest_dir, mode="wb", delete=False
+    ) as f_uncompressed:
+        shutil.copyfileobj(f_compressed, f_uncompressed)
+    return f_uncompressed.name

providers/tests/apache/hive/transfers/test_s3_to_hive.py

@@ -30,7 +30,7 @@
 import pytest
 
 from airflow.exceptions import AirflowException
-from airflow.providers.apache.hive.transfers.s3_to_hive import S3ToHiveOperator
+from airflow.providers.apache.hive.transfers.s3_to_hive import S3ToHiveOperator, uncompress_file
 
 boto3 = pytest.importorskip("boto3")
 moto = pytest.importorskip("moto")
@@ -122,7 +122,7 @@ def _set_fn(self, fn, ext, header):
 
     # Helper method to fetch a file of a
     # certain format (file extension and header)
-    def _get_fn(self, ext, header):
+    def _get_fn(self, ext, header=None):
         key = self._get_key(ext, header)
         return self.file_names[key]
 
@@ -279,3 +279,19 @@ def test_execute_with_select_expression(self, mock_hiveclihook):
                     expression=select_expression,
                     input_serialization=input_serialization,
                 )
+
+    def test_uncompress_file(self):
+        # Testing txt file type
+        with pytest.raises(NotImplementedError, match="^Received .txt format. Only gz and bz2.*"):
+            uncompress_file(
+                **{"input_file_name": None, "file_extension": ".txt", "dest_dir": None},
+            )
+        # Testing gz file type
+        fn_txt = self._get_fn(".txt")
+        fn_gz = self._get_fn(".gz")
+        txt_gz = uncompress_file(fn_gz, ".gz", self.tmp_dir)
+        assert filecmp.cmp(txt_gz, fn_txt, shallow=False), "Uncompressed file does not match original"
+        # Testing bz2 file type
+        fn_bz2 = self._get_fn(".bz2")
+        txt_bz2 = uncompress_file(fn_bz2, ".bz2", self.tmp_dir)
+        assert filecmp.cmp(txt_bz2, fn_txt, shallow=False), "Uncompressed file does not match original"