Fix passing the gzip compression parameter on sftp_to_gcs. (#20553)

apache · Dec 29, 2021 · 3a480f5 · 3a480f5
1 parent c5c18c5
commit 3a480f5
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/airflow/providers/google/cloud/transfers/sftp_to_gcs.py b/airflow/providers/google/cloud/transfers/sftp_to_gcs.py
@@ -168,6 +168,7 @@ def _copy_single_object(
                 object_name=destination_object,
                 filename=tmp.name,
                 mime_type=self.mime_type,
+                gzip=self.gzip,
             )
 
         if self.move_object:

diff --git a/tests/providers/google/cloud/transfers/test_sftp_to_gcs.py b/tests/providers/google/cloud/transfers/test_sftp_to_gcs.py
@@ -84,6 +84,44 @@ def test_execute_copy_single_file(self, sftp_hook, gcs_hook):
             object_name=DESTINATION_PATH_FILE,
             filename=mock.ANY,
             mime_type=DEFAULT_MIME_TYPE,
+            gzip=False,
+        )
+
+        sftp_hook.return_value.delete_file.assert_not_called()
+
+    @mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.GCSHook")
+    @mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.SFTPHook")
+    def test_execute_copy_single_file_with_compression(self, sftp_hook, gcs_hook):
+        task = SFTPToGCSOperator(
+            task_id=TASK_ID,
+            source_path=SOURCE_OBJECT_NO_WILDCARD,
+            destination_bucket=TEST_BUCKET,
+            destination_path=DESTINATION_PATH_FILE,
+            move_object=False,
+            gcp_conn_id=GCP_CONN_ID,
+            sftp_conn_id=SFTP_CONN_ID,
+            delegate_to=DELEGATE_TO,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            gzip=True,
+        )
+        task.execute(None)
+        gcs_hook.assert_called_once_with(
+            gcp_conn_id=GCP_CONN_ID,
+            delegate_to=DELEGATE_TO,
+            impersonation_chain=IMPERSONATION_CHAIN,
+        )
+        sftp_hook.assert_called_once_with(SFTP_CONN_ID)
+
+        sftp_hook.return_value.retrieve_file.assert_called_once_with(
+            os.path.join(SOURCE_OBJECT_NO_WILDCARD), mock.ANY
+        )
+
+        gcs_hook.return_value.upload.assert_called_once_with(
+            bucket_name=TEST_BUCKET,
+            object_name=DESTINATION_PATH_FILE,
+            filename=mock.ANY,
+            mime_type=DEFAULT_MIME_TYPE,
+            gzip=True,
         )
 
         sftp_hook.return_value.delete_file.assert_not_called()
@@ -119,6 +157,7 @@ def test_execute_move_single_file(self, sftp_hook, gcs_hook):
             object_name=DESTINATION_PATH_FILE,
             filename=mock.ANY,
             mime_type=DEFAULT_MIME_TYPE,
+            gzip=False,
         )
 
         sftp_hook.return_value.delete_file.assert_called_once_with(SOURCE_OBJECT_NO_WILDCARD)
@@ -162,12 +201,14 @@ def test_execute_copy_with_wildcard(self, sftp_hook, gcs_hook):
                     object_name="destination_dir/test_object3.json",
                     mime_type=DEFAULT_MIME_TYPE,
                     filename=mock.ANY,
+                    gzip=False,
                 ),
                 mock.call(
                     bucket_name=TEST_BUCKET,
                     object_name="destination_dir/sub_dir/test_object3.json",
                     mime_type=DEFAULT_MIME_TYPE,
                     filename=mock.ANY,
+                    gzip=False,
                 ),
             ]
         )