[AIRFLOW-7069] Fix cloudsql system tests (#7770)

Author: potiuk

airflow/providers/google/cloud/example_dags/example_cloud_sql.py

@@ -202,7 +202,7 @@
     sql_instance_read_replica_create = CloudSQLCreateInstanceOperator(
         project_id=GCP_PROJECT_ID,
         body=read_replica_body,
-        instance=INSTANCE_NAME2,
+        instance=READ_REPLICA_NAME,
         task_id='sql_instance_read_replica_create'
     )
 
@@ -217,13 +217,14 @@
         instance=INSTANCE_NAME,
         task_id='sql_instance_patch_task'
     )
+    # [END howto_operator_cloudsql_patch]
 
     sql_instance_patch_task2 = CloudSQLInstancePatchOperator(
+        project_id=GCP_PROJECT_ID,
         body=patch_body,
         instance=INSTANCE_NAME,
         task_id='sql_instance_patch_task2'
     )
-    # [END howto_operator_cloudsql_patch]
 
     # [START howto_operator_cloudsql_db_create]
     sql_db_create_task = CloudSQLCreateInstanceDatabaseOperator(

airflow/providers/google/cloud/hooks/base.py

@@ -19,7 +19,6 @@
 """
 This module contains a Google Cloud API base hook.
 """
-
 import functools
 import json
 import logging
@@ -39,6 +38,7 @@
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.auth import _cloud_sdk
 from google.auth.environment_vars import CREDENTIALS
+from googleapiclient.errors import HttpError
 from googleapiclient.http import set_user_agent
 
 from airflow import version
@@ -92,13 +92,32 @@ def is_soft_quota_exception(exception: Exception):
     return False
 
 
+def is_operation_in_progress_exception(exception: Exception):
+    """
+    Some of the calls return 429 (too many requests!) or 409 errors (Conflict)
+    in case of operation in progress.
+
+    * Google Cloud SQL
+    """
+    if isinstance(exception, HttpError):
+        return exception.resp.status == 429 or exception.resp.status == 409
+    return False
+
+
 class retry_if_temporary_quota(tenacity.retry_if_exception):  # pylint: disable=invalid-name
     """Retries if there was an exception for exceeding the temporary quote limit."""
 
     def __init__(self):
         super().__init__(is_soft_quota_exception)
 
 
+class retry_if_operation_in_progress(tenacity.retry_if_exception):  # pylint: disable=invalid-name
+    """Retries if there was an exception for exceeding the temporary quote limit."""
+
+    def __init__(self):
+        super().__init__(is_operation_in_progress_exception)
+
+
 RT = TypeVar('RT')  # pylint: disable=invalid-name
 
 
@@ -295,7 +314,7 @@ def scopes(self) -> Sequence[str]:
     @staticmethod
     def quota_retry(*args, **kwargs) -> Callable:
         """
-        A decorator who provides a mechanism to repeat requests in response to exceeding a temporary quote
+        A decorator that provides a mechanism to repeat requests in response to exceeding a temporary quote
         limit.
         """
         def decorator(fun: Callable):
@@ -311,6 +330,26 @@ def decorator(fun: Callable):
             )(fun)
         return decorator
 
