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

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

@@ -28,7 +28,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import types
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job import PipelineJobHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIPipelineJobLink,
@@ -42,7 +42,7 @@ if TYPE_CHECKING:
     from google.cloud.aiplatform import PipelineJob
     from google.cloud.aiplatform.metadata import experiment_resources
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class RunPipelineJobOperator(GoogleCloudBaseOperator):
@@ -112,6 +112,10 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
         "project_id",
         "input_artifacts",
         "impersonation_chain",
+        "template_path",
+        "pipeline_root",
+        "parameter_values",
+        "service_account",
     ]
     operator_extra_links = (VertexAIPipelineJobLink(),)
 
@@ -162,6 +166,13 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
         self.deferrable = deferrable
         self.poll_interval = poll_interval
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         self.log.info("Running Pipeline job")
         pipeline_job_obj: PipelineJob = self.hook.submit_pipeline_job(
@@ -184,8 +195,8 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
         )
         pipeline_job_id = pipeline_job_obj.job_id
         self.log.info("Pipeline job was created. Job id: %s", pipeline_job_id)
-        self.xcom_push(context, key="pipeline_job_id", value=pipeline_job_id)
-        VertexAIPipelineJobLink.persist(context=context, task_instance=self, pipeline_id=pipeline_job_id)
+        context["ti"].xcom_push(key="pipeline_job_id", value=pipeline_job_id)
+        VertexAIPipelineJobLink.persist(context=context, pipeline_id=pipeline_job_id)
 
         if self.deferrable:
             pipeline_job_obj.wait_for_resource_creation()
