This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 14fd8202a5b Honor gzip and mime_type on SFTPToGCS stream path (#72132)
14fd8202a5b is described below
commit 14fd8202a5bf5c43b87cdd57cf941e1b77b9775b
Author: deepinsight coder <[email protected]>
AuthorDate: Thu Aug 27 05:16:04 2026 -0700
Honor gzip and mime_type on SFTPToGCS stream path (#72132)
use_stream already skips the tempfile. It still dropped gzip and
mime_type, and docs never mentioned the flag. Apply those on the
blob writer, pin tests against a tempfile regression, and floor
the sftp extra at 5.3.3 for BlobWriter getfo.
closes: #34995
---
providers/google/README.rst | 2 +-
providers/google/docs/index.rst | 2 +-
.../google/docs/operators/transfer/sftp_to_gcs.rst | 14 ++++
providers/google/pyproject.toml | 2 +-
.../google/cloud/transfers/sftp_to_gcs.py | 21 ++++--
.../system/google/cloud/gcs/example_sftp_to_gcs.py | 12 ++++
.../gcs/resources/openlineage/sftp_to_gcs.json | 36 ++++++++++
.../google/cloud/transfers/test_sftp_to_gcs.py | 80 +++++++++++++++++++++-
8 files changed, 160 insertions(+), 9 deletions(-)
diff --git a/providers/google/README.rst b/providers/google/README.rst
index 0602b6f5b00..40545c6684f 100644
--- a/providers/google/README.rst
+++ b/providers/google/README.rst
@@ -198,7 +198,7 @@ Extra Dependencies
``postgres`` ``apache-airflow-providers-postgres``
``presto`` ``apache-airflow-providers-presto``
``salesforce`` ``apache-airflow-providers-salesforce``
-``sftp`` ``apache-airflow-providers-sftp``
+``sftp`` ``apache-airflow-providers-sftp>=5.3.3``
``ssh`` ``apache-airflow-providers-ssh``
``trino`` ``apache-airflow-providers-trino``
``http`` ``apache-airflow-providers-http``
diff --git a/providers/google/docs/index.rst b/providers/google/docs/index.rst
index 516d2d2ae92..38e66e3af24 100644
--- a/providers/google/docs/index.rst
+++ b/providers/google/docs/index.rst
@@ -259,7 +259,7 @@ Extra Dependencies
``postgres`` ``apache-airflow-providers-postgres``
``presto`` ``apache-airflow-providers-presto``
``salesforce`` ``apache-airflow-providers-salesforce``
-``sftp`` ``apache-airflow-providers-sftp``
+``sftp`` ``apache-airflow-providers-sftp>=5.3.3``
``ssh`` ``apache-airflow-providers-ssh``
``trino`` ``apache-airflow-providers-trino``
``http`` ``apache-airflow-providers-http``
diff --git a/providers/google/docs/operators/transfer/sftp_to_gcs.rst
b/providers/google/docs/operators/transfer/sftp_to_gcs.rst
index 12a5f5f8581..ac3dffc40a4 100644
--- a/providers/google/docs/operators/transfer/sftp_to_gcs.rst
+++ b/providers/google/docs/operators/transfer/sftp_to_gcs.rst
@@ -52,6 +52,20 @@ The following Operator copies a single file.
:start-after: [START howto_operator_sftp_to_gcs_copy_single_file]
:end-before: [END howto_operator_sftp_to_gcs_copy_single_file]
+Streaming large files
+---------------------
+
+Set ``use_stream=True`` to stream from SFTP into GCS without a local tempfile.
+This avoids worker disk usage. Paramiko prefetch (``sftp_prefetch``, default
+``True``) can still queue the whole file in RAM, so for large files also set
+``sftp_prefetch=False``. ``gzip`` and ``mime_type`` apply on the stream path.
+
+.. exampleinclude::
/../../google/tests/system/google/cloud/gcs/example_sftp_to_gcs.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_sftp_to_gcs_copy_single_file_stream]
+ :end-before: [END howto_operator_sftp_to_gcs_copy_single_file_stream]
+
Moving a single file
--------------------
diff --git a/providers/google/pyproject.toml b/providers/google/pyproject.toml
index 7eec07d5df8..a75a49acdc6 100644
--- a/providers/google/pyproject.toml
+++ b/providers/google/pyproject.toml
@@ -205,7 +205,7 @@ dependencies = [
"apache-airflow-providers-salesforce"
]
"sftp" = [
- "apache-airflow-providers-sftp"
+ "apache-airflow-providers-sftp>=5.3.3"
]
"ssh" = [
"apache-airflow-providers-ssh"
diff --git
a/providers/google/src/airflow/providers/google/cloud/transfers/sftp_to_gcs.py
b/providers/google/src/airflow/providers/google/cloud/transfers/sftp_to_gcs.py
index f26fe9f0be2..c7a09d3aed5 100644
---
a/providers/google/src/airflow/providers/google/cloud/transfers/sftp_to_gcs.py
+++
b/providers/google/src/airflow/providers/google/cloud/transfers/sftp_to_gcs.py
@@ -19,6 +19,7 @@
from __future__ import annotations
+import gzip
import os
from collections.abc import Sequence
from functools import cached_property
@@ -59,8 +60,9 @@ class SFTPToGCSOperator(BaseOperator):
:param gcp_conn_id: (Optional) The connection ID used to connect to Google
Cloud.
:param sftp_conn_id: The sftp connection id. The name or identifier for
establishing a connection to the SFTP server.
- :param mime_type: The mime-type string
- :param gzip: Allows for file to be compressed and uploaded as gzip
+ :param mime_type: The mime-type string. Applied on both the tempfile and
stream paths.
+ :param gzip: Compress the object body before upload. Applied on both the
tempfile and stream
+ paths. The object name is unchanged; ``Content-Encoding`` is not set.
:param move_object: When move object is True, the object is moved instead
of copied to the new location. This is the equivalent of a mv command
as opposed to a cp command.
@@ -73,10 +75,13 @@ class SFTPToGCSOperator(BaseOperator):
Service Account Token Creator IAM role to the directly preceding
identity, with first
account from the list granting this role to the originating account
(templated).
:param sftp_prefetch: Whether to enable SFTP prefetch, the default is True.
+ Prefetch can queue the whole file in worker RAM; for large files prefer
+ ``use_stream=True`` and ``sftp_prefetch=False``.
:param use_stream: Determines the transfer method from SFTP to GCS.
When ``False`` (default), the file downloads locally
then uploads (may require significant disk space).
When ``True``, the file streams directly without using local disk.
+ Streaming skips disk but does not bound RAM if prefetch is on.
Defaults to ``False``.
:param fail_on_file_not_exist: If True, operator fails when file does not
exist,
if False, operator will not fail and skips transfer. Default is True.
@@ -185,8 +190,16 @@ class SFTPToGCSOperator(BaseOperator):
if self.use_stream:
dest_bucket = gcs_hook.get_bucket(self.destination_bucket)
dest_blob = dest_bucket.blob(destination_object)
- with dest_blob.open("wb") as write_stream:
- sftp_hook.retrieve_file(source_path, write_stream,
prefetch=self.sftp_prefetch)
+ dest_blob.content_type = self.mime_type
+ # GzipFile.close() flushes the wrapped writer; BlobWriter.flush()
raises unless
+ # ignore_flush is set.
+ open_kwargs = {"ignore_flush": True} if self.gzip else {}
+ with dest_blob.open("wb", **open_kwargs) as write_stream:
+ if self.gzip:
+ with gzip.GzipFile(fileobj=write_stream, mode="wb") as
compressed:
+ sftp_hook.retrieve_file(source_path, compressed,
prefetch=self.sftp_prefetch)
+ else:
+ sftp_hook.retrieve_file(source_path, write_stream,
prefetch=self.sftp_prefetch)
else:
with NamedTemporaryFile("w") as tmp:
sftp_hook.retrieve_file(source_path, tmp.name,
prefetch=self.sftp_prefetch)
diff --git
a/providers/google/tests/system/google/cloud/gcs/example_sftp_to_gcs.py
b/providers/google/tests/system/google/cloud/gcs/example_sftp_to_gcs.py
index db907c9e9bd..319e65c01b4 100644
--- a/providers/google/tests/system/google/cloud/gcs/example_sftp_to_gcs.py
+++ b/providers/google/tests/system/google/cloud/gcs/example_sftp_to_gcs.py
@@ -80,6 +80,17 @@ with DAG(
)
# [END howto_operator_sftp_to_gcs_copy_single_file]
+ # [START howto_operator_sftp_to_gcs_copy_single_file_stream]
+ copy_file_from_sftp_to_gcs_stream = SFTPToGCSOperator(
+ task_id="file-copy-sftp-to-gcs-stream",
+ source_path=f"{FILE_LOCAL_PATH}/{OBJECT_SRC_1}",
+ destination_bucket=BUCKET_NAME,
+ destination_path="streamed/parent-1.bin",
+ use_stream=True,
+ sftp_prefetch=False,
+ )
+ # [END howto_operator_sftp_to_gcs_copy_single_file_stream]
+
# [START howto_operator_sftp_to_gcs_move_single_file_destination]
move_file_from_sftp_to_gcs_destination = SFTPToGCSOperator(
task_id="file-move-sftp-to-gcs-destination",
@@ -123,6 +134,7 @@ with DAG(
unzip_file,
# TEST BODY
copy_file_from_sftp_to_gcs,
+ copy_file_from_sftp_to_gcs_stream,
move_file_from_sftp_to_gcs_destination,
copy_directory_from_sftp_to_gcs,
move_specific_files_from_sftp_to_gcs,
diff --git
a/providers/google/tests/system/google/cloud/gcs/resources/openlineage/sftp_to_gcs.json
b/providers/google/tests/system/google/cloud/gcs/resources/openlineage/sftp_to_gcs.json
index e4ea9c1b349..bc2f09bd794 100644
---
a/providers/google/tests/system/google/cloud/gcs/resources/openlineage/sftp_to_gcs.json
+++
b/providers/google/tests/system/google/cloud/gcs/resources/openlineage/sftp_to_gcs.json
@@ -35,6 +35,42 @@
}
]
},
+ {
+ "eventType": "START",
+ "job": {
+ "name": "example_sftp_to_gcs.file-copy-sftp-to-gcs-stream"
+ },
+ "inputs": [
+ {
+ "namespace": "file://localhost:22",
+ "name": "{{
result.endswith('airflow/providers/google/tests/system/google/cloud/gcs/resources/tmp/tests_sftp_hook_dir/parent-1.bin')
}}"
+ }
+ ],
+ "outputs": [
+ {
+ "namespace": "gs://bucket-example_sftp_to_gcs-{{
env_var('SYSTEM_TESTS_ENV_ID', 'default') }}",
+ "name": "streamed/parent-1.bin"
+ }
+ ]
+ },
+ {
+ "eventType": "COMPLETE",
+ "job": {
+ "name": "example_sftp_to_gcs.file-copy-sftp-to-gcs-stream"
+ },
+ "inputs": [
+ {
+ "namespace": "file://localhost:22",
+ "name": "{{
result.endswith('airflow/providers/google/tests/system/google/cloud/gcs/resources/tmp/tests_sftp_hook_dir/parent-1.bin')
}}"
+ }
+ ],
+ "outputs": [
+ {
+ "namespace": "gs://bucket-example_sftp_to_gcs-{{
env_var('SYSTEM_TESTS_ENV_ID', 'default') }}",
+ "name": "streamed/parent-1.bin"
+ }
+ ]
+ },
{
"eventType": "START",
"job": {
diff --git
a/providers/google/tests/unit/google/cloud/transfers/test_sftp_to_gcs.py
b/providers/google/tests/unit/google/cloud/transfers/test_sftp_to_gcs.py
index 177345d033b..8e246df8bac 100644
--- a/providers/google/tests/unit/google/cloud/transfers/test_sftp_to_gcs.py
+++ b/providers/google/tests/unit/google/cloud/transfers/test_sftp_to_gcs.py
@@ -87,9 +87,11 @@ class TestSFTPToGCSOperator:
sftp_hook.return_value.delete_file.assert_not_called()
+
@mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.NamedTemporaryFile")
@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_stream(self, sftp_hook, gcs_hook):
+ def test_execute_copy_single_file_with_stream(self, sftp_hook, gcs_hook,
named_temp):
+ mime_type = "text/plain"
task = SFTPToGCSOperator(
task_id=TASK_ID,
source_path=SOURCE_OBJECT_NO_WILDCARD,
@@ -99,6 +101,7 @@ class TestSFTPToGCSOperator:
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
+ mime_type=mime_type,
use_stream=True,
)
@@ -108,15 +111,56 @@ class TestSFTPToGCSOperator:
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
+ blob = gcs_hook.return_value.get_bucket.return_value.blob.return_value
+ write_stream = blob.open.return_value.__enter__.return_value
gcs_hook.return_value.get_bucket.assert_called_once_with(TEST_BUCKET)
gcs_hook.return_value.get_bucket.return_value.blob.assert_called_once_with(DESTINATION_PATH_FILE)
+ assert blob.content_type == mime_type
+ blob.open.assert_called_once_with("wb")
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, prefetch=True
+ os.path.join(SOURCE_OBJECT_NO_WILDCARD), write_stream,
prefetch=True
)
+ sink = sftp_hook.return_value.retrieve_file.call_args.args[1]
+ assert not isinstance(sink, str)
+ assert hasattr(sink, "write")
+ named_temp.assert_not_called()
gcs_hook.return_value.upload.assert_not_called()
sftp_hook.return_value.delete_file.assert_not_called()
+
@mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.gzip.GzipFile")
+
@mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.NamedTemporaryFile")
+ @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_stream_and_gzip(
+ self, sftp_hook, gcs_hook, named_temp, gzip_file_cls
+ ):
+ task = SFTPToGCSOperator(
+ task_id=TASK_ID,
+ source_path=SOURCE_OBJECT_NO_WILDCARD,
+ destination_bucket=TEST_BUCKET,
+ destination_path=DESTINATION_PATH_FILE,
+ gcp_conn_id=GCP_CONN_ID,
+ sftp_conn_id=SFTP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ gzip=True,
+ use_stream=True,
+ )
+ compressed = gzip_file_cls.return_value.__enter__.return_value
+ blob = gcs_hook.return_value.get_bucket.return_value.blob.return_value
+ write_stream = blob.open.return_value.__enter__.return_value
+
+ task.execute(None)
+
+ assert blob.content_type == DEFAULT_MIME_TYPE
+ blob.open.assert_called_once_with("wb", ignore_flush=True)
+ gzip_file_cls.assert_called_once_with(fileobj=write_stream, mode="wb")
+ sftp_hook.return_value.retrieve_file.assert_called_once_with(
+ os.path.join(SOURCE_OBJECT_NO_WILDCARD), compressed, prefetch=True
+ )
+ named_temp.assert_not_called()
+ gcs_hook.return_value.upload.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):
@@ -238,6 +282,38 @@ class TestSFTPToGCSOperator:
]
)
+
@mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.NamedTemporaryFile")
+ @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_with_wildcard_and_stream(self, sftp_hook, gcs_hook,
named_temp):
+ sftp_hook.return_value.get_tree_map.return_value = [
+ ["main_dir/test_object3.json",
"main_dir/sub_dir/test_object3.json"],
+ [],
+ [],
+ ]
+
+ task = SFTPToGCSOperator(
+ task_id=TASK_ID,
+ source_path=SOURCE_OBJECT_WILDCARD_FILENAME,
+ destination_bucket=TEST_BUCKET,
+ destination_path=DESTINATION_PATH_DIR,
+ gcp_conn_id=GCP_CONN_ID,
+ sftp_conn_id=SFTP_CONN_ID,
+ use_stream=True,
+ )
+ task.execute(None)
+
+ named_temp.assert_not_called()
+ gcs_hook.return_value.upload.assert_not_called()
+ assert sftp_hook.return_value.retrieve_file.call_count == 2
+ for call in sftp_hook.return_value.retrieve_file.call_args_list:
+ sink = call.args[1]
+ assert not isinstance(sink, str)
+ assert hasattr(sink, "write")
+ blob = gcs_hook.return_value.get_bucket.return_value.blob.return_value
+ assert blob.open.call_count == 2
+ blob.open.assert_called_with("wb")
+
@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_move_with_wildcard(self, sftp_hook, gcs_hook):