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

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

@@ -23,34 +23,19 @@ from collections.abc import MutableSequence, Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from google.protobuf.json_format import MessageToDict
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.triggers.dataplex import (
     DataplexDataProfileJobTrigger,
     DataplexDataQualityJobTrigger,
 )
-from google.protobuf.json_format import MessageToDict
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.protobuf.field_mask_pb2 import FieldMask
 
-from googleapiclient.errors import HttpError
+    from airflow.providers.common.compat.sdk import Context
 
-from airflow.configuration import conf
-from airflow.providers.google.cloud.hooks.dataplex import AirflowDataQualityScanException, DataplexHook
-from airflow.providers.google.cloud.links.dataplex import (
-    DataplexCatalogAspectTypeLink,
-    DataplexCatalogAspectTypesLink,
-    DataplexCatalogEntryGroupLink,
-    DataplexCatalogEntryGroupsLink,
-    DataplexCatalogEntryLink,
-    DataplexCatalogEntryTypeLink,
-    DataplexCatalogEntryTypesLink,
-    DataplexLakeLink,
-    DataplexTaskLink,
-    DataplexTasksLink,
-)
-from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from google.api_core.exceptions import AlreadyExists, GoogleAPICallError, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.api_core.retry import Retry, exponential_sleep_generator
@@ -72,6 +57,23 @@ from google.cloud.dataplex_v1.types import (
     Task,
     Zone,
 )
+from googleapiclient.errors import HttpError
+
+from airflow.configuration import conf
+from airflow.providers.google.cloud.hooks.dataplex import AirflowDataQualityScanException, DataplexHook
+from airflow.providers.google.cloud.links.dataplex import (
+    DataplexCatalogAspectTypeLink,
+    DataplexCatalogAspectTypesLink,
+    DataplexCatalogEntryGroupLink,
+    DataplexCatalogEntryGroupsLink,
+    DataplexCatalogEntryLink,
+    DataplexCatalogEntryTypeLink,
+    DataplexCatalogEntryTypesLink,
+    DataplexLakeLink,
+    DataplexTaskLink,
+    DataplexTasksLink,
+)
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 
 
 class DataplexCreateTaskOperator(GoogleCloudBaseOperator):
@@ -148,6 +150,15 @@ class DataplexCreateTaskOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.asynchronous = asynchronous
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "lake_id": self.lake_id,
+            "task_id": self.dataplex_task_id,
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict:
         hook = DataplexHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -155,7 +166,7 @@ class DataplexCreateTaskOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
         self.log.info("Creating Dataplex task %s", self.dataplex_task_id)
