apache-airflow-providers-google 10.26.0rc1__py3-none-any.whl → 11.0.0rc1__py3-none-any.whl

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

@@ -20,7 +20,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -42,14 +43,6 @@ if TYPE_CHECKING:
 class ModelServiceHook(GoogleBaseHook):
     """Hook for Google Cloud Vertex AI Endpoint Service APIs."""
 
-    def __init__(self, **kwargs):
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
-        super().__init__(**kwargs)
-
     def get_model_service_client(self, region: str | None = None) -> ModelServiceClient:
         """Return ModelServiceClient object."""
         if region and region != "global":

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

@@ -26,7 +26,8 @@ This module contains a Google Cloud Vertex AI hook.
 from __future__ import annotations
 
 import asyncio
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -62,6 +63,7 @@ class PipelineJobHook(GoogleBaseHook):
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._pipeline_job: PipelineJob | None = None
 

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

@@ -17,7 +17,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.videointelligence_v1 import (
@@ -60,14 +61,10 @@ class CloudVideoIntelligenceHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._conn: VideoIntelligenceServiceClient | None = None
 

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

@@ -19,9 +19,10 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from copy import deepcopy
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable, Sequence
+from typing import TYPE_CHECKING, Any, Callable
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.vision_v1 import (
@@ -130,14 +131,10 @@ class CloudVisionHook(GoogleBaseHook):
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
         super().__init__(
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
+            **kwargs,
         )
         self._client = None
 

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.workflows.executions_v1beta import Execution, ExecutionsClient
@@ -41,14 +42,6 @@ class WorkflowsHook(GoogleBaseHook):
     keyword arguments rather than positional.
     """
 
-    def __init__(self, **kwargs):
-        if kwargs.get("delegate_to") is not None:
-            raise RuntimeError(
-                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
-            )
-        super().__init__(**kwargs)
-
     def get_workflows_client(self) -> WorkflowsClient:
         """Return WorkflowsClient object."""
         return WorkflowsClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)

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

@@ -121,7 +121,6 @@ class DataprocLink(BaseOperatorLink):
         # This link is still used into the selected operators
         # - airflow.providers.google.cloud.operators.dataproc.DataprocScaleClusterOperator
         # - airflow.providers.google.cloud.operators.dataproc.DataprocJobBaseOperator
-        # - airflow.providers.google.cloud.operators.dataproc.DataprocSubmitPigJobOperator
         # As soon as we remove reference to this link we might deprecate it by add warning message
         # with `stacklevel=3` below in this method.
         ...

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

@@ -43,6 +43,13 @@ TRANSLATION_LEGACY_MODEL_PREDICT_LINK = (
     + "/locations/{location}/datasets/{dataset_id}/predict;modelId={model_id}?project={project_id}"
 )
 
+TRANSLATION_TRANSLATE_TEXT_BATCH = BASE_LINK + "/storage/browser/{output_uri_prefix}?project={project_id}"
+
+TRANSLATION_NATIVE_DATASET_LINK = (
+    TRANSLATION_BASE_LINK + "/locations/{location}/datasets/{dataset_id}/sentences?project={project_id}"
+)
+TRANSLATION_NATIVE_LIST_LINK = TRANSLATION_BASE_LINK + "/datasets?project={project_id}"
+
 
 class TranslationLegacyDatasetLink(BaseGoogleLink):
     """
@@ -179,3 +186,87 @@ class TranslationLegacyModelPredictLink(BaseGoogleLink):
                 "project_id": project_id,
             },
         )
+
+
+class TranslateTextBatchLink(BaseGoogleLink):
+    """
+    Helper class for constructing Translation results for the text batch translate.
+
+    Provides link to output results.
+
+    """
+
+    name = "Text Translate Batch"
+    key = "translate_text_batch"
+    format_str = TRANSLATION_TRANSLATE_TEXT_BATCH
+
+    @staticmethod
+    def extract_output_uri_prefix(output_config):
+        return output_config["gcs_destination"]["output_uri_prefix"].rpartition("gs://")[-1]
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        project_id: str,
+        output_config: dict,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslateTextBatchLink.key,
+            value={
+                "project_id": project_id,
+                "output_uri_prefix": TranslateTextBatchLink.extract_output_uri_prefix(output_config),
+            },
+        )
+
+
+class TranslationNativeDatasetLink(BaseGoogleLink):
+    """
+    Helper class for constructing Legacy Translation Dataset link.
+
+    Legacy Datasets are created and managed by AutoML API.
+    """
+
+    name = "Translation Native Dataset"
+    key = "translation_naive_dataset"
+    format_str = TRANSLATION_NATIVE_DATASET_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        dataset_id: str,
+        project_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslationNativeDatasetLink.key,
+            value={"location": task_instance.location, "dataset_id": dataset_id, "project_id": project_id},
+        )
+
+
+class TranslationDatasetsListLink(BaseGoogleLink):
+    """
+    Helper class for constructing Translation Datasets List link.
+
+    Both legacy and native datasets are available under this link.
+    """
+
+    name = "Translation Dataset List"
+    key = "translation_dataset_list"
+    format_str = TRANSLATION_DATASET_LIST_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        project_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslationDatasetsListLink.key,
+            value={
+                "project_id": project_id,
+            },
+        )

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

