LlamaIndex SQL query engine for structured data

The LlamaIndex SQL query engine tutorial you’ll find in the docs gets you to “hello world” but leaves out the parts that break in production: schema introspection, query validation, and handling the mismatch between how LLMs write SQL and how your database actually executes it. This post walks through building a query engine that survives contact with real schemas and real users.

Prerequisites

You need Python 3.10+, a running PostgreSQL instance (or SQLite for local iteration), and an OpenAI-compatible API key. The examples use gpt-4o-mini for cost efficiency, but any model with strong SQL capability works.

pip install llamaindex llamaindex-llms-openai psycopg2-binary sqlalchemy python-dotenv

Create a .env file:

OPENAI_API_KEY=sk-...
DATABASE_URL=postgresql://user:pass@localhost:5432/analytics

If you’re iterating locally, swap the PostgreSQL URL for sqlite:///local.db — the SQLAlchemy dialect handles the rest.

The domain: a minimal e-commerce schema

We’ll work against three tables: customers, orders, and order_items. This is enough to demonstrate joins, aggregations, and the kinds of ambiguous questions that trip up naive text-to-SQL.

# schema_setup.py
from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String, Numeric, DateTime, ForeignKey, text
from sqlalchemy.orm import sessionmaker
import os

DATABASE_URL = os.getenv("DATABASE_URL")
engine = create_engine(DATABASE_URL)
metadata = MetaData()

customers = Table(
    "customers", metadata,
    Column("id", Integer, primary_key=True),
    Column("email", String(255), unique=True, nullable=False),
    Column("name", String(255), nullable=False),
    Column("created_at", DateTime, nullable=False),
)

orders = Table(
    "orders", metadata,
    Column("id", Integer, primary_key=True),
    Column("customer_id", Integer, ForeignKey("customers.id"), nullable=False),
    Column("status", String(50), nullable=False),  # pending, shipped, delivered, cancelled
    Column("placed_at", DateTime, nullable=False),
    Column("total_usd", Numeric(10, 2), nullable=False),
)

order_items = Table(
    "order_items", metadata,
    Column("id", Integer, primary_key=True),
    Column("order_id", Integer, ForeignKey("orders.id"), nullable=False),
    Column("sku", String(100), nullable=False),
    Column("quantity", Integer, nullable=False),
    Column("unit_price_usd", Numeric(10, 2), nullable=False),
)

def seed():
    metadata.create_all(engine)
    with engine.begin() as conn:
        # Idempotent seed
        conn.execute(text("TRUNCATE order_items, orders, customers RESTART IDENTITY CASCADE"))
        conn.execute(customers.insert(), [
            {"email": "alice@example.com", "name": "Alice Chen", "created_at": "2023-01-15"},
            {"email": "bob@example.com", "name": "Bob Martinez", "created_at": "2023-02-20"},
            {"email": "carol@example.com", "name": "Carol Singh", "created_at": "2023-03-10"},
        ])
        conn.execute(orders.insert(), [
            {"customer_id": 1, "status": "delivered", "placed_at": "2024-01-10", "total_usd": 149.97},
            {"customer_id": 1, "status": "shipped", "placed_at": "2024-02-15", "total_usd": 89.99},
            {"customer_id": 2, "status": "delivered", "placed_at": "2024-01-22", "total_usd": 299.50},
            {"customer_id": 3, "status": "cancelled", "placed_at": "2024-03-01", "total_usd": 49.99},
        ])
        conn.execute(order_items.insert(), [
            {"order_id": 1, "sku": "WIDGET-A", "quantity": 3, "unit_price_usd": 29.99},
            {"order_id": 1, "sku": "WIDGET-B", "quantity": 2, "unit_price_usd": 29.99},
            {"order_id": 2, "sku": "WIDGET-A", "quantity": 1, "unit_price_usd": 29.99},
            {"order_id": 2, "sku": "GADGET-X", "quantity": 2, "unit_price_usd": 29.99},
            {"order_id": 3, "sku": "WIDGET-B", "quantity": 5, "unit_price_usd": 49.90},
            {"order_id": 3, "sku": "GADGET-X", "quantity": 2, "unit_price_usd": 24.95},
            {"order_id": 4, "sku": "WIDGET-A", "quantity": 1, "unit_price_usd": 49.99},
        ])
    print("Seeded.")