@@ -276,6 +287,13 @@ class GetPipelineJobOperator(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.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = PipelineJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -292,9 +310,7 @@ class GetPipelineJobOperator(GoogleCloudBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            VertexAIPipelineJobLink.persist(
-                context=context, task_instance=self, pipeline_id=self.pipeline_job_id
-            )
+            VertexAIPipelineJobLink.persist(context=context, pipeline_id=self.pipeline_job_id)
             self.log.info("Pipeline job was gotten.")
             return types.PipelineJob.to_dict(result)
         except NotFound:
@@ -408,6 +424,13 @@ class ListPipelineJobOperator(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.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = PipelineJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -424,7 +447,7 @@ class ListPipelineJobOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIPipelineJobListLink.persist(context=context, task_instance=self)
+        VertexAIPipelineJobListLink.persist(context=context)
         return [types.PipelineJob.to_dict(result) for result in results]
 
 

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

@@ -0,0 +1,393 @@
+#
+# 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 Vertex AI Ray operators."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Literal
+
+from google.api_core.exceptions import NotFound
+from google.cloud.aiplatform.vertex_ray.util import resources
+
+from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.google.cloud.hooks.vertex_ai.ray import RayHook
+from airflow.providers.google.cloud.links.vertex_ai import (
+    VertexAIRayClusterLink,
+    VertexAIRayClusterListLink,
+)
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
+if TYPE_CHECKING:
+    from airflow.providers.common.compat.sdk import Context
+
+
+class RayBaseOperator(GoogleCloudBaseOperator):
+    """
+    Base class for Ray operators.
+
+    :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 region that the service belongs to.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields: Sequence[str] = (
+        "location",
+        "gcp_conn_id",
+        "project_id",
+        "impersonation_chain",
+    )
+
+    def __init__(
+        self,
+        project_id: str,
+        location: str,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.location = location
+        self.project_id = project_id
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    @cached_property
+    def hook(self) -> RayHook:
+        return RayHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+
+class CreateRayClusterOperator(RayBaseOperator):
+    """
+    Create a Ray cluster on the Vertex AI.
+
+    :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 region that the service belongs to.
+    :param head_node_type: The head node resource. Resources.node_count must be 1. If not set, default
+        value of Resources() class will be used.
+    :param python_version: Required. Python version for the ray cluster.
+    :param ray_version: Required. Ray version for the ray cluster.
+        Currently only 3 version are available: 2.9.3, 2.33, 2.42. For more information please refer to
+        https://github.com/googleapis/python-aiplatform/blob/main/setup.py#L101
+    :param network: Virtual private cloud (VPC) network. For Ray Client, VPC peering is required to
+        connect to the Ray Cluster managed in the Vertex API service. For Ray Job API, VPC network is not
+        required because Ray Cluster connection can be accessed through dashboard address.
+    :param service_account: Service account to be used for running Ray programs on the cluster.
+    :param cluster_name: This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`.
+        The first character cannot be a number or hyphen.
+    :param worker_node_types: The list of Resources of the worker nodes. The same Resources object should
+        not appear multiple times in the list.
+    :param custom_images: The NodeImages which specifies head node and worker nodes images. All the
+        workers will share the same image. If each Resource has a specific custom image, use
+        `Resources.custom_image` for head/worker_node_type(s). Note that configuring
+        `Resources.custom_image` will override `custom_images` here. Allowlist only.
+    :param enable_metrics_collection: Enable Ray metrics collection for visualization.
+    :param enable_logging: Enable exporting Ray logs to Cloud Logging.
+    :param psc_interface_config: PSC-I config.
+    :param reserved_ip_ranges: A list of names for the reserved IP ranges under the VPC network that can
+        be used for this cluster. If set, we will deploy the cluster within the provided IP ranges.
+        Otherwise, the cluster is deployed to any IP ranges under the provided VPC network.
+        Example: ["vertex-ai-ip-range"].
+    :param labels: The labels with user-defined metadata to organize Ray cluster.
+        Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain
+        lowercase letters, numeric characters, underscores and dashes. International characters are allowed.
+        See https://goo.gl/xmQnxf for more information and examples of labels.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields: Sequence[str] = tuple(
+        {"head_node_type", "worker_node_types"} | set(RayBaseOperator.template_fields)
+    )
+    operator_extra_links = (VertexAIRayClusterLink(),)
+
+    def __init__(
+        self,
+        python_version: str,
+        ray_version: Literal["2.9.3", "2.33", "2.42"],
+        head_node_type: resources.Resources = resources.Resources(),
+        network: str | None = None,
+        service_account: str | None = None,
+        cluster_name: str | None = None,
+        worker_node_types: list[resources.Resources] | None = None,
+        custom_images: resources.NodeImages | None = None,
+        enable_metrics_collection: bool = True,
+        enable_logging: bool = True,
+        psc_interface_config: resources.PscIConfig | None = None,
+        reserved_ip_ranges: list[str] | None = None,
+        labels: dict[str, str] | None = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.head_node_type = head_node_type
+        self.python_version = python_version
+        self.ray_version = ray_version
+        self.network = network
+        self.service_account = service_account
+        self.cluster_name = cluster_name
+        self.worker_node_types = worker_node_types
+        self.custom_images = custom_images
+        self.enable_metrics_collection = enable_metrics_collection
+        self.enable_logging = enable_logging
+        self.psc_interface_config = psc_interface_config
+        self.reserved_ip_ranges = reserved_ip_ranges
+        self.labels = labels
+
+    def execute(self, context: Context):
+        self.log.info("Creating a Ray cluster.")
+        try:
+            cluster_path = self.hook.create_ray_cluster(
+                project_id=self.project_id,
+                location=self.location,
+                head_node_type=self.head_node_type,
+                python_version=self.python_version,
+                ray_version=self.ray_version,
+                network=self.network,
+                service_account=self.service_account,
+                cluster_name=self.cluster_name,
+                worker_node_types=self.worker_node_types,
+                custom_images=self.custom_images,
+                enable_metrics_collection=self.enable_metrics_collection,
+                enable_logging=self.enable_logging,
+                psc_interface_config=self.psc_interface_config,
+                reserved_ip_ranges=self.reserved_ip_ranges,
+                labels=self.labels,
+            )
+            cluster_id = self.hook.extract_cluster_id(cluster_path)
+            context["ti"].xcom_push(
+                key="cluster_id",
+                value=cluster_id,
+            )
+            VertexAIRayClusterLink.persist(
+                context=context, location=self.location, cluster_id=cluster_id, project_id=self.project_id
+            )
+            self.log.info("Ray cluster was created.")
+        except Exception as error:
+            raise AirflowException(error)
+        return cluster_path
+
+
+class ListRayClustersOperator(RayBaseOperator):
+    """
+    List Ray clusters under the currently authenticated project.
+
+    :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 region that the service belongs to.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    operator_extra_links = (VertexAIRayClusterListLink(),)
+
+    def execute(self, context: Context):
+        VertexAIRayClusterListLink.persist(context=context, project_id=self.project_id)
+        self.log.info("Listing Clusters from location %s.", self.location)
+        try:
+            ray_cluster_list = self.hook.list_ray_clusters(
+                project_id=self.project_id,
+                location=self.location,
+            )
+            ray_cluster_dict_list = [
+                self.hook.serialize_cluster_obj(ray_cluster) for ray_cluster in ray_cluster_list
+            ]
+        except Exception as error:
+            raise AirflowException(error)
+        return ray_cluster_dict_list
+
+
+class GetRayClusterOperator(RayBaseOperator):
+    """
+    Get Ray cluster.
+
+    :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 region that the service belongs to.
+    :param cluster_id: Cluster resource ID.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields: Sequence[str] = tuple({"cluster_id"} | set(RayBaseOperator.template_fields))
+    operator_extra_links = (VertexAIRayClusterLink(),)
+
+    def __init__(
+        self,
+        cluster_id: str,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.cluster_id = cluster_id
+
+    def execute(self, context: Context):
+        VertexAIRayClusterLink.persist(
+            context=context,
+            location=self.location,
+            cluster_id=self.cluster_id,
+            project_id=self.project_id,
+        )
+        self.log.info("Getting Cluster: %s", self.cluster_id)
+        try:
+            ray_cluster = self.hook.get_ray_cluster(
+                project_id=self.project_id,
+                location=self.location,
+                cluster_id=self.cluster_id,
+            )
+            self.log.info("Cluster data has been retrieved.")
+            ray_cluster_dict = self.hook.serialize_cluster_obj(ray_cluster)
+            return ray_cluster_dict
+        except NotFound as not_found_err:
+            self.log.info("The Cluster %s does not exist.", self.cluster_id)
+            raise AirflowException(not_found_err)
+
+
+class UpdateRayClusterOperator(RayBaseOperator):
+    """
+    Update Ray cluster (currently support resizing node counts for worker nodes).
+
+    :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 region that the service belongs to.
+    :param cluster_id: Cluster resource ID.
+    :param worker_node_types: The list of Resources of the resized worker nodes. The same Resources
+        object should not appear multiple times in the list.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields: Sequence[str] = tuple(
+        {"cluster_id", "worker_node_types"} | set(RayBaseOperator.template_fields)
+    )
+    operator_extra_links = (VertexAIRayClusterLink(),)
+
+    def __init__(
+        self,
+        cluster_id: str,
+        worker_node_types: list[resources.Resources],
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.cluster_id = cluster_id
+        self.worker_node_types = worker_node_types
+
+    def execute(self, context: Context):
+        VertexAIRayClusterLink.persist(
+            context=context,
+            location=self.location,
+            cluster_id=self.cluster_id,
+            project_id=self.project_id,
+        )
+        self.log.info("Updating a Ray cluster.")
+        try:
+            cluster_path = self.hook.update_ray_cluster(
+                project_id=self.project_id,
+                location=self.location,
+                cluster_id=self.cluster_id,
+                worker_node_types=self.worker_node_types,
+            )
+            self.log.info("Ray cluster %s was updated.", self.cluster_id)
+            return cluster_path
+        except NotFound as not_found_err:
+            self.log.info("The Cluster %s does not exist.", self.cluster_id)
+            raise AirflowException(not_found_err)
+        except Exception as error:
+            raise AirflowException(error)
+
+
+class DeleteRayClusterOperator(RayBaseOperator):
+    """
+    Delete Ray cluster.
+
+    :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 region that the service belongs to.
+    :param cluster_id: Cluster resource ID.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields: Sequence[str] = tuple({"cluster_id"} | set(RayBaseOperator.template_fields))
+
+    def __init__(
+        self,
+        cluster_id: str,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.cluster_id = cluster_id
+
+    def execute(self, context: Context):
+        try:
+            self.log.info("Deleting Ray cluster: %s", self.cluster_id)
+            self.hook.delete_ray_cluster(
+                project_id=self.project_id,
+                location=self.location,
+                cluster_id=self.cluster_id,
+            )
+            self.log.info("Ray cluster was deleted.")
+        except NotFound as not_found_err:
+            self.log.info("The Ray cluster ID %s does not exist.", self.cluster_id)
+            raise AirflowException(not_found_err)

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

@@ -32,7 +32,7 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CloudVideoIntelligenceDetectVideoLabelsOperator(GoogleCloudBaseOperator):

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

@@ -41,7 +41,7 @@ if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 MetaData = Sequence[tuple[str, str]]
@@ -693,7 +693,7 @@ class CloudVisionUpdateProductOperator(GoogleCloudBaseOperator):
             location=self.location,
             product_id=self.product_id,
             project_id=self.project_id,
-            update_mask=self.update_mask,  # type: ignore
+            update_mask=self.update_mask,
             retry=self.retry,
             timeout=self.timeout,
             metadata=self.metadata,

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

@@ -20,9 +20,10 @@ import datetime
 import json
 import re
 import uuid
-from collections.abc import Sequence
+from collections.abc import Collection, Sequence
 from typing import TYPE_CHECKING
 
+import pendulum
 from google.api_core.exceptions import AlreadyExists
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.workflows.executions_v1beta import Execution
@@ -36,12 +37,13 @@ from airflow.providers.google.cloud.links.workflows import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
+    from airflow.sdk import Context
 
 from airflow.utils.hashlib_wrapper import md5
 
@@ -69,7 +71,7 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
     :param metadata: Additional metadata that is provided to the method.
     """
 
-    template_fields: Sequence[str] = ("location", "workflow", "workflow_id")
+    template_fields: Collection[str] = ("location", "workflow", "workflow_id")
     template_fields_renderers = {"workflow": "json"}
     operator_extra_links = (WorkflowsWorkflowDetailsLink(),)
 
@@ -101,7 +103,7 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.force_rerun = force_rerun
 
-    def _workflow_id(self, context):
+    def _workflow_id(self, context: Context) -> str:
         if self.workflow_id and not self.force_rerun:
             # If users provide workflow id then assuring the idempotency
             # is on their side
@@ -114,8 +116,17 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
 
         # We are limited by allowed length of workflow_id so
         # we use hash of whole information
-        exec_date = context["logical_date"].isoformat()
-        base = f"airflow_{self.dag_id}_{self.task_id}_{exec_date}_{hash_base}"
+        if AIRFLOW_V_3_0_PLUS:
+            if dag_run := context.get("dag_run"):
+                run_after = pendulum.instance(dag_run.run_after)
+            else:
+                run_after = pendulum.now("UTC")
+        else:
+            if logical_date := context.get("logical_date"):
+                run_after = pendulum.instance(logical_date)
+            else:
+                run_after = pendulum.now("UTC")
+        base = f"airflow_{self.dag_id}_{self.task_id}_{run_after.isoformat()}_{hash_base}"
         workflow_id = md5(base.encode()).hexdigest()
         return re.sub(r"[:\-+.]", "_", workflow_id)
 
@@ -147,7 +158,6 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -235,7 +245,6 @@ class WorkflowsUpdateWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -368,7 +377,6 @@ class WorkflowsListWorkflowsOperator(GoogleCloudBaseOperator):
 
         WorkflowsListOfWorkflowsLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -434,7 +442,6 @@ class WorkflowsGetWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -505,11 +512,10 @@ class WorkflowsCreateExecutionOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         execution_id = execution.name.split("/")[-1]
-        self.xcom_push(context, key="execution_id", value=execution_id)
+        context["task_instance"].xcom_push(key="execution_id", value=execution_id)
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=execution_id,
@@ -582,7 +588,6 @@ class WorkflowsCancelExecutionOperator(GoogleCloudBaseOperator):
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=self.execution_id,
@@ -661,7 +666,6 @@ class WorkflowsListExecutionsOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -737,7 +741,6 @@ class WorkflowsGetExecutionOperator(GoogleCloudBaseOperator):
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=self.execution_id,

airflow/providers/google/cloud/secrets/secret_manager.py

@@ -23,7 +23,7 @@ from collections.abc import Sequence
 
 from google.auth.exceptions import DefaultCredentialsError
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud._internal_client.secret_manager_client import _SecretManagerClient
 from airflow.providers.google.cloud.utils.credentials_provider import (
     _get_target_principal_and_delegates,
@@ -160,11 +160,12 @@ class CloudSecretManagerBackend(BaseSecretsBackend, LoggingMixin):
 
         return self._get_secret(self.connections_prefix, conn_id)
 
-    def get_variable(self, key: str) -> str | None:
+    def get_variable(self, key: str, team_name: str | None = None) -> str | None:
         """
         Get Airflow Variable from Environment Variable.
 
         :param key: Variable Key
+        :param team_name: Team name associated to the task trying to access the variable (if any)
         :return: Variable Value
         """
         if self.variables_prefix is None: