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

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

@@ -20,9 +20,15 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Literal
 
-from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import GenerativeModelHook
+from google.api_core import exceptions
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import (
+    ExperimentRunHook,
+    GenerativeModelHook,
+)
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 
 if TYPE_CHECKING:
@@ -38,9 +44,8 @@ class TextEmbeddingModelGetEmbeddingsOperator(GoogleCloudBaseOperator):
     :param location: Required. The ID of the Google Cloud location that the
         service belongs to (templated).
     :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 (templated).
-    :param pretrained_model: By default uses the pre-trained model `textembedding-gecko`,
-        optimized for performing text embeddings.
+        to the Vertex AI Generative Model API, in order to elicit a specific response (templated).
+    :param pretrained_model: Required. Model, optimized for performing text embeddings.
     :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
@@ -60,7 +65,7 @@ class TextEmbeddingModelGetEmbeddingsOperator(GoogleCloudBaseOperator):
         project_id: str,
         location: str,
         prompt: str,
-        pretrained_model: str = "textembedding-gecko",
+        pretrained_model: str,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
@@ -88,7 +93,7 @@ class TextEmbeddingModelGetEmbeddingsOperator(GoogleCloudBaseOperator):
         )
 
         self.log.info("Model response: %s", response)
-        self.xcom_push(context, key="model_response", value=response)
+        context["ti"].xcom_push(key="model_response", value=response)
 
         return response
 
@@ -107,10 +112,9 @@ class GenerativeModelGenerateContentOperator(GoogleCloudBaseOperator):
     :param safety_settings: Optional. Per request settings for blocking unsafe content.
     :param tools: Optional. A list of tools available to the model during evaluation, such as a data store.
     :param system_instruction: Optional. An instruction given to the model to guide its behavior.
-    :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
-        supporting prompts with text-only input, including natural language
-        tasks, multi-turn text and code chat, and code generation. It can
-        output text and code.
+    :param pretrained_model: Required. The name of the model to use for content generation,
+        which can be a text-only or multimodal model. For example, `gemini-pro` or
+        `gemini-pro-vision`.
     :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
@@ -134,7 +138,7 @@ class GenerativeModelGenerateContentOperator(GoogleCloudBaseOperator):
         generation_config: dict | None = None,
         safety_settings: dict | None = None,
         system_instruction: str | None = None,
-        pretrained_model: str = "gemini-pro",
+        pretrained_model: str,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
@@ -168,7 +172,7 @@ class GenerativeModelGenerateContentOperator(GoogleCloudBaseOperator):
         )
 
         self.log.info("Model response: %s", response)
-        self.xcom_push(context, key="model_response", value=response)
+        context["ti"].xcom_push(key="model_response", value=response)
 
         return response
 
@@ -218,7 +222,7 @@ class SupervisedFineTuningTrainOperator(GoogleCloudBaseOperator):
         tuned_model_display_name: str | None = None,
         validation_dataset: str | None = None,
         epochs: int | None = None,
-        adapter_size: int | None = None,
+        adapter_size: Literal[1, 4, 8, 16] | None = None,
         learning_rate_multiplier: float | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
@@ -257,8 +261,8 @@ class SupervisedFineTuningTrainOperator(GoogleCloudBaseOperator):
         self.log.info("Tuned Model Name: %s", response.tuned_model_name)
         self.log.info("Tuned Model Endpoint Name: %s", response.tuned_model_endpoint_name)
 
