Implement mirror_partition (#6861)

Author: nadove-ucsc

src/azul/azulclient.py

@@ -68,6 +68,7 @@
     Serializable,
 )
 from azul.plugins import (
+    File,
     RepositoryPlugin,
 )
 from azul.queues import (
@@ -112,6 +113,7 @@ class IndexAction(Action):
 class MirrorAction(Action):
     mirror_source = auto()
     mirror_partition = auto()
+    mirror_file = auto()
 
 
 @attrs.frozen(kw_only=True)
@@ -195,6 +197,21 @@ def mirror_partition_message(self,
             group_id=f'{source.id}:{prefix}'
         )
 
+    def mirror_file_message(self,
+                            catalog: CatalogName,
+                            source: SourceRef,
+                            file: File,
+                            ) -> SQSFifoMessage:
+        return SQSFifoMessage(
+            body={
+                'action': MirrorAction.mirror_file.to_json(),
+                'catalog': catalog,
+                'source': cast(JSON, source.to_json()),
+                'file': file.to_json()
+            },
+            group_id=f'{source.id}:{file.uuid}'
+        )
+
     def local_reindex(self, catalog: CatalogName, prefix: str) -> int:
         notifications = [
             self.notification(bundle_fqid)
@@ -607,7 +624,7 @@ def message(source: SourceRef):
     def mirror_source(self, catalog: CatalogName, source_json: JSON):
         plugin = self.repository_plugin(catalog)
         source = plugin.source_ref_cls.from_json(source_json)
-        source = plugin.partition_source_for_indexing(catalog, source)
+        source = plugin.partition_source_for_mirroring(catalog, source)
 
         def message(prefix: str) -> SQSMessage:
             log.info('Mirroring files in partition %r of source %r from catalog %r',
@@ -617,6 +634,18 @@ def message(prefix: str) -> SQSMessage:
         messages = map(message, source.spec.prefix.partition_prefixes())
         self.queue_mirror_messages(messages)
 
+    def mirror_partition(self, catalog: CatalogName, source_json: JSON, prefix: str):
+        plugin = self.repository_plugin(catalog)
+        source = plugin.source_ref_cls.from_json(source_json)
+
+        def message(file: File) -> SQSMessage:
+            log.info('Mirroring file %r in source %r from catalog %r',
+                     file.uuid, str(source.spec), catalog)
+            return self.mirror_file_message(catalog, source, file)
+
+        messages = map(message, plugin.list_files(source, prefix))
+        self.queue_mirror_messages(messages)
+
 
 class AzulClientError(RuntimeError):
     pass

src/azul/indexer/mirror_controller.py

@@ -20,6 +20,7 @@
 )
 from azul.types import (
     JSON,
+    json_mapping,
     json_str,
 )
 
@@ -40,11 +41,14 @@ def _mirror(self, message: JSON):
         if action is MirrorAction.mirror_source:
             self.client.mirror_source(message['catalog'], message['source'])
         elif action is MirrorAction.mirror_partition:
-            # FIXME: Implement mirror_partition
-            #        https://github.com/DataBiosphere/azul/issues/6861
-            log.info('Would mirror files in partition %r of source %r',
-                     message['prefix'], message['source'])
-            time.sleep(10)
+            self.client.mirror_partition(message['catalog'],
+                                         message['source'],
+                                         message['prefix'])
+        elif action is MirrorAction.mirror_file:
+            # FIXME: Implement mirror_file, mirror_part & finalize_file
+            #        https://github.com/DataBiosphere/azul/issues/6862
+            log.info('Would mirror parts of file %r', json_mapping(message['file'])['uuid'])
+            time.sleep(1)
         else:
             # FIXME: Implement mirror_file, mirror_part & finalize_file
             #        https://github.com/DataBiosphere/azul/issues/6862

src/azul/plugins/__init__.py

@@ -672,6 +672,14 @@ def count_bundles(self, source: SOURCE_SPEC) -> int:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def count_files(self, source: SOURCE_SPEC) -> int:
+        """
+        The total number of files in the given source. The source's prefix
+        may be None.
+        """
+        raise NotImplementedError
+
     def partition_source_for_indexing(self,
                                       catalog: CatalogName,
                                       source: SOURCE_REF
@@ -683,6 +691,17 @@ def partition_source_for_indexing(self,
         """
         return self._partition_source(catalog, source, self.count_bundles)
 
+    def partition_source_for_mirroring(self,
+                                       catalog: CatalogName,
+                                       source: SOURCE_REF
+                                       ) -> SOURCE_REF:
+        """
+        If the source already has a prefix, return the source. Otherwise, return
+        an updated copy of the source with a heuristically computed prefix that
+        should be appropriate for mirroring in the given catalog.
+        """
+        return self._partition_source(catalog, source, self.count_files)
+
     def _partition_source(self,
                           catalog: CatalogName,
                           source: SOURCE_REF,
@@ -730,6 +749,20 @@ def fetch_bundle(self, bundle_fqid: BUNDLE_FQID) -> BUNDLE:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def list_files(self, source: SOURCE_REF, prefix: str) -> list['File']:
+        """
+        List the files in the given source whose digest value starts with the
+        given prefix.
+
+        :param source: A reference to the repository source that contains the
+                       files to list
+
+        :param prefix: A string of lower-case hexadecimal characters
+        """
+
+        raise NotImplementedError
+
     @abstractmethod
     def drs_client(self,
                    authentication: Authentication | None = None

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

@@ -50,6 +50,9 @@
     RepositoryFileDownload,
     RepositoryPlugin,
 )
+from azul.plugins.metadata.hca import (
+    HCAFile,
+)
 from azul.plugins.metadata.hca.bundle import (
     HCABundle,
 )
@@ -171,6 +174,15 @@ def count_bundles(self, source: CannedSourceRef) -> int:
                 if links_id.lower().startswith(prefix)
             )
 
+    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)
+        )
+
     def list_bundles(self,
                      source: CannedSourceRef,
                      prefix: str
@@ -205,6 +217,19 @@ def fetch_bundle(self, bundle_fqid: CannedBundleFQID) -> CannedBundle:
                  time.time() - now, bundle.uuid, bundle.version)
         return bundle
 
+    def list_files(self, source: CannedSourceRef, prefix: str) -> list[HCAFile]:
+        self._assert_source(source)
+        self._assert_partition(source, 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)
+        ]
+
     def _construct_file_url(self, url: furl, file_name: str) -> furl:
         """
         >>> plugin = Plugin(_sources=set())

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

@@ -119,6 +119,9 @@ def _lookup_source_id(self, spec: SimpleSourceSpec) -> str:
     def count_bundles(self, source: SimpleSourceSpec) -> NoReturn:
         assert False, 'DSS is EOL'
 
+    def count_files(self, source: SimpleSourceSpec) -> NoReturn:
+        assert False, 'DSS is EOL'
+
     def list_sources(self,
                      authentication: Authentication | None
                      ) -> list[DSSSourceRef]:
@@ -136,6 +139,9 @@ def list_bundles(self,
     def fetch_bundle(self, bundle_fqid: DSSBundleFQID) -> NoReturn:
         assert False, 'DSS is EOL'
 
+    def list_files(self, source: DSSSourceRef, prefix: str) -> NoReturn:
+        assert False, 'DSS is EOL'
+
     def dss_subscription_query(self, prefix: str) -> JSON:
         return {
             "query": {

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

@@ -42,6 +42,9 @@
     EntityReference,
     EntityType,
 )
+from azul.plugins.metadata.anvil import (
+    AnvilFile,
+)
 from azul.plugins.metadata.anvil.bundle import (
     AnvilBundle,
     EntityLink,
@@ -241,6 +244,15 @@ def _batch_uuid(self,
                               self.batch_uuid_version,
                               self.bundle_uuid_version)
 
+    def count_files(self, source: TDRSourceSpec) -> int:
+        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})
+        ''')
+        return one(self._run_sql(query))['count']
+
     def count_bundles(self, source: TDRSourceSpec) -> int:
         prefix = '' if source.prefix is None else source.prefix.common
         assert not any(map(str.isupper, prefix)), source
@@ -316,6 +328,23 @@ def list_bundles(self,
                                                   batch_prefix=batch_prefix))
         return bundles
 
+    def list_files(self, source: TDRSourceRef, prefix: str) -> list[AnvilFile]:
+        self._assert_source(source)
+        self._assert_partition(source, prefix)
+        batch = self._get_batch(source.spec,
+                                'anvil_file',
+                                prefix,
+                                id_column='file_md5sum')
+        return [
+            AnvilFile(uuid=ref.entity_id,
+                      name=row['file_name'],
+                      version=self._version,
+                      size=row['file_size'],
+                      md5=row['file_md5sum'],
+                      drs_uri=row['file_ref'])
+            for ref, row in batch
+        ]
+
     def _emulate_bundle(self, bundle_fqid: TDRAnvilBundleFQID) -> TDRAnvilBundle:
         if bundle_fqid.table_name == BundleType.primary.value:
             log.info('Bundle %r is a primary bundle', bundle_fqid.uuid)

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

@@ -242,6 +242,20 @@ def count_bundles(self, source: TDRSourceSpec) -> int:
         rows = self._run_sql(query)
         return one(rows)['count']
 
+    def count_files(self, source: TDRSourceSpec) -> int:
+        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})
+            ''')
+            for entity_type, entity_cls in api.entity_types.items()
+            if entity_type.endswith('_file')
+        )
+        rows = self._run_sql(query)
+        return sum(row['count'] for row in rows)
+
     def list_bundles(self,
                      source: TDRSourceRef,
                      prefix: str
@@ -261,6 +275,20 @@ def list_bundles(self,
             for row in current_bundles
         ]
 
+    def list_files(self, source: TDRSourceRef, prefix: str) -> list[HCAFile]:
+        self._assert_source(source)
+        self._assert_partition(source, 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})
+            '''
+            for entity_type, entity_cls in api.entity_types.items()
+            if entity_type.endswith('_file')
+        ))
+        return list(map(TDRHCABundle.file_from_row, rows))
+
     def _query_unique_sorted(self,
                              query: str,
                              group_by: str

test/indexer/test_mirror_controller.py

@@ -17,10 +17,16 @@
 from azul.indexer.mirror_controller import (
     MirrorController,
 )
+from azul.json import (
+    copy_json,
+)
 from azul.logging import (
     configure_test_logging,
     get_test_logger,
 )
+from azul.plugins.metadata.hca import (
+    HCAFile,
+)
 from azul_test_case import (
     DCP2TestCase,
 )
@@ -69,6 +75,7 @@ def test_mirroring(self):
             controller = MirrorController(app=MagicMock())
             controller.mirror(event)
             partition_messages = self._read_queue(self.client.mirror_queue())
+            partition_message = copy_json(partition_messages[0])
             partitions = []
             for message in partition_messages:
                 partitions.append(message.pop('prefix'))
@@ -77,3 +84,22 @@ def test_mirroring(self):
                                       source=self.source.to_json()),
                                  message)
             self.assertEqual(list(self.source.spec.prefix.partition_prefixes()), partitions)
+
+        with self.subTest('mirror_partition'):
+            event = [self._mock_sqs_record(partition_message)]
+            file = HCAFile(uuid='405852c9-a0cc-4cd8-b9ff-7c6296223661',
+                           name='foo.txt',
+                           version=None,
+                           drs_uri=None,
+                           size=0,
+                           content_type='text/plain',
+                           sha256='123')
+            plugin_cls = type(self.client.repository_plugin(self.catalog))
+            with patch.object(plugin_cls, 'list_files', return_value=[file]):
+                controller.mirror(event)
+            file_message = one(self._read_queue(self.client.mirror_queue()))
+            expected_message = dict(action='mirror_file',
+                                    catalog=self.catalog,
+                                    source=self.source.to_json(),
+                                    file=file.to_json())
+            self.assertEqual(expected_message, file_message)