apache-airflow-providers-google 15.1.0rc1__py3-none-any.whl → 19.3.0__py3-none-any.whl

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

@@ -21,7 +21,7 @@ from typing import Any
 
 from googleapiclient.discovery import Resource, build
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 

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

@@ -0,0 +1,196 @@
+#
+# 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.
+"""This module contains a Google Cloud GenAI Generative Model hook."""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, Any
+
+from google import genai
+
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
+if TYPE_CHECKING:
+    from google.genai.types import (
+        ContentListUnion,
+        ContentListUnionDict,
+        CountTokensConfigOrDict,
+        CountTokensResponse,
+        CreateCachedContentConfigOrDict,
+        CreateTuningJobConfigOrDict,
+        EmbedContentConfigOrDict,
+        EmbedContentResponse,
+        GenerateContentConfig,
+        TuningDatasetOrDict,
+        TuningJob,
+    )
+
+
+class GenAIGenerativeModelHook(GoogleBaseHook):
+    """Class for Google Cloud Generative AI Vertex AI hook."""
+
+    def get_genai_client(self, project_id: str, location: str):
+        return genai.Client(
+            vertexai=True,
+            project=project_id,
+            location=location,
+        )
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def embed_content(
+        self,
+        model: str,
+        location: str,
+        contents: ContentListUnion | ContentListUnionDict | list[str],
+        config: EmbedContentConfigOrDict | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> EmbedContentResponse:
+        """
+        Generate embeddings for words, phrases, sentences, and code.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param model: Required. The model to use.
+        :param contents: Optional. The contents to use for embedding.
+        :param config: Optional. Configuration for embeddings.
+        """
+        client = self.get_genai_client(project_id=project_id, location=location)
+
+        resp = client.models.embed_content(model=model, contents=contents, config=config)
+        return resp
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def generate_content(
+        self,
+        location: str,
+        model: str,
+        contents: ContentListUnionDict,
+        generation_config: GenerateContentConfig | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Make an API request to generate content using a model.
+
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param model: Required. The model to use.
+        :param contents: Required. The multi-part content of a message that a user or a program
+            gives to the generative model, in order to elicit a specific response.
+        :param generation_config: Optional. Generation configuration settings.
+        """
+        client = self.get_genai_client(project_id=project_id, location=location)
+        response = client.models.generate_content(
+            model=model,
+            contents=contents,
+            config=generation_config,
+        )
+
+        return response.text
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def supervised_fine_tuning_train(
+        self,
+        source_model: str,
+        location: str,
+        training_dataset: TuningDatasetOrDict,
+        tuning_job_config: CreateTuningJobConfigOrDict | dict[str, Any] | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> TuningJob:
+        """
+        Create a tuning job to adapt model behavior with a labeled dataset.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param source_model: Required. A pre-trained model optimized for performing natural
+            language tasks such as classification, summarization, extraction, content
+            creation, and ideation.
+        :param train_dataset: Required. Cloud Storage URI of your training dataset. The dataset
+            must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples.
+        :param tuning_job_config: Optional. Configuration of the Tuning job to be created.
+        """
+        client = self.get_genai_client(project_id=project_id, location=location)
+
+        tuning_job = client.tunings.tune(
+            base_model=source_model,
+            training_dataset=training_dataset,
+            config=tuning_job_config,
+        )
+
+        # Poll until completion
+        running = {"JOB_STATE_PENDING", "JOB_STATE_RUNNING"}
+        while tuning_job.state in running:
+            time.sleep(60)
+            tuning_job = client.tunings.get(name=tuning_job.name)
+
+        return tuning_job
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def count_tokens(
+        self,
+        location: str,
+        model: str,
+        contents: ContentListUnion | ContentListUnionDict,
+        config: CountTokensConfigOrDict | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> CountTokensResponse:
+        """
+        Use Count Tokens API to calculate the number of input tokens before sending a request to Gemini API.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param contents: Required. The multi-part content of a message that a user or a program
+            gives to the generative model, in order to elicit a specific response.
+        :param model: Required. Model,
+            supporting prompts with text-only input, including natural language
+            tasks, multi-turn text and code chat, and code generation. It can
+            output text and code.
+        :param config: Optional. Configuration for Count Tokens.
+        """
+        client = self.get_genai_client(project_id=project_id, location=location)
+        response = client.models.count_tokens(
+            model=model,
+            contents=contents,
+            config=config,
+        )
+
+        return response
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def create_cached_content(
+        self,
+        model: str,
+        location: str,
+        cached_content_config: CreateCachedContentConfigOrDict | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Create CachedContent to reduce the cost of requests containing repeat content.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param model: Required. The name of the publisher model to use for cached content.
+        :param cached_content_config: Optional. Configuration of the Cached Content.
+        """
+        client = self.get_genai_client(project_id=project_id, location=location)
+        resp = client.caches.create(
+            model=model,
+            config=cached_content_config,
+        )
+
+        return resp.name

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

@@ -30,7 +30,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.auth.transport import requests as google_requests
 
 # not sure why but mypy complains on missing `container_v1` but it is clearly there and is importable
-from google.cloud import exceptions  # type: ignore[attr-defined]
+from google.cloud import exceptions
 from google.cloud.container_v1 import ClusterManagerAsyncClient, ClusterManagerClient
 from google.cloud.container_v1.types import Cluster, Operation
 from kubernetes import client
@@ -38,9 +38,9 @@ from kubernetes_asyncio import client as async_client
 from kubernetes_asyncio.config.kube_config import FileOrData
 
 from airflow import version
-from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook, KubernetesHook
 from airflow.providers.cncf.kubernetes.kube_client import _enable_tcp_keepalive
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,
@@ -498,7 +498,7 @@ class GKEKubernetesAsyncHook(GoogleBaseAsyncHook, AsyncKubernetesHook):
         )
 
     @contextlib.asynccontextmanager
-    async def get_conn(self) -> async_client.ApiClient:  # type: ignore[override]
+    async def get_conn(self) -> async_client.ApiClient:
         kube_client = None
         try:
             kube_client = await self._load_config()

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

@@ -28,8 +28,7 @@ from looker_sdk.rtl import api_settings, auth_session, requests_transport, seria
 from looker_sdk.sdk.api40 import methods as methods40
 from packaging.version import parse as parse_version
 
-from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
 from airflow.version import version
 
 if TYPE_CHECKING:
@@ -39,6 +38,11 @@ if TYPE_CHECKING:
 class LookerHook(BaseHook):
     """Hook for Looker APIs."""
 
+    conn_name_attr = "looker_conn_id"
+    default_conn_name = "looker_default"
+    conn_type = "gcp_looker"
+    hook_name = "Google Looker"
+
     def __init__(
         self,
         looker_conn_id: str,

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

@@ -31,7 +31,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.auth.transport import requests as google_requests
 from google.cloud.managedkafka_v1 import Cluster, ConsumerGroup, ManagedKafkaClient, Topic, types
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 

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

@@ -23,14 +23,15 @@ import contextlib
 import logging
 import random
 import time
-from typing import TYPE_CHECKING, Callable
+from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 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 airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,
     GoogleBaseAsyncHook,
@@ -587,7 +588,7 @@ class MLEngineAsyncHook(GoogleBaseAsyncHook):
                 job = await self.get_job(
                     project_id=project_id,
                     job_id=job_id,
-                    session=session,  #  type: ignore
+                    session=session,  # type: ignore
                 )
                 job = await job.json(content_type=None)
                 self.log.info("Retrieving json_response: %s", job)

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

@@ -82,6 +82,7 @@ class PubSubHook(GoogleBaseHook):
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         enable_message_ordering: bool = False,
+        enable_open_telemetry_tracing: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -90,6 +91,7 @@ class PubSubHook(GoogleBaseHook):
             **kwargs,
         )
         self.enable_message_ordering = enable_message_ordering
+        self.enable_open_telemetry_tracing = enable_open_telemetry_tracing
         self._client = None
 
     def get_conn(self) -> PublisherClient:
@@ -104,6 +106,7 @@ class PubSubHook(GoogleBaseHook):
                 client_info=CLIENT_INFO,
                 publisher_options=PublisherOptions(
                     enable_message_ordering=self.enable_message_ordering,
+                    enable_open_telemetry_tracing=self.enable_open_telemetry_tracing,
                 ),
             )
         return self._client

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

