apache-airflow-providers-google 16.0.0a1__py3-none-any.whl → 16.1.0__py3-none-any.whl

airflow/providers/google/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "15.1.0"
+__version__ = "16.1.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/google/ads/hooks/ads.py

@@ -28,14 +28,18 @@ from google.ads.googleads.errors import GoogleAdsException
 from google.auth.exceptions import GoogleAuthError
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.providers.google.common.hooks.base_google import get_field
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
-    from google.ads.googleads.v19.services.services.customer_service import CustomerServiceClient
-    from google.ads.googleads.v19.services.services.google_ads_service import GoogleAdsServiceClient
-    from google.ads.googleads.v19.services.services.google_ads_service.pagers import SearchPager
-    from google.ads.googleads.v19.services.types.google_ads_service import GoogleAdsRow
+    from google.ads.googleads.v20.services.services.customer_service import CustomerServiceClient
+    from google.ads.googleads.v20.services.services.google_ads_service import GoogleAdsServiceClient
+    from google.ads.googleads.v20.services.services.google_ads_service.pagers import SearchPager
+    from google.ads.googleads.v20.services.types.google_ads_service import GoogleAdsRow
 
 
 class GoogleAdsHook(BaseHook):
@@ -101,6 +105,40 @@ class GoogleAdsHook(BaseHook):
     :param api_version: The Google Ads API version to use.
     """
 
+    conn_name_attr = "google_ads_conn_id"
+    default_conn_name = "google_ads_default"
+    conn_type = "google_ads"
+    hook_name = "Google Ads"
+
+    @classmethod
+    def get_connection_form_widgets(cls) -> dict[str, Any]:
+        """Return connection widgets to add to Google Ads connection form."""
+        from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget
+        from flask_babel import lazy_gettext
+        from wtforms import PasswordField, StringField
+
+        return {
+            "developer_token": StringField(lazy_gettext("Developer token"), widget=BS3TextFieldWidget()),
+            "client_id": StringField(lazy_gettext("OAuth2 Client ID"), widget=BS3TextFieldWidget()),
+            "client_secret": PasswordField(
+                lazy_gettext("OAuth2 Client Secret"), widget=BS3PasswordFieldWidget()
+            ),
+            "refresh_token": PasswordField(
+                lazy_gettext("OAuth2 Refresh Token"), widget=BS3PasswordFieldWidget()
+            ),
+        }
+
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
+        """Return custom UI field behaviour for Google Ads connection."""
+        return {
+            "hidden_fields": ["host", "login", "schema", "port"],
+            "relabeling": {},
+            "placeholders": {
+                "password": "Leave blank (optional)",
+            },
+        }
+
     def __init__(
         self,
         api_version: str | None = None,

airflow/providers/google/ads/operators/ads.py

@@ -24,9 +24,9 @@ from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.ads.hooks.ads import GoogleAdsHook
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/ads/transfers/ads_to_gcs.py

@@ -22,9 +22,9 @@ from operator import attrgetter
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.ads.hooks.ads import GoogleAdsHook
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -29,7 +29,7 @@ import uuid
 from collections.abc import Iterable, Mapping, Sequence
 from copy import deepcopy
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, NoReturn, Union, cast
+from typing import TYPE_CHECKING, Any, Literal, NoReturn, cast, overload
 
 from aiohttp import ClientSession as ClientSession
 from gcloud.aio.bigquery import Job, Table as Table_async
@@ -58,7 +58,11 @@ from pandas_gbq import read_gbq
 from pandas_gbq.gbq import GbqConnector  # noqa: F401 used in ``airflow.contrib.hooks.bigquery``
 from sqlalchemy import create_engine
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import (
+    AirflowException,
+    AirflowOptionalProviderFeatureException,
+    AirflowProviderDeprecationWarning,
+)
 from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.google.cloud.utils.bigquery import bq_cast
@@ -77,13 +81,14 @@ from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
     import pandas as pd
+    import polars as pl
     from google.api_core.page_iterator import HTTPIterator
     from google.api_core.retry import Retry
     from requests import Session
 
 log = logging.getLogger(__name__)
 
-BigQueryJob = Union[CopyJob, QueryJob, LoadJob, ExtractJob]
+BigQueryJob = CopyJob | QueryJob | LoadJob | ExtractJob
 
 
 class BigQueryHook(GoogleBaseHook, DbApiHook):
@@ -275,15 +280,57 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         """
         raise NotImplementedError()
 
