###################################################################### # # File: b2sdk/_internal/transfer/emerge/executor.py # # Copyright 2020 Backblaze Inc. All Rights Reserved. # # License https://www.backblaze.com/using_b2_code.html # ###################################################################### from __future__ import annotations import logging import threading from abc import ABCMeta, abstractmethod from typing import TYPE_CHECKING from b2sdk._internal.encryption.setting import EncryptionSetting from b2sdk._internal.exception import MaxFileSizeExceeded from b2sdk._internal.file_lock import NO_RETENTION_FILE_SETTING, FileRetentionSetting, LegalHold from b2sdk._internal.http_constants import LARGE_FILE_SHA1 from b2sdk._internal.transfer.outbound.large_file_upload_state import LargeFileUploadState from b2sdk._internal.transfer.outbound.upload_source import UploadSourceStream AUTO_CONTENT_TYPE = 'b2/x-auto' logger = logging.getLogger(__name__) if TYPE_CHECKING: from b2sdk._internal.transfer.emerge.planner.part_definition import UploadEmergePartDefinition from b2sdk._internal.transfer.emerge.planner.planner import StreamingEmergePlan class EmergeExecutor: def __init__(self, services): self.services = services def execute_emerge_plan( self, emerge_plan, bucket_id, file_name, content_type, file_info, progress_listener, continue_large_file_id=None, max_queue_size=None, encryption: EncryptionSetting | None = None, file_retention: FileRetentionSetting | None = None, legal_hold: LegalHold | None = None, custom_upload_timestamp: int | None = None, ): if emerge_plan.is_large_file(): execution = LargeFileEmergeExecution( self.services, bucket_id, file_name, content_type, file_info, progress_listener, encryption=encryption, file_retention=file_retention, legal_hold=legal_hold, continue_large_file_id=continue_large_file_id, max_queue_size=max_queue_size, custom_upload_timestamp=custom_upload_timestamp, ) else: if continue_large_file_id is not None: raise ValueError('Cannot resume emerging single part plan.') execution = SmallFileEmergeExecution( self.services, bucket_id, file_name, content_type, file_info, progress_listener, encryption=encryption, file_retention=file_retention, legal_hold=legal_hold, custom_upload_timestamp=custom_upload_timestamp, ) return execution.execute_plan(emerge_plan) class BaseEmergeExecution(metaclass=ABCMeta): DEFAULT_CONTENT_TYPE = AUTO_CONTENT_TYPE def __init__( self, services, bucket_id, file_name, content_type, file_info, progress_listener, encryption: EncryptionSetting | None = None, file_retention: FileRetentionSetting | None = None, legal_hold: LegalHold | None = None, custom_upload_timestamp: int | None = None, ): self.services = services self.bucket_id = bucket_id self.file_name = file_name self.content_type = content_type self.file_info = file_info self.progress_listener = progress_listener self.encryption = encryption self.file_retention = file_retention self.legal_hold = legal_hold self.custom_upload_timestamp = custom_upload_timestamp @abstractmethod def execute_plan(self, emerge_plan): pass class SmallFileEmergeExecution(BaseEmergeExecution): def execute_plan(self, emerge_plan): emerge_parts = list(emerge_plan.emerge_parts) assert len(emerge_parts) == 1 emerge_part = emerge_parts[0] execution_step_factory = SmallFileEmergeExecutionStepFactory(self, emerge_part) execution_step = execution_step_factory.get_execution_step() future = execution_step.execute() return future.result() class LargeFileEmergeExecution(BaseEmergeExecution): MAX_LARGE_FILE_SIZE = 10 * 1000 * 1000 * 1000 * 1000 # 10 TB def __init__( self, services, bucket_id, file_name, content_type, file_info, progress_listener, encryption: EncryptionSetting | None = None, file_retention: FileRetentionSetting | None = None, legal_hold: LegalHold | None = None, continue_large_file_id=None, max_queue_size=None, custom_upload_timestamp: int | None = None, ): super().__init__( services, bucket_id, file_name, content_type, file_info, progress_listener, encryption=encryption, file_retention=file_retention, legal_hold=legal_hold, custom_upload_timestamp=custom_upload_timestamp, ) self.continue_large_file_id = continue_large_file_id self.max_queue_size = max_queue_size self._semaphore = None if self.max_queue_size is not None: self._semaphore = threading.Semaphore(self.max_queue_size) def execute_plan(self, emerge_plan: StreamingEmergePlan): total_length = emerge_plan.get_total_length() encryption = self.encryption if total_length is not None and total_length > self.MAX_LARGE_FILE_SIZE: raise MaxFileSizeExceeded(total_length, self.MAX_LARGE_FILE_SIZE) plan_id = emerge_plan.get_plan_id() file_info = dict(self.file_info) if self.file_info is not None else {} if plan_id is not None: file_info['plan_id'] = plan_id self.progress_listener.set_total_bytes(total_length or 0) emerge_parts_dict = None if total_length is not None: emerge_parts_dict = dict(emerge_plan.enumerate_emerge_parts()) unfinished_file, finished_parts = self._get_unfinished_file_and_parts( self.bucket_id, self.file_name, file_info, self.continue_large_file_id, encryption=encryption, file_retention=self.file_retention, legal_hold=self.legal_hold, emerge_parts_dict=emerge_parts_dict, custom_upload_timestamp=self.custom_upload_timestamp, ) if unfinished_file is None: if self.content_type is None: content_type = self.DEFAULT_CONTENT_TYPE else: content_type = self.content_type unfinished_file = self.services.large_file.start_large_file( self.bucket_id, self.file_name, content_type, file_info, encryption=encryption, file_retention=self.file_retention, legal_hold=self.legal_hold, ) file_id = unfinished_file.file_id large_file_upload_state = LargeFileUploadState(self.progress_listener) part_futures = [] for part_number, emerge_part in emerge_plan.enumerate_emerge_parts(): execution_step_factory = LargeFileEmergeExecutionStepFactory( self, emerge_part, part_number, file_id, large_file_upload_state, finished_parts=finished_parts, # it already knows encryption from BaseMergeExecution being passed as self ) execution_step = execution_step_factory.get_execution_step() future = self._execute_step(execution_step) part_futures.append(future) # Collect the sha1 checksums of the parts as the uploads finish. # If any of them raised an exception, that same exception will # be raised here by result() part_sha1_array = [f.result()['contentSha1'] for f in part_futures] # Finish the large file response = self.services.session.finish_large_file(file_id, part_sha1_array) return self.services.api.file_version_factory.from_api_response(response) def _execute_step(self, execution_step: UploadPartExecutionStep): semaphore = self._semaphore if semaphore is None: return execution_step.execute() else: semaphore.acquire() try: future = execution_step.execute() except: semaphore.release() raise else: future.add_done_callback(lambda f: semaphore.release()) return future def _get_unfinished_file_and_parts( self, bucket_id, file_name, file_info, continue_large_file_id, encryption: EncryptionSetting, file_retention: FileRetentionSetting | None = None, legal_hold: LegalHold | None = None, emerge_parts_dict=None, custom_upload_timestamp: int | None = None, ): if 'listFiles' not in self.services.session.account_info.get_allowed()['capabilities']: return None, {} unfinished_file = None finished_parts = {} if continue_large_file_id is not None: unfinished_file = self.services.large_file.get_unfinished_large_file( bucket_id, continue_large_file_id, prefix=file_name, ) if unfinished_file.file_info != file_info: raise ValueError( 'Cannot manually resume unfinished large file with different file_info' ) finished_parts = { part.part_number: part for part in self.services.large_file.list_parts(continue_large_file_id) } elif 'plan_id' in file_info: assert emerge_parts_dict is not None unfinished_file, finished_parts = self._find_unfinished_file_by_plan_id( bucket_id, file_name, file_info, emerge_parts_dict, encryption, file_retention, legal_hold, custom_upload_timestamp=custom_upload_timestamp, ) elif emerge_parts_dict is not None: unfinished_file, finished_parts = self._match_unfinished_file_if_possible( bucket_id, file_name, file_info, emerge_parts_dict, encryption, file_retention, legal_hold, custom_upload_timestamp=custom_upload_timestamp, ) return unfinished_file, finished_parts def _find_matching_unfinished_file( self, bucket_id, file_name, file_info, emerge_parts_dict, encryption: EncryptionSetting, file_retention: FileRetentionSetting | None = None, legal_hold: LegalHold | None = None, custom_upload_timestamp: int | None = None, check_file_info_without_large_file_sha1: bool | None = False, eager_mode: bool | None = False, ): """ Search for a matching unfinished large file in the specified bucket. In case a matching file is found but has inconsistencies (for example, mismatching file info or encryption settings), mismatches are logged. :param bucket_id: The identifier of the bucket where the unfinished file resides. :param file_name: The name of the file to be matched. :param file_info: Information about the file to be uploaded. :param emerge_parts_dict: A dictionary containing the parts of the file to be emerged. :param encryption: The encryption settings for the file. :param file_retention: The retention settings for the file, if any. :param legal_hold: The legal hold status of the file, if any. :param custom_upload_timestamp: The custom timestamp for the upload, if any. :param check_file_info_without_large_file_sha1: A flag indicating whether the file information should be checked without the `large_file_sha1`. :param eager_mode: A flag indicating whether the first matching file should be returned. :return: A tuple of the best matching unfinished file and its finished parts. If no match is found, returns `None`. """ file_retention = file_retention or NO_RETENTION_FILE_SETTING best_match_file = None best_match_parts = {} best_match_parts_len = 0 for file_ in self.services.large_file.list_unfinished_large_files( bucket_id, prefix=file_name ): if file_.file_name != file_name: logger.debug('Rejecting %s: file name mismatch', file_.file_id) continue if file_.file_info != file_info: if check_file_info_without_large_file_sha1: file_info_without_large_file_sha1 = self._get_file_info_without_large_file_sha1( file_info ) if ( file_info_without_large_file_sha1 != self._get_file_info_without_large_file_sha1(file_.file_info) ): logger.debug( 'Rejecting %s: file info mismatch after dropping `large_file_sha1`', file_.file_id, ) continue else: logger.debug('Rejecting %s: file info mismatch', file_.file_id) continue if encryption is not None and encryption != file_.encryption: logger.debug('Rejecting %s: encryption mismatch', file_.file_id) continue if legal_hold is None: if LegalHold.UNSET != file_.legal_hold: logger.debug('Rejecting %s: legal hold mismatch (not unset)', file_.file_id) continue elif legal_hold != file_.legal_hold: logger.debug('Rejecting %s: legal hold mismatch', file_.file_id) continue if file_retention != file_.file_retention: logger.debug('Rejecting %s: retention mismatch', file_.file_id) continue if ( custom_upload_timestamp is not None and file_.upload_timestamp != custom_upload_timestamp ): logger.debug('Rejecting %s: custom_upload_timestamp mismatch', file_.file_id) continue finished_parts = {} conflict_detected = False for part in self.services.large_file.list_parts(file_.file_id): emerge_part = emerge_parts_dict.get(part.part_number) if emerge_part is None: # something is wrong - we have a part that we don't know about # so we can't resume this upload logger.debug( 'Rejecting %s: part %s not found in emerge parts, giving up.', file_.file_id, part.part_number, ) conflict_detected = True break # Compare part sizes if emerge_part.get_length() != part.content_length: logger.debug( 'Rejecting %s: part %s size mismatch', file_.file_id, part.part_number ) conflict_detected = True break # part size doesn't match - so we reupload # Compare part hashes if emerge_part.is_hashable() and emerge_part.get_sha1() != part.content_sha1: logger.debug( 'Rejecting %s: part %s sha1 mismatch', file_.file_id, part.part_number ) conflict_detected = True break # part.sha1 doesn't match - so we reupload finished_parts[part.part_number] = part if conflict_detected: continue finished_parts_len = len(finished_parts) if best_match_file is None or finished_parts_len > best_match_parts_len: best_match_file = file_ best_match_parts = finished_parts best_match_parts_len = finished_parts_len if eager_mode and best_match_file is not None: break return best_match_file, best_match_parts def _find_unfinished_file_by_plan_id( self, bucket_id, file_name, file_info, emerge_parts_dict, encryption: EncryptionSetting, file_retention: FileRetentionSetting | None = None, legal_hold: LegalHold | None = None, custom_upload_timestamp: int | None = None, ): """ Search for a matching unfinished large file by plan_id in the specified bucket. This function aims to locate a matching unfinished large file using the plan_id and the supplied parameters. It's used to resume an interrupted upload, centralizing the shared logic between `_find_unfinished_file_by_plan_id` and `_match_unfinished_file_if_possible`. In case a matching file is found but has inconsistencies (for example, mismatching file info or encryption settings), the function checks if 'plan_id' is in file_info, as this is a prerequisite. :param bucket_id: The identifier of the bucket where the unfinished file resides. :param file_name: The name of the file to be matched. :param file_info: Information about the file to be uploaded. :param emerge_parts_dict: A dictionary containing the parts of the file to be emerged. :param encryption: The encryption settings for the file. :param file_retention: The retention settings for the file, if any. :param legal_hold: The legal hold status of the file, if any. :param custom_upload_timestamp: The custom timestamp for the upload, if any. :return: A tuple of the best matching unfinished file and its finished parts. If no match is found, it returns `None`. """ if 'plan_id' not in file_info: raise ValueError("The 'plan_id' key must be in file_info dictionary.") return self._find_matching_unfinished_file( bucket_id=bucket_id, file_name=file_name, file_info=file_info, emerge_parts_dict=emerge_parts_dict, encryption=encryption, file_retention=file_retention or NO_RETENTION_FILE_SETTING, legal_hold=legal_hold, custom_upload_timestamp=custom_upload_timestamp, check_file_info_without_large_file_sha1=False, ) @classmethod def _get_file_info_without_large_file_sha1( cls, file_info: dict[str, str] | None, ) -> dict[str, str] | None: if not file_info or LARGE_FILE_SHA1 not in file_info: return file_info out_file_info = dict(file_info) del out_file_info[LARGE_FILE_SHA1] return out_file_info def _match_unfinished_file_if_possible( self, bucket_id, file_name, file_info, emerge_parts_dict, encryption: EncryptionSetting, file_retention: FileRetentionSetting | None = None, legal_hold: LegalHold | None = None, custom_upload_timestamp: int | None = None, ): """ Scan for a suitable unfinished large file in the specified bucket to resume upload. This function examines each unfinished large file for a possible match with the provided parameters. This enables resumption of an interrupted upload by reusing the unfinished file, provided that file's info and additional parameters match. Along with the filename and file info, additional parameters like encryption, file retention, legal hold, custom upload timestamp, and cache control are compared for a match. The 'emerge_parts_dict' is also cross-checked for matching file parts. Function is eager to find a match, and will return the first match it finds. If no match is found, it returns `None`. :param bucket_id: The identifier of the bucket containing the unfinished file. :param file_name: The name of the file to find. :param file_info: Information about the file to be uploaded. :param emerge_parts_dict: A dictionary of the parts of the file to be emerged. :param encryption: The encryption settings for the file. :param file_retention: The retention settings for the file, if applicable. :param legal_hold: The legal hold status of the file, if applicable. :param custom_upload_timestamp: The custom timestamp for the upload, if set. :return: A tuple of the best matching unfinished file and its finished parts. If no match is found, returns `None`. """ logger.debug('Checking for matching unfinished large files for %s...', file_name) file_, finished_parts = self._find_matching_unfinished_file( bucket_id, file_name, file_info, emerge_parts_dict, encryption, file_retention, legal_hold, custom_upload_timestamp, check_file_info_without_large_file_sha1=True, eager_mode=True, ) if file_ is None: logger.debug('No matching unfinished files found.') return None, {} logger.debug( 'Unfinished file %s matches with %i finished parts', file_.file_id, len(finished_parts) ) return file_, finished_parts class BaseExecutionStepFactory(metaclass=ABCMeta): def __init__(self, emerge_execution, emerge_part): self.emerge_execution = emerge_execution self.emerge_part = emerge_part @abstractmethod def create_copy_execution_step(self, copy_range): pass @abstractmethod def create_upload_execution_step(self, stream_opener, stream_length=None, stream_sha1=None): pass def get_execution_step(self): return self.emerge_part.get_execution_step(self) class SmallFileEmergeExecutionStepFactory(BaseExecutionStepFactory): def create_copy_execution_step(self, copy_range): return CopyFileExecutionStep(self.emerge_execution, copy_range) def create_upload_execution_step(self, stream_opener, stream_length=None, stream_sha1=None): return UploadFileExecutionStep( self.emerge_execution, stream_opener, stream_length=stream_length, stream_sha1=stream_sha1, ) class LargeFileEmergeExecutionStepFactory(BaseExecutionStepFactory): def __init__( self, emerge_execution, emerge_part: UploadEmergePartDefinition, part_number, large_file_id, large_file_upload_state, finished_parts=None, ): super().__init__(emerge_execution, emerge_part) self.part_number = part_number self.large_file_id = large_file_id self.large_file_upload_state = large_file_upload_state self.finished_parts = finished_parts or {} def create_copy_execution_step(self, copy_range): return CopyPartExecutionStep( self.emerge_execution, copy_range, self.part_number, self.large_file_id, self.large_file_upload_state, ) def create_upload_execution_step(self, stream_opener, stream_length=None, stream_sha1=None): return UploadPartExecutionStep( self.emerge_execution, stream_opener, self.part_number, self.large_file_id, self.large_file_upload_state, stream_length=stream_length, stream_sha1=stream_sha1, finished_parts=self.finished_parts, ) class BaseExecutionStep(metaclass=ABCMeta): def __init__(self, emerge_execution: BaseEmergeExecution): self.emerge_execution = emerge_execution @abstractmethod def execute(self): pass class CopyFileExecutionStep(BaseExecutionStep): def __init__(self, emerge_execution, copy_source_range): super().__init__(emerge_execution) self.copy_source_range = copy_source_range def execute(self): execution = self.emerge_execution # if content type is not None then we support empty dict as default file info # but if content type is None, then setting empty dict as file info # would result with an error, because default in such case is: copy from source if execution.content_type is not None: file_info = execution.file_info or {} else: file_info = None return execution.services.copy_manager.copy_file( self.copy_source_range, execution.file_name, content_type=execution.content_type, file_info=file_info, destination_bucket_id=execution.bucket_id, progress_listener=execution.progress_listener, destination_encryption=execution.encryption, source_encryption=self.copy_source_range.encryption, file_retention=execution.file_retention, legal_hold=execution.legal_hold, ) class CopyPartExecutionStep(BaseExecutionStep): def __init__( self, emerge_execution, copy_source_range, part_number, large_file_id, large_file_upload_state, finished_parts=None, ): super().__init__(emerge_execution) self.copy_source_range = copy_source_range self.part_number = part_number self.large_file_id = large_file_id self.large_file_upload_state = large_file_upload_state self.finished_parts = finished_parts or {} def execute(self): return self.emerge_execution.services.copy_manager.copy_part( self.large_file_id, self.copy_source_range, self.part_number, self.large_file_upload_state, finished_parts=self.finished_parts, destination_encryption=self.emerge_execution.encryption, source_encryption=self.copy_source_range.encryption, ) class UploadFileExecutionStep(BaseExecutionStep): def __init__(self, emerge_execution, stream_opener, stream_length=None, stream_sha1=None): super().__init__(emerge_execution) self.stream_opener = stream_opener self.stream_length = stream_length self.stream_sha1 = stream_sha1 def execute(self): upload_source = UploadSourceStream( self.stream_opener, stream_length=self.stream_length, stream_sha1=self.stream_sha1, ) execution = self.emerge_execution return execution.services.upload_manager.upload_file( execution.bucket_id, upload_source, execution.file_name, execution.content_type or execution.DEFAULT_CONTENT_TYPE, execution.file_info or {}, execution.progress_listener, encryption=execution.encryption, file_retention=execution.file_retention, legal_hold=execution.legal_hold, custom_upload_timestamp=execution.custom_upload_timestamp, ) class UploadPartExecutionStep(BaseExecutionStep): def __init__( self, emerge_execution, stream_opener, part_number, large_file_id, large_file_upload_state, stream_length=None, stream_sha1=None, finished_parts=None, ): super().__init__(emerge_execution) self.stream_opener = stream_opener self.stream_length = stream_length self.stream_sha1 = stream_sha1 self.part_number = part_number self.large_file_id = large_file_id self.large_file_upload_state = large_file_upload_state self.finished_parts = finished_parts or {} def execute(self): execution = self.emerge_execution upload_source = UploadSourceStream( self.stream_opener, stream_length=self.stream_length, stream_sha1=self.stream_sha1, ) return execution.services.upload_manager.upload_part( execution.bucket_id, self.large_file_id, upload_source, self.part_number, self.large_file_upload_state, finished_parts=self.finished_parts, encryption=execution.encryption, )