@@ -47,6 +47,10 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
     See https://cloud.google.com/secret-manager
     """
 
+    def __init__(self, location: str | None = None, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.location = location
+
     @cached_property
     def client(self):
         """
@@ -54,7 +58,16 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
 
         :return: Secret Manager client.
         """
-        return SecretManagerServiceClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
+        if self.location is not None:
+            return SecretManagerServiceClient(
+                credentials=self.get_credentials(),
+                client_info=CLIENT_INFO,
+                client_options={"api_endpoint": f"secretmanager.{self.location}.rep.googleapis.com"},
+            )
+        return SecretManagerServiceClient(
+            credentials=self.get_credentials(),
+            client_info=CLIENT_INFO,
+        )
 
     def get_conn(self) -> SecretManagerServiceClient:
         """
@@ -64,6 +77,60 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         """
         return self.client
 
+    def _get_parent(self, project_id: str, location: str | None = None) -> str:
+        """
+        Return parent path.
+
+        :param project_id: Required. ID of the GCP project that owns the job.
+        :param location: Optional. Target location.  If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used
+            For more details : https://cloud.google.com/secret-manager/docs/locations
+
+        :return: Parent path.
+        """
+        _location = location or self.location
+        if _location is not None:
+            return self.client.common_location_path(project_id, _location)
+        return self.client.common_project_path(project_id)
+
+    def _get_secret_path(self, project_id: str, secret_id: str, location: str | None = None) -> str:
+        """
+        Return secret path.
+
+        :param project_id: Required. ID of the GCP project that owns the job.
+        :param secret_id: Required. Secret ID for which path is required.
+        :param location: Optional. Target location.  If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used
+            For more details : https://cloud.google.com/secret-manager/docs/locations
+
+        :return: Parent path.
+        """
+        _location = location or self.location
+        if _location is not None:
+            # Google's client library does not provide a method to construct regional secret paths, so constructing manually.
+            return f"projects/{project_id}/locations/{_location}/secrets/{secret_id}"
+        return self.client.secret_path(project_id, secret_id)
+
+    def _get_secret_version_path(
+        self, project_id: str, secret_id: str, secret_version: str, location: str | None = None
+    ) -> str:
+        """
+        Return secret version path.
+
+        :param project_id: Required. ID of the GCP project that owns the job.
+        :param secret_id: Required. Secret ID for which path is required.
+        :param secret_version: Required. Secret version for which path is required.
+        :param location: Optional. Target location.  If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used
+            For more details : https://cloud.google.com/secret-manager/docs/locations
+
+        :return: Parent path.
+        """
+        _location = location or self.location
+        if _location is not None:
+            # Google's client library does not provide a method to construct regional secret version paths, so constructing manually.
+            return (
+                f"projects/{project_id}/locations/{_location}/secrets/{secret_id}/versions/{secret_version}"
+            )
+        return self.client.secret_version_path(project_id, secret_id, secret_version)
+
     @GoogleBaseHook.fallback_to_default_project_id
     def create_secret(
         self,
@@ -73,6 +140,7 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
+        location: str | None = None,
     ) -> Secret:
         """
         Create a secret.
