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

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

@@ -25,8 +25,9 @@ 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 functools import cached_property
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.pubsub_v1.types import (
@@ -40,6 +41,7 @@ from google.cloud.pubsub_v1.types import (
     SchemaSettings,
 )
 
+from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
 from airflow.providers.google.cloud.links.pubsub import PubSubSubscriptionLink, PubSubTopicLink
@@ -51,7 +53,8 @@ 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
+    from airflow.providers.openlineage.extractors import OperatorLineage
 
 
 class PubSubCreateTopicOperator(GoogleCloudBaseOperator):
@@ -182,7 +185,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,
         )
@@ -359,15 +361,18 @@ class PubSubCreateSubscriptionOperator(GoogleCloudBaseOperator):
         self.timeout = timeout
         self.metadata = metadata
         self.impersonation_chain = impersonation_chain
+        self._resolved_subscription_name: str | None = None
 
-    def execute(self, context: Context) -> str:
-        hook = PubSubHook(
+    @cached_property
+    def pubsub_hook(self):
+        return PubSubHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
 
+    def execute(self, context: Context) -> str:
         self.log.info("Creating subscription for topic %s", self.topic)
-        result = hook.create_subscription(
+        result = self.pubsub_hook.create_subscription(
             project_id=self.project_id,
             topic=self.topic,
             subscription=self.subscription,
@@ -389,14 +394,34 @@ class PubSubCreateSubscriptionOperator(GoogleCloudBaseOperator):
         )
 
         self.log.info("Created subscription for topic %s", self.topic)
+
+        # Store resolved subscription for Open Lineage
+        self._resolved_subscription_name = self.subscription or result
+
         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,
+            subscription_id=self._resolved_subscription_name,  # result returns subscription name
+            project_id=self.project_id or self.pubsub_hook.project_id,
         )
         return result
 
+    def get_openlineage_facets_on_complete(self, _) -> OperatorLineage:
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        topic_project_id = self.project_id or self.pubsub_hook.project_id
+        subscription_project_id = self.subscription_project_id or topic_project_id
+
+        return OperatorLineage(
+            inputs=[Dataset(namespace="pubsub", name=f"topic:{topic_project_id}:{self.topic}")],
+            outputs=[
+                Dataset(
+                    namespace="pubsub",
+                    name=f"subscription:{subscription_project_id}:{self._resolved_subscription_name}",
+                )
+            ],
+        )
+
 
 class PubSubDeleteTopicOperator(GoogleCloudBaseOperator):
     """
@@ -693,17 +718,28 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         self.enable_message_ordering = enable_message_ordering
         self.impersonation_chain = impersonation_chain
 
-    def execute(self, context: Context) -> None:
-        hook = PubSubHook(
+    @cached_property
+    def pubsub_hook(self):
+        return PubSubHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
             enable_message_ordering=self.enable_message_ordering,
         )
 
+    def execute(self, context: Context) -> None:
         self.log.info("Publishing to topic %s", self.topic)
-        hook.publish(project_id=self.project_id, topic=self.topic, messages=self.messages)
+        self.pubsub_hook.publish(project_id=self.project_id, topic=self.topic, messages=self.messages)
         self.log.info("Published to topic %s", self.topic)
 
+    def get_openlineage_facets_on_complete(self, _) -> OperatorLineage:
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        project_id = self.project_id or self.pubsub_hook.project_id
+        output_dataset = [Dataset(namespace="pubsub", name=f"topic:{project_id}:{self.topic}")]
+
+        return OperatorLineage(outputs=output_dataset)
+
 
 class PubSubPullOperator(GoogleCloudBaseOperator):
     """
@@ -770,7 +806,7 @@ class PubSubPullOperator(GoogleCloudBaseOperator):
         messages_callback: Callable[[list[ReceivedMessage], Context], Any] | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
-        deferrable: bool = False,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         poll_interval: int = 300,
         **kwargs,
     ) -> None:
@@ -854,3 +890,13 @@ class PubSubPullOperator(GoogleCloudBaseOperator):
         messages_json = [ReceivedMessage.to_dict(m) for m in pulled_messages]
 
         return messages_json
