From e00d7d0216ab9366e6d9d2b19ddfa12e10d402d4 Mon Sep 17 00:00:00 2001
From: Dilawar Singh <dilawar@subcom.tech>
Date: Mon, 10 Oct 2022 20:58:55 +0530
Subject: [PATCH] chore: minor tweaks.

---
 bitia/__main__.py | 28 ++++++++++++++--------------
 bitia/helper.py   |  7 ++++---
 bitia/session.py  |  4 ++--
 3 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/bitia/__main__.py b/bitia/__main__.py
index 36e6752..77d0e6b 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 e6d1b91..f5fad05 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 15cf542..3f624ea 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,
     )
 
-- 
GitLab