apache-airflow-providers-google 18.0.0__py3-none-any.whl → 18.1.0__py3-none-any.whl

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

@@ -50,7 +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.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,
 )
@@ -60,7 +67,6 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseHook,
     get_field,
 )
-from airflow.providers.google.version_compat import BaseHook
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
@@ -1045,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
@@ -1072,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 "",
@@ -1113,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.
@@ -1120,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
 

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

@@ -38,6 +38,7 @@ from typing import TYPE_CHECKING, Any
 
 from google.cloud.storage_transfer_v1 import (
     ListTransferJobsRequest,
+    RunTransferJobRequest,
     StorageTransferServiceAsyncClient,
     TransferJob,
     TransferOperation,
@@ -55,6 +56,7 @@ from airflow.providers.google.common.hooks.base_google import (
 )
 
 if TYPE_CHECKING:
+    from google.api_core import operation_async
     from google.cloud.storage_transfer_v1.services.storage_transfer_service.pagers import (
         ListTransferJobsAsyncPager,
     )
@@ -712,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

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

@@ -51,12 +51,12 @@ from googleapiclient.discovery import Resource, build
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType, beam_options_to_args
+from airflow.providers.common.compat.sdk import timeout
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,
     GoogleBaseAsyncHook,
     GoogleBaseHook,
 )
-from airflow.providers.google.version_compat import timeout
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:

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

@@ -28,7 +28,7 @@ import requests
 from requests import HTTPError
 from tenacity import retry, stop_after_attempt, wait_exponential
 
-from airflow.providers.google.version_compat import BaseHook
+from airflow.providers.common.compat.sdk import BaseHook
 
 
 def _get_field(extras: dict, field_name: str) -> str | None:

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

@@ -912,12 +912,15 @@ class DataprocHook(GoogleBaseHook):
         state = None
         start = time.monotonic()
         while state not in (JobStatus.State.ERROR, JobStatus.State.DONE, JobStatus.State.CANCELLED):
+            self.log.debug("Waiting for job %s to complete", job_id)
             if timeout and start + timeout < time.monotonic():
                 raise AirflowException(f"Timeout: dataproc job {job_id} is not ready after {timeout}s")
+            self.log.debug("Sleeping for %s seconds", wait_time)
             time.sleep(wait_time)
             try:
                 job = self.get_job(project_id=project_id, region=region, job_id=job_id)
                 state = job.status.state
+                self.log.debug("Job %s is in state %s", job_id, state)
             except ServerError as err:
                 self.log.info("Retrying. Dataproc API returned server error when waiting for job: %s", err)
 

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

@@ -28,8 +28,10 @@ import time
 import warnings
 from collections.abc import Callable, Generator, Sequence
 from contextlib import contextmanager
+from datetime import datetime
 from functools import partial
 from io import BytesIO
+from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import IO, TYPE_CHECKING, Any, ParamSpec, TypeVar, cast, overload
 from urllib.parse import urlsplit
@@ -50,12 +52,14 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseAsyncHook,
     GoogleBaseHook,
 )
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.version import version
 
 if TYPE_CHECKING:
-    from datetime import datetime
-
     from aiohttp import ClientSession
     from google.api_core.retry import Retry
     from google.cloud.storage.blob import Blob
@@ -1249,6 +1253,106 @@ class GCSHook(GoogleBaseHook):
 
         self.log.info("Completed successfully.")
 
