querypilot 0.1.0__py3-none-any.whl

querypilot/connectors/base.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+
+from querypilot.core.types import DatabaseSchema
+
+
+class BaseConnector(ABC):
+    def __init__(self, database_url: str, timeout_seconds: int = 10) -> None:
+        self.database_url = database_url
+        self.timeout_seconds = timeout_seconds
+
+    @abstractmethod
+    def get_schema(self) -> DatabaseSchema:
+        raise NotImplementedError
+
+    @abstractmethod
+    def execute_readonly(self, sql: str) -> tuple[list[dict], int]:
+        raise NotImplementedError
+
+    @staticmethod
+    def rows_to_dicts(keys: Sequence[str], rows: Sequence[Sequence[object]]) -> list[dict]:
+        return [dict(zip(keys, row, strict=False)) for row in rows]

querypilot/connectors/postgres.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from sqlalchemy import create_engine, inspect, text
+from sqlalchemy.engine import Engine
+
+from querypilot.connectors.base import BaseConnector
+from querypilot.core.types import ColumnSchema, DatabaseSchema, TableSchema
+
+
+class PostgresConnector(BaseConnector):
+    def __init__(self, database_url: str, timeout_seconds: int = 10) -> None:
+        super().__init__(database_url, timeout_seconds)
+        self.engine: Engine = create_engine(_normalize_postgres_url(database_url))
+
+    def get_schema(self) -> DatabaseSchema:
+        inspector = inspect(self.engine)
+        tables: list[TableSchema] = []
+        for table_name in inspector.get_table_names():
+            columns = [
+                ColumnSchema(
+                    name=column["name"],
+                    type=str(column["type"]),
+                    nullable=bool(column.get("nullable", True)),
+                )
+                for column in inspector.get_columns(table_name)
+            ]
+            tables.append(TableSchema(name=table_name, columns=columns))
+        return DatabaseSchema(dialect="postgres", tables=tables)
+
+    def execute_readonly(self, sql: str) -> tuple[list[dict], int]:
+        with self.engine.connect() as conn:
+            readonly_result = conn.execute(text("SET TRANSACTION READ ONLY"))
+            readonly_result.close()
+            timeout_result = conn.execute(
+                text(f"SET LOCAL statement_timeout = {self.timeout_seconds * 1000}")
+            )
+            timeout_result.close()
+            query_result = conn.execute(text(sql))
+            rows = [dict(row._mapping) for row in query_result.fetchall()]
+        return rows, len(rows)
+
+
+def _normalize_postgres_url(database_url: str) -> str:
+    if database_url.startswith("postgresql://"):
+        return database_url.replace("postgresql://", "postgresql+psycopg://", 1)
+    if database_url.startswith("postgres://"):
+        return database_url.replace("postgres://", "postgresql+psycopg://", 1)
+    return database_url

querypilot/connectors/sqlite.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from sqlalchemy import create_engine, inspect, text
+from sqlalchemy.engine import Engine
+
+from querypilot.connectors.base import BaseConnector
+from querypilot.core.types import ColumnSchema, DatabaseSchema, TableSchema
+
+
+class SQLiteConnector(BaseConnector):
+    def __init__(self, database_url: str, timeout_seconds: int = 10) -> None:
+        super().__init__(database_url, timeout_seconds)
+        self.engine: Engine = create_engine(database_url)
+
+    def get_schema(self) -> DatabaseSchema:
+        inspector = inspect(self.engine)
+        tables: list[TableSchema] = []
+        for table_name in inspector.get_table_names():
+            columns = [
+                ColumnSchema(
+                    name=column["name"],
+                    type=str(column["type"]),
+                    nullable=bool(column.get("nullable", True)),
+                )
+                for column in inspector.get_columns(table_name)
+            ]
+            tables.append(TableSchema(name=table_name, columns=columns))
+        return DatabaseSchema(dialect="sqlite", tables=tables)
+
+    def execute_readonly(self, sql: str) -> tuple[list[dict], int]:
+        with self.engine.connect() as conn:
+            conn.connection.set_progress_handler(
+                _timeout_handler(self.timeout_seconds),
+                1000,
+            )
+            result = conn.execute(text(sql))
+            rows = [dict(row._mapping) for row in result.fetchall()]
+        return rows, len(rows)
+
+
+def _timeout_handler(timeout_seconds: int):
+    import time
+
+    deadline = time.monotonic() + timeout_seconds
+
+    def handler() -> int:
+        return 1 if time.monotonic() > deadline else 0
+
+    return handler

querypilot/core/__init__.py

@@ -0,0 +1 @@
+"""Core QueryPilot package."""

querypilot/core/client.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import TypeAdapter
+
+from querypilot.access import AccessPolicy
+from querypilot.adapters.anthropic import anthropic_tools
+from querypilot.adapters.openai import openai_tools
+from querypilot.audit import AuditMetadata, AuditSink, InMemoryAuditSink, QueryAuditRecord
+from querypilot.connectors.base import BaseConnector
+from querypilot.connectors.postgres import PostgresConnector
+from querypilot.connectors.sqlite import SQLiteConnector
+from querypilot.core.config import QueryPilotConfig, SafetyPolicy
+from querypilot.core.types import (
+    DatabaseSchema,
+    GeneratedSQL,
+    QueryPilotAnswer,
+    QueryResult,
+    SchemaMatch,
+    ValidationResult,
+)
+from querypilot.execution.formatter import explain_result
+from querypilot.execution.runner import execute
+from querypilot.generation.sql_generator import DemoSQLGenerator, SQLGenerator
+from querypilot.schema.search import search_schema as search_schema_impl
+from querypilot.validation.validator import SQLValidator
+
+
+class QueryPilot:
+    def __init__(
+        self,
+        connector: BaseConnector,
+        config: QueryPilotConfig,
+        generator: SQLGenerator | None = None,
+        audit_sink: AuditSink | None = None,
+        audit_metadata: AuditMetadata | None = None,
+    ) -> None:
+        self.connector = connector
+        self.config = config
+        self.generator = generator or DemoSQLGenerator()
+        self.validator = SQLValidator(config)
+        self.audit_sink = audit_sink or InMemoryAuditSink()
+        self.audit_metadata = audit_metadata or AuditMetadata()
+
+    @classmethod
+    def connect(
+        cls,
+        database_url: str,
+        dialect: str = "sqlite",
+        readonly: bool = True,
+        max_rows: int = 100,
+        timeout_seconds: int = 10,
+        max_generation_attempts: int = 2,
+        allowed_tables: list[str] | None = None,
+        blocked_tables: list[str] | None = None,
+        safety_policy: SafetyPolicy | None = None,
+        access_policy: AccessPolicy | None = None,
+        generator: SQLGenerator | None = None,
+        audit_sink: AuditSink | None = None,
+        audit_metadata: AuditMetadata | None = None,
+    ) -> "QueryPilot":
+        config = QueryPilotConfig(
+            dialect=dialect,
+            readonly=readonly,
+            max_rows=max_rows,
+            timeout_seconds=timeout_seconds,
+            max_generation_attempts=max_generation_attempts,
+            allowed_tables=allowed_tables,
+            blocked_tables=blocked_tables,
+            safety_policy=safety_policy or SafetyPolicy(),
+            access_policy=access_policy or AccessPolicy(),
+        )
+        connector = _connector_for(database_url, dialect, timeout_seconds)
+        return cls(
+            connector=connector,
+            config=config,
+            generator=generator,
+            audit_sink=audit_sink,
+            audit_metadata=audit_metadata,
+        )
+
+    def get_schema(self) -> DatabaseSchema:
+        return self.connector.get_schema()
+
+    def search_schema(self, query: str) -> list[SchemaMatch]:
+        matches = search_schema_impl(self.get_schema(), query)
+        self._write_audit_record(
+            operation="schema_search",
+            question=query,
+            row_count=len(matches),
+            executed=True,
+        )
+        return matches
+
+    def generate_sql(self, question: str) -> GeneratedSQL:
+        generated = self.generator.generate(question, self.get_schema(), self.config.max_rows)
+        self._write_audit_record(
+            operation="generate_sql",
+            question=question,
+            sql=generated.sql,
+            error="; ".join(generated.errors) if generated.errors else None,
+            executed=generated.sql is not None,
+        )
+        return generated
+
+    def validate_sql(self, sql: str) -> ValidationResult:
+        validation = self.validator.validate(sql, self.get_schema())
+        record = self._write_audit_record(
+            operation="validate_sql",
+            sql=sql,
+            rewritten_sql=validation.rewritten_sql,
+                validation=validation,
+                valid=validation.valid,
+                access_policy=validation.access_policy,
+                error="; ".join(validation.errors) if validation.errors else None,
+            )
+        validation.audit_id = record.audit_id
+        return validation
+
+    def execute_sql(self, sql: str) -> QueryResult:
+        validation = self.validator.validate(sql, self.get_schema())
+        if not validation.valid or validation.rewritten_sql is None:
+            details = "; ".join(validation.errors) or "unknown validation error"
+            self._write_audit_record(
+                operation="execute_sql",
+                sql=sql,
+                rewritten_sql=validation.rewritten_sql,
+                validation=validation,
+                valid=validation.valid,
+                executed=False,
+                access_policy=validation.access_policy,
+                error=details,
+            )
+            raise ValueError(f"SQL validation failed: {details}")
+        result = execute(self.connector, validation.rewritten_sql)
+        result.rows = self._apply_masking(result.rows, validation)
+        result.access_policy = validation.access_policy
+        record = self._write_audit_record(
+            operation="execute_sql",
+            sql=sql,
+            rewritten_sql=result.sql,
+            validation=validation,
+            valid=True,
+            executed=True,
+            row_count=result.row_count,
+            execution_time_ms=result.execution_time_ms,
+            access_policy=validation.access_policy,
+        )
+        validation.audit_id = record.audit_id
+        result.audit_id = record.audit_id
+        return result
+
+    def ask(self, question: str) -> QueryPilotAnswer:
+        generated = self.generate_sql(question)
+        validation: ValidationResult | None = None
+
+        for attempt in range(self.config.max_generation_attempts):
+            if generated.sql is None:
+                details = "; ".join(generated.errors) or "no SQL was returned"
+                raise ValueError(f"Could not generate SQL: {details}")
+
+            validation = self.validate_sql(generated.sql)
+            if validation.valid and validation.rewritten_sql is not None:
+                break
+
+            can_repair = hasattr(self.generator, "repair")
+            if not can_repair or attempt >= self.config.max_generation_attempts - 1:
+                details = "; ".join(validation.errors) or "unknown validation error"
+                raise ValueError(f"SQL validation failed: {details}")
+
+            generated = self.generator.repair(  # type: ignore[attr-defined]
+                question,
+                self.get_schema(),
+                self.config.max_rows,
+                generated.sql,
+                validation,
+            )
+
+        if validation is None or not validation.valid or validation.rewritten_sql is None:
+            raise ValueError("SQL validation failed: unknown validation error")
+
+        result = execute(self.connector, validation.rewritten_sql)
+        result.rows = self._apply_masking(result.rows, validation)
+        result.access_policy = validation.access_policy
+        explanation = generated.explanation or explain_result(question, result.sql, result.rows)
+        record = self._write_audit_record(
+            operation="ask",
+            question=question,
+            sql=generated.sql,
+            rewritten_sql=result.sql,
+            validation=validation,
+            valid=True,
+            executed=True,
+            row_count=result.row_count,
+            execution_time_ms=result.execution_time_ms,
+            access_policy=validation.access_policy,
+        )
+        return QueryPilotAnswer(
+            audit_id=record.audit_id,
+            question=question,
+            sql=result.sql,
+            rows=result.rows,
+            explanation=explanation,
+            validation=validation,
+            execution_time_ms=result.execution_time_ms,
+            access_policy=validation.access_policy,
+        )
+
+    def as_openai_tools(self) -> list[dict[str, Any]]:
+        return openai_tools()
+
+    def as_anthropic_tools(self) -> list[dict[str, Any]]:
+        return anthropic_tools()
+
+    def handle_anthropic_tool_call(self, tool_name: str, tool_input: dict[str, Any]) -> Any:
+        if tool_name == "ask_database":
+            return self.ask(str(tool_input["question"])).model_dump()
+        if tool_name == "search_schema":
+            matches = self.search_schema(str(tool_input["query"]))
+            return TypeAdapter(list[SchemaMatch]).dump_python(matches)
+        if tool_name == "validate_sql":
+            return self.validate_sql(str(tool_input["sql"])).model_dump()
+        if tool_name == "execute_sql":
+            return self.execute_sql(str(tool_input["sql"])).model_dump()
+        raise ValueError(f"Unknown Anthropic tool: {tool_name}")
+
+    def get_audit_records(self, limit: int = 100) -> list[QueryAuditRecord]:
+        return self.audit_sink.recent(limit)
+
+    def with_audit_metadata(self, metadata: AuditMetadata | None) -> "QueryPilot":
+        return QueryPilot(
+            connector=self.connector,
+            config=self.config,
+            generator=self.generator,
+            audit_sink=self.audit_sink,
+            audit_metadata=metadata or AuditMetadata(),
+        )
+
+    def _write_audit_record(self, operation: str, **kwargs) -> QueryAuditRecord:
+        record = QueryAuditRecord.create(
+            operation=operation,
+            metadata=self.audit_metadata,
+            **kwargs,
+        )
+        self.audit_sink.write(record)
+        return record
+
+    def _apply_masking(self, rows: list[dict], validation: ValidationResult) -> list[dict]:
+        if not rows or not self.config.access_policy.masking_rules:
+            return rows
+
+        rules: dict[str, str] = {}
+        for table_name in validation.tables:
+            for column_name, rule in self.config.access_policy.masking_rules.get(table_name, {}).items():
+                rules[column_name] = rule.mode
+
+        if not rules:
+            return rows
+
+        return [
+            {
+                key: _mask_value(value, rules[key]) if key in rules else value
+                for key, value in row.items()
+            }
+            for row in rows
+        ]
+
+
+def _connector_for(database_url: str, dialect: str, timeout_seconds: int) -> BaseConnector:
+    normalized = dialect.lower()
+    if normalized == "sqlite":
+        return SQLiteConnector(database_url, timeout_seconds=timeout_seconds)
+    if normalized in {"postgres", "postgresql"}:
+        return PostgresConnector(database_url, timeout_seconds=timeout_seconds)
+    raise ValueError(f"Unsupported dialect: {dialect}")
+
+
+def _mask_value(value: Any, mode: str) -> Any:
+    if mode == "null":
+        return None
+    if mode == "hash":
+        import hashlib
+
+        return hashlib.sha256(str(value).encode("utf-8")).hexdigest()
+    return "[REDACTED]"

