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

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

@@ -21,7 +21,7 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import AlreadyExists, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -150,9 +150,17 @@ class ManagedKafkaCreateClusterOperator(ManagedKafkaBaseOperator):
         self.cluster_id = cluster_id
         self.request_id = request_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         self.log.info("Creating an Apache Kafka cluster.")
-        ApacheKafkaClusterLink.persist(context=context, task_instance=self, cluster_id=self.cluster_id)
+        ApacheKafkaClusterLink.persist(context=context)
         try:
             operation = self.hook.create_cluster(
                 project_id=self.project_id,
@@ -227,8 +235,14 @@ class ManagedKafkaListClustersOperator(ManagedKafkaBaseOperator):
         self.filter = filter
         self.order_by = order_by
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterListLink.persist(context=context, task_instance=self)
+        ApacheKafkaClusterListLink.persist(context=context)
         self.log.info("Listing Clusters from location %s.", self.location)
         try:
             cluster_list_pager = self.hook.list_clusters(
@@ -242,8 +256,7 @@ class ManagedKafkaListClustersOperator(ManagedKafkaBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="cluster_page",
                 value=types.ListClustersResponse.to_dict(cluster_list_pager._response),
             )
@@ -285,12 +298,16 @@ class ManagedKafkaGetClusterOperator(ManagedKafkaBaseOperator):
         super().__init__(*args, **kwargs)
         self.cluster_id = cluster_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-        )
+        ApacheKafkaClusterLink.persist(context=context)
         self.log.info("Getting Cluster: %s", self.cluster_id)
         try:
             cluster = self.hook.get_cluster(
@@ -362,12 +379,16 @@ class ManagedKafkaUpdateClusterOperator(ManagedKafkaBaseOperator):
         self.update_mask = update_mask
         self.request_id = request_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-        )
+        ApacheKafkaClusterLink.persist(context=context)
         self.log.info("Updating an Apache Kafka cluster.")
         try:
             operation = self.hook.update_cluster(
@@ -497,14 +518,18 @@ class ManagedKafkaCreateTopicOperator(ManagedKafkaBaseOperator):
         self.topic_id = topic_id
         self.topic = topic
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "topic_id": self.topic_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         self.log.info("Creating an Apache Kafka topic.")
-        ApacheKafkaTopicLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            topic_id=self.topic_id,
-        )
+        ApacheKafkaTopicLink.persist(context=context)
         try:
             topic_obj = self.hook.create_topic(
                 project_id=self.project_id,
@@ -574,8 +599,16 @@ class ManagedKafkaListTopicsOperator(ManagedKafkaBaseOperator):
         self.page_size = page_size
         self.page_token = page_token
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterLink.persist(context=context, task_instance=self, cluster_id=self.cluster_id)
+        ApacheKafkaClusterLink.persist(context=context)
         self.log.info("Listing Topics for cluster %s.", self.cluster_id)
         try:
             topic_list_pager = self.hook.list_topics(
@@ -588,8 +621,7 @@ class ManagedKafkaListTopicsOperator(ManagedKafkaBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="topic_page",
                 value=types.ListTopicsResponse.to_dict(topic_list_pager._response),
             )
@@ -636,13 +668,17 @@ class ManagedKafkaGetTopicOperator(ManagedKafkaBaseOperator):
         self.cluster_id = cluster_id
         self.topic_id = topic_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "topic_id": self.topic_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaTopicLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            topic_id=self.topic_id,
-        )
+        ApacheKafkaTopicLink.persist(context=context)
         self.log.info("Getting Topic: %s", self.topic_id)
         try:
             topic = self.hook.get_topic(
@@ -707,13 +743,17 @@ class ManagedKafkaUpdateTopicOperator(ManagedKafkaBaseOperator):
         self.topic = topic
         self.update_mask = update_mask
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "topic_id": self.topic_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaTopicLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            topic_id=self.topic_id,
-        )
+        ApacheKafkaTopicLink.persist(context=context)
         self.log.info("Updating an Apache Kafka topic.")
         try:
             topic_obj = self.hook.update_topic(
@@ -833,8 +873,16 @@ class ManagedKafkaListConsumerGroupsOperator(ManagedKafkaBaseOperator):
         self.page_size = page_size
         self.page_token = page_token
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterLink.persist(context=context, task_instance=self, cluster_id=self.cluster_id)
+        ApacheKafkaClusterLink.persist(context=context)
         self.log.info("Listing Consumer Groups for cluster %s.", self.cluster_id)
         try:
             consumer_group_list_pager = self.hook.list_consumer_groups(
@@ -847,8 +895,7 @@ class ManagedKafkaListConsumerGroupsOperator(ManagedKafkaBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="consumer_group_page",
                 value=types.ListConsumerGroupsResponse.to_dict(consumer_group_list_pager._response),
             )
@@ -895,13 +942,17 @@ class ManagedKafkaGetConsumerGroupOperator(ManagedKafkaBaseOperator):
         self.cluster_id = cluster_id
         self.consumer_group_id = consumer_group_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "consumer_group_id": self.consumer_group_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaConsumerGroupLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            consumer_group_id=self.consumer_group_id,
-        )
+        ApacheKafkaConsumerGroupLink.persist(context=context)
         self.log.info("Getting Consumer Group: %s", self.consumer_group_id)
         try:
             consumer_group = self.hook.get_consumer_group(
@@ -971,13 +1022,17 @@ class ManagedKafkaUpdateConsumerGroupOperator(ManagedKafkaBaseOperator):
         self.consumer_group = consumer_group
         self.update_mask = update_mask
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "consumer_group_id": self.consumer_group_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaConsumerGroupLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            consumer_group_id=self.consumer_group_id,
-        )
+        ApacheKafkaConsumerGroupLink.persist(context=context)
         self.log.info("Updating an Apache Kafka consumer group.")
         try:
             consumer_group_obj = self.hook.update_consumer_group(

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

@@ -104,7 +104,6 @@ class MLEngineCreateModelOperator(GoogleCloudBaseOperator):
         if project_id:
             MLEngineModelLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 model_id=self.model["name"],
             )

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

@@ -25,8 +25,8 @@ This module contains Google PubSub operators.
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.pubsub_v1.types import (
@@ -183,7 +183,6 @@ class PubSubCreateTopicOperator(GoogleCloudBaseOperator):
         self.log.info("Created topic %s", self.topic)
         PubSubTopicLink.persist(
             context=context,
-            task_instance=self,
             topic_id=self.topic,
             project_id=self.project_id or hook.project_id,
         )
@@ -392,7 +391,6 @@ class PubSubCreateSubscriptionOperator(GoogleCloudBaseOperator):
         self.log.info("Created subscription for topic %s", self.topic)
         PubSubSubscriptionLink.persist(
             context=context,
-            task_instance=self,
             subscription_id=self.subscription or result,  # result returns subscription name
             project_id=self.project_id or hook.project_id,
         )

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

@@ -122,7 +122,6 @@ class SpannerDeployInstanceOperator(GoogleCloudBaseOperator):
         )
         SpannerInstanceLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             project_id=self.project_id or hook.project_id,
         )
@@ -290,7 +289,6 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
         )
         SpannerDatabaseLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             database_id=self.database_id,
             project_id=self.project_id or hook.project_id,
@@ -380,7 +378,6 @@ class SpannerDeployDatabaseInstanceOperator(GoogleCloudBaseOperator):
         )
         SpannerDatabaseLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             database_id=self.database_id,
             project_id=self.project_id or hook.project_id,
