PYTHONSDK-99: Add optional blob checksum calculation to helpers. Also added documentation to helpers.

RachelTucker · RachelTucker · commit a758ef229575 · 2022-02-21T13:32:30.000-07:00
diff --git a/ds3/ds3Helpers.py b/ds3/ds3Helpers.py
@@ -1,4 +1,4 @@
-#   Copyright 2021 Spectra Logic Corporation. All Rights Reserved.
+#   Copyright 2021-2022 Spectra Logic Corporation. All Rights Reserved.
 #   Licensed under the Apache License, Version 2.0 (the "License"). You may not use
 #   this file except in compliance with the License. A copy of the License is located at
 #
@@ -8,14 +8,17 @@
 #   This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 #   CONDITIONS OF ANY KIND, either express or implied. See the License for the
 #   specific language governing permissions and limitations under the License.
-
-import time
 import concurrent.futures
-from .ds3 import *
+import hashlib
+import time
+import zlib
 from os import walk, path
+from platform import system
 from typing import List, Set, Dict
 
-from platform import system
+from .ds3 import *
+
+crc_byte_length = 4
 
 
 class EmptyReader(object):
@@ -87,12 +90,90 @@ def object_name_to_file_path(object_name: str) -> str:
     return object_name
 
 
+def calculate_checksum_header(object_data_stream, checksum_type: str, length: int):
+    bytes_read = 0
+    if checksum_type == 'CRC_32':
+        checksum = 0
+        while bytes_read < length:
+            bytes_to_read = min(1024 * 1024, length - bytes_read)
+            cur_bytes = object_data_stream.read(bytes_to_read)
+            checksum = zlib.crc32(cur_bytes, checksum)
+            bytes_read += bytes_to_read
+        encoded_checksum = base64.b64encode(
+            checksum.to_bytes(crc_byte_length, byteorder='big')).decode()
+        return {'Content-CRC32': encoded_checksum}
+    else:
+        if checksum_type == 'MD5':
+            checksum_calculator = hashlib.md5()
+            header_key = 'Content-MD5'
+        elif checksum_type == 'SHA_256':
+            checksum_calculator = hashlib.sha256()
+            header_key = 'Content-SHA256'
+        elif checksum_type == 'SHA_512':
+            checksum_calculator = hashlib.sha512()
+            header_key = 'Content-SHA512'
+        else:
+            raise 'Not Implemented: calculating checksum type {0} is not currently supported in the SDK helpers'.format(
+                checksum_type)
+        while bytes_read < length:
+            bytes_to_read = min(1024 * 1024, length - bytes_read)
+            cur_bytes = object_data_stream.read(bytes_to_read)
+            checksum_calculator.update(cur_bytes)
+            bytes_read += bytes_to_read
+        encoded_checksum = base64.b64encode(checksum_calculator.digest()).decode('utf-8')
+        return {header_key: encoded_checksum}
+
+
 class Helper(object):
+    """A class that moves data to and from a Black Pearl"""
+
     def __init__(self, client: Client, retry_delay_in_seconds: int = 60):
+        """
+        Parameters
+        ----------
+        client : ds3.Client
+            A python client that is connected to a Black Pearl.
+        retry_delay_in_seconds : int
+            The number of seconds to wait between retrying a call if the Black Pearl is busy and unable to process the
+            previous attempt (aka when BP returns error code 307).
+        """
         self.client = client
         self.retry_delay_in_seconds = retry_delay_in_seconds
 
