Eliminate azul.terra.TDRClient.TDRSource (#5524)

DataBiosphere · Sep 11, 2024 · 0f9ca1a · 0f9ca1a
1 parent 9870c49
commit 0f9ca1a
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 19 deletions.
diff --git a/scripts/post_deploy_tdr.py b/scripts/post_deploy_tdr.py
@@ -25,7 +25,6 @@
 )
 from azul.terra import (
     TDRClient,
-    TDRSourceRef,
     TDRSourceSpec,
 )
 
@@ -91,20 +90,9 @@ def verify_source(self,
                       catalog: CatalogName,
                       source_spec: TDRSourceSpec
                       ) -> None:
-        source = self.tdr.lookup_source(source_spec)
-        log.info('TDR client is authorized for API access to %s.', source_spec)
-        require(source.project == source_spec.subdomain,
-                'Actual Google project of TDR source differs from configured one',
-                source.project, source_spec.subdomain)
-        # Uppercase is standard for multi-regions in the documentation but TDR
-        # returns 'us' in lowercase
-        require(source.location.lower() == config.tdr_source_location.lower(),
-                'Actual storage location of TDR source differs from configured one',
-                source.location, config.tdr_source_location)
-        # FIXME: Eliminate azul.terra.TDRClient.TDRSource
-        #        https://github.com/DataBiosphere/azul/issues/5524
-        ref = TDRSourceRef(id=source.id, spec=source_spec)
         plugin = self.repository_plugin(catalog)
+        ref = plugin.resolve_source(str(source_spec))
+        log.info('TDR client is authorized for API access to %s.', source_spec)
         subgraph_count = sum(plugin.list_partitions(ref).values())
         require(subgraph_count > 0,
                 'Source spec is empty (bad prefix?)', source_spec)

diff --git a/src/azul/plugins/repository/tdr.py b/src/azul/plugins/repository/tdr.py
@@ -193,7 +193,7 @@ def _drs_client(cls,
         return cls._user_authenticated_tdr(authentication).drs_client()
 
     def _lookup_source_id(self, spec: TDRSourceSpec) -> str:
-        return self.tdr.lookup_source(spec).id
+        return self.tdr.lookup_source(spec)
 
     def list_bundles(self,
                      source: TDRSourceRef,

diff --git a/src/azul/terra.py b/src/azul/terra.py
@@ -409,16 +409,29 @@ class TDRSource:
         location: str
 
     @cache
-    def lookup_source(self, source_spec: TDRSourceSpec) -> TDRSource:
+    def lookup_source(self, source_spec: TDRSourceSpec) -> str:
+        """
+        Validate that the repository's reported values for the snapshot's Google
+        project name and storage location match our expectations, and return the
+        snapshot's UUID.
+        """
         source = self._lookup_source(source_spec)
+        actual_project = source['dataProject']
+        require(actual_project == source_spec.subdomain,
+                'Actual Google project of TDR source differs from configured one',
+                actual_project, source_spec.subdomain)
         storage = one(
             resource
             for resource in source['storage']
             if resource['cloudResource'] == 'bigquery'
         )
-        return self.TDRSource(project=source['dataProject'],
-                              id=source['id'],
-                              location=storage['region'])
+        actual_location = storage['region']
+        # Uppercase is standard for multi-regions in the documentation but TDR
+        # returns 'us' in lowercase
+        require(actual_location.lower() == config.tdr_source_location.lower(),
+                'Actual storage location of TDR source differs from configured one',
+                actual_location, config.tdr_source_location)
+        return source['id']
 
     def _retrieve_source(self, source: TDRSourceRef) -> MutableJSON:
         endpoint = self._repository_endpoint('snapshots', source.id)