sqlnow-mcp 0.1.0__py3-none-any.whl

sqlnow_mcp/__init__.py

@@ -0,0 +1,3 @@
+from sqlnow_mcp.cli import main
+
+__all__ = ["main"]

sqlnow_mcp/cli.py

@@ -0,0 +1,545 @@
+from __future__ import annotations
+
+import csv
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+import click
+
+from sqlnow_mcp.config import LocalConfig, PublishConfig, create_session, default_allow_paths
+from sqlnow_mcp.db import DuckDBSession, DuckDBSessionError
+from sqlnow_mcp.metadata import build_metadata_template, generate_stats, write_stats
+from sqlnow_mcp.server import run_local_server, run_publish_server
+
+
+def _session_from_ctx(ctx: click.Context) -> DuckDBSession:
+    return ctx.ensure_object(dict)["session"]
+
+
+def _require_database(ctx: click.Context) -> str:
+    db = ctx.ensure_object(dict).get("database")
+    if not db:
+        raise click.ClickException("--database / -d is required for this command")
+    return db
+
+
+def _open_database(session: DuckDBSession, db_name: str) -> dict[str, Any]:
+    try:
+        return session.use_database(db_name)
+    except DuckDBSessionError as exc:
+        raise click.ClickException(str(exc)) from exc
+
+
+def _make_session(
+    data_dir: Path,
+    allow_paths: tuple[Path, ...],
+    allow_external: bool,
+) -> DuckDBSession:
+    try:
+        return create_session(
+            LocalConfig(
+                data_dir=data_dir,
+                allow_paths=allow_paths,
+                allow_external=allow_external,
+            )
+        )
+    except DuckDBSessionError as exc:
+        raise click.ClickException(str(exc)) from exc
+
+
+@click.group(invoke_without_command=True)
+@click.option(
+    "--mode",
+    type=click.Choice(["local", "publish"], case_sensitive=False),
+    help="Run as MCP server: 'local' (read/write, data directory) or 'publish' (read-only, metadata).",
+)
+@click.option(
+    "--data-dir",
+    type=click.Path(file_okay=False, path_type=Path),
+    default=".",
+    show_default=True,
+    help="Data directory for local mode and dev subcommands (not used by publish startup).",
+)
+@click.option(
+    "--allow-path",
+    "allow_paths",
+    multiple=True,
+    type=click.Path(path_type=Path),
+    help="Extra directory allowed for attach-file in local mode (repeatable; default: home).",
+)
+@click.option(
+    "--allow-external/--no-allow-external",
+    default=True,
+    show_default=True,
+    help="Allow attach-database in local mode (Postgres, SQLite, MySQL).",
+)
+@click.option(
+    "--host",
+    default="127.0.0.1",
+    show_default=True,
+    help="Bind address for HTTP transports (local and publish).",
+)
+@click.option(
+    "--port",
+    default=8000,
+    show_default=True,
+    type=int,
+    help="Port for HTTP transports (local and publish).",
+)
+@click.option(
+    "--transport",
+    type=click.Choice(
+        ["stdio", "http", "streamable-http", "sse"],
+        case_sensitive=False,
+    ),
+    default="stdio",
+    show_default=True,
+    help="MCP transport for local and publish (stdio spawns process; others listen on --host/--port).",
+)
+@click.option(
+    "--database",
+    "-d",
+    help="Active database name (without .db extension).",
+)
+@click.option(
+    "--db",
+    "publish_db",
+    type=click.Path(exists=True, dir_okay=False, path_type=Path),
+    help="DuckDB file for publish mode.",
+)
+@click.option(
+    "--metadata",
+    "metadata_path",
+    type=click.Path(exists=True, dir_okay=False, path_type=Path),
+    help="metadata.yaml for publish mode.",
+)
+@click.option(
+    "--stats",
+    "stats_path",
+    type=click.Path(exists=True, dir_okay=False, path_type=Path),
+    help="Precomputed stats.json for publish mode.",
+)
+@click.option(
+    "--strict-metadata/--no-strict-metadata",
+    default=False,
+    show_default=True,
+    help="Fail publish startup on metadata/DB mismatches.",
+)
+@click.option(
+    "--query-timeout",
+    type=float,
+    default=10.0,
+    show_default=True,
+    help="Per-query timeout in seconds (publish mode).",
+)
+@click.option(
+    "--memory-limit",
+    default="50%",
+    show_default=True,
+    help="DuckDB memory_limit for publish mode (e.g. 50%%, 2GiB).",
+)
+@click.option(
+    "--threads",
+    type=int,
+    default=1,
+    show_default=True,
+    help="DuckDB worker threads for publish mode.",
+)
+@click.option(
+    "--max-temp-directory-size",
+    default=None,
+    help="Cap DuckDB temp spill size for publish mode (e.g. 10GiB).",
+)
+@click.pass_context
+def cli(
+    ctx: click.Context,
+    mode: str | None,
+    data_dir: Path,
+    allow_paths: tuple[Path, ...],
+    allow_external: bool,
+    host: str,
+    port: int,
+    transport: str,
+    database: str | None,
+    publish_db: Path | None,
+    metadata_path: Path | None,
+    stats_path: Path | None,
+    strict_metadata: bool,
+    query_timeout: float,
+    memory_limit: str,
+    threads: int,
+    max_temp_directory_size: str | None,
+) -> None:
+    """sqlnow-mcp — DuckDB MCP server and dev CLI."""
+    effective_allow_paths = default_allow_paths(allow_paths)
+
+    if ctx.invoked_subcommand is None:
+        if mode is None:
+            click.echo(ctx.get_help())
+            return
+        mode_lower = mode.lower()
+        if mode_lower == "publish":
+            if publish_db is None:
+                raise click.ClickException("Publish mode requires --db")
+            if metadata_path is None:
+                raise click.ClickException("Publish mode requires --metadata")
+            cfg = PublishConfig(
+                db_path=publish_db,
+                metadata_path=metadata_path,
+                stats_path=stats_path,
+                host=host,
+                port=port,
+                transport=transport.lower(),  # type: ignore[arg-type]
+                strict_metadata=strict_metadata,
+                query_timeout_sec=query_timeout,
+                memory_limit=memory_limit,
+                threads=threads,
+                max_temp_directory_size=max_temp_directory_size,
+            )
+            try:
+                run_publish_server(cfg)
+            except DuckDBSessionError as exc:
+                raise click.ClickException(str(exc)) from exc
+            return
+        if mode_lower == "local":
+            cfg = LocalConfig(
+                data_dir=data_dir,
+                allow_paths=effective_allow_paths,
+                allow_external=allow_external,
+                host=host,
+                port=port,
+                transport=transport.lower(),  # type: ignore[arg-type]
+            )
+            run_local_server(cfg)
+            return
+        raise click.ClickException(f"Unknown mode: {mode}")
+
+    ctx.ensure_object(dict)["session"] = _make_session(
+        data_dir, effective_allow_paths, allow_external
+    )
+    ctx.ensure_object(dict)["database"] = database
+
+
+@cli.command("create-database")
+@click.argument("name")
+@click.pass_context
+def create_database_cmd(ctx: click.Context, name: str) -> None:
+    """Create a new empty DuckDB database in data-dir."""
+    session = _session_from_ctx(ctx)
+    try:
+        result = session.create_database(name)
+    except DuckDBSessionError as exc:
+        raise click.ClickException(str(exc)) from exc
+    click.echo(json.dumps(result, indent=2))
+
+
+@cli.command("list-databases")
+@click.pass_context
+def list_databases(ctx: click.Context) -> None:
+    """List .db files in data-dir."""
+    session = _session_from_ctx(ctx)
+    dbs = session.list_databases()
+    if not dbs:
+        click.echo("No databases found.")
+        return
+    for db in dbs:
+        click.echo(f"{db['name']}\t{db['size_mb']} MB\t{db['last_modified']}")
+
+
+@cli.command("use-database")
+@click.argument("name")
+@click.pass_context
+def use_database_cmd(ctx: click.Context, name: str) -> None:
+    """Open a database and print available tables."""
+    session = _session_from_ctx(ctx)
+    info = _open_database(session, name)
+    click.echo(json.dumps(info, indent=2))
+
+
+@cli.command("current")
+@click.pass_context
+def current(ctx: click.Context) -> None:
+    """Show session state after opening --database."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    click.echo(json.dumps(session.current_database(), indent=2))
+
+
+@cli.command("attach-file")
+@click.argument("path", type=click.Path(path_type=Path))
+@click.option("--name", help="Table/view name (default: file stem).")
+@click.option(
+    "--mode",
+    "attach_mode",
+    type=click.Choice(["view", "load"], case_sensitive=False),
+    default="view",
+    show_default=True,
+)
+@click.pass_context
+def attach_file_cmd(
+    ctx: click.Context,
+    path: Path,
+    name: str | None,
+    attach_mode: str,
+) -> None:
+    """Attach a CSV or Parquet file as a view or table."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    try:
+        result = session.attach_file(
+            str(path), name=name, mode=attach_mode.lower()  # type: ignore[arg-type]
+        )
+    except DuckDBSessionError as exc:
+        raise click.ClickException(str(exc)) from exc
+    click.echo(json.dumps(result, indent=2))
+
+
+@cli.command("attach-database")
+@click.argument("connection_string")
+@click.option("--name", required=True, help="Catalog name for ATTACH.")
+@click.option(
+    "--table",
+    "tables",
+    multiple=True,
+    help="Limit exposed tables from the external database (repeatable).",
+)
+@click.pass_context
+def attach_database_cmd(
+    ctx: click.Context,
+    connection_string: str,
+    name: str,
+    tables: tuple[str, ...],
+) -> None:
+    """Attach an external database (Postgres, SQLite, MySQL)."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    try:
+        result = session.attach_database(
+            connection_string,
+            name,
+            tables=list(tables) or None,
+        )
+    except DuckDBSessionError as exc:
+        raise click.ClickException(str(exc)) from exc
+    click.echo(json.dumps(result, indent=2))
+
+
+@cli.command("detach")
+@click.argument("name")
+@click.pass_context
+def detach(ctx: click.Context, name: str) -> None:
+    """Detach an external database or drop an attached file view/table."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    try:
+        result = session.detach_source(name)
+    except DuckDBSessionError as exc:
+        raise click.ClickException(str(exc)) from exc
+    click.echo(json.dumps(result, indent=2))
+
+
+@cli.command("list-tables")
+@click.pass_context
+def list_tables_cmd(ctx: click.Context) -> None:
+    """List tables and columns in the active database."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    for table in session.list_tables():
+        cols = ", ".join(f"{c['name']}:{c['type']}" for c in table["columns"])
+        click.echo(f"{table['name']}\t{cols}")
+
+
+@cli.command("describe")
+@click.argument("table_name")
+@click.pass_context
+def describe(ctx: click.Context, table_name: str) -> None:
+    """Describe a table (JSON)."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    try:
+        result = session.describe_table(table_name)
+    except DuckDBSessionError as exc:
+        raise click.ClickException(str(exc)) from exc
+    click.echo(json.dumps(result, indent=2))
+
+
+@cli.command("sample")
+@click.argument("table_name")
+@click.option("--limit", "-n", default=10, show_default=True)
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(path_type=Path),
+    help="Write rows as CSV (default: JSON on stdout).",
+)
+@click.pass_context
+def sample(
+    ctx: click.Context,
+    table_name: str,
+    limit: int,
+    output: Path | None,
+) -> None:
+    """Sample rows from a table."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    result = session.sample_table(table_name, n=limit)
+    rows = result["rows"]
+    if output is None:
+        click.echo(json.dumps(result, indent=2))
+        return
+    headers = list(rows[0].keys()) if rows else []
+    with output.open("w", newline="", encoding="utf-8") as f:
+        if headers:
+            writer = csv.DictWriter(f, fieldnames=headers)
+            writer.writeheader()
+            writer.writerows(rows)
+    click.echo(f"Wrote {len(rows)} rows to {output}")
+
+
+@cli.command("query")
+@click.argument("sql")
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(path_type=Path),
+    help="Write query results as CSV (default: stdout).",
+)
+@click.option("--limit", default=500, show_default=True, help="Max rows to fetch.")
+@click.pass_context
+def query(ctx: click.Context, sql: str, output: Path | None, limit: int) -> None:
+    """Run SQL and write results as CSV to stdout or a file."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    try:
+        if DuckDBSession._is_mutating_sql(sql):
+            result = session.run_mutating_query(sql)
+            click.echo(result["message"])
+            return
+
+        result = session.run_query(sql, limit=limit)
+    except Exception as exc:
+        raise click.ClickException(str(exc)) from exc
+
+    if output is None:
+        writer = csv.writer(sys.stdout)
+        if result["headers"]:
+            writer.writerow(result["headers"])
+        writer.writerows(result["rows"])
+        return
+
+    with output.open("w", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        if result["headers"]:
+            writer.writerow(result["headers"])
+        writer.writerows(result["rows"])
+    click.echo(f"Wrote {len(result['rows'])} rows to {output}")
+
+
+@cli.command("generate-metadata")
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(path_type=Path),
+    help="Write metadata.yaml to this path (default: stdout).",
+)
+@click.option("--title", help="Dataset title (default: derived from database name).")
+@click.option(
+    "--include-attached/--no-include-attached",
+    default=False,
+    show_default=True,
+    help="Include tables from attached databases (default: native tables only).",
+)
+@click.pass_context
+def generate_metadata_cmd(
+    ctx: click.Context,
+    output: Path | None,
+    title: str | None,
+    include_attached: bool,
+) -> None:
+    """Generate a starter metadata.yaml from the active database schema."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    try:
+        text = build_metadata_template(
+            session,
+            title=title,
+            include_attached=include_attached,
+        )
+    except (DuckDBSessionError, ValueError) as exc:
+        raise click.ClickException(str(exc)) from exc
+
+    if output is None:
+        click.echo(text, nl=False)
+        if not text.endswith("\n"):
+            click.echo()
+        return
+
+    output.write_text(text, encoding="utf-8")
+    click.echo(f"Wrote metadata template to {output}")
+
+
+@cli.command("generate-stats")
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(path_type=Path),
+    help="Write stats.json to this path (default: stdout).",
+)
+@click.option(
+    "--sample-threshold",
+    default=100_000,
+    show_default=True,
+    type=int,
+    help="Row count above which SUMMARIZE uses Bernoulli sampling.",
+)
+@click.option(
+    "--include-attached/--no-include-attached",
+    default=False,
+    show_default=True,
+    help="Profile tables from attached databases (default: native tables only).",
+)
+@click.pass_context
+def generate_stats_cmd(
+    ctx: click.Context,
+    output: Path | None,
+    sample_threshold: int,
+    include_attached: bool,
+) -> None:
+    """Generate stats.json with SUMMARIZE profiles for publish mode."""
+    session = _session_from_ctx(ctx)
+    db = _require_database(ctx)
+    _open_database(session, db)
+    try:
+        stats = generate_stats(
+            session,
+            sample_threshold=sample_threshold,
+            include_attached=include_attached,
+        )
+    except (DuckDBSessionError, ValueError) as exc:
+        raise click.ClickException(str(exc)) from exc
+
+    if output is None:
+        click.echo(json.dumps(stats.to_dict(), indent=2, ensure_ascii=False))
+        return
+
+    write_stats(output, stats)
+    click.echo(f"Wrote stats for {len(stats.tables)} table(s) to {output}")
+
+
+def main() -> None:
+    try:
+        cli(standalone_mode=True)
+    except click.ClickException as exc:
+        click.echo(str(exc), err=True)
+        sys.exit(1)

sqlnow_mcp/config.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+from sqlnow_mcp.db import DuckDBSession, DuckDBSessionError
+from sqlnow_mcp.metadata import (
+    MetadataContext,
+    load_metadata,
+    load_stats,
+    log_validation_report,
+    validate_metadata,
+)
+from sqlnow_mcp.resource_limits import resolve_memory_limit
+
+FORBIDDEN_ALLOW_ROOTS = {Path("/"), Path()}
+
+Transport = Literal["stdio", "http", "streamable-http", "sse"]
+
+
+def default_allow_paths(explicit: tuple[Path, ...]) -> tuple[Path, ...]:
+    return explicit if explicit else (Path.home(),)
+
+
+def reject_unsafe_allow_paths(paths: tuple[Path, ...]) -> None:
+    for path in paths:
+        resolved = path.resolve()
+        if resolved in FORBIDDEN_ALLOW_ROOTS:
+            raise DuckDBSessionError(
+                f"Refusing unsafe --allow-path {path!s} (would allow any file on disk)"
+            )
+
+
+@dataclass(frozen=True)
+class LocalConfig:
+    data_dir: Path
+    allow_paths: tuple[Path, ...]
+    allow_external: bool
+    host: str = "127.0.0.1"
+    port: int = 8000
+    transport: Transport = "stdio"
+
+
+@dataclass(frozen=True)
+class PublishConfig:
+    db_path: Path
+    metadata_path: Path
+    stats_path: Path | None = None
+    host: str = "127.0.0.1"
+    port: int = 8000
+    transport: Transport = "stdio"
+    strict_metadata: bool = False
+    query_timeout_sec: float = 10.0
+    memory_limit: str = "50%"
+    threads: int = 1
+    max_temp_directory_size: str | None = None
+
+
+def create_session(cfg: LocalConfig) -> DuckDBSession:
+    effective = default_allow_paths(cfg.allow_paths)
+    reject_unsafe_allow_paths(effective)
+    return DuckDBSession(
+        data_dir=cfg.data_dir,
+        allow_paths=effective,
+        allow_external=cfg.allow_external,
+    )
+
+
+def create_publish_session(cfg: PublishConfig) -> tuple[DuckDBSession, MetadataContext]:
+    db_path = cfg.db_path.resolve()
+    if not db_path.exists():
+        raise DuckDBSessionError(f"Database not found: {db_path}")
+
+    memory_limit = resolve_memory_limit(cfg.memory_limit)
+    session = DuckDBSession(
+        data_dir=db_path.parent,
+        allow_paths=(),
+        allow_external=False,
+        read_only=True,
+        query_timeout_sec=cfg.query_timeout_sec,
+    )
+    session.open_publish_database(
+        db_path,
+        memory_limit=memory_limit,
+        threads=cfg.threads,
+        max_temp_directory_size=cfg.max_temp_directory_size,
+    )
+
+    meta = load_metadata(cfg.metadata_path.resolve())
+    stats = load_stats(cfg.stats_path.resolve() if cfg.stats_path else None)
+    meta_ctx = MetadataContext(meta, stats)
+
+    report = validate_metadata(meta, session, strict=cfg.strict_metadata)
+    log_validation_report(report)
+    if cfg.strict_metadata and report.has_errors():
+        messages = "; ".join(issue.message for issue in report.errors)
+        raise DuckDBSessionError(f"Metadata validation failed: {messages}")
+
+    return session, meta_ctx