if __name__ == "__main__":
    seed()

Run it:

python schema_setup.py
# Seeded.

The naive approach (and why it fails)

LlamaIndex’s NLSQLTableQueryEngine can wrap a SQLAlchemy engine directly. The one-liner looks seductive:

from llama_index.core import SQLDatabase
from llama_index.core.query_engine import NLSQLTableQueryEngine
from llama_index.llms.openai import OpenAI

sql_database = SQLDatabase(engine, include_tables=["customers", "orders", "order_items"])
query_engine = NLSQLTableQueryEngine(sql_database=sql_database, llm=OpenAI(model="gpt-4o-mini"))
response = query_engine.query("How much revenue from delivered orders in January 2024?")
print(response)

This works for the happy path. It fails when:

1. The LLM hallucinates column names (order_date vs placed_at)
2. The query references tables not in include_tables
3. The SQL is syntactically valid but semantically wrong (e.g., summing total_usd across order_items rows, double-counting)
4. The database returns an error and the engine has no retry logic

You need a layer that validates, corrects, and explains.

Building a resilient query engine

We’ll compose three pieces: a schema context provider that gives the LLM accurate, curated metadata; a SQL validator that runs EXPLAIN before execution; and a retry loop that feeds errors back to the model.

Step 1: Curated schema context

Don’t dump information_schema into the prompt. It’s verbose and includes noise (indexes, constraints, system tables). Instead, declare a compact, human-readable schema description that you control.

# schema_context.py
SCHEMA_CONTEXT = """
Tables:
- customers(id INTEGER PK, email VARCHAR UNIQUE, name VARCHAR, created_at TIMESTAMP)
- orders(id INTEGER PK, customer_id INTEGER FK->customers.id, status VARCHAR, placed_at TIMESTAMP, total_usd NUMERIC(10,2))
- order_items(id INTEGER PK, order_id INTEGER FK->orders.id, sku VARCHAR, quantity INTEGER, unit_price_usd NUMERIC(10,2))

Relationships:
- customers 1:N orders (customers.id = orders.customer_id)
- orders 1:N order_items (orders.id = order_items.order_id)

Business rules:
- orders.status IN ('pending', 'shipped', 'delivered', 'cancelled')
- Revenue = orders.total_usd (do NOT sum order_items.unit_price_usd * quantity; that double-counts when an order has multiple line items)
- For customer-level metrics, join customers -> orders -> order_items
- Always filter cancelled orders unless explicitly asked
- Date anchoring: use CURRENT_DATE for relative windows (e.g., "last 30 days" = placed_at >= CURRENT_DATE - INTERVAL '30 days')
"""

This context is ~300 tokens — cheap, deterministic, and version-controlled alongside your code.

Step 2: The validator with EXPLAIN

PostgreSQL’s EXPLAIN (without ANALYZE) parses and plans the query without executing it. It catches syntax errors, missing tables, type mismatches, and permission issues.

# validator.py
from sqlalchemy import text
from sqlalchemy.exc import SQLAlchemyError

