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

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

@@ -26,19 +26,26 @@ import os
 import shutil
 import time
 import warnings
-from collections.abc import Generator, Sequence
+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, Callable, TypeVar, cast, overload
+from typing import IO, TYPE_CHECKING, Any, ParamSpec, TypeVar, cast, overload
 from urllib.parse import urlsplit
 
+# Make mypy happy by importing as aliases
+import google.cloud.storage as storage
 from gcloud.aio.storage import Storage
-from requests import Session
+from google.api_core.exceptions import GoogleAPICallError, NotFound
+from google.cloud.exceptions import GoogleCloudError
+from google.cloud.storage.retry import DEFAULT_RETRY
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
+from airflow.providers.common.compat.sdk import timezone
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (
@@ -46,23 +53,13 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseAsyncHook,
     GoogleBaseHook,
 )
-from airflow.typing_compat import ParamSpec
-from airflow.utils import timezone
 from airflow.version import version
-from google.api_core.exceptions import GoogleAPICallError, NotFound
-
-# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
-from google.cloud import storage  # type: ignore[attr-defined]
-from google.cloud.exceptions import GoogleCloudError
-from google.cloud.storage.retry import DEFAULT_RETRY
 
 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
+    from requests import Session
 
 
 RT = TypeVar("RT")
@@ -136,16 +133,16 @@ def _fallback_object_url_to_object_name_and_bucket_name(
 
             return func(self, *args, **kwargs)
 
-        return cast(Callable[FParams, RT], _inner_wrapper)
+        return cast("Callable[FParams, RT]", _inner_wrapper)
 
-    return cast(Callable[[T], T], _wrapper)
+    return cast("Callable[[T], T]", _wrapper)
 
 
 # A fake bucket to use in functions decorated by _fallback_object_url_to_object_name_and_bucket_name.
 # This allows the 'bucket' argument to be of type str instead of str | None,
 # making it easier to type hint the function body without dealing with the None
 # case that can never happen at runtime.
-PROVIDE_BUCKET: str = cast(str, None)
+PROVIDE_BUCKET: str = cast("str", None)
 
 
 class GCSHook(GoogleBaseHook):
@@ -359,11 +356,10 @@ class GCSHook(GoogleBaseHook):
                     )
                     self.log.info("File downloaded to %s", filename)
                     return filename
-                else:
-                    get_hook_lineage_collector().add_input_asset(
-                        context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name}
-                    )
-                    return blob.download_as_bytes()
+                get_hook_lineage_collector().add_input_asset(
+                    context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name}
+                )
+                return blob.download_as_bytes()
 
             except GoogleCloudError:
                 if attempt == num_max_attempts - 1:
@@ -375,8 +371,7 @@ class GCSHook(GoogleBaseHook):
                         num_max_attempts,
                     )
                     raise
-        else:
-            raise NotImplementedError  # should not reach this, but makes mypy happy
+        raise NotImplementedError  # should not reach this, but makes mypy happy
 
     def download_as_byte_array(
         self,
@@ -551,13 +546,13 @@ class GCSHook(GoogleBaseHook):
         if cache_control:
             blob.cache_control = cache_control
 
-        if filename and data:
+        if filename is not None and data is not None:
             raise ValueError(
                 "'filename' and 'data' parameter provided. Please "
                 "specify a single parameter, either 'filename' for "
                 "local file uploads or 'data' for file content uploads."
             )
-        elif filename:
+        if filename is not None:
             if not mime_type:
                 mime_type = "application/octet-stream"
             if gzip:
@@ -577,7 +572,7 @@ class GCSHook(GoogleBaseHook):
             if gzip:
                 os.remove(filename)
             self.log.info("File %s uploaded to %s in %s bucket", filename, object_name, bucket_name)
-        elif data:
+        elif data is not None:
             if not mime_type:
                 mime_type = "text/plain"
             if gzip:
@@ -598,7 +593,13 @@ class GCSHook(GoogleBaseHook):
             context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name}
         )
 
-    def exists(self, bucket_name: str, object_name: str, retry: Retry = DEFAULT_RETRY) -> bool:
+    def exists(
+        self,
+        bucket_name: str,
+        object_name: str,
+        retry: Retry = DEFAULT_RETRY,
+        user_project: str | None = None,
+    ) -> bool:
         """
         Check for the existence of a file in Google Cloud Storage.
 
@@ -606,9 +607,11 @@ class GCSHook(GoogleBaseHook):
         :param object_name: The name of the blob_name to check in the Google cloud
             storage bucket.
         :param retry: (Optional) How to retry the RPC
+        :param user_project: The identifier of the Google Cloud project to bill for the request.
+            Required for Requester Pays buckets.
         """
         client = self.get_conn()
-        bucket = client.bucket(bucket_name)
+        bucket = client.bucket(bucket_name, user_project=user_project)
         blob = bucket.blob(blob_name=object_name)
         return blob.exists(retry=retry)
 
@@ -625,7 +628,7 @@ class GCSHook(GoogleBaseHook):
 
     def is_updated_after(self, bucket_name: str, object_name: str, ts: datetime) -> bool:
         """
-        Check if an blob_name is updated in Google Cloud Storage.
+        Check if a blob_name is updated in Google Cloud Storage.
 
         :param bucket_name: The Google Cloud Storage bucket where the object is.
         :param object_name: The name of the object to check in the Google cloud
@@ -645,7 +648,7 @@ class GCSHook(GoogleBaseHook):
         self, bucket_name: str, object_name: str, min_ts: datetime, max_ts: datetime
     ) -> bool:
         """