+
+    def get_openlineage_facets_on_complete(self, _) -> OperatorLineage:
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        output_dataset = [
+            Dataset(namespace="pubsub", name=f"subscription:{self.project_id}:{self.subscription}")
+        ]
+
+        return OperatorLineage(outputs=output_dataset)

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

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+from functools import cached_property
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
@@ -29,7 +30,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
+    from airflow.providers.openlineage.extractors import OperatorLineage
 
 
 class SpannerDeployInstanceOperator(GoogleCloudBaseOperator):
@@ -122,7 +124,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,
         )
@@ -255,6 +256,13 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @cached_property
+    def hook(self) -> SpannerHook:
+        return SpannerHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
     def _validate_inputs(self) -> None:
         if self.project_id == "":
             raise AirflowException("The required parameter 'project_id' is empty")
@@ -266,10 +274,6 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
             raise AirflowException("The required parameter 'query' is empty")
 
     def execute(self, context: Context):
-        hook = SpannerHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
         if isinstance(self.query, str):
             queries = [x.strip() for x in self.query.split(";")]
             self.sanitize_queries(queries)
@@ -281,8 +285,8 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
             self.instance_id,
             self.database_id,
         )
-        self.log.info(queries)
-        hook.execute_dml(
+        self.log.info("Executing queries: %s", queries)
+        result_rows_count_per_query = self.hook.execute_dml(
             project_id=self.project_id,
             instance_id=self.instance_id,
             database_id=self.database_id,
@@ -290,11 +294,11 @@ 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,
+            project_id=self.project_id or self.hook.project_id,
         )
+        return result_rows_count_per_query
 
     @staticmethod
     def sanitize_queries(queries: list[str]) -> None:
@@ -306,6 +310,17 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
         if queries and queries[-1] == "":
             queries.pop()
 
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage | None:
+        """Build a generic OpenLineage facet, aligned with SQL-based operators."""
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+
+        return get_openlineage_facets_with_sql(
+            hook=self.hook,
+            sql=self.query,
+            conn_id=self.gcp_conn_id,
+            database=self.database_id,
+        )
+
 
 class SpannerDeployDatabaseInstanceOperator(GoogleCloudBaseOperator):
     """
@@ -380,7 +395,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 +510,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

@@ -35,7 +35,7 @@ if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.speech_v1.types import RecognitionConfig
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CloudSpeechToTextRecognizeSpeechOperator(GoogleCloudBaseOperator):
@@ -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

@@ -34,7 +34,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
 
 
 class StackdriverListAlertPoliciesOperator(GoogleCloudBaseOperator):
@@ -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

@@ -35,7 +35,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]]
@@ -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

@@ -36,7 +36,7 @@ if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.texttospeech_v1.types import AudioConfig, SynthesisInput, VoiceSelectionParams
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CloudTextToSpeechSynthesizeOperator(GoogleCloudBaseOperator):
@@ -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:
@@ -55,7 +56,7 @@ if TYPE_CHECKING:
     )
     from google.cloud.translate_v3.types.translation_service import Glossary, GlossaryInputConfig
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CloudTranslateTextOperator(GoogleCloudBaseOperator):
@@ -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")
@@ -1630,8 +1654,8 @@ class TranslateListGlossariesOperator(GoogleCloudBaseOperator):
             raise AirflowException(e)
 
         result_ids = []
-        for glossary_item in results_pager:
-            glossary_item = type(glossary_item).to_dict(glossary_item)
+        for glossary_item_raw in results_pager:
+            glossary_item = type(glossary_item_raw).to_dict(glossary_item_raw)
             glossary_id = hook.extract_object_id(glossary_item)
             result_ids.append(glossary_id)
         self.log.info("Fetching the glossaries list complete. Glossary id-s: %s", result_ids)

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

@@ -34,7 +34,7 @@ from airflow.providers.google.common.links.storage import FileDetailsLink
 if TYPE_CHECKING:
     from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CloudTranslateSpeechOperator(GoogleCloudBaseOperator):
@@ -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,