apache-airflow-providers-google 10.18.0rc2__py3-none-any.whl → 10.19.0__py3-none-any.whl

airflow/providers/google/__init__.py

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

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

@@ -1588,7 +1588,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         job_id: str,
         project_id: str = PROVIDE_PROJECT_ID,
         location: str | None = None,
-    ) -> CopyJob | QueryJob | LoadJob | ExtractJob | UnknownJob:
+    ) -> BigQueryJob | UnknownJob:
         """Retrieve a BigQuery job.
 
         .. seealso:: https://cloud.google.com/bigquery/docs/reference/v2/jobs
@@ -1596,8 +1596,8 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         :param job_id: The ID of the job. The ID must contain only letters (a-z, A-Z),
             numbers (0-9), underscores (_), or dashes (-). The maximum length is 1,024
             characters.
-        :param project_id: Google Cloud Project where the job is running
-        :param location: location the job is running
+        :param project_id: Google Cloud Project where the job is running.
+        :param location: Location where the job is running.
         """
         client = self.get_client(project_id=project_id, location=location)
         job = client.get_job(job_id=job_id, project=project_id, location=location)
@@ -2849,15 +2849,16 @@ class BigQueryCursor(BigQueryBaseCursor):
         return -1
 
     def execute(self, operation: str, parameters: dict | None = None) -> None:
-        """Execute a BigQuery query, and return the job ID.
+        """Execute a BigQuery query, and update the BigQueryCursor description.
 
         :param operation: The query to execute.
         :param parameters: Parameters to substitute into the query.
         """
         sql = _bind_parameters(operation, parameters) if parameters else operation
         self.flush_results()
-        self.job_id = self._run_query(sql)
-
+        job = self._run_query(sql)
+        self.job_id = job.job_id
+        self.location = self.location or job.location
         query_results = self._get_query_result()
         if "schema" in query_results:
             self.description = _format_schema_for_description(query_results["schema"])
@@ -2997,15 +2998,15 @@ class BigQueryCursor(BigQueryBaseCursor):
         self,
         sql,
         location: str | None = None,
-    ) -> str:
-        """Run job query."""
+    ) -> BigQueryJob:
+        """Run a job query and return the job instance."""
         if not self.project_id:
             raise ValueError("The project_id should be set")
 
         configuration = self._prepare_query_configuration(sql)
         job = self.hook.insert_job(configuration=configuration, project_id=self.project_id, location=location)
 
-        return job.job_id
+        return job
 
     def _prepare_query_configuration(
         self,
@@ -3357,7 +3358,7 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
 
     async def _get_job(
         self, job_id: str | None, project_id: str = PROVIDE_PROJECT_ID, location: str | None = None
-    ) -> CopyJob | QueryJob | LoadJob | ExtractJob | UnknownJob:
+    ) -> BigQueryJob | UnknownJob:
         """
         Get BigQuery job by its ID, project ID and location.
 

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

@@ -21,6 +21,9 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+from deprecated import deprecated
+
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
@@ -44,6 +47,13 @@ AUTOML_MODEL_PREDICT_LINK = (
 )
 
 