+    @staticmethod
+    def operation_in_progress_retry(*args, **kwargs) -> Callable:
+        """
+        A decorator that provides a mechanism to repeat requests in response to
+        operation in progress (HTTP 409)
+        limit.
+        """
+        def decorator(fun: Callable):
+            default_kwargs = {
+                'wait': tenacity.wait_exponential(multiplier=1, max=300),
+                'retry': retry_if_operation_in_progress(),
+                'before': tenacity.before_log(log, logging.DEBUG),
+                'after': tenacity.after_log(log, logging.DEBUG),
+            }
+            default_kwargs.update(**kwargs)
+            return tenacity.retry(
+                *args, **default_kwargs
+            )(fun)
+        return decorator
+
     @staticmethod
     def fallback_to_default_project_id(func: Callable[..., RT]) -> Callable[..., RT]:
         """

airflow/providers/google/cloud/hooks/cloud_sql.py

@@ -56,7 +56,7 @@
 UNIX_PATH_MAX = 108
 
 # Time to sleep between active checks of the operation results
-TIME_TO_SLEEP_IN_SECONDS = 1
+TIME_TO_SLEEP_IN_SECONDS = 20
 
 
 class CloudSqlOperationStatus:
@@ -113,14 +113,13 @@ def get_instance(self, instance: str, project_id: Optional[str] = None) -> Dict:
         :return: A Cloud SQL instance resource.
         :rtype: dict
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        return self.get_conn().instances().get(  # pylint: disable=no-member
+        return self.get_conn().instances().get(  # noqa # pylint: disable=no-member
             project=project_id,
             instance=instance
         ).execute(num_retries=self.num_retries)
 
     @CloudBaseHook.fallback_to_default_project_id
+    @CloudBaseHook.operation_in_progress_retry()
     def create_instance(self, body: Dict, project_id: Optional[str] = None) -> None:
         """
         Creates a new Cloud SQL instance.
