ubunye-engine 0.1.1__py3-none-any.whl

ubunye/__init__.py

@@ -0,0 +1,10 @@
+"""Ubunye Engine package."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("ubunye-engine")
+except PackageNotFoundError:
+    __version__ = "unknown"
+
+__all__ = ["core", "config", "cli", "plugins", "backends"]

ubunye/backends/__init__.py

@@ -0,0 +1 @@
+from .spark_backend import SparkBackend as SparkBackend

ubunye/backends/spark_backend.py

@@ -0,0 +1,127 @@
+"""
+Spark backend implementation for Ubunye.
+
+- Lazily imports pyspark so users can install the package without Spark.
+- Provides a simple lifecycle: start() / stop().
+- Supports context-manager usage: `with SparkBackend(...) as backend: ...`
+- Guards against double-starts and exposes the effective Spark conf for debugging.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Optional
+
+from ubunye.core.interfaces import Backend
+
+if TYPE_CHECKING:  # only for type checkers; no runtime dependency on pyspark
+    from pyspark.sql import SparkSession
+
+
+class SparkBackend(Backend):
+    """Creates and manages a SparkSession for a Ubunye run.
+
+    Parameters
+    ----------
+    app_name : str
+        Spark application name (appears in Spark UI/history server).
+    conf : Optional[Dict[str, str]]
+        Spark configuration key-values (e.g., {"spark.master": "yarn"}).
+
+    Notes
+    -----
+    - pyspark is imported lazily inside `start()` to keep installation lightweight.
+    - `spark` property is only valid after `start()` (or inside the context manager).
+    """
+
+    def __init__(self, app_name: str = "ubunye", conf: Optional[Dict[str, str]] = None) -> None:
+        self._spark: Optional["SparkSession"] = None
+        self._app_name = app_name
+        self._conf = dict(conf or {})
+
+    # -------------------------
+    # Lifecycle
+    # -------------------------
+    def start(self) -> None:
+        """Create (or reuse) a SparkSession.
+
+        Safe to call multiple times; a second call is a no-op if a session already exists.
+        """
+        if self._spark is not None:
+            return  # already started
+
+        # Lazy import to avoid hard dependency during pip install
+        from pyspark.sql import SparkSession  # type: ignore
+
+        builder = SparkSession.builder.appName(self._app_name)
+        for k, v in self._conf.items():
+            builder = builder.config(k, v)
+        self._spark = builder.getOrCreate()
+
+    def stop(self) -> None:
+        """Stop the SparkSession if running."""
+        if self._spark is not None:
+            try:
+                self._spark.stop()
+            finally:
+                self._spark = None
+
+    # Context manager support
+    def __enter__(self) -> "SparkBackend":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.stop()
+
+    def __del__(self) -> None:
+        # Best-effort cleanup; not guaranteed to run in all interpreter shutdown scenarios.
+        try:
+            self.stop()
+        except Exception:
+            # Avoid noisy destructor exceptions at process teardown.
+            pass
+
+    # -------------------------
+    # Properties
+    # -------------------------
+    @property
+    def spark(self) -> "SparkSession":
+        """Return the active SparkSession (after `start()`).
+
+        Raises
+        ------
+        RuntimeError
+            If `start()` has not been called.
+        """
+        if self._spark is None:
+            raise RuntimeError(
+                "Spark session not started. Call start() first or use context manager."
+            )
+        return self._spark
+
+    @property
+    def is_spark(self) -> bool:
+        """Whether this backend is Spark-based (always True here)."""
+        return True
+
+    @property
+    def app_name(self) -> str:
+        """Configured Spark app name."""
+        return self._app_name
+
+    @property
+    def conf_input(self) -> Dict[str, str]:
+        """The configuration dict passed into this backend at construction time."""
+        return dict(self._conf)
+
+    @property
+    def conf_effective(self) -> Dict[str, str]:
+        """The *effective* Spark configuration currently applied to the session.
+
+        Returns an empty dict if the session hasn't been started yet.
+        """
+        if self._spark is None:
+            return {}
+        sc = self._spark.sparkContext
+        # Convert list[tuple[str,str]] to dict[str,str]
+        return {k: v for (k, v) in sc.getConf().getAll()}

