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

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

@@ -33,7 +33,7 @@ from google.cloud.bigquery_datatransfer_v1 import (
 )
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook, get_object_id
 from airflow.providers.google.cloud.links.bigquery_dts import BigQueryDataTransferConfigLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
@@ -43,7 +43,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 def _get_transfer_config_details(config_transfer_name: str):
@@ -134,7 +134,6 @@ class BigQueryCreateDataTransferOperator(GoogleCloudBaseOperator):
         transfer_config = _get_transfer_config_details(response.name)
         BigQueryDataTransferConfigLink.persist(
             context=context,
-            task_instance=self,
             region=transfer_config["region"],
             config_id=transfer_config["config_id"],
             project_id=transfer_config["project_id"],
@@ -142,7 +141,7 @@ class BigQueryCreateDataTransferOperator(GoogleCloudBaseOperator):
 
         result = TransferConfig.to_dict(response)
         self.log.info("Created DTS transfer config %s", get_object_id(result))
-        self.xcom_push(context, key="transfer_config_id", value=get_object_id(result))
+        context["ti"].xcom_push(key="transfer_config_id", value=get_object_id(result))
         # don't push AWS secret in XCOM
         result.get("params", {}).pop("secret_access_key", None)
         result.get("params", {}).pop("access_key_id", None)
@@ -329,7 +328,6 @@ class BigQueryDataTransferServiceStartTransferRunsOperator(GoogleCloudBaseOperat
         transfer_config = _get_transfer_config_details(response.runs[0].name)
         BigQueryDataTransferConfigLink.persist(
             context=context,
-            task_instance=self,
             region=transfer_config["region"],
             config_id=transfer_config["config_id"],
             project_id=transfer_config["project_id"],
@@ -337,7 +335,7 @@ class BigQueryDataTransferServiceStartTransferRunsOperator(GoogleCloudBaseOperat
 
         result = StartManualTransferRunsResponse.to_dict(response)
         run_id = get_object_id(result["runs"][0])
-        self.xcom_push(context, key="run_id", value=run_id)
+        context["ti"].xcom_push(key="run_id", value=run_id)
 
         if not self.deferrable:
             # Save as attribute for further use by OpenLineage

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

@@ -20,11 +20,11 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import google.api_core.exceptions
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.bigtable import BigtableHook
 from airflow.providers.google.cloud.links.bigtable import (
     BigtableClusterLink,
@@ -40,7 +40,7 @@ if TYPE_CHECKING:
     from google.cloud.bigtable import enums
     from google.cloud.bigtable.column_family import GarbageCollectionRule
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class BigtableValidationMixin:
@@ -142,6 +142,13 @@ class BigtableCreateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -155,7 +162,7 @@ class BigtableCreateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
                 "The instance '%s' already exists in this project. Consider it as created",
                 self.instance_id,
             )
-            BigtableInstanceLink.persist(context=context, task_instance=self)
+            BigtableInstanceLink.persist(context=context)
             return
         try:
             hook.create_instance(
@@ -171,7 +178,7 @@ class BigtableCreateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
                 cluster_storage_type=self.cluster_storage_type,
                 timeout=self.timeout,
             )
-            BigtableInstanceLink.persist(context=context, task_instance=self)
+            BigtableInstanceLink.persist(context=context)
         except google.api_core.exceptions.GoogleAPICallError as e:
             self.log.error("An error occurred. Exiting.")
             raise e
@@ -240,6 +247,13 @@ class BigtableUpdateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -258,7 +272,7 @@ class BigtableUpdateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
                 instance_labels=self.instance_labels,
                 timeout=self.timeout,
             )
-            BigtableInstanceLink.persist(context=context, task_instance=self)
+            BigtableInstanceLink.persist(context=context)
         except google.api_core.exceptions.GoogleAPICallError as e:
             self.log.error("An error occurred. Exiting.")
             raise e
@@ -414,6 +428,13 @@ class BigtableCreateTableOperator(GoogleCloudBaseOperator, BigtableValidationMix
                 return False
         return True
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -431,7 +452,7 @@ class BigtableCreateTableOperator(GoogleCloudBaseOperator, BigtableValidationMix
                 initial_split_keys=self.initial_split_keys,
                 column_families=self.column_families,
             )
-            BigtableTablesLink.persist(context=context, task_instance=self)
+            BigtableTablesLink.persist(context=context)
         except google.api_core.exceptions.AlreadyExists:
             if not self._compare_column_families(hook, instance):
                 raise AirflowException(
@@ -575,6 +596,14 @@ class BigtableUpdateClusterOperator(GoogleCloudBaseOperator, BigtableValidationM
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -586,7 +615,7 @@ class BigtableUpdateClusterOperator(GoogleCloudBaseOperator, BigtableValidationM
 
         try:
             hook.update_cluster(instance=instance, cluster_id=self.cluster_id, nodes=self.nodes)
-            BigtableClusterLink.persist(context=context, task_instance=self)
+            BigtableClusterLink.persist(context=context)
         except google.api_core.exceptions.NotFound:
             raise AirflowException(
                 f"Dependency: cluster '{self.cluster_id}' does not exist for instance '{self.instance_id}'."

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

@@ -19,9 +19,11 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 from google.api_core.gapic_v1.method import DEFAULT
 
-from airflow.models import BaseOperator
+from airflow.providers.google.version_compat import BaseOperator
 
 
 class GoogleCloudBaseOperator(BaseOperator):
@@ -36,3 +38,21 @@ class GoogleCloudBaseOperator(BaseOperator):
         """
         memo[id(DEFAULT)] = DEFAULT
         return super().__deepcopy__(memo)
+
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        """
+        Override this method to include parameters for link formatting in extra links.
+
+        For example; most of the links on the Google provider require `project_id` and `location` in the Link.
+        To be not repeat; you can override this function and return something like the following:
+
+        .. code-block:: python
+
+            {
+                "project_id": self.project_id,
+                "location": self.location,
+            }
+
+        """
+        return {}

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

@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING
 from google.cloud.batch_v1 import Job, Task
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_batch import CloudBatchHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.cloud_batch import CloudBatchJobFinishedTrigger
@@ -31,7 +31,7 @@ from airflow.providers.google.cloud.triggers.cloud_batch import CloudBatchJobFin
 if TYPE_CHECKING:
     from google.api_core import operation
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CloudBatchSubmitJobOperator(GoogleCloudBaseOperator):
@@ -244,7 +244,7 @@ class CloudBatchListTasksOperator(GoogleCloudBaseOperator):
     :param job_name: Required. The name of the job for which to list tasks.
     :param gcp_conn_id: The connection ID used to connect to Google Cloud.
     :param filter: The filter based on which to list the jobs. If left empty, all the jobs are listed.
-    :param group_name: The name of the group that owns the task. By default it's `group0`.
+    :param group_name: The name of the group that owns the task. By default, it's `group0`.
     :param limit: The number of tasks to list.
         If left empty, all the tasks matching the filter will be returned.
     :param impersonation_chain: Optional service account to impersonate using short-term

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

@@ -30,7 +30,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_build import CloudBuildHook
 from airflow.providers.google.cloud.links.cloud_build import (
     CloudBuildLink,
@@ -48,7 +48,7 @@ from airflow.utils.helpers import exactly_one
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 REGEX_REPO_PATH = re.compile(r"^/(?P<project_id>[^/]+)/(?P<repo_name>[^/]+)[\+/]*(?P<branch_name>[^:]+)?")
@@ -108,6 +108,12 @@ class CloudBuildCancelBuildOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.cancel_build(
@@ -119,14 +125,12 @@ class CloudBuildCancelBuildOperator(GoogleCloudBaseOperator):
             location=self.location,
         )
 
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 build_id=result.id,
             )
         return Build.to_dict(result)
@@ -210,6 +214,12 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
             if self.build_raw.endswith(".json"):
                 self.build = json.loads(file.read())
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -225,7 +235,7 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             location=self.location,
         )
-        self.xcom_push(context, key="id", value=self.id_)
+        context["task_instance"].xcom_push(key="id", value=self.id_)
         if not self.wait:
             return Build.to_dict(
                 hook.get_build(id_=self.id_, project_id=self.project_id, location=self.location)
@@ -251,9 +261,7 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
             if project_id:
                 CloudBuildLink.persist(
                     context=context,
-                    task_instance=self,
                     project_id=project_id,
-                    region=self.location,
                     build_id=cloud_build_instance_result.id,
                 )
             return Build.to_dict(cloud_build_instance_result)
@@ -269,9 +277,7 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
             if project_id:
                 CloudBuildLink.persist(
                     context=context,
-                    task_instance=self,
                     project_id=project_id,
-                    region=self.location,
                     build_id=event["id_"],
                 )
             return event["instance"]
@@ -336,6 +342,12 @@ class CloudBuildCreateBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.create_build_trigger(
@@ -346,21 +358,17 @@ class CloudBuildCreateBuildTriggerOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             location=self.location,
         )
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildTriggerDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 trigger_id=result.id,
             )
             CloudBuildTriggersListLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
             )
         return BuildTrigger.to_dict(result)
 
@@ -419,6 +427,12 @@ class CloudBuildDeleteBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         hook.delete_build_trigger(
@@ -433,9 +447,7 @@ class CloudBuildDeleteBuildTriggerOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudBuildTriggersListLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
             )
 
 
@@ -493,6 +505,12 @@ class CloudBuildGetBuildOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.get_build(
@@ -507,9 +525,7 @@ class CloudBuildGetBuildOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudBuildLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 build_id=result.id,
             )
         return Build.to_dict(result)
@@ -569,6 +585,12 @@ class CloudBuildGetBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.get_build_trigger(
@@ -583,9 +605,7 @@ class CloudBuildGetBuildTriggerOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudBuildTriggerDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 trigger_id=result.id,
             )
         return BuildTrigger.to_dict(result)
@@ -649,6 +669,12 @@ class CloudBuildListBuildTriggersOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         results = hook.list_build_triggers(
@@ -664,9 +690,7 @@ class CloudBuildListBuildTriggersOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudBuildTriggersListLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
             )
         return [BuildTrigger.to_dict(result) for result in results]
 
@@ -729,6 +753,12 @@ class CloudBuildListBuildsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         results = hook.list_builds(
@@ -743,7 +773,8 @@ class CloudBuildListBuildsOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildListLink.persist(
-                context=context, task_instance=self, project_id=project_id, region=self.location
+                context=context,
+                project_id=project_id,
             )
         return [Build.to_dict(result) for result in results]
 
@@ -805,6 +836,12 @@ class CloudBuildRetryBuildOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.retry_build(
@@ -817,14 +854,12 @@ class CloudBuildRetryBuildOperator(GoogleCloudBaseOperator):
             location=self.location,
         )
 
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 build_id=result.id,
             )
         return Build.to_dict(result)
@@ -891,6 +926,12 @@ class CloudBuildRunBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.run_build_trigger(
@@ -903,14 +944,12 @@ class CloudBuildRunBuildTriggerOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             location=self.location,
         )
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 build_id=result.id,
             )
         return Build.to_dict(result)
@@ -974,6 +1013,12 @@ class CloudBuildUpdateBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.update_build_trigger(
@@ -985,14 +1030,12 @@ class CloudBuildUpdateBuildTriggerOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             location=self.location,
         )
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildTriggerDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 trigger_id=result.id,
             )
         return BuildTrigger.to_dict(result)