@@ -496,7 +493,6 @@ class SpannerUpdateDatabaseInstanceOperator(GoogleCloudBaseOperator):
             )
         SpannerDatabaseLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             database_id=self.database_id,
             project_id=self.project_id or hook.project_id,

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

@@ -117,7 +117,6 @@ class CloudSpeechToTextRecognizeSpeechOperator(GoogleCloudBaseOperator):
         if self.audio.uri:
             FileDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 # Slice from: "gs://{BUCKET_NAME}/{FILE_NAME}" to: "{BUCKET_NAME}/{FILE_NAME}"
                 uri=self.audio.uri[5:],
                 project_id=self.project_id or hook.project_id,

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

@@ -145,7 +145,6 @@ class StackdriverListAlertPoliciesOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
         return [AlertPolicy.to_dict(policy) for policy in result]
@@ -228,7 +227,6 @@ class StackdriverEnableAlertPoliciesOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -311,7 +309,6 @@ class StackdriverDisableAlertPoliciesOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -394,7 +391,6 @@ class StackdriverUpsertAlertOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -580,7 +576,6 @@ class StackdriverListNotificationChannelsOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
         return [NotificationChannel.to_dict(channel) for channel in channels]
@@ -666,7 +661,6 @@ class StackdriverEnableNotificationChannelsOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -751,7 +745,6 @@ class StackdriverDisableNotificationChannelsOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -838,7 +831,6 @@ class StackdriverUpsertNotificationChannelOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 

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