ubunye/cli/lineage.py

@@ -0,0 +1,304 @@
+"""CLI commands for inspecting run lineage records.
+
+Mounted on the main app as a sub-command group:
+
+    ubunye lineage show    -d DIR -u USECASE -p PKG -t TASK [--run-id ID]
+    ubunye lineage list    -d DIR -u USECASE -p PKG -t TASK [-n 10]
+    ubunye lineage compare -d DIR -u USECASE -p PKG -t TASK --run-id1 ID1 --run-id2 ID2
+    ubunye lineage search  -d DIR [--status error] [--since 2025-01-01]
+    ubunye lineage trace   -d DIR -u USECASE -p PKG -t TASK
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+from ubunye.lineage.context import RunContext
+from ubunye.lineage.storage import FileSystemLineageStore
+
+lineage_app = typer.Typer(name="lineage", help="Inspect run lineage records.", add_completion=False)
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+_LINE = "-" * 60
+
+
+def _store(lineage_dir: str) -> FileSystemLineageStore:
+    return FileSystemLineageStore(lineage_dir)
+
+
+def _task_path(usecase: str, package: str, task: str) -> str:
+    return f"{usecase}/{package}/{task}"
+
+
+def _fmt_row(ctx: RunContext) -> str:
+    duration = f"{ctx.duration_sec:.1f}s" if ctx.duration_sec is not None else "-"
+    in_rows = sum(s.row_count or 0 for s in ctx.inputs)
+    out_rows = sum(s.row_count or 0 for s in ctx.outputs)
+    return (
+        f"{ctx.run_id[:8]}  {ctx.started_at[:19]}  "
+        f"{ctx.status:<8}  {duration:>7}  "
+        f"in:{in_rows}  out:{out_rows}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# show
+# ---------------------------------------------------------------------------
+
+
+@lineage_app.command("show")
+def show(
+    usecase_dir: Path = typer.Option(
+        ..., "-d", "--usecase-dir", help="Root directory of pipelines."
+    ),
+    usecase: str = typer.Option(..., "-u", "--usecase"),
+    package: str = typer.Option(..., "-p", "--package"),
+    task: str = typer.Option(..., "-t", "--task"),
+    run_id: Optional[str] = typer.Option(
+        None, "--run-id", help="Specific run ID (default: latest)."
+    ),
+    lineage_dir: str = typer.Option(".ubunye/lineage", "--lineage-dir"),
+):
+    """Show a run record as formatted JSON (latest or specific run)."""
+    store = _store(str(usecase_dir / lineage_dir))
+    tp = _task_path(usecase, package, task)
+
+    try:
+        if run_id:
+            ctx = store.load(tp, run_id)
+        else:
+            runs = store.list_runs(tp, n=1)
+            if not runs:
+                typer.secho(f"No lineage records found for task '{tp}'.", fg=typer.colors.YELLOW)
+                raise typer.Exit(code=1)
+            ctx = runs[0]
+    except FileNotFoundError as e:
+        typer.secho(str(e), fg=typer.colors.RED, err=True)
+        raise typer.Exit(code=1)
+
+    typer.echo(json.dumps(ctx.to_dict(), indent=2))
+
+
+# ---------------------------------------------------------------------------
+# list
+# ---------------------------------------------------------------------------
+
+
+@lineage_app.command("list")
+def list_runs(
+    usecase_dir: Path = typer.Option(..., "-d", "--usecase-dir"),
+    usecase: str = typer.Option(..., "-u", "--usecase"),
+    package: str = typer.Option(..., "-p", "--package"),
+    task: str = typer.Option(..., "-t", "--task"),
+    n: int = typer.Option(10, "-n", "--n", help="Number of recent runs to show."),
+    lineage_dir: str = typer.Option(".ubunye/lineage", "--lineage-dir"),
+):
+    """List recent runs for a task (newest first)."""
+    store = _store(str(usecase_dir / lineage_dir))
+    tp = _task_path(usecase, package, task)
+    runs = store.list_runs(tp, n=n)
+
+    if not runs:
+        typer.secho(f"No lineage records found for '{tp}'.", fg=typer.colors.YELLOW)
+        return
+
+    typer.echo(_LINE)
+    typer.echo(f"Task: {tp}")
+    typer.echo(_LINE)
+    typer.echo(f"{'RUN_ID':9}  {'STARTED':19}  {'STATUS':<8}  {'DUR':>7}  ROWS")
+    typer.echo(_LINE)
+    for ctx in runs:
+        color = (
+            typer.colors.GREEN
+            if ctx.status == "success"
+            else (typer.colors.RED if ctx.status == "error" else typer.colors.YELLOW)
+        )
+        typer.secho(_fmt_row(ctx), fg=color)
+
+
+# ---------------------------------------------------------------------------
+# compare
+# ---------------------------------------------------------------------------
+
+
+@lineage_app.command("compare")
+def compare(
+    usecase_dir: Path = typer.Option(..., "-d", "--usecase-dir"),
+    usecase: str = typer.Option(..., "-u", "--usecase"),
+    package: str = typer.Option(..., "-p", "--package"),
+    task: str = typer.Option(..., "-t", "--task"),
+    run_id1: str = typer.Option(..., "--run-id1", help="First run ID."),
+    run_id2: str = typer.Option(..., "--run-id2", help="Second run ID."),
+    lineage_dir: str = typer.Option(".ubunye/lineage", "--lineage-dir"),
+):
+    """Diff two run records — highlight changes in hashes, row counts, and status."""
+    store = _store(str(usecase_dir / lineage_dir))
+    tp = _task_path(usecase, package, task)
+
+    try:
+        a = store.load(tp, run_id1)
+        b = store.load(tp, run_id2)
+    except FileNotFoundError as e:
+        typer.secho(str(e), fg=typer.colors.RED, err=True)
+        raise typer.Exit(code=1)
+
+    typer.echo(_LINE)
+    typer.echo(f"Comparing runs for '{tp}'")
+    typer.echo(f"  A: {a.run_id[:8]}  {a.started_at[:19]}  {a.status}")
+    typer.echo(f"  B: {b.run_id[:8]}  {b.started_at[:19]}  {b.status}")
+    typer.echo(_LINE)
+
+    def _cmp(label: str, va: object, vb: object) -> None:
+        if va == vb:
+            typer.echo(f"  {label}: {va}  (unchanged)")
+        else:
+            typer.secho(f"  {label}: {va} → {vb}  CHANGED", fg=typer.colors.YELLOW)
+
+    _cmp("status", a.status, b.status)
+    _cmp("duration_sec", a.duration_sec, b.duration_sec)
+    _cmp("config_hash", a.config_hash, b.config_hash)
+
+    # Compare outputs step by step
+    a_out = {s.name: s for s in a.outputs}
+    b_out = {s.name: s for s in b.outputs}
+    all_names = sorted(set(a_out) | set(b_out))
+    for name in all_names:
+        sa = a_out.get(name)
+        sb = b_out.get(name)
+        typer.echo(f"  Output '{name}':")
+        _cmp("    row_count", sa.row_count if sa else None, sb.row_count if sb else None)
+        _cmp("    schema_hash", sa.schema_hash if sa else None, sb.schema_hash if sb else None)
+        _cmp("    data_hash", sa.data_hash if sa else None, sb.data_hash if sb else None)
+
+
+# ---------------------------------------------------------------------------
+# search
+# ---------------------------------------------------------------------------
+
+
+@lineage_app.command("search")
+def search(
+    usecase_dir: Path = typer.Option(..., "-d", "--usecase-dir"),
+    task: Optional[str] = typer.Option(None, "-t", "--task", help="Restrict to this task name."),
+    usecase: Optional[str] = typer.Option(
+        None, "-u", "--usecase", help="Restrict to this usecase."
+    ),
+    package: Optional[str] = typer.Option(
+        None, "-p", "--package", help="Restrict to this package."
+    ),
+    status: Optional[str] = typer.Option(
+        None, "--status", help="Filter by status: success|error|running."
+    ),
+    since: Optional[str] = typer.Option(
+        None, "--since", help="Only runs started on or after this ISO date."
+    ),
+    lineage_dir: str = typer.Option(".ubunye/lineage", "--lineage-dir"),
+):
+    """Search all recorded runs across tasks with optional filters."""
+    store = _store(str(usecase_dir / lineage_dir))
+
+    tp = None
+    if usecase and package and task:
+        tp = _task_path(usecase, package, task)
+    elif usecase and package:
+        tp = f"{usecase}/{package}"
+    elif usecase:
+        tp = usecase
+
+    runs = store.search(task_path=tp, status=status, since=since)
+
+    if not runs:
+        typer.secho("No matching lineage records found.", fg=typer.colors.YELLOW)
+        return
+
+    typer.echo(_LINE)
+    typer.echo(f"{'RUN_ID':9}  {'TASK PATH':40}  {'STARTED':19}  {'STATUS':<8}  DUR")
+    typer.echo(_LINE)
+    for ctx in runs:
+        color = (
+            typer.colors.GREEN
+            if ctx.status == "success"
+            else (typer.colors.RED if ctx.status == "error" else typer.colors.YELLOW)
+        )
+        duration = f"{ctx.duration_sec:.1f}s" if ctx.duration_sec is not None else "-"
+        typer.secho(
+            f"{ctx.run_id[:8]}  {ctx.task_path:<40}  {ctx.started_at[:19]}  "
+            f"{ctx.status:<8}  {duration}",
+            fg=color,
+        )
+
+
+# ---------------------------------------------------------------------------
+# trace
+# ---------------------------------------------------------------------------
+
+
+@lineage_app.command("trace")
+def trace(
+    usecase_dir: Path = typer.Option(..., "-d", "--usecase-dir"),
+    usecase: str = typer.Option(..., "-u", "--usecase"),
+    package: str = typer.Option(..., "-p", "--package"),
+    task: str = typer.Option(..., "-t", "--task"),
+    run_id: Optional[str] = typer.Option(None, "--run-id", help="Specific run (default: latest)."),
+    lineage_dir: str = typer.Option(".ubunye/lineage", "--lineage-dir"),
+):
+    """Print the input → transform → output data flow graph for a run."""
+    store = _store(str(usecase_dir / lineage_dir))
+    tp = _task_path(usecase, package, task)
+
+    try:
+        if run_id:
+            ctx = store.load(tp, run_id)
+        else:
+            runs = store.list_runs(tp, n=1)
+            if not runs:
+                typer.secho(f"No lineage records for '{tp}'.", fg=typer.colors.YELLOW)
+                raise typer.Exit(code=1)
+            ctx = runs[0]
+    except FileNotFoundError as e:
+        typer.secho(str(e), fg=typer.colors.RED, err=True)
+        raise typer.Exit(code=1)
+
+    typer.echo()
+    typer.secho(f"Lineage trace: {ctx.task_path}", bold=True)
+    typer.echo(f"Run:     {ctx.run_id}  [{ctx.status}]  {ctx.started_at[:19]}")
+    typer.echo(f"Version: {ctx.model} v{ctx.version}")
+    typer.echo()
+
+    # INPUTS
+    typer.secho("  INPUTS", fg=typer.colors.CYAN)
+    for step in ctx.inputs:
+        typer.echo(f"    [{step.format}] {step.name}")
+        typer.echo(f"      location : {step.location}")
+        if step.row_count is not None:
+            typer.echo(f"      rows     : {step.row_count:,}")
+        if step.schema_hash:
+            typer.echo(f"      schema   : {step.schema_hash}")
+        if step.data_hash:
+            typer.echo(f"      data     : {step.data_hash}")
+
+    # TRANSFORM
+    typer.echo()
+    typer.secho("  TRANSFORM", fg=typer.colors.CYAN)
+    typer.echo("    (transformations.py)")
+
+    # OUTPUTS
+    typer.echo()
+    typer.secho("  OUTPUTS", fg=typer.colors.CYAN)
+    for step in ctx.outputs:
+        typer.echo(f"    [{step.format}] {step.name}")
+        typer.echo(f"      location : {step.location}")
+        if step.row_count is not None:
+            typer.echo(f"      rows     : {step.row_count:,}")
+        if step.schema_hash:
+            typer.echo(f"      schema   : {step.schema_hash}")
+        if step.data_hash:
+            typer.echo(f"      data     : {step.data_hash}")
+    typer.echo()