pretensor 0.1.0__py3-none-any.whl

pretensor/mcp/tools/validate_sql.py

@@ -0,0 +1,150 @@
+"""MCP ``validate_sql`` tool payload — sqlglot + graph-backed SQL validation."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+from pretensor.mcp.tool_registry import McpTool
+from pretensor.validation.query_validator import QueryValidator
+
+from ..service_registry import (
+    load_registry,
+    open_store_for_entry,
+    release_store,
+    resolve_registry_entry,
+)
+
+__all__ = ["create_tool", "validate_sql_payload"]
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_DIALECT = "postgres"
+
+
+def validate_sql_payload(
+    graph_dir: Path,
+    *,
+    sql: str,
+    database: str,
+    dialect: str = _DEFAULT_DIALECT,
+) -> dict[str, Any]:
+    """Build JSON-serializable payload for ``validate_sql``.
+
+    Args:
+        graph_dir: Root directory of the Pretensor graph store.
+        sql: SQL statement to validate.
+        database: ``connection_name`` or logical ``database`` key.
+        dialect: sqlglot dialect (default ``postgres``).
+    """
+    if not sql.strip():
+        return {"error": "Missing or empty `sql`"}
+    if not database.strip():
+        return {"error": "Missing or empty `database`"}
+
+    reg = load_registry(graph_dir)
+    entry = resolve_registry_entry(reg, database)
+    if entry is None:
+        return {"error": f"No registry entry matches database {database!r}"}
+
+    try:
+        store = open_store_for_entry(entry)
+    except Exception:
+        # Keep the on-disk graph path and engine error server-side; the client
+        # sees a stable, connection-named message only.
+        logger.exception(
+            "validate_sql: failed to open graph for connection %r",
+            entry.connection_name,
+        )
+        return {
+            "error": (
+                f"Could not open the graph for database {database!r}. "
+                "The graph may be missing or corrupt — re-run indexing."
+            )
+        }
+    try:
+        validator = QueryValidator(
+            store,
+            connection_name=entry.connection_name,
+            database_key=entry.database,
+            dialect=dialect,
+        )
+        result = validator.validate(sql)
+    finally:
+        release_store(store)
+
+    return {
+        "valid": result.valid,
+        "dialect": dialect,
+        "syntax_errors": result.syntax_errors,
+        "missing_tables": result.missing_tables,
+        "missing_columns": result.missing_columns,
+        "invalid_joins": [
+            {
+                "message": j.message,
+                "left_table": j.left_table,
+                "right_table": j.right_table,
+            }
+            for j in result.invalid_joins
+        ],
+        "suggestions": result.suggestions,
+    }
+
+
+def create_tool(graph_dir: Path) -> McpTool:
+    from ._timed import timed_tool
+
+    async def _handle(args: dict) -> dict:
+        sql_s = str(args.get("sql", ""))
+        db_t = str(args.get("database", "")).strip()
+        dialect_raw = args.get("dialect")
+        dialect_s = str(dialect_raw).strip() if dialect_raw is not None else "postgres"
+        if not dialect_s:
+            dialect_s = "postgres"
+        if not sql_s.strip():
+            return {"error": "Missing `sql`"}
+        if not db_t:
+            return {"error": "Missing `database`"}
+        with timed_tool("validate_sql", graph_dir, database=db_t, dialect=dialect_s):
+            return validate_sql_payload(
+                graph_dir, sql=sql_s, database=db_t, dialect=dialect_s
+            )
+
+    return McpTool(
+        name="validate_sql",
+        description=(
+            "Validate a SQL statement against the indexed graph before executing it. "
+            "Catches: unknown tables, unknown columns, invalid joins (pairs not in "
+            "FK_REFERENCES or INFERRED_JOIN edges), sqlglot parse errors. Returns "
+            "fuzzy suggestions for misspelled identifiers. Read-only.\n\n"
+            "Examples:\n"
+            "  sql='SELECT * FROM custmer' -> missing_tables=['custmer'], "
+            "suggestions=['customer']\n"
+            "  sql='SELECT c.foo FROM customer c' -> "
+            "missing_columns=[{'table':'customer','column':'foo'}]\n"
+            "  sql='SELECT * FROM orders o JOIN customer c ON o.x = c.y' -> "
+            "invalid_joins=[...] when no FK or inferred edge exists between columns"
+        ),
+        input_schema={
+            "type": "object",
+            "properties": {
+                "sql": {
+                    "type": "string",
+                    "description": "SQL statement to validate",
+                },
+                "database": {
+                    "type": "string",
+                    "description": "Connection name or logical database",
+                },
+                "dialect": {
+                    "type": "string",
+                    "default": "postgres",
+                    "description": "sqlglot dialect (default 'postgres')",
+                },
+            },
+            "required": ["sql", "database"],
+            "additionalProperties": False,
+        },
+        handler=_handle,
+    )

pretensor/observability.py

@@ -0,0 +1,203 @@
+"""Process-wide logging configuration and timing helpers."""
+
+from __future__ import annotations
+
+import logging
+import time
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Iterator, Literal, TypeVar
+
+__all__ = [
+    "LogFormat",
+    "LogLevel",
+    "configure_logging",
+    "log_timed_operation",
+    "run_timed_async",
+]
+
+LogLevel = Literal["debug", "info", "warning", "error"]
+LogFormat = Literal["text", "json"]
+
+_LOG_LEVELS: dict[LogLevel, int] = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+}
+_PRETENSOR_HANDLER_ATTR = "_pretensor_logging_handler"
+_T = TypeVar("_T")
+
+
+class _JsonFormatter(logging.Formatter):
+    """Render log records as one JSON object per line."""
+
+    def format(self, record: logging.LogRecord) -> str:
+        payload: dict[str, Any] = {
+            "timestamp": datetime.fromtimestamp(
+                record.created, tz=timezone.utc
+            ).isoformat(),
+            "level": record.levelname.lower(),
+            "logger": record.name,
+            "event": getattr(record, "event", None),
+            "status": getattr(record, "status", None),
+            "message": record.getMessage(),
+        }
+        if hasattr(record, "duration_ms"):
+            payload["duration_ms"] = getattr(record, "duration_ms")
+        for key, value in record.__dict__.items():
+            if key.startswith("_"):
+                continue
+            if key in payload:
+                continue
+            if key in ("args", "msg", "exc_info", "exc_text", "stack_info"):
+                continue
+            if key in (
+                "name",
+                "levelname",
+                "levelno",
+                "pathname",
+                "filename",
+                "module",
+                "lineno",
+                "funcName",
+                "created",
+                "msecs",
+                "relativeCreated",
+                "thread",
+                "threadName",
+                "processName",
+                "process",
+                "taskName",
+                "asctime",
+            ):
+                continue
+            payload[key] = value
+        if record.exc_info:
+            payload["exception"] = self.formatException(record.exc_info)
+        return _json_dumps(payload)
+
+
+def _json_dumps(payload: dict[str, Any]) -> str:
+    import json
+
+    return json.dumps(payload, ensure_ascii=False, default=str)
+
+
+def configure_logging(
+    *,
+    level: LogLevel = "info",
+    log_format: LogFormat = "text",
+    log_file: Path | None = None,
+) -> None:
+    """Configure root logging handlers for CLI and MCP code paths."""
+    root = logging.getLogger()
+    target_level = _LOG_LEVELS[level]
+    root.setLevel(target_level)
+    logging.getLogger("pretensor").setLevel(logging.NOTSET)
+    _remove_pretensor_handlers(root)
+
+    formatter: logging.Formatter
+    if log_format == "json":
+        formatter = _JsonFormatter()
+    else:
+        formatter = logging.Formatter(
+            "%(asctime)s %(levelname)s %(name)s %(message)s",
+            datefmt="%Y-%m-%dT%H:%M:%S%z",
+        )
+
+    stream_handler = logging.StreamHandler()
+    stream_handler.setFormatter(formatter)
+    stream_handler.addFilter(_build_handler_filter(target_level))
+    setattr(stream_handler, _PRETENSOR_HANDLER_ATTR, True)
+    root.addHandler(stream_handler)
+
+    if log_file is not None:
+        log_file.parent.mkdir(parents=True, exist_ok=True)
+        file_handler = logging.FileHandler(log_file, encoding="utf-8")
+        file_handler.setFormatter(formatter)
+        file_handler.addFilter(_build_handler_filter(target_level))
+        setattr(file_handler, _PRETENSOR_HANDLER_ATTR, True)
+        root.addHandler(file_handler)
+
+
+def _remove_pretensor_handlers(root: logging.Logger) -> None:
+    for handler in list(root.handlers):
+        if not getattr(handler, _PRETENSOR_HANDLER_ATTR, False):
+            continue
+        root.removeHandler(handler)
+        handler.close()
+
+
+def _build_handler_filter(min_level: int) -> Callable[[logging.LogRecord], bool]:
+    def _handler_filter(record: logging.LogRecord) -> bool:
+        if record.name.startswith("pretensor"):
+            if record.levelno < min_level:
+                return False
+            if (
+                min_level == logging.INFO
+                and record.levelno == logging.INFO
+                and not hasattr(record, "event")
+            ):
+                return False
+            return True
+        return record.levelno >= logging.WARNING
+
+    return _handler_filter
+
+
+@contextmanager
+def log_timed_operation(
+    logger: logging.Logger,
+    *,
+    event: str,
+    level: int = logging.INFO,
+    **fields: Any,
+) -> Iterator[None]:
+    """Log operation latency and status with structured fields."""
+    started = time.perf_counter()
+    try:
+        yield
+    except Exception:
+        duration_ms = (time.perf_counter() - started) * 1000
+        logger.log(
+            level,
+            "%s failed in %.2fms",
+            event,
+            duration_ms,
+            extra={
+                "event": event,
+                "status": "error",
+                "duration_ms": duration_ms,
+                **fields,
+            },
+            exc_info=True,
+        )
+        raise
+    duration_ms = (time.perf_counter() - started) * 1000
+    logger.log(
+        level,
+        "%s completed in %.2fms",
+        event,
+        duration_ms,
+        extra={
+            "event": event,
+            "status": "ok",
+            "duration_ms": duration_ms,
+            **fields,
+        },
+    )
+
+
+async def run_timed_async(
+    logger: logging.Logger,
+    *,
+    event: str,
+    callback: Callable[[], Awaitable[_T]],
+    level: int = logging.INFO,
+    **fields: Any,
+) -> _T:
+    """Await ``callback`` and emit a structured timing log."""
+    with log_timed_operation(logger, event=event, level=level, **fields):
+        return await callback()

pretensor/quickstart/README.md

@@ -0,0 +1,29 @@
+# Pretensor quickstart database
+
+A throwaway Postgres instance pre-loaded with the Pagila sample schema (DVD
+rental store), sized for `pretensor quickstart`. Listens on **localhost:55432**.
+
+## Use it
+
+The recommended path is `pretensor quickstart`, which starts the container,
+indexes the DB into `./.pretensor/`, and prints an MCP config snippet.
+
+To run the container directly from a clone of the repo:
+
+```bash
+docker compose -f src/pretensor/quickstart/docker-compose.yml up -d
+# DSN: postgresql://postgres:postgres@localhost:55432/pagila
+docker compose -f src/pretensor/quickstart/docker-compose.yml down -v
+```
+
+## Contents
+
+- 8 tables across `public` and `staff` schemas
+- ~35 seed rows total — enough for `query`, `context`, and `traverse` to
+  return non-empty results
+- FK chain: `payment → rental → inventory → film` plus `rental → customer`
+
+The schema and seed live next to this README as `pagila_ddl.sql` and
+`pagila_data.sql` — a curated subset of upstream Pagila. They ship inside
+the wheel so `pretensor quickstart` works after a plain `pip install`
+without needing a clone of the repo.

pretensor/quickstart/__init__.py

@@ -0,0 +1,6 @@
+"""Quickstart assets shipped with the wheel.
+
+Holds the docker-compose file and pagila SQL fixtures consumed by
+``pretensor quickstart``. Resolved at runtime via ``importlib.resources`` so
+they work the same in editable installs and in pip-installed wheels.
+"""

pretensor/quickstart/docker-compose.yml

@@ -0,0 +1,18 @@
+services:
+  pretensor-quickstart-db:
+    image: postgres:16-alpine
+    container_name: pretensor-quickstart-db
+    environment:
+      POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: postgres
+      POSTGRES_DB: pagila
+    ports:
+      - "55432:5432"
+    volumes:
+      - ./pagila_ddl.sql:/docker-entrypoint-initdb.d/01_pagila_ddl.sql:ro
+      - ./pagila_data.sql:/docker-entrypoint-initdb.d/02_pagila_data.sql:ro
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U postgres -d pagila"]
+      interval: 2s
+      timeout: 3s
+      retries: 30

pretensor/quickstart/pagila_data.sql

@@ -0,0 +1,63 @@
+-- Minimal Pagila seed for the hello-world quickstart.
+-- Curated subset: enough rows for `query` / `context` / `traverse` to return
+-- non-empty results, while staying small enough to commit to the repo.
+
+INSERT INTO public.film (title, description, release_year, rental_rate, length, rating) VALUES
+    ('Academy Dinosaur', 'A Boring Saga of a Dinosaur', 2006, 0.99, 86, 'PG'),
+    ('Ace Goldfinger',   'An Astounding Epistle in Hong Kong', 2006, 4.99, 48, 'G'),
+    ('Adaptation Holes', 'A Astounding Reflection of a Lumberjack', 2006, 2.99, 50, 'NC-17'),
+    ('Affair Prejudice', 'A Fanciful Documentary of a Frisbee', 2006, 2.99, 117, 'G'),
+    ('African Egg',      'A Fast-Paced Documentary of a Pastry Chef', 2006, 2.99, 130, 'G');
+
+INSERT INTO public.actor (first_name, last_name, birth_date) VALUES
+    ('Penelope', 'Guiness', '1970-04-12'),
+    ('Nick',     'Wahlberg', '1971-01-30'),
+    ('Ed',       'Chase',    '1955-07-08'),
+    ('Jennifer', 'Davis',    '1980-09-15'),
+    ('Johnny',   'Lollobrigida', '1948-11-22');
+
+INSERT INTO public.film_actor (actor_id, film_id, role) VALUES
+    (1, 1, 'lead'),
+    (2, 1, 'support'),
+    (3, 2, 'lead'),
+    (4, 3, 'lead'),
+    (5, 4, 'support');
+
+INSERT INTO public.inventory (film_id, store_id, condition, acquired_at) VALUES
+    (1, 1, 'good', '2024-01-15'),
+    (1, 1, 'good', '2024-01-15'),
+    (2, 1, 'fair', '2024-02-01'),
+    (3, 2, 'good', '2024-02-10'),
+    (4, 1, 'good', '2024-03-05');
+
+INSERT INTO public.customer (first_name, last_name, email) VALUES
+    ('Mary',    'Smith',   'mary.smith@example.com'),
+    ('Patricia','Johnson', 'patricia.j@example.com'),
+    ('Linda',   'Williams','linda.w@example.com');
+
+INSERT INTO public.rental (inventory_id, customer_id, rental_date, return_date) VALUES
+    (1, 1, '2024-05-24 22:53:30+00', '2024-05-26 22:04:30+00'),
+    (2, 2, '2024-05-24 22:54:33+00', '2024-05-28 19:40:33+00'),
+    (3, 1, '2024-05-25 11:30:37+00', '2024-06-01 14:25:37+00'),
+    (4, 3, '2024-05-26 09:15:00+00', '2024-05-30 12:00:00+00'),
+    (5, 2, '2024-05-27 13:42:11+00', NULL);
+
+INSERT INTO public.payment (rental_id, customer_id, amount, method) VALUES
+    (1, 1, 2.99, 'cash'),
+    (2, 2, 4.99, 'card'),
+    (3, 1, 0.99, 'card'),
+    (4, 3, 2.99, 'cash'),
+    (5, 2, 4.99, 'card');
+
+INSERT INTO staff.staff (first_name, last_name, email, store_id, username) VALUES
+    ('Mike', 'Hillyer',  'mike@example.com',  1, 'mike'),
+    ('Jon',  'Stephens', 'jon@example.com',   2, 'jon');
+
+-- Reset sequences so subsequent inserts (if any) start past seeded ids.
+SELECT setval(pg_get_serial_sequence('public.film',      'film_id'),      (SELECT MAX(film_id)      FROM public.film));
+SELECT setval(pg_get_serial_sequence('public.actor',     'actor_id'),     (SELECT MAX(actor_id)     FROM public.actor));
+SELECT setval(pg_get_serial_sequence('public.inventory', 'inventory_id'), (SELECT MAX(inventory_id) FROM public.inventory));
+SELECT setval(pg_get_serial_sequence('public.customer',  'customer_id'),  (SELECT MAX(customer_id)  FROM public.customer));
+SELECT setval(pg_get_serial_sequence('public.rental',    'rental_id'),    (SELECT MAX(rental_id)    FROM public.rental));
+SELECT setval(pg_get_serial_sequence('public.payment',   'payment_id'),   (SELECT MAX(payment_id)   FROM public.payment));
+SELECT setval(pg_get_serial_sequence('staff.staff',      'staff_id'),     (SELECT MAX(staff_id)     FROM staff.staff));

pretensor/quickstart/pagila_ddl.sql

@@ -0,0 +1,92 @@
+-- Minimal Pagila-derived schema for E2E testing
+-- 2 schemas: public, staff
+-- 7 tables in public: film, actor, film_actor, inventory, customer, rental, payment
+-- 1 table in staff: staff
+-- FK chains: inventory→film, rental→inventory, rental→customer, payment→rental
+
+CREATE SCHEMA IF NOT EXISTS staff;
+
+-- ============================================================
+-- public schema
+-- ============================================================
+
+CREATE TABLE public.film (
+    film_id       SERIAL PRIMARY KEY,
+    title         VARCHAR(255)   NOT NULL,
+    description   TEXT,
+    release_year  SMALLINT,
+    rental_rate   NUMERIC(4,2)   NOT NULL DEFAULT 4.99,
+    length        SMALLINT,
+    rating        VARCHAR(10)    DEFAULT 'G',
+    last_update   TIMESTAMPTZ    NOT NULL DEFAULT NOW()
+);
+
+CREATE TABLE public.actor (
+    actor_id    SERIAL PRIMARY KEY,
+    first_name  VARCHAR(45)  NOT NULL,
+    last_name   VARCHAR(45)  NOT NULL,
+    birth_date  DATE,
+    active      BOOLEAN      NOT NULL DEFAULT TRUE,
+    last_update TIMESTAMPTZ  NOT NULL DEFAULT NOW()
+);
+
+CREATE TABLE public.film_actor (
+    actor_id    INT          NOT NULL REFERENCES public.actor(actor_id),
+    film_id     INT          NOT NULL REFERENCES public.film(film_id),
+    role        VARCHAR(50),
+    created_at  TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
+    PRIMARY KEY (actor_id, film_id)
+);
+
+CREATE TABLE public.inventory (
+    inventory_id  SERIAL PRIMARY KEY,
+    film_id       INT          NOT NULL REFERENCES public.film(film_id),
+    store_id      SMALLINT     NOT NULL DEFAULT 1,
+    condition     VARCHAR(20)  NOT NULL DEFAULT 'good',
+    acquired_at   DATE,
+    last_update   TIMESTAMPTZ  NOT NULL DEFAULT NOW()
+);
+
+CREATE TABLE public.customer (
+    customer_id  SERIAL PRIMARY KEY,
+    first_name   VARCHAR(45)   NOT NULL,
+    last_name    VARCHAR(45)   NOT NULL,
+    email        VARCHAR(100),
+    active       BOOLEAN       NOT NULL DEFAULT TRUE,
+    created_at   TIMESTAMPTZ   NOT NULL DEFAULT NOW(),
+    last_update  TIMESTAMPTZ   NOT NULL DEFAULT NOW()
+);
+
+CREATE TABLE public.rental (
+    rental_id     SERIAL PRIMARY KEY,
+    inventory_id  INT          NOT NULL REFERENCES public.inventory(inventory_id),
+    customer_id   INT          NOT NULL REFERENCES public.customer(customer_id),
+    rental_date   TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
+    return_date   TIMESTAMPTZ,
+    staff_id      SMALLINT     NOT NULL DEFAULT 1,
+    last_update   TIMESTAMPTZ  NOT NULL DEFAULT NOW()
+);
+
+CREATE TABLE public.payment (
+    payment_id   SERIAL PRIMARY KEY,
+    rental_id    INT           NOT NULL REFERENCES public.rental(rental_id),
+    customer_id  INT           NOT NULL REFERENCES public.customer(customer_id),
+    amount       NUMERIC(5,2)  NOT NULL,
+    payment_date TIMESTAMPTZ   NOT NULL DEFAULT NOW(),
+    method       VARCHAR(20)   NOT NULL DEFAULT 'cash'
+);
+
+-- ============================================================
+-- staff schema (multi-schema coverage)
+-- ============================================================
+
+CREATE TABLE staff.staff (
+    staff_id    SERIAL PRIMARY KEY,
+    first_name  VARCHAR(45)   NOT NULL,
+    last_name   VARCHAR(45)   NOT NULL,
+    email       VARCHAR(100),
+    store_id    SMALLINT      NOT NULL DEFAULT 1,
+    active      BOOLEAN       NOT NULL DEFAULT TRUE,
+    username    VARCHAR(50)   NOT NULL,
+    last_update TIMESTAMPTZ   NOT NULL DEFAULT NOW()
+);

pretensor/search/__init__.py

@@ -0,0 +1,6 @@
+"""Keyword search over graph metadata (FTS5)."""
+
+from pretensor.search.base import BaseSearchIndex, SearchResult
+from pretensor.search.index import KeywordSearchIndex, SearchIndex
+
+__all__ = ["BaseSearchIndex", "KeywordSearchIndex", "SearchIndex", "SearchResult"]

pretensor/search/base.py

@@ -0,0 +1,80 @@
+"""Abstract base for search index implementations."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+from pretensor.core.registry import GraphRegistry
+
+__all__ = ["BaseSearchIndex", "SearchResult"]
+
+
+@dataclass(frozen=True, slots=True)
+class SearchResult:
+    """One ranked hit from a search index."""
+
+    node_type: str
+    name: str
+    database_name: str
+    connection_name: str
+    description: str
+    score: float
+
+
+class BaseSearchIndex(ABC):
+    """Abstract search index over graph metadata.
+
+    Concrete implementations may use keyword (BM25), vector (embeddings),
+    or hybrid strategies.  OSS ships ``KeywordSearchIndex``; plugins can register
+    a ``HybridSearchIndex`` subclass.
+    """
+
+    @abstractmethod
+    def search(
+        self,
+        q: str,
+        *,
+        db: str | None = None,
+        limit: int = 10,
+    ) -> list[SearchResult]:
+        """Search indexed metadata.
+
+        Args:
+            q: Free-text query string.
+            db: Optional connection or logical database name filter.
+            limit: Maximum number of results to return.
+
+        Returns:
+            Ranked list of search results.
+        """
+
+    @abstractmethod
+    def similar(
+        self,
+        name: str,
+        *,
+        db: str | None = None,
+        limit: int = 10,
+    ) -> list[SearchResult]:
+        """Return nodes related to the named node.
+
+        OSS implementations use graph proximity (FK / join edges).  Other
+        implementations may additionally use vector similarity.
+
+        Args:
+            name: Fully-qualified node name, e.g. ``schema.table``.
+            db: Optional connection or logical database name filter.
+            limit: Maximum number of results to return.
+
+        Returns:
+            Related nodes ranked by proximity or similarity.
+        """
+
+    @abstractmethod
+    def index_graph(self, registry: GraphRegistry) -> None:
+        """Build or rebuild the search index from all graphs in *registry*.
+
+        Args:
+            registry: Loaded registry describing available graph files.
+        """