fixup! Implement mirror_partition (#6861)

Author: nadove-ucsc

src/azul/plugins/__init__.py

@@ -676,7 +676,8 @@ def count_bundles(self, source: SOURCE_SPEC) -> int:
     def count_files(self, source: SOURCE_SPEC) -> int:
         """
         The total number of files in the given source. The source's prefix
-        may be None.
+        may be None, indicating that the source hasn't been partitioned yet and
+        that this method should count all files in the source.
         """
         raise NotImplementedError
 

src/azul/plugins/repository/canned/__init__.py

@@ -176,12 +176,16 @@ def count_bundles(self, source: CannedSourceRef) -> int:
 
     def count_files(self, source: SimpleSourceSpec) -> int:
         staging_area = self.staging_area(source.name)
-        prefix = '' if source.prefix is None else source.prefix.common
-        return sum(
-            1
-            for descriptor in staging_area.descriptors.values()
-            if descriptor.content['sha256'].startswith(prefix)
-        )
+        if source.prefix is None:
+            return len(staging_area.descriptors)
+        else:
+            prefix = source.prefix.common
+            assert not any(map(str.isupper, prefix)), source
+            return sum(
+                1
+                for descriptor in staging_area.descriptors.values()
+                if descriptor.content['sha256'].lower().startswith(prefix)
+            )
 
     def list_bundles(self,
                      source: CannedSourceRef,
@@ -220,14 +224,15 @@ def fetch_bundle(self, bundle_fqid: CannedBundleFQID) -> CannedBundle:
     def list_files(self, source: CannedSourceRef, prefix: str) -> list[HCAFile]:
         self._assert_source(source)
         self._assert_partition(source, prefix)
+        assert not any(map(str.isupper, prefix)), prefix
         staging_area = self.staging_area(source.spec.name)
         return [
             HCAFile.from_descriptor(descriptor.content,
                                     uuid=file_uuid,
                                     name=descriptor.content['file_name'],
                                     drs_uri=None)
             for file_uuid, descriptor in staging_area.descriptors.items()
-            if descriptor.content['sha256'].startswith(prefix)
+            if descriptor.content['sha256'].lower().startswith(prefix)
         ]
 
     def _construct_file_url(self, url: furl, file_name: str) -> furl:

src/azul/plugins/repository/tdr_anvil/__init__.py

@@ -245,12 +245,16 @@ def _batch_uuid(self,
                               self.bundle_uuid_version)
 
     def count_files(self, source: TDRSourceSpec) -> int:
+        if source.prefix is None:
+            prefix = ''
+        else:
+            prefix = source.prefix.common
+            assert not any(map(str.isupper, prefix)), source
         query = f'''
         SELECT COUNT(*) AS count
         FROM {backtick(self._full_table_name(source, 'anvil_file'))}
-        ''' + ('' if source.prefix is None else f'''
-        WHERE starts_with(file_md5sum, {source.prefix.common!r})
-        ''')
+        WHERE STARTS_WITH(LOWER(file_md5sum), {prefix!r})
+        '''
         return one(self._run_sql(query))['count']
 
     def count_bundles(self, source: TDRSourceSpec) -> int:
@@ -334,7 +338,7 @@ def list_files(self, source: TDRSourceRef, prefix: str) -> list[AnvilFile]:
         batch = self._get_batch(source.spec,
                                 'anvil_file',
                                 prefix,
-                                id_column='file_md5sum')
+                                key_column='file_md5sum')
         return [
             AnvilFile(uuid=ref.entity_id,
                       name=row['file_name'],
@@ -562,13 +566,14 @@ def _get_batch(self,
                    table_name: str,
                    batch_prefix: str,
                    *,
-                   id_column: str
+                   key_column: str
                    ) -> Iterable[tuple[EntityReference, BigQueryRow]]:
+        assert batch_prefix.islower(), batch_prefix
         columns = self._columns(table_name)
         for row in self._run_sql(f'''
             SELECT {', '.join(sorted(columns))}
             FROM {backtick(self._full_table_name(source, table_name))}
-            WHERE STARTS_WITH(LOWER({id_column}), {batch_prefix!r})
+            WHERE STARTS_WITH(LOWER({key_column}), {batch_prefix!r})
         '''):
             ref = EntityReference(entity_type=table_name, entity_id=row['datarepo_row_id'])
             yield ref, row
@@ -579,7 +584,7 @@ def _get_bundle_batch(self,
         return self._get_batch(bundle_fqid.source.spec,
                                bundle_fqid.table_name,
                                bundle_fqid.batch_prefix,
-                               id_column='datarepo_row_id')
+                               key_column='datarepo_row_id')
 
     def _bundle_entity(self, bundle_fqid: TDRAnvilBundleFQID) -> KeyReference:
         source = bundle_fqid.source

src/azul/plugins/repository/tdr_hca/__init__.py

@@ -243,13 +243,18 @@ def count_bundles(self, source: TDRSourceSpec) -> int:
         return one(rows)['count']
 
     def count_files(self, source: TDRSourceSpec) -> int:
+        if source.prefix is None:
+            prefix = ''
+        else:
+            prefix = source.prefix.common
+            assert not any(map(str.isupper, prefix)), source
         query = ' UNION ALL '.join(
             f'''
             SELECT COUNT(*) AS count
             FROM {backtick(self._full_table_name(source, entity_type))}
-            ''' + ('' if source.prefix is None else f'''
-            WHERE STARTS_WITH(JSON_EXTRACT_SCALAR(descriptor, "$.sha256"), {source.prefix.common!r})
-            ''')
+            WHERE STARTS_WITH(LOWER(JSON_EXTRACT_SCALAR(descriptor, "$.sha256")),
+                              {prefix!r})
+            '''
             for entity_type, entity_cls in api.entity_types.items()
             if entity_type.endswith('_file')
         )
@@ -278,11 +283,13 @@ def list_bundles(self,
     def list_files(self, source: TDRSourceRef, prefix: str) -> list[HCAFile]:
         self._assert_source(source)
         self._assert_partition(source, prefix)
+        assert not any(map(str.isupper, prefix)), prefix
         rows = self._run_sql(' UNION ALL '.join(
             f'''
             SELECT {', '.join(TDRHCABundle.data_columns)}
             FROM {backtick(self._full_table_name(source.spec, entity_type))}
-            WHERE STARTS_WITH(JSON_EXTRACT_SCALAR(descriptor, "$.sha256"), {prefix!r})
+            WHERE STARTS_WITH(LOWER(JSON_EXTRACT_SCALAR(descriptor, "$.sha256")),
+                              {prefix!r})
             '''
             for entity_type, entity_cls in api.entity_types.items()
             if entity_type.endswith('_file')