jarzul/use dce progress api for progress bar (#88)

* Use the DCE Progress API to display the progress bar, instead of reading logs

* Create a singleton holding a common rich_console to use in both the logging configuration and when creating a rich Progress

This allows for all logs to go through rich and making regular logs behave correctly out-of-the box with the rich Progress bar

* Update the CLI Progress logic to separate the logic constructing state from the events received and the actual rendering using rich's Progress bar

* Test the logic handling the progress state from received events

* Use rich_console.is_terminal instead of sys.stderr.isatty()

Author: JulienArzul

src/databao_cli/features/build.py

@@ -6,9 +6,10 @@
 
 def build_impl(project_layout: ProjectLayout, domain: str, should_index: bool) -> list[BuildDatasourceResult]:
     dce_project_dir = project_layout.domains_dir / domain
-    manager = DatabaoContextDomainManager(domain_dir=dce_project_dir)
 
-    datasources = manager.get_configured_datasource_list()
-    with cli_progress(total=len(datasources), label="Building datasources"):
-        results: list[BuildDatasourceResult] = manager.build_context(datasource_ids=None, should_index=should_index)
+    with cli_progress() as progress:
+        results: list[BuildDatasourceResult] = DatabaoContextDomainManager(domain_dir=dce_project_dir).build_context(
+            datasource_ids=None, should_index=should_index, progress=progress
+        )
+
     return results

src/databao_cli/features/index.py

@@ -13,9 +13,7 @@ def index_impl(
 
     manager = DatabaoContextDomainManager(domain_dir=dce_project_dir)
 
-    total = len(datasource_ids) if datasource_ids is not None else len(manager.get_configured_datasource_list())
-
-    with cli_progress(total=total, label="Indexing datasources"):
-        results: list[IndexDatasourceResult] = manager.index_built_contexts(datasource_ids=datasource_ids)
+    with cli_progress() as progress:
+        results: list[IndexDatasourceResult] = manager.index_built_contexts(datasource_ids=datasource_ids, progress=progress)
 
     return results

src/databao_cli/shared/log/cli_progress.py

@@ -1,107 +1,165 @@
 from __future__ import annotations
 
-import logging
-import re
-import sys
 from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Any
-
-# Log patterns emitted by databao_context_engine.build_sources.build_runner
-_BUILD_START_RE = re.compile(r'^Found datasource of type ".*" with name (.+)$')
-_INDEX_START_RE = re.compile(r"^Indexing datasource (.+)$")
-_ENRICH_START_RE = re.compile(r"^Enriching context for datasource (.+)$")
-_SKIP_RE = re.compile(r"^Skipping disabled datasource (.+)$")
-_FAIL_RE = re.compile(r"^Failed to build source at \((.+?)\)")
-_FAIL_ENRICH_RE = re.compile(r"^Failed to enrich context for datasource \((.+?)\)")
-
-
-class _ProgressTrackingHandler(logging.Handler):
-    """Intercepts databao_context_engine log messages to drive a Rich progress bar.
-
-    The library processes datasources sequentially. It logs "Found datasource..."
-    at the START of each one. We advance the progress bar when we detect that
-    a new datasource has started (meaning the previous one finished), and once
-    more when the context manager exits (for the last datasource).
-    """
-
-    def __init__(
-        self,
-        progress: Any,
-        overall_task: Any,
-        datasource_task: Any,
-    ) -> None:
-        super().__init__()
+from dataclasses import dataclass, field
+
+from databao_context_engine.progress.progress import (
+    ProgressCallback,
+    ProgressEvent,
+    ProgressKind,
+    ProgressStep,
+)
+from rich.progress import (
+    BarColumn,
+    Progress,
+    SpinnerColumn,
+    TaskID,
+    TaskProgressColumn,
+    TextColumn,
+)
+from rich.table import Column
+
+from databao_cli.shared.log.console import rich_console
+
+
+def _noop_progress_cb(_: ProgressEvent) -> None:
+    return
+
+
+@dataclass
+class _StepProgress:
+    step: ProgressStep
+    units_completed: int
+    units_total: int
+
+
+@dataclass
+class _DatasourceProgress:
+    datasource_id: str
+    step_plan: tuple[ProgressStep, ...] = ()
+    completed_steps: set[ProgressStep] = field(default_factory=set)
+    current_step_progress: _StepProgress | None = None
+    finished: bool = False
+
+
+@dataclass
+class _OperationProgress:
+    total_datasources: int | None = None
+    completed_datasource_ids: list[str] = field(default_factory=list)
+    current_datasource: _DatasourceProgress | None = None
+
+
+def _compute_percent(ds: _DatasourceProgress) -> float | None:
+    if ds.finished:
+        return 100.0
+    if not ds.step_plan:
+        return None
+    fraction = 0.0
+    p = ds.current_step_progress
+    if p is not None and p.units_total > 0:
+        fraction = p.units_completed / p.units_total
+    return ((len(ds.completed_steps) + fraction) / len(ds.step_plan)) * 100.0
+
+
+def _apply_event(state: _OperationProgress, ev: ProgressEvent) -> _OperationProgress:
+    match ev.kind:
+        case ProgressKind.OPERATION_STARTED:
+            state.total_datasources = ev.operation_total
+
+        case ProgressKind.OPERATION_FINISHED:
+            pass
+
+        case ProgressKind.DATASOURCE_STARTED:
+            if state.total_datasources is None:
+                state.total_datasources = ev.datasource_total
+            state.current_datasource = _DatasourceProgress(datasource_id=ev.datasource_id or "datasource")
+
+        case ProgressKind.DATASOURCE_STEP_PLAN_SET:
+            if state.current_datasource is None:
+                state.current_datasource = _DatasourceProgress(datasource_id=ev.datasource_id or "datasource")
+            state.current_datasource.step_plan = ev.step_plan or ()
+
+        case ProgressKind.DATASOURCE_STEP_COMPLETED:
+            ds = state.current_datasource
+            if ds is not None and ev.step is not None and ev.step not in ds.completed_steps:
+                ds.completed_steps.add(ev.step)
+                if ds.current_step_progress is not None and ds.current_step_progress.step == ev.step:
+                    ds.current_step_progress = None
+
+        case ProgressKind.DATASOURCE_STEP_PROGRESS:
+            ds = state.current_datasource
+            if ds is not None and ev.step is not None and ev.step not in ds.completed_steps:
+                current = ds.current_step_progress
+                new_completed = ev.current_units_completed or 0
+                if current is not None and current.step == ev.step:
+                    new_completed = max(current.units_completed, new_completed)
+                ds.current_step_progress = _StepProgress(
+                    step=ev.step,
+                    units_completed=new_completed,
+                    units_total=ev.current_units_total or 0,
+                )
+
+        case ProgressKind.DATASOURCE_FINISHED:
+            ds = state.current_datasource
+            if ds is not None:
+                ds.finished = True
+                state.completed_datasource_ids.append(ds.datasource_id)
+
+    return state
+
+
+class _ProgressRenderer:
+    def __init__(self, progress: Progress) -> None:
         self._progress = progress
-        self._overall_task = overall_task
-        self._datasource_task = datasource_task
-        self._has_active = False  # whether a datasource is currently being processed
-
-    def emit(self, record: logging.LogRecord) -> None:
-        msg = record.getMessage()
-
-        # Datasource processing started
-        m = _BUILD_START_RE.match(msg) or _INDEX_START_RE.match(msg) or _ENRICH_START_RE.match(msg)
-        if m:
-            if self._has_active:
-                # Previous datasource finished — advance
-                self._progress.advance(self._overall_task)
-            self._has_active = True
-            name = m.group(1)
-            self._progress.update(self._datasource_task, description=f"  [dim]{name}[/dim]")
+        self._overall_task_id: TaskID | None = None
+        self._datasource_task_id: TaskID | None = None
+
+    def render(self, state: _OperationProgress) -> None:
+        self._render_overall(state)
+        self._render_datasource(state.current_datasource)
+
+    def _render_overall(self, state: _OperationProgress) -> None:
+        completed = len(state.completed_datasource_ids)
+        description = f"Datasources {completed}/{state.total_datasources}" if state.total_datasources else "Datasources"
+        if self._overall_task_id is None:
+            self._overall_task_id = self._progress.add_task(description, total=state.total_datasources, completed=completed)
+        else:
+            self._progress.update(
+                self._overall_task_id,
+                description=description,
+                total=state.total_datasources,
+                completed=float(completed),
+            )
+
+    def _render_datasource(self, ds: _DatasourceProgress | None) -> None:
+        if ds is None:
             return
 
-        # Datasource skipped (immediately done)
-        if _SKIP_RE.match(msg):
-            self._progress.advance(self._overall_task)
-            return
-
-        # Datasource failed (after "Found datasource", so active is already True)
-        if _FAIL_RE.match(msg) or _FAIL_ENRICH_RE.match(msg):
-            if self._has_active:
-                self._progress.advance(self._overall_task)
-                self._has_active = False
-            return
+        description = f"    {ds.datasource_id}"
+        percent = _compute_percent(ds)
 
-    def finish(self) -> None:
-        """Advance for the last datasource that was being processed."""
-        if self._has_active:
-            self._progress.advance(self._overall_task)
-            self._has_active = False
+        if self._datasource_task_id is None:
+            self._datasource_task_id = self._progress.add_task(
+                description,
+                total=100.0 if percent is not None else None,
+                completed=int(percent or 0),
+            )
+        else:
+            self._progress.update(
+                self._datasource_task_id,
+                description=description,
+                total=100.0 if percent is not None else None,
+                completed=percent or 0.0,
+            )
 
 
 @contextmanager
-def cli_progress(total: int | None = None, label: str = "Datasources") -> Iterator[None]:
-    """Show a Rich progress bar during build/index operations.
-
-    Intercepts ``databao_context_engine`` log messages to track per-datasource progress.
-    Redirects library logging through Rich for clean TTY output.
-
-    Args:
-        total: Number of datasources to process.
-        label: Label for the overall progress bar.
-    """
-    try:
-        from rich.console import Console
-        from rich.logging import RichHandler
-        from rich.progress import (
-            BarColumn,
-            MofNCompleteColumn,
-            Progress,
-            SpinnerColumn,
-            TextColumn,
-        )
-        from rich.table import Column
-    except ImportError:
-        yield
-        return
-
-    console = Console(stderr=True)
-
+def cli_progress() -> Iterator[ProgressCallback]:
     # Rich's is_terminal already checks isatty(), NO_COLOR, TERM=dumb, etc.
     # This prevents progress bar ANSI output from breaking pexpect-based e2e tests.
-    if not console.is_terminal:
-        yield
+    if not rich_console.is_terminal:
+        yield _noop_progress_cb
         return
 
     progress = Progress(
@@ -111,49 +169,16 @@ def cli_progress(total: int | None = None, label: str = "Datasources") -> Iterat
             table_column=Column(width=50, overflow="ellipsis", no_wrap=True),
         ),
         BarColumn(),
-        MofNCompleteColumn(),
+        TaskProgressColumn(),
         transient=True,
-        console=console,
-    )
-
-    overall_task = progress.add_task(label, total=total)
-    datasource_task = progress.add_task("  [dim]starting…[/dim]", total=None)
-
-    # --- logging setup ---
-    engine_logger = logging.getLogger("databao_context_engine")
-    cli_logger = logging.getLogger("databao_cli")
-
-    prev_engine = (list(engine_logger.handlers), engine_logger.propagate)
-    prev_cli = (list(cli_logger.handlers), cli_logger.propagate)
-
-    def _is_console_handler(h: logging.Handler) -> bool:
-        return isinstance(h, logging.StreamHandler) and getattr(h, "stream", None) in (sys.stderr, sys.stdout)
-
-    rich_handler = RichHandler(
-        console=console,
-        show_time=False,
-        show_level=True,
-        show_path=False,
-        rich_tracebacks=False,
+        console=rich_console,
     )
 
-    tracker = _ProgressTrackingHandler(progress, overall_task, datasource_task)
-    tracker.setLevel(logging.DEBUG)
-
-    try:
-        for lgr in (engine_logger, cli_logger):
-            kept = [h for h in lgr.handlers if not _is_console_handler(h)]
-            lgr.handlers = [*kept, rich_handler]
-            lgr.propagate = False
-
-        engine_logger.addHandler(tracker)
+    state = _OperationProgress()
+    renderer = _ProgressRenderer(progress)
 
-        with progress:
-            yield
+    def on_event(ev: ProgressEvent) -> None:
+        renderer.render(_apply_event(state, ev))
 
-        tracker.finish()
-    finally:
-        engine_logger.handlers = prev_engine[0]
-        engine_logger.propagate = prev_engine[1]
-        cli_logger.handlers = prev_cli[0]
-        cli_logger.propagate = prev_cli[1]
+    with progress:
+        yield on_event

src/databao_cli/shared/log/console.py

@@ -0,0 +1,3 @@
+from rich.console import Console
+
+rich_console = Console(stderr=True)