diff --git a/config.yaml b/config.yaml index ceb12a0..31c6571 100755 --- a/config.yaml +++ b/config.yaml @@ -23,7 +23,7 @@ rabbitmq: storage: backend: $STORAGE_BACKEND|s3 # The type of storage to use {s3, azure} - bucket: $STORAGE_BUCKET|"pyinfra-test-bucket" # The bucket / container to pull files specified in queue requests from + bucket: "STORAGE_BUCKET_NAME|STORAGE_AZURECONTAINERNAME|pyinfra-test-bucket" # The bucket / container to pull files specified in queue requests from s3: endpoint: $STORAGE_ENDPOINT|"http://127.0.0.1:9000" diff --git a/pyinfra/config.py b/pyinfra/config.py index 06c9cae..c306e19 100644 --- a/pyinfra/config.py +++ b/pyinfra/config.py @@ -1,7 +1,9 @@ """Implements a config object with dot-indexing syntax.""" - +import os +from itertools import chain from envyaml import EnvYAML +from funcy import first, juxt, butlast, last from pyinfra.locations import CONFIG_FILE @@ -38,3 +40,15 @@ class Config: CONFIG = Config(CONFIG_FILE) + + +def parse_disjunction_string(disjunction_string): + def try_parse_env_var(disjunction_string): + try: + return os.environ[disjunction_string] + except KeyError: + return None + + options = disjunction_string.split("|") + identifiers, fallback_value = juxt(butlast, last)(options) + return first(chain(filter(try_parse_env_var, identifiers), [fallback_value])) diff --git a/pyinfra/visitor.py b/pyinfra/visitor.py index 9e0a30f..0399789 100644 --- a/pyinfra/visitor.py +++ b/pyinfra/visitor.py @@ -5,7 +5,7 @@ import logging from operator import itemgetter from typing import Callable -from pyinfra.config import CONFIG +from pyinfra.config import CONFIG, parse_disjunction_string from pyinfra.exceptions import DataLoadingFailure from pyinfra.storage.storage import Storage @@ -23,11 +23,12 @@ def get_response_object_name(body): def get_object_descriptor(body): - return {"bucket_name": CONFIG.storage.bucket, "object_name": get_object_name(body)} + return {"bucket_name": parse_disjunction_string(CONFIG.storage.bucket), "object_name": get_object_name(body)} def get_response_object_descriptor(body): - return {"bucket_name": CONFIG.storage.bucket, "object_name": get_response_object_name(body)} + return {"bucket_name": parse_disjunction_string(CONFIG.storage.bucket), + "object_name": get_response_object_name(body)} class ResponseStrategy(abc.ABC): diff --git a/requirements.txt b/requirements.txt index 64f0c27..ba7589e 100755 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,4 @@ testcontainers==3.4.2 docker-compose==1.29.2 tqdm==4.62.3 pytest~=7.0.1 +funcy==1.17 diff --git a/scripts/manage_minio.py b/scripts/manage_minio.py index eaa921e..88e1127 100644 --- a/scripts/manage_minio.py +++ b/scripts/manage_minio.py @@ -5,7 +5,7 @@ from pathlib import Path from tqdm import tqdm -from pyinfra.config import CONFIG +from pyinfra.config import CONFIG, parse_disjunction_string from pyinfra.storage.storages import get_s3_storage @@ -51,7 +51,7 @@ def add_file_compressed(storage, bucket_name, dossier_id, path) -> None: if __name__ == "__main__": storage = get_s3_storage() - bucket_name = CONFIG.storage.bucket + bucket_name = parse_disjunction_string(CONFIG.storage.bucket) if not storage.has_bucket(bucket_name): storage.make_bucket(bucket_name) diff --git a/scripts/mock_client.py b/scripts/mock_client.py index e7699fa..0b456d9 100644 --- a/scripts/mock_client.py +++ b/scripts/mock_client.py @@ -3,7 +3,7 @@ import json import pika -from pyinfra.config import CONFIG +from pyinfra.config import CONFIG, parse_disjunction_string from pyinfra.storage.storages import get_s3_storage @@ -53,7 +53,7 @@ def build_message_bodies(analyse_container_type): return message_dict storage = get_s3_storage() - for bucket_name, pdf_name in storage.get_all_object_names(CONFIG.storage.bucket): + for bucket_name, pdf_name in storage.get_all_object_names(parse_disjunction_string(CONFIG.storage.bucket)): if "pdf" not in pdf_name: continue file_id = pdf_name.split(".")[0]