Fix: Verbatim PFB handover fails when using a common prefix (#6843)

Author: nadove-ucsc

src/azul/service/manifest_service.py

@@ -102,6 +102,9 @@
 from azul.deployment import (
     aws,
 )
+from azul.indexer import (
+    SourceSpec,
+)
 from azul.indexer.document import (
     DocumentType,
     FieldPath,
@@ -2051,6 +2054,9 @@ class ReplicaKeys:
         """
         hub_id: str
         replica_ids: list[str]
+        # We need to track whether sources use a common prefix to avoid dangling
+        # references in PFB relations
+        source_spec: str
 
     def _replica_keys(self) -> Iterable[ReplicaKeys]:
         request = self._create_request()
@@ -2065,14 +2071,16 @@ def _replica_keys(self) -> Iterable[ReplicaKeys]:
                 for document_id in always_iterable(inner_entity['document_id'])
             ]
             yield self.ReplicaKeys(hub_id=hit['entity_id'],
-                                   replica_ids=document_ids)
+                                   replica_ids=document_ids,
+                                   source_spec=one(hit['sources'])['spec'])
 
-    def _all_replicas(self) -> Iterable[JSON]:
+    def _all_replicas(self) -> Iterable[tuple[JSON, set[str]]]:
         emitted_replica_ids = set()
         page_size = 100
         for page in chunked(self._replica_keys(), page_size):
             num_replicas = 0
             num_new_replicas = 0
+            sources = {keys.source_spec for keys in page}
             for replica in self._join_replicas(page):
                 num_replicas += 1
                 # A single replica may have many hubs. To prevent replicas from
@@ -2081,7 +2089,7 @@ def _all_replicas(self) -> Iterable[JSON]:
                 replica_id = replica.meta.id
                 if replica_id not in emitted_replica_ids:
                     num_new_replicas += 1
-                    yield replica.to_dict()
+                    yield replica.to_dict(), sources
                     emitted_replica_ids.add(replica_id)
             log.info('Found %d replicas (%d already emitted) from page of %d hubs',
                      num_replicas, num_replicas - num_new_replicas, len(page))
@@ -2119,7 +2127,7 @@ def create_file(self) -> tuple[str, str | None]:
         fd, path = mkstemp(suffix=f'.{self.file_name_extension()}')
         os.close(fd)
         with open(path, 'w') as f:
-            for replica in self._all_replicas():
+            for replica, _ in self._all_replicas():
                 entry = {
                     'value': replica['contents'],
                     'type': replica['replica_type']
@@ -2143,7 +2151,14 @@ def file_name_extension(cls):
     def format(cls) -> ManifestFormat:
         return ManifestFormat.verbatim_pfb
 
-    def _include_relations(self) -> bool:
+    @property
+    def included_fields(self) -> list[FieldPath]:
+        return [
+            *super().included_fields,
+            ('sources',)
+        ]
+
+    def _include_relations(self, sources: Iterable[str]) -> bool:
         # Terra will reject the handover if the manifest includes
         # dangling relations, i.e., if any entity references another
         # entity that isn't included in the manifest. There are three
@@ -2171,14 +2186,17 @@ def _include_relations(self) -> bool:
         #    See https://github.com/DataBiosphere/azul/issues/4440
         #
         # (1) can only occur when orphans are included, and (2) and (3)
-        # can only occur when orphans are *not* included. Because (1)
-        # can only occur on lower deployments, we make the inclusion of
-        # relations conditional on avoiding (2) and (3).
+        # can only occur when orphans are *not* included.
         #
-        return config.enable_verbatim_relations and self.include_orphans
+        complete = not any(
+            SourceSpec.parse_prefix(source)[1].common
+            for source in sources
+        )
+        return config.enable_verbatim_relations and self.include_orphans and complete
 
     def create_file(self) -> tuple[str, str | None]:
-        replicas = list(self._all_replicas())
+        replicas, sources = zip(*self._all_replicas())
+        sources = set().union(*sources)
         plugin = self.metadata_plugin
         replica_schemas = plugin.verbatim_pfb_schema(replicas)
         # Ensure field order is consistent for unit tests
@@ -2192,10 +2210,11 @@ def create_file(self) -> tuple[str, str | None]:
 
         def pfb_entities():
             yield pfb_metadata_entity
+            include_relations = self._include_relations(sources)
             for replica in replicas:
                 id = plugin.verbatim_pfb_entity_id(replica)
                 entity = avro_pfb.PFBEntity.for_replica(id, dict(replica))
-                if self._include_relations():
+                if include_relations:
                     relations = plugin.verbatim_pfb_relations(replica)
                     entity_relations = [
                         PFBRelation(dst_name=replica_type, dst_id=entity_id)