apache-airflow-providers-google 16.0.0a1__py3-none-any.whl → 16.0.0rc1__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.0.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

@@ -101,6 +101,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/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, NoReturn, Union, cast, overload
 
 from aiohttp import ClientSession as ClientSession
 from gcloud.aio.bigquery import Job, Table as Table_async
@@ -57,8 +57,13 @@ from googleapiclient.discovery import build
 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 typing_extensions import Literal
 
-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,6 +82,7 @@ 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
@@ -275,15 +281,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 +347,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 +1992,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/gcs.py

@@ -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

@@ -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,

airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py

@@ -185,42 +185,6 @@ class AutoMLHook(GoogleBaseHook, OperationHelper):
             model_encryption_spec_key_name=model_encryption_spec_key_name,
         )
 
-    @deprecated(
-        planned_removal_date="June 15, 2025",
-        category=AirflowProviderDeprecationWarning,
-        reason="Deprecation of AutoMLText API",
-    )
-    def get_auto_ml_text_training_job(
-        self,
-        display_name: str,
-        prediction_type: str,
-        multi_label: bool = False,
-        sentiment_max: int = 10,
-        project: str | None = None,
-        location: str | None = None,
-        labels: dict[str, str] | None = None,
-        training_encryption_spec_key_name: str | None = None,
-        model_encryption_spec_key_name: str | None = None,
-    ) -> AutoMLTextTrainingJob:
-        """
-        Return AutoMLTextTrainingJob object.
-
-        WARNING: Text creation API is deprecated since September 15, 2024
-        (https://cloud.google.com/vertex-ai/docs/tutorials/text-classification-automl/overview).
-        """
-        return AutoMLTextTrainingJob(
-            display_name=display_name,
-            prediction_type=prediction_type,
-            multi_label=multi_label,
-            sentiment_max=sentiment_max,
-            project=project,
-            location=location,
-            credentials=self.get_credentials(),
-            labels=labels,
-            training_encryption_spec_key_name=training_encryption_spec_key_name,
-            model_encryption_spec_key_name=model_encryption_spec_key_name,
-        )
-
     def get_auto_ml_video_training_job(
         self,
         display_name: str,

airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py

@@ -25,14 +25,12 @@ from typing import TYPE_CHECKING
 
 import vertexai
 from vertexai.generative_models import GenerativeModel
-from vertexai.language_models import TextEmbeddingModel, TextGenerationModel
+from vertexai.language_models import TextEmbeddingModel
 from vertexai.preview.caching import CachedContent
 from vertexai.preview.evaluation import EvalResult, EvalTask
 from vertexai.preview.generative_models import GenerativeModel as preview_generative_model
 from vertexai.preview.tuning import sft
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 if TYPE_CHECKING:
@@ -43,16 +41,6 @@ if TYPE_CHECKING:
 class GenerativeModelHook(GoogleBaseHook):
     """Hook for Google Cloud Vertex AI Generative Model APIs."""
 
-    @deprecated(
-        planned_removal_date="April 09, 2025",
-        use_instead="GenerativeModelHook.get_generative_model",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_text_generation_model(self, pretrained_model: str):
-        """Return a Model Garden Model object based on Text Generation."""
-        model = TextGenerationModel.from_pretrained(pretrained_model)
-        return model
-
     def get_text_embedding_model(self, pretrained_model: str):
         """Return a Model Garden Model object based on Text Embedding."""
         model = TextEmbeddingModel.from_pretrained(pretrained_model)
@@ -100,59 +88,6 @@ class GenerativeModelHook(GoogleBaseHook):
         cached_context_model = preview_generative_model.from_cached_content(cached_content)
         return cached_context_model
 
-    @deprecated(
-        planned_removal_date="April 09, 2025",
-        use_instead="GenerativeModelHook.generative_model_generate_content",
-        category=AirflowProviderDeprecationWarning,
-    )
-    @GoogleBaseHook.fallback_to_default_project_id
-    def text_generation_model_predict(
-        self,
-        prompt: str,
-        pretrained_model: str,
-        temperature: float,
-        max_output_tokens: int,
-        top_p: float,
-        top_k: int,
-        location: str,
-        project_id: str = PROVIDE_PROJECT_ID,
-    ) -> str:
-        """
-        Use the Vertex AI PaLM API to generate natural language text.
-
-        :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 prompt: Required. Inputs or queries that a user or a program gives
-            to the Vertex AI PaLM API, in order to elicit a specific response.
-        :param pretrained_model: A pre-trained model optimized for performing natural
-            language tasks such as classification, summarization, extraction, content
-            creation, and ideation.
-        :param temperature: Temperature controls the degree of randomness in token
-            selection.
-        :param max_output_tokens: Token limit determines the maximum amount of text
-            output.
-        :param top_p: Tokens are selected from most probable to least until the sum
-            of their probabilities equals the top_p value. Defaults to 0.8.
-        :param top_k: A top_k of 1 means the selected token is the most probable
-            among all tokens.
-        """
-        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
-
-        parameters = {
-            "temperature": temperature,
-            "max_output_tokens": max_output_tokens,
-            "top_p": top_p,
-            "top_k": top_k,
-        }
-
-        model = self.get_text_generation_model(pretrained_model)
-
-        response = model.predict(
-            prompt=prompt,
-            **parameters,
-        )
-        return response.text
-
     @GoogleBaseHook.fallback_to_default_project_id
     def text_embedding_model_get_embeddings(
         self,

airflow/providers/google/cloud/hooks/vertex_ai/ray.py

@@ -23,7 +23,7 @@ import dataclasses
 from typing import Any
 
 import vertex_ray
-from google._upb._message import ScalarMapContainer
+from google._upb._message import ScalarMapContainer  # type: ignore[attr-defined]
 from google.cloud import aiplatform
 from google.cloud.aiplatform.vertex_ray.util import resources
 from google.cloud.aiplatform_v1 import (

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

@@ -0,0 +1,59 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.google.cloud.links.base import BaseGoogleLink
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if TYPE_CHECKING:
+    from airflow.models import BaseOperator
+    from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.utils.context import Context
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk.execution_time.xcom import XCom
+else:
+    from airflow.models.xcom import XCom  # type: ignore[no-redef]
+
+
+class CloudRunJobLoggingLink(BaseGoogleLink):
+    """Helper class for constructing Cloud Run Job Logging link."""
+
+    name = "Cloud Run Job Logging"
+    key = "log_uri"
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance: BaseOperator,
+        log_uri: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=CloudRunJobLoggingLink.key,
+            value=log_uri,
+        )
+
+    def get_link(
+        self,
+        operator: BaseOperator,
+        *,
+        ti_key: TaskInstanceKey,
+    ) -> str:
+        return XCom.get_value(key=self.key, ti_key=ti_key)

airflow/providers/google/cloud/log/gcs_task_handler.py

@@ -61,11 +61,11 @@ class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
     remote_base: str
     base_log_folder: Path = attrs.field(converter=Path)
     delete_local_copy: bool
-    project_id: str
+    project_id: str | None = None
 
-    gcp_key_path: str | None
-    gcp_keyfile_dict: dict | None
-    scopes: Collection[str] | None
+    gcp_key_path: str | None = None
+    gcp_keyfile_dict: dict | None = None
+    scopes: Collection[str] | None = _DEFAULT_SCOPESS
 
     processors = ()