model2data-studio 0.1.0__py3-none-any.whl

model2data_studio/__init__.py

@@ -0,0 +1,17 @@
+"""The model2data studio command line client.
+
+A schema is designed once, so the studio itself is a place you visit and leave.
+This is the part that stays: it pulls a saved project and writes its data, which
+is what turns "I made some fixtures" into "CI rebuilds the fixtures every run"
+and a committed, quietly rotting CSV into something regenerated from the model
+it came from.
+
+Deliberately thin. It holds no parser, no generator and no schema knowledge —
+the studio's API does all of that, and this only knows how to ask.
+"""
+
+from __future__ import annotations
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__"]

model2data_studio/archive.py

@@ -0,0 +1,72 @@
+"""Unpacking what the studio sends back.
+
+Its own module because two callers write these files: the CLI (`generate
+--out`) and the MCP server (`write_data`). The path check below is the reason
+it is worth sharing rather than copying -- an archive is the classic way to
+write outside the directory you were pointed at, and the directory is somebody
+repository.
+"""
+
+from __future__ import annotations
+
+import io
+import shutil
+import zipfile
+from pathlib import Path
+
+from .client import StudioError
+
+
+def unpack(
+    archive: bytes, out: Path, *, force: bool, force_hint: str = "Pass --force to replace them."
+) -> list[Path]:
+    """Write the archive's files into `out`, or refuse before writing any.
+
+    Every target is resolved and checked first. A refusal halfway through would
+    leave a directory that is neither the previous export nor the new one, and
+    in CI that is a build that half-succeeds against half-old fixtures.
+
+    `force_hint` ends the refusal, because the way to say "do it anyway" is the
+    caller's word: a flag to the CLI, an argument to an MCP tool.
+    """
+    try:
+        bundle = zipfile.ZipFile(io.BytesIO(archive))
+    except zipfile.BadZipFile as error:
+        raise StudioError(
+            f"The studio returned something that isn't an archive: {error}"
+        ) from error
+
+    with bundle:
+        targets = [
+            (member, target_for(out, member.filename))
+            for member in bundle.infolist()
+            if not member.is_dir()
+        ]
+        if not force:
+            existing = [str(target) for _, target in targets if target.exists()]
+            if existing:
+                raise StudioError(
+                    "Refusing to overwrite "
+                    + ", ".join(existing[:5])
+                    + (f" and {len(existing) - 5} more" if len(existing) > 5 else "")
+                    + f". {force_hint}"
+                )
+        for member, target in targets:
+            target.parent.mkdir(parents=True, exist_ok=True)
+            with bundle.open(member) as source, target.open("wb") as sink:
+                shutil.copyfileobj(source, sink)
+    return sorted(target for _, target in targets)
+
+
+def target_for(out: Path, member: str) -> Path:
+    """Where one archive entry lands, refusing anything that escapes `out`.
+
+    The archive comes from the studio, but an archive is still the classic way
+    to write outside the directory you were pointed at, and `--out` in CI is a
+    checkout of somebody's repository.
+    """
+    root = out.resolve()
+    target = (root / member).resolve()
+    if target != root and root not in target.parents:
+        raise StudioError(f"The archive contains a path outside {out}: {member!r}")
+    return target

model2data_studio/cli.py