+@deprecated(
+    reason=(
+        "Class `AutoMLDatasetLink` has been deprecated and will be removed after 31.12.2024. "
+        "Please use `TranslationLegacyDatasetLink` from `airflow/providers/google/cloud/links/translate.py` instead."
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLDatasetLink(BaseGoogleLink):
     """Helper class for constructing AutoML Dataset link."""
 
@@ -65,6 +75,13 @@ class AutoMLDatasetLink(BaseGoogleLink):
         )
 
 
+@deprecated(
+    reason=(
+        "Class `AutoMLDatasetListLink` has been deprecated and will be removed after 31.12.2024. "
+        "Please use `TranslationDatasetListLink` from `airflow/providers/google/cloud/links/translate.py` instead."
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLDatasetListLink(BaseGoogleLink):
     """Helper class for constructing AutoML Dataset List link."""
 
@@ -87,6 +104,13 @@ class AutoMLDatasetListLink(BaseGoogleLink):
         )
 
 
+@deprecated(
+    reason=(
+        "Class `AutoMLModelLink` has been deprecated and will be removed after 31.12.2024. "
+        "Please use `TranslationLegacyModelLink` from `airflow/providers/google/cloud/links/translate.py` instead."
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLModelLink(BaseGoogleLink):
     """Helper class for constructing AutoML Model link."""
 
@@ -114,6 +138,13 @@ class AutoMLModelLink(BaseGoogleLink):
         )
 
 
+@deprecated(
+    reason=(
+        "Class `AutoMLModelTrainLink` has been deprecated and will be removed after 31.12.2024. "
+        "Please use `TranslationLegacyModelTrainLink` from `airflow/providers/google/cloud/links/translate.py` instead."
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLModelTrainLink(BaseGoogleLink):
     """Helper class for constructing AutoML Model Train link."""
 
@@ -138,6 +169,13 @@ class AutoMLModelTrainLink(BaseGoogleLink):
         )
 
 
+@deprecated(
+    reason=(
+        "Class `AutoMLModelPredictLink` has been deprecated and will be removed after 31.12.2024. "
+        "Please use `TranslationLegacyModelPredictLink` from `airflow/providers/google/cloud/links/translate.py` instead."
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLModelPredictLink(BaseGoogleLink):
     """Helper class for constructing AutoML Model Predict link."""
 

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

@@ -0,0 +1,180 @@
+# 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 Google Translate links."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.google.cloud.links.base import BASE_LINK, BaseGoogleLink
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+TRANSLATION_BASE_LINK = BASE_LINK + "/translation"
+TRANSLATION_LEGACY_DATASET_LINK = (
+    TRANSLATION_BASE_LINK + "/locations/{location}/datasets/{dataset_id}/sentences?project={project_id}"
+)
+TRANSLATION_DATASET_LIST_LINK = TRANSLATION_BASE_LINK + "/datasets?project={project_id}"
+TRANSLATION_LEGACY_MODEL_LINK = (
+    TRANSLATION_BASE_LINK
+    + "/locations/{location}/datasets/{dataset_id}/evaluate;modelId={model_id}?project={project_id}"
+)
+TRANSLATION_LEGACY_MODEL_TRAIN_LINK = (
+    TRANSLATION_BASE_LINK + "/locations/{location}/datasets/{dataset_id}/train?project={project_id}"
+)
+TRANSLATION_LEGACY_MODEL_PREDICT_LINK = (
+    TRANSLATION_BASE_LINK
+    + "/locations/{location}/datasets/{dataset_id}/predict;modelId={model_id}?project={project_id}"
+)
+
+
+class TranslationLegacyDatasetLink(BaseGoogleLink):
+    """
+    Helper class for constructing Legacy Translation Dataset link.
+
+    Legacy Datasets are created and managed by AutoML API.
+    """
+
+    name = "Translation Legacy Dataset"
+    key = "translation_legacy_dataset"
+    format_str = TRANSLATION_LEGACY_DATASET_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        dataset_id: str,
+        project_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslationLegacyDatasetLink.key,
+            value={"location": task_instance.location, "dataset_id": dataset_id, "project_id": project_id},
+        )
+
+
+class TranslationDatasetListLink(BaseGoogleLink):
+    """Helper class for constructing Translation Dataset List 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=TranslationDatasetListLink.key,
+            value={
+                "project_id": project_id,
+            },
+        )
+
+
+class TranslationLegacyModelLink(BaseGoogleLink):
+    """
+    Helper class for constructing Translation Legacy Model link.
+
+    Legacy Models are created and managed by AutoML API.
+    """
+
+    name = "Translation Legacy Model"
+    key = "translation_legacy_model"
+    format_str = TRANSLATION_LEGACY_MODEL_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        dataset_id: str,
+        model_id: str,
+        project_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslationLegacyModelLink.key,
+            value={
+                "location": task_instance.location,
+                "dataset_id": dataset_id,
+                "model_id": model_id,
+                "project_id": project_id,
+            },
+        )
+
+
+class TranslationLegacyModelTrainLink(BaseGoogleLink):
+    """
+    Helper class for constructing Translation Legacy Model Train link.
+
+    Legacy Models are created and managed by AutoML API.
+    """
+
+    name = "Translation Legacy Model Train"
+    key = "translation_legacy_model_train"
+    format_str = TRANSLATION_LEGACY_MODEL_TRAIN_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        project_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslationLegacyModelTrainLink.key,
+            value={
+                "location": task_instance.location,
+                "dataset_id": task_instance.model["dataset_id"],
+                "project_id": project_id,
+            },
+        )
+
+
+class TranslationLegacyModelPredictLink(BaseGoogleLink):
+    """
+    Helper class for constructing Translation Legacy Model Predict link.
+
+    Legacy Models are created and managed by AutoML API.
+    """
+
+    name = "Translation Legacy Model Predict"
+    key = "translation_legacy_model_predict"
+    format_str = TRANSLATION_LEGACY_MODEL_PREDICT_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        model_id: str,
+        project_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=TranslationLegacyModelPredictLink.key,
+            value={
+                "location": task_instance.location,
+                "dataset_id": task_instance.model.dataset_id,
+                "model_id": model_id,
+                "project_id": project_id,
+            },
+        )

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

@@ -180,8 +180,7 @@ class StackdriverTaskHandler(logging.Handler):
         """
         message = self.format(record)
         ti = None
-        # todo: remove ctx_indiv_trigger is not None check when min airflow version >= 2.6
-        if ctx_indiv_trigger is not None and getattr(record, ctx_indiv_trigger.name, None):
+        if getattr(record, ctx_indiv_trigger.name, None):
             ti = getattr(record, "task_instance", None)  # trigger context
         labels = self._get_labels(ti)
         self._transport.send(record, message, resource=self.resource, labels=labels)

airflow/providers/google/cloud/openlineage/BigQueryErrorRunFacet.json

@@ -0,0 +1,30 @@
+{
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "$defs": {
+      "BigQueryErrorRunFacet": {
+        "allOf": [
+          {
+            "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunFacet"
+          },
+          {
+            "type": "object",
+            "properties": {
+              "clientError": {
+                "type": "string"
+              },
+              "parserError": {
+                "type": "string"
+              }
+            }
+          }
+        ],
+        "type": "object"
+      }
+    },
+    "type": "object",
+    "properties": {
+      "bigQuery_error": {
+        "$ref": "#/$defs/BigQueryErrorRunFacet"
+      }
+    }
+  }

airflow/providers/google/cloud/openlineage/BigQueryJobRunFacet.json

@@ -0,0 +1,37 @@
+{
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "$defs": {
+      "BigQueryJobRunFacet": {
+        "allOf": [
+          {
+            "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunFacet"
+          },
+          {
+            "type": "object",
+            "properties": {
+              "cached": {
+                "type": "boolean"
+              },
+              "billedBytes": {
+                "type": "int",
+                "example": 321
+              },
+              "properties": {
+                "type": "string"
+              }
+            },
+            "required": [
+              "cached"
+            ]
+          }
+        ],
+        "type": "object"
+      }
+    },
+    "type": "object",
+    "properties": {
+      "bigQueryJob": {
+        "$ref": "#/$defs/BigQueryJobRunFacet"
+      }
+    }
+  }

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

@@ -0,0 +1,16 @@
+# 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.