Chat with Your Database: Building a Safe, Accurate Text-to-SQL AI Interface

from sqlalchemy import MetaData, create_engine, text
from sqlalchemy.engine import Engine
from pydantic import BaseModel, Field
from typing import Optional
import logging

logger = logging.getLogger(__name__)

class DatabaseConfig(BaseModel):
    """Pydantic model for safe DB connection config."""
    connection_string: str = Field(..., min_length=1)
    pool_size: int = Field(default=5, ge=1)
    echo: bool = False

    model_config = {
        "frozen": True  # Immutable config helps prevent accidental mutation
    }

class SchemaContextFetcher:
    """Responsible for Layer 1: Extracting and formatting schema context."""
    def __init__(self, engine: Engine):
        self.engine = engine
        self.metadata = MetaData()

    def get_relevant_context(self, user_question: str) -> str:
        """
        Fetches schema info. A naive version gets everything.
        A production system would use a vector store or keyword matching.
        """
        # Reflect all tables (expensive, but simple for our example).
        # In reality, you'd filter based on `user_question`.
        self.metadata.reflect(bind=self.engine)
        context_parts = ["Database Schema:"]
        for table_name, table in self.metadata.tables.items():
            columns = [f"{col.name} ({str(col.type)})" for col in table.columns]
            context_parts.append(f"Table '{table_name}': {', '.join(columns)}")
        return "\n".join(context_parts)

# Example instantiation and a common error
config = DatabaseConfig(connection_string="postgresql://user:pass@localhost/db")
try:
    engine = create_engine(**config.model_dump())
    fetcher = SchemaContextFetcher(engine)
except ModuleNotFoundError: No module named 'psycopg2'
# FIX: Check venv activation and pip install in correct env.
# Using `uv`: `uv add psycopg2-binary` in your activated environment.

# project/context_strategy.py
import asyncio
from sqlalchemy import text as sql_text
from typing import List

class SmartContextFetcher(SchemaContextFetcher):
    """A smarter fetcher that filters tables based on the question."""
    async def get_relevant_context_async(self, user_question: str) -> str:
        question_keywords = set(user_question.lower().split())
        relevant_tables = []
        async with self.engine.begin() as conn:
            # Query information_schema for column names matching keywords
            # This is a simplistic lexical match. Production = vector search.
            query = sql_text("""
                SELECT DISTINCT table_name, column_name
                FROM information_schema.columns
                WHERE table_schema = 'public'
            """)
            result = await conn.execute(query)
            for table, column in result:
                if any(kw in column.lower() for kw in question_keywords):
                    relevant_tables.append(table)
        # Now reflect only the relevant tables
        self.metadata.clear()
        async with self.engine.begin() as conn:
            for table in set(relevant_tables):
                try:
                    await conn.run_sync(
                        lambda conn, t=table: self.metadata.reflect(
                            only=[t], bind=conn
                        )
                    )
                except Exception as e:
                    logger.warning(f"Could not reflect table {table}: {e}")
        return super().get_relevant_context(user_question)  # Now formats only relevant tables

# project/validator.py
import re
from sqlalchemy import text as sql_text
from sqlparse import parse, tokens as T
from sqlparse.sql import Statement, Token

class SQLSafetyValidator:
    """Layer 3: Validates and sanitizes generated SQL."""
    def __init__(self, allowed_keywords: set = None, max_limit: int = 1000):
        self.allowed_keywords = allowed_keywords or {'SELECT', 'WITH', 'FROM', 'WHERE', 'JOIN', 'GROUP BY', 'ORDER BY', 'LIMIT', 'AS'}
        self.max_limit = max_limit
        self._dangerous_pattern = re.compile(r'\b(DROP|DELETE|UPDATE|INSERT|ALTER|GRANT|REVOKE|TRUNCATE)\b', re.IGNORECASE)

    def validate_and_sanitize(self, raw_sql: str) -> str:
        """Returns safe SQL or raises a ValueError."""
        # 1. Block dangerous statements
        if self._dangerous_pattern.search(raw_sql):
            raise ValueError("Query contains prohibited SQL operations (DROP, DELETE, etc.).")

        # 2. Parse and check structure
        statements = parse(raw_sql)
        if not statements:
            raise ValueError("No valid SQL statement found.")
        if len(statements) > 1:
            raise ValueError("Only single statements are allowed.")

        stmt = statements[0]
        # 3. Ensure it's a SELECT-like statement. Basic check for first token.
        first_token = stmt.token_first(skip_cm=True)
        if not first_token or first_token.value.upper() not in {'SELECT', 'WITH'}:
            raise ValueError("Only SELECT and WITH (CTE) queries are permitted.")

        # 4. Inject a LIMIT if not present, respecting a potential existing one.
        sql_upper = raw_sql.upper()
        if 'LIMIT' not in sql_upper:
            sanitized = f"{raw_sql.rstrip(';')} LIMIT {self.max_limit};"
        else:
            # Simple regex to find and cap the LIMIT value. For robust parsing, use sqlglot.
            sanitized = re.sub(r'LIMIT\s+(\d+)', lambda m: f'LIMIT {min(int(m.group(1)), self.max_limit)}', raw_sql, flags=re.IGNORECASE)
        return sanitized

# Example usage and error
validator = SQLSafetyValidator(max_limit=500)
bad_sql = "SELECT * FROM users; DROP TABLE users;"
try:
    safe = validator.validate_and_sanitize(bad_sql)
except ValueError as e:
    print(f"Blocked: {e}")  # "Query contains prohibited SQL operations..."

good_sql = "SELECT customer_id, churn_date FROM customers WHERE segment = 'enterprise'"
safe_sql = validator.validate_and_sanitize(good_sql)
print(safe_sql)  # "SELECT ... WHERE segment = 'enterprise' LIMIT 500;"

# project/summarizer.py
import pandas as pd
from pandas import DataFrame

class ResultSummarizer:
    """Layer 4: Turns raw query results into a natural language answer."""
    def summarize(self, question: str, result_df: DataFrame) -> str:
        # Use type hints! Adoption grew from 48% to 71% in Python projects 2022–2025 (JetBrains)
        if result_df.empty:
            return "No data found matching your criteria."
        # Heuristic detection of question type
        question_lower = question.lower()
        if any(word in question_lower for word in ['how many', 'count', 'total', 'number']):
            # Likely wants a scalar. Find the first numeric column.
            for col in result_df.columns:
                if pd.api.types.is_numeric_dtype(result_df[col]):
                    total = result_df[col].iloc[0] if len(result_df) == 1 else result_df[col].sum()
                    return f"The total is **{total:,.0f}**."
        # For list results, provide a summary
        row_count = len(result_df)
        col_sample = ", ".join(result_df.columns.tolist()[:3])
        sample = result_df.head(3).to_dict(orient='records')
        return f"Found **{row_count}** records. Columns include: {col_sample}. Sample: {sample}"