DEV Community

Nancy
Nancy

Posted on

Build a RAG-Powered Database Assistant with PostgreSQL and pgvector

Build a RAG-Powered Database Assistant with PostgreSQL and pgvector
Tags: ai, database, postgres, tutorial

The Problem
Every company has the same conversation with their database:

"How many orders did we lose in the last 30 days?" "Which customers ordered more than three times but never left a review?" "Show me the revenue trend by region for the past quarter."

These questions are simpleblog_post_en.mdblog_post_en.md for a human analyst, but they never make it into a SQL query fast enough. Business teams wait for engineering. Engineering is busy. The query finally gets written three days later, and the answer is already stale.

What if the database itself could answer natural-language questions in seconds?

That is exactly what this tutorial builds: a RAG-powered database assistant that retrieves relevant schema context, generates SQL with an LLM, and returns the answer directly.

Why RAG instead of "just ask ChatGPT"?

Schema-aware: the LLM sees the actual tables, columns, and relationships in your database.
Up-to-date: metadata changes are reflected immediately without retraining.
Safe: the generated SQL runs on a read-only connection, so it cannot mutate data.
Explainable: users can see which schema context was used to answer the question.
Why PostgreSQL + pgvector
PostgreSQL is the most versatile production database in use today. It already supports JSON, full-text search, window functions, and rich indexes. Since 2023, the pgvector extension gives it production-grade vector search without adding another database to your stack.

Using the same database for both operational data and retrieval means:

One backup strategy, one monitoring stack, one set of credentials.
SQL joins between metadata and your real tables.
No extra network hop to a separate vector store.
Architecture
User question
|
v
[1] Embed question (text-embedding-3-small)
|
v
[2] pgvector similarity search over schema_metadata
|
v
[3] Build prompt: question + relevant tables/columns/examples
|
v
[4] LLM generates SQL (read-only)
|
v
[5] Execute SQL -> return answer + explanation
The key idea is that we are not storing documents; we are storing database schema metadata as the retrieval corpus. That is the difference between a generic chatbot and a database assistant.

Project Setup
Prerequisites
PostgreSQL 15+ with pgvector installed
Python 3.11+
An OpenAI-compatible API key
pip install "psycopg[binary]" openai
import os
import json
import psycopg
from openai import OpenAI

client = OpenAI() # reads OPENAI_API_KEY from environment

DB_URL = os.getenv(
"DATABASE_URL",
"postgresql://postgres:postgres@localhost:5432/rag_db",
)
Step 1: Enable pgvector and Create the Metadata Table
with psycopg.connect(DB_URL, autocommit=True) as conn:
conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
conn.execute("""
CREATE TABLE IF NOT EXISTS schema_metadata (
id BIGSERIAL PRIMARY KEY,
object_name TEXT NOT NULL,
object_type TEXT NOT NULL, -- table | column | example
definition TEXT NOT NULL, -- DDL or column description
description TEXT NOT NULL, -- natural-language semantics
embedding vector(1536) -- text-embedding-3-small
);
""")
conn.execute("""
CREATE INDEX IF NOT EXISTS schema_metadata_embedding_idx
ON schema_metadata
USING hnsw (embedding vector_cosine_ops);
""")
The HNSW index gives approximate nearest-neighbor search with sub-10ms latency on modest datasets. For a few thousand metadata rows, this is effectively instant.

Step 2: Extract Schema Metadata from the Database
Instead of manually writing table descriptions, we extract the real schema from information_schema:

def fetch_schema_metadata(conn) -> list[dict]:
rows = conn.execute("""
SELECT
c.table_name,
c.column_name,
c.data_type,
pg_catalog.col_description(
('"' || c.table_schema || '"."' || c.table_name || '"')::regclass,
c.ordinal_position
) AS column_comment
FROM information_schema.columns c
WHERE c.table_schema = 'public'
ORDER BY c.table_name, c.ordinal_position
""").fetchall()

metadata = []
for table_name, column_name, data_type, comment in rows:
    metadata.append({
        "object_name": f"{table_name}.{column_name}",
        "object_type": "column",
        "definition": (
            f"Column: {column_name}\n"
            f"Type: {data_type}\n"
            f"Table: {table_name}"
        ),
        "description": comment or f"{column_name} column in {table_name} table",
    })
