Merge pull request #313 from Backblaze/new-b2-api-params

New B2Api parameters: save_to_buffer_size, check_download_hash

Author: ppolewicz

CHANGELOG.md

@@ -20,6 +20,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Add parameters to set the min/max part size for large file upload/copy methods
 * Add CopySourceTooBig exception
 * Add an option to set a custom file version class to FileVersionFactory
+* Add an option for B2Api to turn off hash checking for downloaded files
+* Add an option for B2Api to set write buffer size for DownloadedFile.save_to method
 
 ### Fixed
 * Fix copying objects larger than 1TB

CONTRIBUTING.md

@@ -88,6 +88,10 @@ To run just integration tests:
     export B2_TEST_APPLICATION_KEY_ID=your_app_key_id
     nox -s integration-3.10
 
+To run tests by keyword expressions:
+    
+    nox -s unit-3.10 -- -k keyword
+
 ## Documentation
 
 To build the documentation and watch for changes (including the source code):

b2sdk/api.py

@@ -60,6 +60,8 @@ def __init__(
         max_upload_workers: Optional[int] = None,
         max_copy_workers: Optional[int] = None,
         max_download_workers: Optional[int] = None,
+        save_to_buffer_size: Optional[int] = None,
+        check_download_hash: bool = True,
     ):
         """
         Initialize Services object using given session.
@@ -68,6 +70,8 @@ def __init__(
         :param max_upload_workers: a number of upload threads
         :param max_copy_workers: a number of copy threads
         :param max_download_workers: maximum number of download threads
+        :param save_to_buffer_size: buffer size to use when writing files using DownloadedFile.save_to
+        :param check_download_hash: whether to check hash of downloaded files. Can be disabled for files with internal checksums, for example, or to forcefully retrieve objects with corrupted payload or hash value
         """
         self.api = api
         self.session = api.session
@@ -77,7 +81,10 @@ def __init__(
         )
         self.copy_manager = self.COPY_MANAGER_CLASS(services=self, max_workers=max_copy_workers)
         self.download_manager = self.DOWNLOAD_MANAGER_CLASS(
-            services=self, max_workers=max_download_workers
+            services=self,
+            max_workers=max_download_workers,
+            write_buffer_size=save_to_buffer_size,
+            check_hash=check_download_hash,
         )
         self.emerger = Emerger(self)
 
