Add mirror URI column to compact manifest (#7151, PR #7576)

Author: dsotirho-ucsc

src/azul/indexer/mirror_file_service.py

@@ -164,7 +164,17 @@ def _storage(self) -> StorageService:
             bucket = aws.mirror_bucket
         return StorageService(bucket)
 
-    def get_mirror_url(self, file: File) -> str:
+    def mirror_uri(self, file: File) -> str:
+        """
+        Speculative S3 URI of the given file. No check is performed to see if
+        the file is currently mirrored, so there is no guarantee that requests
+        to the URI will succeed.
+        """
+        return str(furl(scheme='s3',
+                        netloc=self._storage.bucket_name,
+                        path=self.mirror_object_key(file)))
+
+    def mirror_url(self, file: File) -> str:
         return self._storage.get_presigned_url(key=self.mirror_object_key(file),
                                                file_name=file.name,
                                                content_type=file.content_type)

src/azul/indexer/mirror_service.py

@@ -131,19 +131,29 @@ class BaseMirrorService:
     def _queues(self) -> Queues:
         return Queues()
 
+    @classmethod
+    def may_mirror(cls, catalog: CatalogName, file_size: int = 0) -> bool:
+        """
+        Test whether it makes sense to request the mirroring of files from the
+        given catalog if they are of the given size or larger. If this method
+        returns True, such files may or may not be mirrored. If this method
+        returns False, the service will definitely refuse to mirror such files,
+        although it may accept smaller files.
+        """
+        if config.enable_mirroring:
+            max_size = config.catalogs[catalog].mirror_limit
+            return max_size is None or file_size <= max_size
+        else:
+            return False
+
     def mirror_sources(self,
                        catalog: CatalogName,
                        sources: Iterable[tuple[SourceRef, SourceConfig]]
                        ):
-        mirror_limit = config.catalogs[catalog].mirror_limit
-        if mirror_limit is not None and mirror_limit < 0:
-            log.info('Not mirroring any files in catalog %r because the file '
-                     'size limit is negative', catalog)
-        else:
-
+        if self.may_mirror(catalog):
             def messages():
-                for source, cfg in sources:
-                    if cfg.mirror:
+                for source, source_config in sources:
+                    if source_config.mirror:
                         log.info('Mirroring files in source %r from catalog %r',
                                  str(source.spec), catalog)
                         yield MirrorSourceAction(catalog=catalog, source=source)
@@ -153,6 +163,9 @@ def messages():
                                  str(source.spec), catalog)
 
             self._queue_messages(messages())
