Fix Cloud SQL 409 operationInProgress on import/export operations

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

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
@@ -354,11 +364,54 @@ def export_instance(self, instance: str, body: dict, project_id: str):
         operation_name = response["name"]
         return operation_name
 
+    @GoogleBaseHook.operation_in_progress_retry()
+    def _submit_import(self, instance: str, body: dict, project_id: str) -> str:
+        """
+        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. 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 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.
+        :return: The name of the accepted import operation.
+        """
+        try:
+            response = (
+                self.get_conn()
+                .instances()
+                .import_(project=project_id, instance=instance, body=body)
+                .execute(num_retries=self.num_retries)
+            )
+            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
+            # 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
     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
@@ -367,14 +420,8 @@ def import_instance(self, instance: str, body: dict, project_id: str) -> None:
             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:
-            response = (
-                self.get_conn()
-                .instances()
-                .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)
         except HttpError as ex:
             raise AirflowException(f"Importing instance {instance} failed: {ex.content}")

providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py

@@ -165,10 +165,65 @@ 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``. ``_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.
+        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_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"),