Implement mirror_file, mirror_part & finalize_file (#6862)

Author: nadove-ucsc

src/azul/azulclient.py

@@ -64,11 +64,15 @@
 from azul.indexer.index_service import (
     IndexService,
 )
+from azul.indexer.mirror_service import (
+    MirrorService,
+)
 from azul.json import (
     Serializable,
 )
 from azul.plugins import (
     File,
+    MetadataPlugin,
     RepositoryPlugin,
 )
 from azul.queues import (
@@ -113,6 +117,8 @@ class MirrorAction(Action):
     mirror_source = auto()
     mirror_partition = auto()
     mirror_file = auto()
+    mirror_part = auto()
+    finalize_file = auto()
 
 
 @attrs.frozen(kw_only=True)
@@ -122,6 +128,9 @@ class AzulClient(SignatureHelper, HasCachedHttpClient):
     def repository_plugin(self, catalog: CatalogName) -> RepositoryPlugin:
         return self.index_service.repository_plugin(catalog)
 
+    def metadata_plugin(self, catalog: CatalogName) -> MetadataPlugin:
+        return self.index_service.metadata_plugin(catalog)
+
     def notification(self, bundle_fqid: SourcedBundleFQID) -> JSON:
         """
         Generate an indexer notification for the given bundle.
@@ -197,6 +206,38 @@ def mirror_file_message(self,
             'file': file.to_json()
         })
 
+    def mirror_part_message(self,
+                            catalog: CatalogName,
+                            source: SourceRef,
+                            file: File,
+                            part: File.PartRange,
+                            upload_id: str
+                            ) -> SQSFifoMessage:
+        key = self.mirror_service.mirror_object_key(file)
+        return SQSFifoMessage({
+            'action': MirrorAction.mirror_part.to_json(),
+            'catalog': catalog,
+            'source': source.to_json(),
+            'file': file.to_json(),
+            'part': part.to_json(),
+            'upload_id': upload_id
+        }, group_id=key)
+
+    def finalize_file_message(self,
+                              catalog: CatalogName,
+                              source: SourceRef,
+                              file: File,
+                              upload_id: str
+                              ) -> SQSFifoMessage:
+        key = self.mirror_service.mirror_object_key(file)
+        return SQSFifoMessage({
+            'action': MirrorAction.finalize_file.to_json(),
+            'catalog': catalog,
+            'source': source.to_json(),
+            'file': file.to_json(),
+            'upload_id': upload_id
+        }, group_id=key)
+
     def local_reindex(self, catalog: CatalogName, prefix: str) -> int:
         notifications = [
             self.notification(bundle_fqid)
@@ -475,6 +516,10 @@ def group_key(fqid: SourcedBundleFQID):
     def index_service(self) -> IndexService:
         return IndexService()
 
+    @cached_property
+    def mirror_service(self) -> MirrorService:
+        return MirrorService()
+
     def delete_all_indices(self, catalog: CatalogName):
         self.index_service.delete_indices(catalog)
 
@@ -632,6 +677,26 @@ def message(file: File) -> SQSMessage:
         messages = map(message, plugin.list_files(source, prefix))
         self.queue_mirror_messages(messages)
 
+    def mirror_file(self,
+                    catalog: CatalogName,
+                    source_json: JSON,
+                    file_json: JSON,
+                    part_size: int):
+        source = self.repository_plugin(catalog).source_ref_cls.from_json(source_json)
+        file = self.load_file(catalog, file_json)
+        upload_id = self.mirror_service.begin_mirroring_file(file)
+
+        def message(part: File.PartRange) -> SQSMessage:
+            log.info('Mirroring part %r of file %r in source %r from catalog %r ',
+                     part.part_number, file.uuid, str(source.spec), catalog)
+            return self.mirror_part_message(catalog, source, file, part, upload_id)
+
+        messages = map(message, file.partition(part_size))
+        self.queue_mirror_messages(messages)
+
+    def load_file(self, catalog: CatalogName, file: JSON) -> File:
+        return self.metadata_plugin(catalog).file_class().from_json(file)
+
 
 class AzulClientError(RuntimeError):
     pass

src/azul/indexer/mirror_controller.py

@@ -19,6 +19,12 @@
 from azul.indexer.action_controller import (
     ActionController,
 )
+from azul.indexer.mirror_service import (
+    MirrorService,
+)
+from azul.plugins import (
+    File,
+)
 
 log = logging.getLogger(__name__)
 
@@ -29,6 +35,10 @@ class MirrorController(ActionController[MirrorAction]):
     def client(self) -> AzulClient:
         return AzulClient()
 
+    @property
+    def service(self) -> MirrorService:
+        return self.client.mirror_service
+
     def mirror(self, event: Iterable[SQSRecord]):
         for record in event:
             message = json.loads(record.body)
@@ -43,11 +53,22 @@ def mirror(self, event: Iterable[SQSRecord]):
                                                  message['source'],
                                                  message['prefix'])
                 elif action is MirrorAction.mirror_file:
-                    log.info('Would mirror parts of file %r', message['file']['uuid'])
-                    time.sleep(1)
+                    self.client.mirror_file(message['catalog'],
+                                            message['source'],
+                                            message['file'],
+                                            part_size=...)
+                elif action is MirrorAction.mirror_part:
+                    catalog = message['catalog']
+                    file = self.client.load_file(catalog, message['file'])
+                    part = File.PartRange.from_json(message['part'])
+                    self.service.mirror_file_part(catalog, file, part, message['upload_id'])
+                elif action is MirrorAction.finalize_file:
+                    file = self.client.load_file(message['catalog'], message['file'])
+                    # We rely on the FIFO mechanics of the mirror queue to
+                    # ensure that all parts have been successfully uploaded
+                    # before the finalization message is handled
+                    self.service.finish_mirroring_file(file, message['upload_id'], etags=None)
                 else:
-                    # FIXME: Implement mirror_file, mirror_part & finalize_file
-                    #        https://github.com/DataBiosphere/azul/issues/6862
                     assert False, action
             except BaseException:
                 log.warning(f'Worker failed to handle message {message}.', exc_info=True)

src/azul/indexer/mirror_service.py

@@ -96,6 +96,9 @@ def get_mirror_url(self, file: File) -> str:
     def mirror_object_key(self, file: File) -> str:
         return self._file_key('file', file)
 
+    def info_object_key(self, file: File) -> str:
+        return self._file_key('info', file)
+
     def _file_key(self, prefix: str, file: File) -> str:
         digest, digest_type = file.digest()
         assert all(c in string.hexdigits for c in digest), R(