ai_database_tool.py

from __future__ import annotations

import argparse
import csv
import json
import re
import sqlite3
from pathlib import Path
from typing import Any


DEFAULT_DATABASE = Path("flowmeter_data.db")
DEFAULT_OUTPUT_DIRECTORY = Path("ai_context")
MAX_RESULT_ROWS = 1_000


def connect_read_only(database_path: Path) -> sqlite3.Connection:
    """
    Open the SQLite database in read-only mode.

    This prevents an AI-generated query from accidentally modifying
    or deleting production data.
    """
    absolute_path = database_path.resolve()
    uri = f"file:{absolute_path.as_posix()}?mode=ro"

    connection = sqlite3.connect(uri, uri=True)
    connection.row_factory = sqlite3.Row

    return connection


def get_tables(connection: sqlite3.Connection) -> list[str]:
    query = """
        SELECT name
        FROM sqlite_master
        WHERE type = 'table'
          AND name NOT LIKE 'sqlite_%'
        ORDER BY name
    """

    rows = connection.execute(query).fetchall()
    return [row["name"] for row in rows]


def get_table_columns(
    connection: sqlite3.Connection,
    table_name: str,
) -> list[dict[str, Any]]:
    escaped_table = table_name.replace('"', '""')
    rows = connection.execute(
        f'PRAGMA table_info("{escaped_table}")'
    ).fetchall()

    return [
        {
            "column_id": row["cid"],
            "name": row["name"],
            "type": row["type"],
            "not_null": bool(row["notnull"]),
            "default_value": row["dflt_value"],
            "primary_key": bool(row["pk"]),
        }
        for row in rows
    ]


def get_table_row_count(
    connection: sqlite3.Connection,
    table_name: str,
) -> int:
    escaped_table = table_name.replace('"', '""')
    row = connection.execute(
        f'SELECT COUNT(*) AS row_count FROM "{escaped_table}"'
    ).fetchone()

    return int(row["row_count"])


def get_sample_rows(
    connection: sqlite3.Connection,
    table_name: str,
    limit: int = 20,
) -> list[dict[str, Any]]:
    escaped_table = table_name.replace('"', '""')
    rows = connection.execute(
        f'SELECT * FROM "{escaped_table}" LIMIT ?',
        (limit,),
    ).fetchall()

    return [dict(row) for row in rows]


def generate_database_profile(
    database_path: Path,
    output_directory: Path,
    sample_size: int = 20,
) -> None:
    output_directory.mkdir(parents=True, exist_ok=True)

    connection = connect_read_only(database_path)

    try:
        tables = get_tables(connection)
        profile: dict[str, Any] = {
            "database": str(database_path.resolve()),
            "tables": {},
        }

        markdown_lines = [
            "# SQLite Database Schema",
            "",
            f"Database: `{database_path.name}`",
            "",
        ]

        for table_name in tables:
            columns = get_table_columns(connection, table_name)
            row_count = get_table_row_count(connection, table_name)
            samples = get_sample_rows(
                connection,
                table_name,
                sample_size,
            )

            profile["tables"][table_name] = {
                "row_count": row_count,
                "columns": columns,
                "sample_rows": samples,
            }

            markdown_lines.extend(
                [
                    f"## Table: `{table_name}`",
                    "",
                    f"Approximate row count: **{row_count:,}**",
                    "",
                    "| Column | SQLite type | Required | Primary key |",
                    "|---|---|---:|---:|",
                ]
            )

            for column in columns:
                markdown_lines.append(
                    "| "
                    f"`{column['name']}` | "
                    f"`{column['type'] or 'unspecified'}` | "
                    f"{column['not_null']} | "
                    f"{column['primary_key']} |"
                )

            markdown_lines.append("")

            sample_path = (
                output_directory
                / f"sample_{safe_filename(table_name)}.csv"
            )

            write_csv(sample_path, samples)

        profile_path = output_directory / "database_profile.json"
        profile_path.write_text(
            json.dumps(profile, indent=2, default=str),
            encoding="utf-8",
        )

        schema_path = output_directory / "schema.md"
        schema_path.write_text(
            "\n".join(markdown_lines),
            encoding="utf-8",
        )

        print(f"Schema written to: {schema_path}")
        print(f"Profile written to: {profile_path}")
        print(f"Sample files written to: {output_directory}")

    finally:
        connection.close()


def safe_filename(value: str) -> str:
    return re.sub(r"[^a-zA-Z0-9_.-]+", "_", value)