-        Check if an blob_name is updated in Google Cloud Storage.
+        Check if a blob_name is updated in Google Cloud Storage.
 
         :param bucket_name: The Google Cloud Storage bucket where the object is.
         :param object_name: The name of the object to check in the Google cloud
@@ -666,7 +669,7 @@ class GCSHook(GoogleBaseHook):
 
     def is_updated_before(self, bucket_name: str, object_name: str, ts: datetime) -> bool:
         """
-        Check if an blob_name is updated before given time in Google Cloud Storage.
+        Check if a blob_name is updated before given time in Google Cloud Storage.
 
         :param bucket_name: The Google Cloud Storage bucket where the object is.
         :param object_name: The name of the object to check in the Google cloud
@@ -719,6 +722,14 @@ class GCSHook(GoogleBaseHook):
 
         self.log.info("Blob %s deleted.", object_name)
 
+    def get_bucket(self, bucket_name: str) -> storage.Bucket:
+        """
+        Get a bucket object from the Google Cloud Storage.
+
+        :param bucket_name: name of the bucket
+        """
+        return self.get_conn().bucket(bucket_name)
+
     def delete_bucket(self, bucket_name: str, force: bool = False, user_project: str | None = None) -> None:
         """
         Delete a bucket object from the Google Cloud Storage.
@@ -1237,6 +1248,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,
@@ -1487,5 +1598,5 @@ class GCSAsyncHook(GoogleBaseAsyncHook):
         token = await self.get_token(session=session)
         return Storage(
             token=token,
-            session=cast(Session, session),
+            session=cast("Session", session),
         )

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

@@ -23,11 +23,12 @@ import base64
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.kms_v1 import KeyManagementServiceClient
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 

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

@@ -25,6 +25,14 @@ import time
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+from google.api_core.exceptions import NotFound
+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
+from google.cloud.container_v1 import ClusterManagerAsyncClient, ClusterManagerClient
+from google.cloud.container_v1.types import Cluster, Operation
 from kubernetes import client
 from kubernetes_asyncio import client as async_client
 from kubernetes_asyncio.config.kube_config import FileOrData
@@ -39,14 +47,6 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseAsyncHook,
     GoogleBaseHook,
 )
-from google.api_core.exceptions import NotFound
-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.container_v1 import ClusterManagerAsyncClient, ClusterManagerClient
-from google.cloud.container_v1.types import Cluster, Operation
 
 if TYPE_CHECKING:
     import google.auth.credentials
@@ -64,11 +64,13 @@ class GKEClusterConnection:
         ssl_ca_cert: str,
         credentials: google.auth.credentials.Credentials,
         enable_tcp_keepalive: bool = False,
+        use_dns_endpoint: bool = False,
     ):
         self._cluster_url = cluster_url
         self._ssl_ca_cert = ssl_ca_cert
         self._credentials = credentials
         self.enable_tcp_keepalive = enable_tcp_keepalive
+        self.use_dns_endpoint = use_dns_endpoint
 
     def get_conn(self) -> client.ApiClient:
         configuration = self._get_config()
@@ -86,12 +88,13 @@ class GKEClusterConnection:
             api_key_prefix={"authorization": "Bearer"},
             api_key={"authorization": self._get_token(self._credentials)},
         )
-        configuration.ssl_ca_cert = FileOrData(
-            {
-                "certificate-authority-data": self._ssl_ca_cert,
-            },
-            file_key_name="certificate-authority",
-        ).as_file()
+        if not self.use_dns_endpoint:
+            configuration.ssl_ca_cert = FileOrData(
+                {
+                    "certificate-authority-data": self._ssl_ca_cert,
+                },
+                file_key_name="certificate-authority",
+            ).as_file()
         return configuration
 
     @staticmethod
@@ -349,8 +352,7 @@ class GKEHook(GoogleBaseHook):
             or node_pools_autoscaled
         ):
             return True
-        else:
-            return False
+        return False
 
 
 class GKEAsyncHook(GoogleBaseAsyncHook):
@@ -417,6 +419,7 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
         cluster_url: str,
         ssl_ca_cert: str,
         enable_tcp_keepalive: bool = False,
+        use_dns_endpoint: bool = False,
         *args,
         **kwargs,
     ):
@@ -424,6 +427,7 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
         self._cluster_url = cluster_url
         self._ssl_ca_cert = ssl_ca_cert
         self.enable_tcp_keepalive = enable_tcp_keepalive
+        self.use_dns_endpoint = use_dns_endpoint
 
     def get_conn(self) -> client.ApiClient:
         return GKEClusterConnection(
@@ -431,6 +435,7 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
             ssl_ca_cert=self._ssl_ca_cert,
             credentials=self.get_credentials(),
             enable_tcp_keepalive=self.enable_tcp_keepalive,
+            use_dns_endpoint=self.use_dns_endpoint,
         ).get_conn()
 
     def apply_from_yaml_file(
@@ -493,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

@@ -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.hooks.base import BaseHook
+from airflow.providers.common.compat.sdk import BaseHook
 from airflow.version import version
 
 if TYPE_CHECKING:
@@ -39,6 +39,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,