Change CloudDatastoreExportEntitiesLink to StorageLink

Author: MaksYermak

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

@@ -52,30 +52,6 @@ def persist(
         )
 
 
-class CloudDatastoreExportEntitiesLink(BaseGoogleLink):
-    """Helper class for constructing Cloud Datastore Export Entities Link"""
-
-    name = "Export Entities"
-    key = "export_conf"
-    format_str = DATASTORE_EXPORT_ENTITIES_LINK
-
-    @staticmethod
-    def persist(
-        context: "Context",
-        task_instance,
-        output_url: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=CloudDatastoreExportEntitiesLink.key,
-            value={
-                "project_id": task_instance.project_id,
-                "bucket_name": task_instance.bucket,
-                "export_name": output_url.split('/')[3],
-            },
-        )
-
-
 class CloudDatastoreEntitiesLink(BaseGoogleLink):
     """Helper class for constructing Cloud Datastore Entities Link"""
 

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

@@ -25,9 +25,9 @@
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.links.datastore import (
     CloudDatastoreEntitiesLink,
-    CloudDatastoreExportEntitiesLink,
     CloudDatastoreImportExportLink,
 )
+from airflow.providers.google.common.links.storage import StorageLink
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -78,7 +78,7 @@ class CloudDatastoreExportEntitiesOperator(BaseOperator):
         'labels',
         'impersonation_chain',
     )
-    operator_extra_links = (CloudDatastoreExportEntitiesLink(),)
+    operator_extra_links = (StorageLink(),)
 
     def __init__(
         self,
@@ -138,10 +138,10 @@ def execute(self, context: 'Context') -> dict:
         state = result['metadata']['common']['state']
         if state != 'SUCCESSFUL':
             raise AirflowException(f'Operation failed: result={result}')
-        CloudDatastoreExportEntitiesLink.persist(
+        StorageLink.persist(
             context=context,
             task_instance=self,
-            output_url=result['response']['outputUrl'],
+            uri=f"{self.bucket}/{result['response']['outputUrl'].split('/')[3]}",
         )
         return result
 

airflow/providers/google/provider.yaml

@@ -905,7 +905,6 @@ extra-links:
   - airflow.providers.google.cloud.operators.cloud_composer.CloudComposerEnvironmentsLink
   - airflow.providers.google.cloud.links.dataflow.DataflowJobLink
   - airflow.providers.google.cloud.links.datastore.CloudDatastoreImportExportLink
-  - airflow.providers.google.cloud.links.datastore.CloudDatastoreExportEntitiesLink
   - airflow.providers.google.cloud.links.datastore.CloudDatastoreEntitiesLink
   - airflow.providers.google.common.links.storage.StorageLink