@@ -88,12 +156,20 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         :param retry: Optional. Designation of what errors, if any, should be retried.
         :param timeout: Optional. The timeout for this request.
         :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        :param location: Optional. Location where secret should be created. Used for creating regional secret.  If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used
+            For more details : https://cloud.google.com/secret-manager/docs/locations
         :return: Secret object.
         """
-        _secret = secret or {"replication": {"automatic": {}}}
+        if not secret:
+            _secret: dict | Secret = {}
+            if (location or self.location) is None:
+                _secret["replication"] = {"automatic": {}}
+        else:
+            _secret = secret
+
         response = self.client.create_secret(
             request={
-                "parent": f"projects/{project_id}",
+                "parent": self._get_parent(project_id, location),
                 "secret_id": secret_id,
                 "secret": _secret,
             },
@@ -113,6 +189,7 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
+        location: str | None = None,
     ) -> SecretVersion:
         """
         Add a version to the secret.
@@ -128,11 +205,13 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         :param retry: Optional. Designation of what errors, if any, should be retried.
         :param timeout: Optional. The timeout for this request.
         :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        :param location: Optional. Location where secret is located. Used for adding version to regional secret.  If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used
+            For more details : https://cloud.google.com/secret-manager/docs/locations
         :return: Secret version object.
         """
         response = self.client.add_secret_version(
             request={
-                "parent": f"projects/{project_id}/secrets/{secret_id}",
+                "parent": self._get_secret_path(project_id, secret_id, location),
                 "payload": secret_payload,
             },
             retry=retry,
@@ -152,6 +231,7 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
+        location: str | None = None,
     ) -> ListSecretsPager:
         """
         List secrets.
