Repository: DataBiosphere/terra-notebook-utils (language: python)

Issue:
Add blobstore tests
related to #294 
Add extensible support for cloud storage
- Terra will eventually support Azure.
- A general DRS client should support cloud copies to major cloud providers.

┆Issue is synchronized with this [Jira Story](https://ucsc-cgl.atlassian.net/browse/TLP-2083)



Discussion / hints:

Fix patch:
diff --git a/terra_notebook_utils/blobstore/copy_client.py b/terra_notebook_utils/blobstore/copy_client.py
new file mode 100644
index 00000000..44a2fc36
--- /dev/null
+++ b/terra_notebook_utils/blobstore/copy_client.py
@@ -0,0 +1,103 @@
+import os
+import logging
+import multiprocessing
+from concurrent.futures import ProcessPoolExecutor
+from typing import Union
+
+from getm import default_chunk_size
+from getm.concurrent.collections import ConcurrentHeap
+
+from terra_notebook_utils.blobstore.gs import GSBlobStore, GSBlob
+from terra_notebook_utils.blobstore.url import URLBlobStore, URLBlob
+from terra_notebook_utils.blobstore.local import LocalBlobStore, LocalBlob
+from terra_notebook_utils.blobstore import BlobstoreChecksumError
+
+
+logger = logging.getLogger(__name__)
+
+AnyBlobStore = Union[GSBlobStore, URLBlobStore, LocalBlobStore]
+AnyBlob = Union[GSBlob, URLBlob, LocalBlob]
+CloudBlob = GSBlob
+
+def _finalize_copy(src_blob: AnyBlob, dst_blob: AnyBlob):
+    logger.info(f"Copied {src_blob.url} to {dst_blob.url}")
+
+def _download(src_blob: AnyBlob, dst_blob: LocalBlob):
+    logger.debug(f"Starting download {src_blob.url} to {dst_blob.url}")
+    dirname = os.path.dirname(dst_blob.url)
+    if dirname:
+        os.makedirs(dirname, exist_ok=True)
+    # The download methods for each Blob is expected to compute checksums
+    src_blob.download(dst_blob.url)
+    _finalize_copy(src_blob, dst_blob)
+
+def _copy_intra_cloud(src_blob: AnyBlob, dst_blob: AnyBlob):
+    # In general it is not necesary to compute checksums for intra cloud copies. The Storage services will do that for
+    # us.
+    # However, S3Etags depend on the part layout. Either ensure source and destination part layouts are the same, or
+    # compute the destination S3Etag on the fly.
+    logger.debug(f"Starting intra-cloud {src_blob.url} to {dst_blob.url}")
+    assert isinstance(src_blob, type(dst_blob))
+    dst_blob.copy_from(src_blob)  # type: ignore
+    if src_blob.cloud_native_checksum() != dst_blob.cloud_native_checksum():
+        logger.error(f"Checksum failed for {src_blob.url} to {dst_blob.url}")
+        dst_blob.delete()
+        raise BlobstoreChecksumError()
+    _finalize_copy(src_blob, dst_blob)
+
+def _copy_oneshot_passthrough(src_blob: AnyBlob, dst_blob: CloudBlob):
+    logger.debug(f"Starting oneshot passthrough {src_blob.url} to {dst_blob.url}")
+    data = src_blob.get()
+    dst_blob.put(data)
+    if not dst_blob.Hash(data).matches(dst_blob.cloud_native_checksum()):
+        logger.error(f"Checksum failed for {src_blob.url} to {dst_blob.url}")
+        dst_blob.delete()
+        raise BlobstoreChecksumError()
+    _finalize_copy(src_blob, dst_blob)
+
+def _copy_multipart_passthrough(src_blob: AnyBlob, dst_blob: CloudBlob):
+    logger.debug(f"Starting multipart passthrough {src_blob.url} to {dst_blob.url}")
+    cs = dst_blob.Hash()
+    with dst_blob.multipart_writer() as writer:
+        for part in src_blob.iter_content():
+            writer.put_part(part)
+            cs.update(part.data)
+    if not cs.matches(dst_blob.cloud_native_checksum()):
+        logger.error(f"Checksum failed for {src_blob.url} to {dst_blob.url}")
+        dst_blob.delete()
+        raise BlobstoreChecksumError()
+    _finalize_copy(src_blob, dst_blob)
+
+class CopyClient:
+    multipart_threshold = default_chunk_size
+
+    def __init__(self, concurrency: int=multiprocessing.cpu_count()):
+        self._executor = ProcessPoolExecutor(max_workers=concurrency)
+        self._queue = ConcurrentHeap(self._executor, concurrency)
+
+    def copy(self, src_blob: AnyBlob, dst_blob: AnyBlob):
+        size = src_blob.size()
+        priority = -size
+        if isinstance(dst_blob, LocalBlob):
+            self._queue.priority_put(priority, _download, src_blob, dst_blob)
+        elif isinstance(src_blob, type(dst_blob)):
+            self._queue.priority_put(priority, _copy_intra_cloud, src_blob, dst_blob)
+        else:
+            if size <= self.multipart_threshold:
+                self._queue.priority_put(priority, _copy_oneshot_passthrough, src_blob, dst_blob)
+            else:
+                self._queue.priority_put(priority, _copy_multipart_passthrough, src_blob, dst_blob)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        try:
+            for _ in self._queue:
+                pass
+        finally:
+            self._executor.shutdown()
+
+def copy(src_blob: AnyBlob, dst_blob: AnyBlob):
+    with CopyClient() as client:
+        client.copy(src_blob, dst_blob)


Test patch:
diff --git a/tests/infra/__init__.py b/tests/infra/__init__.py
index 4c531b7c..e8503188 100644
--- a/tests/infra/__init__.py
+++ b/tests/infra/__init__.py
@@ -1,5 +1,7 @@
 import io
 import os
+import logging
+import requests
 import warnings
 
 from google.cloud.storage import Client
@@ -18,6 +20,9 @@ def setUp(self):
         # Suppress unclosed socket warnings
         warnings.simplefilter("ignore", ResourceWarning)
 
+        # Suppress urllib3 warnings
+        logging.getLogger(requests.packages.urllib3.__package__).setLevel(logging.ERROR)
+
 def upload_data(uri: str, data: bytes):
     if uri.startswith("gs://"):
         bucket, key = uri[5:].split("/", 1)
diff --git a/tests/test_copy_client.py b/tests/test_copy_client.py
new file mode 100644
index 00000000..d6f24dd7
--- /dev/null
+++ b/tests/test_copy_client.py
@@ -0,0 +1,104 @@
+#!/usr/bin/env python
+import os
+import sys
+import logging
+import tempfile
+import unittest
+from unittest import mock
+from uuid import uuid4
+from typing import Dict, Iterable, Optional
+
+from getm import default_chunk_size
+
+pkg_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))  # noqa
+sys.path.insert(0, pkg_root)  # noqa
+
+from terra_notebook_utils import WORKSPACE_BUCKET
+from terra_notebook_utils.blobstore.gs import GSBlobStore, GSBlob
+from terra_notebook_utils.blobstore.local import LocalBlobStore, LocalBlob
+from terra_notebook_utils.blobstore import BlobStore, copy_client
+
+from tests import infra
+
+
+gs_blobstore = GSBlobStore(infra.get_env("TNU_BLOBSTORE_TEST_GS_BUCKET"))
+local_test_tempdir = tempfile.TemporaryDirectory()
+local_test_bucket = local_test_tempdir.name
+local_blobstore = LocalBlobStore(local_test_tempdir.name)
+
+logging.basicConfig(stream=sys.stderr, level=logging.INFO)
+copy_client.logger.setLevel(logging.DEBUG)
+
+class TestData:
+    def __init__(self, oneshot_size: int=7, multipart_size: int=2 * default_chunk_size + 1):
+        self.oneshot_size = oneshot_size
+        self.multipart_size = multipart_size
+
+        self._oneshot: Optional[bytes] = None
+        self._multipart: Optional[bytes] = None
+        self._oneshot_key = f"{uuid4()}"
+        self._multipart_key = f"{uuid4()}"
+        self._uploaded: Dict[BlobStore, bool] = dict()
+
+    @property
+    def oneshot(self):
+        self._oneshot = self._oneshot or os.urandom(self.oneshot_size)
+        return self._oneshot
+
+    @property
+    def multipart(self):
+        self._multipart = self._multipart or os.urandom(self.multipart_size)
+        return self._multipart
+
+    def uploaded(self, dests: Iterable[BlobStore]):
+        oneshot = dict(key=self._oneshot_key, data=self.oneshot)
+        multipart = dict(key=self._multipart_key, data=self.multipart)
+        for bs in dests:
+            if not self._uploaded.get(bs):
+                for data_config in [oneshot, multipart]:
+                    bs.blob(data_config['key']).put(data_config['data'])
+                self._uploaded[bs] = True
+        return oneshot, multipart
+
+test_data = TestData()
+
+class TestCopyClient(infra.SuppressWarningsMixin, unittest.TestCase):
+    def test_copy(self):
+        src_blob = mock.MagicMock()
+        dst_blob = mock.MagicMock()
+        with mock.patch("terra_notebook_utils.blobstore.copy_client.copy") as copy_method:
+            copy_client.copy(src_blob, dst_blob)
+            copy_method.assert_called_once()
+
+    def test_copy_client(self):
+        with self.subTest("should work"):
+            expected_data_map, completed_keys = self._do_blobstore_copies()
+            self.assertEqual(len(expected_data_map), len(completed_keys))
+            for blob, expected_data in expected_data_map.items():
+                with self.subTest(blob.url):
+                    self.assertEqual(blob.get(), expected_data)
+
+    def _do_blobstore_copies(self,
+                             src_blobstores=(local_blobstore, gs_blobstore),
+                             dst_blobstores=(local_blobstore, gs_blobstore),
+                             ignore_missing_checksums=True,
+                             compute_checksums=False):
+        oneshot, multipart = test_data.uploaded([local_blobstore, gs_blobstore])
+        expected_data_map = dict()
+        completed_keys = list()
+        with copy_client.CopyClient() as client:
+            for src_bs in src_blobstores:
+                for dst_bs in dst_blobstores:
+                    for data_bundle in (oneshot, multipart):
+                        src_blob = src_bs.blob(data_bundle['key'])
+                        dst_blob = dst_bs.blob(os.path.join(f"{uuid4()}", f"{uuid4()}"))
+                        completed_keys.append(dst_blob.key)
+                        if compute_checksums:
+                            client.copy_compute_checksums(src_blob, dst_blob)
+                        else:
+                            client.copy(src_blob, dst_blob)
+                        expected_data_map[dst_blob] = data_bundle['data']
+        return expected_data_map, completed_keys
+
+if __name__ == '__main__':
+    unittest.main()