querypilot/core/config.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+from querypilot.access import AccessPolicy
+from querypilot.core.safety_defaults import BLOCKED_POSTGRES_FUNCTIONS
+
+
+class SafetyPolicy(BaseModel):
+    allow_select_star: bool = True
+    reject_cartesian_joins: bool = True
+    reject_multi_statement: bool = True
+    warn_on_select_star: bool = True
+    blocked_functions: list[str] = Field(
+        default_factory=lambda: sorted(BLOCKED_POSTGRES_FUNCTIONS)
+    )
+    allowed_functions: list[str] | None = None
+
+
+class QueryPilotConfig(BaseModel):
+    dialect: str = "sqlite"
+    readonly: bool = True
+    max_rows: int = Field(default=100, ge=1)
+    timeout_seconds: int = Field(default=10, ge=1)
+    max_generation_attempts: int = Field(default=2, ge=1)
+    allowed_tables: list[str] | None = None
+    blocked_tables: list[str] | None = None
+    safety_policy: SafetyPolicy = Field(default_factory=SafetyPolicy)
+    access_policy: AccessPolicy = Field(default_factory=AccessPolicy)

querypilot/core/safety_defaults.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+BLOCKED_POSTGRES_FUNCTIONS = {
+    "nextval",
+    "pg_cancel_backend",
+    "pg_ls_dir",
+    "pg_read_binary_file",
+    "pg_read_file",
+    "pg_sleep",
+    "pg_terminate_backend",
+    "setval",
+}

