From 30709e3ec9e7a1aa8d45c8cf66eff74469afc3ae Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Sat, 25 Jul 2026 14:43:16 +0530 Subject: [PATCH] Fail deferred Cloud Composer tasks when the GCP operation errors --- .../google/cloud/triggers/cloud_composer.py | 7 ++- .../cloud/triggers/test_cloud_composer.py | 59 +++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/providers/google/src/airflow/providers/google/cloud/triggers/cloud_composer.py b/providers/google/src/airflow/providers/google/cloud/triggers/cloud_composer.py index 040e35f2a3859..7269a4a88a2f4 100644 --- a/providers/google/src/airflow/providers/google/cloud/triggers/cloud_composer.py +++ b/providers/google/src/airflow/providers/google/cloud/triggers/cloud_composer.py @@ -80,9 +80,12 @@ async def run(self): while True: operation = await self.gcp_hook.get_operation(operation_name=self.operation_name) if operation.done: + # A long-running operation signals failure as done=True with ``error`` set; while + # it is still running neither ``error`` nor ``response`` is populated. Checking + # ``error`` only in the not-done branch would therefore never see a failure. + if operation.error.message: + raise AirflowException(f"Cloud Composer Environment error: {operation.error.message}") break - elif operation.error.message: - raise AirflowException(f"Cloud Composer Environment error: {operation.error.message}") await asyncio.sleep(self.pooling_period_seconds) yield TriggerEvent( { diff --git a/providers/google/tests/unit/google/cloud/triggers/test_cloud_composer.py b/providers/google/tests/unit/google/cloud/triggers/test_cloud_composer.py index f19fb6e2ca414..2e016158bf09a 100644 --- a/providers/google/tests/unit/google/cloud/triggers/test_cloud_composer.py +++ b/providers/google/tests/unit/google/cloud/triggers/test_cloud_composer.py @@ -18,15 +18,18 @@ from __future__ import annotations from datetime import datetime +from types import SimpleNamespace from unittest import mock import pytest from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.models import Connection +from airflow.providers.common.compat.sdk import AirflowException from airflow.providers.google.cloud.triggers.cloud_composer import ( CloudComposerAirflowCLICommandTrigger, CloudComposerDAGRunTrigger, + CloudComposerExecutionTrigger, CloudComposerExternalTaskTrigger, ) from airflow.triggers.base import TriggerEvent @@ -59,6 +62,23 @@ "output_end": True, "exit_info": {"exit_code": 0, "error": ""}, } +TEST_OPERATION_NAME = "test_operation_name" +TEST_POOLING_PERIOD_SECONDS = 30 +TEST_OPERATION_ERROR_MESSAGE = "Environment creation failed: quota exceeded" + + +def make_operation(*, done: bool, error_message: str = ""): + """ + Build a stand-in for a ``google.longrunning.Operation``. + + Per the LRO contract, a finished operation has ``done=True`` and exactly one of + ``error`` or ``response`` set; while it is still running neither is set. + """ + return SimpleNamespace( + name=TEST_OPERATION_NAME, + done=done, + error=SimpleNamespace(message=error_message), + ) @pytest.fixture @@ -127,6 +147,45 @@ def external_task_trigger(mock_conn): ) +@pytest.fixture +@mock.patch( + "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.get_connection", + return_value=Connection(conn_id="test_conn"), +) +def execution_trigger(mock_conn): + return CloudComposerExecutionTrigger( + project_id=TEST_PROJECT_ID, + region=TEST_LOCATION, + operation_name=TEST_OPERATION_NAME, + gcp_conn_id=TEST_GCP_CONN_ID, + impersonation_chain=TEST_IMPERSONATION_CHAIN, + pooling_period_seconds=TEST_POOLING_PERIOD_SECONDS, + ) + + +class TestCloudComposerExecutionTrigger: + @pytest.mark.asyncio + @mock.patch("airflow.providers.google.cloud.hooks.cloud_composer.CloudComposerAsyncHook.get_operation") + async def test_run_raises_when_operation_finished_with_error(self, mock_get_operation, execution_trigger): + mock_get_operation.return_value = make_operation( + done=True, error_message=TEST_OPERATION_ERROR_MESSAGE + ) + + with pytest.raises(AirflowException, match=TEST_OPERATION_ERROR_MESSAGE): + await execution_trigger.run().asend(None) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.google.cloud.hooks.cloud_composer.CloudComposerAsyncHook.get_operation") + async def test_run_yields_event_when_operation_finished_without_error( + self, mock_get_operation, execution_trigger + ): + mock_get_operation.return_value = make_operation(done=True) + + actual_event = await execution_trigger.run().asend(None) + + assert actual_event == TriggerEvent({"operation_name": TEST_OPERATION_NAME, "operation_done": True}) + + class TestCloudComposerAirflowCLICommandTrigger: def test_serialize(self, cli_command_trigger): actual_data = cli_command_trigger.serialize()