duckhaven-cli 0.1.0__py3-none-any.whl

dh/commands/query.py

@@ -0,0 +1,239 @@
+"""`dh sql` and `dh query` — running SQL and working with the runs.
+
+`dh sql` is the one deliberately verb-first command in the tree.
+"""
+
+from __future__ import annotations
+
+import sys
+
+import typer
+
+from dh import context, execute, repl
+from dh.errors import ConflictError
+
+app = typer.Typer(name="query", help="Run SQL and inspect past runs.")
+
+
+def _read_sql(query: str | None, file: typer.FileText | None, stdin: bool) -> str:
+    """The SQL, from exactly one of the three ways of supplying it."""
+    given = [source for source in (query, file, stdin or None) if source]
+    if len(given) != 1:
+        raise ConflictError(
+            "no_sql",
+            "Give the SQL exactly one way: --query, --file, or --stdin.",
+        )
+    if query:
+        return query
+    if file:
+        return file.read()
+    return sys.stdin.read()
+
+
+def _run(
+    ctx: typer.Context,
+    *,
+    query: str | None,
+    file: typer.FileText | None,
+    stdin: bool,
+    no_wait: bool,
+    timeout: str,
+    limit: int | None,
+    fetch_all: bool,
+    agent: str | None,
+) -> None:
+    cli = context.of(ctx)
+    settings = cli.settings()
+    workspace = settings.require("workspace")
+    if not (query or file or stdin) and repl.is_interactive():
+        # No SQL and a terminal: the user wants a shell, not a usage error.
+        raise typer.Exit(repl.run(cli, settings, agent=agent))
+    sql = _read_sql(query, file, stdin)
+    budget = execute.parse_duration(timeout)
+
+    with cli.client(settings) as client:
+        chosen = execute.resolve_agent(client, agent or settings.get("agent"))
+        submitted = execute.submit(
+            client,
+            workspace,
+            sql,
+            agent=chosen,
+            catalog=settings.get("catalog"),
+            # The same budget bounds the client's wait and the server's execution.
+            # Two timeouts would be a contract people get wrong; one cannot drift.
+            timeout_s=budget,
+        )
+        if no_wait:
+            cli.note(f"Submitted {submitted['id']}.")
+            cli.emit(submitted)
+            return
+
+        finished = execute.wait(client, submitted["id"], timeout=budget)
+        execute.raise_for_status(finished)
+        page = execute.fetch_rows(client, finished["id"], limit=None if fetch_all else limit)
+
+    if page["truncated"]:
+        cli.note(f"Showing {len(page['rows'])} of {page['total']} rows. Use --all for the rest.")
+    cli.emit({"columns": page["columns"], "rows": page["rows"], "total": page["total"]})
+
+
+_QUERY = typer.Option(None, "--query", "-q", help="SQL to run.")
+_FILE = typer.Option(None, "--file", "-f", help="Read the SQL from a file.")
+_STDIN = typer.Option(False, "--stdin", "-i", help="Read the SQL from standard input.")
+_NO_WAIT = typer.Option(False, "--no-wait", help="Print the query id and exit without waiting.")
+_TIMEOUT = typer.Option("20m", "--timeout", help="Give up after this long, e.g. 30s, 20m, 1h.")
+_LIMIT = typer.Option(None, "--limit", help="Stop after this many rows.")
+_ALL = typer.Option(False, "--all", help="Fetch every row, ignoring --limit.")
+_AGENT = typer.Option(None, "--agent", help="Agent name or id to run on.")
+
+
+def sql(
+    ctx: typer.Context,
+    query: str = _QUERY,
+    file: typer.FileText = _FILE,
+    stdin: bool = _STDIN,
+    no_wait: bool = _NO_WAIT,
+    timeout: str = _TIMEOUT,
+    limit: int = _LIMIT,
+    fetch_all: bool = _ALL,
+    agent: str = _AGENT,
+) -> None:
+    """Run SQL and print the results.
+
+    Submits to the workspace's one-shot query route, waits for the run to reach a
+    terminal state, then fetches every page of results. Ctrl-C cancels the query
+    on the server rather than orphaning it.
+
+    A query that runs and fails exits 6, so a pipeline can tell bad SQL from a
+    broken CLI or an unreachable server.
+    """
+    _run(
+        ctx,
+        query=query,
+        file=file,
+        stdin=stdin,
+        no_wait=no_wait,
+        timeout=timeout,
+        limit=limit,
+        fetch_all=fetch_all,
+        agent=agent,
+    )
+
+
+@app.command("run")
+def run_query(
+    ctx: typer.Context,
+    query: str = _QUERY,
+    file: typer.FileText = _FILE,
+    stdin: bool = _STDIN,
+    no_wait: bool = _NO_WAIT,
+    timeout: str = _TIMEOUT,
+    limit: int = _LIMIT,
+    fetch_all: bool = _ALL,
+    agent: str = _AGENT,
+) -> None:
+    """Run SQL and print the results. The noun-first spelling of `dh sql`."""
+    _run(
+        ctx,
+        query=query,
+        file=file,
+        stdin=stdin,
+        no_wait=no_wait,
+        timeout=timeout,
+        limit=limit,
+        fetch_all=fetch_all,
+        agent=agent,
+    )
+
+
+@app.command("get")
+def get_query(ctx: typer.Context, query_id: str) -> None:
+    """One run's status, timings and error."""
+    cli = context.of(ctx)
+    with cli.client() as client:
+        cli.emit(client.get(f"queries/{query_id}"))
+
+
+@app.command("rows")
+def query_rows(
+    ctx: typer.Context,
+    query_id: str,
+    limit: int = _LIMIT,
+    fetch_all: bool = _ALL,
+) -> None:
+    """The results of a finished run, following the cursor to the end."""
+    cli = context.of(ctx)
+    with cli.client() as client:
+        page = execute.fetch_rows(client, query_id, limit=None if fetch_all else limit)
+    if page["truncated"]:
+        cli.note(f"Showing {len(page['rows'])} of {page['total']} rows. Use --all for the rest.")
+    cli.emit({"columns": page["columns"], "rows": page["rows"], "total": page["total"]})
+
+
+@app.command("profile")
+def query_profile(ctx: typer.Context, query_id: str) -> None:
+    """The execution profile captured for a finished run, if there was one."""
+    cli = context.of(ctx)
+    with cli.client() as client:
+        cli.emit(client.get(f"queries/{query_id}/profile"))
+
+
+@app.command("cancel")
+def cancel_query(ctx: typer.Context, query_id: str) -> None:
+    """Ask the agent to stop a running query.
+
+    Idempotent: cancelling a finished query succeeds and changes nothing.
+    """
+    cli = context.of(ctx)
+    with cli.client() as client:
+        client.delete(f"queries/{query_id}")
+    cli.note(f"Cancelled {query_id}.")
+
+
+@app.command("list")
+def list_queries(
+    ctx: typer.Context,
+    status: list[str] = typer.Option(
+        None, "--status", help="Repeatable: queued, running, done, failed, cancelled."
+    ),
+    statement_type: list[str] = typer.Option(None, "--statement-type", help="Repeatable."),
+    since: str = typer.Option(None, "--since", help="ISO-8601 lower bound on start time."),
+    until: str = typer.Option(None, "--until", help="ISO-8601 upper bound on start time."),
+    origin: str = typer.Option(None, "--origin", help="interactive, session, schedule, elastic."),
+    session: str = typer.Option(None, "--session", help="Only statements from this session."),
+    agent: str = typer.Option(None, "--agent", help="Only runs on this agent id."),
+    user: str = typer.Option(None, "--user", help="Only this user's runs (admin for others)."),
+    search: str = typer.Option(None, "--search", "-q", help="Free text over the SQL."),
+    slower_than: int = typer.Option(None, "--slower-than", help="Milliseconds."),
+    sort: str = typer.Option(None, "--sort", help="started_at or duration."),
+    direction: str = typer.Option(None, "--dir", help="asc or desc."),
+    all_workspaces: bool = typer.Option(False, "--all-workspaces", help="Admin only."),
+    limit: int = typer.Option(None, "--limit", help="Rows per page."),
+    fetch_all: bool = typer.Option(False, "--all", help="Walk every page."),
+) -> None:
+    """The query log, newest first. Doubles as the audit trail.
+
+    Filtering and sorting happen server-side over the whole history before the
+    page is cut, so narrowing here is not the same as filtering what you were
+    shown. Omitting a filter means "everything": `--status` has no default.
+    """
+    cli = context.of(ctx)
+    settings = cli.settings()
+    params = {
+        "status": status or None,
+        "statement_type": statement_type or None,
+        "since": since,
+        "until": until,
+        "origin": origin,
+        "session_id": session,
+        "agent_id": agent,
+        "user_id": user,
+        "q": search,
+        "slower_than_ms": slower_than,
+        "sort": sort,
+        "dir": direction,
+        "all_workspaces": all_workspaces or None,
+    }
+    path = f"workspaces/{settings.require('workspace')}/queries"
+    with cli.client(settings) as client:
+        cli.page(client, path, params=params, limit=limit, fetch_all=fetch_all)