@@ -133,17 +132,16 @@ def create_instance(self, body: Dict, project_id: Optional[str] = None) -> None:
         :type project_id: str
         :return: None
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        response = self.get_conn().instances().insert(  # pylint: disable=no-member
+        response = self.get_conn().instances().insert(  # noqa # pylint: disable=no-member
             project=project_id,
             body=body
         ).execute(num_retries=self.num_retries)
         operation_name = response["name"]
-        self._wait_for_operation_to_complete(project_id=project_id,
+        self._wait_for_operation_to_complete(project_id=project_id,  # type:ignore
                                              operation_name=operation_name)
 
     @CloudBaseHook.fallback_to_default_project_id
+    @CloudBaseHook.operation_in_progress_retry()
     def patch_instance(self, body: Dict, instance: str, project_id: Optional[str] = None) -> None:
         """
         Updates settings of a Cloud SQL instance.
@@ -161,18 +159,17 @@ def patch_instance(self, body: Dict, instance: str, project_id: Optional[str] =
         :type project_id: str
         :return: None
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        response = self.get_conn().instances().patch(  # pylint: disable=no-member
+        response = self.get_conn().instances().patch(  # noqa # pylint: disable=no-member
             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,
+        self._wait_for_operation_to_complete(project_id=project_id,  # type:ignore
                                              operation_name=operation_name)
 
     @CloudBaseHook.fallback_to_default_project_id
+    @CloudBaseHook.operation_in_progress_retry()
     def delete_instance(self, instance: str, project_id: Optional[str] = None) -> None:
         """
         Deletes a Cloud SQL instance.
@@ -184,14 +181,12 @@ def delete_instance(self, instance: str, project_id: Optional[str] = None) -> No
         :type instance: str
         :return: None
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        response = self.get_conn().instances().delete(  # pylint: disable=no-member
+        response = self.get_conn().instances().delete(  # noqa # pylint: disable=no-member
             project=project_id,
             instance=instance,
         ).execute(num_retries=self.num_retries)
         operation_name = response["name"]
-        self._wait_for_operation_to_complete(project_id=project_id,
+        self._wait_for_operation_to_complete(project_id=project_id,  # type:ignore
                                              operation_name=operation_name)
 
     @CloudBaseHook.fallback_to_default_project_id
@@ -210,15 +205,14 @@ def get_database(self, instance: str, database: str, project_id: Optional[str] =
             https://cloud.google.com/sql/docs/mysql/admin-api/v1beta4/databases#resource.
         :rtype: dict
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        return self.get_conn().databases().get(  # pylint: disable=no-member
+        return self.get_conn().databases().get(  # noqa # pylint: disable=no-member
             project=project_id,
             instance=instance,
             database=database
         ).execute(num_retries=self.num_retries)
 
     @CloudBaseHook.fallback_to_default_project_id
+    @CloudBaseHook.operation_in_progress_retry()
     def create_database(self, instance: str, body: Dict, project_id: Optional[str] = None) -> None:
         """
         Creates a new database inside a Cloud SQL instance.
@@ -233,18 +227,17 @@ def create_database(self, instance: str, body: Dict, project_id: Optional[str] =
         :type project_id: str
         :return: None
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        response = self.get_conn().databases().insert(  # pylint: disable=no-member
+        response = self.get_conn().databases().insert(  # noqa # pylint: disable=no-member
             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,
+        self._wait_for_operation_to_complete(project_id=project_id,  # type:ignore
                                              operation_name=operation_name)
 
     @CloudBaseHook.fallback_to_default_project_id
+    @CloudBaseHook.operation_in_progress_retry()
     def patch_database(
         self,
         instance: str,
@@ -270,19 +263,18 @@ def patch_database(
         :type project_id: str
         :return: None
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        response = self.get_conn().databases().patch(  # pylint: disable=no-member
+        response = self.get_conn().databases().patch(  # noqa # pylint: disable=no-member
             project=project_id,
             instance=instance,
             database=database,
             body=body
         ).execute(num_retries=self.num_retries)
         operation_name = response["name"]
-        self._wait_for_operation_to_complete(project_id=project_id,
+        self._wait_for_operation_to_complete(project_id=project_id,  # type:ignore
                                              operation_name=operation_name)
 
     @CloudBaseHook.fallback_to_default_project_id
+    @CloudBaseHook.operation_in_progress_retry()
     def delete_database(self, instance: str, database: str, project_id: Optional[str] = None) -> None:
         """
         Deletes a database from a Cloud SQL instance.
@@ -296,18 +288,17 @@ def delete_database(self, instance: str, database: str, project_id: Optional[str
         :type project_id: str
         :return: None
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        response = self.get_conn().databases().delete(  # pylint: disable=no-member
+        response = self.get_conn().databases().delete(  # noqa # pylint: disable=no-member
             project=project_id,
             instance=instance,
             database=database
         ).execute(num_retries=self.num_retries)
         operation_name = response["name"]
-        self._wait_for_operation_to_complete(project_id=project_id,
+        self._wait_for_operation_to_complete(project_id=project_id,  # type:ignore
                                              operation_name=operation_name)
 
     @CloudBaseHook.fallback_to_default_project_id
+    @CloudBaseHook.operation_in_progress_retry()
     def export_instance(self, instance: str, body: Dict, project_id: Optional[str] = None) -> None:
         """
         Exports data from a Cloud SQL instance to a Cloud Storage bucket as a SQL dump
@@ -324,21 +315,14 @@ def export_instance(self, instance: str, body: Dict, project_id: Optional[str] =
         :type project_id: str
         :return: None
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
-        try:
-            response = self.get_conn().instances().export(  # pylint: disable=no-member
-                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(
-                'Exporting instance {} failed: {}'.format(instance, ex.content)
-            )
+        response = self.get_conn().instances().export(  # noqa # pylint: disable=no-member
+            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,  # type:ignore
+                                             operation_name=operation_name)
 
     @CloudBaseHook.fallback_to_default_project_id
     def import_instance(self, instance: str, body: Dict, project_id: Optional[str] = None) -> None:
@@ -357,16 +341,14 @@ def import_instance(self, instance: str, body: Dict, project_id: Optional[str] =
         :type project_id: str
         :return: None
         """
-        if not project_id:
-            raise ValueError("The project_id should be set")
         try:
-            response = self.get_conn().instances().import_(  # pylint: disable=no-member
+            response = self.get_conn().instances().import_(  # noqa # pylint: disable=no-member
                 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,
+            self._wait_for_operation_to_complete(project_id=project_id,  # type: ignore
                                                  operation_name=operation_name)
         except HttpError as ex:
             raise AirflowException(
@@ -388,7 +370,7 @@ def _wait_for_operation_to_complete(self, project_id: str, operation_name: str)
             raise ValueError("The project_id should be set")
         service = self.get_conn()
         while True:
-            operation_response = service.operations().get(  # pylint: disable=no-member
+            operation_response = service.operations().get(  # noqa # pylint: disable=no-member
                 project=project_id,
                 operation=operation_name,
             ).execute(num_retries=self.num_retries)