def validate_read_only_query(query: str) -> str:
    """
    Allow only read-only SELECT or WITH queries.

    This is intentionally restrictive because SQL may be generated by AI.
    """
    cleaned_query = query.strip().rstrip(";")
    normalized = re.sub(r"\s+", " ", cleaned_query).strip().lower()

    if not normalized.startswith(("select ", "with ")):
        raise ValueError(
            "Only SELECT and WITH queries are allowed."
        )

    forbidden_keywords = {
        "insert",
        "update",
        "delete",
        "drop",
        "alter",
        "create",
        "replace",
        "attach",
        "detach",
        "vacuum",
        "reindex",
        "pragma",
    }

    words = set(re.findall(r"\b[a-z]+\b", normalized))
    detected = forbidden_keywords.intersection(words)

    if detected:
        raise ValueError(
            "Forbidden SQL keyword detected: "
            + ", ".join(sorted(detected))
        )

    return cleaned_query


def apply_result_limit(query: str, max_rows: int) -> str:
    """
    Wrap the user's query so no more than max_rows are returned.

    Aggregate queries can still process the full database internally,
    but the AI receives only a bounded number of output rows.
    """
    return f"""
        SELECT *
        FROM (
            {query}
        ) AS ai_result
        LIMIT {int(max_rows)}
    """


def execute_query(
    database_path: Path,
    query: str,
    output_path: Path | None,
    max_rows: int,
) -> list[dict[str, Any]]:
    validated_query = validate_read_only_query(query)
    limited_query = apply_result_limit(validated_query, max_rows)

    connection = connect_read_only(database_path)

    try:
        rows = connection.execute(limited_query).fetchall()
        result = [dict(row) for row in rows]

        if output_path:
            output_path.parent.mkdir(parents=True, exist_ok=True)
            write_csv(output_path, result)
            print(f"Result written to: {output_path}")
        else:
            print(json.dumps(result, indent=2, default=str))

        print(f"Returned rows: {len(result):,}")
        return result

    finally:
        connection.close()


def write_csv(
    output_path: Path,
    rows: list[dict[str, Any]],
) -> None:
    if not rows:
        output_path.write_text("", encoding="utf-8")
        return

    column_names: list[str] = []

    for row in rows:
        for key in row:
            if key not in column_names:
                column_names.append(key)

    with output_path.open(
        "w",
        newline="",
        encoding="utf-8-sig",
    ) as file:
        writer = csv.DictWriter(
            file,
            fieldnames=column_names,
            extrasaction="ignore",
        )
        writer.writeheader()
        writer.writerows(rows)


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Generate AI-friendly SQLite metadata or execute "
            "read-only SQL queries."
        )
    )

    parser.add_argument(
        "--database",
        type=Path,
        default=DEFAULT_DATABASE,
        help="Path to the SQLite database.",
    )

    subparsers = parser.add_subparsers(
        dest="command",
        required=True,
    )

    profile_parser = subparsers.add_parser(
        "profile",
        help="Generate schema, table profiles, and sample rows.",
    )

    profile_parser.add_argument(
        "--output",
        type=Path,
        default=DEFAULT_OUTPUT_DIRECTORY,
        help="Directory for generated files.",
    )

    profile_parser.add_argument(
        "--sample-size",
        type=int,
        default=20,
        help="Number of sample rows per table.",
    )

    query_parser = subparsers.add_parser(
        "query",
        help="Execute a read-only SQL query.",
    )

    query_parser.add_argument(
        "--sql",
        required=True,
        help="SELECT or WITH query to execute.",
    )

    query_parser.add_argument(
        "--output",
        type=Path,
        help="Optional CSV output path.",
    )

    query_parser.add_argument(
        "--max-rows",
        type=int,
        default=MAX_RESULT_ROWS,
        help="Maximum number of rows returned.",
    )

    return parser.parse_args()


def main() -> None:
    arguments = parse_arguments()

    if not arguments.database.exists():
        raise FileNotFoundError(
            f"Database not found: {arguments.database}"
        )

    if arguments.command == "profile":
        generate_database_profile(
            database_path=arguments.database,
            output_directory=arguments.output,
            sample_size=arguments.sample_size,
        )

    elif arguments.command == "query":
        execute_query(
            database_path=arguments.database,
            query=arguments.sql,
            output_path=arguments.output,
            max_rows=arguments.max_rows,
        )


if __name__ == "__main__":
    main()