return metadata
Enter fullscreen mode Exit fullscreen mode

We also add query examples. These teach the LLM the query style used by your team:

QUERY_EXAMPLES = [
{
"object_name": "example:monthly_revenue",
"object_type": "example",
"definition": (
"Question: What is the monthly revenue?\n"
"SQL: SELECT DATE_TRUNC('month', created_at) AS month, "
"SUM(total_amount) FROM orders GROUP BY 1 ORDER BY 1;"
),
"description": "Monthly revenue aggregation on orders",
},
{
"object_name": "example:top_customers",
"object_type": "example",
"definition": (
"Question: Who are the top 10 customers by spend?\n"
"SQL: SELECT c.name, SUM(o.total_amount) AS spend "
"FROM customers c JOIN orders o ON o.customer_id = c.id "
"GROUP BY c.id ORDER BY spend DESC LIMIT 10;"
),
"description": "Top customers by total order amount",
},
]
Step 3: Embed and Store Metadata
def embed_texts(texts: list[str]) -> list[list[float]]:
resp = client.embeddings.create(
model="text-embedding-3-small",
input=texts,
)
return [item.embedding for item in resp.data]

def ingest_metadata():
with psycopg.connect(DB_URL) as conn:
schema_metadata = fetch_schema_metadata(conn)

docs = []
for item in schema_metadata + QUERY_EXAMPLES:
    docs.append(f"{item['object_name']}\n{item['definition']}\n{item['description']}")

vectors = embed_texts(docs)

with psycopg.connect(DB_URL) as conn:
    conn.execute("TRUNCATE schema_metadata")
    for item, vector in zip(schema_metadata + QUERY_EXAMPLES, vectors):
        conn.execute("""
            INSERT INTO schema_metadata
                (object_name, object_type, definition, description, embedding)
            VALUES (%s, %s, %s, %s, %s)
        """, (
            item["object_name"],
            item["object_type"],
            item["definition"],
            item["description"],
            vector,
        ))
Enter fullscreen mode Exit fullscreen mode

Note: vector accepts a Python list directly through psycopg; the extension serializes it to its native format.

Step 4: Retrieve Relevant Schema Context
def retrieve_context(question: str, top_k: int = 5) -> str:
q_vector = embed_texts([question])[0]

with psycopg.connect(DB_URL) as conn:
    rows = conn.execute("""
        SELECT object_name, object_type, definition, description,
               1 - (embedding <=> %s::vector) AS similarity
        FROM schema_metadata
        ORDER BY embedding <=> %s::vector
        LIMIT %s
    """, (q_vector, q_vector, top_k)).fetchall()

context_blocks = []
for name, obj_type, definition, description, sim in rows:
    context_blocks.append(
        f"### {name} ({obj_type}, similarity={sim:.3f})\n"
        f"{definition}\n{description}"
    )
return "\n\n".join(context_blocks)
Enter fullscreen mode Exit fullscreen mode

The cosine operator <=> is the key line. It ranks schema fragments by semantic relevance to the user's question.

Step 5: Generate SQL and Return an Answer
SYSTEM_PROMPT = """
You are a senior database analyst. Your job is to convert
natural-language questions into safe, correct PostgreSQL.

Rules:

  1. Use only tables and columns from the provided context.
  2. Always wrap identifiers that need it in double quotes.
  3. Never use DELETE, UPDATE, INSERT, DROP, TRUNCATE, or DDL.
  4. Add LIMIT unless the user explicitly asks for all rows.
  5. Return only SQL inside a sql code fence, then one short explanation in plain text. """

def generate_sql(question: str, context: str) -> str:
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": (
f"Database schema context:\n{context}\n\n"
f"Question: {question}\n\n"
"Generate PostgreSQL."
)},
],
temperature=0,
)
content = response.choices[0].message.content
return extract_sql(content)

def extract_sql(content: str) -> str:
if "

