Implement mirror_file, mirror_part & finalize_file (#6862, PR #7043)

Author: achave11-ucsc

.mypy.ini

@@ -52,7 +52,8 @@ modules =
     azul.indexer.field,
     azul.indexer.action_controller,
     azul.indexer.mirror_controller,
-    azul.service.app_controller
+    azul.service.app_controller,
+    azul.indexer.mirror_service
 packages =
     azul.openapi
 

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

scripts/mirror_file.py

@@ -5,8 +5,6 @@
 """
 import argparse
 import logging
-import math
-import string
 import sys
 
 from azul import (
@@ -17,24 +15,16 @@
 from azul.args import (
     AzulArgumentHelpFormatter,
 )
-from azul.azulclient import (
-    AzulClient,
-)
-from azul.deployment import (
-    aws,
-)
-from azul.drs import (
-    AccessMethod,
-)
 from azul.http import (
     http_client,
 )
+from azul.indexer.mirror_service import (
+    FilePart,
+    MirrorService,
+)
 from azul.logging import (
     configure_script_logging,
 )
-from azul.plugins import (
-    File,
-)
 from azul.plugins.metadata.hca import (
     HCAFile,
 )
@@ -47,9 +37,6 @@
 from azul.service.source_service import (
     SourceService,
 )
-from azul.service.storage_service import (
-    StorageService,
-)
 
 log = logging.getLogger(__name__)
 
@@ -70,63 +57,30 @@ def get_file(catalog: CatalogName, file_uuid: str) -> HCAFile:
     return file
 
 
-def get_download_url(catalog: CatalogName, file: File) -> str:
-    drs_uri = file.drs_uri
-    drs = AzulClient().repository_plugin(catalog).drs_client()
-    access = drs.get_object(drs_uri, AccessMethod.gs)
-    assert access.method is AccessMethod.https, access
-    return access.url
-
-
-def object_key(file: File) -> str:
-    digest, digest_type = file.digest()
-    assert all(c in string.hexdigits for c in digest), R(
-        'Expected a hexadecimal digest', digest)
-    return f'file/{digest.lower()}.{digest_type}'
-
-
 def mirror_file(catalog: CatalogName, file_uuid: str, part_size: int) -> str:
     assert config.enable_mirroring, R('Mirroring must be enabled')
     assert config.is_tdr_enabled(catalog), R('Only TDR catalogs are supported')
-    assert config.is_hca_enabled(catalog), R('Only HCA catalogs are supported')
-    # https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html
-    assert 5 * 2 ** 20 <= part_size <= 5 * 2 ** 30, R(
-        'Invalid part size', part_size)
     file = get_file(catalog, file_uuid)
-    download_url = get_download_url(catalog, file)
-    key = object_key(file)
-    storage = StorageService(bucket_name=aws.mirror_bucket)
-    upload = storage.create_multipart_upload(key, content_type=file.content_type)
-
-    total_size = file.size
-    part_count = math.ceil(total_size / part_size)
-    assert part_count <= 10000, R('Part size is too small for this file',
-                                  total_size, part_size, part_count)
+    service = MirrorService()
+    upload_id = service.begin_mirroring_file(file)
 
-    def file_part(part_number: int) -> str:
-        start = part_number * part_size
-        end = min((part_number + 1) * part_size - 1, total_size)
-        response = http.request('GET',
-                                download_url,
-                                headers={'Range': f'bytes={start}-{end}'})
-        if response.status == 206:
-            return storage.upload_multipart_part(response.data,
-                                                 part_number + 1,
-                                                 upload)
-        else:
-            raise RuntimeError('Unexpected response from repository', response.status)
+    def mirror_parts():
+        part = FilePart.first(file, part_size)
+        while part is not None:
+            yield service.mirror_file_part(catalog, file, part, upload_id)
+            part = part.next(file)
 
-    etags = list(map(file_part, range(part_count)))
-    storage.complete_multipart_upload(upload, etags)
-    return storage.get_presigned_url(key, file.name)
+    etags = list(mirror_parts())
+    service.finish_mirroring_file(file, upload_id, etags=etags)
+    return service.get_mirror_url(file)
 
 
 def main(argv):
     parser = argparse.ArgumentParser(description=__doc__,
                                      formatter_class=AzulArgumentHelpFormatter)
     parser.add_argument('-c', '--catalog', default=config.default_catalog)
     parser.add_argument('-f', '--file-uuid')
-    parser.add_argument('-p', '--part-size', type=int, default=50 * 2 ** 20)
+    parser.add_argument('-p', '--part-size', type=int, default=FilePart.default_size)
     args = parser.parse_args(argv)
     signed_url = mirror_file(args.catalog, args.file_uuid, args.part_size)
     print(signed_url)

src/azul/azulclient.py

@@ -24,6 +24,7 @@
 )
 from typing import (
     Self,
+    Sequence,
     TYPE_CHECKING,
     cast,
 )
@@ -41,7 +42,6 @@
 from azul import (
     CatalogName,
     R,
-    cache,
     cached_property,
     config,
 )
@@ -64,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 (
@@ -114,15 +119,19 @@ class MirrorAction(Action):
     mirror_source = auto()
     mirror_partition = auto()
     mirror_file = auto()
+    mirror_part = auto()
+    finalize_file = auto()
 
 
 @attrs.frozen(kw_only=True)
 class AzulClient(SignatureHelper, HasCachedHttpClient):
     num_workers: int = 16
 
-    @cache
     def repository_plugin(self, catalog: CatalogName) -> RepositoryPlugin:
-        return RepositoryPlugin.load(catalog).create(catalog)
+        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:
         """
@@ -212,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)
@@ -490,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)
 
@@ -646,6 +695,77 @@ 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)
+
+        file_is_large = file.size > 10 * 1024 ** 2
+        deployment_is_stable = (config.deployment.is_stable
+                                and not config.deployment.is_unit_test
+                                and catalog not in config.integration_test_catalogs)
+        if file_is_large and not deployment_is_stable:
+            log.info('Not mirroring file %r (%d bytes) to save cost',
+                     file.uuid, file.size)
+        else:
+            # Ensure we test with multiple parts on lower deployments
+            part_size = FilePart.default_size if deployment_is_stable else FilePart.min_size
+            if file.size <= part_size:
+                log.info('Mirroring file %r via standard upload', file.uuid)
+                self.mirror_service.mirror_file(catalog, file)
+                log.info('Successfully mirrored file %r via standard upload', file.uuid)
+            else:
+                log.info('Mirroring file %r via multi-part upload', file.uuid)
+                upload_id = self.mirror_service.begin_mirroring_file(file)
+                first_part = FilePart.first(file, part_size)
+                etag = self.mirror_service.mirror_file_part(catalog, file, first_part, upload_id)
+                next_part = first_part.next(file)
+                assert next_part is not None
+                messages = [self.mirror_part_message(catalog, file, next_part, upload_id, [etag])]
+                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)
+        assert len(etags) > 0
+        self.mirror_service.finish_mirroring_file(file, upload_id, etags)
+        log.info('Successfully mirrored file %r via multi-part upload', file.uuid)
+
+    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': [