-        DataplexTaskLink.persist(context=context, task_instance=self)
+        DataplexTaskLink.persist(context=context)
 
         try:
             operation = hook.create_task(
@@ -349,6 +360,14 @@ class DataplexListTasksOperator(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,
+            "lake_id": self.lake_id,
+            "region": self.region,
+        }
+
     def execute(self, context: Context) -> list[dict]:
         hook = DataplexHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -356,7 +375,7 @@ class DataplexListTasksOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
         self.log.info("Listing Dataplex tasks from lake %s", self.lake_id)
-        DataplexTasksLink.persist(context=context, task_instance=self)
+        DataplexTasksLink.persist(context=context)
 
         tasks = hook.list_tasks(
             project_id=self.project_id,
@@ -428,6 +447,15 @@ class DataplexGetTaskOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "lake_id": self.lake_id,
+            "task_id": self.dataplex_task_id,
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict:
         hook = DataplexHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -435,7 +463,7 @@ class DataplexGetTaskOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
         self.log.info("Retrieving Dataplex task %s", self.dataplex_task_id)
-        DataplexTaskLink.persist(context=context, task_instance=self)
+        DataplexTaskLink.persist(context=context)
 
         task = hook.get_task(
             project_id=self.project_id,
@@ -446,7 +474,7 @@ class DataplexGetTaskOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        DataplexTasksLink.persist(context=context, task_instance=self)
+        DataplexTasksLink.persist(context=context)
         return Task.to_dict(task)
 
 
@@ -520,6 +548,14 @@ class DataplexCreateLakeOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.asynchronous = asynchronous
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "lake_id": self.lake_id,
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict:
         hook = DataplexHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -564,10 +600,7 @@ class DataplexCreateLakeOperator(GoogleCloudBaseOperator):
                 if lake["state"] != "CREATING":
                     break
                 time.sleep(time_to_wait)
-        DataplexLakeLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexLakeLink.persist(context=context)
         return Lake.to_dict(lake)
 
 
@@ -623,6 +656,14 @@ class DataplexDeleteLakeOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "lake_id": self.lake_id,
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = DataplexHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -640,7 +681,7 @@ class DataplexDeleteLakeOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        DataplexLakeLink.persist(context=context, task_instance=self)
+        DataplexLakeLink.persist(context=context)
         hook.wait_for_operation(timeout=self.timeout, operation=operation)
         self.log.info("Dataplex lake %s deleted successfully!", self.lake_id)
 
@@ -1041,11 +1082,11 @@ class DataplexRunDataQualityScanOperator(GoogleCloudBaseOperator):
         """
         job_state = event["job_state"]
         job_id = event["job_id"]
-        if job_state == DataScanJob.State.FAILED:
+        if job_state == DataScanJob.State.FAILED.name:  # type: ignore
             raise AirflowException(f"Job failed:\n{job_id}")
-        if job_state == DataScanJob.State.CANCELLED:
+        if job_state == DataScanJob.State.CANCELLED.name:  # type: ignore
             raise AirflowException(f"Job was cancelled:\n{job_id}")
-        if job_state == DataScanJob.State.SUCCEEDED:
+        if job_state == DataScanJob.State.SUCCEEDED.name:  # type: ignore
             job = event["job"]
             if not job["data_quality_result"]["passed"]:
                 if self.fail_on_dq_failure:
@@ -1219,11 +1260,11 @@ class DataplexGetDataQualityScanResultOperator(GoogleCloudBaseOperator):
         job_state = event["job_state"]
         job_id = event["job_id"]
         job = event["job"]
-        if job_state == DataScanJob.State.FAILED:
+        if job_state == DataScanJob.State.FAILED.name:  # type: ignore
             raise AirflowException(f"Job failed:\n{job_id}")
-        if job_state == DataScanJob.State.CANCELLED:
+        if job_state == DataScanJob.State.CANCELLED.name:  # type: ignore
             raise AirflowException(f"Job was cancelled:\n{job_id}")
-        if job_state == DataScanJob.State.SUCCEEDED:
+        if job_state == DataScanJob.State.SUCCEEDED.name:  # type: ignore
             if not job["data_quality_result"]["passed"]:
                 if self.fail_on_dq_failure:
                     raise AirflowDataQualityScanException(
@@ -1598,12 +1639,12 @@ class DataplexRunDataProfileScanOperator(GoogleCloudBaseOperator):
                 result_timeout=self.result_timeout,
             )
 
-            if job.state == DataScanJob.State.FAILED:
+            if job.state == DataScanJob.State.FAILED.name:  # type: ignore
                 raise AirflowException(f"Data Profile job failed: {job_id}")
-            if job.state == DataScanJob.State.SUCCEEDED:
+            if job.state == DataScanJob.State.SUCCEEDED.name:  # type: ignore
                 self.log.info("Data Profile job executed successfully.")
             else:
-                self.log.info("Data Profile job execution returned status: %s", job.status)
+                self.log.info("Data Profile job execution returned status: %s", job.state)
 
         return job_id
 
@@ -1616,11 +1657,11 @@ class DataplexRunDataProfileScanOperator(GoogleCloudBaseOperator):
         """
         job_state = event["job_state"]
         job_id = event["job_id"]
-        if job_state == DataScanJob.State.FAILED:
+        if job_state == DataScanJob.State.FAILED.name:  # type: ignore
             raise AirflowException(f"Job failed:\n{job_id}")
-        if job_state == DataScanJob.State.CANCELLED:
+        if job_state == DataScanJob.State.CANCELLED.name:  # type: ignore
             raise AirflowException(f"Job was cancelled:\n{job_id}")
-        if job_state == DataScanJob.State.SUCCEEDED:
+        if job_state == DataScanJob.State.SUCCEEDED.name:  # type: ignore
             self.log.info("Data Profile job executed successfully.")
         return job_id
 
@@ -2177,6 +2218,13 @@ class DataplexCatalogBaseOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "project_id": self.project_id,
+        }
+
 
 class DataplexCatalogCreateEntryGroupOperator(DataplexCatalogBaseOperator):
     """
@@ -2228,12 +2276,15 @@ class DataplexCatalogCreateEntryGroupOperator(DataplexCatalogBaseOperator):
         self.entry_group_configuration = entry_group_configuration
         self.validate_request = validate_request
 
-    def execute(self, context: Context):
-        DataplexCatalogEntryGroupLink.persist(
-            context=context,
-            task_instance=self,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_group_id": self.entry_group_id,
+        }
 
+    def execute(self, context: Context):
+        DataplexCatalogEntryGroupLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating a Create Dataplex Catalog EntryGroup request.")
         else:
@@ -2314,11 +2365,15 @@ class DataplexCatalogGetEntryGroupOperator(DataplexCatalogBaseOperator):
         super().__init__(*args, **kwargs)
         self.entry_group_id = entry_group_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_group_id": self.entry_group_id,
+        }
+
     def execute(self, context: Context):
-        DataplexCatalogEntryGroupLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogEntryGroupLink.persist(context=context)
         self.log.info(
             "Retrieving Dataplex Catalog EntryGroup %s.",
             self.entry_group_id,
@@ -2460,10 +2515,7 @@ class DataplexCatalogListEntryGroupsOperator(DataplexCatalogBaseOperator):
         self.order_by = order_by
 
     def execute(self, context: Context):
-        DataplexCatalogEntryGroupsLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogEntryGroupsLink.persist(context=context)
         self.log.info(
             "Listing Dataplex Catalog EntryGroup from location %s.",
             self.location,
@@ -2481,8 +2533,7 @@ class DataplexCatalogListEntryGroupsOperator(DataplexCatalogBaseOperator):
                 metadata=self.metadata,
             )
             self.log.info("EntryGroup on page: %s", entry_group_on_page)
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="entry_group_page",
                 value=ListEntryGroupsResponse.to_dict(entry_group_on_page._response),
             )
@@ -2552,12 +2603,15 @@ class DataplexCatalogUpdateEntryGroupOperator(DataplexCatalogBaseOperator):
         self.update_mask = update_mask
         self.validate_request = validate_request
 
-    def execute(self, context: Context):
-        DataplexCatalogEntryGroupLink.persist(
-            context=context,
-            task_instance=self,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_group_id": self.entry_group_id,
+        }
 
+    def execute(self, context: Context):
+        DataplexCatalogEntryGroupLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update Dataplex Catalog EntryGroup request.")
         else:
@@ -2642,12 +2696,15 @@ class DataplexCatalogCreateEntryTypeOperator(DataplexCatalogBaseOperator):
         self.entry_type_configuration = entry_type_configuration
         self.validate_request = validate_request
 
-    def execute(self, context: Context):
-        DataplexCatalogEntryTypeLink.persist(
-            context=context,
-            task_instance=self,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_type_id": self.entry_type_id,
+        }
 
+    def execute(self, context: Context):
+        DataplexCatalogEntryTypeLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating a Create Dataplex Catalog EntryType request.")
         else:
@@ -2728,11 +2785,15 @@ class DataplexCatalogGetEntryTypeOperator(DataplexCatalogBaseOperator):
         super().__init__(*args, **kwargs)
         self.entry_type_id = entry_type_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_type_id": self.entry_type_id,
+        }
+
     def execute(self, context: Context):
-        DataplexCatalogEntryTypeLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogEntryTypeLink.persist(context=context)
         self.log.info(
             "Retrieving Dataplex Catalog EntryType %s.",
             self.entry_type_id,
@@ -2874,10 +2935,7 @@ class DataplexCatalogListEntryTypesOperator(DataplexCatalogBaseOperator):
         self.order_by = order_by
 
     def execute(self, context: Context):
-        DataplexCatalogEntryTypesLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogEntryTypesLink.persist(context=context)
         self.log.info(
             "Listing Dataplex Catalog EntryType from location %s.",
             self.location,
@@ -2895,8 +2953,7 @@ class DataplexCatalogListEntryTypesOperator(DataplexCatalogBaseOperator):
                 metadata=self.metadata,
             )
             self.log.info("EntryType on page: %s", entry_type_on_page)
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="entry_type_page",
                 value=ListEntryTypesResponse.to_dict(entry_type_on_page._response),
             )
@@ -2966,12 +3023,15 @@ class DataplexCatalogUpdateEntryTypeOperator(DataplexCatalogBaseOperator):
         self.update_mask = update_mask
         self.validate_request = validate_request
 
-    def execute(self, context: Context):
-        DataplexCatalogEntryTypeLink.persist(
-            context=context,
-            task_instance=self,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_type_id": self.entry_type_id,
+        }
 
+    def execute(self, context: Context):
+        DataplexCatalogEntryTypeLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update Dataplex Catalog EntryType request.")
         else:
@@ -3056,12 +3116,15 @@ class DataplexCatalogCreateAspectTypeOperator(DataplexCatalogBaseOperator):
         self.aspect_type_configuration = aspect_type_configuration
         self.validate_request = validate_request
 
-    def execute(self, context: Context):
-        DataplexCatalogAspectTypeLink.persist(
-            context=context,
-            task_instance=self,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "aspect_type_id": self.aspect_type_id,
+        }
 
+    def execute(self, context: Context):
+        DataplexCatalogAspectTypeLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating a Create Dataplex Catalog AspectType request.")
         else:
@@ -3142,11 +3205,15 @@ class DataplexCatalogGetAspectTypeOperator(DataplexCatalogBaseOperator):
         super().__init__(*args, **kwargs)
         self.aspect_type_id = aspect_type_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "aspect_type_id": self.aspect_type_id,
+        }
+
     def execute(self, context: Context):
-        DataplexCatalogAspectTypeLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogAspectTypeLink.persist(context=context)
         self.log.info(
             "Retrieving Dataplex Catalog AspectType %s.",
             self.aspect_type_id,
@@ -3221,10 +3288,7 @@ class DataplexCatalogListAspectTypesOperator(DataplexCatalogBaseOperator):
         self.order_by = order_by
 
     def execute(self, context: Context):
-        DataplexCatalogAspectTypesLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogAspectTypesLink.persist(context=context)
         self.log.info(
             "Listing Dataplex Catalog AspectType from location %s.",
             self.location,
@@ -3242,8 +3306,7 @@ class DataplexCatalogListAspectTypesOperator(DataplexCatalogBaseOperator):
                 metadata=self.metadata,
             )
             self.log.info("AspectType on page: %s", aspect_type_on_page)
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="aspect_type_page",
                 value=ListAspectTypesResponse.to_dict(aspect_type_on_page._response),
             )
@@ -3313,12 +3376,15 @@ class DataplexCatalogUpdateAspectTypeOperator(DataplexCatalogBaseOperator):
         self.update_mask = update_mask
         self.validate_request = validate_request
 
-    def execute(self, context: Context):
-        DataplexCatalogAspectTypeLink.persist(
-            context=context,
-            task_instance=self,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "aspect_type_id": self.aspect_type_id,
+        }
 
+    def execute(self, context: Context):
+        DataplexCatalogAspectTypeLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update Dataplex Catalog AspectType request.")
         else:
@@ -3491,12 +3557,16 @@ class DataplexCatalogCreateEntryOperator(DataplexCatalogBaseOperator):
                 f"Missing required fields in Entry configuration: {', '.join(missing_fields)}. "
             )
 
-    def execute(self, context: Context):
-        DataplexCatalogEntryLink.persist(
-            context=context,
-            task_instance=self,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_id": self.entry_id,
+            "entry_group_id": self.entry_group_id,
+        }
 
+    def execute(self, context: Context):
+        DataplexCatalogEntryLink.persist(context=context)
         self._validate_fields(self.entry_configuration)
         try:
             entry = self.hook.create_entry(
@@ -3597,11 +3667,16 @@ class DataplexCatalogGetEntryOperator(DataplexCatalogBaseOperator):
         self.aspect_types = aspect_types
         self.paths = paths
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_id": self.entry_id,
+            "entry_group_id": self.entry_group_id,
+        }
+
     def execute(self, context: Context):
-        DataplexCatalogEntryLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogEntryLink.persist(context=context)
         self.log.info(
             "Retrieving Dataplex Catalog Entry %s.",
             self.entry_id,
@@ -3699,11 +3774,15 @@ class DataplexCatalogListEntriesOperator(DataplexCatalogBaseOperator):
         self.page_token = page_token
         self.filter_by = filter_by
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_group_id": self.entry_group_id,
+        }
+
     def execute(self, context: Context):
-        DataplexCatalogEntryGroupLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogEntryGroupLink.persist(context=context)
         self.log.info(
             "Listing Dataplex Catalog Entry from location %s.",
             self.location,
@@ -3721,8 +3800,7 @@ class DataplexCatalogListEntriesOperator(DataplexCatalogBaseOperator):
                 metadata=self.metadata,
             )
             self.log.info("Entries on page: %s", entries_on_page)
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="entry_page",
                 value=ListEntriesResponse.to_dict(entries_on_page._response),
             )
@@ -3819,8 +3897,7 @@ class DataplexCatalogSearchEntriesOperator(DataplexCatalogBaseOperator):
                 metadata=self.metadata,
             )
             self.log.info("Entries on page: %s", entries_on_page)
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="entry_page",
                 value=SearchEntriesResponse.to_dict(entries_on_page._response),
             )
@@ -3901,11 +3978,16 @@ class DataplexCatalogLookupEntryOperator(DataplexCatalogBaseOperator):
         self.aspect_types = aspect_types
         self.paths = paths
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_id": self.entry_id,
+            "entry_group_id": self.entry_group_id,
+        }
+
     def execute(self, context: Context):
-        DataplexCatalogEntryLink.persist(
-            context=context,
-            task_instance=self,
-        )
+        DataplexCatalogEntryLink.persist(context=context)
         self.log.info(
             "Looking for Dataplex Catalog Entry %s.",
             self.entry_id,
@@ -3961,7 +4043,7 @@ class DataplexCatalogUpdateEntryOperator(DataplexCatalogBaseOperator):
     :param entry_configuration: Required. The updated configuration body of the Entry.
     :param allow_missing: Optional. If set to true and entry doesn't exist, the service will create it.
     :param delete_missing_aspects: Optional. If set to true and the aspect_keys specify aspect
-        ranges, the service deletes any existing aspects from that range that weren't provided
+        ranges, the service deletes any existing aspects from that range that were not provided
         in the request.
     :param aspect_keys: Optional. The map keys of the Aspects which the service should modify.
         It supports the following syntax:
@@ -4020,12 +4102,16 @@ class DataplexCatalogUpdateEntryOperator(DataplexCatalogBaseOperator):
         self.delete_missing_aspects = delete_missing_aspects
         self.aspect_keys = aspect_keys
 
-    def execute(self, context: Context):
-        DataplexCatalogEntryLink.persist(
-            context=context,
-            task_instance=self,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            **super().extra_links_params,
+            "entry_id": self.entry_id,
+            "entry_group_id": self.entry_group_id,
+        }
 
+    def execute(self, context: Context):
+        DataplexCatalogEntryLink.persist(context=context)
         try:
             entry = self.hook.update_entry(
                 location=self.location,

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

@@ -28,7 +28,7 @@ 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
 
 
 class DataprepGetJobsForJobGroupOperator(GoogleCloudBaseOperator):
@@ -113,7 +113,6 @@ class DataprepGetJobGroupOperator(GoogleCloudBaseOperator):
         if self.project_id:
             DataprepJobGroupLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=self.project_id,
                 job_group_id=int(self.job_group_id),
             )
@@ -170,7 +169,6 @@ class DataprepRunJobGroupOperator(GoogleCloudBaseOperator):
         if self.project_id and job_group_id:
             DataprepJobGroupLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=self.project_id,
                 job_group_id=int(job_group_id),
             )
@@ -230,7 +228,6 @@ class DataprepCopyFlowOperator(GoogleCloudBaseOperator):
         if self.project_id and copied_flow_id:
             DataprepFlowLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=self.project_id,
                 flow_id=int(copied_flow_id),
             )
@@ -303,7 +300,6 @@ class DataprepRunFlowOperator(GoogleCloudBaseOperator):
             job_group_id = response["data"][0]["id"]
             DataprepJobGroupLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=self.project_id,
                 job_group_id=int(job_group_id),
             )