apache-airflow-providers-google 10.2.0rc1__py3-none-any.whl → 10.3.0rc1__py3-none-any.whl

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

@@ -39,16 +39,18 @@ from typing import Any, Sequence
 from urllib.parse import quote_plus
 
 import httpx
+from aiohttp import ClientSession
+from gcloud.aio.auth import AioSession, Token
 from googleapiclient.discovery import Resource, build
 from googleapiclient.errors import HttpError
-
-from airflow.exceptions import AirflowException
+from requests import Session
 
 # Number of retries - used by googleapiclient method calls to perform retries
 # For requests that are "retriable"
+from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
 from airflow.models import Connection
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook, get_field
 from airflow.providers.mysql.hooks.mysql import MySqlHook
 from airflow.providers.postgres.hooks.postgres import PostgresHook
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -300,8 +302,7 @@ class CloudSQLHook(GoogleBaseHook):
         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) -> None:
+    def export_instance(self, instance: str, body: dict, project_id: str):
         """
         Exports data from a Cloud SQL instance to a Cloud Storage bucket as a SQL dump
         or CSV file.
@@ -321,7 +322,7 @@ class CloudSQLHook(GoogleBaseHook):
             .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 operation_name
 
     @GoogleBaseHook.fallback_to_default_project_id
     def import_instance(self, instance: str, body: dict, project_id: str) -> None:
@@ -376,6 +377,7 @@ class CloudSQLHook(GoogleBaseHook):
         except HttpError as ex:
             raise AirflowException(f"Cloning of instance {instance} failed: {ex.content}")
 
+    @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
     ) -> None:
@@ -412,6 +414,42 @@ CLOUD_SQL_PROXY_VERSION_DOWNLOAD_URL = (
 )
 
 
+class CloudSQLAsyncHook(GoogleBaseAsyncHook):
+    """Class to get asynchronous hook for Google Cloud SQL."""
+
+    sync_hook_class = CloudSQLHook
+
+    async def _get_conn(self, session: Session, url: str):
+        scopes = [
+            "https://www.googleapis.com/auth/cloud-platform",
+            "https://www.googleapis.com/auth/sqlservice.admin",
+        ]
+
+        async with Token(scopes=scopes) as token:
+            session_aio = AioSession(session)
+            headers = {
+                "Authorization": f"Bearer {await token.get()}",
+            }
+        return await session_aio.get(url=url, headers=headers)
+
+    async def get_operation_name(self, project_id: str, operation_name: str, session):
+        url = f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{project_id}/operations/{operation_name}"
+        return await self._get_conn(url=str(url), session=session)
+
+    async def get_operation(self, project_id: str, operation_name: str):
+        async with ClientSession() as session:
+            try:
+                operation = await self.get_operation_name(
+                    project_id=project_id,
+                    operation_name=operation_name,
+                    session=session,
+                )
+                operation = await operation.json(content_type=None)
+            except HttpError as e:
+                raise e
+            return operation
+
+
 class CloudSqlProxyRunner(LoggingMixin):
     """
     Downloads and runs cloud-sql-proxy as subprocess of the Python process.
@@ -840,12 +878,13 @@ class CloudSQLDatabaseHook(BaseHook):
 
     @staticmethod
     def _generate_unique_path() -> str:
-        """
-        We are not using mkdtemp here as the path generated with mkdtemp
-        can be close to 60 characters and there is a limitation in
-        length of socket path to around 100 characters in total.
-        We append project/location/instance to it later and postgres
-        appends its own prefix, so we chose a shorter "${tempdir()}[8 random characters]".
+        """Generate a unique path.
+
+        We don't using mkdtemp here since it can generate paths close to 60
+        characters. We append project/location/instance to the path, Postgres
+        will then appends its own prefix, making the resulting path exceed the
+        100 character length limitation of a socket path. This generates a
+        shorter path ``${tempdir()}[8 random characters]``.
         """
         random.seed()
         while True:
@@ -926,9 +965,10 @@ class CloudSQLDatabaseHook(BaseHook):
         return instance_specification
 
     def create_connection(self) -> Connection:
-        """
-        Create Connection object, according to whether it uses proxy, TCP, UNIX sockets, SSL.
-        Connection ID will be randomly generated.
+        """Create a connection.
+
+        Connection ID will be randomly generated according to whether it uses
+        proxy, TCP, UNIX sockets, SSL.
         """
         uri = self._generate_connection_uri()
         connection = Connection(conn_id=self.db_conn_id, uri=uri)