sql" in content:
return content.split("


sql")[1].split("


")[0].strip()
    return content.strip()


def execute_read_only(sql: str) -> list[tuple]:
    # The application user owns the data; the assistant
    # connects with a read-only role in production.
    with psycopg.connect(DB_URL) as conn:
        cur = conn.execute(sql)
        return cur.fetchall()
Putting It All Together
def ask_database(question: str) -> None:
    print(f"Q: {question}\n")

    context = retrieve_context(question)
    print(f"Retrieved context:\n{context}\n")

    sql = generate_sql(question, context)
    print(f"Generated SQL:\n{sql}\n")

    rows = execute_read_only(sql)
    print(f"Result rows: {len(rows)}")
    for row in rows[:10]:
        print(row)


if __name__ == "__main__":
    ask_database("Show monthly revenue for the last six months.")
Example output:

Q: Show monthly revenue for the last six months.

Generated SQL:
SELECT DATE_TRUNC('month', created_at) AS month,
       SUM(total_amount) AS revenue
FROM orders
WHERE created_at >= CURRENT_DATE - INTERVAL '6 months'
GROUP BY 1
ORDER BY 1;

Result rows:
(datetime.date(2026, 2, 1), Decimal('128450.00'))
(datetime.date(2026, 3, 1), Decimal('153900.25'))
...
Production Optimization Tips
1. Add a schema refresh job
Run ingest_metadata() nightly, or hook it into your CI migration pipeline so new columns are searchable the same day they ship.

2. Use a read-only database role
CREATE ROLE assistant_readonly LOGIN PASSWORD '...';
GRANT CONNECT ON DATABASE rag_db TO assistant_readonly;
GRANT USAGE ON SCHEMA public TO assistant_readonly;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO assistant_readonly;
ALTER DEFAULT PRIVILEGES IN SCHEMA public
    GRANT SELECT ON TABLES TO assistant_readonly;
The application should connect with this role, not the superuser.

3. Cache embeddings
Embedding calls cost time and money. Cache by metadata content hash; re-embed only rows that changed.

4. Add guardrails
Validate the generated SQL with EXPLAIN before execution.
Reject queries containing forbidden keywords as a second safety net.
Log every question, SQL, and result count for auditability.
5. Let the LLM explain the result
After execution, send the row summary back to the LLM and ask for a concise, business-friendly interpretation. This is what makes the tool feel useful to non-technical teams.

When Not to Use This Pattern
Ad-hoc data exploration on massive tables: the LLM will write SELECT * or miss an index. Keep a human in the loop.
Financial reporting with strict audit requirements: use a governed BI tool instead.
Low-latency APIs: LLM inference adds seconds. Pre-generate SQL for known questions or cache results.
Complex domain semantics: if your schema names are cryptic, invest in comments first; retrieval is only as good as the metadata.
Conclusion
This pattern is not about replacing SQL. It is about making the database accessible to people who should not need to learn SQL to get an answer.

The stack is deliberately boring:

PostgreSQL for storage and search
pgvector for vector similarity
OpenAI-compatible embeddings and chat completions
A Python script that ties them together
The result is a database assistant that understands your schema, respects your query style, and answers business questions in seconds. Start with one table and one example query. Add more metadata as you trust it, and your team will stop waiting for engineering to write the next report.

Found this useful? Follow me for more practical AI + database engineering posts.
Enter fullscreen mode Exit fullscreen mode

Top comments (1)

Collapse
 
merbayerp profile image
Mustafa ERBAY

Nice practical walkthrough — especially the idea of treating schema metadata as the retrieval corpus rather than ordinary documents.

One production concern is that a read-only role plus a forbidden-keyword check is not a complete SQL safety boundary. A generated SELECT can still exhaust resources, access unintended tables, or invoke permitted functions. I would add PostgreSQL AST validation, table/function allowlists, read-only transactions, statement and lock timeouts, row limits, and estimated-cost controls.

I would also include foreign keys, constraints, table-level semantics, tenant rules, and sensitive-column classifications in the retrieval corpus. Column-level metadata with top_k=5 may not provide enough context for reliable multi-table joins.

Small implementation detail: Psycopg normally needs pgvector’s type adapter registration rather than assuming an ordinary Python list will always serialize directly as a vector.

Still a strong foundation and a very useful demo — it just needs a stricter distinction between “read-only” and genuinely “safe to execute.”