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 dd4d2f9ef7c Add CloudSQLNoOperationInProgressSensor for parallel admin
ops (#68151)
dd4d2f9ef7c is described below
commit dd4d2f9ef7c8e4cc12d9b10dcc3334af5dee9fb6
Author: deepinsight coder <[email protected]>
AuthorDate: Fri Jul 31 13:23:48 2026 -0700
Add CloudSQLNoOperationInProgressSensor for parallel admin ops (#68151)
---
.../google/docs/operators/cloud/cloud_sql.rst | 39 ++++++
providers/google/provider.yaml | 3 +
.../providers/google/cloud/hooks/cloud_sql.py | 32 +++++
.../providers/google/cloud/sensors/cloud_sql.py | 147 +++++++++++++++++++++
.../providers/google/cloud/triggers/cloud_sql.py | 81 +++++++++++-
.../airflow/providers/google/get_provider_info.py | 4 +
.../google/cloud/cloud_sql/example_cloud_sql.py | 11 ++
.../unit/google/cloud/hooks/test_cloud_sql.py | 20 +++
.../unit/google/cloud/sensors/test_cloud_sql.py | 127 ++++++++++++++++++
.../unit/google/cloud/triggers/test_cloud_sql.py | 82 +++++++++++-
10 files changed, 544 insertions(+), 2 deletions(-)
diff --git a/providers/google/docs/operators/cloud/cloud_sql.rst
b/providers/google/docs/operators/cloud/cloud_sql.rst
index dd1ea693048..544ab207157 100644
--- a/providers/google/docs/operators/cloud/cloud_sql.rst
+++ b/providers/google/docs/operators/cloud/cloud_sql.rst
@@ -359,6 +359,45 @@ as shown in the example:
:start-after: [START howto_operator_cloudsql_import_gcs_permissions]
:end-before: [END howto_operator_cloudsql_import_gcs_permissions]
+.. _howto/operator:CloudSQLNoOperationInProgressSensor:
+
+CloudSQLNoOperationInProgressSensor
+-----------------------------------
+
+Cloud SQL serializes administrative operations per instance: only one import,
export, backup or
+similar operation can run against an instance at a time. Submitting another
while one is in flight
+fails immediately with HTTP 409 ``operationInProgress``. This commonly affects
DAGs that fan out
+multiple
:class:`~airflow.providers.google.cloud.operators.cloud_sql.CloudSQLImportInstanceOperator`
+/
:class:`~airflow.providers.google.cloud.operators.cloud_sql.CloudSQLExportInstanceOperator`
tasks
+against the same instance in parallel.
+
+Use
+:class:`~airflow.providers.google.cloud.sensors.cloud_sql.CloudSQLNoOperationInProgressSensor`
+to wait until the instance has no administrative operation in progress before
submitting the next
+one. The sensor polls ``sqladmin.operations.list`` for the instance and
succeeds once no operation
+is in a non-terminal (``PENDING`` / ``RUNNING``) state. It supports deferrable
mode and fails fast
+on HTTP 403/404 (the instance is missing or access is denied).
+
+.. code-block:: python
+
+ from airflow.providers.google.cloud.sensors.cloud_sql import (
+ CloudSQLNoOperationInProgressSensor,
+ )
+
+ wait_for_slot = CloudSQLNoOperationInProgressSensor(
+ task_id="wait_for_slot",
+ instance="my-cloudsql-pg",
+ poke_interval=60,
+ timeout=2 * 60 * 60,
+ deferrable=True,
+ )
+
+ wait_for_slot >> import_data
+
+The sensor is best-effort: it reduces the chance of a 409 but cannot guarantee
exclusivity, since a
+new operation (for example one triggered from the console or by an automated
backup) could start
+between the sensor passing and the operator submitting.
+
.. _howto/operator:CloudSQLCreateInstanceOperator:
CloudSQLCreateInstanceOperator
diff --git a/providers/google/provider.yaml b/providers/google/provider.yaml
index 4a1fe1d3589..e1e6d841c6e 100644
--- a/providers/google/provider.yaml
+++ b/providers/google/provider.yaml
@@ -653,6 +653,9 @@ sensors:
- integration-name: Google Bigtable
python-modules:
- airflow.providers.google.cloud.sensors.bigtable
+ - integration-name: Google Cloud SQL
+ python-modules:
+ - airflow.providers.google.cloud.sensors.cloud_sql
- integration-name: Managed Service for Apache Airflow
python-modules:
- airflow.providers.google.cloud.sensors.cloud_composer
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
b/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
index d4d5f18134f..57889e96c2b 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
@@ -92,6 +92,15 @@ class CloudSqlOperationStatus:
UNKNOWN = "UNKNOWN"
+# Statuses that mean an administrative operation is still in flight on the
instance. Cloud SQL
+# serializes admin operations per instance, so a new import/export submitted
while one of these is
+# active fails with HTTP 409 ``operationInProgress``. Keying off an explicit
set (rather than
+# ``status != DONE``) avoids treating UNKNOWN/unexpected statuses as
in-progress and poking forever.
+CLOUD_SQL_NON_TERMINAL_STATUSES = frozenset(
+ {CloudSqlOperationStatus.PENDING, CloudSqlOperationStatus.RUNNING}
+)
+
+
class CloudSQLHook(GoogleBaseHook):
"""
Hook for Google Cloud SQL APIs.
@@ -429,6 +438,29 @@ class CloudSQLHook(GoogleBaseHook):
.execute(num_retries=self.num_retries)
)
+ @GoogleBaseHook.fallback_to_default_project_id
+ def list_operations(self, instance: str, project_id: str, max_results: int
| None = None) -> list[dict]:
+ """
+ List administrative operations for a Cloud SQL instance.
+
+ Must be called with keyword arguments because ``project_id`` is
injected by the
+ ``fallback_to_default_project_id`` decorator.
+
+ :param instance: Name of the Cloud SQL instance whose operations are
listed.
+ :param project_id: Project ID of the project that contains the
instance.
+ :param max_results: Optional maximum number of operations to return
per page.
+ :return: The list of operation resources for the instance (may be
empty).
+ """
+ response = (
+ self.get_conn()
+ .operations()
+ .list(project=project_id, instance=instance,
maxResults=max_results)
+ .execute(num_retries=self.num_retries)
+ )
+ # ``operations.list`` already filters server-side by ``instance``;
keep a defensive
+ # client-side filter on ``targetId`` in case the API ever returns
broader results.
+ return [op for op in response.get("items", []) if op.get("targetId")
== instance]
+
@GoogleBaseHook.fallback_to_default_project_id
def _wait_for_operation_to_complete(
self, project_id: str, operation_name: str, time_to_sleep: int =
TIME_TO_SLEEP_IN_SECONDS
diff --git
a/providers/google/src/airflow/providers/google/cloud/sensors/cloud_sql.py
b/providers/google/src/airflow/providers/google/cloud/sensors/cloud_sql.py
new file mode 100644
index 00000000000..50f0576e10a
--- /dev/null
+++ b/providers/google/src/airflow/providers/google/cloud/sensors/cloud_sql.py
@@ -0,0 +1,147 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""This module contains Google Cloud SQL sensors."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from googleapiclient.errors import HttpError
+
+from airflow.providers.common.compat.sdk import AirflowException,
BaseSensorOperator, conf
+from airflow.providers.google.cloud.hooks.cloud_sql import (
+ CLOUD_SQL_NON_TERMINAL_STATUSES,
+ CloudSQLHook,
+)
+from airflow.providers.google.cloud.triggers.cloud_sql import
CloudSQLNoOperationInProgressTrigger
+from airflow.providers.google.common.hooks.base_google import
PROVIDE_PROJECT_ID
+
+if TYPE_CHECKING:
+ from airflow.providers.common.compat.sdk import Context
+
+
+class CloudSQLOperationError(AirflowException):
+ """Raised when the Cloud SQL operations check fails (for example, the
instance is missing or access is denied)."""
+
+
+class CloudSQLNoOperationInProgressSensor(BaseSensorOperator):
+ """
+ Wait until a Cloud SQL instance has no administrative operation in
progress.
+
+ Cloud SQL serializes administrative operations per instance: only one
import, export, backup or
+ similar operation can run at a time. Submitting another while one is in
flight fails with HTTP
+ 409 ``operationInProgress``. Place this sensor upstream of
+
:class:`~airflow.providers.google.cloud.operators.cloud_sql.CloudSQLImportInstanceOperator`
/
+
:class:`~airflow.providers.google.cloud.operators.cloud_sql.CloudSQLExportInstanceOperator`
+ (or between mutually exclusive admin operators) to serialize work against
the same instance.
+
+ The sensor is operation-agnostic: it polls ``sqladmin.operations.list``
for the instance and
+ succeeds once no operation is in a non-terminal (PENDING/RUNNING) state.
It is best-effort -- a
+ new operation could still start between the sensor passing and the next
operator submitting.
+
+ .. seealso::
+ For more information on how to use this sensor, take a look at the
guide:
+ :ref:`howto/operator:CloudSQLNoOperationInProgressSensor`
+
+ :param instance: Name of the Cloud SQL instance to watch.
+ :param project_id: Optional, Google Cloud Project ID. If not provided the
default project is used.
+ :param gcp_conn_id: The connection ID used to connect to Google Cloud.
+ :param api_version: API version used (e.g. v1beta4).
+ :param impersonation_chain: Optional service account to impersonate using
short-term
+ credentials, or chained list of accounts required to get the
access_token of the last
+ account in the list, which will be impersonated in the request.
+ :param deferrable: Run the sensor in deferrable mode.
+ """
+
+ template_fields: Sequence[str] = (
+ "project_id",
+ "instance",
+ "impersonation_chain",
+ )
+ ui_color = "#D4ECEA"
+
+ def __init__(
+ self,
+ *,
+ instance: str,
+ project_id: str = PROVIDE_PROJECT_ID,
+ gcp_conn_id: str = "google_cloud_default",
+ api_version: str = "v1beta4",
+ impersonation_chain: str | Sequence[str] | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.instance = instance
+ self.project_id = project_id
+ self.gcp_conn_id = gcp_conn_id
+ self.api_version = api_version
+ self.impersonation_chain = impersonation_chain
+ self.deferrable = deferrable
+
+ def _get_hook(self) -> CloudSQLHook:
+ return CloudSQLHook(
+ api_version=self.api_version,
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+ def poke(self, context: Context) -> bool:
+ hook = self._get_hook()
+ try:
+ operations = hook.list_operations(project_id=self.project_id,
instance=self.instance)
+ except HttpError as e:
+ if e.resp.status in (403, 404):
+ # Instance missing or access denied - surface the
misconfiguration instead of poking.
+ raise CloudSQLOperationError(
+ f"Cloud SQL operations.list failed for instance
{self.instance}: {e}"
+ )
+ raise
+ in_progress = [op for op in operations if op.get("status") in
CLOUD_SQL_NON_TERMINAL_STATUSES]
+ if in_progress:
+ self.log.info(
+ "%s operation(s) still in progress on instance %s.",
len(in_progress), self.instance
+ )
+ return False
+ return True
+
+ def execute(self, context: Context) -> None:
+ """Run on the worker and defer using the trigger when in deferrable
mode."""
+ if self.deferrable:
+ if not self.poke(context=context):
+ self.defer(
+ timeout=timedelta(seconds=self.timeout),
+ trigger=CloudSQLNoOperationInProgressTrigger(
+ instance=self.instance,
+ project_id=self.project_id,
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ poke_interval=int(self.poke_interval),
+ api_version=self.api_version,
+ ),
+ method_name="execute_complete",
+ )
+ else:
+ super().execute(context)
+
+ def execute_complete(self, context: Context, event: dict | None = None) ->
None:
+ """Act as a callback for when the trigger fires."""
+ if event and event.get("status") in ("failed", "error"):
+ raise CloudSQLOperationError(event["message"])
+ self.log.info("No administrative operation in progress on instance
%s.", self.instance)
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/cloud_sql.py
b/providers/google/src/airflow/providers/google/cloud/triggers/cloud_sql.py
index dd8a482d92d..5b23ea62732 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/cloud_sql.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/cloud_sql.py
@@ -23,8 +23,13 @@ import asyncio
from collections.abc import Sequence
from asgiref.sync import sync_to_async
+from googleapiclient.errors import HttpError
-from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLAsyncHook,
CloudSqlOperationStatus
+from airflow.providers.google.cloud.hooks.cloud_sql import (
+ CLOUD_SQL_NON_TERMINAL_STATUSES,
+ CloudSQLAsyncHook,
+ CloudSqlOperationStatus,
+)
from airflow.providers.google.common.hooks.base_google import
PROVIDE_PROJECT_ID
from airflow.triggers.base import BaseTrigger, TriggerEvent
@@ -117,3 +122,77 @@ class CloudSQLExportTrigger(BaseTrigger):
"message": str(e),
}
)
+
+
+class CloudSQLNoOperationInProgressTrigger(BaseTrigger):
+ """
+ Trigger that waits until a Cloud SQL instance has no administrative
operation in progress.
+
+ Polls ``sqladmin.operations.list`` for the target instance and fires once
no operation is in a
+ non-terminal state (PENDING/RUNNING). Fails fast on 403/404 (the instance
is missing or access
+ is denied) rather than polling until timeout.
+ """
+
+ def __init__(
+ self,
+ instance: str,
+ project_id: str = PROVIDE_PROJECT_ID,
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ poke_interval: int = 20,
+ api_version: str = "v1beta4",
+ ):
+ super().__init__()
+ self.instance = instance
+ self.project_id = project_id
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+ self.poke_interval = poke_interval
+ self.api_version = api_version
+ self.hook = CloudSQLAsyncHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+ def serialize(self):
+ return (
+
"airflow.providers.google.cloud.triggers.cloud_sql.CloudSQLNoOperationInProgressTrigger",
+ {
+ "instance": self.instance,
+ "project_id": self.project_id,
+ "gcp_conn_id": self.gcp_conn_id,
+ "impersonation_chain": self.impersonation_chain,
+ "poke_interval": self.poke_interval,
+ "api_version": self.api_version,
+ },
+ )
+
+ async def run(self):
+ try:
+ sync_hook = await
self.hook.get_sync_hook(api_version=self.api_version)
+ while True:
+ # No async ``operations.list`` exists on the hook, so run the
sync call in a thread.
+ operations = await sync_to_async(sync_hook.list_operations)(
+ project_id=self.project_id, instance=self.instance
+ )
+ in_progress = [op for op in operations if op.get("status") in
CLOUD_SQL_NON_TERMINAL_STATUSES]
+ if not in_progress:
+ yield TriggerEvent({"instance": self.instance, "status":
"success"})
+ return
+ self.log.info(
+ "%s operation(s) still in progress on instance %s,
sleeping for %s seconds.",
+ len(in_progress),
+ self.instance,
+ self.poke_interval,
+ )
+ await asyncio.sleep(self.poke_interval)
+ except HttpError as e:
+ if e.resp.status in (403, 404):
+ # Instance missing or access denied - no point retrying.
+ yield TriggerEvent({"status": "failed", "message": str(e)})
+ return
+ self.log.exception("Error listing operations for instance %s.",
self.instance)
+ yield TriggerEvent({"status": "failed", "message": str(e)})
+ except Exception as e:
+ self.log.exception("Error listing operations for instance %s.",
self.instance)
+ yield TriggerEvent({"status": "failed", "message": str(e)})
diff --git a/providers/google/src/airflow/providers/google/get_provider_info.py
b/providers/google/src/airflow/providers/google/get_provider_info.py
index 8806a6a31e0..8a1024cd714 100644
--- a/providers/google/src/airflow/providers/google/get_provider_info.py
+++ b/providers/google/src/airflow/providers/google/get_provider_info.py
@@ -721,6 +721,10 @@ def get_provider_info():
"integration-name": "Google Bigtable",
"python-modules":
["airflow.providers.google.cloud.sensors.bigtable"],
},
+ {
+ "integration-name": "Google Cloud SQL",
+ "python-modules":
["airflow.providers.google.cloud.sensors.cloud_sql"],
+ },
{
"integration-name": "Managed Service for Apache Airflow",
"python-modules":
["airflow.providers.google.cloud.sensors.cloud_composer"],
diff --git
a/providers/google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py
b/providers/google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py
index 0abc5d9f806..eda31bd43c9 100644
--- a/providers/google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py
+++ b/providers/google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py
@@ -46,6 +46,7 @@ from airflow.providers.google.cloud.operators.gcs import (
GCSDeleteBucketOperator,
GCSObjectCreateAclEntryOperator,
)
+from airflow.providers.google.cloud.sensors.cloud_sql import
CloudSQLNoOperationInProgressSensor
try:
from airflow.sdk import TriggerRule
@@ -228,6 +229,15 @@ with DAG(
)
# [END howto_operator_cloudsql_import_gcs_permissions]
+ # Cloud SQL serializes admin operations per instance, so wait until the
export above has
+ # finished (no operation in progress) before submitting the import to
avoid a 409.
+ # [START howto_sensor_cloudsql_no_operation_in_progress]
+ sql_wait_no_operation_task = CloudSQLNoOperationInProgressSensor(
+ instance=INSTANCE_NAME,
+ task_id="sql_wait_no_operation_task",
+ )
+ # [END howto_sensor_cloudsql_no_operation_in_progress]
+
# [START howto_operator_cloudsql_import]
sql_import_task = CloudSQLImportInstanceOperator(
body=import_body, instance=INSTANCE_NAME, task_id="sql_import_task"
@@ -288,6 +298,7 @@ with DAG(
>> sql_export_task
>> sql_export_def_task
>> sql_gcp_add_object_permission_task
+ >> sql_wait_no_operation_task
>> sql_import_task
>> sql_instance_clone
>> sql_db_delete_task
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py
b/providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py
index 54480d17118..518f509e522 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py
@@ -543,6 +543,26 @@ class TestGcpSqlHookDefaultProjectId:
get_method.assert_called_once_with(project="gcp-project",
operation="operation_id")
execute_method.assert_called_once_with(num_retries=self.cloudsql_hook.num_retries)
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn")
+ def test_list_operations(self, mock_get_conn):
+ operations_method = mock_get_conn.return_value.operations
+ list_method = operations_method.return_value.list
+ execute_method = list_method.return_value.execute
+ execute_method.return_value = {
+ "items": [
+ {"name": "op1", "status": "DONE", "targetId": "my-instance"},
+ {"name": "op2", "status": "RUNNING", "targetId":
"other-instance"},
+ ]
+ }
+
+ result = self.cloudsql_hook.list_operations(project_id="gcp-project",
instance="my-instance")
+
+ # Only operations whose targetId matches the instance are returned.
+ assert result == [{"name": "op1", "status": "DONE", "targetId":
"my-instance"}]
+ operations_method.assert_called_once()
+ list_method.assert_called_once_with(project="gcp-project",
instance="my-instance", maxResults=None)
+
execute_method.assert_called_once_with(num_retries=self.cloudsql_hook.num_retries)
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.build")
def test_get_conn_obj_caching(self, mock_build):
self.cloudsql_hook._authorize = mock.MagicMock()
diff --git a/providers/google/tests/unit/google/cloud/sensors/test_cloud_sql.py
b/providers/google/tests/unit/google/cloud/sensors/test_cloud_sql.py
new file mode 100644
index 00000000000..4f28c9482d0
--- /dev/null
+++ b/providers/google/tests/unit/google/cloud/sensors/test_cloud_sql.py
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+
+import httplib2
+import pytest
+from googleapiclient.errors import HttpError
+
+from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
+from airflow.providers.google.cloud.sensors.cloud_sql import
CloudSQLNoOperationInProgressSensor
+from airflow.providers.google.cloud.triggers.cloud_sql import
CloudSQLNoOperationInProgressTrigger
+
+SENSOR_PATH = "airflow.providers.google.cloud.sensors.cloud_sql.{}"
+TASK_ID = "test_no_op_in_progress"
+INSTANCE = "test-instance"
+PROJECT_ID = "test-project"
+
+
+def _http_error(status: int) -> HttpError:
+ return HttpError(resp=httplib2.Response({"status": status}),
content=b"error content")
+
+
+class TestCloudSQLNoOperationInProgressSensor:
+ @mock.patch(SENSOR_PATH.format("CloudSQLHook"), autospec=True)
+ def test_poke_returns_true_when_no_in_progress_operations(self, mock_hook):
+ mock_hook.return_value.list_operations.return_value = [
+ {"name": "op1", "status": "DONE", "targetId": INSTANCE},
+ ]
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID
+ )
+ assert sensor.poke(context={}) is True
+ mock_hook.return_value.list_operations.assert_called_once_with(
+ project_id=PROJECT_ID, instance=INSTANCE
+ )
+
+ @mock.patch(SENSOR_PATH.format("CloudSQLHook"), autospec=True)
+ def test_poke_returns_false_when_operation_in_progress(self, mock_hook):
+ mock_hook.return_value.list_operations.return_value = [
+ {"name": "op1", "status": "RUNNING", "targetId": INSTANCE},
+ {"name": "op2", "status": "PENDING", "targetId": INSTANCE},
+ ]
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID
+ )
+ assert sensor.poke(context={}) is False
+
+ @pytest.mark.parametrize("status", [403, 404])
+ @mock.patch(SENSOR_PATH.format("CloudSQLHook"), autospec=True)
+ def test_poke_fails_fast_on_403_404(self, mock_hook, status):
+ mock_hook.return_value.list_operations.side_effect =
_http_error(status)
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID
+ )
+ with pytest.raises(AirflowException, match="operations.list failed"):
+ sensor.poke(context={})
+
+ @mock.patch(SENSOR_PATH.format("CloudSQLHook"), autospec=True)
+ def test_poke_reraises_other_http_errors(self, mock_hook):
+ mock_hook.return_value.list_operations.side_effect = _http_error(500)
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID
+ )
+ with pytest.raises(HttpError):
+ sensor.poke(context={})
+
+ @mock.patch(SENSOR_PATH.format("CloudSQLHook"), autospec=True)
+ def test_execute_defers_when_deferrable_and_not_idle(self, mock_hook):
+ mock_hook.return_value.list_operations.return_value = [
+ {"name": "op1", "status": "RUNNING", "targetId": INSTANCE},
+ ]
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID,
deferrable=True
+ )
+ with pytest.raises(TaskDeferred) as exc:
+ sensor.execute(context={})
+ assert isinstance(exc.value.trigger,
CloudSQLNoOperationInProgressTrigger)
+ assert exc.value.method_name == "execute_complete"
+
+ @mock.patch(SENSOR_PATH.format("CloudSQLHook"), autospec=True)
+ def test_execute_does_not_defer_when_idle(self, mock_hook):
+ mock_hook.return_value.list_operations.return_value = [
+ {"name": "op1", "status": "DONE", "targetId": INSTANCE},
+ ]
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID,
deferrable=True
+ )
+ assert sensor.execute(context={}) is None
+
+ @mock.patch(SENSOR_PATH.format("BaseSensorOperator.execute"),
autospec=True)
+ @mock.patch(SENSOR_PATH.format("CloudSQLHook"), autospec=True)
+ def test_execute_non_deferrable_delegates_to_super(self, mock_hook,
mock_super_execute):
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID,
deferrable=False
+ )
+ sensor.execute(context={})
+ mock_super_execute.assert_called_once()
+
+ def test_execute_complete_success(self):
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID
+ )
+ assert sensor.execute_complete(context={}, event={"instance":
INSTANCE, "status": "success"}) is None
+
+ @pytest.mark.parametrize("status", ["failed", "error"])
+ def test_execute_complete_failure_raises(self, status):
+ sensor = CloudSQLNoOperationInProgressSensor(
+ task_id=TASK_ID, instance=INSTANCE, project_id=PROJECT_ID
+ )
+ with pytest.raises(AirflowException, match="boom"):
+ sensor.execute_complete(context={}, event={"status": status,
"message": "boom"})
diff --git
a/providers/google/tests/unit/google/cloud/triggers/test_cloud_sql.py
b/providers/google/tests/unit/google/cloud/triggers/test_cloud_sql.py
index f5de135c47d..fe8e501ef27 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_cloud_sql.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_cloud_sql.py
@@ -20,12 +20,20 @@ import asyncio
import logging
from unittest import mock
+import httplib2
import pytest
+from googleapiclient.errors import HttpError
from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLHook
-from airflow.providers.google.cloud.triggers.cloud_sql import
CloudSQLExportTrigger
+from airflow.providers.google.cloud.triggers.cloud_sql import (
+ CloudSQLExportTrigger,
+ CloudSQLNoOperationInProgressTrigger,
+)
from airflow.triggers.base import TriggerEvent
+INSTANCE = "test-instance"
+NO_OP_CLASSPATH =
"airflow.providers.google.cloud.triggers.cloud_sql.CloudSQLNoOperationInProgressTrigger"
+
CLASSPATH =
"airflow.providers.google.cloud.triggers.cloud_sql.CloudSQLExportTrigger"
TASK_ID = "test_task"
TEST_POLL_INTERVAL = 10
@@ -197,3 +205,75 @@ class TestCloudSQLExportTrigger:
# Verify the default universe branch not being called
mock_async_get_op.assert_not_called()
task.cancel()
+
+
[email protected]
+def no_op_trigger():
+ return CloudSQLNoOperationInProgressTrigger(
+ instance=INSTANCE,
+ project_id=PROJECT_ID,
+ impersonation_chain=None,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ poke_interval=TEST_POLL_INTERVAL,
+ api_version=API_VERSION,
+ )
+
+
+class TestCloudSQLNoOperationInProgressTrigger:
+ def test_serialization(self, no_op_trigger, sync_hook_mock):
+ classpath, kwargs = no_op_trigger.serialize()
+ assert classpath == NO_OP_CLASSPATH
+ assert kwargs == {
+ "instance": INSTANCE,
+ "project_id": PROJECT_ID,
+ "impersonation_chain": None,
+ "gcp_conn_id": TEST_GCP_CONN_ID,
+ "poke_interval": TEST_POLL_INTERVAL,
+ "api_version": API_VERSION,
+ }
+
+ @pytest.mark.asyncio
+ async def test_run_success_when_no_operation_in_progress(self,
no_op_trigger, sync_hook_mock):
+ sync_hook_mock.list_operations.return_value = [
+ {"name": "op1", "status": "DONE", "targetId": INSTANCE},
+ ]
+ actual = await no_op_trigger.run().asend(None)
+ assert actual == TriggerEvent({"instance": INSTANCE, "status":
"success"})
+
+ @pytest.mark.asyncio
+ async def test_run_success_when_only_terminal_operations(self,
no_op_trigger, sync_hook_mock):
+ # Operations in terminal states (DONE) do not block; the trigger keys
off the non-terminal set.
+ sync_hook_mock.list_operations.return_value = [
+ {"name": "op-done", "status": "DONE", "targetId": INSTANCE},
+ {"name": "op-unknown", "status": "UNKNOWN", "targetId": INSTANCE},
+ ]
+ actual = await no_op_trigger.run().asend(None)
+ assert actual == TriggerEvent({"instance": INSTANCE, "status":
"success"})
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.cloud_sql.asyncio.sleep",
new_callable=mock.AsyncMock
+ )
+ async def test_run_sleeps_while_operation_in_progress(self, mock_sleep,
no_op_trigger, sync_hook_mock):
+ sync_hook_mock.list_operations.side_effect = [
+ [{"name": "op1", "status": "RUNNING", "targetId": INSTANCE}],
+ [{"name": "op2", "status": "DONE", "targetId": INSTANCE}],
+ ]
+ actual = await no_op_trigger.run().asend(None)
+ assert actual == TriggerEvent({"instance": INSTANCE, "status":
"success"})
+ mock_sleep.assert_awaited_once_with(TEST_POLL_INTERVAL)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("status", [403, 404])
+ async def test_run_fails_fast_on_403_404(self, status, no_op_trigger,
sync_hook_mock):
+ sync_hook_mock.list_operations.side_effect = HttpError(
+ resp=httplib2.Response({"status": status}), content=b"denied or
missing"
+ )
+ actual = await no_op_trigger.run().asend(None)
+ assert actual.payload["status"] == "failed"
+
+ @pytest.mark.asyncio
+ async def test_run_fails_on_generic_exception(self, no_op_trigger,
sync_hook_mock):
+ sync_hook_mock.list_operations.side_effect = Exception("boom")
+ actual = await no_op_trigger.run().asend(None)
+ assert actual == TriggerEvent({"status": "failed", "message": "boom"})