+        else:
+            log.info('Not mirroring any files in catalog %r because the file '
+                     'size limit is negative', catalog)
 
     def mirror_file(self, catalog: CatalogName, source: SourceRef, file: File):
         self._queue_messages([MirrorFileAction(catalog=catalog,
@@ -238,17 +251,16 @@ def _list_public_source_ids(self, catalog: CatalogName) -> set[str]:
     def _(self, a: MirrorPartitionAction) -> Iterable[MirrorAction]:
         plugin = self._repository_plugin(a.catalog)
         files = plugin.list_files(a.source, a.prefix)
-        max_size = config.catalogs[a.catalog].mirror_limit
         for file in files:
             assert file.size is not None, R('File size unknown', file)
-            if max_size is not None and file.size > max_size:
-                log.info('Not mirroring file to save cost: %r', file)
-            else:
+            if self.may_mirror(a.catalog, file.size):
                 log.debug('Queueing file %r', file)
                 yield MirrorFileAction(catalog=a.catalog,
                                        source=a.source,
                                        prefix=a.prefix,
                                        file=file)
+            else:
+                log.info('Not mirroring file to save cost: %r', file)
         log.info('Queued %d files in partition %r of source %r in catalog %r',
                  len(files), a.prefix, str(a.source), a.catalog)
 

src/azul/plugins/metadata/anvil/__init__.py

@@ -359,6 +359,10 @@ def recurse(mapping: MetadataPlugin._FieldMapping, path: FieldPath):
         # Above, we already configured these two fields to be omitted from the
         # manifest since they are not informative to the user.
         result[('contents', 'files')]['file_url'] = 'files.azul_url'
+        # FIXME: AnVIL uses uncommon encoding for MD5 digests
+        #        https://github.com/DataBiosphere/azul/issues/7154
+        if False:
+            result[('contents', 'files')]['file_mirror_uri'] = 'files.azul_mirror_uri'
         return result
 
     primary_keys_by_table = {

src/azul/plugins/metadata/hca/__init__.py

@@ -381,7 +381,8 @@ def manifest_config(self) -> ManifestConfig:
                 'sha256': 'file_sha256',
                 'content-type': 'file_content_type',
                 'drs_uri': 'file_drs_uri',
-                'file_url': 'file_azul_url'
+                'file_url': 'file_azul_url',
+                'file_mirror_uri': 'file_mirror_uri',
             },
             ('contents', 'cell_suspensions'): {
                 'document_id': 'cell_suspension.provenance.document_id',

src/azul/service/manifest_service.py

@@ -107,6 +107,12 @@
     FieldTypes,
     null_str,
 )
+from azul.indexer.mirror_file_service import (
+    BaseMirrorFileService,
+)
+from azul.indexer.mirror_service import (
+    BaseMirrorService,
+)
 from azul.json import (
     copy_json,
 )
@@ -814,6 +820,10 @@ def repository_plugin(self) -> RepositoryPlugin:
     def metadata_plugin(self) -> MetadataPlugin:
         return self.service.metadata_plugin(self.catalog)
 
+    @cached_property
+    def mirror_file_service(self) -> BaseMirrorFileService:
+        return BaseMirrorFileService(catalog=self.catalog)
+
     @classmethod
     @abstractmethod
     def file_name_extension(cls) -> str:
@@ -1083,7 +1093,7 @@ def convert(field_name, field_value):
             try:
                 field_type = field_types[field_name]
             except KeyError:
-                if field_name == 'file_url':
+                if field_name in ('file_url', 'file_mirror_uri'):
                     field_type = null_str
                 else:
                     raise
@@ -1145,6 +1155,13 @@ def _azul_file_url(self,
                                           fetch=False,
                                           **args))
 
+    def _azul_mirror_uri(self, file: JSON) -> str | None:
+        file = self.metadata_plugin.file_class.from_index(file)
+        if BaseMirrorService.may_mirror(self.catalog, file.size):
+            return self.mirror_file_service.mirror_uri(file)
+        else:
+            return None
+
     @cached_property
     def manifest_content_hash(self) -> int:
         log.debug('Computing content hash for manifest using filters %r ...', self.filters)
@@ -1671,7 +1688,10 @@ def write_page_to(self,
                     entities = self._get_entities(field_path, doc)
                     if field_path == ('contents', 'files'):
                         file = copy_json(one(entities))
-                        file['file_url'] = self._azul_file_url(file)
+                        if 'file_url' in column_mapping:
+                            file['file_url'] = self._azul_file_url(file)
+                        if 'file_mirror_uri' in column_mapping:
+                            file['file_mirror_uri'] = self._azul_mirror_uri(file)
                         entities = [file]
                     self._extract_fields(field_path=field_path,
                                          entities=entities,
@@ -1684,7 +1704,10 @@ def write_page_to(self,
                             for related_file in file['related_files']:
                                 related_row = {}
                                 file.update(related_file)
-                                file['file_url'] = self._azul_file_url(file)
+                                if 'file_url' in column_mapping:
+                                    file['file_url'] = self._azul_file_url(file)
+                                if 'file_mirror_uri' in column_mapping:
+                                    file['file_mirror_uri'] = self._azul_mirror_uri(file)
                                 self._extract_fields(field_path=field_path,
                                                      entities=[file],
                                                      column_mapping=column_mapping,

src/azul/service/repository_controller.py

@@ -275,7 +275,7 @@ def download_file(self,
             assert request_index == 0, request_index
             download = MirrorFileDownload(
                 file=file,
-                location=mirror_service.get_mirror_url(file),
+                location=mirror_service.mirror_url(file),
                 replica=replica,
                 token=token
             )

test/indexer/test_mirror_controller.py

@@ -227,13 +227,13 @@ def test_files_not_mirrored(self):
 
         catalog = config.catalogs[self.catalog]
 
-        def patch_max_file_size(size):
+        def patch_mirror_limit(size):
             return patch.dict(config.catalogs, {
                 self.catalog: attrs.evolve(catalog, mirror_limit=size)
             })
 
         with self.subTest(mirror_limit=-1):
-            with patch_max_file_size(-1):
+            with patch_mirror_limit(-1):
                 messages = self._mirror_sources()
                 self.assertEqual([], messages)
 
@@ -243,5 +243,5 @@ def patch_max_file_size(size):
                                    size=self._file.size + 1)
             source_message = self._test_mirror_sources()
             partition_message = self._test_mirror_source(source_message)
-            with patch_max_file_size(self._file.size):
+            with patch_mirror_limit(self._file.size):
                 self._test_mirror_partition(partition_message, [too_big, self._file])

test/integration_test.py

@@ -468,7 +468,7 @@ class Catalog:
                     # If test_mirroring is run for the catalog, ensure that the
                     # source is not flagged as no_mirror so that we can test
                     # downloading a mirrored file
-                    mirror=config.enable_mirroring and catalog.mirror_limit >= 0
+                    mirror=self.azul_client.mirror_service.may_mirror(catalog.name)
                 )
                 ma_source = self._select_source(catalog.name, public=False)
                 if ma_source is not None:
@@ -1744,7 +1744,7 @@ def _test_mirroring(self, *, delete: bool):
             catalogs = [
                 c.name
                 for c in config.catalogs.values()
-                if c.is_integration_test_catalog and c.mirror_limit >= 0
+                if c.is_integration_test_catalog and mirror_service.may_mirror(c.name)
             ]
             sources_by_catalog = {
                 catalog: [self._select_source(catalog, public=True, mirror=True)]

test/service/test_manifest.py

@@ -135,6 +135,7 @@
 )
 from service import (
     DocumentCloningTestCase,
+    MirrorTestCase,
     StorageServiceTestCase,
     WebServiceTestCase,
 )
@@ -340,7 +341,10 @@ class DCP1ManifestTestCase(ManifestTestCase, DCP1CannedBundleTestCase):
     pass
 
 
-class TestManifests(DCP1ManifestTestCase):
+class TestManifests(DCP1ManifestTestCase, MirrorTestCase):
+
+    def _mirror_uri(self, sha256: str):
+        return f's3://{self.mirror_bucket}/file/{sha256}.sha256'
 
     def run(self,
             result: Optional[unittest.result.TestResult] = None
@@ -500,6 +504,10 @@ def test_compact_manifest(self):
              self._file_url('f2b6c6f0-8d25-4aae-b255-1974cc110cfe',
                             '2018-09-14T12:33:43.720332Z')),
 
+            ('file_mirror_uri',
+             self._mirror_uri('2f6866c4ede92123f90dd15fb180fac56e33309b8fd3f4f52f263ed2f8af2f16'),
+             self._mirror_uri('3125f2f86092798b85be93fbc66f4e733e9aec0929b558589c06929627115582')),
+
             ('cell_suspension.provenance.document_id',
              '',
              '0037c9eb-8038-432f-8d9d-13ee094e54ab || aaaaaaaa-8038-432f-8d9d-13ee094e54ab'),