@@ -936,9 +976,9 @@ class CloudSQLDatabaseHook(BaseHook):
         return connection
 
     def get_sqlproxy_runner(self) -> CloudSqlProxyRunner:
-        """
-        Retrieve Cloud SQL Proxy runner. It is used to manage the proxy
-        lifecycle per task.
+        """Retrieve Cloud SQL Proxy runner.
+
+        It is used to manage the proxy lifecycle per task.
 
         :return: The Cloud SQL Proxy runner.
         """
@@ -956,9 +996,10 @@ class CloudSQLDatabaseHook(BaseHook):
         )
 
     def get_database_hook(self, connection: Connection) -> PostgresHook | MySqlHook:
-        """
-        Retrieve database hook. This is the actual Postgres or MySQL database hook
-        that uses proxy or connects directly to the Google Cloud SQL database.
+        """Retrieve database hook.
+
+        This is the actual Postgres or MySQL database hook that uses proxy or
+        connects directly to the Google Cloud SQL database.
         """
         if self.database_type == "postgres":
             db_hook: PostgresHook | MySqlHook = PostgresHook(connection=connection, schema=self.database)
@@ -986,7 +1027,10 @@ class CloudSQLDatabaseHook(BaseHook):
         self.sql_proxy_tcp_port = self.reserved_tcp_socket.getsockname()[1]
 
     def free_reserved_port(self) -> None:
-        """Free TCP port. Makes it immediately ready to be used by Cloud SQL Proxy."""
+        """Free TCP port.
+
+        Makes it immediately ready to be used by Cloud SQL Proxy.
+        """
         if self.reserved_tcp_socket:
             self.reserved_tcp_socket.close()
             self.reserved_tcp_socket = None

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

@@ -39,7 +39,7 @@ TIME_TO_SLEEP_IN_SECONDS = 10
 
 
 class GcpTransferJobsStatus:
-    """Class with Google Cloud Transfer jobs statuses."""
+    """Google Cloud Transfer job status."""
 
     ENABLED = "ENABLED"
     DISABLED = "DISABLED"
@@ -47,7 +47,7 @@ class GcpTransferJobsStatus:
 
 
 class GcpTransferOperationStatus:
-    """Class with Google Cloud Transfer operations statuses."""
+    """Google Cloud Transfer operation status."""
 
     IN_PROGRESS = "IN_PROGRESS"
     PAUSED = "PAUSED"