-    def put_objects(self, put_objects: List[HelperPutObject], bucket: str, max_threads: int = 5) -> str:
+    def get_checksum_type(self, bucket_name: str) -> str:
+        bucket_response = self.client.get_bucket_spectra_s3(GetBucketSpectraS3Request(bucket_name=bucket_name))
+
+        data_policy_id = bucket_response.result['DataPolicyId']
+        policy_response = self.client.get_data_policy_spectra_s3(
+            GetDataPolicySpectraS3Request(data_policy_id=data_policy_id))
+
+        return policy_response.result['ChecksumType']
+
+    def put_objects(self, put_objects: List[HelperPutObject], bucket: str, max_threads: int = 5,
+                    calculate_checksum: bool = False) -> str:
+        """
+        Puts a list of objects to a Black Pearl bucket.
+
+        Parameters
+        ----------
+        put_objects : List[HelperPutObject]
+            The list of objects to put into the BP bucket.
+        bucket : str
+            The name of the bucket where the objects are being landed.
+        max_threads : int
+            The number of concurrent objects being transferred at once (default 5).
+        calculate_checksum : bool
+            Weather or not the client calculates the object checksum before sending it to the BP (default False). The BP
+            also calculates the checksum and compares it with the value the client calculates. The object put will fail
+            if the client and BP checksums do not match. Note that calculating the checksum is processor intensive, and
+            it also requires two reads of the object (first to calculate checksum, and secondly to send the data). The
+            type of checksum calculated is determined by the data policy associated with the bucket.
+        """
+        # If calculating checksum, then determine the checksum type from the data policy
+        checksum_type = None
+        if calculate_checksum is True:
+            checksum_type = self.get_checksum_type(bucket_name=bucket)
+
         ds3_put_objects: List[Ds3PutObject] = []
         put_objects_map: Dict[str, HelperPutObject] = dict()
         for entry in put_objects:
@@ -137,22 +218,56 @@ def put_objects(self, put_objects: List[HelperPutObject], bucket: str, max_threa
                             blob_set.remove(cur_blob)
                             put_object = put_objects_map[cur_blob.name]
 
-                            executor.submit(self.put_blob, bucket, put_object, cur_blob.length, cur_blob.offset, job_id)
+                            executor.submit(self.put_blob, bucket, put_object, cur_blob.length, cur_blob.offset, job_id,
+                                            checksum_type)
 
         return job_id
 
-    def put_blob(self, bucket: str, put_object: HelperPutObject, length: int, offset: int, job_id: str):
+    def put_blob(self, bucket: str, put_object: HelperPutObject, length: int, offset: int, job_id: str,
+                 checksum_type: str = None):
+        headers = None
+        if checksum_type is not None:
+            checksum_stream = put_object.get_data_stream(offset=offset)
+            headers = calculate_checksum_header(object_data_stream=checksum_stream,
+                                                checksum_type=checksum_type,
+                                                length=length)
+            checksum_stream.close()
+
         stream = put_object.get_data_stream(offset)
         self.client.put_object(PutObjectRequest(bucket_name=bucket,
                                                 object_name=put_object.object_name,
                                                 length=length,
                                                 stream=stream,
                                                 offset=offset,
-                                                job=job_id))
+                                                job=job_id,
+                                                headers=headers))
         stream.close()
 
     def put_all_objects_in_directory(self, source_dir: str, bucket: str, objects_per_bp_job: int = 1000,
-                                     max_threads: int = 5) -> List[str]:
+                                     max_threads: int = 5, calculate_checksum: bool = False) -> List[str]:
+        """
+        Puts all files and subdirectories to a Black Pearl bucket.
+
+        Parameters
+        ----------
+        source_dir : str
+            The directory that contains all the files and subdirectories to be put to the BP. Note that subdirectories
+            will be represented by zero length objects whose name ends with the file path separator.
+        bucket : str
+            The name of the bucket where the files are being landed.
+        objects_per_bp_job : int
+            The number of objects per BP job (default 1000). A directory may contain more objects than a BP job can hold
+            (max 500,000). In order to put all objects in a very large directory, multiple BP jobs may need to be
+            created. This determines how many objects to bundle per BP job.
+        max_threads : int
+            The number of concurrent objects being transferred at once (default 5).
+        calculate_checksum : bool
+            Weather or not the client calculates the object checksum before sending it to the BP. The BP also calculates
+            the checksum and compares it with the value the client calculates. The object put will fail if the client
+            and BP checksums do not match. Note that calculating the checksum is processor intensive, and it also
+            requires two reads of the object (first to calculate checksum, and secondly to send the data). The type of
+            checksum calculated is determined by the data policy associated with the bucket.
+        """
         obj_list: List[HelperPutObject] = list()
         job_list: List[str] = list()
         for root, dirs, files in walk(top=source_dir):
