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

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,
         )
@@ -188,13 +189,12 @@ class SpannerDeleteInstanceOperator(GoogleCloudBaseOperator):
         )
         if hook.get_instance(project_id=self.project_id, instance_id=self.instance_id):
             return hook.delete_instance(project_id=self.project_id, instance_id=self.instance_id)
-        else:
-            self.log.info(
-                "Instance '%s' does not exist in project '%s'. Aborting delete.",
-                self.instance_id,
-                self.project_id,
-            )
-            return True
+        self.log.info(
+            "Instance '%s' does not exist in project '%s'. Aborting delete.",
+            self.instance_id,
+            self.project_id,
+        )
+        return True
 
 
 class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
@@ -256,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")
@@ -267,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)
@@ -282,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,
@@ -291,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:
@@ -307,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):
     """
@@ -381,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,
@@ -401,14 +414,12 @@ class SpannerDeployDatabaseInstanceOperator(GoogleCloudBaseOperator):
                 database_id=self.database_id,
                 ddl_statements=self.ddl_statements,
             )
-        else:
-            self.log.info(
-                "The database '%s' in project '%s' and instance '%s'"
-                " already exists. Nothing to do. Exiting.",
-                self.database_id,
-                self.project_id,
-                self.instance_id,
-            )
+        self.log.info(
+            "The database '%s' in project '%s' and instance '%s' already exists. Nothing to do. Exiting.",
+            self.database_id,
+            self.project_id,
+            self.instance_id,
+        )
         return True
 
 
@@ -497,21 +508,19 @@ class SpannerUpdateDatabaseInstanceOperator(GoogleCloudBaseOperator):
                 f"and instance '{self.instance_id}' is missing. "
                 f"Create the database first before you can update it."
             )
-        else:
-            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,
-            )
-            return hook.update_database(
-                project_id=self.project_id,
-                instance_id=self.instance_id,
-                database_id=self.database_id,
-                ddl_statements=self.ddl_statements,
-                operation_id=self.operation_id,
-            )
+        SpannerDatabaseLink.persist(
+            context=context,
+            instance_id=self.instance_id,
+            database_id=self.database_id,
+            project_id=self.project_id or hook.project_id,
+        )
+        return hook.update_database(
+            project_id=self.project_id,
+            instance_id=self.instance_id,
+            database_id=self.database_id,
+            ddl_statements=self.ddl_statements,
+            operation_id=self.operation_id,
+        )
 
 
 class SpannerDeleteDatabaseInstanceOperator(GoogleCloudBaseOperator):
@@ -590,7 +599,6 @@ class SpannerDeleteDatabaseInstanceOperator(GoogleCloudBaseOperator):
                 self.instance_id,
             )
             return True
-        else:
-            return hook.delete_database(
-                project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id
-            )
+        return hook.delete_database(
+            project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id
+        )

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

@@ -22,19 +22,21 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.protobuf.json_format import MessageToDict
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.speech_to_text import CloudSpeechToTextHook, RecognitionAudio
 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.common.links.storage import FileDetailsLink
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.protobuf.json_format import MessageToDict
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.speech_v1.types import RecognitionConfig
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class CloudSpeechToTextRecognizeSpeechOperator(GoogleCloudBaseOperator):
     """
@@ -115,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

