optimization: prefix filtering via storage API for get_all_object_names

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

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

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)

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)

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)

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

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)