@@ -0,0 +1,452 @@
+"""The commands themselves.
+
+Three, and only one of them matters: `generate` is the reason this exists. It
+pulls a saved project, asks the studio to regenerate it, and unpacks the archive
+where a build expects to find seeds. `projects` and `whoami` are there so the
+first five minutes of setting that up are not guesswork about ids and keys.
+
+Everything is resolved from the saved project first and overridden by a flag
+second, so a pipeline that passes nothing but `--project` and `--out` keeps
+producing whatever the studio last previewed — the settings live with the model,
+not scattered across a YAML file that drifts away from it.
+"""
+
+from __future__ import annotations
+
+import csv
+import datetime as dt
+import enum
+import tempfile
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, NoReturn
+
+import typer
+
+from . import __version__, export
+from .archive import unpack as _unpack
+from .client import DEFAULT_URL, StudioClient, StudioError
+
+KEY_ENV = "MODEL2DATA_STUDIO_KEY"
+URL_ENV = "MODEL2DATA_STUDIO_URL"
+
+app = typer.Typer(
+    add_completion=False,
+    no_args_is_help=True,
+    help="Generate a model2data studio project's data from the command line.",
+)
+
+# Swapped in the tests for one that hands the client an `httpx.MockTransport`.
+# It is the only seam: every command below goes through it, so all of them can
+# be run end to end without a studio to talk to.
+client_factory: Callable[[str, str], StudioClient] = StudioClient
+
+
+def _print_version(value: bool) -> None:
+    if value:
+        typer.echo(f"model2data-studio {__version__}")
+        raise typer.Exit()
+
+
+@app.callback()
+def _root(
+    version: bool = typer.Option(
+        False,
+        "--version",
+        callback=_print_version,
+        is_eager=True,
+        help="Print the version and exit.",
+    ),
+) -> None:
+    """Generate a model2data studio project's data from the command line."""
+
+
+class Format(str, enum.Enum):
+    csv = "csv"
+    dbt = "dbt"
+
+
+class Adapter(str, enum.Enum):
+    duckdb = "duckdb"
+    postgres = "postgres"
+
+
+KeyOption = typer.Option(
+    None,
+    "--key",
+    envvar=KEY_ENV,
+    help="Studio API key. Create one in the studio's account menu.",
+)
+UrlOption = typer.Option(
+    DEFAULT_URL,
+    "--url",
+    envvar=URL_ENV,
+    help="The studio to talk to. Set this for a self-hosted one.",
+)
+
+
+def _die(message: str) -> NoReturn:
+    typer.secho(message, fg=typer.colors.RED, err=True)
+    raise typer.Exit(1)
+
+
+@contextmanager
+def _studio(url: str, key: str | None) -> Iterator[StudioClient]:
+    """A client for the duration of one command, with its failures printed.
+
+    Every `StudioError` raised inside the block — by the client or by the
+    command's own checks — leaves the process with a message and a non-zero
+    exit, which is the only behaviour a CI step reacts to.
+    """
+    if not key:
+        _die(
+            f"No API key. Pass --key, or set {KEY_ENV} to one created in the studio's account menu."
+        )
+    client = client_factory(url, key)
+    try:
+        yield client
+    except StudioError as error:
+        _die(str(error))
+    finally:
+        client.close()
+
+
+def _table(headers: Sequence[str], rows: Sequence[Sequence[str]]) -> str:
+    """Columns padded to their widest cell. Deliberately not `rich`: the output
+    of a CI step is a log file, and box drawing in one helps nobody."""
+    widths = (
+        [
+            max(len(header), *(len(row[index]) for row in rows))
+            for index, header in enumerate(headers)
+        ]
+        if rows
+        else [len(header) for header in headers]
+    )
+    lines = [
+        "  ".join(header.ljust(widths[index]) for index, header in enumerate(headers)).rstrip()
+    ]
+    lines.append("  ".join("-" * width for width in widths))
+    lines.extend(
+        "  ".join(str(cell).ljust(widths[index]) for index, cell in enumerate(row)).rstrip()
+        for row in rows
+    )
+    return "\n".join(lines)
+
+
+def _day(timestamp: str | None) -> str:
+    """The date out of an ISO instant. A saved project's clock time is noise in
+    a listing whose only question is "which one was I last working on"."""
+    return (timestamp or "")[:10] or "-"
+
+
+def _words(key: str) -> str:
+    """`maxRows` or `max_rows` as "max rows", so a plan's limits read as prose
+    whatever the API decides to call its fields."""
+    spaced = "".join(
+        f" {character.lower()}" if character.isupper() else character for character in key
+    )
+    return spaced.replace("_", " ").strip()
+
+
+def _limit(value: Any) -> str:
+    # null means unlimited everywhere in this API — the plan table, the project
+    # allowance, all of it.
+    if value is None:
+        return "unlimited"
+    if isinstance(value, bool):
+        return "yes" if value else "no"
+    return f"{value:,}" if isinstance(value, int) else str(value)
+
+
+@app.command()
+def projects(key: str | None = KeyOption, url: str = UrlOption) -> None:
+    """List the saved projects this key can see."""
+    with _studio(url, key) as client:
+        saved = client.list_projects()
+
+    if not saved:
+        typer.echo("No saved projects. Create one in the studio first.")
+        return
+    typer.echo(
+        _table(
+            ("id", "name", "rows", "adapter", "updated"),
+            [
+                (
+                    str(project.get("id", "")),
+                    str(project.get("name", "")),
+                    str(project.get("rows", "")),
+                    # Absent on a studio old enough to predate the column.
+                    str(project.get("adapter", "-")),
+                    _day(project.get("updatedAt")),
+                )
+                for project in saved
+            ],
+        )
+    )
+
+
+@app.command()
+def whoami(key: str | None = KeyOption, url: str = UrlOption) -> None:
+    """Show which plan this key is on, and what the account has used."""
+    with _studio(url, key) as client:
+        plan = client.get_plan()
+        usage = client.get_usage()
+
+    typer.echo(f"Plan: {plan.get('name') if plan else usage.get('plan', 'unknown')}")
+    for field, value in (plan or {}).items():
+        if field != "name":
+            typer.echo(f"  {_words(field)}: {_limit(value)}")
+
+    allowance = usage.get("projects") or {}
+    typer.echo(f"Projects: {allowance.get('used', 0)} of {_limit(allowance.get('limit'))}")
+    typer.echo(f"Generations: {usage.get('generations', 0)}")
+
+
+@app.command()
+def mcp() -> None:
+    """Run the MCP server over stdio, for an editor's agent (needs the mcp extra).
+
+    The same server as `model2data-studio-mcp`, reachable through the package's
+    own name: a client that starts a PyPI server runs `uvx <package> ...`, which
+    can only call the command named after the package.
+    """
+    from . import mcp_server
+
+    try:
+        server = mcp_server.build()
+    except StudioError as error:
+        _die(str(error))
+    server.run()
+
+
+@app.command()
+def generate(
+    project_id: str = typer.Option(..., "--project", help="Project id, from `projects`."),
+    out: Path = typer.Option(..., "--out", help="Directory to unpack the archive into."),
+    export_format: Format = typer.Option(
+        Format.csv, "--format", help="csv seeds, or a dbt project."
+    ),
+    rows: int | None = typer.Option(
+        None, "--rows", min=1, help="Override the project's row count."
+    ),
+    rows_for: list[str] | None = typer.Option(
+        None, "--rows-for", metavar="TABLE=N", help="Override one table's row count. Repeatable."
+    ),
+    seed: int | None = typer.Option(
+        None, "--seed", help="Pin the data. Defaults to the project's."
+    ),
+    as_of: str | None = typer.Option(
+        None,
+        "--as-of",
+        metavar="YYYY-MM-DD",
+        help="Pin the day dates and timestamps are anchored on. Defaults to the project's, "
+        "then today.",
+    ),
+    locale: str | None = typer.Option(None, "--locale", help="Override the project's locale."),
+    business_hours: bool | None = typer.Option(
+        None,
+        "--business-hours/--no-business-hours",
+        help="Override the project's business-hours weighting. Pro only.",
+    ),
+    growth: float | None = typer.Option(
+        None,
+        "--growth",
+        min=-1.0,
+        help="Override the project's growth trend across the window. Pro only.",
+    ),
+    seasonality: float | None = typer.Option(
+        None,
+        "--seasonality",
+        min=0.0,
+        max=1.0,
+        help="Override the project's seasonal peak, 0.0-1.0. Pro only.",
+    ),
+    skew: float | None = typer.Option(
+        None,
+        "--skew",
+        min=0.0,
+        max=1.0,
+        help="Override how unevenly child rows spread over their parents, 0.0-1.0. Pro only.",
+    ),
+    name: str | None = typer.Option(None, "--name", help="Override the dbt project name."),
+    adapter: Adapter | None = typer.Option(
+        None, "--adapter", help="Warehouse the dbt project targets."
+    ),
+    force: bool = typer.Option(False, "--force", help="Overwrite files already in --out."),
+    check: bool = typer.Option(
+        False,
+        "--check",
+        help="Compare against --out instead of writing to it. Exits 1 if they differ.",
+    ),
+    key: str | None = KeyOption,
+    url: str = UrlOption,
+) -> None:
+    """Regenerate a saved project's data and write it to a directory."""
+    with _studio(url, key) as client:
+        # Parsed before anything is fetched, so a typo in a flag costs a
+        # message rather than a generation.
+        overrides = _row_overrides(rows_for)
+        # Checked before the fetch, like `--rows-for`, so a typo costs a
+        # message rather than a request -- and with the flag's own name in it,
+        # which is the only reason this is not left to `export.resolve`.
+        if as_of is not None:
+            _anchor(as_of)
+        saved = client.get_project(project_id)
+        # Every fallback the resolution makes -- the seed, the anchor day, one
+        # time-profile field against the project's other two -- lives in
+        # `export.resolve`, so the MCP server generating from the same project
+        # with the same arguments gets the same rows.
+        resolved = export.resolve(
+            saved,
+            export.Overrides(
+                rows=rows,
+                row_overrides=overrides,
+                seed=seed,
+                as_of=as_of,
+                locale=locale,
+                business_hours=business_hours,
+                growth=growth,
+                seasonality=seasonality,
+                skew=skew,
+                project_name=name,
+                adapter=adapter.value if adapter is not None else None,
+            ),
+        )
+        body = resolved.body
+        resolved_seed = resolved.seed
+        picked_seed = resolved.picked_seed
+        anchor = resolved.anchor
+        picked_anchor = resolved.picked_anchor
+
+        archive = client.export(export_format.value, body)
+
+        if check:
+            # Unpacked into a scratch directory rather than --out: --check
+            # reports on the committed files, it never touches them. Still
+            # inside `_studio`'s block: `_unpack` raises `StudioError` too, and
+            # only that block turns one into a clean message instead of a
+            # traceback.
+            with tempfile.TemporaryDirectory() as scratch:
+                fresh = Path(scratch)
+                unpacked = _unpack(archive, fresh, force=True)
+                _print_row_summary(fresh, export_format)
+                changed, missing, extra = _diff_trees(fresh, out)
+
+            if not (changed or missing or extra):
+                typer.echo(f"Up to date: {len(unpacked)} files")
+                return
+            for label, paths in (("Changed", changed), ("Missing", missing), ("Extra", extra)):
+                for path in paths:
+                    typer.echo(f"{label}: {path}")
+            if picked_anchor:
+                # The likeliest reason a check that passed yesterday fails
+                # today, so it is said next to the diff rather than left for
+                # someone to work out from a column of dates.
+                typer.echo(
+                    f"Dates were anchored on today, {anchor:%Y-%m-%d}. If the committed files "
+                    f"were generated on another day, pass --as-of that day to compare fairly."
+                )
+            raise typer.Exit(1)
+
+        written = _unpack(archive, out, force=force)
+
+    typer.echo(f"Wrote {len(written)} file{'' if len(written) == 1 else 's'} to {out}:")
+    for path in written:
+        typer.echo(f"  {path}")
+    _print_row_summary(out, export_format)
+    if picked_seed:
+        typer.echo(
+            f"This project has no saved seed, so seed {resolved_seed} was picked at "
+            f"random. Pass --seed {resolved_seed} to generate these same rows again."
+        )
+    if picked_anchor:
+        typer.echo(
+            f"Dates were anchored on today, {anchor:%Y-%m-%d}. Pass --as-of {anchor:%Y-%m-%d} "
+            f"to generate these same rows again on another day, or pin the day in the "
+            f"studio's generation settings."
+        )
+
+
+def _anchor(as_of: str | None) -> dt.datetime:
+    """`export.anchor_from`, with the flag's name in the error.
+
+    The rule itself is shared with the MCP server; only the wording is the
+    CLI's, because "--as-of wants a day" is a sentence about a terminal.
+    """
+    try:
+        return export.anchor_from(as_of)
+    except StudioError:
+        raise StudioError(f"--as-of wants a day as YYYY-MM-DD, not {as_of!r}.") from None
+
+
+def _row_overrides(values: Sequence[str] | None) -> dict[str, int]:
+    overrides: dict[str, int] = {}
+    for value in values or []:
+        table, separator, count = value.partition("=")
+        if not table or not separator or not count.isdigit() or int(count) < 1:
+            raise StudioError(f"--rows-for wants TABLE=N with N at least 1, not {value!r}.")
+        overrides[table] = int(count)
+    return overrides
+
+
+def _diff_trees(generated: Path, out: Path) -> tuple[list[str], list[str], list[str]]:
+    """How a freshly generated tree differs from what `--out` already holds.
+
+    Three ways two trees can disagree, each meaning something different in CI:
+    a file in both with different bytes is a schema change nobody regenerated
+    fixtures for; one only in the fresh generation is new; one only in `--out`
+    is what a table or column dropped from the schema leaves behind, still
+    committed. `out` need not exist at all — a first run — in which case
+    everything generated is reported missing.
+    """
+    fresh = {path.relative_to(generated) for path in generated.rglob("*") if path.is_file()}
+    committed = (
+        {path.relative_to(out) for path in out.rglob("*") if path.is_file()}
+        if out.exists()
+        else set()
+    )
+
+    changed = sorted(
+        str(rel)
+        for rel in fresh & committed
+        if (generated / rel).read_bytes() != (out / rel).read_bytes()
+    )
+    missing = sorted(str(rel) for rel in fresh - committed)
+    extra = sorted(str(rel) for rel in committed - fresh)
+    return changed, missing, extra
+
+
+def _csv_row_count(path: Path) -> int:
+    """Rows in a CSV, minus its header.
+
+    `csv.reader` rather than counting newlines: a free-text column can carry
+    one inside a quoted value, which would count one row as two.
+    """
+    with path.open(newline="") as handle:
+        return sum(1 for _ in csv.reader(handle)) - 1
+
+
+def _row_counts(root: Path, export_format: Format) -> list[tuple[str, int]]:
+    """Table name and row count, read back out of the generated CSVs.
+
+    A csv export writes one seed per table at the top level of `--out`; a dbt
+    project nests them under `seeds/raw/`, alongside the models and tests that
+    reference them. Read back off disk rather than carried down from
+    generation itself, so this reports what actually landed there.
+    """
+    pattern = "**/seeds/raw/*.csv" if export_format == Format.dbt else "*.csv"
+    return sorted((csv_path.stem, _csv_row_count(csv_path)) for csv_path in root.glob(pattern))
+
+
+def _print_row_summary(root: Path, export_format: Format) -> None:
+    """A per-table row count, so a CI log shows what was generated without
+    anyone having to open the seeds themselves."""
+    counts = _row_counts(root, export_format)
+    if not counts:
+        return
+    typer.echo(_table(("table", "rows"), [(name, f"{count:,}") for name, count in counts]))
+
+
+if __name__ == "__main__":  # pragma: no cover
+    app()