diff --git a/pyinfra/file_descriptor_manager.py b/pyinfra/file_descriptor_manager.py index ab2f18e..1ca8292 100644 --- a/pyinfra/file_descriptor_manager.py +++ b/pyinfra/file_descriptor_manager.py @@ -87,3 +87,12 @@ class FileDescriptorManager: def build_storage_upload_info(analysis_payload, request_metadata): storage_upload_info = {**request_metadata, "id": analysis_payload["metadata"].get("id", 0)} return storage_upload_info + + def get_path_prefix(self, queue_item_body): + + prefix = "/".join( + itemgetter("dossierId", "fileId")( + self.build_file_descriptor(queue_item_body, end="input"), + ) + ) + return prefix diff --git a/pyinfra/storage/adapters/adapter.py b/pyinfra/storage/adapters/adapter.py index 57cff28..7eb31a9 100644 --- a/pyinfra/storage/adapters/adapter.py +++ b/pyinfra/storage/adapters/adapter.py @@ -30,5 +30,5 @@ class StorageAdapter(ABC): raise NotImplementedError @abstractmethod - def get_all_object_names(self, bucket_name): + def get_all_object_names(self, bucket_name, prefix=None): raise NotImplementedError diff --git a/pyinfra/storage/adapters/azure.py b/pyinfra/storage/adapters/azure.py index b3d498c..074a9d6 100644 --- a/pyinfra/storage/adapters/azure.py +++ b/pyinfra/storage/adapters/azure.py @@ -58,7 +58,7 @@ class AzureStorageAdapter(StorageAdapter): blobs = container_client.list_blobs() container_client.delete_blobs(*blobs) - def get_all_object_names(self, bucket_name): + def get_all_object_names(self, bucket_name, prefix=None): container_client = self.__provide_container_client(bucket_name) - blobs = container_client.list_blobs() + blobs = container_client.list_blobs(name_starts_with=prefix) return map(attrgetter("name"), blobs) diff --git a/pyinfra/storage/adapters/s3.py b/pyinfra/storage/adapters/s3.py index 2fe5dcc..a7729df 100644 --- a/pyinfra/storage/adapters/s3.py +++ b/pyinfra/storage/adapters/s3.py @@ -53,6 +53,6 @@ class S3StorageAdapter(StorageAdapter): for obj in objects: self.__client.remove_object(bucket_name, obj.object_name) - def get_all_object_names(self, bucket_name): - objs = self.__client.list_objects(bucket_name, recursive=True) + def get_all_object_names(self, bucket_name, prefix=None): + objs = self.__client.list_objects(bucket_name, recursive=True, prefix=prefix) return map(attrgetter("object_name"), objs) diff --git a/pyinfra/storage/storage.py b/pyinfra/storage/storage.py index a0026c3..2bc4eff 100644 --- a/pyinfra/storage/storage.py +++ b/pyinfra/storage/storage.py @@ -40,5 +40,5 @@ class Storage: def clear_bucket(self, bucket_name): return self.__adapter.clear_bucket(bucket_name) - def get_all_object_names(self, bucket_name): - return self.__adapter.get_all_object_names(bucket_name) + def get_all_object_names(self, bucket_name, prefix=None): + return self.__adapter.get_all_object_names(bucket_name, prefix=prefix) diff --git a/pyinfra/visitor/downloader.py b/pyinfra/visitor/downloader.py index 230e392..ec5573e 100644 --- a/pyinfra/visitor/downloader.py +++ b/pyinfra/visitor/downloader.py @@ -1,3 +1,4 @@ +import logging from functools import partial from funcy import compose @@ -7,6 +8,8 @@ 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): @@ -17,19 +20,36 @@ class Downloader: def __call__(self, queue_item_body): return self.download(queue_item_body) - def get_names_of_objects_by_pattern(self, file_pattern: str): - matches_pattern = flift(file_pattern) - page_object_names = compose(matches_pattern, self.storage.get_all_object_names)(self.bucket_name) - return page_object_names + 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( + self.get_pattern_filter(queue_item_body), + self.get_names_of_all_associated_objects, + )(queue_item_body) + + logger.debug(f"Completed filtering.") + + 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 download(self, queue_item_body): + 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): file_pattern = self.file_descriptor_manager.build_input_matcher(queue_item_body) - - page_object_names = self.get_names_of_objects_by_pattern(file_pattern) - objects = self.download_and_decompress_object(page_object_names) - - return objects + matches_pattern = flift(file_pattern) + return matches_pattern diff --git a/test/storage/adapter_mock.py b/test/storage/adapter_mock.py index 5443ade..6b60a27 100644 --- a/test/storage/adapter_mock.py +++ b/test/storage/adapter_mock.py @@ -26,5 +26,5 @@ class StorageAdapterMock(StorageAdapter): def clear_bucket(self, bucket_name): return self.__client.clear_bucket(bucket_name) - def get_all_object_names(self, bucket_name): + def get_all_object_names(self, bucket_name, prefix=None): return self.__client.get_all_object_names(bucket_name)