Implement mirror_file, mirror_part & finalize_file (#6862)

Author: nadove-ucsc

lambdas/indexer/.chalice/config.json.template.py

@@ -57,7 +57,7 @@
                     {
                         indexer.mirror.name: {
                             'reserved_concurrency': config.mirroring_concurrency,
-                            'lambda_memory_size': 256,
+                            'lambda_memory_size': 512,
                             'lambda_timeout': config.mirror_lambda_timeout
                         },
                     }

scripts/mirror.py

@@ -1,7 +1,6 @@
 """
 Copy all files from the public sources in a catalog to the current deployment's
-mirroring bucket. The actual file-copying is not yet implemented, so all this
-currently does is send messages to the indexer app that don't do anything.
+mirroring bucket.
 """
 import argparse
 import logging

src/azul/azulclient.py

@@ -24,6 +24,7 @@
 )
 from typing import (
     Self,
+    Sequence,
     TYPE_CHECKING,
     cast,
 )
@@ -63,11 +64,16 @@
 from azul.indexer.index_service import (
     IndexService,
 )
+from azul.indexer.mirror_service import (
+    FilePart,
+    MirrorService,
+)
 from azul.json import (
     Serializable,
 )
 from azul.plugins import (
     File,
+    MetadataPlugin,
     RepositoryPlugin,
 )
 from azul.queues import (
@@ -113,6 +119,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 +130,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.
@@ -210,6 +221,42 @@ def mirror_file_message(self,
             group_id=f'{source.id}:{file.uuid}'
         )
 
+    def mirror_part_message(self,
+                            catalog: CatalogName,
+                            file: File,
+                            part: FilePart,
+                            upload_id: str,
+                            etags: Sequence[str] = ()
+                            ) -> SQSFifoMessage:
+        return SQSFifoMessage(
+            body={
+                'catalog': catalog,
+                'file': file.to_json(),
+                'upload_id': upload_id,
+                'action': MirrorAction.mirror_part.to_json(),
+                'part': part.to_json(),
+                'etags': etags
+            },
+            group_id=self.mirror_service.mirror_object_key(file)
+        )
+
+    def finalize_file_message(self,
+                              catalog: CatalogName,
+                              file: File,
+                              upload_id: str,
+                              etags: Sequence[str]
+                              ) -> SQSFifoMessage:
+        return SQSFifoMessage(
+            body={
+                'catalog': catalog,
+                'file': file.to_json(),
+                'upload_id': upload_id,
+                'action': MirrorAction.finalize_file.to_json(),
+                'etags': etags
+            },
+            group_id=self.mirror_service.mirror_object_key(file)
+        )
+
     def local_reindex(self, catalog: CatalogName, prefix: str) -> int:
         notifications = [
             self.notification(bundle_fqid)
@@ -488,6 +535,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)
 
@@ -644,6 +695,61 @@ def message(file: File) -> SQSMessage:
         messages = map(message, plugin.list_files(source, prefix))
         self.queue_mirror_messages(messages)
 
+    def mirror_file(self,
+                    catalog: CatalogName,
+                    file_json: JSON
+                    ):
+        file = self.load_file(catalog, file_json)
+        assert file.size is not None, R('File size unknown', file)
+        part_size = FilePart.default_size
+        if file.size <= part_size:
+            log.info('Mirroring file %r via single-part upload', file.uuid)
+            self.mirror_service.mirror_file(catalog, file)
+        else:
+            log.info('Mirroring file %r via multi-part upload', file.uuid)
+            upload_id = self.mirror_service.begin_mirroring_file(file)
+            part = FilePart.head(file, part_size)
+            messages = [self.mirror_part_message(catalog, file, part, upload_id)]
+            self.queue_mirror_messages(messages)
+
+    def mirror_file_part(self,
+                         catalog: CatalogName,
+                         file_json: JSON,
+                         part_json: JSON,
+                         upload_id: str,
+                         etags: Sequence[str]
+                         ):
+        file = self.load_file(catalog, file_json)
+        part = FilePart.from_json(part_json)
+        etag = self.mirror_service.mirror_file_part(catalog, file, part, upload_id)
+        etags = [*etags, etag]
+        next_part = part.next(file)
+        if next_part is None:
+            log.info('File %r fully uploaded in %d parts', file.uuid, len(etags))
+            message = self.finalize_file_message(catalog,
+                                                 file,
+                                                 upload_id,
+                                                 etags)
+        else:
+            message = self.mirror_part_message(catalog,
+                                               file,
+                                               next_part,
+                                               upload_id,
+                                               etags)
+        self.queue_mirror_messages([message])
+
+    def finalize_file(self,
+                      catalog: CatalogName,
+                      file_json: JSON,
+                      upload_id: str,
+                      etags: Sequence[str]
+                      ):
+        file = self.load_file(catalog, file_json)
+        self.mirror_service.finish_mirroring_file(file, upload_id, etags)
+
+    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/lambda_iam_policy.py

@@ -113,6 +113,20 @@
                 },
             ] if config.enable_log_forwarding else []
         ),