-    def get_pandas_df(
+    def _get_pandas_df(
         self,
         sql: str,
         parameters: Iterable | Mapping[str, Any] | None = None,
         dialect: str | None = None,
         **kwargs,
     ) -> pd.DataFrame:
+        if dialect is None:
+            dialect = "legacy" if self.use_legacy_sql else "standard"
+
+        credentials, project_id = self.get_credentials_and_project_id()
+
+        return read_gbq(sql, project_id=project_id, dialect=dialect, credentials=credentials, **kwargs)
+
+    def _get_polars_df(self, sql, parameters=None, dialect=None, **kwargs) -> pl.DataFrame:
+        try:
+            import polars as pl
+        except ImportError:
+            raise AirflowOptionalProviderFeatureException(
+                "Polars is not installed. Please install it with `pip install polars`."
+            )
+
+        if dialect is None:
+            dialect = "legacy" if self.use_legacy_sql else "standard"
+
+        credentials, project_id = self.get_credentials_and_project_id()
+
+        pandas_df = read_gbq(sql, project_id=project_id, dialect=dialect, credentials=credentials, **kwargs)
+        return pl.from_pandas(pandas_df)
+
+    @overload
+    def get_df(
+        self, sql, parameters=None, dialect=None, *, df_type: Literal["pandas"] = "pandas", **kwargs
+    ) -> pd.DataFrame: ...
+
+    @overload
+    def get_df(
+        self, sql, parameters=None, dialect=None, *, df_type: Literal["polars"], **kwargs
+    ) -> pl.DataFrame: ...
+
+    def get_df(
+        self,
+        sql,
+        parameters=None,
+        dialect=None,
+        *,
+        df_type: Literal["pandas", "polars"] = "pandas",
+        **kwargs,
+    ) -> pd.DataFrame | pl.DataFrame:
         """
-        Get a Pandas DataFrame for the BigQuery results.
+        Get a DataFrame for the BigQuery results.
 
         The DbApiHook method must be overridden because Pandas doesn't support
         PEP 249 connections, except for SQLite.
@@ -299,12 +346,19 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
             defaults to use `self.use_legacy_sql` if not specified
         :param kwargs: (optional) passed into pandas_gbq.read_gbq method
         """
-        if dialect is None:
-            dialect = "legacy" if self.use_legacy_sql else "standard"
+        if df_type == "polars":
+            return self._get_polars_df(sql, parameters, dialect, **kwargs)
 
-        credentials, project_id = self.get_credentials_and_project_id()
+        if df_type == "pandas":
+            return self._get_pandas_df(sql, parameters, dialect, **kwargs)
 
-        return read_gbq(sql, project_id=project_id, dialect=dialect, credentials=credentials, **kwargs)
+    @deprecated(
+        planned_removal_date="November 30, 2025",
+        use_instead="airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_df",
+        category=AirflowProviderDeprecationWarning,
+    )
+    def get_pandas_df(self, sql, parameters=None, dialect=None, **kwargs):
+        return self._get_pandas_df(sql, parameters, dialect, **kwargs)
 
     @GoogleBaseHook.fallback_to_default_project_id
     def table_exists(self, dataset_id: str, table_id: str, project_id: str) -> bool:
@@ -1937,74 +1991,6 @@ def _escape(s: str) -> str:
     return e
 
 
-@deprecated(
-    planned_removal_date="April 01, 2025",
-    use_instead="airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.split_tablename",
-    category=AirflowProviderDeprecationWarning,
-)
-def split_tablename(
-    table_input: str, default_project_id: str, var_name: str | None = None
-) -> tuple[str, str, str]:
-    if "." not in table_input:
-        raise ValueError(f"Expected table name in the format of <dataset>.<table>. Got: {table_input}")
-
-    if not default_project_id:
-        raise ValueError("INTERNAL: No default project is specified")
-
-    def var_print(var_name):
-        if var_name is None:
-            return ""
-        return f"Format exception for {var_name}: "
-
-    if table_input.count(".") + table_input.count(":") > 3:
-        raise ValueError(f"{var_print(var_name)}Use either : or . to specify project got {table_input}")
-    cmpt = table_input.rsplit(":", 1)
-    project_id = None
-    rest = table_input
-    if len(cmpt) == 1:
-        project_id = None
-        rest = cmpt[0]
-    elif len(cmpt) == 2 and cmpt[0].count(":") <= 1:
-        if cmpt[-1].count(".") != 2:
-            project_id = cmpt[0]
-            rest = cmpt[1]
-    else:
-        raise ValueError(
-            f"{var_print(var_name)}Expect format of (<project:)<dataset>.<table>, got {table_input}"
-        )
-
-    cmpt = rest.split(".")
-    if len(cmpt) == 3:
-        if project_id:
-            raise ValueError(f"{var_print(var_name)}Use either : or . to specify project")
-        project_id = cmpt[0]
-        dataset_id = cmpt[1]
-        table_id = cmpt[2]
-
-    elif len(cmpt) == 2:
-        dataset_id = cmpt[0]
-        table_id = cmpt[1]
-    else:
-        raise ValueError(
-            f"{var_print(var_name)}Expect format of (<project.|<project:)<dataset>.<table>, got {table_input}"
-        )
-
-    # Exclude partition from the table name
-    table_id = table_id.split("$")[0]
-
-    if project_id is None:
-        if var_name is not None:
-            log.info(
-                'Project is not included in %s: %s; using project "%s"',
-                var_name,
-                table_input,
-                default_project_id,
-            )
-        project_id = default_project_id
-
-    return project_id, dataset_id, table_id
-
-
 def _cleanse_time_partitioning(
     destination_dataset_table: str | None, time_partitioning_in: dict | None
 ) -> dict:  # if it is a partitioned table ($ is in the table name) add partition load option

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,7 +50,6 @@ 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.cloud.hooks.secret_manager import (
     GoogleCloudSecretManagerHook,
@@ -61,6 +60,11 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseHook,
     get_field,
 )
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
@@ -847,8 +851,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

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

@@ -32,8 +32,9 @@ from google.cloud.datacatalog import (
     TagTemplateField,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 if TYPE_CHECKING:
@@ -41,6 +42,13 @@ if TYPE_CHECKING:
     from google.protobuf.field_mask_pb2 import FieldMask
 
 
+@deprecated(
+    planned_removal_date="January 30, 2026",
+    use_instead="airflow.providers.google.cloud.hooks.dataplex.DataplexHook",
+    reason="The Data Catalog will be discontinued on January 30, 2026 "
+    "in favor of Dataplex Universal Catalog.",
+    category=AirflowProviderDeprecationWarning,
+)
 class CloudDataCatalogHook(GoogleBaseHook):
     """
     Hook for Google Cloud Data Catalog Service.

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

@@ -27,9 +27,9 @@ import subprocess
 import time
 import uuid
 import warnings
-from collections.abc import Generator, Sequence
+from collections.abc import Callable, Generator, Sequence
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 from google.cloud.dataflow_v1beta3 import (
     GetJobRequest,

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

@@ -515,7 +515,7 @@ class DataplexHook(GoogleBaseHook, OperationHelper):
             fields are non-required and omitted in the request body, their values are emptied.
         :param allow_missing: Optional. If set to true and entry doesn't exist, the service will create it.
         :param delete_missing_aspects: Optional. If set to true and the aspect_keys specify aspect
-            ranges, the service deletes any existing aspects from that range that weren't provided
+            ranges, the service deletes any existing aspects from that range that were not provided
             in the request.
         :param aspect_keys: Optional. The map keys of the Aspects which the service should modify.
             It supports the following syntax:

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

@@ -28,7 +28,10 @@ import requests
 from requests import HTTPError
 from tenacity import retry, stop_after_attempt, wait_exponential
 
-from airflow.hooks.base import BaseHook
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
 
 
 def _get_field(extras: dict, field_name: str) -> str | None:

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

@@ -26,12 +26,12 @@ 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 functools import partial
 from io import BytesIO
 from tempfile import NamedTemporaryFile
-from typing import IO, TYPE_CHECKING, Any, Callable, TypeVar, cast, overload
+from typing import IO, TYPE_CHECKING, Any, TypeVar, cast, overload
 from urllib.parse import urlsplit
 
 from gcloud.aio.storage import Storage
@@ -549,13 +549,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."
             )
-        if filename:
+        if filename is not None:
             if not mime_type:
                 mime_type = "application/octet-stream"
             if gzip:
@@ -575,7 +575,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:

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

@@ -29,7 +29,11 @@ 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
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
 from airflow.version import version
 
 if TYPE_CHECKING:
@@ -39,6 +43,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/mlengine.py

@@ -23,7 +23,8 @@ 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