@@ -20,9 +20,10 @@ from __future__ import annotations
 import logging
 import os
 import shutil
+from collections.abc import Collection
 from functools import cached_property
 from pathlib import Path
-from typing import TYPE_CHECKING, Collection
+from typing import TYPE_CHECKING
 
 # 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]

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

@@ -20,8 +20,9 @@ from __future__ import annotations
 
 import logging
 import warnings
+from collections.abc import Collection
 from functools import cached_property
-from typing import TYPE_CHECKING, Collection
+from typing import TYPE_CHECKING
 from urllib.parse import urlencode
 
 from google.cloud import logging as gcp_logging
@@ -29,13 +30,18 @@ from google.cloud.logging import Resource
 from google.cloud.logging.handlers.transports import BackgroundThreadTransport, Transport
 from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2Client
 from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesResponse
+from packaging.version import Version
 
+from airflow import __version__ as airflow_version
 from airflow.exceptions import RemovedInAirflow3Warning
 from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.utils.log.trigger_handler import ctx_indiv_trigger
 from airflow.utils.types import NOTSET, ArgNotSet
 
+AIRFLOW_VERSION = Version(airflow_version)
+AIRFLOW_V_3_0_PLUS = Version(AIRFLOW_VERSION.base_version) >= Version("3.0.0")
+
 if TYPE_CHECKING:
     from google.auth.credentials import Credentials
 
@@ -82,7 +88,7 @@ class StackdriverTaskHandler(logging.Handler):
 
     LABEL_TASK_ID = "task_id"
     LABEL_DAG_ID = "dag_id"
-    LABEL_EXECUTION_DATE = "execution_date"
+    LABEL_LOGICAL_DATE = "logical_date" if AIRFLOW_V_3_0_PLUS else "execution_date"
     LABEL_TRY_NUMBER = "try_number"
     LOG_VIEWER_BASE_URL = "https://console.cloud.google.com/logs/viewer"
     LOG_NAME = "Google Stackdriver"
@@ -338,7 +344,9 @@ class StackdriverTaskHandler(logging.Handler):
         return {
             cls.LABEL_TASK_ID: ti.task_id,
             cls.LABEL_DAG_ID: ti.dag_id,
-            cls.LABEL_EXECUTION_DATE: str(ti.execution_date.isoformat()),
+            cls.LABEL_LOGICAL_DATE: str(ti.logical_date.isoformat())
+            if AIRFLOW_V_3_0_PLUS
+            else str(ti.execution_date.isoformat()),
             cls.LABEL_TRY_NUMBER: str(ti.try_number),
         }
 

airflow/providers/google/cloud/openlineage/utils.py