dh/commands/saved.py

@@ -0,0 +1,213 @@
+"""`dh saved-query`, `dh schedule` and `dh search`.
+
+Saved queries and schedules are the two places where a CLI beats the worksheet
+UI outright: both are things people want in version control and applied from CI,
+not clicked into a browser.
+"""
+
+from __future__ import annotations
+
+import typer
+
+from dh import context
+from dh.errors import ConflictError
+
+saved_app = typer.Typer(name="saved-query", help="Named SQL saved in the workspace.")
+schedule_app = typer.Typer(name="schedule", help="Cron schedules for saved queries.")
+
+
+def _paged(cli, path: str, *, params=None, limit=None, fetch_all=False) -> None:
+    """List through the shared envelope, whichever shape the endpoint returns."""
+    with cli.client() as client:
+        cli.page(client, path, params=params, limit=limit, fetch_all=fetch_all)
+
+
+# --- Saved queries ---------------------------------------------------------
+
+
+@saved_app.command("list")
+def list_saved(
+    ctx: typer.Context,
+    limit: int = typer.Option(None, "--limit"),
+    fetch_all: bool = typer.Option(False, "--all"),
+) -> None:
+    """The workspace's saved queries, newest first."""
+    cli = context.of(ctx)
+    workspace = cli.settings().require("workspace")
+    _paged(cli, f"workspaces/{workspace}/saved-queries", limit=limit, fetch_all=fetch_all)
+
+
+@saved_app.command("create")
+def create_saved(
+    ctx: typer.Context,
+    name: str = typer.Argument(..., help="Name to save under."),
+    query: str = typer.Option(None, "--query", "-q", help="SQL to save."),
+    file: typer.FileText = typer.Option(None, "--file", "-f", help="Read the SQL from a file."),
+    agent: str = typer.Option(None, "--agent", help="Default agent id for runs."),
+) -> None:
+    """Save SQL under a name, replacing any query already using that name.
+
+    Overwrite-by-name is the server's behaviour and is deliberate: saving over
+    "report" updates that query rather than accumulating "report v2".
+    """
+    cli = context.of(ctx)
+    settings = cli.settings()
+    if bool(query) == bool(file):
+        raise ConflictError("no_sql", "Give the SQL exactly one way: --query or --file.")
+    body = {"name": name, "sql": query or file.read()}
+    if agent:
+        body["default_agent_id"] = agent
+    with cli.client(settings) as client:
+        cli.emit(
+            client.post(f"workspaces/{settings.require('workspace')}/saved-queries", json=body)
+        )
+
+
+@saved_app.command("update")
+def update_saved(
+    ctx: typer.Context,
+    saved_query_id: str,
+    name: str = typer.Option(None, "--name"),
+    query: str = typer.Option(None, "--query", "-q"),
+    file: typer.FileText = typer.Option(None, "--file", "-f"),
+    agent: str = typer.Option(None, "--agent"),
+) -> None:
+    """Change a saved query's name, SQL or default agent. Omitted fields are left alone."""
+    cli = context.of(ctx)
+    settings = cli.settings()
+    body = {
+        k: v
+        for k, v in {
+            "name": name,
+            "sql": query or (file.read() if file else None),
+            "default_agent_id": agent,
+        }.items()
+        if v is not None
+    }
+    if not body:
+        raise ConflictError("nothing_to_update", "Give at least one of --name, --query or --agent.")
+    path = f"workspaces/{settings.require('workspace')}/saved-queries/{saved_query_id}"
+    with cli.client(settings) as client:
+        cli.emit(client.patch(path, json=body))
+
+
+@saved_app.command("delete")
+def delete_saved(ctx: typer.Context, saved_query_id: str) -> None:
+    """Delete a saved query."""
+    cli = context.of(ctx)
+    settings = cli.settings()
+    path = f"workspaces/{settings.require('workspace')}/saved-queries/{saved_query_id}"
+    with cli.client(settings) as client:
+        client.delete(path)
+    cli.note(f"Deleted saved query {saved_query_id}.")
+
+
+# --- Schedules -------------------------------------------------------------
+
+
+@schedule_app.command("list")
+def list_schedules(ctx: typer.Context) -> None:
+    """The workspace's schedules."""
+    cli = context.of(ctx)
+    workspace = cli.settings().require("workspace")
+    _paged(cli, f"workspaces/{workspace}/schedules")
+
+
+@schedule_app.command("create")
+def create_schedule(
+    ctx: typer.Context,
+    saved_query_id: str = typer.Argument(..., help="The saved query to run."),
+    cron: str = typer.Option(..., "--cron", help="Five-field cron expression."),
+    agent: str = typer.Option(None, "--agent", help="Agent id to run on."),
+    disabled: bool = typer.Option(False, "--disabled", help="Create it paused."),
+) -> None:
+    """Schedule a saved query to run on a cron expression."""
+    cli = context.of(ctx)
+    settings = cli.settings()
+    body = {
+        "job_type": "saved_query",
+        "saved_query_id": saved_query_id,
+        "cron": cron,
+        "enabled": not disabled,
+    }
+    if agent:
+        body["agent_id"] = agent
+    with cli.client(settings) as client:
+        cli.emit(client.post(f"workspaces/{settings.require('workspace')}/schedules", json=body))
+
+
+@schedule_app.command("update")
+def update_schedule(
+    ctx: typer.Context,
+    schedule_id: str,
+    cron: str = typer.Option(None, "--cron"),
+    agent: str = typer.Option(None, "--agent"),
+    enabled: bool = typer.Option(None, "--enabled/--disabled"),
+) -> None:
+    """Change a schedule's cron, agent, or whether it runs at all."""
+    cli = context.of(ctx)
+    settings = cli.settings()
+    body = {
+        k: v
+        for k, v in {"cron": cron, "agent_id": agent, "enabled": enabled}.items()
+        if v is not None
+    }
+    if not body:
+        raise ConflictError(
+            "nothing_to_update", "Give at least one of --cron, --agent, --enabled or --disabled."
+        )
+    path = f"workspaces/{settings.require('workspace')}/schedules/{schedule_id}"
+    with cli.client(settings) as client:
+        cli.emit(client.patch(path, json=body))
+
+
+@schedule_app.command("delete")
+def delete_schedule(ctx: typer.Context, schedule_id: str) -> None:
+    """Delete a schedule. The saved query it ran is left alone."""
+    cli = context.of(ctx)
+    settings = cli.settings()
+    with cli.client(settings) as client:
+        client.delete(f"workspaces/{settings.require('workspace')}/schedules/{schedule_id}")
+    cli.note(f"Deleted schedule {schedule_id}.")
+
+
+@schedule_app.command("runs")
+def schedule_runs(
+    ctx: typer.Context,
+    schedule_id: str = typer.Argument(None, help="Omit for every schedule's runs."),
+    limit: int = typer.Option(None, "--limit"),
+    fetch_all: bool = typer.Option(False, "--all"),
+) -> None:
+    """Runs produced by a schedule, or by all of them."""
+    cli = context.of(ctx)
+    workspace = cli.settings().require("workspace")
+    path = (
+        f"workspaces/{workspace}/schedules/{schedule_id}/runs"
+        if schedule_id
+        else f"workspaces/{workspace}/schedule-runs"
+    )
+    _paged(cli, path, limit=limit, fetch_all=fetch_all)
+
+
+# --- Search ----------------------------------------------------------------
+
+
+def search(
+    ctx: typer.Context,
+    query: str = typer.Argument(..., help="Text to look for."),
+    limit: int = typer.Option(None, "--limit"),
+) -> None:
+    """Find catalogs, schemas, tables and saved queries by name.
+
+    A truncated report rather than a page: `limit` caps how many come back and
+    `has_more` says whether it cut. There is no cursor to walk -- narrow the
+    query instead.
+    """
+    cli = context.of(ctx)
+    settings = cli.settings()
+    with cli.client(settings) as client:
+        body = client.get(
+            f"workspaces/{settings.require('workspace')}/search",
+            params={"q": query, "limit": limit},
+        )
+    cli.emit(body.get("items", body), has_more=bool(body.get("has_more")))

