61 lines
2.4 KiB
Python
61 lines
2.4 KiB
Python
import logging
|
|
from functools import partial
|
|
|
|
from funcy import compose
|
|
|
|
from pyinfra.file_descriptor_manager import FileDescriptorManager
|
|
from pyinfra.storage.storage import Storage
|
|
from pyinfra.utils.encoding import decompress
|
|
from pyinfra.utils.func import flift
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Downloader:
|
|
def __init__(self, storage: Storage, bucket_name, file_descriptor_manager: FileDescriptorManager):
|
|
self.storage = storage
|
|
self.bucket_name = bucket_name
|
|
self.file_descriptor_manager = file_descriptor_manager
|
|
|
|
def __call__(self, queue_item_body):
|
|
return self.download(queue_item_body)
|
|
|
|
def download(self, queue_item_body):
|
|
names_of_relevant_objects = self.get_names_of_objects_by_pattern(queue_item_body)
|
|
objects = self.download_and_decompress_object(names_of_relevant_objects)
|
|
|
|
return objects
|
|
|
|
def get_names_of_objects_by_pattern(self, queue_item_body):
|
|
logger.debug(f"Filtering objects in bucket {self.bucket_name} by pattern...")
|
|
|
|
names_of_relevant_objects = compose(
|
|
list,
|
|
self.get_pattern_filter(queue_item_body),
|
|
self.get_names_of_all_associated_objects,
|
|
)(queue_item_body)
|
|
|
|
logger.debug(f"Found {len(names_of_relevant_objects)} objects matching filter.")
|
|
|
|
return names_of_relevant_objects
|
|
|
|
def download_and_decompress_object(self, object_names):
|
|
download = partial(self.storage.get_object, self.bucket_name)
|
|
return map(compose(decompress, download), object_names)
|
|
|
|
def get_names_of_all_associated_objects(self, queue_item_body):
|
|
prefix = self.file_descriptor_manager.get_path_prefix(queue_item_body)
|
|
# TODO: performance tests for the following situations:
|
|
# 1) dossier with very many files
|
|
# 2) prefix matches very many files, independent of dossier cardinality
|
|
yield from self.storage.get_all_object_names(self.bucket_name, prefix=prefix)
|
|
|
|
def get_pattern_filter(self, queue_item_body):
|
|
print("<<<<<<<<<<<<<<<<<<<<<<<<<<<", queue_item_body)
|
|
file_pattern = self.file_descriptor_manager.build_input_matcher(queue_item_body)
|
|
print(f"{file_pattern=}")
|
|
# input()
|
|
logger.debug(f"Filtering pattern: {file_pattern if len(file_pattern) <= 120 else (file_pattern[:120]+'...')}")
|
|
matches_pattern = flift(file_pattern)
|
|
return matches_pattern
|