Skip to content

Commit 3a480f5

Browse files
authored
Fix passing the gzip compression parameter on sftp_to_gcs. (#20553)
1 parent c5c18c5 commit 3a480f5

File tree

2 files changed

+42
-0
lines changed

2 files changed

+42
-0
lines changed

β€Žairflow/providers/google/cloud/transfers/sftp_to_gcs.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ def _copy_single_object(
168168
object_name=destination_object,
169169
filename=tmp.name,
170170
mime_type=self.mime_type,
171+
gzip=self.gzip,
171172
)
172173

173174
if self.move_object:

β€Žtests/providers/google/cloud/transfers/test_sftp_to_gcs.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,44 @@ def test_execute_copy_single_file(self, sftp_hook, gcs_hook):
8484
object_name=DESTINATION_PATH_FILE,
8585
filename=mock.ANY,
8686
mime_type=DEFAULT_MIME_TYPE,
87+
gzip=False,
88+
)
89+
90+
sftp_hook.return_value.delete_file.assert_not_called()
91+
92+
@mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.GCSHook")
93+
@mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.SFTPHook")
94+
def test_execute_copy_single_file_with_compression(self, sftp_hook, gcs_hook):
95+
task = SFTPToGCSOperator(
96+
task_id=TASK_ID,
97+
source_path=SOURCE_OBJECT_NO_WILDCARD,
98+
destination_bucket=TEST_BUCKET,
99+
destination_path=DESTINATION_PATH_FILE,
100+
move_object=False,
101+
gcp_conn_id=GCP_CONN_ID,
102+
sftp_conn_id=SFTP_CONN_ID,
103+
delegate_to=DELEGATE_TO,
104+
impersonation_chain=IMPERSONATION_CHAIN,
105+
gzip=True,
106+
)
107+
task.execute(None)
108+
gcs_hook.assert_called_once_with(
109+
gcp_conn_id=GCP_CONN_ID,
110+
delegate_to=DELEGATE_TO,
111+
impersonation_chain=IMPERSONATION_CHAIN,
112+
)
113+
sftp_hook.assert_called_once_with(SFTP_CONN_ID)
114+
115+
sftp_hook.return_value.retrieve_file.assert_called_once_with(
116+
os.path.join(SOURCE_OBJECT_NO_WILDCARD), mock.ANY
117+
)
118+
119+
gcs_hook.return_value.upload.assert_called_once_with(
120+
bucket_name=TEST_BUCKET,
121+
object_name=DESTINATION_PATH_FILE,
122+
filename=mock.ANY,
123+
mime_type=DEFAULT_MIME_TYPE,
124+
gzip=True,
87125
)
88126

89127
sftp_hook.return_value.delete_file.assert_not_called()
@@ -119,6 +157,7 @@ def test_execute_move_single_file(self, sftp_hook, gcs_hook):
119157
object_name=DESTINATION_PATH_FILE,
120158
filename=mock.ANY,
121159
mime_type=DEFAULT_MIME_TYPE,
160+
gzip=False,
122161
)
123162

124163
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):
162201
object_name="destination_dir/test_object3.json",
163202
mime_type=DEFAULT_MIME_TYPE,
164203
filename=mock.ANY,
204+
gzip=False,
165205
),
166206
mock.call(
167207
bucket_name=TEST_BUCKET,
168208
object_name="destination_dir/sub_dir/test_object3.json",
169209
mime_type=DEFAULT_MIME_TYPE,
170210
filename=mock.ANY,
211+
gzip=False,
171212
),
172213
]
173214
)

0 commit comments

Comments
 (0)