dh/commands/session.py

@@ -0,0 +1,210 @@
+"""`dh session` and the interactive REPL.
+
+Two different jobs, so two different mechanisms.
+
+`dh session *` operates a session's lifecycle from *separate* CLI invocations --
+open now, run statements later, close tomorrow -- so it is plain REST over the
+session routes; a connection object cannot outlive the process that made it.
+
+The REPL is the opposite: one long-lived process holding one session, which is
+exactly what `duckhaven-sql-connector` is for. It handles elastic cold-start
+waiting, retries and row pagination, so the REPL does not reimplement any of it.
+
+Sessions are disabled by default on the server and the routes answer **404** when
+they are. Rendered as "not found" that would be actively misleading, so it is
+translated -- and the REPL falls back to one-shot execution rather than refusing
+to start.
+"""
+
+from __future__ import annotations
+
+import time
+
+import typer
+
+from dh import context, execute
+from dh.errors import DhError, NotFoundError
+
+app = typer.Typer(name="session", help="Stateful SQL sessions, for dbt, dlt and the REPL.")
+
+#: Long enough to cover an elastic cold start, which is what the wait is for.
+_DEFAULT_WAIT = 300.0
+
+#: How long the *server* is asked to hold the request before answering.
+#:
+#: It must stay under the client's socket timeout: a longer server-side block
+#: aborts the request socket-side while the server goes on to open a session
+#: nobody holds. The connector in this workspace documents the same hazard. Any
+#: remaining wait is spent polling instead, which is interruptible and leaves the
+#: session id in hand.
+_SERVER_HOLD_S = 10.0
+
+
+class SessionsDisabled(DhError):
+    """The session surface is switched off on this deployment."""
+
+
+def _translate(exc: DhError) -> DhError:
+    """Turn the disabled-surface 404 into something actionable."""
+    if isinstance(exc, NotFoundError) and "not enabled" in exc.message.lower():
+        return SessionsDisabled(
+            "sessions_disabled",
+            "SQL sessions are not enabled on this server. An operator turns them on with "
+            "DH_SQL_SESSIONS_ENABLED=true; `dh sql` works without them.",
+        )
+    return exc
+
+
+@app.command("open")
+def open_session(
+    ctx: typer.Context,
+    agent: str = typer.Option(None, "--agent", help="Agent name or id to bind the session to."),
+    catalog: str = typer.Option(None, "--catalog"),
+    wait: float = typer.Option(_DEFAULT_WAIT, "--wait", help="Seconds to wait for compute."),
+    no_wait: bool = typer.Option(
+        False, "--no-wait", help="Return a pending session instead of holding the request."
+    ),
+) -> None:
+    """Open a session and print its id.
+
+    Opening can take a while when compute has to start, so `--no-wait` hands back
+    a pending session to poll with `dh session get` instead of holding the
+    request open.
+    """
+    cli = context.of(ctx)
+    settings = cli.settings()
+    # `--no-wait` must not hold the request at all; otherwise ask the server for a
+    # short hold and poll out the rest ourselves.
+    body: dict[str, object] = {
+        "wait_timeout_s": 0.0 if no_wait else min(wait, _SERVER_HOLD_S),
+        "on_wait_timeout": "continue",
+    }
+    if catalog or settings.get("catalog"):
+        body["catalog"] = catalog or settings.get("catalog")
+    with cli.client(settings) as client:
+        chosen = execute.resolve_agent(client, agent or settings.get("agent"))
+        if chosen:
+            body["agent_id"] = chosen
+        try:
+            session = client.post(
+                f"workspaces/{settings.require('workspace')}/sql/sessions", json=body
+            )
+        except DhError as exc:
+            raise _translate(exc) from exc
+        if not no_wait:
+            session = _await_open(cli, client, session, deadline=wait)
+    cli.note(f"Session {session['id']} is {session.get('status')}.")
+    cli.emit(session)
+
+
+def _await_open(cli, client, session: dict, *, deadline: float) -> dict:
+    """Poll a pending session to `open`, or hand back whatever it reached.
+
+    `on_wait_timeout=continue` means the server answers with a session still
+    starting rather than failing, so the id is already in hand -- an interrupt or
+    a timeout here leaves something the caller can close, not an orphan.
+    """
+    started = time.monotonic()
+    delay = 0.5
+    while session.get("status") in ("pending", "opening"):
+        if time.monotonic() - started >= deadline:
+            cli.note(
+                f"Session {session['id']} is still {session.get('status')}; "
+                "poll it with `dh session get`."
+            )
+            return session
+        time.sleep(min(delay, 2.0))
+        delay *= 1.5
+        session = client.get(f"sql/sessions/{session['id']}")
+    return session
+
+
+@app.command("list")
+def list_sessions(
+    ctx: typer.Context,
+    status: list[str] = typer.Option(None, "--status", help="Repeatable."),
+    fetch_all: bool = typer.Option(False, "--all"),
+) -> None:
+    """The workspace's sessions, newest first. The audit list."""
+    cli = context.of(ctx)
+    settings = cli.settings()
+    path = f"workspaces/{settings.require('workspace')}/sql/sessions"
+    params = {"status": status or None}
+    with cli.client(settings) as client:
+        cli.page(client, path, params=params, fetch_all=fetch_all, translate=_translate)
+
+
+@app.command("get")
+def get_session(ctx: typer.Context, session_id: str) -> None:
+    """One session's status and the agent holding it."""
+    cli = context.of(ctx)
+    with cli.client() as client:
+        try:
+            cli.emit(client.get(f"sql/sessions/{session_id}"))
+        except DhError as exc:
+            raise _translate(exc) from exc
+
+
+@app.command("close")
+def close_session(ctx: typer.Context, session_id: str) -> None:
+    """Close a session and release its connection.
+
+    Idempotent. Everything session-local -- temp tables, `SET`s -- is gone
+    afterwards.
+    """
+    cli = context.of(ctx)
+    with cli.client() as client:
+        try:
+            client.delete(f"sql/sessions/{session_id}")
+        except DhError as exc:
+            raise _translate(exc) from exc
+    cli.note(f"Closed session {session_id}.")
+
+
+@app.command("exec")
+def exec_statement(
+    ctx: typer.Context,
+    session_id: str,
+    query: str = typer.Option(None, "--query", "-q"),
+    file: typer.FileText = typer.Option(None, "--file", "-f"),
+    timeout: str = typer.Option("20m", "--timeout"),
+    limit: int = typer.Option(None, "--limit"),
+    fetch_all: bool = typer.Option(False, "--all"),
+) -> None:
+    """Run one statement on an existing session's connection.
+
+    Unlike `dh sql`, this reuses a held connection, so temp tables, `SET`s and
+    attached catalogs persist between statements.
+    """
+    from dh.commands.query import _read_sql
+
+    cli = context.of(ctx)
+    sql = _read_sql(query, file, False)
+    budget = execute.parse_duration(timeout)
+    with cli.client() as client:
+        try:
+            statement = client.post(
+                f"sql/sessions/{session_id}/statements",
+                json={"sql": sql, "timeout_s": budget},
+            )
+        except DhError as exc:
+            raise _translate(exc) from exc
+        finished = execute.wait(client, statement["id"], timeout=budget)
+        execute.raise_for_status(finished)
+        page = execute.fetch_rows(client, finished["id"], limit=None if fetch_all else limit)
+    cli.emit({"columns": page["columns"], "rows": page["rows"], "total": page["total"]})
+
+
+@app.command("statements")
+def list_statements(
+    ctx: typer.Context, session_id: str, fetch_all: bool = typer.Option(False, "--all")
+) -> None:
+    """A session's statements in execution order.
+
+    Ascending, unlike the newest-first history feeds: a session is one workload
+    read top to bottom, so a `dbt run` reads as the sequence it actually was.
+    """
+    cli = context.of(ctx)
+    path = f"sql/sessions/{session_id}/statements"
+    with cli.client() as client:
+        cli.page(client, path, fetch_all=fetch_all, translate=_translate)