This is an automated email from the ASF dual-hosted git repository.
shahar1 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 91353550246 Honor download_num_attempts and upload_num_attempts in
GCSTimeSpanFileTransformOperator (#72953)
91353550246 is described below
commit 91353550246cc27bd1a225bbdbe68e16b34ed538
Author: Bingqin Wang <[email protected]>
AuthorDate: Sat Oct 3 02:26:57 2026 -0500
Honor download_num_attempts and upload_num_attempts in
GCSTimeSpanFileTransformOperator (#72953)
---
.../providers/google/cloud/operators/gcs.py | 24 +++-
.../tests/unit/google/cloud/operators/test_gcs.py | 135 +++++++++++++++++++++
2 files changed, 156 insertions(+), 3 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/gcs.py
b/providers/google/src/airflow/providers/google/cloud/operators/gcs.py
index 59e077e2406..9265f23ac0a 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/gcs.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/gcs.py
@@ -20,6 +20,7 @@
from __future__ import annotations
import datetime
+import logging
import subprocess
import sys
import warnings
@@ -30,6 +31,7 @@ from tempfile import NamedTemporaryFile, TemporaryDirectory
from typing import TYPE_CHECKING
import pendulum
+import tenacity
from airflow.providers.google.common.hooks.base_google import
PROVIDE_PROJECT_ID
@@ -985,6 +987,16 @@ class
GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
# Get storage client once (storage.Client is thread-safe for
concurrent requests).
client = source_hook.get_conn()
+ # A transient ``GoogleCloudError`` retries the whole download.
``reraise`` keeps the
+ # original error after the last attempt, for the
``download_continue_on_fail``
+ # handling below. The waits (2, 4, 8 s, ...) match
``GCSHook.download``.
+ @tenacity.retry(
+ stop=tenacity.stop_after_attempt(self.download_num_attempts),
+ wait=tenacity.wait_exponential(multiplier=2, max=60),
+ retry=tenacity.retry_if_exception_type(GoogleCloudError),
+ before_sleep=tenacity.before_sleep_log(self.log,
logging.WARNING),
+ reraise=True,
+ )
def _download(blob_name: str):
bucket = client.bucket(bucket_name=self.source_bucket)
@@ -1065,6 +1077,14 @@ class
GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
# Get storage client once (storage.Client is thread-safe for
concurrent requests).
client = destination_hook.get_conn()
+ # Same retry policy as the downloads, for ``upload_num_attempts``.
+ @tenacity.retry(
+ stop=tenacity.stop_after_attempt(self.upload_num_attempts),
+ wait=tenacity.wait_exponential(multiplier=2, max=60),
+ retry=tenacity.retry_if_exception_type(GoogleCloudError),
+ before_sleep=tenacity.before_sleep_log(self.log,
logging.WARNING),
+ reraise=True,
+ )
def _upload(upload_file: Path):
bucket = client.bucket(bucket_name=self.destination_bucket)
@@ -1077,9 +1097,7 @@ class
GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
blob = bucket.blob(blob_name=upload_file_name,
chunk_size=self.chunk_size)
- blob.upload_from_filename(
- filename=str(upload_file),
- )
+ blob.upload_from_filename(filename=str(upload_file))
return upload_file_name
diff --git a/providers/google/tests/unit/google/cloud/operators/test_gcs.py
b/providers/google/tests/unit/google/cloud/operators/test_gcs.py
index 4d53d254e49..fcd17462ead 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_gcs.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_gcs.py
@@ -521,6 +521,16 @@ class TestGCSTimeSpanFileTransformOperator:
mock_blob = mock_bucket.blob.return_value
return mock_client, mock_bucket, mock_blob
+ @staticmethod
+ def _setup_transform_process(mock_subprocess):
+ mock_proc = mock.MagicMock()
+ mock_proc.returncode = 0
+ mock_proc.stdout.readline = lambda: b""
+ mock_proc.wait.return_value = None
+ mock_subprocess.Popen.return_value.__enter__.return_value = mock_proc
+ mock_subprocess.PIPE = "pipe"
+ mock_subprocess.STDOUT = "stdout"
+
@mock.patch("airflow.providers.google.cloud.operators.gcs.TemporaryDirectory")
@mock.patch("airflow.providers.google.cloud.operators.gcs.subprocess")
@mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
@@ -1151,6 +1161,131 @@ class TestGCSTimeSpanFileTransformOperator:
op.execute(context=context)
mock_blob.download_to_filename.assert_not_called()
+ @pytest.mark.parametrize(
+ ("num_attempts", "outcomes", "succeeds", "expected_sleeps"),
+ [
+ (1, [GoogleCloudError("fail")], False, []),
+ (2, [GoogleCloudError("fail"), GoogleCloudError("fail")], False,
[2]),
+ (3, [GoogleCloudError("fail"), None], True, [2]),
+ ],
+ )
+ @mock.patch("tenacity.nap.time.sleep")
+
@mock.patch("airflow.providers.google.cloud.operators.gcs.TemporaryDirectory")
+ @mock.patch("airflow.providers.google.cloud.operators.gcs.subprocess")
+ @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
+ def test_download_honors_download_num_attempts(
+ self,
+ mock_hook,
+ mock_subprocess,
+ mock_tempdir,
+ mock_sleep,
+ num_attempts,
+ outcomes,
+ succeeds,
+ expected_sleeps,
+ ):
+ timespan_start = datetime(2015, 2, 1, tzinfo=timezone.utc)
+ context = {
+ "logical_date": timespan_start,
+ "data_interval_start": timespan_start,
+ "data_interval_end": timespan_start + timedelta(hours=1),
+ "ti": mock.Mock(),
+ "task": mock.MagicMock(),
+ }
+ mock_tempdir.return_value.__enter__.side_effect = ["source",
"destination"]
+ mock_hook.return_value.list_by_timespan.return_value = ["file1"]
+ _, _, mock_blob = self._setup_gcs_client_chain(mock_hook)
+ mock_blob.download_to_filename.side_effect = outcomes
+ self._setup_transform_process(mock_subprocess)
+
+ op = GCSTimeSpanFileTransformOperator(
+ task_id=TASK_ID,
+ source_bucket="bucket",
+ source_prefix="prefix",
+ source_gcp_conn_id="",
+ destination_bucket="dest",
+ destination_prefix="dest",
+ destination_gcp_conn_id="",
+ transform_script="script.py",
+ download_num_attempts=num_attempts,
+ )
+
+ with (
+ mock.patch.object(Path, "glob") as path_glob,
+ mock.patch.object(Path, "is_file", return_value=True),
+ ):
+ path_glob.return_value.__iter__.return_value = []
+ if succeeds:
+ op.execute(context=context)
+ else:
+ with pytest.raises(GoogleCloudError):
+ op.execute(context=context)
+
+ assert mock_blob.download_to_filename.call_count == len(outcomes)
+ assert [call.args[0] for call in mock_sleep.call_args_list] ==
expected_sleeps
+
+ @pytest.mark.parametrize(
+ ("num_attempts", "outcomes", "succeeds", "expected_sleeps"),
+ [
+ (1, [GoogleCloudError("fail")], False, []),
+ (2, [GoogleCloudError("fail"), None], True, [2]),
+ ],
+ )
+ @mock.patch("tenacity.nap.time.sleep")
+
@mock.patch("airflow.providers.google.cloud.operators.gcs.TemporaryDirectory")
+ @mock.patch("airflow.providers.google.cloud.operators.gcs.subprocess")
+ @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
+ def test_upload_honors_upload_num_attempts(
+ self,
+ mock_hook,
+ mock_subprocess,
+ mock_tempdir,
+ mock_sleep,
+ num_attempts,
+ outcomes,
+ succeeds,
+ expected_sleeps,
+ ):
+ timespan_start = datetime(2015, 2, 1, tzinfo=timezone.utc)
+ context = {
+ "logical_date": timespan_start,
+ "data_interval_start": timespan_start,
+ "data_interval_end": timespan_start + timedelta(hours=1),
+ "ti": mock.Mock(),
+ "task": mock.MagicMock(),
+ }
+ mock_tempdir.return_value.__enter__.side_effect = ["source",
"destination"]
+ mock_hook.return_value.list_by_timespan.return_value = []
+ _, _, mock_blob = self._setup_gcs_client_chain(mock_hook)
+ mock_blob.upload_from_filename.side_effect = outcomes
+ self._setup_transform_process(mock_subprocess)
+
+ op = GCSTimeSpanFileTransformOperator(
+ task_id=TASK_ID,
+ source_bucket="bucket",
+ source_prefix="prefix",
+ source_gcp_conn_id="",
+ destination_bucket="dest",
+ destination_prefix="dest",
+ destination_gcp_conn_id="",
+ transform_script="script.py",
+ upload_num_attempts=num_attempts,
+ )
+
+ with (
+ mock.patch.object(Path, "glob") as path_glob,
+ mock.patch.object(Path, "is_file", return_value=True),
+ ):
+ path_glob.return_value.__iter__.return_value =
[Path("destination/file1")]
+ if succeeds:
+ assert op.execute(context=context) == ["dest/file1"]
+ else:
+ with pytest.raises(GoogleCloudError):
+ op.execute(context=context)
+
+ assert mock_blob.upload_from_filename.call_count == len(outcomes)
+ assert [call.args[0] for call in mock_sleep.call_args_list] ==
expected_sleeps
+
class TestGCSDeleteBucketOperator:
@mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
