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 3d42610a26e Fix Cloud Composer create env operator (#64956)
3d42610a26e is described below
commit 3d42610a26e4a8e49e2cf3d9eb6fc5ad712f180e
Author: Ulada Zakharava <[email protected]>
AuthorDate: Mon Apr 27 01:20:06 2026 +0000
Fix Cloud Composer create env operator (#64956)
---
.../google/docs/operators/cloud/cloud_composer.rst | 14 ++
.../google/cloud/operators/cloud_composer.py | 44 +++++-
.../cloud/composer/example_cloud_composer.py | 69 +++++++++
.../google/cloud/operators/test_cloud_composer.py | 170 +++++++++++++++++++++
4 files changed, 296 insertions(+), 1 deletion(-)
diff --git a/providers/google/docs/operators/cloud/cloud_composer.rst
b/providers/google/docs/operators/cloud/cloud_composer.rst
index 7a0b1ad9349..8a95c3d9121 100644
--- a/providers/google/docs/operators/cloud/cloud_composer.rst
+++ b/providers/google/docs/operators/cloud/cloud_composer.rst
@@ -48,6 +48,10 @@ A simple environment configuration can look as followed:
With this configuration we can create the environment:
:class:`~airflow.providers.google.cloud.operators.cloud_composer.CloudComposerCreateEnvironmentOperator`
+The create operator only succeeds after the Composer environment reaches the
``RUNNING`` state.
+If the long-running create operation finishes but the environment remains in
another state such as
+``ERROR`` or ``CREATING``, the task fails so downstream tasks do not run
against an unusable environment.
+
.. exampleinclude::
/../../google/tests/system/google/cloud/composer/example_cloud_composer.py
:language: python
:dedent: 4
@@ -63,6 +67,16 @@ or you can define the same operator in the deferrable mode:
:start-after: [START
howto_operator_create_composer_environment_deferrable_mode]
:end-before: [END
howto_operator_create_composer_environment_deferrable_mode]
+For retry-heavy system tests, you can clean up a failed environment before
retrying the create task.
+The example below only deletes environments that are already in the ``ERROR``
state and leaves other
+states untouched.
+
+.. exampleinclude::
/../../google/tests/system/google/cloud/composer/example_cloud_composer.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_operator_composer_retry_cleanup]
+ :end-before: [END howto_operator_composer_retry_cleanup]
+
Get an environment
------------------
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/cloud_composer.py
b/providers/google/src/airflow/providers/google/cloud/operators/cloud_composer.py
index c426e7d4264..80281cede49 100644
---
a/providers/google/src/airflow/providers/google/cloud/operators/cloud_composer.py
+++
b/providers/google/src/airflow/providers/google/cloud/operators/cloud_composer.py
@@ -65,6 +65,10 @@ class CloudComposerEnvironmentsLink(BaseGoogleLink):
format_str = CLOUD_COMPOSER_ENVIRONMENTS_LINK
+class ComposerEnvironmentNotRunningError(RuntimeError):
+ """Raised when a Composer environment is not in the RUNNING state."""
+
+
class CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
"""
Create a new environment.
@@ -137,6 +141,22 @@ class
CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
"environment_id": self.environment_id,
}
+ def _raise_if_environment_not_running(self, environment: Environment) ->
None:
+ env_state = Environment.State(environment.state)
+ self.log.info("Environment state: %s", env_state.name)
+
+ self.log.info(
+ "Composer environment final state: raw=%s, name=%s, value=%s",
+ environment.state,
+ env_state.name,
+ env_state.value,
+ )
+
+ if env_state != Environment.State.RUNNING:
+ raise ComposerEnvironmentNotRunningError(
+ f"Composer environment is not in RUNNING state. Current state:
{env_state.name}"
+ )
+
def execute(self, context: Context):
hook = CloudComposerHook(
gcp_conn_id=self.gcp_conn_id,
@@ -162,7 +182,22 @@ class
CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
context["ti"].xcom_push(key="operation_id",
value=result.operation.name)
if not self.deferrable:
- environment = hook.wait_for_operation(timeout=self.timeout,
operation=result)
+ self.log.info("Waiting for the create environment operation to
complete...")
+ hook.wait_for_operation(timeout=self.timeout, operation=result)
+ self.log.info(
+ "Create environment operation completed. Checking the
state of the environment..."
+ )
+ environment = hook.get_environment(
+ project_id=self.project_id,
+ region=self.region,
+ environment_id=self.environment_id,
+ retry=self.retry,
+ timeout=self.timeout,
+ metadata=self.metadata,
+ )
+ self._raise_if_environment_not_running(environment)
+
+ self.log.info("Environment created and is in RUNNING state")
return Environment.to_dict(environment)
self.defer(
trigger=CloudComposerExecutionTrigger(
@@ -184,6 +219,10 @@ class
CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+ self.log.info("Environment already exists. Checking its state...")
+ self._raise_if_environment_not_running(environment)
+
+ self.log.info("Environment already exists and is in RUNNING state")
return Environment.to_dict(environment)
def execute_complete(self, context: Context, event: dict):
@@ -201,6 +240,9 @@ class
CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+ self._raise_if_environment_not_running(env)
+
+ self.log.info("Environment created and is in RUNNING state")
return Environment.to_dict(env)
raise AirflowException(f"Unexpected error in the operation:
{event['operation_name']}")
diff --git
a/providers/google/tests/system/google/cloud/composer/example_cloud_composer.py
b/providers/google/tests/system/google/cloud/composer/example_cloud_composer.py
index 110f99251e1..31d455ef1bf 100644
---
a/providers/google/tests/system/google/cloud/composer/example_cloud_composer.py
+++
b/providers/google/tests/system/google/cloud/composer/example_cloud_composer.py
@@ -17,9 +17,13 @@
# under the License.
from __future__ import annotations
+import logging
import os
from datetime import datetime, timedelta
+from typing import Any
+from google.api_core.exceptions import NotFound
+from google.cloud.orchestration.airflow.service_v1.types import Environment
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
@@ -33,6 +37,7 @@ else:
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.google.cloud.hooks.cloud_composer import
CloudComposerHook
from airflow.providers.google.cloud.operators.cloud_composer import (
CloudComposerCreateEnvironmentOperator,
CloudComposerDeleteEnvironmentOperator,
@@ -86,6 +91,7 @@ UPDATE_MASK = {"paths": ["labels.label"]}
# [END howto_operator_composer_update_environment]
COMMAND = "dags list -o json --verbose"
+log = logging.getLogger(__name__)
@task(task_id="get_project_number")
@@ -104,6 +110,63 @@ def get_project_number():
raise exc
+# [START howto_operator_composer_retry_cleanup]
+def cleanup_failed_environment_before_retry(context: dict[str, Any]) -> None:
+ task = context["task"]
+
+ hook = CloudComposerHook(
+ gcp_conn_id=task.gcp_conn_id,
+ impersonation_chain=task.impersonation_chain,
+ )
+
+ environment_id = task.environment_id
+ project_id = task.project_id
+ region = task.region
+
+ log.info(
+ "Retry cleanup started for Composer env. project_id=%s region=%s
environment_id=%s",
+ project_id,
+ region,
+ environment_id,
+ )
+
+ try:
+ environment = hook.get_environment(
+ project_id=project_id,
+ region=region,
+ environment_id=environment_id,
+ )
+ except NotFound:
+ log.info("Environment does not exist. Nothing to clean up.")
+ return
+
+ env_state = Environment.State(environment.state)
+ log.info(
+ "Environment exists before retry. state=%s value=%s",
+ env_state.name,
+ env_state.value,
+ )
+
+ if env_state != Environment.State.ERROR:
+ log.info(
+ "Skipping cleanup before retry because environment is not in ERROR
state. current_state=%s",
+ env_state.name,
+ )
+ return
+
+ log.info("Deleting Composer environment %s in ERROR state before retry.",
environment_id)
+ delete_operation = hook.delete_environment(
+ project_id=project_id,
+ region=region,
+ environment_id=environment_id,
+ )
+ hook.wait_for_operation(operation=delete_operation)
+ log.info("Environment %s deleted before retry.", environment_id)
+
+
+# [END howto_operator_composer_retry_cleanup]
+
+
with DAG(
DAG_ID,
schedule="@once",
@@ -126,6 +189,9 @@ with DAG(
region=REGION,
environment_id=ENVIRONMENT_ID,
environment=ENVIRONMENT,
+ retries=1,
+ retry_delay=timedelta(minutes=1),
+ on_retry_callback=cleanup_failed_environment_before_retry,
)
# [END howto_operator_create_composer_environment]
@@ -137,6 +203,9 @@ with DAG(
environment_id=ENVIRONMENT_ID_ASYNC,
environment=ENVIRONMENT,
deferrable=True,
+ retries=1,
+ retry_delay=timedelta(minutes=1),
+ on_retry_callback=cleanup_failed_environment_before_retry,
)
# [END howto_operator_create_composer_environment_deferrable_mode]
diff --git
a/providers/google/tests/unit/google/cloud/operators/test_cloud_composer.py
b/providers/google/tests/unit/google/cloud/operators/test_cloud_composer.py
index 3642bff87b7..d7d5af767d8 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_cloud_composer.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_cloud_composer.py
@@ -19,7 +19,9 @@ from __future__ import annotations
from unittest import mock
import pytest
+from google.api_core.exceptions import AlreadyExists
from google.api_core.gapic_v1.method import DEFAULT
+from google.cloud.orchestration.airflow.service_v1.types import Environment
from airflow.providers.common.compat.sdk import TaskDeferred
from airflow.providers.google.cloud.operators.cloud_composer import (
@@ -31,6 +33,7 @@ from airflow.providers.google.cloud.operators.cloud_composer
import (
CloudComposerRunAirflowCLICommandOperator,
CloudComposerTriggerDAGRunOperator,
CloudComposerUpdateEnvironmentOperator,
+ ComposerEnvironmentNotRunningError,
)
from airflow.providers.google.cloud.triggers.cloud_composer import (
CloudComposerAirflowCLICommandTrigger,
@@ -77,9 +80,16 @@ TEST_COMPOSER_AIRFLOW_VERSION = 1
class TestCloudComposerCreateEnvironmentOperator:
+ @staticmethod
+ def _mock_environment(state: int) -> Environment:
+ return Environment(state=state)
+
@mock.patch(COMPOSER_STRING.format("Environment.to_dict"))
@mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
def test_execute(self, mock_hook, to_dict_mode) -> None:
+ environment = self._mock_environment(Environment.State.RUNNING)
+ mock_hook.return_value.get_environment.return_value = environment
+
op = CloudComposerCreateEnvironmentOperator(
task_id=TASK_ID,
project_id=TEST_GCP_PROJECT,
@@ -104,6 +114,14 @@ class TestCloudComposerCreateEnvironmentOperator:
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
+ mock_hook.return_value.get_environment.assert_called_once_with(
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ )
@mock.patch(COMPOSER_STRING.format("Environment.to_dict"))
@mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
@@ -128,6 +146,158 @@ class TestCloudComposerCreateEnvironmentOperator:
assert isinstance(exc.value.trigger, CloudComposerExecutionTrigger)
assert exc.value.method_name == GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
+ @mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
+ def test_execute_raises_when_operation_finishes_in_error_state(self,
mock_hook) -> None:
+ mock_hook.return_value.get_environment.return_value =
self._mock_environment(Environment.State.ERROR)
+
+ op = CloudComposerCreateEnvironmentOperator(
+ task_id=TASK_ID,
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ environment=TEST_ENVIRONMENT,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ )
+
+ with pytest.raises(ComposerEnvironmentNotRunningError, match="Current
state: ERROR"):
+ op.execute(mock.MagicMock())
+
+ @mock.patch(COMPOSER_STRING.format("Environment.to_dict"))
+ @mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
+ def test_execute_rechecks_environment_state_after_operation_completion(
+ self, mock_hook, to_dict_mode
+ ) -> None:
+ mock_hook.return_value.wait_for_operation.return_value =
self._mock_environment(
+ Environment.State.STATE_UNSPECIFIED
+ )
+ mock_hook.return_value.get_environment.return_value =
self._mock_environment(
+ Environment.State.RUNNING
+ )
+
+ op = CloudComposerCreateEnvironmentOperator(
+ task_id=TASK_ID,
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ environment=TEST_ENVIRONMENT,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ )
+
+ op.execute(mock.MagicMock())
+
+ mock_hook.return_value.wait_for_operation.assert_called_once()
+ mock_hook.return_value.get_environment.assert_called_once_with(
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ )
+
+ @mock.patch(COMPOSER_STRING.format("Environment.to_dict"))
+ @mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
+ def test_execute_existing_running_environment_succeeds(self, mock_hook,
to_dict_mode) -> None:
+ mock_hook.return_value.create_environment.side_effect =
AlreadyExists("already exists")
+ mock_hook.return_value.get_environment.return_value =
self._mock_environment(
+ Environment.State.RUNNING
+ )
+
+ op = CloudComposerCreateEnvironmentOperator(
+ task_id=TASK_ID,
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ environment=TEST_ENVIRONMENT,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ )
+
+ op.execute(mock.MagicMock())
+
+ mock_hook.return_value.get_environment.assert_called_once_with(
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ )
+
+ @mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
+ def test_execute_existing_environment_in_error_state_fails(self,
mock_hook) -> None:
+ mock_hook.return_value.create_environment.side_effect =
AlreadyExists("already exists")
+ mock_hook.return_value.get_environment.return_value =
self._mock_environment(Environment.State.ERROR)
+
+ op = CloudComposerCreateEnvironmentOperator(
+ task_id=TASK_ID,
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ environment=TEST_ENVIRONMENT,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ )
+
+ with pytest.raises(ComposerEnvironmentNotRunningError, match="Current
state: ERROR"):
+ op.execute(mock.MagicMock())
+
+ @mock.patch(COMPOSER_STRING.format("Environment.to_dict"))
+ @mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
+ def test_execute_complete_returns_running_environment(self, mock_hook,
to_dict_mode) -> None:
+ mock_hook.return_value.get_environment.return_value =
self._mock_environment(
+ Environment.State.RUNNING
+ )
+
+ op = CloudComposerCreateEnvironmentOperator(
+ task_id=TASK_ID,
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ environment=TEST_ENVIRONMENT,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ deferrable=True,
+ )
+
+ op.execute_complete(mock.MagicMock(), {"operation_done": True,
"operation_name": "test-operation"})
+
+ @mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
+ def test_execute_complete_raises_when_environment_not_running(self,
mock_hook) -> None:
+ mock_hook.return_value.get_environment.return_value =
self._mock_environment(
+ Environment.State.CREATING
+ )
+
+ op = CloudComposerCreateEnvironmentOperator(
+ task_id=TASK_ID,
+ project_id=TEST_GCP_PROJECT,
+ region=TEST_GCP_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ environment=TEST_ENVIRONMENT,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ retry=TEST_RETRY,
+ timeout=TEST_TIMEOUT,
+ metadata=TEST_METADATA,
+ deferrable=True,
+ )
+
+ with pytest.raises(ComposerEnvironmentNotRunningError, match="Current
state: CREATING"):
+ op.execute_complete(
+ mock.MagicMock(), {"operation_done": True, "operation_name":
"test-operation"}
+ )
+
class TestCloudComposerDeleteEnvironmentOperator:
@mock.patch(COMPOSER_STRING.format("CloudComposerHook"))
