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

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

@@ -111,10 +111,9 @@ class CloudDataFusionRestartInstanceOperator(GoogleCloudBaseOperator):
         project_id = resource_path_to_dict(resource_name=instance["name"])["projects"]
         DataFusionInstanceLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
             instance_name=self.instance_name,
-            location=self.location,
+            region=self.location,
         )
 
 
@@ -269,10 +268,9 @@ class CloudDataFusionCreateInstanceOperator(GoogleCloudBaseOperator):
         project_id = resource_path_to_dict(resource_name=instance["name"])["projects"]
         DataFusionInstanceLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
             instance_name=self.instance_name,
-            location=self.location,
+            region=self.location,
         )
         return instance
 
@@ -358,10 +356,9 @@ class CloudDataFusionUpdateInstanceOperator(GoogleCloudBaseOperator):
         project_id = resource_path_to_dict(resource_name=instance["name"])["projects"]
         DataFusionInstanceLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
             instance_name=self.instance_name,
-            location=self.location,
+            region=self.location,
         )
 
 
@@ -429,10 +426,9 @@ class CloudDataFusionGetInstanceOperator(GoogleCloudBaseOperator):
         project_id = resource_path_to_dict(resource_name=instance["name"])["projects"]
         DataFusionInstanceLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
             instance_name=self.instance_name,
-            location=self.location,
+            region=self.location,
         )
         return instance
 
@@ -519,7 +515,6 @@ class CloudDataFusionCreatePipelineOperator(GoogleCloudBaseOperator):
         )
         DataFusionPipelineLink.persist(
             context=context,
-            task_instance=self,
             uri=instance["serviceEndpoint"],
             pipeline_name=self.pipeline_name,
             namespace=self.namespace,
@@ -693,7 +688,6 @@ class CloudDataFusionListPipelinesOperator(GoogleCloudBaseOperator):
 
         DataFusionPipelinesLink.persist(
             context=context,
-            task_instance=self,
             uri=service_endpoint,
             namespace=self.namespace,
         )
@@ -813,7 +807,6 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
 
         DataFusionPipelineLink.persist(
             context=context,
-            task_instance=self,
             uri=instance["serviceEndpoint"],
             pipeline_name=self.pipeline_name,
             namespace=self.namespace,
@@ -943,7 +936,6 @@ class CloudDataFusionStopPipelineOperator(GoogleCloudBaseOperator):
 
         DataFusionPipelineLink.persist(
             context=context,
-            task_instance=self,
             uri=instance["serviceEndpoint"],
             pipeline_name=self.pipeline_name,
             namespace=self.namespace,

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

@@ -150,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,
@@ -157,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(
@@ -351,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,
@@ -358,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,
@@ -430,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,
@@ -437,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,
@@ -448,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)
 
 
@@ -522,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,
@@ -566,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)
 
 
@@ -625,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,
@@ -642,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)
 
@@ -2179,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):
     """
@@ -2230,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:
@@ -2316,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,
@@ -2462,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,
@@ -2483,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),
             )
@@ -2554,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:
@@ -2644,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:
@@ -2730,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,
@@ -2876,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,
@@ -2897,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),
             )
@@ -2968,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:
@@ -3058,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:
@@ -3144,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,
@@ -3223,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,
@@ -3244,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),
             )
@@ -3315,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:
@@ -3493,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(
@@ -3599,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,
@@ -3701,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,
@@ -3723,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),
             )
@@ -3821,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),
             )
@@ -3903,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,
@@ -3963,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:
@@ -4022,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

@@ -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),
             )