@@ -168,11 +248,13 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         :param retry: Optional. Designation of what errors, if any, should be retried.
         :param timeout: Optional. The timeout for this request.
         :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        :param location: Optional. The regional secrets stored in the provided location will be listed.  If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used
+            For more details : https://cloud.google.com/secret-manager/docs/locations
         :return: Secret List object.
         """
         response = self.client.list_secrets(
             request={
-                "parent": f"projects/{project_id}",
+                "parent": self._get_parent(project_id, location),
                 "page_size": page_size,
                 "page_token": page_token,
                 "filter": secret_filter,
@@ -185,18 +267,22 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         return response
 
     @GoogleBaseHook.fallback_to_default_project_id
-    def secret_exists(self, project_id: str, secret_id: str) -> bool:
+    def secret_exists(self, project_id: str, secret_id: str, location: str | None = None) -> bool:
         """
         Check whether secret exists.
 
         :param project_id: Required. ID of the GCP project that owns the job.
             If set to ``None`` or missing, the default project_id from the GCP connection is used.
         :param secret_id: Required. ID of the secret to find.
+        :param location: Optional. Location where secret is expected to be stored regionally.  If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used
+            For more details : https://cloud.google.com/secret-manager/docs/locations
         :return: True if the secret exists, False otherwise.
         """
         secret_filter = f"name:{secret_id}"
-        secret_name = self.client.secret_path(project_id, secret_id)
-        for secret in self.list_secrets(project_id=project_id, page_size=100, secret_filter=secret_filter):
+        secret_name = self._get_secret_path(project_id, secret_id, location)
+        for secret in self.list_secrets(
+            project_id=project_id, page_size=100, secret_filter=secret_filter, location=location
+        ):
             if secret.name.split("/")[-1] == secret_id:
                 self.log.info("Secret %s exists.", secret_name)
                 return True
@@ -212,6 +298,7 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
+        location: str | None = None,
     ) -> AccessSecretVersionResponse:
         """
         Access a secret version.
@@ -227,11 +314,13 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         :param retry: Optional. Designation of what errors, if any, should be retried.
         :param timeout: Optional. The timeout for this request.
         :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        :param location: Optional. Location where secret is stored regionally.  If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used
+            For more details : https://cloud.google.com/secret-manager/docs/locations
         :return: Access secret version response object.
         """
         response = self.client.access_secret_version(
             request={
-                "name": self.client.secret_version_path(project_id, secret_id, secret_version),
+                "name": self._get_secret_version_path(project_id, secret_id, secret_version, location),
             },
             retry=retry,
             timeout=timeout,
@@ -248,6 +337,7 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
+        location: str | None = None,
     ) -> None:
         """
         Delete a secret.
@@ -262,9 +352,11 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         :param retry: Optional. Designation of what errors, if any, should be retried.
         :param timeout: Optional. The timeout for this request.
         :param metadata: Optional. Strings which should be sent along with the request as metadata.