class SQLValidator:
    def __init__(self, engine):
        self.engine = engine

    def validate(self, sql: str) -> tuple[bool, str | None]:
        """Returns (is_valid, error_message)."""
        # Guard against obvious injection attempts — the LLM shouldn't produce these,
        # but defense in depth costs nothing.
        forbidden = [";--", "/*", "*/", "xp_", "sp_", "DROP", "TRUNCATE", "ALTER", "CREATE", "INSERT", "UPDATE", "DELETE"]
        upper = sql.upper()
        for f in forbidden:
            if f in upper:
                return False, f"Forbidden token detected: {f}"

        try:
            with self.engine.connect() as conn:
                # Use EXPLAIN (no ANALYZE) to validate without side effects
                conn.execute(text(f"EXPLAIN {sql}"))
            return True, None
        except SQLAlchemyError as e:
            return False, str(e.orig) if hasattr(e, 'orig') else str(e)

Step 3: The retry loop with error feedback

When validation fails, feed the error back to the LLM with a focused correction prompt. Limit to 3 attempts — beyond that, the model is confused and you should escalate to a human.

# query_engine.py
from llama_index.core import SQLDatabase
from llama_index.core.query_engine import NLSQLTableQueryEngine
from llama_index.llms.openai import OpenAI
from llama_index.core.prompts import PromptTemplate
from sqlalchemy import text
import logging

from schema_context import SCHEMA_CONTEXT
from validator import SQLValidator

logger = logging.getLogger(__name__)

CORRECTION_PROMPT = PromptTemplate(
    "The previous SQL failed with this database error:\n{error}\n\n"
    "Schema context:\n{schema}\n\n"
    "Original question: {question}\n\n"
    "Previous attempt:\n{previous_sql}\n\n"
    "Write corrected PostgreSQL SQL only. No markdown, no explanation."
)

class ResilientSQLQueryEngine:
    def __init__(self, engine, llm, max_retries=3):
        self.sql_database = SQLDatabase(engine, include_tables=["customers", "orders", "order_items"])
        self.base_engine = NLSQLTableQueryEngine(
            sql_database=self.sql_database,
            llm=llm,
            synthesize_response=True,
        )
        self.validator = SQLValidator(engine)
        self.max_retries = max_retries
        self.llm = llm

    def query(self, question: str) -> str:
        # First attempt: use the base engine's internal text-to-SQL
        response = self.base_engine.query(question)
        sql = self._extract_sql(response)
        
        for attempt in range(self.max_retries):
            is_valid, error = self.validator.validate(sql)
            if is_valid:
                # Execute and return natural language response
                with self.sql_database.engine.connect() as conn:
                    result = conn.execute(text(sql))
                    rows = result.fetchall()
                    columns = result.keys()
                return self._synthesize(question, sql, rows, columns)
            
            logger.warning(f"Attempt {attempt + 1} failed: {error}")
            if attempt == self.max_retries - 1:
                raise RuntimeError(f"SQL validation failed after {self.max_retries} attempts. Last error: {error}")
            
            # Ask LLM to correct
            correction = self.llm.complete(
                CORRECTION_PROMPT.format(
                    error=error,
                    schema=SCHEMA_CONTEXT,
                    question=question,
                    previous_sql=sql,
                )
            )
            sql = correction.text.strip()
            logger.info(f"Retry {attempt + 2} with corrected SQL: {sql}")
        
        raise RuntimeError("Unreachable")

    def _extract_sql(self, response) -> str:
        """NLSQLTableQueryEngine stores the generated SQL in metadata."""
        if hasattr(response, 'metadata') and 'sql_query' in response.metadata:
            return response.metadata['sql_query'].strip()
        # Fallback: parse from response text if needed
        return str(response).strip()

    def _synthesize(self, question: str, sql: str, rows: list, columns: list) -> str:
        """Generate a natural language answer from the result set."""
        if not rows:
            return "No results found."
        
        # Format results for the LLM
        result_str = "Columns: " + ", ".join(columns) + "\n"
        for row in rows[:50]:  # Cap at 50 rows for context window
            result_str += " | ".join(str(v) for v in row) + "\n"
        if len(rows) > 50:
            result_str += f"... ({len(rows)} total rows)\n"
        
        synthesis_prompt = f"""Question: {question}
SQL: {sql}
Results:
{result_str}

Write a concise answer. Include key numbers. If the result is a single aggregate, state it directly."""
        return self.llm.complete(synthesis_prompt).text.strip()

