diff --git a/cli.py b/cli.py index c21acfa..64addf5 100644 --- a/cli.py +++ b/cli.py @@ -3,9 +3,12 @@ import json import click from query_performance_gauge import run_measuring_session - +from utils import compose_config @click.command() -@click.option("--config", required=True, type=click.File()) -def measure_performance(config): - run_measuring_session(json.load(config)) +@click.option("--config", type=click.File()) +@click.option("--credentials", type=click.File()) +@click.option("--queries", type=click.File()) +def measure_performance(config, credentials, queries): + config = compose_config(config, credentials, queries) + run_measuring_session(config) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..89cb246 --- /dev/null +++ b/utils.py @@ -0,0 +1,29 @@ +from typing import Union +from io import TextIOWrapper +import json + + +def compose_config( + config: Union[TextIOWrapper, None], + credentials: Union[TextIOWrapper, None], + queries: Union[TextIOWrapper, None], +) -> dict: + """ + Receive the CLI arguments and compose a session config. + + :param config: file pointer to a full config file. + :param credentials: file pointer to a credentials file. + :param queries: file pointer to a queries file. + :return: a dict with the composed configuration. + """ + + if config is not None: + return json.load(config) + + if credentials is None or queries is None: + raise ValueError( + "You need to provide both --credentials and --queries arguments." + ) + + # Merge both into one dict so it follows the same structure as a full config file + return {**json.load(credentials), **json.load(queries)}