Change CloudDatastoreExportEntitiesLink to StorageLink

apache · Apr 25, 2022 · 544d658 · 544d658
1 parent b3cc2f5
commit 544d658
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 29 deletions.
diff --git a/airflow/providers/google/cloud/links/datastore.py b/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"""
 

diff --git a/airflow/providers/google/cloud/operators/datastore.py b/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
 

diff --git a/airflow/providers/google/provider.yaml b/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