From f706c4128d0de5df94430e9d7b5e15de9b87b218 Mon Sep 17 00:00:00 2001 From: radhwene Date: Sat, 13 Jun 2026 12:24:16 +0200 Subject: [PATCH 1/2] Fix Cloud SQL 409 operationInProgress on import/export operations Apply the existing operation_in_progress_retry() policy to CloudSQLHook.import_instance and export_instance, the only two admin methods that lacked it. Also re-raise operationInProgress HttpError un-wrapped from import_instance so the retry decorator can see it; terminal HttpErrors still get the friendly AirflowException message. --- .../providers/google/cloud/hooks/cloud_sql.py | 25 +++++++++++++ .../unit/google/cloud/hooks/test_cloud_sql.py | 35 +++++++++++++++++-- 2 files changed, 58 insertions(+), 2 deletions(-) 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 57889e96c2b42..8b7239e269465 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 @@ -66,6 +66,7 @@ GoogleBaseAsyncHook, GoogleBaseHook, get_field, + is_operation_in_progress_exception, ) from airflow.utils.log.logging_mixin import LoggingMixin @@ -333,10 +334,19 @@ def delete_database(self, instance: str, database: str, project_id: str) -> None self._wait_for_operation_to_complete(project_id=project_id, operation_name=operation_name) @GoogleBaseHook.fallback_to_default_project_id + @GoogleBaseHook.operation_in_progress_retry() def export_instance(self, instance: str, body: dict, project_id: str): """ Export data from a Cloud SQL instance to a Cloud Storage bucket as a SQL dump or CSV file. + Cloud SQL serializes administrative operations per instance, so submitting an export while + another admin operation is already running on the same instance fails with HTTP 409 + ``operationInProgress``. The ``operation_in_progress_retry`` decorator retries the submit on + HTTP 409/429 with exponential backoff (capped at 300s between attempts) until the API accepts + it. This makes ``export_instance`` consistent with the other admin methods on this hook that + already carry the decorator (``create_instance``, ``patch_instance``, ``delete_instance``, + ``create_database``, ``patch_database``, ``delete_database``). + :param instance: Database instance ID of the Cloud SQL instance. This does not include the project ID. :param body: The request body, as described in @@ -355,10 +365,19 @@ def export_instance(self, instance: str, body: dict, project_id: str): return operation_name @GoogleBaseHook.fallback_to_default_project_id + @GoogleBaseHook.operation_in_progress_retry() def import_instance(self, instance: str, body: dict, project_id: str) -> None: """ Import data into a Cloud SQL instance from a SQL dump or CSV file in Cloud Storage. + Cloud SQL serializes administrative operations per instance, so submitting an import while + another admin operation is already running on the same instance fails with HTTP 409 + ``operationInProgress``. The ``operation_in_progress_retry`` decorator retries the submit on + HTTP 409/429 with exponential backoff (capped at 300s between attempts) until the API accepts + it. This makes ``import_instance`` consistent with the other admin methods on this hook that + already carry the decorator (``create_instance``, ``patch_instance``, ``delete_instance``, + ``create_database``, ``patch_database``, ``delete_database``). + :param instance: Database instance ID. This does not include the project ID. :param body: The request body, as described in @@ -377,6 +396,12 @@ def import_instance(self, instance: str, body: dict, project_id: str) -> None: operation_name = response["name"] self._wait_for_operation_to_complete(project_id=project_id, operation_name=operation_name) except HttpError as ex: + # ``operation_in_progress_retry`` retries on the raw ``HttpError`` (status 409/429), not on + # ``AirflowException``. Re-raise operation-in-progress errors unchanged so the decorator can + # see and retry them; otherwise the 409 would be wrapped below and the retry never triggers. + # Genuinely terminal HttpErrors still get the friendly message. + if is_operation_in_progress_exception(ex): + raise raise AirflowException(f"Importing instance {instance} failed: {ex.content}") @GoogleBaseHook.fallback_to_default_project_id 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 518f509e522b5..6b8f1bdb87efa 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 @@ -165,10 +165,41 @@ def test_instance_export_with_in_progress_retry(self, wait_for_operation_to_comp ), {"name": "operation_id"}, ] - with pytest.raises(HttpError): - self.cloudsql_hook.export_instance(project_id="example-project", instance="instance", body={}) + # First submit returns 429 (one of the two operation-in-progress codes recognised by + # ``is_operation_in_progress_exception``); ``operation_in_progress_retry`` retries and the + # second submit succeeds, returning the operation name. The import test below covers 409. + result = self.cloudsql_hook.export_instance( + project_id="example-project", instance="instance", body={} + ) + assert result == "operation_id" + assert export_method.call_count == 2 + assert execute_method.call_count == 2 wait_for_operation_to_complete.assert_not_called() + @mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn") + @mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete") + def test_instance_import_with_in_progress_retry(self, wait_for_operation_to_complete, get_conn): + import_method = get_conn.return_value.instances.return_value.import_ + execute_method = import_method.return_value.execute + execute_method.side_effect = [ + HttpError( + resp=httplib2.Response({"status": 409}), + content=b"operationInProgress", + ), + {"name": "operation_id"}, + ] + wait_for_operation_to_complete.return_value = None + # First submit returns 409 ``operationInProgress``. ``import_instance`` re-raises it past its + # friendly-message wrapper (instead of converting it to AirflowException), so + # ``operation_in_progress_retry`` sees the raw HttpError, retries, and the second submit + # succeeds; the resulting operation is awaited exactly once. + self.cloudsql_hook.import_instance(project_id="example-project", instance="instance", body={}) + assert import_method.call_count == 2 + assert execute_method.call_count == 2 + wait_for_operation_to_complete.assert_called_once_with( + project_id="example-project", operation_name="operation_id" + ) + @mock.patch( "airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_credentials_and_project_id", return_value=(mock.MagicMock(), "example-project"), From fa839504bada1e691501ce37eb33786a343e45b3 Mon Sep 17 00:00:00 2001 From: radhwene Date: Sat, 13 Jun 2026 12:25:27 +0200 Subject: [PATCH 2/2] Limit Cloud SQL 409 retry to the import submit call Address review feedback on #68361: operation_in_progress_retry() previously wrapped the whole import_instance, including the operation-status polling. A retryable 429 raised during polling re-ran the method and re-submitted an import that was already accepted, importing the same data twice. The submit now lives in _submit_import, which alone carries the retry decorator; import_instance waits outside the retry scope, so a polling failure fails the task instead of re-submitting. Adds a regression test asserting exactly one submit when polling raises a retryable error. --- generated/known_airflow_exceptions.txt | 2 +- .../providers/google/cloud/hooks/cloud_sql.py | 48 ++++++++++++++----- .../unit/google/cloud/hooks/test_cloud_sql.py | 26 +++++++++- 3 files changed, 61 insertions(+), 15 deletions(-) diff --git a/generated/known_airflow_exceptions.txt b/generated/known_airflow_exceptions.txt index e170208c3678d..16c59ee6f8180 100644 --- a/generated/known_airflow_exceptions.txt +++ b/generated/known_airflow_exceptions.txt @@ -214,7 +214,7 @@ providers/google/src/airflow/providers/google/cloud/hooks/cloud_build.py::3 providers/google/src/airflow/providers/google/cloud/hooks/cloud_composer.py::5 providers/google/src/airflow/providers/google/cloud/hooks/cloud_memorystore.py::5 providers/google/src/airflow/providers/google/cloud/hooks/cloud_run.py::1 -providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py::32 +providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py::33 providers/google/src/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py::5 providers/google/src/airflow/providers/google/cloud/hooks/compute.py::6 providers/google/src/airflow/providers/google/cloud/hooks/compute_ssh.py::6 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 8b7239e269465..7b9ac0415c508 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 @@ -364,27 +364,23 @@ def export_instance(self, instance: str, body: dict, project_id: str): operation_name = response["name"] return operation_name - @GoogleBaseHook.fallback_to_default_project_id @GoogleBaseHook.operation_in_progress_retry() - def import_instance(self, instance: str, body: dict, project_id: str) -> None: + def _submit_import(self, instance: str, body: dict, project_id: str) -> str: """ - Import data into a Cloud SQL instance from a SQL dump or CSV file in Cloud Storage. + Submit an import request for a Cloud SQL instance, retrying while an operation is in progress. Cloud SQL serializes administrative operations per instance, so submitting an import while another admin operation is already running on the same instance fails with HTTP 409 ``operationInProgress``. The ``operation_in_progress_retry`` decorator retries the submit on HTTP 409/429 with exponential backoff (capped at 300s between attempts) until the API accepts - it. This makes ``import_instance`` consistent with the other admin methods on this hook that - already carry the decorator (``create_instance``, ``patch_instance``, ``delete_instance``, - ``create_database``, ``patch_database``, ``delete_database``). + it. The decorator deliberately wraps only this submit call: a 409/429 here guarantees the + import was rejected, so repeating it is safe. - :param instance: Database instance ID. This does not include the - project ID. + :param instance: Database instance ID. This does not include the project ID. :param body: The request body, as described in https://cloud.google.com/sql/docs/mysql/admin-api/v1beta4/instances/import#request-body - :param project_id: Project ID of the project that contains the instance. If set - to None or missing, the default project_id from the Google Cloud connection is used. - :return: None + :param project_id: Project ID of the project that contains the instance. + :return: The name of the accepted import operation. """ try: response = ( @@ -393,8 +389,7 @@ def import_instance(self, instance: str, body: dict, project_id: str) -> None: .import_(project=project_id, instance=instance, body=body) .execute(num_retries=self.num_retries) ) - operation_name = response["name"] - self._wait_for_operation_to_complete(project_id=project_id, operation_name=operation_name) + return response["name"] except HttpError as ex: # ``operation_in_progress_retry`` retries on the raw ``HttpError`` (status 409/429), not on # ``AirflowException``. Re-raise operation-in-progress errors unchanged so the decorator can @@ -404,6 +399,33 @@ def import_instance(self, instance: str, body: dict, project_id: str) -> None: raise raise AirflowException(f"Importing instance {instance} failed: {ex.content}") + @GoogleBaseHook.fallback_to_default_project_id + def import_instance(self, instance: str, body: dict, project_id: str) -> None: + """ + Import data into a Cloud SQL instance from a SQL dump or CSV file in Cloud Storage. + + The submit call is retried on HTTP 409 ``operationInProgress`` / 429 (see ``_submit_import``), + consistent with the other admin methods on this hook that carry the decorator + (``create_instance``, ``patch_instance``, ``delete_instance``, ``create_database``, + ``patch_database``, ``delete_database``). The status polling that follows runs outside the + retry scope: once the import was accepted, a retryable polling error must fail the task + rather than restart the method, because re-submitting an accepted import would import the + same data twice. + + :param instance: Database instance ID. This does not include the + project ID. + :param body: The request body, as described in + https://cloud.google.com/sql/docs/mysql/admin-api/v1beta4/instances/import#request-body + :param project_id: Project ID of the project that contains the instance. If set + to None or missing, the default project_id from the Google Cloud connection is used. + :return: None + """ + operation_name = self._submit_import(instance=instance, body=body, project_id=project_id) + try: + self._wait_for_operation_to_complete(project_id=project_id, operation_name=operation_name) + except HttpError as ex: + raise AirflowException(f"Importing instance {instance} failed: {ex.content}") + @GoogleBaseHook.fallback_to_default_project_id def clone_instance(self, instance: str, body: dict, project_id: str) -> None: """ 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 6b8f1bdb87efa..78f8f5b919101 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 @@ -189,7 +189,7 @@ def test_instance_import_with_in_progress_retry(self, wait_for_operation_to_comp {"name": "operation_id"}, ] wait_for_operation_to_complete.return_value = None - # First submit returns 409 ``operationInProgress``. ``import_instance`` re-raises it past its + # First submit returns 409 ``operationInProgress``. ``_submit_import`` re-raises it past its # friendly-message wrapper (instead of converting it to AirflowException), so # ``operation_in_progress_retry`` sees the raw HttpError, retries, and the second submit # succeeds; the resulting operation is awaited exactly once. @@ -200,6 +200,30 @@ def test_instance_import_with_in_progress_retry(self, wait_for_operation_to_comp project_id="example-project", operation_name="operation_id" ) + @mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn") + @mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete") + def test_instance_import_does_not_resubmit_when_polling_fails( + self, wait_for_operation_to_complete, get_conn + ): + import_method = get_conn.return_value.instances.return_value.import_ + execute_method = import_method.return_value.execute + execute_method.return_value = {"name": "operation_id"} + wait_for_operation_to_complete.side_effect = HttpError( + resp=httplib2.Response({"status": 429}), + content=b"rate limited", + ) + # The submit succeeds, then the operation-status polling raises a retryable 429. The retry + # scope must not include the polling: re-running ``import_instance`` would re-submit an + # import that was already accepted and import the same data twice. The task must fail + # instead, with exactly one submit on record. + with pytest.raises(AirflowException, match="Importing instance instance failed"): + self.cloudsql_hook.import_instance(project_id="example-project", instance="instance", body={}) + import_method.assert_called_once_with(body={}, instance="instance", project="example-project") + execute_method.assert_called_once_with(num_retries=5) + wait_for_operation_to_complete.assert_called_once_with( + project_id="example-project", operation_name="operation_id" + ) + @mock.patch( "airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_credentials_and_project_id", return_value=(mock.MagicMock(), "example-project"),