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

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

@@ -807,7 +807,6 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocClusterLink.persist(
                 context=context,
-                operator=self,
                 cluster_id=self.cluster_name,
                 project_id=project_id,
                 region=self.region,
@@ -1174,7 +1173,6 @@ class DataprocStartClusterOperator(_DataprocStartStopClusterBaseOperator):
         cluster = super().execute(context)
         DataprocClusterLink.persist(
             context=context,
-            operator=self,
             cluster_id=self.cluster_name,
             project_id=self._get_project_id(),
             region=self.region,
@@ -1355,7 +1353,11 @@ class DataprocJobBaseOperator(GoogleCloudBaseOperator):
             self.log.info("Job %s submitted successfully.", job_id)
             # Save data required for extra links no matter what the job status will be
             DataprocLink.persist(
-                context=context, task_instance=self, url=DATAPROC_JOB_LINK_DEPRECATED, resource=job_id
+                context=context,
+                url=DATAPROC_JOB_LINK_DEPRECATED,
+                resource=job_id,
+                region=self.region,
+                project_id=self.project_id,
             )
 
             if self.deferrable:
@@ -1459,7 +1461,6 @@ class DataprocCreateWorkflowTemplateOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocWorkflowTemplateLink.persist(
                 context=context,
-                operator=self,
                 workflow_template_id=self.template["id"],
                 region=self.region,
                 project_id=project_id,
@@ -1571,7 +1572,6 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocWorkflowLink.persist(
                 context=context,
-                operator=self,
                 workflow_id=workflow_id,
                 region=self.region,
                 project_id=project_id,
@@ -1727,7 +1727,6 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
         if project_id:
             DataprocWorkflowLink.persist(
                 context=context,
-                operator=self,
                 workflow_id=workflow_id,
                 region=self.region,
                 project_id=project_id,
@@ -1901,7 +1900,6 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocJobLink.persist(
                 context=context,
-                operator=self,
                 job_id=new_job_id,
                 region=self.region,
                 project_id=project_id,
@@ -2074,7 +2072,6 @@ class DataprocUpdateClusterOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocClusterLink.persist(
                 context=context,
-                operator=self,
                 cluster_id=self.cluster_name,
                 project_id=project_id,
                 region=self.region,
@@ -2373,7 +2370,6 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
             # Persist the link earlier so users can observe the progress
             DataprocBatchLink.persist(
                 context=context,
-                operator=self,
                 project_id=self.project_id,
                 region=self.region,
                 batch_id=self.batch_id,
@@ -2410,7 +2406,6 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
 
         DataprocBatchLink.persist(
             context=context,
-            operator=self,
             project_id=self.project_id,
             region=self.region,
             batch_id=batch_id,
@@ -2493,13 +2488,13 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
         link = DATAPROC_BATCH_LINK.format(region=self.region, project_id=self.project_id, batch_id=batch_id)
         if state == Batch.State.FAILED:
             raise AirflowException(
-                f"Batch job {batch_id} failed with error: {state_message}\nDriver Logs: {link}"
+                f"Batch job {batch_id} failed with error: {state_message}.\nDriver logs: {link}"
             )
         if state in (Batch.State.CANCELLED, Batch.State.CANCELLING):
-            raise AirflowException(f"Batch job {batch_id} was cancelled. Driver logs: {link}")
+            raise AirflowException(f"Batch job {batch_id} was cancelled.\nDriver logs: {link}")
         if state == Batch.State.STATE_UNSPECIFIED:
-            raise AirflowException(f"Batch job {batch_id} unspecified. Driver logs: {link}")
-        self.log.info("Batch job %s completed. Driver logs: %s", batch_id, link)
+            raise AirflowException(f"Batch job {batch_id} unspecified.\nDriver logs: {link}")
+        self.log.info("Batch job %s completed.\nDriver logs: %s", batch_id, link)
 
     def retry_batch_creation(
         self,
@@ -2723,7 +2718,6 @@ class DataprocGetBatchOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocBatchLink.persist(
                 context=context,
-                operator=self,
                 project_id=project_id,
                 region=self.region,
                 batch_id=self.batch_id,
@@ -2806,7 +2800,7 @@ class DataprocListBatchesOperator(GoogleCloudBaseOperator):
         )
         project_id = self.project_id or hook.project_id
         if project_id:
-            DataprocBatchesListLink.persist(context=context, operator=self, project_id=project_id)
+            DataprocBatchesListLink.persist(context=context, project_id=project_id)
         return [Batch.to_dict(result) for result in results]
 
 

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

@@ -21,7 +21,7 @@ from __future__ import annotations
 
 import time
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import AlreadyExists
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -32,26 +32,17 @@ from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
+from airflow.providers.google.cloud.links.base import BaseGoogleLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.links.storage import StorageLink
 
 if TYPE_CHECKING:
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.models import BaseOperator
     from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.google.version_compat import BaseOperator
     from airflow.utils.context import Context
 
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseOperatorLink
-    from airflow.sdk.execution_time.xcom import XCom
-else:
-    from airflow.models import XCom  # type: ignore[no-redef]
-    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
-
-
 BASE_LINK = "https://console.cloud.google.com"
 METASTORE_BASE_LINK = BASE_LINK + "/dataproc/metastore/services/{region}/{service_id}"
 METASTORE_BACKUP_LINK = METASTORE_BASE_LINK + "/backups/{resource}?project={project_id}"
@@ -61,97 +52,50 @@ METASTORE_IMPORT_LINK = METASTORE_BASE_LINK + "/imports/{resource}?project={proj
 METASTORE_SERVICE_LINK = METASTORE_BASE_LINK + "/config?project={project_id}"
 
 
-class DataprocMetastoreLink(BaseOperatorLink):
+class DataprocMetastoreLink(BaseGoogleLink):
     """Helper class for constructing Dataproc Metastore resource link."""
 
     name = "Dataproc Metastore"
     key = "conf"
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: (
-            DataprocMetastoreCreateServiceOperator
-            | DataprocMetastoreGetServiceOperator
-            | DataprocMetastoreRestoreServiceOperator
-            | DataprocMetastoreUpdateServiceOperator
-            | DataprocMetastoreListBackupsOperator
-            | DataprocMetastoreExportMetadataOperator
-        ),
-        url: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataprocMetastoreLink.key,
-            value={
-                "region": task_instance.region,
-                "service_id": task_instance.service_id,
-                "project_id": task_instance.project_id,
-                "url": url,
-            },
-        )
-
     def get_link(
         self,
         operator: BaseOperator,
         *,
         ti_key: TaskInstanceKey,
     ) -> str:
-        conf = XCom.get_value(key=self.key, ti_key=ti_key)
-        return (
-            conf["url"].format(
-                region=conf["region"],
-                service_id=conf["service_id"],
-                project_id=conf["project_id"],
-            )
-            if conf
-            else ""
+        conf = self.get_config(operator, ti_key)
+        if not conf:
+            return ""
+
+        return conf["url"].format(
+            region=conf["region"],
+            service_id=conf["service_id"],
+            project_id=conf["project_id"],
         )
 
 
-class DataprocMetastoreDetailedLink(BaseOperatorLink):
+class DataprocMetastoreDetailedLink(BaseGoogleLink):
     """Helper class for constructing Dataproc Metastore detailed resource link."""
 
     name = "Dataproc Metastore resource"
     key = "config"
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: (
-            DataprocMetastoreCreateBackupOperator | DataprocMetastoreCreateMetadataImportOperator
-        ),
-        url: str,
-        resource: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataprocMetastoreDetailedLink.key,
-            value={
-                "region": task_instance.region,
-                "service_id": task_instance.service_id,
-                "project_id": task_instance.project_id,
-                "url": url,
-                "resource": resource,
-            },
-        )
-
     def get_link(
         self,
         operator: BaseOperator,
         *,
         ti_key: TaskInstanceKey,
     ) -> str:
-        conf = XCom.get_value(key=self.key, ti_key=ti_key)
-        return (
-            conf["url"].format(
-                region=conf["region"],
-                service_id=conf["service_id"],
-                project_id=conf["project_id"],
-                resource=conf["resource"],
-            )
-            if conf
-            else ""
+        conf = self.get_config(operator, ti_key)
+        if not conf:
+            return ""
+
+        return conf["url"].format(
+            region=conf["region"],
+            service_id=conf["service_id"],
+            project_id=conf["project_id"],
+            resource=conf["resource"],
         )
 
 
@@ -231,6 +175,14 @@ class DataprocMetastoreCreateBackupOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "service_id": self.service_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict:
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -263,7 +215,7 @@ class DataprocMetastoreCreateBackupOperator(GoogleCloudBaseOperator):
                 metadata=self.metadata,
             )
         DataprocMetastoreDetailedLink.persist(
-            context=context, task_instance=self, url=METASTORE_BACKUP_LINK, resource=self.backup_id
+            context=context, url=METASTORE_BACKUP_LINK, resource=self.backup_id
         )
         return Backup.to_dict(backup)
 
@@ -344,6 +296,14 @@ class DataprocMetastoreCreateMetadataImportOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "service_id": self.service_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -364,7 +324,7 @@ class DataprocMetastoreCreateMetadataImportOperator(GoogleCloudBaseOperator):
         self.log.info("Metadata import %s created successfully", self.metadata_import_id)
 
         DataprocMetastoreDetailedLink.persist(
-            context=context, task_instance=self, url=METASTORE_IMPORT_LINK, resource=self.metadata_import_id
+            context=context, url=METASTORE_IMPORT_LINK, resource=self.metadata_import_id
         )
         return MetadataImport.to_dict(metadata_import)
 
@@ -437,6 +397,14 @@ class DataprocMetastoreCreateServiceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "service_id": self.service_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict:
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -465,7 +433,7 @@ class DataprocMetastoreCreateServiceOperator(GoogleCloudBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-        DataprocMetastoreLink.persist(context=context, task_instance=self, url=METASTORE_SERVICE_LINK)
+        DataprocMetastoreLink.persist(context=context, url=METASTORE_SERVICE_LINK)
         return Service.to_dict(service)
 
 
@@ -689,6 +657,14 @@ class DataprocMetastoreExportMetadataOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "service_id": self.service_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -708,9 +684,9 @@ class DataprocMetastoreExportMetadataOperator(GoogleCloudBaseOperator):
         metadata_export = self._wait_for_export_metadata(hook)
         self.log.info("Metadata from service %s exported successfully", self.service_id)
 
-        DataprocMetastoreLink.persist(context=context, task_instance=self, url=METASTORE_EXPORT_LINK)
+        DataprocMetastoreLink.persist(context=context, url=METASTORE_EXPORT_LINK)
         uri = self._get_uri_from_destination(MetadataExport.to_dict(metadata_export)["destination_gcs_uri"])
-        StorageLink.persist(context=context, task_instance=self, uri=uri, project_id=self.project_id)
+        StorageLink.persist(context=context, uri=uri, project_id=self.project_id)
         return MetadataExport.to_dict(metadata_export)
 
     def _get_uri_from_destination(self, destination_uri: str):
@@ -799,6 +775,14 @@ class DataprocMetastoreGetServiceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "service_id": self.service_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict:
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -812,7 +796,7 @@ class DataprocMetastoreGetServiceOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        DataprocMetastoreLink.persist(context=context, task_instance=self, url=METASTORE_SERVICE_LINK)
+        DataprocMetastoreLink.persist(context=context, url=METASTORE_SERVICE_LINK)
         return Service.to_dict(result)
 
 
@@ -880,6 +864,14 @@ class DataprocMetastoreListBackupsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "service_id": self.service_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> list[dict]:
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -897,7 +889,7 @@ class DataprocMetastoreListBackupsOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        DataprocMetastoreLink.persist(context=context, task_instance=self, url=METASTORE_BACKUPS_LINK)
+        DataprocMetastoreLink.persist(context=context, url=METASTORE_BACKUPS_LINK)
         return [Backup.to_dict(backup) for backup in backups]
 
 
@@ -981,6 +973,14 @@ class DataprocMetastoreRestoreServiceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "service_id": self.service_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -1004,7 +1004,7 @@ class DataprocMetastoreRestoreServiceOperator(GoogleCloudBaseOperator):
         )
         self._wait_for_restore_service(hook)
         self.log.info("Service %s restored from backup %s", self.service_id, self.backup_id)
-        DataprocMetastoreLink.persist(context=context, task_instance=self, url=METASTORE_SERVICE_LINK)
+        DataprocMetastoreLink.persist(context=context, url=METASTORE_SERVICE_LINK)
 
     def _wait_for_restore_service(self, hook: DataprocMetastoreHook):
         """
@@ -1107,6 +1107,14 @@ class DataprocMetastoreUpdateServiceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "service_id": self.service_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -1126,4 +1134,4 @@ class DataprocMetastoreUpdateServiceOperator(GoogleCloudBaseOperator):
         )
         hook.wait_for_operation(self.timeout, operation)
         self.log.info("Service %s updated successfully", self.service.get("name"))
-        DataprocMetastoreLink.persist(context=context, task_instance=self, url=METASTORE_SERVICE_LINK)
+        DataprocMetastoreLink.persist(context=context, url=METASTORE_SERVICE_LINK)

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

@@ -136,7 +136,6 @@ class CloudDatastoreExportEntitiesOperator(GoogleCloudBaseOperator):
             raise AirflowException(f"Operation failed: result={result}")
         StorageLink.persist(
             context=context,
-            task_instance=self,
             uri=f"{self.bucket}/{result['response']['outputUrl'].split('/')[3]}",
             project_id=self.project_id or ds_hook.project_id,
         )
@@ -211,6 +210,12 @@ class CloudDatastoreImportEntitiesOperator(GoogleCloudBaseOperator):
         self.project_id = project_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         self.log.info("Importing data from Cloud Storage bucket %s", self.bucket)
         ds_hook = DatastoreHook(
@@ -231,8 +236,7 @@ class CloudDatastoreImportEntitiesOperator(GoogleCloudBaseOperator):
         state = result["metadata"]["common"]["state"]
         if state != "SUCCESSFUL":
             raise AirflowException(f"Operation failed: result={result}")
-
-        CloudDatastoreImportExportLink.persist(context=context, task_instance=self)
+        CloudDatastoreImportExportLink.persist(context=context)
         return result
 
 
@@ -282,6 +286,12 @@ class CloudDatastoreAllocateIdsOperator(GoogleCloudBaseOperator):
         self.project_id = project_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> list:
         hook = DatastoreHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -291,7 +301,7 @@ class CloudDatastoreAllocateIdsOperator(GoogleCloudBaseOperator):
             partial_keys=self.partial_keys,
             project_id=self.project_id,
         )
-        CloudDatastoreEntitiesLink.persist(context=context, task_instance=self)
+        CloudDatastoreEntitiesLink.persist(context=context)
         return keys
 
 
@@ -398,6 +408,12 @@ class CloudDatastoreCommitOperator(GoogleCloudBaseOperator):
         self.project_id = project_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict:
         hook = DatastoreHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -407,7 +423,7 @@ class CloudDatastoreCommitOperator(GoogleCloudBaseOperator):
             body=self.body,
             project_id=self.project_id,
         )
-        CloudDatastoreEntitiesLink.persist(context=context, task_instance=self)
+        CloudDatastoreEntitiesLink.persist(context=context)
         return response