+        *(
+            [
+                {
+                    'Effect': 'Allow',
+                    'Action': [
+                        's3:PutObject',
+                        's3:GetObject',
+                    ],
+                    'Resource': [
+                        f'arn:aws:s3:::{aws.mirror_bucket}/*'
+                    ]
+                }
+            ] if config.enable_mirroring else []
+        ),
         {
             'Effect': 'Allow',
             'Action': [

src/azul/indexer/mirror_controller.py

@@ -1,5 +1,4 @@
 import logging
-import time
 from typing import (
     Iterable,
 )
@@ -18,9 +17,11 @@
 from azul.indexer.action_controller import (
     ActionController,
 )
+from azul.indexer.mirror_service import (
+    MirrorService,
+)
 from azul.types import (
     JSON,
-    json_mapping,
     json_str,
 )
 
@@ -33,6 +34,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]):
         self._handle_events(event, self._mirror)
 
@@ -45,11 +50,18 @@ def _mirror(self, message: JSON):
                                          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)
+            self.client.mirror_file(message['catalog'],
+                                    message['file'])
+        elif action is MirrorAction.mirror_part:
+            self.client.mirror_file_part(message['catalog'],
+                                         message['file'],
+                                         message['part'],
+                                         message['upload_id'],
+                                         message['etags'])
+        elif action is MirrorAction.finalize_file:
+            self.client.finalize_file(message['catalog'],
+                                      message['file'],
+                                      message['upload_id'],
+                                      message['etags'])
         else:
-            # FIXME: Implement mirror_file, mirror_part & finalize_file
-            #        https://github.com/DataBiosphere/azul/issues/6862
             assert False, action

src/azul/indexer/mirror_service.py

@@ -124,6 +124,16 @@ def _storage(self) -> StorageService:
     def repository_plugin(self, catalog: CatalogName) -> RepositoryPlugin:
         return RepositoryPlugin.load(catalog).create(catalog)
 
+    def mirror_file(self, catalog: CatalogName, file: File):
+        """
+        Upload the file in a single request. For larger files, use
+        :meth:`begin_mirroring_file` instead.
+        """
+        file_content = self._download(catalog, file)
+        self._storage.put(object_key=self.mirror_object_key(file),
+                          data=file_content,
+                          content_type=file.content_type)
+
     def begin_mirroring_file(self, file: File) -> str:
         """
         Initiate a multipart upload of the file's content and return the upload
@@ -137,10 +147,11 @@ def mirror_file_part(self,
                          catalog: CatalogName,
                          file: File,
                          part: FilePart,
-                         upload_id: str):
+                         upload_id: str
+                         ) -> str:
         """
         Upload a part of a file to a multipart upload begun with
-        :meth:`begin_mirroring_file`
+        :meth:`begin_mirroring_file` and return the uploaded part's ETag.
         """
         upload = self._get_upload(file, upload_id)
         file_content = self._download(catalog, file, part)

test/indexer/test_mirror_controller.py

@@ -1,3 +1,4 @@
+import hashlib
 from unittest.mock import (
     MagicMock,
     PropertyMock,
@@ -14,9 +15,15 @@
 from azul import (
     config,
 )
+from azul.deployment import (
+    aws,
+)
 from azul.indexer.mirror_controller import (
     MirrorController,
 )
+from azul.indexer.mirror_service import (
+    MirrorService,
+)
 from azul.json import (
     copy_json,
 )
@@ -30,6 +37,9 @@
 from azul_test_case import (
     DCP2TestCase,
 )
+from service import (
+    S3TestCase,
+)
 from sqs_test_case import (
     WorkQueueTestCase,
 )
@@ -43,7 +53,7 @@ def setUpModule():
 
 
 @mock_aws
-class TestMirrorController(DCP2TestCase, WorkQueueTestCase):
+class TestMirrorController(DCP2TestCase, WorkQueueTestCase, S3TestCase):
 
     @classmethod
     def _patch_enable_mirroring(cls):
@@ -60,6 +70,11 @@ def setUp(self):
         super().setUp()
         self._create_mock_queues(config.mirror_queue.name,
                                  config.mirror_queue.to_fail.name)
+        self._create_test_bucket(self.bucket)
+
+    @property
+    def bucket(self) -> str:
+        return aws.mirror_bucket
 
     def test_mirroring(self):
         with self.subTest('remote_mirror'):
@@ -85,15 +100,17 @@ def test_mirroring(self):
                                  message)
             self.assertEqual(list(self.source.spec.prefix.partition_prefixes()), partitions)
 
+        file_contents = b'lorem ipsum dolor sit\n'
+
         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,
+                           drs_uri='drs://fake-domain.lan/foo',
+                           size=len(file_contents),
                            content_type='text/plain',
-                           sha256='123')
+                           sha256=hashlib.sha256(file_contents).hexdigest())
             plugin_cls = type(self.client.repository_plugin(self.catalog))
             with patch.object(plugin_cls, 'list_files', return_value=[file]):
                 controller.mirror(event)
@@ -103,3 +120,12 @@ def test_mirroring(self):
                                     source=self.source.to_json(),
                                     file=file.to_json())
             self.assertEqual(expected_message, file_message)
+
+        with self.subTest('mirror_file'):
+            event = [self._mock_sqs_record(file_message)]
+            with patch.object(MirrorService, '_download', return_value=file_contents):
+                controller.mirror(event)
+            response = self._s3.get_object(Bucket=self.bucket,
+                                           Key=controller.service.mirror_object_key(file))
+            mirrored_file_contents = response['Body'].read()
+            self.assertEqual(mirrored_file_contents, file_contents)