@@ -104,9 +104,7 @@ NEGATIVE_STATUSES = {GcpTransferOperationStatus.FAILED, GcpTransferOperationStat
 
 
 def gen_job_name(job_name: str) -> str:
-    """
-    Adds unique suffix to job name. If suffix already exists, updates it.
-    Suffix — current timestamp.
+    """Add a unique suffix to the job name.
 
     :param job_name:
     :return: job_name with suffix
@@ -116,11 +114,10 @@ def gen_job_name(job_name: str) -> str:
 
 
 class CloudDataTransferServiceHook(GoogleBaseHook):
-    """
-    Hook for Google Storage Transfer Service.
+    """Google Storage Transfer Service functionalities.
 
-    All the methods in the hook where project_id is used must be called with
-    keyword arguments rather than positional.
+    All methods in the hook with *project_id* in the signature must be called
+    with keyword arguments rather than positional.
     """
 
     def __init__(
@@ -130,10 +127,11 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
+        if "delegate_to" in kwargs:
             raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
+                "The `delegate_to` parameter has been deprecated before and "
+                "finally removed in this version of Google Provider. You MUST "
+                "convert it to `impersonate_chain`."
             )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
@@ -143,8 +141,7 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         self._conn = None
 
     def get_conn(self) -> Resource:
-        """
-        Retrieves connection to Google Storage Transfer service.
+        """Retrieve connection to Google Storage Transfer service.
 
         :return: Google Storage Transfer service object
         """
@@ -156,13 +153,11 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         return self._conn
 
     def create_transfer_job(self, body: dict) -> dict:
-        """
-        Creates a transfer job that runs periodically.
+        """Create a transfer job that runs periodically.
 
-        :param body: (Required) A request body, as described in
+        :param body: (Required) The request body, as described in
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/patch#request-body
-        :return: transfer job.
-            See:
+        :return: The transfer job. See:
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs#TransferJob
         """
         body = self._inject_project_id(body, BODY, PROJECT_ID)
@@ -200,9 +195,7 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def get_transfer_job(self, job_name: str, project_id: str) -> dict:
-        """
-        Gets the latest state of a long-running operation in Google Storage
-        Transfer Service.
+        """Get latest state of a long-running Google Storage Transfer Service job.
 
         :param job_name: (Required) Name of the job to be fetched
         :param project_id: (Optional) the ID of the project that owns the Transfer
@@ -218,9 +211,9 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         )
 
     def list_transfer_job(self, request_filter: dict | None = None, **kwargs) -> list[dict]:
-        """
-        Lists long-running operations in Google Storage Transfer
-        Service that match the specified filter.
+        """List long-running operations in Google Storage Transfer Service.
+
+        A filter can be specified to match only certain entries.
 
         :param request_filter: (Required) A request filter, as described in
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/list#body.QUERY_PARAMETERS.filter
@@ -252,8 +245,7 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def enable_transfer_job(self, job_name: str, project_id: str) -> dict:
-        """
-        New transfers will be performed based on the schedule.
+        """Make new transfers be performed based on the schedule.
 
         :param job_name: (Required) Name of the job to be updated
         :param project_id: (Optional) the ID of the project that owns the Transfer
@@ -276,8 +268,7 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         )
 
     def update_transfer_job(self, job_name: str, body: dict) -> dict:
-        """
-        Updates a transfer job that runs periodically.
+        """Update a transfer job that runs periodically.
 
         :param job_name: (Required) Name of the job to be updated
         :param body: A request body, as described in
@@ -294,11 +285,11 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def delete_transfer_job(self, job_name: str, project_id: str) -> None:
-        """
-        Deletes a transfer job. This is a soft delete. After a transfer job is
-        deleted, the job and all the transfer executions are subject to garbage
-        collection. Transfer jobs become eligible for garbage collection
-        30 days after soft delete.
+        """Delete a transfer job.
+
+        This is a soft delete. After a transfer job is deleted, the job and all
+        the transfer executions are subject to garbage collection. Transfer jobs
+        become eligible for garbage collection 30 days after soft delete.
 
         :param job_name: (Required) Name of the job to be deleted
         :param project_id: (Optional) the ID of the project that owns the Transfer
@@ -320,21 +311,19 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         )
 
     def cancel_transfer_operation(self, operation_name: str) -> None:
-        """
-        Cancels an transfer operation in Google Storage Transfer Service.
+        """Cancel a transfer operation in Google Storage Transfer Service.
 
         :param operation_name: Name of the transfer operation.
         """
         self.get_conn().transferOperations().cancel(name=operation_name).execute(num_retries=self.num_retries)
 
     def get_transfer_operation(self, operation_name: str) -> dict:
-        """
-        Gets an transfer operation in Google Storage Transfer Service.
+        """Get a transfer operation in Google Storage Transfer Service.
 
         :param operation_name: (Required) Name of the transfer operation.
         :return: transfer operation
-            See:
-            https://cloud.google.com/storage-transfer/docs/reference/rest/v1/Operation
+
+        .. seealso:: https://cloud.google.com/storage-transfer/docs/reference/rest/v1/Operation
         """
         return (
             self.get_conn()
@@ -344,18 +333,15 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         )
 
     def list_transfer_operations(self, request_filter: dict | None = None, **kwargs) -> list[dict]:
-        """
-        Gets an transfer operation in Google Storage Transfer Service.
+        """Get a transfer operation in Google Storage Transfer Service.
 
         :param request_filter: (Required) A request filter, as described in
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/list#body.QUERY_PARAMETERS.filter
             With one additional improvement:
-
-            * project_id is optional if you have a project id defined
-              in the connection
-              See: :doc:`/connections/gcp`
-
         :return: transfer operation
+
+        The ``project_id`` parameter is optional if you have a project ID
+        defined in the connection. See: :doc:`/connections/gcp`
         """
         # To preserve backward compatibility
         # TODO: remove one day
@@ -390,16 +376,14 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         return operations
 
     def pause_transfer_operation(self, operation_name: str) -> None:
-        """
-        Pauses an transfer operation in Google Storage Transfer Service.
+        """Pause a transfer operation in Google Storage Transfer Service.
 
         :param operation_name: (Required) Name of the transfer operation.
         """
         self.get_conn().transferOperations().pause(name=operation_name).execute(num_retries=self.num_retries)
 
     def resume_transfer_operation(self, operation_name: str) -> None:
-        """
-        Resumes an transfer operation in Google Storage Transfer Service.
+        """Resume a transfer operation in Google Storage Transfer Service.
 
         :param operation_name: (Required) Name of the transfer operation.
         """
@@ -411,17 +395,14 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         expected_statuses: set[str] | None = None,
         timeout: float | timedelta | None = None,
     ) -> None:
