diff --git a/bitia/__main__.py b/bitia/__main__.py index 36e6752dc5301730b21b727674c33f9d59e0d233..77d0e6b02ea5fd00f57aa12ab8dc0725b873f521 100644 --- a/bitia/__main__.py +++ b/bitia/__main__.py @@ -13,7 +13,7 @@ from enum import Enum from rich.progress import track import bitia.helper as bhelper -from bitia.logger import logger, cprint, set_logger_level +from bitia.logger import logger, cprint, set_logger_level, console from bitia import version as bversion import bitia.pipeline as bpipeline @@ -81,10 +81,13 @@ def create_remote_container( @app.command("list-container") @session -def list_remote_container( - user_input, server: str = bconfig.get_server() -) -> T.List[str]: +def list_containers(user_input): """List the remote server associated with the pipeline.""" + for container in _list_remote_container(user_input): + cprint(container) + + +def _list_remote_container(user_input) -> T.List[str]: pipeline = bpipeline.user_input_to_pipeline(user_input) logger.debug(f"sha256 of `{pipeline.zipfile}` is `{pipeline.checksum}`") logger.info(f"Finding container for {user_input}, sha256sum={pipeline.checksum}") @@ -94,9 +97,7 @@ def list_remote_container( params=dict(pipeline_sha256=pipeline.checksum), ) res.raise_for_status() - containers = res.json()["containers"].split(",") - cprint(containers) - return containers + return res.json()["containers"].split(",") @app.command("logs") @@ -113,7 +114,7 @@ def stream_log(user_input, server: str = bconfig.get_server()): ) res.raise_for_status() for line in res.iter_lines(): - print(line.decode().rstrip()) + cprint(line.decode().rstrip()) @app.command("submit") @@ -126,7 +127,7 @@ def submit_pipeline(user_input, *, rerun: bool = False, output_lines: T.List[str """ res = create_remote_container(user_input, recreate=rerun, output_lines=output_lines) logger.info("Remote container: %s", res) - containers = list_remote_container(user_input) + containers = _list_remote_container(user_input) cprint(f"{containers}") return containers @@ -136,11 +137,10 @@ def submit_pipeline(user_input, *, rerun: bool = False, output_lines: T.List[str def run_user_input(user_input, *, rerun: bool = False, output_lines: T.List[str] = []): """Run a pipeline""" create_remote_container(user_input, recreate=rerun, output_lines=output_lines) - containers = list_remote_container(user_input) - return [ - bhelper.log_container(container, server=bconfig.get_server()) - for container in containers - ] + containers = _list_remote_container(user_input) + for container in containers: + for _bl in bhelper.log_container(container, server=bconfig.get_server()): + print(_bl.decode().rstrip()) @app.command("checksum") diff --git a/bitia/helper.py b/bitia/helper.py index e6d1b91117f6b43f4583950b608ed6735fabc775..f5fad05427e9dd33193771e60ea55f7a42604d60 100644 --- a/bitia/helper.py +++ b/bitia/helper.py @@ -23,12 +23,13 @@ from bitia.checksumdir import dirhash from bitia.logger import logger -def log_container(container: str, server: str): +def log_container(container: str, server: str, *, timestamps: bool = False): assert ( container ), "Failed to determine the container that is runnning the pipeline. There is probably a bug in server end." - for line in bsession.fetch_logs(container, server=server): - print(line.decode().rstrip()) + for line in bsession.fetch_logs(container, server=server, + timestamps=timestamps): + yield line def _check_server_status(server: str) -> int: diff --git a/bitia/session.py b/bitia/session.py index 15cf542e113dc5a8ac4cb3eac9201ecc9bd5b65e..3f624ea8c4dcca5c7801fbe3326c72e2b5f68e84 100644 --- a/bitia/session.py +++ b/bitia/session.py @@ -9,12 +9,12 @@ g_session = requests.Session() SESSION_PICKLE_FILE = bconfig.bitia_dir() / ".session.pickle" -def fetch_logs(container: str, *, server): +def fetch_logs(container: str, *, server, timestamps: bool = True): """Fetch logs from a container.""" logger.info(f"Fetching logs for container `{container}`") return get( f"{server}/container/logs", - params=dict(container=container), + params=dict(container=container, timestamps=timestamps), stream=True, )