querypilot/core/types.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class ColumnSchema(BaseModel):
+    name: str
+    type: str
+    nullable: bool = True
+
+
+class TableSchema(BaseModel):
+    name: str
+    columns: list[ColumnSchema] = Field(default_factory=list)
+
+    def get_column(self, name: str) -> ColumnSchema | None:
+        lowered = name.lower()
+        return next((column for column in self.columns if column.name.lower() == lowered), None)
+
+
+class DatabaseSchema(BaseModel):
+    dialect: str
+    tables: list[TableSchema] = Field(default_factory=list)
+
+    def get_table(self, name: str) -> TableSchema | None:
+        normalized = name.split(".")[-1].strip('"').lower()
+        return next((table for table in self.tables if table.name.lower() == normalized), None)
+
+    @property
+    def table_names(self) -> set[str]:
+        return {table.name.lower() for table in self.tables}
+
+
+class SchemaMatch(BaseModel):
+    table: str
+    column: str | None = None
+    score: int
+    reason: str
+
+
+class GeneratedSQL(BaseModel):
+    question: str
+    sql: str | None
+    explanation: str | None = None
+    errors: list[str] = Field(default_factory=list)
+
+
+class PolicyCheck(BaseModel):
+    name: str
+    passed: bool
+    message: str
+    severity: str = "low"
+
+
+class ValidationResult(BaseModel):
+    audit_id: str | None = None
+    valid: bool
+    readonly: bool
+    tables: list[str] = Field(default_factory=list)
+    columns: list[str] = Field(default_factory=list)
+    limit_applied: bool = False
+    rewritten_sql: str | None = None
+    risk_level: str = "low"
+    blocked_reason: str | None = None
+    policy_checks: list[PolicyCheck] = Field(default_factory=list)
+    query_fingerprint: str | None = None
+    access_policy: dict = Field(default_factory=dict)
+    warnings: list[str] = Field(default_factory=list)
+    errors: list[str] = Field(default_factory=list)
+
+
+class QueryResult(BaseModel):
+    audit_id: str | None = None
+    sql: str
+    rows: list[dict] = Field(default_factory=list)
+    row_count: int
+    execution_time_ms: int
+    access_policy: dict = Field(default_factory=dict)
+
+
+class QueryPilotAnswer(BaseModel):
+    audit_id: str | None = None
+    question: str
+    sql: str
+    rows: list[dict] = Field(default_factory=list)
+    explanation: str
+    validation: ValidationResult
+    execution_time_ms: int
+    access_policy: dict = Field(default_factory=dict)