-        """
-        Waits until the job reaches the expected state.
+        """Wait until the job reaches the expected state.
 
-        :param job: Transfer job
-            See:
+        :param job: The transfer job to wait for. See:
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs#TransferJob
-        :param expected_statuses: State that is expected
-            See:
+        :param expected_statuses: The expected state. See:
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferOperations#Status
-        :param timeout: Time in which the operation must end in seconds. If not specified, defaults to 60
-            seconds.
+        :param timeout: Time in which the operation must end in seconds. If not
+            specified, defaults to 60 seconds.
         """
         expected_statuses = (
             {GcpTransferOperationStatus.SUCCESS} if not expected_statuses else expected_statuses
@@ -458,20 +439,15 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
     def operations_contain_expected_statuses(
         operations: list[dict], expected_statuses: set[str] | str
     ) -> bool:
-        """
-        Checks whether the operation list has an operation with the
-        expected status, then returns true
-        If it encounters operations in FAILED or ABORTED state
-        throw :class:`airflow.exceptions.AirflowException`.
+        """Check whether an operation exists with the expected status.
 
         :param operations: (Required) List of transfer operations to check.
-        :param expected_statuses: (Required) status that is expected
-            See:
+        :param expected_statuses: (Required) The expected status. See:
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferOperations#Status
-        :return: If there is an operation with the expected state
-            in the operation list, returns true,
-        :raises: airflow.exceptions.AirflowException If it encounters operations
-            with a state in the list,
+        :return: If there is an operation with the expected state in the
+            operation list, returns true,
+        :raises AirflowException: If it encounters operations with state FAILED
+            or ABORTED in the list.
         """
         expected_statuses_set = (
             {expected_statuses} if isinstance(expected_statuses, str) else set(expected_statuses)

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

@@ -57,29 +57,25 @@ T = TypeVar("T", bound=Callable)
 def process_line_and_extract_dataflow_job_id_callback(
     on_new_job_id_callback: Callable[[str], None] | None
 ) -> Callable[[str], None]:
-    """
-    Returns callback which triggers function passed as `on_new_job_id_callback` when Dataflow job_id is found.
-    To be used for `process_line_callback` in
+    """Build callback that triggers the specified function.
+
+    The returned callback is intended to be used as ``process_line_callback`` in
     :py:class:`~airflow.providers.apache.beam.hooks.beam.BeamCommandRunner`.
 
     :param on_new_job_id_callback: Callback called when the job ID is known
     """
 
-    def _process_line_and_extract_job_id(
-        line: str,
-        # on_new_job_id_callback: Callable[[str], None] | None
-    ) -> None:
+    def _process_line_and_extract_job_id(line: str) -> None:
         # Job id info: https://goo.gl/SE29y9.
+        if on_new_job_id_callback is None:
+            return
         matched_job = JOB_ID_PATTERN.search(line)
-        if matched_job:
-            job_id = matched_job.group("job_id_java") or matched_job.group("job_id_python")
-            if on_new_job_id_callback:
-                on_new_job_id_callback(job_id)
-
-    def wrap(line: str):
-        return _process_line_and_extract_job_id(line)
+        if matched_job is None:
+            return
+        job_id = matched_job.group("job_id_java") or matched_job.group("job_id_python")
+        on_new_job_id_callback(job_id)
 
-    return wrap
+    return _process_line_and_extract_job_id
 
 
 def _fallback_variable_parameter(parameter_name: str, variable_key_name: str) -> Callable[[T], T]:

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

@@ -604,9 +604,9 @@ class DataformHook(GoogleBaseHook):
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> InstallNpmPackagesResponse:
-        """
-        Installs npm dependencies in the provided workspace. Requires "package.json"
-        to be created in workspace.
+        """Install NPM dependencies in the provided workspace.
+
+        Requires "package.json" to be created in the workspace.
 
         :param project_id: Required. The ID of the Google Cloud project where workspace located.
         :param region: Required. The ID of the Google Cloud region where workspace located.