Skip to content

Commit e75094d

Browse files
mitre88potiuk
mitre88
and
potiuk
authored
Resolve GCSFileTransformOperator destination fallbacks after rendering (#70488)
* Resolve GCSFileTransformOperator destination fallbacks after rendering destination_bucket/destination_object are template fields, so resolving their 'or source_*' fallbacks in __init__ reads the unrendered value. Move the fallback resolution to execute() where the rendered values are available, keep get_openlineage_facets_on_complete resolving fallbacks locally (it can run outside execute), and remove the operator from the validate_operators_init exemptions list. Related: #70296 * Add tests for the GCSFileTransformOperator destination fallback Cover the omitted-destination path in execute() and get_openlineage_facets_on_start(): the source values are changed after construction (as template rendering does) and the destination must follow them. Both tests fail against the previous __init__-time fallback. Generated-by: Claude Opus 5 --------- Co-authored-by: Dr Alex Mitre <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]>
1 parent af1c66d commit e75094d

3 files changed

Lines changed: 61 additions & 5 deletions

File tree

‎providers/google/src/airflow/providers/google/cloud/operators/gcs.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -692,15 +692,17 @@ def __init__(
692692
super().__init__(**kwargs)
693693
self.source_bucket = source_bucket
694694
self.source_object = source_object
695-
self.destination_bucket = destination_bucket or self.source_bucket
696-
self.destination_object = destination_object or self.source_object
695+
self.destination_bucket = destination_bucket
696+
self.destination_object = destination_object
697697

698698
self.gcp_conn_id = gcp_conn_id
699699
self.transform_script = transform_script
700700
self.output_encoding = sys.getdefaultencoding()
701701
self.impersonation_chain = impersonation_chain
702702

703703
def execute(self, context: Context) -> None:
704+
self.destination_bucket = self.destination_bucket or self.source_bucket
705+
self.destination_object = self.destination_object or self.source_object
704706
hook = GCSHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
705707

706708
with NamedTemporaryFile() as source_file, NamedTemporaryFile() as destination_file:
@@ -747,8 +749,8 @@ def get_openlineage_facets_on_start(self):
747749
name=self.source_object,
748750
)
749751
output_dataset = Dataset(
750-
namespace=f"gs://{self.destination_bucket}",
751-
name=self.destination_object,
752+
namespace=f"gs://{self.destination_bucket or self.source_bucket}",
753+
name=self.destination_object or self.source_object,
752754
)
753755

754756
return OperatorLineage(inputs=[input_dataset], outputs=[output_dataset])

‎providers/google/tests/unit/google/cloud/operators/test_gcs.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -426,6 +426,61 @@ def test_get_openlineage_facets_on_start(self):
426426
assert lineage.inputs[0] == expected_input
427427
assert lineage.outputs[0] == expected_output
428428

429+
@mock.patch("airflow.providers.google.cloud.operators.gcs.NamedTemporaryFile")
430+
@mock.patch("airflow.providers.google.cloud.operators.gcs.subprocess")
431+
@mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
432+
def test_execute_destination_falls_back_to_rendered_source(
433+
self, mock_hook, mock_subprocess, mock_tempfile
434+
):
435+
source = "source"
436+
destination = "destination"
437+
mock1 = mock.Mock()
438+
mock2 = mock.Mock()
439+
mock1.name = source
440+
mock2.name = destination
441+
mock_tempfile.return_value.__enter__.side_effect = [mock1, mock2]
442+
443+
mock_proc = mock.MagicMock()
444+
mock_proc.returncode = 0
445+
mock_proc.stdout.readline = lambda: b""
446+
mock_proc.wait.return_value = None
447+
mock_subprocess.Popen.return_value.__enter__.return_value = mock_proc
448+
449+
op = GCSFileTransformOperator(
450+
task_id=TASK_ID,
451+
source_bucket="{{ var.value.bucket }}",
452+
source_object="{{ var.value.object }}",
453+
transform_script="script.py",
454+
)
455+
# Template rendering happens after __init__; the destination fallback must see the rendered values.
456+
op.source_bucket = TEST_BUCKET
457+
op.source_object = "rendered.txt"
458+
459+
op.execute(context=mock.MagicMock())
460+
461+
mock_hook.return_value.download.assert_called_once_with(
462+
bucket_name=TEST_BUCKET, object_name="rendered.txt", filename=source
463+
)
464+
mock_hook.return_value.upload.assert_called_with(
465+
bucket_name=TEST_BUCKET, object_name="rendered.txt", filename=destination
466+
)
467+
468+
def test_get_openlineage_facets_on_start_destination_falls_back_to_rendered_source(self):
469+
operator = GCSFileTransformOperator(
470+
task_id=TASK_ID,
471+
source_bucket="{{ var.value.bucket }}",
472+
source_object="{{ var.value.object }}",
473+
transform_script="/path/to_script",
474+
)
475+
operator.source_bucket = TEST_BUCKET
476+
operator.source_object = "folder/a.txt"
477+
478+
lineage = operator.get_openlineage_facets_on_start()
479+
480+
expected = Dataset(namespace=f"gs://{TEST_BUCKET}", name="folder/a.txt")
481+
assert lineage.inputs == [expected]
482+
assert lineage.outputs == [expected]
483+
429484

430485
class TestGCSTimeSpanFileTransformOperatorDateInterpolation:
431486
def test_execute(self):

‎scripts/ci/prek/validate_operators_init_exemptions.txt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,5 @@ providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_tran
1212
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocCreateClusterOperator
1313
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocSubmitJobOperator
1414
providers/google/src/airflow/providers/google/cloud/operators/functions.py::CloudFunctionDeployFunctionOperator
15-
providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSFileTransformOperator
1615
providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py::BigQueryDataTransferServiceTransferRunSensor
1716
providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py::CloudComposerExternalTaskSensor

0 commit comments

Comments
 (0)