@@ -27,6 +27,7 @@ if TYPE_CHECKING:
     from airflow.providers.common.compat.openlineage.facet import Dataset
 
 from airflow.providers.common.compat.openlineage.facet import (
+    BaseFacet,
     ColumnLineageDatasetFacet,
     DocumentationDatasetFacet,
     Fields,
@@ -41,50 +42,82 @@ BIGQUERY_NAMESPACE = "bigquery"
 BIGQUERY_URI = "bigquery"
 
 
-def get_facets_from_bq_table(table: Table) -> dict[Any, Any]:
+def get_facets_from_bq_table(table: Table) -> dict[str, BaseFacet]:
     """Get facets from BigQuery table object."""
-    facets = {
-        "schema": SchemaDatasetFacet(
+    facets: dict[str, BaseFacet] = {}
+    if table.schema:
+        facets["schema"] = SchemaDatasetFacet(
             fields=[
                 SchemaDatasetFacetFields(
-                    name=field.name, type=field.field_type, description=field.description
+                    name=schema_field.name, type=schema_field.field_type, description=schema_field.description
                 )
-                for field in table.schema
+                for schema_field in table.schema
             ]
-        ),
-        "documentation": DocumentationDatasetFacet(description=table.description or ""),
-    }
+        )
+    if table.description:
+        facets["documentation"] = DocumentationDatasetFacet(description=table.description)
 
     return facets
 
 
 def get_identity_column_lineage_facet(
-    field_names: list[str],
+    dest_field_names: list[str],
     input_datasets: list[Dataset],
-) -> ColumnLineageDatasetFacet:
+) -> dict[str, ColumnLineageDatasetFacet]:
     """
-    Get column lineage facet.
-
-    Simple lineage will be created, where each source column corresponds to single destination column
-    in each input dataset and there are no transformations made.
+    Get column lineage facet for identity transformations.
+
+    This function generates a simple column lineage facet, where each destination column
+    consists of source columns of the same name from all input datasets that have that column.
+    The lineage assumes there are no transformations applied, meaning the columns retain their
+    identity between the source and destination datasets.
+
+    Args:
+        dest_field_names: A list of destination column names for which lineage should be determined.
+        input_datasets: A list of input datasets with schema facets.
+
+    Returns:
+        A dictionary containing a single key, `columnLineage`, mapped to a `ColumnLineageDatasetFacet`.
+         If no column lineage can be determined, an empty dictionary is returned - see Notes below.
+
+    Notes:
+        - If any input dataset lacks a schema facet, the function immediately returns an empty dictionary.
+        - If any field in the source dataset's schema is not present in the destination table,
+          the function returns an empty dictionary. The destination table can contain extra fields, but all
+          source columns should be present in the destination table.
+        - If none of the destination columns can be matched to input dataset columns, an empty
+          dictionary is returned.
+        - Extra columns in the destination table that do not exist in the input datasets are ignored and
+          skipped in the lineage facet, as they cannot be traced back to a source column.
+        - The function assumes there are no transformations applied, meaning the columns retain their
+          identity between the source and destination datasets.
     """
-    if field_names and not input_datasets:
-        raise ValueError("When providing `field_names` You must provide at least one `input_dataset`.")
+    fields_sources: dict[str, list[Dataset]] = {}
+    for ds in input_datasets:
+        if not ds.facets or "schema" not in ds.facets:
+            return {}
+        for schema_field in ds.facets["schema"].fields:  # type: ignore[attr-defined]
+            if schema_field.name not in dest_field_names:
+                return {}
+            fields_sources[schema_field.name] = fields_sources.get(schema_field.name, []) + [ds]
+
+    if not fields_sources:
+        return {}
 
     column_lineage_facet = ColumnLineageDatasetFacet(
         fields={
-            field: Fields(
+            field_name: Fields(
                 inputFields=[
-                    InputField(namespace=dataset.namespace, name=dataset.name, field=field)
-                    for dataset in input_datasets
+                    InputField(namespace=dataset.namespace, name=dataset.name, field=field_name)
+                    for dataset in source_datasets
                 ],
                 transformationType="IDENTITY",
                 transformationDescription="identical",
             )
-            for field in field_names
+            for field_name, source_datasets in fields_sources.items()
         }
     )
-    return column_lineage_facet
+    return {"columnLineage": column_lineage_facet}
 
 
 @define

airflow/providers/google/cloud/operators/automl.py

@@ -21,8 +21,9 @@ from __future__ import annotations
 
 import ast
 import warnings
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Sequence, Tuple, cast
+from typing import TYPE_CHECKING, cast
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.automl_v1beta1 import (
@@ -53,7 +54,7 @@ if TYPE_CHECKING:
 
     from airflow.utils.context import Context
 
-MetaData = Sequence[Tuple[str, str]]
+MetaData = Sequence[tuple[str, str]]
 
 
 def _raise_exception_for_deprecated_operator(
@@ -85,7 +86,7 @@ class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
     :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLTabularTrainingJobOperator`,
     :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLVideoTrainingJobOperator`,
     :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLImageTrainingJobOperator`,
-    :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLTextTrainingJobOperator`,
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator`,
     instead.
 
     .. seealso::
@@ -155,7 +156,7 @@ class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
                     "CreateAutoMLTabularTrainingJobOperator",
                     "CreateAutoMLVideoTrainingJobOperator",
                     "CreateAutoMLImageTrainingJobOperator",
-                    "CreateAutoMLTextTrainingJobOperator",
+                    "SupervisedFineTuningTrainOperator",
                 ],
             )
         hook = CloudAutoMLHook(