+    def _sync_to_local_dir_delete_stale_local_files(self, current_gcs_objects: List[Path], local_dir: Path):
+        current_gcs_keys = {key.resolve() for key in current_gcs_objects}
+
+        for item in local_dir.rglob("*"):
+            if item.is_file():
+                if item.resolve() not in current_gcs_keys:
+                    self.log.debug("Deleting stale local file: %s", item)
+                    item.unlink()
+        # Clean up empty directories
+        for root, dirs, _ in os.walk(local_dir, topdown=False):
+            for d in dirs:
+                dir_path = os.path.join(root, d)
+                if not os.listdir(dir_path):
+                    self.log.debug("Deleting stale empty directory: %s", dir_path)
+                    os.rmdir(dir_path)
+
+    def _sync_to_local_dir_if_changed(self, blob: Blob, local_target_path: Path):
+        should_download = False
+        download_msg = ""
+        if not local_target_path.exists():
+            should_download = True
+            download_msg = f"Local file {local_target_path} does not exist."
+        else:
+            local_stats = local_target_path.stat()
+            # Reload blob to get fresh metadata, including size and updated time
+            blob.reload()
+
+            if blob.size != local_stats.st_size:
+                should_download = True
+                download_msg = (
+                    f"GCS object size ({blob.size}) and local file size ({local_stats.st_size}) differ."
+                )
+
+            gcs_last_modified = blob.updated
+            if (
+                not should_download
+                and gcs_last_modified
+                and local_stats.st_mtime < gcs_last_modified.timestamp()
+            ):
+                should_download = True
+                download_msg = f"GCS object last modified ({gcs_last_modified}) is newer than local file last modified ({datetime.fromtimestamp(local_stats.st_mtime, tz=timezone.utc)})."
+
+        if should_download:
+            self.log.debug("%s Downloading %s to %s", download_msg, blob.name, local_target_path.as_posix())
+            self.download(
+                bucket_name=blob.bucket.name, object_name=blob.name, filename=str(local_target_path)
+            )
+        else:
+            self.log.debug(
+                "Local file %s is up-to-date with GCS object %s. Skipping download.",
+                local_target_path.as_posix(),
+                blob.name,
+            )
+
+    def sync_to_local_dir(
+        self,
+        bucket_name: str,
+        local_dir: str | Path,
+        prefix: str | None = None,
+        delete_stale: bool = False,
+    ) -> None:
+        """
+        Download files from a GCS bucket to a local directory.
+
+        It will download all files from the given ``prefix`` and create the corresponding
+        directory structure in the ``local_dir``.
+
+        If ``delete_stale`` is ``True``, it will delete all local files that do not exist in the GCS bucket.
+
+        :param bucket_name: The name of the GCS bucket.
+        :param local_dir: The local directory to which the files will be downloaded.
+        :param prefix: The prefix of the files to be downloaded.
+        :param delete_stale: If ``True``, deletes local files that don't exist in the bucket.
+        """
+        prefix = prefix or ""
+        local_dir_path = Path(local_dir)
+        self.log.debug("Downloading data from gs://%s/%s to %s", bucket_name, prefix, local_dir_path)
+
+        gcs_bucket = self.get_bucket(bucket_name)
+        local_gcs_objects = []
+
+        for blob in gcs_bucket.list_blobs(prefix=prefix):
+            # GCS lists "directories" as objects ending with a slash. We should skip them.
+            if blob.name.endswith("/"):
+                continue
+
+            blob_path = Path(blob.name)
+            local_target_path = local_dir_path.joinpath(blob_path.relative_to(prefix))
+            if not local_target_path.parent.exists():
+                local_target_path.parent.mkdir(parents=True, exist_ok=True)
+                self.log.debug("Created local directory: %s", local_target_path.parent)
+
+            self._sync_to_local_dir_if_changed(blob=blob, local_target_path=local_target_path)
+            local_gcs_objects.append(local_target_path)
+
+        if delete_stale:
+            self._sync_to_local_dir_delete_stale_local_files(
+                current_gcs_objects=local_gcs_objects, local_dir=local_dir_path
+            )
+
     def sync(
         self,
         source_bucket: str,

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/looker.py

@@ -29,7 +29,7 @@ from looker_sdk.sdk.api40 import methods as methods40
 from packaging.version import parse as parse_version
 
 from airflow.exceptions import AirflowException
-from airflow.providers.google.version_compat import BaseHook
+from airflow.providers.common.compat.sdk import BaseHook
 from airflow.version import version
 
 if TYPE_CHECKING:

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

@@ -31,6 +31,7 @@ from airflow.exceptions 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
@@ -38,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."""
@@ -427,3 +430,45 @@ class SpannerHook(GoogleBaseHook, DbApiHook):
             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