-        self.xcom_push(context, key="tuned_model_name", value=response.tuned_model_name)
-        self.xcom_push(context, key="tuned_model_endpoint_name", value=response.tuned_model_endpoint_name)
+        context["ti"].xcom_push(key="tuned_model_name", value=response.tuned_model_name)
+        context["ti"].xcom_push(key="tuned_model_endpoint_name", value=response.tuned_model_endpoint_name)
 
         result = {
             "tuned_model_name": response.tuned_model_name,
@@ -278,10 +282,9 @@ class CountTokensOperator(GoogleCloudBaseOperator):
         service belongs to (templated).
     :param contents: Required. The multi-part content of a message that a user or a program
         gives to the generative model, in order to elicit a specific response.
-    :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
-        supporting prompts with text-only input, including natural language
-        tasks, multi-turn text and code chat, and code generation. It can
-        output text and code.
+    :param pretrained_model: Required. Model, supporting prompts with text-only input,
+        including natural language tasks, multi-turn text and code chat,
+        and code generation. It can output text and code.
     :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
@@ -301,7 +304,7 @@ class CountTokensOperator(GoogleCloudBaseOperator):
         project_id: str,
         location: str,
         contents: list,
-        pretrained_model: str = "gemini-pro",
+        pretrained_model: str,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
@@ -329,8 +332,8 @@ class CountTokensOperator(GoogleCloudBaseOperator):
         self.log.info("Total tokens: %s", response.total_tokens)
         self.log.info("Total billable characters: %s", response.total_billable_characters)
 
-        self.xcom_push(context, key="total_tokens", value=response.total_tokens)
-        self.xcom_push(context, key="total_billable_characters", value=response.total_billable_characters)
+        context["ti"].xcom_push(key="total_tokens", value=response.total_tokens)
+        context["ti"].xcom_push(key="total_billable_characters", value=response.total_billable_characters)
 
 
 class RunEvaluationOperator(GoogleCloudBaseOperator):
@@ -470,8 +473,8 @@ class CreateCachedContentOperator(GoogleCloudBaseOperator):
         project_id: str,
         location: str,
         model_name: str,
-        system_instruction: str | None = None,
-        contents: list | None = None,
+        system_instruction: Any | None = None,
+        contents: list[Any] | None = None,
         ttl_hours: float = 1,
         display_name: str | None = None,
         gcp_conn_id: str = "google_cloud_default",
@@ -582,3 +585,68 @@ class GenerateFromCachedContentOperator(GoogleCloudBaseOperator):
         self.log.info("Cached Content Response: %s", cached_content_text)
 
         return cached_content_text
+
+
+class DeleteExperimentRunOperator(GoogleCloudBaseOperator):
+    """
+    Use the Rapid Evaluation API to evaluate a model.
+
+    :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 experiment_name: Required. The name of the evaluation experiment.
+    :param experiment_run_name: Required. The specific run name or ID for this experiment.
+    :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 = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "experiment_name",
+        "experiment_run_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        experiment_name: str,
+        experiment_run_name: str,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.experiment_name = experiment_name
+        self.experiment_run_name = experiment_run_name
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> None:
+        self.hook = ExperimentRunHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            self.hook.delete_experiment_run(
+                project_id=self.project_id,
+                location=self.location,
+                experiment_name=self.experiment_name,
+                experiment_run_name=self.experiment_run_name,
+            )
+        except exceptions.NotFound:
+            raise AirflowException(f"Experiment Run with name {self.experiment_run_name} not found")
+
+        self.log.info("Deleted experiment run: %s", self.experiment_run_name)

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

@@ -257,10 +257,8 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
         hyperparameter_tuning_job_id = hyperparameter_tuning_job.name
         self.log.info("Hyperparameter Tuning job was created. Job id: %s", hyperparameter_tuning_job_id)
 
-        self.xcom_push(context, key="hyperparameter_tuning_job_id", value=hyperparameter_tuning_job_id)
-        VertexAITrainingLink.persist(
-            context=context, task_instance=self, training_id=hyperparameter_tuning_job_id
-        )
+        context["ti"].xcom_push(key="hyperparameter_tuning_job_id", value=hyperparameter_tuning_job_id)
+        VertexAITrainingLink.persist(context=context, training_id=hyperparameter_tuning_job_id)
 
         if self.deferrable:
             self.defer(
@@ -355,9 +353,7 @@ class GetHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            VertexAITrainingLink.persist(
-                context=context, task_instance=self, training_id=self.hyperparameter_tuning_job_id
-            )
+            VertexAITrainingLink.persist(context=context, training_id=self.hyperparameter_tuning_job_id)
             self.log.info("Hyperparameter tuning job was gotten.")
             return types.HyperparameterTuningJob.to_dict(result)
         except NotFound:
@@ -487,6 +483,12 @@ class ListHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = HyperparameterTuningJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -503,5 +505,5 @@ class ListHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIHyperparameterTuningJobListLink.persist(context=context, task_instance=self)
+        VertexAIHyperparameterTuningJobListLink.persist(context=context)
         return [types.HyperparameterTuningJob.to_dict(result) for result in results]

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

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -161,6 +161,13 @@ class GetModelOperator(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 = ModelServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -179,8 +186,8 @@ class GetModelOperator(GoogleCloudBaseOperator):
             )
             self.log.info("Model found. Model ID: %s", self.model_id)
 
-            self.xcom_push(context, key="model_id", value=self.model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=self.model_id)
+            context["ti"].xcom_push(key="model_id", value=self.model_id)
+            VertexAIModelLink.persist(context=context, model_id=self.model_id)
             return Model.to_dict(model)
         except NotFound:
             self.log.info("The Model ID %s does not exist.", self.model_id)
@@ -257,7 +264,12 @@ class ExportModelOperator(GoogleCloudBaseOperator):
                 metadata=self.metadata,
             )
             hook.wait_for_operation(timeout=self.timeout, operation=operation)
-            VertexAIModelExportLink.persist(context=context, task_instance=self)
+            VertexAIModelExportLink.persist(
+                context=context,
+                output_config=self.output_config,
+                model_id=self.model_id,
+                project_id=self.project_id,
+            )
             self.log.info("Model was exported.")
         except NotFound:
             self.log.info("The Model ID %s does not exist.", self.model_id)
@@ -335,6 +347,12 @@ class ListModelsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = ModelServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -352,7 +370,7 @@ class ListModelsOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIModelListLink.persist(context=context, task_instance=self)
+        VertexAIModelListLink.persist(context=context)
         return [Model.to_dict(result) for result in results]
 
 
@@ -407,6 +425,13 @@ class UploadModelOperator(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 = ModelServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -428,8 +453,8 @@ class UploadModelOperator(GoogleCloudBaseOperator):
         model_id = hook.extract_model_id(model_resp)
         self.log.info("Model was uploaded. Model ID: %s", model_id)
 
-        self.xcom_push(context, key="model_id", value=model_id)
-        VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+        context["ti"].xcom_push(key="model_id", value=model_id)
+        VertexAIModelLink.persist(context=context, model_id=model_id)
         return model_resp
 
 
@@ -553,6 +578,13 @@ class SetDefaultVersionOnModelOperator(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 = ModelServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -571,7 +603,7 @@ class SetDefaultVersionOnModelOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIModelLink.persist(context=context, task_instance=self, model_id=self.model_id)
+        VertexAIModelLink.persist(context=context, model_id=self.model_id)
         return Model.to_dict(updated_model)
 
 
@@ -627,6 +659,13 @@ class AddVersionAliasesOnModelOperator(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 = ModelServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -645,7 +684,7 @@ class AddVersionAliasesOnModelOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIModelLink.persist(context=context, task_instance=self, model_id=self.model_id)
+        VertexAIModelLink.persist(context=context, model_id=self.model_id)
         return Model.to_dict(updated_model)
 
 
@@ -701,6 +740,13 @@ class DeleteVersionAliasesOnModelOperator(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 = ModelServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -721,7 +767,7 @@ class DeleteVersionAliasesOnModelOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIModelLink.persist(context=context, task_instance=self, model_id=self.model_id)
+        VertexAIModelLink.persist(context=context, model_id=self.model_id)
         return Model.to_dict(updated_model)
 
 

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

@@ -166,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(
@@ -188,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()
@@ -280,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,
@@ -296,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:
@@ -412,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,
@@ -428,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

@@ -188,12 +188,13 @@ class CreateRayClusterOperator(RayBaseOperator):
                 labels=self.labels,
             )
             cluster_id = self.hook.extract_cluster_id(cluster_path)
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="cluster_id",
                 value=cluster_id,
             )
-            VertexAIRayClusterLink.persist(context=context, task_instance=self, cluster_id=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)
@@ -220,7 +221,7 @@ class ListRayClustersOperator(RayBaseOperator):
     operator_extra_links = (VertexAIRayClusterListLink(),)
 
     def execute(self, context: Context):
-        VertexAIRayClusterListLink.persist(context=context, task_instance=self)
+        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(
@@ -268,8 +269,9 @@ class GetRayClusterOperator(RayBaseOperator):
     def execute(self, context: Context):
         VertexAIRayClusterLink.persist(
             context=context,
-            task_instance=self,
+            location=self.location,
             cluster_id=self.cluster_id,
+            project_id=self.project_id,
         )
         self.log.info("Getting Cluster: %s", self.cluster_id)
         try:
@@ -325,8 +327,9 @@ class UpdateRayClusterOperator(RayBaseOperator):
     def execute(self, context: Context):
         VertexAIRayClusterLink.persist(
             context=context,
-            task_instance=self,
+            location=self.location,
             cluster_id=self.cluster_id,
+            project_id=self.project_id,
         )
         self.log.info("Updating a Ray cluster.")
         try:

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

@@ -147,7 +147,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 +234,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 +366,6 @@ class WorkflowsListWorkflowsOperator(GoogleCloudBaseOperator):
 
         WorkflowsListOfWorkflowsLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -434,7 +431,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 +501,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 +577,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 +655,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 +730,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/sensors/bigquery.py

@@ -31,7 +31,7 @@ from airflow.providers.google.cloud.triggers.bigquery import (
     BigQueryTableExistenceTrigger,
     BigQueryTablePartitionExistenceTrigger,
 )
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -28,7 +28,12 @@ from google.cloud.bigquery_datatransfer_v1 import TransferState
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

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

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import google.api_core.exceptions
 from google.cloud.bigtable import enums
@@ -30,7 +30,12 @@ from airflow.providers.google.cloud.hooks.bigtable import BigtableHook
 from airflow.providers.google.cloud.links.bigtable import BigtableTablesLink
 from airflow.providers.google.cloud.operators.bigtable import BigtableValidationMixin
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -89,6 +94,13 @@ class BigtableTableReplicationCompletedSensor(BaseSensorOperator, BigtableValida
         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 poke(self, context: Context) -> bool:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -119,5 +131,5 @@ class BigtableTableReplicationCompletedSensor(BaseSensorOperator, BigtableValida
             return False
 
         self.log.info("Table '%s' is replicated.", self.table_id)
-        BigtableTablesLink.persist(context=context, task_instance=self)
+        BigtableTablesLink.persist(context=context)
         return True