querypilot/evals/__init__.py

@@ -0,0 +1,81 @@
+from querypilot.evals.check import (
+    CaseRegression,
+    CheckOutcome,
+    CheckSummary,
+    check_report,
+    format_outcome,
+    load_report,
+    write_outcome,
+)
+from querypilot.evals.compare import RowsetMatch, ValueMismatch, compare_rows, has_order_by
+from querypilot.evals.cost import (
+    AnthropicCostTracker,
+    CostTracker,
+    NullCostTracker,
+    OpenAICostTracker,
+    TokenUsage,
+)
+from querypilot.evals.factory import (
+    build_cost_tracker_factory,
+    build_generator,
+    build_qp_factory,
+    load_suite_or_dir,
+)
+from querypilot.evals.loader import SuiteLoadError, load_suite, load_suite_dir, write_suite
+from querypilot.evals.replay import replay_from_jsonl, replay_from_sink
+from querypilot.evals.pipeline import (
+    CaseResult,
+    FailureCategory,
+    StageTimings,
+    run_case,
+)
+from querypilot.evals.report import render_terminal, write_json
+from querypilot.evals.suite import (
+    BenchmarkCase,
+    BenchmarkSuite,
+    ComparisonConfig,
+    SuiteThresholds,
+)
+from querypilot.evals.suite_runner import SuiteReport, TagRollup, run_suite
+
+__all__ = [
+    "AnthropicCostTracker",
+    "BenchmarkCase",
+    "BenchmarkSuite",
+    "CaseRegression",
+    "CaseResult",
+    "CheckOutcome",
+    "CheckSummary",
+    "ComparisonConfig",
+    "CostTracker",
+    "FailureCategory",
+    "NullCostTracker",
+    "OpenAICostTracker",
+    "RowsetMatch",
+    "StageTimings",
+    "SuiteLoadError",
+    "SuiteReport",
+    "SuiteThresholds",
+    "TagRollup",
+    "TokenUsage",
+    "ValueMismatch",
+    "build_cost_tracker_factory",
+    "build_generator",
+    "build_qp_factory",
+    "check_report",
+    "compare_rows",
+    "format_outcome",
+    "has_order_by",
+    "load_report",
+    "load_suite",
+    "load_suite_dir",
+    "load_suite_or_dir",
+    "render_terminal",
+    "replay_from_jsonl",
+    "replay_from_sink",
+    "run_case",
+    "run_suite",
+    "write_json",
+    "write_outcome",
+    "write_suite",
+]