@@ -20,6 +20,9 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel
+
 from airflow.providers.google.cloud.hooks.stackdriver import StackdriverHook
 from airflow.providers.google.cloud.links.stackdriver import (
     StackdriverNotificationsLink,
@@ -27,13 +30,12 @@ from airflow.providers.google.cloud.links.stackdriver import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class StackdriverListAlertPoliciesOperator(GoogleCloudBaseOperator):
     """
@@ -143,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]
@@ -226,7 +227,6 @@ class StackdriverEnableAlertPoliciesOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -309,7 +309,6 @@ class StackdriverDisableAlertPoliciesOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -392,7 +391,6 @@ class StackdriverUpsertAlertOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -578,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]
@@ -664,7 +661,6 @@ class StackdriverEnableNotificationChannelsOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -749,7 +745,6 @@ class StackdriverDisableNotificationChannelsOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -836,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

@@ -22,19 +22,21 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.exceptions import AlreadyExists
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.tasks_v2.types import Queue, Task
+
 from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook
 from airflow.providers.google.cloud.links.cloud_tasks import CloudTasksLink, CloudTasksQueueLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from google.api_core.exceptions import AlreadyExists
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.tasks_v2.types import Queue, Task
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 MetaData = Sequence[tuple[str, str]]
 
@@ -135,7 +137,6 @@ class CloudTasksQueueCreateOperator(GoogleCloudBaseOperator):
                 metadata=self.metadata,
             )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -234,7 +235,6 @@ class CloudTasksQueueUpdateOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -317,7 +317,6 @@ class CloudTasksQueueGetOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -404,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,
         )
@@ -562,7 +560,6 @@ class CloudTasksQueuePurgeOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -645,7 +642,6 @@ class CloudTasksQueuePauseOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -728,7 +724,6 @@ class CloudTasksQueueResumeOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -828,7 +823,6 @@ class CloudTasksTaskCreateOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=task.name,
         )
@@ -921,7 +915,6 @@ class CloudTasksTaskGetOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=task.name,
         )
@@ -1014,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}",
@@ -1188,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

@@ -23,19 +23,21 @@ from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.hooks.text_to_speech import CloudTextToSpeechHook
 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.common.links.storage import FileDetailsLink
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.texttospeech_v1.types import AudioConfig, SynthesisInput, VoiceSelectionParams
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class CloudTextToSpeechSynthesizeOperator(GoogleCloudBaseOperator):
     """
@@ -148,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

@@ -22,6 +22,9 @@ from __future__ import annotations
 from collections.abc import MutableMapping, MutableSequence, Sequence
 from typing import TYPE_CHECKING, cast
 
+from google.api_core.exceptions import GoogleAPICallError
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook, TranslateHook
 from airflow.providers.google.cloud.links.translate import (
@@ -34,12 +37,10 @@ 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
-from google.api_core.exceptions import GoogleAPICallError
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.translate_v3.types import (
         BatchDocumentInputConfig,
@@ -55,6 +56,8 @@ if TYPE_CHECKING:
     )
     from google.cloud.translate_v3.types.translation_service import Glossary, GlossaryInputConfig
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class CloudTranslateTextOperator(GoogleCloudBaseOperator):
     """
@@ -392,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,
         )
@@ -478,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
 
@@ -554,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")
@@ -575,7 +576,7 @@ class TranslateDatasetsListOperator(GoogleCloudBaseOperator):
         return result_ids
 
 
-class TranslateImportDataOperator(GoogleCloudBaseOperator):
+class TranslateImportDataOperator(GoogleCloudBaseOperator, DatasetImportDataResultsCheckHelper):
     """
     Import data to the translation dataset.
 
@@ -602,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] = (
@@ -627,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)
@@ -639,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,
@@ -655,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):
@@ -819,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
 
@@ -896,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")
@@ -1139,11 +1168,10 @@ 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,
             )
-        return cast(dict, type(doc_translation_result).to_dict(doc_translation_result))
+        return cast("dict", type(doc_translation_result).to_dict(doc_translation_result))
 
 
 class TranslateDocumentBatchOperator(GoogleCloudBaseOperator):
@@ -1302,13 +1330,12 @@ 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,
         )
         result = hook.wait_for_operation_result(batch_document_translate_operation)
         self.log.info("Batch document translation job finished")
-        return cast(dict, type(result).to_dict(result))
+        return cast("dict", type(result).to_dict(result))
 
 
 class TranslateCreateGlossaryOperator(GoogleCloudBaseOperator):
@@ -1409,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
 
@@ -1608,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")
@@ -1628,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)