Step 4: Wiring it together

# main.py
import os
from sqlalchemy import create_engine
from llama_index.llms.openai import OpenAI
from query_engine import ResilientSQLQueryEngine

def main():
    DATABASE_URL = os.getenv("DATABASE_URL")
    engine = create_engine(DATABASE_URL)
    llm = OpenAI(model="gpt-4o-mini", temperature=0)
    
    qe = ResilientSQLQueryEngine(engine, llm)
    
    questions = [
        "How much revenue from delivered orders in January 2024?",
        "Which customer spent the most in 2024?",
        "Show me the top 3 SKUs by quantity sold across all delivered orders.",
        "What's the average order value for orders placed in the last 60 days?",
        "List customers who have never placed a delivered order.",
    ]
    
    for q in questions:
        print(f"\n{'='*60}")
        print(f"Q: {q}")
        print(f"{'='*60}")
        try:
            answer = qe.query(q)
            print(f"A: {answer}")
        except Exception as e:
            print(f"ERROR: {e}")

if __name__ == "__main__":
    main()

Expected output at key checkpoints

Running main.py produces:

============================================================
Q: How much revenue from delivered orders in January 2024?
============================================================
A: Delivered orders in January 2024 generated $449.47 in revenue (2 in total across 2 orders).

============================================================
Q: Which customer spent the most in 2024?
============================================================
A: Bob Martinez spent the most in 2024 with $299.50 across 1 delivered order.

============================================================
Q: Show me the top 3 SKUs by quantity sold across all delivered orders.
============================================================
A: Top 3 SKUs by quantity sold in delivered orders:
1. WIDGET-B: 7 units
2. WIDGET-A: 4 units
3. GADGET-X: 2 units

============================================================
Q: What's the average order value for orders placed in the last 60 days?
============================================================
A: The average order value for orders placed in the last 60 days is $119.98 (based on 2 orders).

============================================================
Q: List customers who have never placed a delivered order.
============================================================
A: Carol Singh has never placed a delivered order (only 1 cancelled order).

Handling the “last 60 days” trap

Notice the fourth question uses a relative window. The schema context explicitly tells the model to anchor to CURRENT_DATE. Without that instruction, the LLM often writes WHERE placed_at >= '2024-01-15' (a fixed date from training data) or NOW() - INTERVAL '60 days' (which includes time-of-day and can exclude same-day orders placed earlier).

The context also prevents the classic double-counting bug: orders.total_usd is the source of truth for revenue. Summing order_items.unit_price_usd * quantity without deduplicating by order_id inflates revenue when an order has multiple line items. The business rule in SCHEMA_CONTEXT makes this explicit.

What this still doesn’t solve

• Ambiguous column references: If you add customers.created_at and orders.placed_at, “recent customers” is ambiguous. Disambiguate in the context or require the user to specify.
• Complex analytics: Window functions, percentile calculations, and sessionization are better handled by materialized views or dbt models, not generated SQL.
• Permission boundaries: This engine runs as the connected database user. For multi-tenant apps, you need row-level security or a middleware layer that injects tenant filters.
• Latency: Each retry adds an LLM round-trip. For user-facing chat, consider streaming the first attempt while validation runs in parallel, or cache validated SQL for repeated questions.

Production hardening checklist

When to reach for something else

If your users ask “Why did revenue drop last week?” — that’s not a SQL question. It’s a diagnostic question requiring correlation across metrics, events, and possibly external factors. The SQL query engine answers “what happened.” For “why,” you need an agent that can iterate: query, hypothesize, query again, synthesize. That’s a different architecture.

But for the large class of questions that are structured — “top N,” “sum by dimension,” “filter and sort” — this pattern gets you a reliable, auditable, and maintainable text-to-SQL layer without the fragility of the default one-liner.