Refactor list_bundles methods

DataBiosphere · Nov 7, 2024 · ec2e7a2 · ec2e7a2
1 parent 96747c3
commit ec2e7a2
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 41 deletions.
diff --git a/src/azul/azulclient.py b/src/azul/azulclient.py
@@ -225,7 +225,11 @@ def list_bundles(self,
             source = plugin.resolve_source(source)
         else:
             assert isinstance(source, SourceRef), source
-        return plugin.list_bundles(source, prefix)
+        log.info('Listing bundles with prefix %r in source %r.', prefix, source)
+        bundle_fqids = plugin.list_bundles(source, prefix)
+        log.info('There are %i bundle(s) with prefix %r in source %r.',
+                 len(bundle_fqids), prefix, source)
+        return bundle_fqids
 
     @property
     def sqs(self):

diff --git a/src/azul/plugins/repository/canned/__init__.py b/src/azul/plugins/repository/canned/__init__.py
@@ -64,9 +64,6 @@
 from azul.types import (
     JSON,
 )
-from azul.uuids import (
-    validate_uuid_prefix,
-)
 from humancellatlas.data.metadata.helpers.staging_area import (
     CannedStagingAreaFactory,
     StagingArea,
@@ -179,18 +176,14 @@ def list_bundles(self,
                      prefix: str
                      ) -> list[CannedBundleFQID]:
         self._assert_source(source)
-        validate_uuid_prefix(prefix)
-        log.info('Listing bundles with prefix %r in source %r.', prefix, source)
-        bundle_fqids = []
         staging_area = self.staging_area(source.spec.name)
-        for link in staging_area.links.values():
-            if link.uuid.startswith(prefix):
-                bundle_fqids.append(CannedBundleFQID(source=source,
-                                                     uuid=link.uuid,
-                                                     version=link.version))
-        log.info('There are %i bundle(s) with prefix %r in source %r.',
-                 len(bundle_fqids), prefix, source)
-        return bundle_fqids
+        return [
+            CannedBundleFQID(source=source,
+                             uuid=link.uuid,
+                             version=link.version)
+            for link in staging_area.links.values()
+            if link.uuid.startswith(prefix)
+        ]
 
     def fetch_bundle(self, bundle_fqid: CannedBundleFQID) -> CannedBundle:
         self._assert_source(bundle_fqid.source)

diff --git a/src/azul/plugins/repository/tdr.py b/src/azul/plugins/repository/tdr.py
@@ -191,17 +191,6 @@ def _drs_client(cls,
     def _lookup_source_id(self, spec: TDRSourceSpec) -> str:
         return self.tdr.lookup_source(spec)
 
-    def list_bundles(self,
-                     source: TDRSourceRef,
-                     prefix: str
-                     ) -> list[TDRBundleFQID]:
-        self._assert_source(source)
-        log.info('Listing bundles with prefix %r in source %r.', prefix, source)
-        bundle_fqids = self._list_bundles(source, prefix)
-        log.info('There are %i bundle(s) with prefix %r in source %r.',
-                 len(bundle_fqids), prefix, source)
-        return bundle_fqids
-
     def fetch_bundle(self, bundle_fqid: TDRBundleFQID) -> TDR_BUNDLE:
         self._assert_source(bundle_fqid.source)
         now = time.time()
@@ -223,13 +212,6 @@ def _run_sql(self, query):
     def _full_table_name(self, source: TDRSourceSpec, table_name: str) -> str:
         return source.qualify_table(table_name)
 
-    @abstractmethod
-    def _list_bundles(self,
-                      source: TDRSourceRef,
-                      prefix: str
-                      ) -> list[TDRBundleFQID]:
-        raise NotImplementedError
-
     @abstractmethod
     def _emulate_bundle(self, bundle_fqid: TDRBundleFQID) -> TDR_BUNDLE:
         raise NotImplementedError

diff --git a/src/azul/plugins/repository/tdr_anvil/__init__.py b/src/azul/plugins/repository/tdr_anvil/__init__.py
@@ -182,10 +182,11 @@ def count_bundles(self, source: TDRSourceSpec) -> int:
         ''')
         return sum(row['count'] for row in rows)
 
-    def _list_bundles(self,
-                      source: TDRSourceRef,
-                      prefix: str
-                      ) -> list[TDRAnvilBundleFQID]:
+    def list_bundles(self,
+                     source: TDRSourceRef,
+                     prefix: str
+                     ) -> list[TDRAnvilBundleFQID]:
+        self._assert_source(source)
         spec = source.spec
         primary = BundleType.primary.value
         supplementary = BundleType.supplementary.value

diff --git a/src/azul/plugins/repository/tdr_hca/__init__.py b/src/azul/plugins/repository/tdr_hca/__init__.py
@@ -288,10 +288,11 @@ def count_bundles(self, source: TDRSourceSpec) -> int:
         rows = self._run_sql(query)
         return one(rows)['count']
 
-    def _list_bundles(self,
-                      source: TDRSourceRef,
-                      prefix: str
-                      ) -> list[TDRBundleFQID]:
+    def list_bundles(self,
+                     source: TDRSourceRef,
+                     prefix: str
+                     ) -> list[TDRBundleFQID]:
+        self._assert_source(source)
         current_bundles = self._query_unique_sorted(f'''
             SELECT links_id, version
             FROM {backtick(self._full_table_name(source.spec, 'links'))}