[a 2/2] Add mirror URI to service /index/file response (#7624)

hannes-ucsc · hannes-ucsc · commit 6a5ae2e6dd25 · 2025-12-18T16:40:34.000-08:00
Expose the file mirror service as a property of the mirror service. Consolidate conditional URI emission as a service method. Consider the file's source when deciding whether to emit a URI in the /index/file response. Previously, the source of a file was only considered when emitting the URI in manifests.
diff --git a/src/azul/__init__.py b/src/azul/__init__.py
@@ -24,11 +24,14 @@
 from typing import (
     Any,
     BinaryIO,
+    Callable,
     ClassVar,
+    Hashable,
     IO,
     Literal,
     NotRequired,
     Self,
+    TYPE_CHECKING,
     TextIO,
     TypedDict,
     final,
@@ -82,7 +85,19 @@
 
 lru_cache = functools.lru_cache
 
-cache = functools.cache
+if TYPE_CHECKING:
+    # Work around https://github.com/python/typeshed/issues/15139
+    @final
+    class CacheWrapper[_T]:
+
+        def __call__(self, *args: Hashable, **kwargs: Hashable) -> _T:
+            ...
+
+
+    def cache[_T](f: Callable[..., _T], /) -> CacheWrapper[_T]:  # noqa: E303
+        ...
+else:
+    cache = functools.cache
 
 
 def cache_per_thread(f, /):
diff --git a/src/azul/indexer/mirror_service.py b/src/azul/indexer/mirror_service.py
@@ -41,6 +41,7 @@
     SourceSpec,
 )
 from azul.indexer.mirror_file_service import (
+    BaseMirrorFileService,
     FilePart,
     MirrorFileService,
     SchemaUrlFunc,
@@ -58,6 +59,7 @@
     SourceService,
 )
 from azul.types import (
+    JSON,
     json_element_strings,
 )
 
@@ -128,6 +130,10 @@ class BaseMirrorService:
     Service for queuing mirroring work, e.g., sending action messages.
     """
 
+    @cache
+    def file_service(self, catalog: CatalogName) -> BaseMirrorFileService:
+        return BaseMirrorFileService(catalog=catalog)
+
     @cached_property
     def _queues(self) -> Queues:
         return Queues()
@@ -153,8 +159,7 @@ def may_mirror_files_from_source(self,
         else:
             return False
 
-    @classmethod
-    def may_mirror(cls, catalog: CatalogName, file_size: int = 0) -> bool:
+    def may_mirror(self, 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
@@ -168,6 +173,38 @@ def may_mirror(cls, catalog: CatalogName, file_size: int = 0) -> bool:
         else:
             return False
 
+    def mirror_uri(self,
+                   catalog: CatalogName,
+                   source: SourceSpec,
+                   file_cls: type[File],
+                   file_json: JSON
+                   ) -> str | None:
+        """
+        Return the the URI of the mirror copy of the given file from the given
+        catalog. If this method returns None, the file was not mirrored, and no
+        such URI exists. Otherwise, a mirror copy of the file may or may not
+        exist under the returned URI.
+
+        :param catalog: The catalog the file is part of
+
+        :param source: The source of the file
+
+        :param file_cls: The type of the file. This parameter is needed in order
+                         to avoid deserializing a file from a source that was
+                         configured to not be mirrored because the file metadata
+                         in that source is incomplete or broken
+
+        :param file_json: the index representation of the file
+        """
+        if self.may_mirror_files_from_source(catalog, source):
+            file = file_cls.from_index(file_json)
+            if self.may_mirror(catalog, 0 if file.size is None else file.size):
+                return self.file_service(catalog).mirror_uri(file)
+            else:
+                return None
+        else:
+            return None
+
     def mirror_sources(self,
                        catalog: CatalogName,
                        sources: Iterable[tuple[SourceRef, SourceConfig]]
@@ -217,7 +254,7 @@ class MirrorService(BaseMirrorService):
     _schema_url_func: SchemaUrlFunc
 
     @cache
-    def _file_service(self, catalog: CatalogName) -> MirrorFileService:
+    def file_service(self, catalog: CatalogName) -> MirrorFileService:
         return MirrorFileService(catalog=catalog,
                                  schema_url_func=self._schema_url_func)
 
@@ -285,7 +322,7 @@ def _(self, a: MirrorPartitionAction) -> Iterable[MirrorAction]:
     @_mirror.register
     def _(self, a: MirrorFileAction) -> Iterable[MirrorAction]:
         assert a.file.size is not None, R('File size unknown', a.file)
-        service = self._file_service(a.catalog)
+        service = self.file_service(a.catalog)
         if service.info_exists(a.file):
             log.info('File is already mirrored, skipping upload: %r', a.file)
         elif service.file_exists(a.file):
@@ -318,7 +355,7 @@ def _(self, a: MirrorFileAction) -> Iterable[MirrorAction]:
     @_mirror.register
     def _(self, a: MirrorPartAction) -> Iterable[MirrorAction]:
         log.info('Uploading part #%d of file %r', a.part.index, a.file)
-        service = self._file_service(a.catalog)
+        service = self.file_service(a.catalog)
         # Hashers are mutable so we need to make a copy
         hasher = a.hasher.copy()
         etag = service.mirror_file_part(a.file, a.part, a.upload_id, hasher)
@@ -348,7 +385,7 @@ def _(self, a: MirrorPartAction) -> Iterable[MirrorAction]:
     @_mirror.register
     def _(self, a: FinalizeFileAction) -> Iterable[MirrorAction]:
         assert len(a.etags) > 0
-        service = self._file_service(a.catalog)
+        service = self.file_service(a.catalog)
         service.finish_mirroring_file(file=a.file,
                                       upload_id=a.upload_id,
                                       etags=a.etags,
diff --git a/src/azul/plugins/metadata/anvil/service/response.py b/src/azul/plugins/metadata/anvil/service/response.py
@@ -14,6 +14,10 @@
 from azul import (
     cached_property,
 )
+from azul.indexer import (
+    SourceRef,
+    SourceSpec,
+)
 from azul.json import (
     copy_any_json,
     copy_json,
@@ -171,15 +175,18 @@ def choose_entry(_term) -> AnyMutableJSON:
         }
 
     def _make_hit(self, es_hit: JSON) -> MutableJSON:
-        sources, bundles = es_hit['sources'], es_hit['bundles']
+        contents = json_mapping(es_hit['contents'])
+        sources = json_element_mappings(es_hit['sources'])
+        bundles = json_element_mappings(es_hit['bundles'])
+        source = SourceRef[SourceSpec].from_json(one(sources)).spec
         return {
             'entryId': json_str(es_hit['entity_id']),
             # Note that there is a brittle coupling that must be maintained
             # between the `sources` and `bundles` field paths here and the
             # renamed fields in `Plugin.manifest_config`.
-            'sources': list(map(self._make_source, json_element_mappings(sources))),
-            'bundles': list(map(self._make_bundle, json_element_mappings(bundles))),
-            **self._make_contents(json_mapping(es_hit['contents']))
+            'sources': list(map(self._make_source, sources)),
+            'bundles': list(map(self._make_bundle, bundles)),
+            **self._make_contents(source, contents)
         }
 
     def _make_source(self, es_source: JSON) -> MutableJSON:
@@ -199,26 +206,31 @@ def _make_bundle(self, es_bundle: JSON) -> MutableJSON:
             self._special_fields.bundle_version.name_in_hit: json_str(es_bundle['version'])
         }
 
-    def _make_contents(self, es_contents: JSON) -> MutableJSON:
+    def _make_contents(self, source: SourceSpec, es_contents: JSON) -> MutableJSON:
         return {
             inner_entity_type: (
-                [self._pivotal_entity(inner_entity_type, json_mapping(one(inner_entities)))]
+                [
+                    self._pivotal_entity(source,
+                                         inner_entity_type,
+                                         json_mapping(one(inner_entities)))
+                ]
                 if inner_entity_type == self.entity_type else
                 list(map(partial(self._non_pivotal_entity, inner_entity_type), inner_entities))
             )
             for inner_entity_type, inner_entities in json_item_sequences(es_contents)
         }
 
     def _pivotal_entity(self,
+                        source: SourceSpec,
                         inner_entity_type: str,
-                        inner_entity: JSON
+                        inner_entity: JSON,
                         ) -> MutableJSON:
         inner_entity = copy_json(inner_entity)
         if inner_entity_type == 'files':
             inner_entity['azul_url'] = self._file_url(uuid=json_str(inner_entity['document_id']),
                                                       version=json_str(inner_entity['version']),
                                                       drs_uri=optional(json_str, inner_entity['drs_uri']))
-            inner_entity['azul_mirror_uri'] = self._file_mirror_uri(inner_entity)
+            inner_entity['azul_mirror_uri'] = self._file_mirror_uri(source, inner_entity)
             inner_entity.pop('version', None)
         return inner_entity
 
diff --git a/src/azul/plugins/metadata/hca/service/response.py b/src/azul/plugins/metadata/hca/service/response.py
@@ -21,6 +21,10 @@
 from azul import (
     cached_property,
 )
+from azul.indexer import (
+    SourceRef,
+    SourceSpec,
+)
 from azul.plugins import (
     SpecialFields,
 )
@@ -345,7 +349,7 @@ def make_dates(self, entry) -> MutableJSONs:
             for dates in entry['contents']['dates']
         ]
 
-    def make_projects(self, entry) -> MutableJSONs:
+    def make_projects(self, source: SourceSpec, entry) -> MutableJSONs:
         projects = []
         contents = entry['contents']
         for project in contents['projects']:
@@ -374,34 +378,36 @@ def make_projects(self, entry) -> MutableJSONs:
                     for key in list(publication.keys()):
                         publication[to_camel_case(key)] = publication.pop(key)
                 translated_project['supplementaryLinks'] = project.get('supplementary_links', [None])
-                translated_project['matrices'] = self.make_matrices_(contents['matrices'])
-                translated_project['contributedAnalyses'] = self.make_matrices_(contents['contributed_analyses'])
+                translated_project['matrices'] = self.make_matrices_(source,
+                                                                     contents['matrices'])
+                translated_project['contributedAnalyses'] = self.make_matrices_(source,
+                                                                                contents['contributed_analyses'])
                 translated_project['accessions'] = project.get('accessions', [None])
             projects.append(translated_project)
         return projects
 
     # FIXME: Move this to during aggregation
     #        https://github.com/DataBiosphere/azul/issues/2415
 
-    def make_matrices_(self, matrices: JSONs) -> JSON:
+    def make_matrices_(self, source: SourceSpec, matrices: JSONs) -> JSON:
         files: list[JSON] = []
         if matrices:
             for file in json_element_mappings(one(matrices)['file']):
                 translated_file = {
-                    **self.make_file(file),
+                    **self.make_file(source, file),
                     'strata': json_str(file['strata'])
                 }
                 files.append(translated_file)
         return make_stratification_tree(files)
 
-    def make_files(self, entry: JSON) -> JSONs:
+    def make_files(self, source: SourceSpec, entry: JSON) -> JSONs:
         files = []
         for _file in json_element_mappings(json_mapping(entry['contents'])['files']):
-            translated_file = self.make_file(_file)
+            translated_file = self.make_file(source, _file)
             files.append(translated_file)
         return files
 
-    def make_file(self, file: JSON) -> JSON:
+    def make_file(self, source: SourceSpec, file: JSON) -> JSON:
         translated_file = {
             'contentDescription': file.get('content_description'),
             'format': file.get('file_format'),
@@ -417,7 +423,7 @@ def make_file(self, file: JSON) -> JSON:
             'azul_url': self._file_url(uuid=json_str(file['uuid']),
                                        version=json_str(file['version']),
                                        drs_uri=optional(json_str, file['drs_uri'])),
-            'azul_mirror_uri': self._file_mirror_uri(file),
+            'azul_mirror_uri': self._file_mirror_uri(source, file),
         }
         return translated_file
 
@@ -511,10 +517,11 @@ def make_hits(self, hits: JSONs) -> list[SummarizedHit | CompleteHit]:
         return list(map(self.make_hit, hits))
 
     def make_hit(self, es_hit) -> SummarizedHit | CompleteHit:
+        source: SourceSpec = SourceRef.from_json(one(es_hit['sources'])).spec
         hit = Hit(protocols=self.make_protocols(es_hit),
                   entryId=es_hit['entity_id'],
                   sources=self.make_sources(es_hit),
-                  projects=self.make_projects(es_hit),
+                  projects=self.make_projects(source, es_hit),
                   samples=self.make_samples(es_hit),
                   specimens=self.make_specimens(es_hit),
                   cellLines=self.make_cell_lines(es_hit),
@@ -525,7 +532,7 @@ def make_hit(self, es_hit) -> SummarizedHit | CompleteHit:
         if self.entity_type in ('files', 'bundles'):
             complete_hit = cast(CompleteHit, hit)
             complete_hit['bundles'] = self.make_bundles(es_hit)
-            complete_hit['files'] = self.make_files(es_hit)
+            complete_hit['files'] = self.make_files(source, es_hit)
             return complete_hit
         else:
             summarized_hit = cast(SummarizedHit, hit)
diff --git a/src/azul/service/manifest_service.py b/src/azul/service/manifest_service.py
@@ -109,9 +109,6 @@
     FieldTypes,
     null_str,
 )
-from azul.indexer.mirror_file_service import (
-    BaseMirrorFileService,
-)
 from azul.indexer.mirror_service import (
     BaseMirrorService,
 )
@@ -827,10 +824,6 @@ def metadata_plugin(self) -> MetadataPlugin:
     def mirror_service(self) -> BaseMirrorService:
         return BaseMirrorService()
 
-    @cached_property
-    def mirror_file_service(self) -> BaseMirrorFileService:
-        return BaseMirrorFileService(catalog=self.catalog)
-
     @classmethod
     @abstractmethod
     def file_name_extension(cls) -> str:
@@ -1167,15 +1160,9 @@ def _azul_file_url(self,
                                           fetch=False,
                                           **args))
 
-    def _azul_mirror_uri(self, file: JSON, source: SourceSpec) -> str | None:
-        if self.mirror_service.may_mirror_files_from_source(self.catalog, source):
-            file = self.metadata_plugin.file_class.from_index(file)
-            if self.mirror_service.may_mirror(self.catalog, file.size):
-                return self.mirror_file_service.mirror_uri(file)
-            else:
-                return None
-        else:
-            return None
+    def _azul_mirror_uri(self, source: SourceSpec, file: JSON) -> str | None:
+        file_cls = self.metadata_plugin.file_class
+        return self.mirror_service.mirror_uri(self.catalog, source, file_cls, file)
 
     @cached_property
     def manifest_content_hash(self) -> int:
@@ -1707,7 +1694,7 @@ def write_page_to(self,
                         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, source)
+                            file['file_mirror_uri'] = self._azul_mirror_uri(source, file)
                         entities = [file]
                     self._extract_fields(field_path=field_path,
                                          entities=entities,
@@ -1723,7 +1710,7 @@ def write_page_to(self,
                                 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, source)
+                                    file['file_mirror_uri'] = self._azul_mirror_uri(source, file)
                                 self._extract_fields(field_path=field_path,
                                                      entities=[file],
                                                      column_mapping=column_mapping,
diff --git a/src/azul/service/repository_controller.py b/src/azul/service/repository_controller.py
@@ -44,7 +44,6 @@
     pass_thru_bool,
 )
 from azul.indexer.mirror_file_service import (
-    BaseMirrorFileService,
     MirrorFileDownload,
 )
 from azul.plugins import (
@@ -85,9 +84,6 @@ class RepositoryController(ServiceAppController):
     def service(self) -> RepositoryService:
         return RepositoryService()
 
-    def mirror_service(self, catalog: CatalogName) -> BaseMirrorFileService:
-        return self.service.mirror_service(catalog)
-
     def repository_plugin(self, catalog: CatalogName) -> RepositoryPlugin:
         return self.service.repository_plugin(catalog)
 
@@ -262,7 +258,7 @@ def download_file(self,
         plugin = self.repository_plugin(catalog)
 
         if config.enable_mirroring:
-            mirror_service = self.mirror_service(catalog)
+            mirror_service = self.service.mirror_service.file_service(catalog)
             is_mirrored = mirror_service.info_exists(file)
         else:
             mirror_service, is_mirrored = None, False
diff --git a/src/azul/service/repository_service.py b/src/azul/service/repository_service.py
diff --git a/test/indexer/test_mirror_controller.py b/test/indexer/test_mirror_controller.py
