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)
@@ -185,6 +187,18 @@ def mirror_partition_message(self,
             'prefix': prefix
         })
 
+    def mirror_file_message(self,
+                            catalog: CatalogName,
+                            source: SourceRef,
+                            file: File,
+                            ) -> SQSFifoMessage:
+        return SQSFifoMessage({
+            'action': MirrorAction.mirror_file.to_json(),
+            'catalog': catalog,
+            'source': cast(JSON, source.to_json()),
+            'file': file.to_json()
+        })
+
     def local_reindex(self, catalog: CatalogName, prefix: str) -> int:
         notifications = [
             self.notification(bundle_fqid)
@@ -607,6 +621,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, prefix: str):
+        plugin = self.repository_plugin(catalog)
+        source = plugin.source_ref_cls.from_json(source)
+
+        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

@@ -39,11 +39,12 @@ def mirror(self, event: Iterable[SQSRecord]):
                 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:
+                    log.info('Would mirror parts of file %r', 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

@@ -719,6 +719,21 @@ 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 UUID starts with the given
+        prefix.
+
+        :param source: A reference to the repository source that contains the
+                       files to list
+
+        :param prefix: A string of at most eight 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,
 )
@@ -199,6 +202,18 @@ 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()
+        ]
+
     def _construct_file_url(self, url: furl, file_name: str) -> furl:
         """
         >>> plugin = Plugin(_sources=set())

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

@@ -136,6 +136,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,
@@ -314,6 +317,19 @@ 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)
+        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 self._get_batch(source.spec, 'anvil_file', prefix)
+        ]
+
     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

@@ -259,6 +259,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({entity_type}_id, {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)