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

airflow/providers/google/__init__.py

@@ -29,11 +29,11 @@ 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.9.0"
+    "2.10.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-google:{__version__}` needs Apache Airflow 2.9.0+"
+        f"The package `apache-airflow-providers-google:{__version__}` needs Apache Airflow 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/dataflow.py

@@ -185,7 +185,67 @@ class DataflowJobType:
     JOB_TYPE_STREAMING = "JOB_TYPE_STREAMING"
 
 
-class _DataflowJobsController(LoggingMixin):
+class DataflowJobTerminalStateHelper(LoggingMixin):
+    """Helper to define and validate the dataflow job terminal state."""
+
+    @staticmethod
+    def expected_terminal_state_is_allowed(expected_terminal_state):
+        job_allowed_terminal_states = DataflowJobStatus.TERMINAL_STATES | {
+            DataflowJobStatus.JOB_STATE_RUNNING
+        }
+        if expected_terminal_state not in job_allowed_terminal_states:
+            raise AirflowException(
+                f"Google Cloud Dataflow job's expected terminal state "
+                f"'{expected_terminal_state}' is invalid."
+                f" The value should be any of the following: {job_allowed_terminal_states}"
+            )
+        return True
+
+    @staticmethod
+    def expected_terminal_state_is_valid_for_job_type(expected_terminal_state, is_streaming: bool):
+        if is_streaming:
+            invalid_terminal_state = DataflowJobStatus.JOB_STATE_DONE
+            job_type = "streaming"
+        else:
+            invalid_terminal_state = DataflowJobStatus.JOB_STATE_DRAINED
+            job_type = "batch"
+
+        if expected_terminal_state == invalid_terminal_state:
+            raise AirflowException(
+                f"Google Cloud Dataflow job's expected terminal state cannot be {invalid_terminal_state} while it is a {job_type} job"
+            )
+        return True
+
+    def job_reached_terminal_state(self, job, wait_until_finished=None, custom_terminal_state=None) -> bool:
+        """
+        Check the job reached terminal state, if job failed raise exception.
+
+        :return: True if job is done.
+        :raise: Exception
+        """
+        current_state = job["currentState"]
+        is_streaming = job.get("type") == DataflowJobType.JOB_TYPE_STREAMING
+        expected_terminal_state = (
+            DataflowJobStatus.JOB_STATE_RUNNING if is_streaming else DataflowJobStatus.JOB_STATE_DONE
+        )
+        if custom_terminal_state is not None:
+            expected_terminal_state = custom_terminal_state
+        self.expected_terminal_state_is_allowed(expected_terminal_state)
+        self.expected_terminal_state_is_valid_for_job_type(expected_terminal_state, is_streaming=is_streaming)
+        if current_state == expected_terminal_state:
+            if expected_terminal_state == DataflowJobStatus.JOB_STATE_RUNNING and wait_until_finished:
+                return False
+            return True
+        if current_state in DataflowJobStatus.AWAITING_STATES:
+            return wait_until_finished is False
+        self.log.debug("Current job: %s", job)
+        raise AirflowException(
+            f"Google Cloud Dataflow job {job['name']} is in an unexpected terminal state: {current_state}, "
+            f"expected terminal state: {expected_terminal_state}"
+        )
+
+
+class _DataflowJobsController(DataflowJobTerminalStateHelper):
     """
     Interface for communication with Google Cloud Dataflow API.
 
@@ -462,7 +522,10 @@ class _DataflowJobsController(LoggingMixin):
         """Wait for result of submitted job."""
         self.log.info("Start waiting for done.")
         self._refresh_jobs()
-        while self._jobs and not all(self._check_dataflow_job_state(job) for job in self._jobs):
+        while self._jobs and not all(
+            self.job_reached_terminal_state(job, self._wait_until_finished, self._expected_terminal_state)
+            for job in self._jobs
+        ):
             self.log.info("Waiting for done. Sleep %s s", self._poll_sleep)
             time.sleep(self._poll_sleep)
             self._refresh_jobs()
@@ -1295,8 +1358,7 @@ class DataflowHook(GoogleBaseHook):
             location=location,
         )
         job = job_controller.fetch_job_by_id(job_id)
-
-        return job_controller._check_dataflow_job_state(job)
+        return job_controller.job_reached_terminal_state(job)
 
     @GoogleBaseHook.fallback_to_default_project_id
     def create_data_pipeline(
@@ -1425,7 +1487,7 @@ class DataflowHook(GoogleBaseHook):
         return f"projects/{project_id}/locations/{location}"
 
 
-class AsyncDataflowHook(GoogleBaseAsyncHook):
+class AsyncDataflowHook(GoogleBaseAsyncHook, DataflowJobTerminalStateHelper):
     """Async hook class for dataflow service."""
 
     sync_hook_class = DataflowHook

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,