diff --git a/config.yaml b/config.yaml index 75a0408..acad3be 100755 --- a/config.yaml +++ b/config.yaml @@ -31,6 +31,8 @@ service: name: $SERVICE_NAME|pyinfra-service-v1 # Name of the service in the kubernetes cluster storage_backend: $STORAGE_BACKEND|s3 # The storage to pull files to be processed from analysis_endpoint: $ANALYSIS_ENDPOINT|"http://127.0.0.1:5000" + bucket_name: $STORAGE_BUCKET_NAME|"pyinfra-test-bucket" # TODO Maybe refactor! + target_file_extension: $TARGET_FILE_EXTENSION|"pdf" # Defines type of file pulled from storage probing_webserver: host: $PROBING_WEBSERVER_HOST|"0.0.0.0" # Probe webserver address diff --git a/pyinfra/core.py b/pyinfra/core.py index 63a83ae..c66d65a 100644 --- a/pyinfra/core.py +++ b/pyinfra/core.py @@ -9,7 +9,7 @@ from pyinfra.exceptions import DataLoadingFailure, AnalysisFailure, ProcessingFa from pyinfra.utils.file import dossier_id_and_file_id_to_compressed_storage_pdf_object_name -def make_storage_data_loader(storage): +def make_storage_data_loader(storage, bucket_name): def get_object_name(payload: dict) -> str: dossier_id, file_id = itemgetter("dossierId", "fileId")(payload) object_name = dossier_id_and_file_id_to_compressed_storage_pdf_object_name(dossier_id, file_id) @@ -18,7 +18,7 @@ def make_storage_data_loader(storage): def download(payload): object_name = get_object_name(payload) logging.debug(f"Downloading {object_name}...") - data = storage.get_object(object_name) + data = storage.get(bucket_name, object_name) logging.debug(f"Downloaded {object_name}.") return data diff --git a/pyinfra/storage/adapters/s3.py b/pyinfra/storage/adapters/s3.py index 26bdd68..a93c5d3 100644 --- a/pyinfra/storage/adapters/s3.py +++ b/pyinfra/storage/adapters/s3.py @@ -40,7 +40,7 @@ class S3StorageAdapter(StorageAdapter): response.release_conn() def get_all(self, bucket_name): - for obj in self.__client.list_objects(bucket_name): + for obj in self.__client.list_objects(bucket_name, recursive=True): logger.debug(f"Downloading '{obj.object_name}'...") yield self.get(bucket_name, obj.object_name) @@ -51,5 +51,5 @@ class S3StorageAdapter(StorageAdapter): self.__client.remove_object(bucket_name, obj.object_name) def list_bucket_files(self, bucket_name): - objs = self.__client.list_objects(bucket_name) + objs = self.__client.list_objects(bucket_name, recursive=True) return zip(repeat(bucket_name), map(attrgetter("object_name"), objs)) diff --git a/pyinfra/utils/file.py b/pyinfra/utils/file.py index a74c222..4068df4 100644 --- a/pyinfra/utils/file.py +++ b/pyinfra/utils/file.py @@ -2,8 +2,12 @@ import os +from pyinfra.config import CONFIG -def produce_compressed_storage_pdf_object_name(path_no_ext, ext="pdf"): + +def produce_compressed_storage_pdf_object_name(path_no_ext, ext=None): + if not ext: + ext = CONFIG.service.target_file_extension return f"{path_no_ext}.ORIGIN.{ext}.gz" diff --git a/scripts/manage_minio.py b/scripts/manage_minio.py index 56fb50e..8c7bf4a 100644 --- a/scripts/manage_minio.py +++ b/scripts/manage_minio.py @@ -8,6 +8,7 @@ from pyinfra.config import CONFIG from pyinfra.storage.storages import get_s3_storage from pyinfra.test.unit_tests.conftest import bucket_name +from pyinfra.utils.file import dossier_id_and_file_id_to_compressed_storage_pdf_object_name def parse_args(): @@ -29,7 +30,7 @@ def parse_args(): def add_file_compressed(storage, bucket_name, dossier_id, path) -> None: - path_gz = str(Path(dossier_id) / (Path(path).name + ".gz")) + path_gz = dossier_id_and_file_id_to_compressed_storage_pdf_object_name(dossier_id, Path(path).stem) with open(path, "rb") as f: data = gzip.compress(f.read()) @@ -40,6 +41,8 @@ if __name__ == "__main__": storage = get_s3_storage() bucket_name = CONFIG.test.bucket + if not storage.bucket_exists(bucket_name): + storage.make_bucket(bucket_name) args = parse_args() diff --git a/scripts/mock_client.py b/scripts/mock_client.py index 53b4dde..1ae07aa 100644 --- a/scripts/mock_client.py +++ b/scripts/mock_client.py @@ -7,8 +7,10 @@ from pyinfra.storage.storages import get_s3_storage def build_message_bodies(): storage = get_s3_storage() - for dossier_id, pdf_name in storage.list_bucket_files(CONFIG.test.bucket): + for bucket_name, pdf_name in storage.list_bucket_files(CONFIG.test.bucket): + print(pdf_name) file_id = pdf_name.split(".")[0] + dossier_id, file_id = file_id.split("/") yield json.dumps({"dossierId": dossier_id, "fileId": file_id}) diff --git a/src/serve.py b/src/serve.py index a9d9e3b..f7b6d00 100644 --- a/src/serve.py +++ b/src/serve.py @@ -40,7 +40,7 @@ def republish(channel, body, n_current_attempts): def make_callback(): - load_data = make_storage_data_loader(get_storage()) + load_data = make_storage_data_loader(get_storage(), CONFIG.service.bucket_name) analyze_file = make_analyzer(CONFIG.service.analysis_endpoint) json_wrapped_body_processor = make_payload_processor(load_data, analyze_file)