@@ -162,7 +277,8 @@ def put_all_objects_in_directory(self, source_dir: str, bucket: str, objects_per
                 size = os.path.getsize(obj_path)
                 obj_list.append(HelperPutObject(object_name=obj_name, file_path=obj_path, size=size))
                 if len(obj_list) >= objects_per_bp_job:
-                    job_list.append(self.put_objects(obj_list, bucket, max_threads=max_threads))
+                    job_list.append(self.put_objects(
+                        obj_list, bucket, max_threads=max_threads, calculate_checksum=calculate_checksum))
                     obj_list = []
 
             for name in dirs:
@@ -171,15 +287,29 @@ def put_all_objects_in_directory(self, source_dir: str, bucket: str, objects_per
                     path.join(path.normpath(path.relpath(path=dir_path, start=source_dir)), ""))
                 obj_list.append(HelperPutObject(object_name=dir_name, file_path=dir_path, size=0))
                 if len(obj_list) >= objects_per_bp_job:
-                    job_list.append(self.put_objects(obj_list, bucket, max_threads=max_threads))
+                    job_list.append(self.put_objects(
+                        obj_list, bucket, max_threads=max_threads, calculate_checksum=calculate_checksum))
                     obj_list = []
 
         if len(obj_list) > 0:
-            job_list.append(self.put_objects(obj_list, bucket, max_threads=max_threads))
+            job_list.append(self.put_objects(
+                obj_list, bucket, max_threads=max_threads, calculate_checksum=calculate_checksum))
 
         return job_list
 
     def get_objects(self, get_objects: List[HelperGetObject], bucket: str, max_threads: int = 5) -> str:
+        """
+        Retrieves a list of objects from a Black Pearl bucket.
+
+        Parameters
+        ----------
+        get_objects : List[HelperGetObject]
+            The list of objects to retrieve from the BP bucket.
+        bucket : str
+            The name of the bucket where the objects are being retrieved from.
+        max_threads : int
+            The number of concurrent objects being transferred at once (default 5).
+        """
         ds3_get_objects: List[Ds3GetObject] = []
         get_objects_map: Dict[str, HelperGetObject] = dict()
         for entry in get_objects:
@@ -240,6 +370,22 @@ def get_blob(self, bucket: str, get_object: HelperGetObject, offset: int, job_id
 
     def get_all_files_in_bucket(self, destination_dir: str, bucket: str, objects_per_bp_job: int = 1000,
                                 max_threads: int = 5) -> List[str]:
+        """
+        Retrieves all objects from a Black Pearl bucket.
+
+        Parameters
+        ----------
+        destination_dir : str
+            The directory where all the objects will be landed.
+        bucket : str
+            The name of the bucket where the objects are being retrieved from.
+        objects_per_bp_job : int
+            The number of objects per BP job (default 1000). A bucket may contain more objects than a BP job can hold
+            (max 500,000). In order to put all objects in a very large bucket, multiple BP jobs may need to be created.
+            This determines how many objects to bundle per BP job.
+        max_threads : int
+            The number of concurrent objects being transferred at once (default 5).
+        """
         truncated: str = 'true'
         marker = ""
         job_ids: List[str] = []
diff --git a/samples/puttingDataWithCRC32.py b/samples/puttingDataWithCRC32.py
@@ -1,4 +1,4 @@
-#   Copyright 2014-2017 Spectra Logic Corporation. All Rights Reserved.
+#   Copyright 2014-2022 Spectra Logic Corporation. All Rights Reserved.
 #   Licensed under the Apache License, Version 2.0 (the "License"). You may not use
 #   this file except in compliance with the License. A copy of the License is located at
 #
@@ -9,13 +9,13 @@
 #   CONDITIONS OF ANY KIND, either express or implied. See the License for the
 #   specific language governing permissions and limitations under the License.
 
-from ds3 import ds3
+import base64
 import os
 import time
-import sys
-import base64
 import zlib
 
+from ds3 import ds3
+
 client = ds3.createClientFromEnv()
 
 crc_byte_length = 4
@@ -83,8 +83,18 @@ def pathForResource(resourceName):
                 localFileName = "resources/" + obj['Name']
                 objectDataStream = open(localFileName, "rb")
                 objectDataStream.seek(int(obj['Offset']), 0)
-                objectChunk = objectDataStream.read(int(obj['Length']))
-                checksum = zlib.crc32(objectChunk)
+
+                # Calculate the rolling checksum by loading the file 1 MiB at a time. This allows
+                # the calculation for very large blobs that otherwise cannot be loaded fully into memory.
+                bytesRead = 0
+                checksum = 0
+                blobLength = int(obj['Length'])
+                while bytesRead < blobLength:
+                    bytesToRead = min(1024 * 1024, blobLength - bytesRead)
+                    curBytes = objectDataStream.read(bytesToRead)
+                    checksum = zlib.crc32(curBytes, checksum)
+                    bytesRead += bytesToRead
+
                 encodedChecksum = base64.b64encode(
                     checksum.to_bytes(crc_byte_length, byteorder='big')).decode()
                 objectDataStream.seek(int(obj['Offset']), 0)
diff --git a/samples/puttingDataWithMD5.py b/samples/puttingDataWithMD5.py
@@ -80,9 +80,19 @@ def pathForResource(resourceName):
                 localFileName = "resources/" + obj['Name']
                 objectDataStream = open(localFileName, "rb")
                 objectDataStream.seek(int(obj['Offset']), 0)
-                objectChunk = objectDataStream.read(int(obj['Length']))
-                checksum = hashlib.md5(objectChunk)
-                encodedChecksum = base64.b64encode(checksum.digest()).decode('utf-8')
+
+                # Calculate the rolling checksum by loading the file 1 MiB at a time. This allows
+                # the calculation for very large blobs that otherwise cannot be loaded fully into memory.
+                bytesRead = 0
+                blobLength = int(obj['Length'])
+                checksumCalculator = hashlib.md5()
+                while bytesRead < blobLength:
+                    bytesToRead = min(1024 * 1024, blobLength - bytesRead)
+                    curBytes = objectDataStream.read(bytesToRead)
+                    checksumCalculator.update(curBytes)
+                    bytesRead += bytesToRead
+
+                encodedChecksum = base64.b64encode(checksumCalculator.digest()).decode('utf-8')
                 objectDataStream.seek(int(obj['Offset']), 0)
                 client.put_object(ds3.PutObjectRequest(bucketName,
                                                        obj['Name'],
diff --git a/tests/helpersTests.py b/tests/helpersTests.py
@@ -32,7 +32,7 @@ def create_files_in_directory(directory: str, num_files: int, root_dir: str,
             obj_name = ds3Helpers.file_path_to_object_store_name(os.path.join(os.path.relpath(directory, root_dir), ""))
             put_objects.append(ds3Helpers.HelperPutObject(object_name=obj_name, file_path=directory, size=0))
 
-    # create an empty sub directory
+    # create an empty subdirectory
     if include_dirs:
         dir_path = os.path.join(directory, 'empty-dir')
         os.mkdir(path=dir_path)
@@ -155,7 +155,7 @@ def test_put_and_get_objects(self):
     def test_put_and_get_all_objects_in_directory(self):
         bucket = f'ds3-python3-sdk-test-{uuid.uuid1()}'
 
-        # create temporary directory with some files and sub directories
+        # create temporary directory with some files and subdirectories
         source = tempfile.TemporaryDirectory(prefix="ds3-python3-sdk-src-")
 
         put_objects = create_files_in_directory(directory=source.name, num_files=5, root_dir=source.name)
@@ -204,6 +204,68 @@ def test_put_and_get_all_objects_in_directory(self):
         destination.cleanup()
         client.delete_bucket_spectra_s3(ds3.DeleteBucketSpectraS3Request(bucket_name=bucket, force=True))
 
+    def test_put_all_objects_in_directory_with_md5_checksum(self):
+        self.put_all_objects_in_directory_with_checksum(checksum_type='MD5')
+
+    def test_put_all_objects_in_directory_with_crc32_checksum(self):
+        self.put_all_objects_in_directory_with_checksum(checksum_type='CRC_32')
+
+    def test_put_all_objects_in_directory_with_sha_256_checksum(self):
+        self.put_all_objects_in_directory_with_checksum(checksum_type='SHA_256')
+
+    def test_put_all_objects_in_directory_with_sha_512_checksum(self):
+        self.put_all_objects_in_directory_with_checksum(checksum_type='SHA_512')
+
+    def put_all_objects_in_directory_with_checksum(self, checksum_type: str):
+        bucket = f'ds3-python3-sdk-test-{uuid.uuid1()}'
+        # checksum_type = 'MD5'
+
+        # create the BP client
+        client = ds3.createClientFromEnv()
+
+        # create a data policy
+        data_policy = client.put_data_policy_spectra_s3(ds3.PutDataPolicySpectraS3Request(
+            name=f'sdk-test-{checksum_type}', checksum_type=checksum_type, end_to_end_crc_required=True))
+        data_policy_id = data_policy.result['Id']
+
+        # fetch existing storage domain
+        storage_domain = client.get_storage_domains_spectra_s3(ds3.GetStorageDomainsSpectraS3Request())
+        storage_domain_id = storage_domain.result['StorageDomainList'][0]['Id']
+        print("test")
+
+        data_persistence_rule = client.put_data_persistence_rule_spectra_s3(
+            ds3.PutDataPersistenceRuleSpectraS3Request(data_policy_id=data_policy_id, isolation_level='standard',
+                                                       storage_domain_id=storage_domain_id, type='permanent'))
+
+        # create temporary directory with some files and subdirectories
+        source = tempfile.TemporaryDirectory(prefix="ds3-python3-sdk-src-")
+
+        put_objects = create_files_in_directory(directory=source.name, num_files=5, root_dir=source.name,
+                                                include_dirs=True)
+
+        # create the BP helper and perform the put all objects call
+        client.put_bucket_spectra_s3(ds3.PutBucketSpectraS3Request(name=bucket, data_policy_id=data_policy_id))
+
+        helpers = ds3Helpers.Helper(client=client)
+        job_ids = helpers.put_all_objects_in_directory(source_dir=source.name, bucket=bucket, calculate_checksum=True)
+        self.assertGreaterEqual(len(job_ids), 1, "received at least one job id")
+
+        # verify all the files and directories are on the BP
+        for put_object in put_objects:
+            head_obj = client.head_object(ds3.HeadObjectRequest(bucket_name=bucket, object_name=put_object.object_name))
+            self.assertNotEqual(head_obj.result, "DOESNTEXIST")
+
+        # cleanup
+        source.cleanup()
+
+        client.delete_bucket_spectra_s3(ds3.DeleteBucketSpectraS3Request(bucket_name=bucket, force=True))
+
+        client.delete_data_persistence_rule_spectra_s3(
+            ds3.DeleteDataPersistenceRuleSpectraS3Request(data_persistence_rule_id=data_persistence_rule.result['Id']))
+
+        client.delete_data_policy_spectra_s3(
+            ds3.DeleteDataPolicySpectraS3Request(data_policy_id=data_policy_id))
+
 
 if __name__ == '__main__':
     unittest.main()
