apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

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

@@ -0,0 +1,109 @@
+#
+# 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 collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from google.cloud.logging_v2.services.config_service_v2 import ConfigServiceV2Client
+from google.cloud.logging_v2.types import (
+    CreateSinkRequest,
+    DeleteSinkRequest,
+    GetSinkRequest,
+    ListSinksRequest,
+    LogSink,
+    UpdateSinkRequest,
+)
+
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
+if TYPE_CHECKING:
+    from google.protobuf.field_mask_pb2 import FieldMask
+
+
+class CloudLoggingHook(GoogleBaseHook):
+    """
+    Hook for Google Cloud Logging Log Sinks API.
+
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :param impersonation_chain: Optional service account to impersonate.
+    """
+
+    def __init__(
+        self,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, **kwargs)
+        self._client: ConfigServiceV2Client | None = None
+
+    def get_conn(self) -> ConfigServiceV2Client:
+        """Return the Google Cloud Logging Config client."""
+        if not self._client:
+            self._client = ConfigServiceV2Client(credentials=self.get_credentials(), client_info=CLIENT_INFO)
+        return self._client
+
+    def get_parent(self, project_id):
+        return f"projects/{project_id}"
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def create_sink(
+        self, sink: LogSink | dict, unique_writer_identity: bool = True, project_id: str = PROVIDE_PROJECT_ID
+    ) -> LogSink:
+        if isinstance(sink, dict):
+            sink = LogSink(**sink)
+        request = CreateSinkRequest(
+            parent=self.get_parent(project_id), sink=sink, unique_writer_identity=unique_writer_identity
+        )
+        return self.get_conn().create_sink(request=request)
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def get_sink(self, sink_name: str, project_id: str = PROVIDE_PROJECT_ID) -> LogSink:
+        request = GetSinkRequest(sink_name=f"projects/{project_id}/sinks/{sink_name}")
+        return self.get_conn().get_sink(request=request)
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def list_sinks(self, page_size: int | None = None, project_id: str = PROVIDE_PROJECT_ID) -> list[LogSink]:
+        request = ListSinksRequest(parent=self.get_parent(project_id), page_size=page_size)
+        return list(self.get_conn().list_sinks(request=request))
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def delete_sink(self, sink_name: str, project_id: str = PROVIDE_PROJECT_ID) -> None:
+        request = DeleteSinkRequest(sink_name=f"projects/{project_id}/sinks/{sink_name}")
+        self.get_conn().delete_sink(request=request)
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def update_sink(
+        self,
+        sink_name: str,
+        sink: LogSink | dict,
+        unique_writer_identity: bool,
+        update_mask: FieldMask | dict,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> LogSink:
+        if isinstance(sink, dict):
+            sink = LogSink(**sink)
+        request = UpdateSinkRequest(
+            sink_name=f"projects/{project_id}/sinks/{sink_name}",
+            sink=sink,
+            unique_writer_identity=unique_writer_identity,
+            update_mask=update_mask,
+        )
+        return self.get_conn().update_sink(request=request)

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

@@ -31,9 +31,6 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow import version
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core import path_template
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -47,6 +44,10 @@ from google.cloud.redis_v1 import (
     OutputConfig,
 )
 
+from airflow import version
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask

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

@@ -21,9 +21,6 @@ import itertools
 from collections.abc import Iterable, Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.cloud.run_v2 import (
     CreateJobRequest,
     CreateServiceRequest,
@@ -41,7 +38,15 @@ from google.cloud.run_v2 import (
     ServicesClient,
     UpdateJobRequest,
 )
-from google.longrunning import operations_pb2  # type: ignore[attr-defined]
+from google.longrunning import operations_pb2
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import (
+    PROVIDE_PROJECT_ID,
+    GoogleBaseAsyncHook,
+    GoogleBaseHook,
+)
 
 if TYPE_CHECKING:
     from google.api_core import operation
@@ -158,7 +163,7 @@ class CloudRunHook(GoogleBaseHook):
         return list(itertools.islice(jobs, limit))
 
 
-class CloudRunAsyncHook(GoogleBaseHook):
+class CloudRunAsyncHook(GoogleBaseAsyncHook):
     """
     Async hook for the Google Cloud Run service.
 
@@ -173,6 +178,8 @@ class CloudRunAsyncHook(GoogleBaseHook):
         account from the list granting this role to the originating account.
     """
 
+    sync_hook_class = CloudRunHook
+
     def __init__(
         self,
         gcp_conn_id: str = "google_cloud_default",
@@ -182,16 +189,16 @@ class CloudRunAsyncHook(GoogleBaseHook):
         self._client: JobsAsyncClient | None = None
         super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, **kwargs)
 
-    def get_conn(self):
+    async def get_conn(self):
         if self._client is None:
-            self._client = JobsAsyncClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
+            sync_hook = await self.get_sync_hook()
+            self._client = JobsAsyncClient(credentials=sync_hook.get_credentials(), client_info=CLIENT_INFO)
 
         return self._client
 
     async def get_operation(self, operation_name: str) -> operations_pb2.Operation:
-        return await self.get_conn().get_operation(
-            operations_pb2.GetOperationRequest(name=operation_name), timeout=120
-        )
+        conn = await self.get_conn()
+        return await conn.get_operation(operations_pb2.GetOperationRequest(name=operation_name), timeout=120)
 
 
 class CloudRunServiceHook(GoogleBaseHook):
@@ -257,7 +264,7 @@ class CloudRunServiceHook(GoogleBaseHook):
         return operation.result()
 
 
-class CloudRunServiceAsyncHook(GoogleBaseHook):
+class CloudRunServiceAsyncHook(GoogleBaseAsyncHook):
     """
     Async hook for the Google Cloud Run services.
 
@@ -272,6 +279,8 @@ class CloudRunServiceAsyncHook(GoogleBaseHook):
         account from the list granting this role to the originating account.
     """
 
+    sync_hook_class = CloudRunServiceHook
+
     def __init__(
         self,
         gcp_conn_id: str = "google_cloud_default",

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

@@ -38,7 +38,7 @@ from inspect import signature
 from pathlib import Path
 from subprocess import PIPE, Popen
 from tempfile import NamedTemporaryFile, _TemporaryFileWrapper, gettempdir
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 from urllib.parse import quote_plus
 
 import httpx
@@ -50,8 +50,14 @@ from googleapiclient.errors import HttpError
 # 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.version_compat import AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+    from airflow.sdk import Connection
+else:
+    from airflow.models import Connection  # type: ignore[assignment,attr-defined,no-redef]
+
+from airflow.providers.common.compat.sdk import BaseHook
 from airflow.providers.google.cloud.hooks.secret_manager import (
     GoogleCloudSecretManagerHook,
 )
@@ -64,10 +70,10 @@ from airflow.providers.google.common.hooks.base_google import (
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
+    from google.cloud.secretmanager_v1 import AccessSecretVersionResponse
     from requests import Session
 
     from airflow.providers.common.sql.hooks.sql import DbApiHook
-    from google.cloud.secretmanager_v1 import AccessSecretVersionResponse
 
 UNIX_PATH_MAX = 108
 
@@ -502,7 +508,7 @@ class CloudSqlProxyRunner(LoggingMixin):
     :param project_id: Optional id of the Google Cloud project to connect to - it overwrites
         default project id taken from the Google Cloud connection.
     :param sql_proxy_version: Specific version of SQL proxy to download
-        (for example 'v1.13'). By default latest version is downloaded.
+        (for example 'v1.13'). By default, latest version is downloaded.
     :param sql_proxy_binary_path: If specified, then proxy will be
         used from the path specified rather than dynamically generated. This means
         that if the binary is not present in that path it will also be downloaded.
@@ -634,35 +640,32 @@ class CloudSqlProxyRunner(LoggingMixin):
         self._download_sql_proxy_if_needed()
         if self.sql_proxy_process:
             raise AirflowException(f"The sql proxy is already running: {self.sql_proxy_process}")
-        else:
-            command_to_run = [self.sql_proxy_path]
-            command_to_run.extend(self.command_line_parameters)
-            self.log.info("Creating directory %s", self.cloud_sql_proxy_socket_directory)
-            Path(self.cloud_sql_proxy_socket_directory).mkdir(parents=True, exist_ok=True)
-            command_to_run.extend(self._get_credential_parameters())
-            self.log.info("Running the command: `%s`", " ".join(command_to_run))
-
-            self.sql_proxy_process = Popen(command_to_run, stdin=PIPE, stdout=PIPE, stderr=PIPE)
-            self.log.info("The pid of cloud_sql_proxy: %s", self.sql_proxy_process.pid)
-            while True:
-                line = (
-                    self.sql_proxy_process.stderr.readline().decode("utf-8")
-                    if self.sql_proxy_process.stderr
-                    else ""
-                )
-                return_code = self.sql_proxy_process.poll()
-                if line == "" and return_code is not None:
-                    self.sql_proxy_process = None
-                    raise AirflowException(
-                        f"The cloud_sql_proxy finished early with return code {return_code}!"
-                    )
-                if line != "":
-                    self.log.info(line)
-                if "googleapi: Error" in line or "invalid instance name:" in line:
-                    self.stop_proxy()
-                    raise AirflowException(f"Error when starting the cloud_sql_proxy {line}!")
-                if "Ready for new connections" in line:
-                    return
+        command_to_run = [self.sql_proxy_path]
+        command_to_run.extend(self.command_line_parameters)
+        self.log.info("Creating directory %s", self.cloud_sql_proxy_socket_directory)
+        Path(self.cloud_sql_proxy_socket_directory).mkdir(parents=True, exist_ok=True)
+        command_to_run.extend(self._get_credential_parameters())
+        self.log.info("Running the command: `%s`", " ".join(command_to_run))
+
+        self.sql_proxy_process = Popen(command_to_run, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+        self.log.info("The pid of cloud_sql_proxy: %s", self.sql_proxy_process.pid)
+        while True:
+            line = (
+                self.sql_proxy_process.stderr.readline().decode("utf-8")
+                if self.sql_proxy_process.stderr
+                else ""
+            )
+            return_code = self.sql_proxy_process.poll()
+            if line == "" and return_code is not None:
+                self.sql_proxy_process = None
+                raise AirflowException(f"The cloud_sql_proxy finished early with return code {return_code}!")
+            if line != "":
+                self.log.info(line)
+            if "googleapi: Error" in line or "invalid instance name:" in line:
+                self.stop_proxy()
+                raise AirflowException(f"Error when starting the cloud_sql_proxy {line}!")
+            if "Ready for new connections" in line:
+                return
 
     def stop_proxy(self) -> None:
         """
@@ -672,10 +675,9 @@ class CloudSqlProxyRunner(LoggingMixin):
         """
         if not self.sql_proxy_process:
             raise AirflowException("The sql proxy is not started yet")
-        else:
-            self.log.info("Stopping the cloud_sql_proxy pid: %s", self.sql_proxy_process.pid)
-            self.sql_proxy_process.kill()
-            self.sql_proxy_process = None
+        self.log.info("Stopping the cloud_sql_proxy pid: %s", self.sql_proxy_process.pid)
+        self.sql_proxy_process.kill()
+        self.sql_proxy_process = None
         # Cleanup!
         self.log.info("Removing the socket directory: %s", self.cloud_sql_proxy_socket_directory)
         shutil.rmtree(self.cloud_sql_proxy_socket_directory, ignore_errors=True)
@@ -691,7 +693,7 @@ class CloudSqlProxyRunner(LoggingMixin):
             self.log.info("Skipped removing proxy - it was not downloaded: %s", self.sql_proxy_path)
         if os.path.isfile(self.credentials_path):
             self.log.info("Removing generated credentials file %s", self.credentials_path)
-            # Here file cannot be delete by concurrent task (each task has its own copy)
+            # Here file cannot be deleted by concurrent task (each task has its own copy)
             os.remove(self.credentials_path)
 
     def get_proxy_version(self) -> str | None:
@@ -704,8 +706,7 @@ class CloudSqlProxyRunner(LoggingMixin):
         matched = re.search("[Vv]ersion (.*?);", result)
         if matched:
             return matched.group(1)
-        else:
-            return None
+        return None
 
     def get_socket_path(self) -> str:
         """
@@ -852,8 +853,8 @@ class CloudSQLDatabaseHook(BaseHook):
             self.user = self._get_iam_db_login()
             self.password = self._generate_login_token(service_account=self.cloudsql_connection.login)
         else:
-            self.user = self.cloudsql_connection.login
-            self.password = self.cloudsql_connection.password
+            self.user = cast("str", self.cloudsql_connection.login)
+            self.password = cast("str", self.cloudsql_connection.password)
         self.public_ip = self.cloudsql_connection.host
         self.public_port = self.cloudsql_connection.port
         self.ssl_cert = ssl_cert
@@ -908,10 +909,9 @@ class CloudSQLDatabaseHook(BaseHook):
         secret_data = json.loads(base64.b64decode(secret.payload.data))
         if cert_name in secret_data:
             return secret_data[cert_name]
-        else:
-            raise AirflowException(
-                "Invalid secret format. Expected dictionary with keys: `sslcert`, `sslkey`, `sslrootcert`"
-            )
+        raise AirflowException(
+            "Invalid secret format. Expected dictionary with keys: `sslcert`, `sslkey`, `sslrootcert`"
+        )
 
     def _set_temporary_ssl_file(
         self, cert_name: str, cert_path: str | None = None, cert_value: str | None = None
@@ -1051,15 +1051,26 @@ class CloudSQLDatabaseHook(BaseHook):
     def _quote(value) -> str | None:
         return quote_plus(value) if value else None
 
-    def _generate_connection_uri(self) -> str:
+    def _reserve_port(self):
         if self.use_proxy:
             if self.sql_proxy_use_tcp:
                 if not self.sql_proxy_tcp_port:
                     self.reserve_free_tcp_port()
             if not self.sql_proxy_unique_path:
                 self.sql_proxy_unique_path = self._generate_unique_path()
+
+    def _generate_connection_uri(self) -> str:
+        self._reserve_port()
         if not self.database_type:
             raise ValueError("The database_type should be set")
+        if not self.user:
+            raise AirflowException("The login parameter needs to be set in connection")
+        if not self.public_ip:
+            raise AirflowException("The location parameter needs to be set in connection")
+        if not self.password:
+            raise AirflowException("The password parameter needs to be set in connection")
+        if not self.database:
+            raise AirflowException("The database parameter needs to be set in connection")
 
         database_uris = CONNECTION_URIS[self.database_type]
         ssl_spec = None
@@ -1078,14 +1089,6 @@ class CloudSQLDatabaseHook(BaseHook):
                 ssl_spec = {"cert": self.sslcert, "key": self.sslkey, "ca": self.sslrootcert}
             else:
                 format_string = public_uris["non-ssl"]
-        if not self.user:
-            raise AirflowException("The login parameter needs to be set in connection")
-        if not self.public_ip:
-            raise AirflowException("The location parameter needs to be set in connection")
-        if not self.password:
-            raise AirflowException("The password parameter needs to be set in connection")
-        if not self.database:
-            raise AirflowException("The database parameter needs to be set in connection")
 
         connection_uri = format_string.format(
             user=quote_plus(self.user) if self.user else "",
@@ -1109,6 +1112,8 @@ class CloudSQLDatabaseHook(BaseHook):
         return connection_uri
 
     def _get_instance_socket_name(self) -> str:
+        if self.project_id is None:
+            raise ValueError("The project_id should not be none")
         return self.project_id + ":" + self.location + ":" + self.instance
 
     def _get_sqlproxy_instance_specification(self) -> str:
@@ -1117,6 +1122,69 @@ class CloudSQLDatabaseHook(BaseHook):
             instance_specification += f"=tcp:{self.sql_proxy_tcp_port}"
         return instance_specification
 
+    def _generate_connection_parameters(self) -> dict:
+        self._reserve_port()
+        if not self.database_type:
+            raise ValueError("The database_type should be set")
+        if not self.user:
+            raise AirflowException("The login parameter needs to be set in connection")
+        if not self.public_ip:
+            raise AirflowException("The location parameter needs to be set in connection")
+        if not self.password:
+            raise AirflowException("The password parameter needs to be set in connection")
+        if not self.database:
+            raise AirflowException("The database parameter needs to be set in connection")
+
+        connection_parameters = {}
+
+        connection_parameters["conn_type"] = self.database_type
+        connection_parameters["login"] = self.user
+        connection_parameters["password"] = self.password
+        connection_parameters["schema"] = self.database
+        connection_parameters["extra"] = {}
+
+        database_uris = CONNECTION_URIS[self.database_type]
+        if self.use_proxy:
+            proxy_uris = database_uris["proxy"]
+            if self.sql_proxy_use_tcp:
+                connection_parameters["host"] = "127.0.0.1"
+                connection_parameters["port"] = self.sql_proxy_tcp_port
+            else:
+                socket_path = f"{self.sql_proxy_unique_path}/{self._get_instance_socket_name()}"
+                if "localhost" in proxy_uris["socket"]:
+                    connection_parameters["host"] = "localhost"
+                    connection_parameters["extra"].update({"unix_socket": socket_path})
+                else:
+                    connection_parameters["host"] = socket_path
+        else:
+            public_uris = database_uris["public"]
+            if self.use_ssl:
+                connection_parameters["host"] = self.public_ip
+                connection_parameters["port"] = self.public_port
+                if "ssl_spec" in public_uris["ssl"]:
+                    connection_parameters["extra"].update(
+                        {
+                            "ssl": json.dumps(
+                                {"cert": self.sslcert, "key": self.sslkey, "ca": self.sslrootcert}
+                            )
+                        }
+                    )
+                else:
+                    connection_parameters["extra"].update(
+                        {
+                            "sslmode": "verify-ca",
+                            "sslcert": self.sslcert,
+                            "sslkey": self.sslkey,
+                            "sslrootcert": self.sslrootcert,
+                        }
+                    )
+            else:
+                connection_parameters["host"] = self.public_ip
+                connection_parameters["port"] = self.public_port
+        if connection_parameters.get("extra"):
+            connection_parameters["extra"] = json.dumps(connection_parameters["extra"])
+        return connection_parameters
+
     def create_connection(self) -> Connection:
         """
         Create a connection.
@@ -1124,8 +1192,11 @@ class CloudSQLDatabaseHook(BaseHook):
         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)
+        if AIRFLOW_V_3_1_PLUS:
+            kwargs = self._generate_connection_parameters()
+        else:
+            kwargs = {"uri": self._generate_connection_uri()}
+        connection = Connection(conn_id=self.db_conn_id, **kwargs)
         self.log.info("Creating connection %s", self.db_conn_id)
         return connection
 
@@ -1141,6 +1212,8 @@ class CloudSQLDatabaseHook(BaseHook):
             raise ValueError("Proxy runner can only be retrieved in case of use_proxy = True")
         if not self.sql_proxy_unique_path:
             raise ValueError("The sql_proxy_unique_path should be set")
+        if self.project_id is None:
+            raise ValueError("The project_id should not be None")
         return CloudSqlProxyRunner(
             path_prefix=self.sql_proxy_unique_path,
             instance_specification=self._get_sqlproxy_instance_specification(),
@@ -1177,9 +1250,9 @@ class CloudSQLDatabaseHook(BaseHook):
                 raise ValueError("The db_hook should be set")
             if not isinstance(self.db_hook, PostgresHook):
                 raise ValueError(f"The db_hook should be PostgresHook and is {type(self.db_hook)}")
-            conn = getattr(self.db_hook, "conn")
-            if conn and conn.notices:
-                for output in self.db_hook.conn.notices:
+            conn = getattr(self.db_hook, "conn", None)
+            if conn and hasattr(conn, "notices") and conn.notices:
+                for output in conn.notices:
                     self.log.info(output)
 
     def reserve_free_tcp_port(self) -> None:
@@ -1205,8 +1278,7 @@ class CloudSQLDatabaseHook(BaseHook):
 
         if self.database_type == "postgres":
             return self.cloudsql_connection.login.split(".gserviceaccount.com")[0]
-        else:
-            return self.cloudsql_connection.login.split("@")[0]
+        return self.cloudsql_connection.login.split("@")[0]
 
     def _generate_login_token(self, service_account) -> str:
         """Generate an IAM login token for Cloud SQL and return the token."""
@@ -1215,7 +1287,7 @@ class CloudSQLDatabaseHook(BaseHook):
         cloud_sql_hook = CloudSQLHook(api_version="v1", gcp_conn_id=self.gcp_conn_id)
 
         with cloud_sql_hook.provide_authorized_gcloud():
-            proc = subprocess.run(cmd, capture_output=True)
+            proc = subprocess.run(cmd, check=False, capture_output=True)
 
         if proc.returncode != 0:
             stderr_last_20_lines = "\n".join(proc.stderr.decode().strip().splitlines()[-20:])

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

@@ -36,6 +36,14 @@ from copy import deepcopy
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
+from google.cloud.storage_transfer_v1 import (
+    ListTransferJobsRequest,
+    RunTransferJobRequest,
+    StorageTransferServiceAsyncClient,
+    TransferJob,
+    TransferOperation,
+)
+from google.protobuf.json_format import MessageToDict
 from googleapiclient.discovery import Resource, build
 from googleapiclient.errors import HttpError
 
@@ -46,21 +54,14 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseAsyncHook,
     GoogleBaseHook,
 )
-from google.api_core import protobuf_helpers
-from google.cloud.storage_transfer_v1 import (
-    ListTransferJobsRequest,
-    StorageTransferServiceAsyncClient,
-    TransferJob,
-    TransferOperation,
-)
 
 if TYPE_CHECKING:
-    from proto import Message
-
+    from google.api_core import operation_async
     from google.cloud.storage_transfer_v1.services.storage_transfer_service.pagers import (
         ListTransferJobsAsyncPager,
     )
-    from google.longrunning import operations_pb2  # type: ignore[attr-defined]
+    from google.longrunning import operations_pb2
+    from proto import Message
 
 log = logging.getLogger(__name__)
 
@@ -219,7 +220,7 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
                     return (
                         self.get_conn().transferJobs().create(body=body).execute(num_retries=self.num_retries)
                     )
-                elif transfer_job.get(STATUS) == GcpTransferJobsStatus.DISABLED:
+                if transfer_job.get(STATUS) == GcpTransferJobsStatus.DISABLED:
                     return self.enable_transfer_job(job_name=job_name, project_id=body.get(PROJECT_ID))
             else:
                 raise e
@@ -604,7 +605,7 @@ class CloudDataTransferServiceAsyncHook(GoogleBaseAsyncHook):
         self,
         request_filter: dict | None = None,
         **kwargs,
-    ) -> list[TransferOperation]:
+    ) -> list[dict[str, Any]]:
         """
         Get a transfer operation in Google Storage Transfer Service.
 
@@ -661,7 +662,12 @@ class CloudDataTransferServiceAsyncHook(GoogleBaseAsyncHook):
             )
 
         transfer_operations = [
-            protobuf_helpers.from_any_pb(TransferOperation, op.metadata) for op in operations
+            MessageToDict(
+                getattr(op, "_pb", op),
+                preserving_proto_field_name=True,
+                use_integers_for_enums=True,
+            )
+            for op in operations
         ]
 
         return transfer_operations
@@ -678,7 +684,7 @@ class CloudDataTransferServiceAsyncHook(GoogleBaseAsyncHook):
 
     @staticmethod
     async def operations_contain_expected_statuses(
-        operations: list[TransferOperation], expected_statuses: set[str] | str
+        operations: list[dict[str, Any]], expected_statuses: set[str] | str
     ) -> bool:
         """
         Check whether an operation exists with the expected status.
@@ -697,7 +703,7 @@ class CloudDataTransferServiceAsyncHook(GoogleBaseAsyncHook):
         if not operations:
             return False
 
-        current_statuses = {operation.status.name for operation in operations}
+        current_statuses = {TransferOperation.Status(op["metadata"]["status"]).name for op in operations}
 
         if len(current_statuses - expected_statuses_set) != len(current_statuses):
             return True
@@ -708,3 +714,17 @@ class CloudDataTransferServiceAsyncHook(GoogleBaseAsyncHook):
                 f"Expected: {', '.join(expected_statuses_set)}"
             )
         return False
+
+    async def run_transfer_job(self, job_name: str) -> operation_async.AsyncOperation:
+        """
+        Run Google Storage Transfer Service job.
+
+        :param job_name: (Required) Name of the job to run.
+        """
+        client = await self.get_conn()
+        request = RunTransferJobRequest(
+            job_name=job_name,
+            project_id=self.project_id,
+        )
+        operation = await client.run_transfer_job(request=request)
+        return operation