@@ -137,7 +137,6 @@ class CloudTasksQueueCreateOperator(GoogleCloudBaseOperator):
                 metadata=self.metadata,
             )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -236,7 +235,6 @@ class CloudTasksQueueUpdateOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -319,7 +317,6 @@ class CloudTasksQueueGetOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -406,7 +403,6 @@ class CloudTasksQueuesListOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksLink.persist(
-            operator_instance=self,
             context=context,
             project_id=self.project_id or hook.project_id,
         )
@@ -564,7 +560,6 @@ class CloudTasksQueuePurgeOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -647,7 +642,6 @@ class CloudTasksQueuePauseOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -730,7 +724,6 @@ class CloudTasksQueueResumeOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -830,7 +823,6 @@ class CloudTasksTaskCreateOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=task.name,
         )
@@ -923,7 +915,6 @@ class CloudTasksTaskGetOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=task.name,
         )
@@ -1016,7 +1007,6 @@ class CloudTasksTasksListOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=f"projects/{self.project_id or hook.project_id}/"
             f"locations/{self.location}/queues/{self.queue_name}",
@@ -1190,7 +1180,6 @@ class CloudTasksTaskRunOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=task.name,
         )

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

@@ -150,7 +150,6 @@ class CloudTextToSpeechSynthesizeOperator(GoogleCloudBaseOperator):
             )
             FileDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 uri=f"{self.target_bucket_name}/{self.target_filename}",
                 project_id=cloud_storage_hook.project_id,
             )

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

@@ -37,6 +37,7 @@ from airflow.providers.google.cloud.links.translate import (
     TranslationNativeDatasetLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.cloud.operators.vertex_ai.dataset import DatasetImportDataResultsCheckHelper
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
@@ -394,7 +395,6 @@ class TranslateTextBatchOperator(GoogleCloudBaseOperator):
         self.log.info("Translate text batch job started.")
         TranslateTextBatchLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
             output_config=self.output_config,
         )
@@ -480,15 +480,15 @@ class TranslateCreateDatasetOperator(GoogleCloudBaseOperator):
         result = hook.wait_for_operation_result(result_operation)
         result = type(result).to_dict(result)
         dataset_id = hook.extract_object_id(result)
-        self.xcom_push(context, key="dataset_id", value=dataset_id)
+        context["ti"].xcom_push(key="dataset_id", value=dataset_id)
         self.log.info("Dataset creation complete. The dataset_id: %s.", dataset_id)
 
         project_id = self.project_id or hook.project_id
         TranslationNativeDatasetLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=dataset_id,
             project_id=project_id,
+            location=self.location,
         )
         return result
 
@@ -556,7 +556,6 @@ class TranslateDatasetsListOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         TranslationDatasetsListLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
         )
         self.log.info("Requesting datasets list")
@@ -577,7 +576,7 @@ class TranslateDatasetsListOperator(GoogleCloudBaseOperator):
         return result_ids
 
 
