dataplat 0.1.0__py3-none-any.whl

dataplat/cli/db/__init__.py

@@ -0,0 +1,342 @@
+"""Database query CLI commands."""
+
+from __future__ import annotations
+
+import csv
+import json
+import re
+import sys
+from enum import Enum
+from shutil import get_terminal_size
+from time import perf_counter
+from typing import cast
+
+import typer
+from psycopg.abc import Query
+from rich import box
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+from rich.table import Table
+
+from dataplat.cli.db._common import (
+    ConnCliParams,
+    DatabaseOption,
+    EngineOption,
+    EnvPrefixOption,
+    HostOption,
+    PasswordOption,
+    PortOption,
+    SslmodeOption,
+    TargetOption,
+    UserOption,
+    db_session,
+    resolve_params_or_exit,
+)
+from dataplat.cli.db.dbt_orphans import app as dbt_orphans_app
+from dataplat.cli.db.describe import app as describe_app
+from dataplat.cli.db.long_queries import kill_command, long_queries_command
+from dataplat.cli.db.role import app as role_app
+from dataplat.cli.db.top_tables import top_tables_command
+from dataplat.services.db.connection import SqlEngine
+
+app = typer.Typer(
+    name="db",
+    help="Database query commands",
+    no_args_is_help=True,
+)
+app.add_typer(dbt_orphans_app, name="dbt-orphans")
+app.add_typer(describe_app, name="describe")
+app.add_typer(role_app, name="role")
+app.command(
+    "top-tables",
+    help="Rank largest tables in schemas matching a prefix (default: dev_*).",
+)(top_tables_command)
+app.command(
+    "long-queries",
+    help="Show long-running (and recently failed) queries per target.",
+)(long_queries_command)
+app.command(
+    "kill",
+    help="Cancel or terminate running queries by PID.",
+)(kill_command)
+
+console = Console()
+err_console = Console(stderr=True)
+
+
+class OutputFormat(str, Enum):
+    """Result rendering for ``dp db query``."""
+
+    table = "table"
+    csv = "csv"
+    json = "json"
+
+
+_SQL_COMMENT_RE = re.compile(r"(--[^\n]*)|(/\*.*?\*/)", re.DOTALL)
+_WRITE_KEYWORD_RE = re.compile(r"\b(insert|update|delete|merge)\b", re.IGNORECASE)
+_READ_FIRST_KEYWORDS = {"select", "show", "explain", "table", "values"}
+
+
+def _strip_sql_comments(sql: str) -> str:
+    return _SQL_COMMENT_RE.sub(" ", sql)
+
+
+def _classify_sql(sql: str) -> str:
+    """Classify a statement as ``read`` or ``write``.
+
+    Conservative: a WITH query containing any data-modifying keyword is
+    treated as a write even if the keyword only appears in a literal.
+    """
+    cleaned = _strip_sql_comments(sql).strip()
+    if not cleaned:
+        return "read"
+    first = cleaned.split(None, 1)[0].lower().rstrip(";")
+    if first in _READ_FIRST_KEYWORDS:
+        return "read"
+    if first == "with":
+        return "write" if _WRITE_KEYWORD_RE.search(cleaned) else "read"
+    return "write"
+
+
+def _supports_live_query_progress() -> bool:
+    return isinstance(console, Console) and console.is_terminal and not console.is_dumb_terminal
+
+
+def _load_sql(sql: str | None) -> str:
+    if sql and sql.strip():
+        return sql
+    if not sys.stdin.isatty():
+        return sys.stdin.read()
+    console.print(
+        "[red]Error: SQL is required. Provide it as an argument or via stdin.[/red]"
+    )
+    raise typer.Exit(code=1)
+
+
+def _rows_fit_terminal(row_count: int) -> bool:
+    return row_count <= _max_rows_for_terminal()
+
+
+def _max_rows_for_terminal() -> int:
+    term_lines = get_terminal_size((120, 30)).lines
+    max_rows = term_lines - 6
+    return max(1, max_rows)
+
+
+def _default_page_limit() -> int:
+    return min(100, _max_rows_for_terminal())
+
+
+def _format_cell(value: object, max_length: int) -> str:
+    if value is None:
+        return ""
+    text = str(value)
+    if max_length > 0 and len(text) > max_length:
+        return text[: max_length - 1] + "…"
+    return text
+
+
+def _render_rows(columns: list[str], rows: list[tuple], start_index: int = 1) -> None:
+    display_rows = rows
+    if rows and not _rows_fit_terminal(len(rows)):
+        max_rows = _max_rows_for_terminal()
+        display_rows = rows[:max_rows]
+        console.print(
+            "[yellow]Result too large for terminal. "
+            f"Showing top {len(display_rows)} rows.[/yellow]"
+        )
+
+    term_width = get_terminal_size((120, 30)).columns
+    available_width = max(40, term_width - 6)
+    per_col = max(12, min(40, available_width // max(1, len(columns) + 1)))
+    max_cell_length = max(40, per_col * 4)
+
+    table = Table(
+        show_header=True,
+        header_style="bold cyan",
+        box=box.SIMPLE,
+        show_lines=True,
+    )
+    table.add_column("#", style="dim", justify="right", no_wrap=True, width=4)
+    for column in columns:
+        table.add_column(column, overflow="fold", max_width=per_col)
+
+    for idx, row in enumerate(display_rows, start=start_index):
+        table.add_row(
+            str(idx), *[_format_cell(value, max_cell_length) for value in row]
+        )
+
+    console.print(table)
+    console.print(f"[dim]Rows: {len(rows)}[/dim]")
+
+
+def _emit_csv(columns: list[str], rows: list[tuple]) -> None:
+    writer = csv.writer(sys.stdout)
+    writer.writerow(columns)
+    writer.writerows(rows)
+
+
+def _emit_json(columns: list[str], rows: list[tuple]) -> None:
+    payload = [dict(zip(columns, row, strict=False)) for row in rows]
+    typer.echo(json.dumps(payload, indent=2, default=str))
+
+
+def _confirm_write(sql_text: str, allow_write: bool) -> None:
+    """Gate data-modifying statements behind --write or an interactive confirm."""
+    if allow_write:
+        return
+    preview = " ".join(sql_text.split())
+    if len(preview) > 120:
+        preview = preview[:119] + "…"
+    if sys.stdin.isatty():
+        console.print(f"[yellow]This statement can modify data:[/yellow] {preview}")
+        if typer.confirm("Continue?", default=False):
+            return
+        raise typer.Exit(code=1)
+    console.print(
+        "[red]Error: this statement can modify data. "
+        "Pass --write to run it non-interactively.[/red]"
+    )
+    raise typer.Exit(code=1)
+
+
+def _execute_query(
+    *,
+    sql: str | None,
+    conn_cli: ConnCliParams,
+    limit: int | None,
+    page: int,
+    write: bool,
+    fmt: OutputFormat,
+) -> None:
+    sql_text = _load_sql(sql)
+    if _classify_sql(sql_text) == "write":
+        _confirm_write(sql_text, write)
+
+    cleaned = sql_text.strip().rstrip(";")
+    stripped = _strip_sql_comments(cleaned).strip()
+    first_keyword = stripped.split(None, 1)[0].lower() if stripped else ""
+    is_select = first_keyword in {"select", "with"} and _classify_sql(cleaned) == "read"
+
+    # --limit 0 disables the LIMIT/OFFSET wrapper entirely.
+    paginate = is_select and (limit is None or limit > 0)
+    safe_limit = 1
+    offset = 0
+    if paginate:
+        effective_limit = limit if limit is not None else _default_page_limit()
+        safe_limit = max(1, effective_limit)
+        safe_page = max(1, page)
+        offset = (safe_page - 1) * safe_limit
+        # Newlines guard against a trailing `-- comment` swallowing the paren.
+        sql_text = (
+            f"SELECT * FROM (\n{cleaned}\n) AS dna_sql"
+            f" LIMIT {safe_limit + 1} OFFSET {offset}"
+        )
+
+    conn_params = resolve_params_or_exit(conn_cli)
+    # Decorative output goes to stderr for machine-readable formats.
+    note = console if fmt == OutputFormat.table else err_console
+
+    with db_session(conn_params) as conn, conn.cursor() as cursor:
+        started = perf_counter()
+        rows: list[tuple] = []
+        columns: list[str] = []
+        has_result_set = False
+
+        if _supports_live_query_progress():
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[cyan]Running query...[/cyan]"),
+                TimeElapsedColumn(),
+                console=console,
+                transient=True,
+            ) as progress:
+                progress.add_task("query", total=None)
+                cursor.execute(cast(Query, sql_text))
+                if cursor.description:
+                    has_result_set = True
+                    rows = cursor.fetchall()
+                    columns = [desc.name for desc in cursor.description]
+        else:
+            cursor.execute(cast(Query, sql_text))
+            if cursor.description:
+                has_result_set = True
+                rows = cursor.fetchall()
+                columns = [desc.name for desc in cursor.description]
+
+        if has_result_set:
+            visible = rows
+            more = False
+            if paginate and len(rows) > safe_limit:
+                visible = rows[:safe_limit]
+                more = True
+            if fmt == OutputFormat.csv:
+                _emit_csv(columns, visible)
+            elif fmt == OutputFormat.json:
+                _emit_json(columns, visible)
+            else:
+                start_index = offset + 1 if paginate else 1
+                _render_rows(columns, visible, start_index=start_index)
+            if more:
+                note.print(
+                    "[yellow]More rows available. "
+                    f"Use --page {page + 1} to fetch the next page.[/yellow]"
+                )
+        else:
+            note.print(f"[green]✓ {cursor.rowcount} rows affected[/green]")
+        elapsed = perf_counter() - started
+        note.print(f"[dim]Execution time: {elapsed:.3f}s[/dim]")
+
+
+@app.command("query")
+def query(
+    sql: str | None = typer.Argument(
+        None, help="SQL to execute. If omitted, SQL is read from stdin."
+    ),
+    target: str | None = TargetOption,
+    engine: SqlEngine | None = EngineOption,
+    user: str | None = UserOption,
+    password: str | None = PasswordOption,
+    database: str | None = DatabaseOption,
+    host: str | None = HostOption,
+    port: int | None = PortOption,
+    sslmode: str | None = SslmodeOption,
+    env_prefix: str | None = EnvPrefixOption,
+    limit: int | None = typer.Option(
+        None,
+        "--limit",
+        "-n",
+        help=(
+            "Max rows per page for SELECT queries "
+            "(default: min(100, terminal height); 0 = no limit)"
+        ),
+    ),
+    page: int = typer.Option(1, "--page", help="Page number for SELECT queries"),
+    fmt: OutputFormat = typer.Option(
+        OutputFormat.table, "--format", help="Output format: table, csv, or json."
+    ),
+    write: bool = typer.Option(
+        False,
+        "--write",
+        help="Allow statements that modify data without prompting.",
+    ),
+) -> None:
+    """Run ad-hoc SQL against Postgres or Redshift."""
+    _execute_query(
+        sql=sql,
+        conn_cli=ConnCliParams(
+            target=target,
+            engine=engine,
+            user=user,
+            password=password,
+            database=database,
+            host=host,
+            port=port,
+            sslmode=sslmode,
+            env_prefix=env_prefix,
+        ),
+        limit=limit,
+        page=page,
+        write=write,
+        fmt=fmt,
+    )

dataplat/cli/db/_common.py

@@ -0,0 +1,130 @@
+"""Shared Typer options and connection plumbing for db commands.
+
+Every db command takes the same connection surface: a named ``--target``
+plus low-level overrides. Declaring the options here keeps flags, shorts,
+and help text identical across commands, and ``db_session`` centralizes the
+resolve/connect/error-exit dance that used to be copy-pasted per command.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any
+
+import psycopg
+import typer
+from rich.console import Console
+
+from dataplat.core.errors import DataplatError
+from dataplat.services.db.connection import (
+    DbConnectionParams,
+    SqlEngine,
+    resolve_connection_params,
+)
+from dataplat.services.db.targets import default_target_name, resolve_target
+
+console = Console()
+
+TargetOption = typer.Option(
+    None,
+    "--target",
+    "-t",
+    help="Named DB target from DP_TARGETS (default: DP_DEFAULT_TARGET). "
+    "Sets engine and env prefix.",
+)
+EngineOption = typer.Option(
+    None,
+    "--engine",
+    "-e",
+    help="postgresql or redshift. Overrides the target/<PREFIX>_ENGINE default.",
+)
+UserOption = typer.Option(None, "--user", "-u", help="Database username.")
+PasswordOption = typer.Option(
+    None,
+    "--password",
+    help="Database password (prefer <PREFIX>_PASSWORD env var).",
+)
+DatabaseOption = typer.Option(None, "--database", "-d", help="Database name.")
+HostOption = typer.Option(None, "--host", "-H", help="Database host.")
+PortOption = typer.Option(None, "--port", help="Database port.")
+SslmodeOption = typer.Option(
+    None, "--sslmode", help="SSL mode (e.g., require, prefer, disable)."
+)
+EnvPrefixOption = typer.Option(
+    None,
+    "--env-prefix",
+    help="Env var prefix for connection settings (default: the target's prefix).",
+)
+JsonOption = typer.Option(False, "--json", help="Emit JSON instead of tables.")
+YesOption = typer.Option(
+    False, "--yes", "-y", help="Skip the confirmation prompt."
+)
+
+
+def limit_option(default: int, help_text: str) -> Any:
+    """A per-command row cap with the shared --limit/-n spelling."""
+    return typer.Option(default, "--limit", "-n", help=help_text)
+
+
+@dataclass
+class ConnCliParams:
+    """Raw connection-related CLI values, before resolution."""
+
+    target: str | None = None
+    engine: SqlEngine | None = None
+    user: str | None = None
+    password: str | None = None
+    database: str | None = None
+    host: str | None = None
+    port: int | None = None
+    sslmode: str | None = None
+    env_prefix: str | None = None
+
+    def resolve(self) -> DbConnectionParams:
+        """Resolve to concrete connection params.
+
+        Precedence: explicit flags > target-derived defaults > env vars.
+        """
+        env_prefix = self.env_prefix
+        engine = self.engine
+        target_name = self.target
+        if target_name is None and env_prefix is None:
+            target_name = default_target_name()
+        if target_name:
+            target = resolve_target(target_name)
+            if env_prefix is None:
+                env_prefix = target.env_prefix
+            if engine is None:
+                engine = target.engine
+        return resolve_connection_params(
+            engine=engine,
+            env_prefix=env_prefix,
+            user=self.user,
+            password=self.password,
+            host=self.host,
+            port=self.port,
+            database=self.database,
+            sslmode=self.sslmode,
+        )
+
+
+def resolve_params_or_exit(params: ConnCliParams) -> DbConnectionParams:
+    """Resolve connection params, printing a friendly error on failure."""
+    try:
+        return params.resolve()
+    except DataplatError as exc:
+        console.print(f"[red]Error: {exc}[/red]")
+        raise typer.Exit(code=1) from exc
+
+
+@contextmanager
+def db_session(params: DbConnectionParams) -> Iterator[psycopg.Connection]:
+    """Open a connection; translate psycopg errors into a clean exit."""
+    try:
+        with psycopg.connect(**params.as_psycopg_kwargs()) as conn:  # type: ignore[arg-type]
+            yield conn
+    except psycopg.Error as exc:
+        console.print(f"[red]Database error: {exc}[/red]")
+        raise typer.Exit(code=1) from exc

dataplat/cli/db/_report.py

@@ -0,0 +1,180 @@
+"""Shared Rich rendering helpers for the db command group's report style.
+
+Used by ``describe``, ``role``, and ``top-tables`` so all printed reports
+share one visual language: numbered sections with captions, a rounded title
+card, HORIZONTALS tables, and consistent size/row formatting.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+from rich import box as _box
+from rich.align import Align
+from rich.console import Console
+from rich.padding import Padding
+from rich.panel import Panel
+from rich.table import Table
+
+DASH = "[dim]—[/dim]"
+
+
+def dim(text: str) -> str:
+    return f"[dim]{text}[/dim]"
+
+
+def green(text: str) -> str:
+    return f"[green]{text}[/green]"
+
+
+def fmt_size(n: int | None, *, colored: bool = True) -> str:
+    """Human-friendly byte count; wrapped in [green] when colored."""
+    if n is None:
+        return DASH if colored else "—"
+    size = float(n)
+    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
+        if size < 1024:
+            text = f"{size:.0f} {unit}" if unit == "B" else f"{size:.1f} {unit}"
+            return green(text) if colored else text
+        size /= 1024
+    text = f"{size:.1f} PiB"
+    return green(text) if colored else text
+
+
+def fmt_size_plain(n: int | None) -> str:
+    """Uncolored byte size — for use inside already-styled strings."""
+    return fmt_size(n, colored=False)
+
+
+def fmt_rows(n: int | None, *, colored: bool = True) -> str:
+    if n is None or n < 0:
+        return DASH if colored else "—"
+    text = f"{n:,}"
+    return green(text) if colored else text
+
+
+def indent(renderable: Any, cols: int = 2) -> Padding:
+    """Indent a renderable by the body gutter."""
+    return Padding(renderable, (0, 0, 0, cols))
+
+
+class SectionCounter:
+    """Monotonically increasing section number, reset per report."""
+
+    __slots__ = ("_n",)
+
+    def __init__(self) -> None:
+        self._n = 0
+
+    def next(self) -> int:
+        self._n += 1
+        return self._n
+
+
+def print_section_heading(
+    console: Console, counter: SectionCounter, title: str, caption: str
+) -> None:
+    """Print `  N. Title` + caption below it, with the report spacing rules.
+
+    Caller is responsible for printing a blank line between caption and body
+    (the body itself may begin with a blank line if desired).
+    """
+    n = counter.next()
+    if n > 1:
+        # Two blank lines between end-of-section body and next heading.
+        console.print()
+        console.print()
+    console.print(f"  [bold cyan]{n}.[/bold cyan] [bold]{title}[/bold]")
+    console.print(f"    [dim italic]{caption}[/dim italic]")
+    console.print()
+
+
+def metadata_grid(pairs: Iterable[tuple[str, str]], *, two_column: bool) -> Table:
+    """Label/value grid used inside the title card.
+
+    Labels are right-aligned and dim; values default color.
+    """
+    grid = Table.grid(expand=False)
+    items = [(k, v) for k, v in pairs if v]
+    if not items:
+        grid.add_column()
+        return grid
+
+    if two_column and len(items) > 1:
+        # Two logical columns: [label, gap, value, wide gap, label, gap, value]
+        grid.add_column(justify="right", style="dim", no_wrap=True)
+        grid.add_column(width=2)
+        grid.add_column(justify="left")
+        grid.add_column(width=6)
+        grid.add_column(justify="right", style="dim", no_wrap=True)
+        grid.add_column(width=2)
+        grid.add_column(justify="left")
+        # Split items into left/right columns (roughly balanced).
+        mid = (len(items) + 1) // 2
+        left = items[:mid]
+        right = items[mid:]
+        for i in range(mid):
+            lk, lv = left[i]
+            if i < len(right):
+                rk, rv = right[i]
+                grid.add_row(lk, "", lv, "", rk, "", rv)
+            else:
+                grid.add_row(lk, "", lv, "", "", "", "")
+    else:
+        grid.add_column(justify="right", style="dim", no_wrap=True)
+        grid.add_column(width=2)
+        grid.add_column(justify="left")
+        for k, v in items:
+            grid.add_row(k, "", v)
+    return grid
+
+
+def title_card(
+    console: Console,
+    *,
+    title: str,
+    subtitle: str,
+    metadata: list[tuple[str, str]],
+) -> Align | Panel:
+    """Build the cover card renderable."""
+    width = console.size.width or 80
+    two_column = width >= 100
+
+    inner = Table.grid(expand=True)
+    inner.add_column(justify="center")
+    inner.add_row(f"[bold]{title}[/bold]")
+    inner.add_row(f"[dim italic]{subtitle}[/dim italic]")
+    if metadata:
+        inner.add_row("")
+        grid = metadata_grid(metadata, two_column=two_column)
+        inner.add_row(Align.center(grid))
+
+    panel_width: int | None = None
+    if width >= 120:
+        panel_width = 120
+    panel = Panel(
+        inner,
+        box=_box.ROUNDED,
+        border_style="cyan",
+        padding=(1, 4),
+        expand=(panel_width is None),
+        width=panel_width,
+    )
+    if panel_width is not None:
+        return Align.center(panel)
+    return panel
+
+
+def report_table(*, zebra: bool = False) -> Table:
+    """Base Table configured with the report's paper feel."""
+    row_styles = ["", "on grey11"] if zebra else None
+    return Table(
+        box=_box.HORIZONTALS,
+        show_header=True,
+        header_style="bold",
+        padding=(0, 2),
+        pad_edge=False,
+        row_styles=row_styles,
+        show_edge=False,
+    )