@@ -118,6 +125,8 @@ def __init__(
         max_copy_workers: Optional[int] = None,
         api_config: B2HttpApiConfig = DEFAULT_HTTP_API_CONFIG,
         max_download_workers: Optional[int] = None,
+        save_to_buffer_size: Optional[int] = None,
+        check_download_hash: bool = True,
     ):
         """
         Initialize the API using the given account info.
@@ -132,6 +141,8 @@ def __init__(
         :param max_copy_workers: a number of copy threads
         :param api_config:
         :param max_download_workers: maximum number of download threads
+        :param save_to_buffer_size: buffer size to use when writing files using DownloadedFile.save_to
+        :param check_download_hash: whether to check hash of downloaded files. Can be disabled for files with internal checksums, for example, or to forcefully retrieve objects with corrupted payload or hash value
         """
         self.session = self.SESSION_CLASS(
             account_info=account_info, cache=cache, api_config=api_config
@@ -143,6 +154,8 @@ def __init__(
             max_upload_workers=max_upload_workers,
             max_copy_workers=max_copy_workers,
             max_download_workers=max_download_workers,
+            save_to_buffer_size=save_to_buffer_size,
+            check_download_hash=check_download_hash,
         )
 
     @property

b2sdk/transfer/inbound/download_manager.py

@@ -44,7 +44,7 @@ class DownloadManager(TransferManager, ThreadPoolMixin, metaclass=B2TraceMetaAbs
     PARALLEL_DOWNLOADER_CLASS = staticmethod(ParallelDownloader)
     SIMPLE_DOWNLOADER_CLASS = staticmethod(SimpleDownloader)
 
-    def __init__(self, **kwargs):
+    def __init__(self, write_buffer_size: Optional[int] = None, check_hash: bool = True, **kwargs):
         """
         Initialize the DownloadManager using the given services object.
         """
@@ -54,15 +54,21 @@ def __init__(self, **kwargs):
             self.PARALLEL_DOWNLOADER_CLASS(
                 min_part_size=self.DEFAULT_MIN_PART_SIZE,
                 min_chunk_size=self.MIN_CHUNK_SIZE,
-                max_chunk_size=self.MAX_CHUNK_SIZE,
+                max_chunk_size=max(self.MAX_CHUNK_SIZE, write_buffer_size or 0),
+                align_factor=write_buffer_size,
                 thread_pool=self._thread_pool,
+                check_hash=check_hash,
             ),
             self.SIMPLE_DOWNLOADER_CLASS(
                 min_chunk_size=self.MIN_CHUNK_SIZE,
-                max_chunk_size=self.MAX_CHUNK_SIZE,
+                max_chunk_size=max(self.MAX_CHUNK_SIZE, write_buffer_size or 0),
+                align_factor=write_buffer_size,
                 thread_pool=self._thread_pool,
+                check_hash=check_hash,
             ),
         ]
+        self.write_buffer_size = write_buffer_size
+        self.check_hash = check_hash
 
     def download_file_from_url(
         self,
@@ -98,4 +104,6 @@ def download_file_from_url(
                 response=response,
                 encryption=encryption,
                 progress_listener=progress_listener,
+                write_buffer_size=self.write_buffer_size,
+                check_hash=self.check_hash,
             )

b2sdk/transfer/inbound/downloaded_file.py

@@ -44,9 +44,10 @@ class MtimeUpdatedFile(io.IOBase):
        #  'some_local_path' has the mod_time set according to metadata in B2
     """
 
-    def __init__(self, path_, mod_time_millis: int, mode='wb+'):
+    def __init__(self, path_, mod_time_millis: int, mode='wb+', buffering=None):
         self.path_ = path_
         self.mode = mode
+        self.buffering = buffering if buffering is not None else -1
         self.mod_time_to_set = mod_time_millis
         self.file = None
 
@@ -69,7 +70,7 @@ def tell(self):
         return self.file.tell()
 
     def __enter__(self):
-        self.file = open(self.path_, self.mode)
+        self.file = open(self.path_, self.mode, buffering=self.buffering)
         self.write = self.file.write
         self.read = self.file.read
         return self
@@ -93,6 +94,8 @@ def __init__(
         response: Response,
         encryption: Optional[EncryptionSetting],
         progress_listener: AbstractProgressListener,
+        write_buffer_size=None,
+        check_hash=True,
     ):
         self.download_version = download_version
         self.download_manager = download_manager
@@ -101,13 +104,18 @@ def __init__(
         self.encryption = encryption
         self.progress_listener = progress_listener
         self.download_strategy = None
+        self.write_buffer_size = write_buffer_size
+        self.check_hash = check_hash
 
     def _validate_download(self, bytes_read, actual_sha1):
         if self.range_ is None:
             if bytes_read != self.download_version.content_length:
                 raise TruncatedOutput(bytes_read, self.download_version.content_length)
 
-            if self.download_version.content_sha1 != 'none' and actual_sha1 != self.download_version.content_sha1:
+            if (
+                self.check_hash and self.download_version.content_sha1 != 'none' and
+                actual_sha1 != self.download_version.content_sha1
+            ):
                 raise ChecksumMismatch(
                     checksum_type='sha1',
                     expected=self.download_version.content_sha1,
@@ -158,6 +166,9 @@ def save_to(self, path_, mode='wb+', allow_seeking=True):
                               (parallel strategies) will be discarded.
         """
         with MtimeUpdatedFile(
-            path_, mod_time_millis=self.download_version.mod_time_millis, mode=mode
+            path_,
+            mod_time_millis=self.download_version.mod_time_millis,
+            mode=mode,
+            buffering=self.write_buffer_size,
         ) as file:
             self.save(file, allow_seeking=allow_seeking)

b2sdk/transfer/inbound/downloader/abstract.py

@@ -8,6 +8,8 @@
 #
 ######################################################################
 
+import hashlib
+
 from abc import abstractmethod
 from concurrent.futures import ThreadPoolExecutor
 from io import IOBase
@@ -22,36 +24,64 @@
 from b2sdk.encryption.setting import EncryptionSetting
 
 
+class EmptyHasher:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def update(self, data):
+        pass
+
+    def digest(self):
+        return b''
+
+    def hexdigest(self):
+        return ''
+
+    def copy(self):
+        return self
+
+
 class AbstractDownloader(metaclass=B2TraceMetaAbstract):
 
     REQUIRES_SEEKING = True
     DEFAULT_THREAD_POOL_CLASS = staticmethod(ThreadPoolExecutor)
+    DEFAULT_ALIGN_FACTOR = 4096
 
     def __init__(
         self,
         thread_pool: Optional[ThreadPoolExecutor] = None,
-        force_chunk_size=None,
-        min_chunk_size=None,
-        max_chunk_size=None,
+        force_chunk_size: Optional[int] = None,
+        min_chunk_size: Optional[int] = None,
+        max_chunk_size: Optional[int] = None,
+        align_factor: Optional[int] = None,
+        check_hash: bool = True,
         **kwargs
     ):
+        align_factor = align_factor or self.DEFAULT_ALIGN_FACTOR
         assert force_chunk_size is not None or (
             min_chunk_size is not None and max_chunk_size is not None and
-            0 < min_chunk_size <= max_chunk_size
+            0 < min_chunk_size <= max_chunk_size and max_chunk_size >= align_factor
         )
         self._min_chunk_size = min_chunk_size
         self._max_chunk_size = max_chunk_size
         self._forced_chunk_size = force_chunk_size
+        self._align_factor = align_factor
+        self._check_hash = check_hash
         self._thread_pool = thread_pool if thread_pool is not None \
             else self.DEFAULT_THREAD_POOL_CLASS()
         super().__init__(**kwargs)
 
-    def _get_chunk_size(self, content_length):
+    def _get_hasher(self):
+        if self._check_hash:
+            return hashlib.sha1()
+        return EmptyHasher()
+
+    def _get_chunk_size(self, content_length: Optional[int]):
         if self._forced_chunk_size is not None:
             return self._forced_chunk_size
-        ideal = content_length // 1000
+        ideal = max(content_length // 1000, self._align_factor)
         non_aligned = min(max(ideal, self._min_chunk_size), self._max_chunk_size)
-        aligned = non_aligned // 4096 * 4096
+        aligned = non_aligned // self._align_factor * self._align_factor
         return aligned
 
     @classmethod

b2sdk/transfer/inbound/downloader/parallel.py

@@ -12,7 +12,6 @@
 from io import IOBase
 from typing import Optional
 import logging
-import hashlib
 import queue
 import threading
 
@@ -99,7 +98,7 @@ def download(
 
         first_part = parts_to_download[0]
 
-        hasher = hashlib.sha1()
+        hasher = self._get_hasher()
 
         with WriterThread(file, max_queue_depth=len(parts_to_download) * 2) as writer:
             self._get_parts(
@@ -116,7 +115,10 @@ def download(
 
         # At this point the hasher already consumed the data until the end of first stream.
         # Consume the rest of the file to complete the hashing process
-        self._finish_hashing(first_part, file, hasher, download_version.content_length)
+        if self._check_hash:
+            # we skip hashing if we would not check it - hasher object is actually a EmptyHasher instance
+            # but we avoid here reading whole file (except for the first part) from disk again
+            self._finish_hashing(first_part, file, hasher, download_version.content_length)
 
         return bytes_written, hasher.hexdigest()
 

b2sdk/transfer/inbound/downloader/simple.py

@@ -10,7 +10,6 @@
 
 from io import IOBase
 from typing import Optional
-import hashlib
 import logging
 
 from requests.models import Response
@@ -38,7 +37,8 @@ def _download(
         actual_size = self._get_remote_range(response, download_version).size()
         chunk_size = self._get_chunk_size(actual_size)
 
-        digest = hashlib.sha1()
+        digest = self._get_hasher()
+
         bytes_read = 0
         for data in response.iter_content(chunk_size=chunk_size):
             file.write(data)