-class TranslateImportDataOperator(GoogleCloudBaseOperator):
+class TranslateImportDataOperator(GoogleCloudBaseOperator, DatasetImportDataResultsCheckHelper):
     """
     Import data to the translation dataset.
 
@@ -604,6 +603,7 @@ class TranslateImportDataOperator(GoogleCloudBaseOperator):
         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).
+    :param raise_for_empty_result: Raise an error if no additional data has been populated after the import.
     """
 
     template_fields: Sequence[str] = (
@@ -629,6 +629,7 @@ class TranslateImportDataOperator(GoogleCloudBaseOperator):
         retry: Retry | _MethodDefault = DEFAULT,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
+        raise_for_empty_result: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -641,9 +642,21 @@ class TranslateImportDataOperator(GoogleCloudBaseOperator):
         self.retry = retry
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
+        self.raise_for_empty_result = raise_for_empty_result
 
     def execute(self, context: Context):
         hook = TranslateHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+        initial_dataset_size = self._get_number_of_ds_items(
+            dataset=hook.get_dataset(
+                dataset_id=self.dataset_id,
+                project_id=self.project_id,
+                location=self.location,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            ),
+            total_key_name="example_count",
+        )
         self.log.info("Importing data to dataset...")
         operation = hook.import_dataset_data(
             dataset_id=self.dataset_id,
@@ -657,12 +670,27 @@ class TranslateImportDataOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         TranslationNativeDatasetLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=self.dataset_id,
             project_id=project_id,
+            location=self.location,
         )
         hook.wait_for_operation_done(operation=operation, timeout=self.timeout)
+
+        result_dataset_size = self._get_number_of_ds_items(
+            dataset=hook.get_dataset(
+                dataset_id=self.dataset_id,
+                project_id=self.project_id,
+                location=self.location,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            ),
+            total_key_name="example_count",
+        )
+        if self.raise_for_empty_result:
+            self._raise_for_empty_import_result(self.dataset_id, initial_dataset_size, result_dataset_size)
         self.log.info("Importing data finished!")
+        return {"total_imported": int(result_dataset_size) - int(initial_dataset_size)}
 
 
 class TranslateDeleteDatasetOperator(GoogleCloudBaseOperator):
@@ -821,16 +849,16 @@ class TranslateCreateModelOperator(GoogleCloudBaseOperator):
         result = hook.wait_for_operation_result(operation=result_operation)
         result = type(result).to_dict(result)
         model_id = hook.extract_object_id(result)
-        self.xcom_push(context, key="model_id", value=model_id)
+        context["ti"].xcom_push(key="model_id", value=model_id)
         self.log.info("Model creation complete. The model_id: %s.", model_id)
 
         project_id = self.project_id or hook.project_id
         TranslationModelLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=self.dataset_id,
             model_id=model_id,
             project_id=project_id,
+            location=self.location,
         )
         return result
 
@@ -898,7 +926,6 @@ class TranslateModelsListOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         TranslationModelsListLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
         )
         self.log.info("Requesting models list")
@@ -1141,7 +1168,6 @@ class TranslateDocumentOperator(GoogleCloudBaseOperator):
         if self.document_output_config:
             TranslateResultByOutputConfigLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=self.project_id or hook.project_id,
                 output_config=self.document_output_config,
             )
@@ -1304,7 +1330,6 @@ class TranslateDocumentBatchOperator(GoogleCloudBaseOperator):
         self.log.info("Batch document translation job started.")
         TranslateResultByOutputConfigLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
             output_config=self.output_config,
         )
@@ -1411,7 +1436,7 @@ class TranslateCreateGlossaryOperator(GoogleCloudBaseOperator):
         result = type(result).to_dict(result)
 
         glossary_id = hook.extract_object_id(result)
-        self.xcom_push(context, key="glossary_id", value=glossary_id)
+        context["ti"].xcom_push(key="glossary_id", value=glossary_id)
         self.log.info("Glossary creation complete. The glossary_id: %s.", glossary_id)
         return result
 
@@ -1610,7 +1635,6 @@ class TranslateListGlossariesOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         TranslationGlossariesListLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
         )
         self.log.info("Requesting glossaries list")

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

@@ -173,7 +173,6 @@ class CloudTranslateSpeechOperator(GoogleCloudBaseOperator):
         if self.audio.uri:
             FileDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 # Slice from: "gs://{BUCKET_NAME}/{FILE_NAME}" to: "{BUCKET_NAME}/{FILE_NAME}"
                 uri=self.audio.uri[5:],
                 project_id=self.project_id or translate_hook.project_id,