Encapsulate file mirroring in service class

Author: nadove-ucsc

scripts/mirror_file.py

@@ -5,7 +5,6 @@
 """
 import argparse
 import logging
-import string
 import sys
 
 from azul import (
@@ -16,24 +15,15 @@
 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 (
+    MirrorService,
+)
 from azul.logging import (
     configure_script_logging,
 )
-from azul.plugins import (
-    File,
-)
 from azul.plugins.metadata.hca import (
     HCAFile,
 )
@@ -46,9 +36,6 @@
 from azul.service.source_service import (
     SourceService,
 )
-from azul.service.storage_service import (
-    StorageService,
-)
 
 log = logging.getLogger(__name__)
 
@@ -69,45 +56,18 @@ 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')
     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)
-
-    def upload_part(part: File.PartRange) -> str:
-        response = http.request('GET',
-                                download_url,
-                                headers={'Range': f'bytes={part.start}-{part.end}'})
-        if response.status == 206:
-            return storage.upload_multipart_part(response.data,
-                                                 part.part_number + 1,
-                                                 upload)
-        else:
-            raise RuntimeError('Unexpected response from repository', response.status)
-
-    etags = list(map(upload_part, file.partition(part_size)))
-    storage.complete_multipart_upload(upload, etags)
-    return storage.get_presigned_url(key, file.name)
+    service = MirrorService()
+    upload_id = service.begin_mirroring_file(file)
+    etags = [
+        service.mirror_file_part(catalog, file, part, upload_id)
+        for part in file.partition(part_size)
+    ]
+    service.finish_mirroring_file(file, upload_id, etags=etags)
+    return service.get_mirror_url(file)
 
 
 def main(argv):

src/azul/indexer/mirror_service.py

@@ -0,0 +1,114 @@
+import string
+from typing import (
+    TYPE_CHECKING,
+)
+
+from azul import (
+    CatalogName,
+    cache,
+    cached_property,
+    config,
+)
+from azul.deployment import (
+    R,
+    aws,
+)
+from azul.drs import (
+    AccessMethod,
+)
+from azul.http import (
+    HasCachedHttpClient,
+)
+from azul.plugins import (
+    File,
+    RepositoryPlugin,
+)
+from azul.plugins.metadata.hca import (
+    HCAFile,
+)
+from azul.service.storage_service import (
+    StorageService,
+)
+
+if TYPE_CHECKING:
+    from mypy_boto3_s3.service_resource import (
+        MultipartUpload,
+    )
+
+
+class MirrorService(HasCachedHttpClient):
+
+    @cached_property
+    def _storage(self) -> StorageService:
+        return StorageService(bucket_name=aws.mirror_bucket)
+
+    @cache
+    def repository_plugin(self, catalog: CatalogName) -> RepositoryPlugin:
+        return RepositoryPlugin.load(catalog).create(catalog)
+
+    def begin_mirroring_file(self, file: File) -> str:
+        """
+        Initiate a multipart upload of the given file and return the upload ID.
+        """
+        assert isinstance(file, HCAFile), R('Only HCA catalogs are supported')
+        upload = self._storage.create_multipart_upload(object_key=self.mirror_object_key(file),
+                                                       content_type=file.content_type)
+        return upload.upload_id
+
+    def mirror_file_part(self,
+                         catalog: CatalogName,
+                         file: File,
+                         part: File.PartRange,
+                         upload_id: str):
+        """
+        Upload a part of a file to a multipart upload begun with
+        :meth:`begin_mirroring_file`
+        """
+        download_url = self._get_repository_url(catalog, file)
+        upload = self._get_upload(file, upload_id)
+        response = self._http_client.request('GET',
+                                             download_url,
+                                             headers={'Range': f'bytes={part.start}-{part.end}'})
+        if response.status == 206:
+            return self._storage.upload_multipart_part(response.data,
+                                                       part.part_number + 1,
+                                                       upload)
+        else:
+            raise RuntimeError('Unexpected response from repository', response.status)
+
+    def finish_mirroring_file(self,
+                              file: File,
+                              upload_id: str,
+                              *,
+                              etags: list[str] | None):
+        """
+        Complete a multipart upload begun with :meth:`begin_mirroring_file`.
+        If ETags are not provided, the caller is responsible for ensuring that
+        all previous calls to :meth:`mirror_file_part` were successful.
+        """
+        upload = self._get_upload(file, upload_id)
+        self._storage.complete_multipart_upload(upload, etags)
+
+    def get_mirror_url(self, file: File) -> str:
+        return self._storage.get_presigned_url(key=self.mirror_object_key(file),
+                                               file_name=file.name)
+
+    def mirror_object_key(self, file: File) -> str:
+        return self._file_key('file', 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(
+            'Expected a hexadecimal digest', digest)
+        return f'{prefix}/{digest.lower()}.{digest_type}'
+
+    def _get_repository_url(self, catalog: CatalogName, file: File):
+        assert config.is_tdr_enabled(catalog), R('Only TDR catalogs are supported')
+        drs = self.repository_plugin(catalog).drs_client(authentication=None)
+        access = drs.get_object(file.drs_uri, AccessMethod.gs)
+        assert access.method is AccessMethod.https, access
+        return access.url
+
+    def _get_upload(self, file: File, upload_id: str) -> 'MultipartUpload':
+        return self._storage.load_multipart_upload(object_key=self.mirror_object_key(file),
+                                                   upload_id=upload_id)