+        :param location: Optional. Location where secret is stored regionally. If set to ``None`` or missing, the location provided for GoogleCloudSecretHook instantiation is used.
+            For more details : https://cloud.google.com/secret-manager/docs/locations
         :return: Access secret version response object.
         """
-        name = self.client.secret_path(project_id, secret_id)
+        name = self._get_secret_path(project_id, secret_id, location)
         self.client.delete_secret(
             request={"name": name},
             retry=retry,

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

@@ -19,17 +19,19 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Callable, NamedTuple
+from collections import OrderedDict
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, NamedTuple
 
 from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
 from google.cloud.spanner_v1.client import Client
 from sqlalchemy import create_engine
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field
+from airflow.providers.openlineage.sqlparser import DatabaseInfo
 
 if TYPE_CHECKING:
     from google.cloud.spanner_v1.database import Database
@@ -37,6 +39,8 @@ if TYPE_CHECKING:
     from google.cloud.spanner_v1.transaction import Transaction
     from google.longrunning.operations_grpc_pb2 import Operation
 
+    from airflow.models.connection import Connection
+
 
 class SpannerConnectionParams(NamedTuple):
     """Information about Google Spanner connection parameters."""
@@ -388,7 +392,7 @@ class SpannerHook(GoogleBaseHook, DbApiHook):
         database_id: str,
         queries: list[str],
         project_id: str,
-    ) -> None:
+    ) -> list[int]:
         """
         Execute an arbitrary DML query (INSERT, UPDATE, DELETE).
 
@@ -398,12 +402,73 @@ class SpannerHook(GoogleBaseHook, DbApiHook):
         :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner
             database. If set to None or missing, the default project_id from the Google Cloud connection
             is used.
+        :return: list of numbers of affected rows by DML query
         """
-        self._get_client(project_id=project_id).instance(instance_id=instance_id).database(
-            database_id=database_id
-        ).run_in_transaction(lambda transaction: self._execute_sql_in_transaction(transaction, queries))
+        db = (
+            self._get_client(project_id=project_id)
+            .instance(instance_id=instance_id)
+            .database(database_id=database_id)
+        )
+
+        def _tx_runner(tx: Transaction) -> dict[str, int]:
+            return self._execute_sql_in_transaction(tx, queries)
+
+        result = db.run_in_transaction(_tx_runner)
+
+        result_rows_count_per_query = []
+        for i, (sql, rc) in enumerate(result.items(), start=1):
+            if not sql.startswith("SELECT"):
+                preview = sql if len(sql) <= 300 else sql[:300] + "…"
+                self.log.info("[DML %d/%d] affected rows=%d | %s", i, len(result), rc, preview)
+                result_rows_count_per_query.append(rc)
+        return result_rows_count_per_query
 
     @staticmethod
-    def _execute_sql_in_transaction(transaction: Transaction, queries: list[str]):
+    def _execute_sql_in_transaction(transaction: Transaction, queries: list[str]) -> dict[str, int]:
+        counts: OrderedDict[str, int] = OrderedDict()
         for sql in queries:
-            transaction.execute_update(sql)
+            rc = transaction.execute_update(sql)
+            counts[sql] = rc
+        return counts
+
+    def _get_openlineage_authority_part(self, connection: Connection) -> str | None:
+        """Build Spanner-specific authority part for OpenLineage. Returns {project}/{instance}."""
+        extras = connection.extra_dejson
+        project_id = extras.get("project_id")
+        instance_id = extras.get("instance_id")
+
+        if not project_id or not instance_id:
+            return None
+
+        return f"{project_id}/{instance_id}"
+
+    def get_openlineage_database_dialect(self, connection: Connection) -> str:
+        """Return database dialect for OpenLineage."""
+        return "spanner"
+
+    def get_openlineage_database_info(self, connection: Connection) -> DatabaseInfo:
+        """Return Spanner specific information for OpenLineage."""
+        extras = connection.extra_dejson
+        database_id = extras.get("database_id")
+
+        return DatabaseInfo(
+            scheme=self.get_openlineage_database_dialect(connection),
+            authority=self._get_openlineage_authority_part(connection),
+            database=database_id,
+            information_schema_columns=[
+                "table_schema",
+                "table_name",
+                "column_name",
+                "ordinal_position",
+                "spanner_type",
+            ],
+        )
+
+    def get_openlineage_default_schema(self) -> str | None:
+        """
+        Spanner expose 'public' or '' schema depending on dialect(Postgres vs GoogleSQL).
+
+        SQLAlchemy dialect for Spanner does not expose default schema, so we return None
+        to follow the same approach.
+        """
+        return None