AI Agents · Lesson

Schema Understanding and Injection

Extracting and formatting DB schema for LLM context: tables, columns, relations.

Schema Understanding and Injection is a free AI Agents lesson on CoddyKit — lesson 2 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the AI Agents learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

Why Schema Context Matters

The LLM knows SQL syntax but knows nothing about your database. Without schema context, it will hallucinate table and column names.

Schema injection means programmatically extracting your DB structure and including it in every prompt — making the LLM aware of your exact tables, columns, and types.

Querying INFORMATION_SCHEMA

All major relational databases expose metadata through INFORMATION_SCHEMA. You can query it to get every table, column name, and data type without touching application code.

This works in PostgreSQL, MySQL, SQL Server, and SQLite (with minor differences).

import psycopg2

def get_schema(conn):
    query = '''
        SELECT table_name, column_name, data_type
        FROM information_schema.columns
        WHERE table_schema = 'public'
        ORDER BY table_name, ordinal_position
    '''
    with conn.cursor() as cur:
        cur.execute(query)
        return cur.fetchall()

Grouping Columns by Table

The raw INFORMATION_SCHEMA result is a flat list of rows. Group them by table name to build a structured representation that is easier to format into a prompt.

from collections import defaultdict

def build_schema_dict(conn):
    rows = get_schema(conn)
    schema = defaultdict(list)
    for table_name, column_name, data_type in rows:
        schema[table_name].append({
            'name': column_name,
            'type': data_type
        })
    return dict(schema)

# Result:
# {
#   'users': [{'name': 'id', 'type': 'integer'}, {'name': 'email', 'type': 'character varying'}],
#   'orders': [{'name': 'id', 'type': 'integer'}, {'name': 'user_id', 'type': 'integer'}]
# }

Formatting Schema for LLM Prompts

The LLM reads the schema as plain text. Use a concise, readable format: one table per line with column names and types in parentheses.

Including primary keys (PK) and foreign keys (FK) helps the LLM write correct JOIN statements.

def format_schema_for_prompt(schema_dict, pk_info=None, fk_info=None):
    lines = []
    for table, columns in schema_dict.items():
        col_parts = []
        for col in columns:
            label = col['name']
            if pk_info and (table, col['name']) in pk_info:
                label += ' PK'
            if fk_info and (table, col['name']) in fk_info:
                label += f' FK->{fk_info[(table, col["name"])]}'
            col_parts.append(f"{label} ({col['type']})")
        lines.append(f"Table {table}: {', '.join(col_parts)}")
    return '\n'.join(lines)

# Output:
# Table users: id PK (integer), email (varchar), created_at (timestamp)
# Table orders: id PK (integer), user_id FK->users.id (integer), total (float)

if __name__ == '__main__':
    demo_schema = {'users': [{'name': 'id', 'type': 'integer'}, {'name': 'email', 'type': 'varchar'}]}
    demo_pk = {('users', 'id')}
    print(format_schema_for_prompt(demo_schema, pk_info=demo_pk))

Including Primary and Foreign Keys

Foreign key relationships are the most important part of schema context — they tell the LLM how to write JOINs. Query information_schema.table_constraints and key_column_usage to extract them.

def get_foreign_keys(conn):
    query = '''
        SELECT
            kcu.table_name,
            kcu.column_name,
            ccu.table_name AS foreign_table,
            ccu.column_name AS foreign_column
        FROM information_schema.table_constraints AS tc
        JOIN information_schema.key_column_usage AS kcu
            ON tc.constraint_name = kcu.constraint_name
        JOIN information_schema.constraint_column_usage AS ccu
            ON ccu.constraint_name = tc.constraint_name
        WHERE tc.constraint_type = 'FOREIGN KEY'
    '''
    with conn.cursor() as cur:
        cur.execute(query)
        return {
            (row[0], row[1]): f'{row[2]}.{row[3]}'
            for row in cur.fetchall()
        }

if __name__ == '__main__':
    class FakeCursor:
        def __enter__(self): return self
        def __exit__(self, *a): return False
        def execute(self, query): pass
        def fetchall(self):
            return [('orders', 'user_id', 'users', 'id')]
    class FakeConn:
        def cursor(self): return FakeCursor()

    fks = get_foreign_keys(FakeConn())
    print('Foreign keys found:')
    for (table, col), ref in fks.items():
        print(f'  {table}.{col} -> {ref}')

Schema Compression: The Problem

A real enterprise database may have 200+ tables. If you inject the full schema, you will exceed GPT-4's context window and waste money on tokens.

A 200-table schema with 20 columns each is roughly 40,000+ tokens — too expensive to send on every query.

def estimate_schema_tokens(schema_dict):
    text = format_schema_for_prompt(schema_dict)
    # Rough estimate: 1 token per 4 characters
    estimated_tokens = len(text) // 4
    print(f'Tables: {len(schema_dict)}')
    print(f'Estimated schema tokens: {estimated_tokens}')
    return estimated_tokens

# 200 tables * 15 columns * 25 chars/col = 75,000 chars = ~18,750 tokens
# Plus user question + system prompt = easily over context limit

Schema Compression: Selective Injection

The most effective compression strategy: only inject tables relevant to the question. Use a two-phase approach — first ask the LLM which tables it needs, then inject only those schemas.

def select_relevant_tables(question, all_table_names, n=5):
    table_list = ', '.join(all_table_names)
    prompt = f'''Database tables: {table_list}

Question: {question}

List the {n} most relevant table names as a JSON array.
Example: ["users", "orders", "products"]'''

    response = llm_call(prompt)
    import json
    return json.loads(response)

def compressed_schema(question, conn):
    all_tables = list(build_schema_dict(conn).keys())
    relevant = select_relevant_tables(question, all_tables)
    full_schema = build_schema_dict(conn)
    return {t: full_schema[t] for t in relevant if t in full_schema}

Schema Compression: Excluding Noise Columns

Many tables have audit columns like created_at, updated_at, deleted_at, version, created_by that are rarely relevant to business queries. Strip them to reduce token count.

AUDIT_COLUMNS = {
    'created_at', 'updated_at', 'deleted_at', 'created_by',
    'updated_by', 'version', 'is_deleted', 'modified_at'
}

def compress_schema(schema_dict, exclude_audit=True):
    compressed = {}
    for table, columns in schema_dict.items():
        # Skip internal/system tables
        if table.startswith('_') or table.startswith('pg_'):
            continue
        if exclude_audit:
            columns = [c for c in columns if c['name'] not in AUDIT_COLUMNS]
        if columns:  # only include if columns remain
            compressed[table] = columns
    return compressed

if __name__ == '__main__':
    demo_schema = {
        'users': [{'name': 'id', 'type': 'INT'}, {'name': 'email', 'type': 'VARCHAR'}, {'name': 'created_at', 'type': 'TIMESTAMP'}],
        'pg_stat': [{'name': 'x', 'type': 'INT'}],
    }
    compressed = compress_schema(demo_schema)
    print('Tables kept:', list(compressed.keys()))
    print('users columns after compression:', [c['name'] for c in compressed['users']])

Adding Table Descriptions

Column names alone are not always self-explanatory. Adding natural language descriptions of what each table represents dramatically improves SQL generation quality.

Store descriptions in a config file or as PostgreSQL table comments.

TABLE_DESCRIPTIONS = {
    'users': 'Registered app users with authentication info',
    'orders': 'Customer purchase orders',
    'order_items': 'Individual line items within an order',
    'products': 'Product catalog with pricing',
    'payments': 'Payment transactions linked to orders'
}

def format_schema_with_descriptions(schema_dict):
    lines = []
    for table, columns in schema_dict.items():
        desc = TABLE_DESCRIPTIONS.get(table, '')
        col_str = ', '.join(f"{c['name']} ({c['type']})" for c in columns)
        if desc:
            lines.append(f"Table {table} ({desc}): {col_str}")
        else:
            lines.append(f"Table {table}: {col_str}")
    return '\n'.join(lines)

if __name__ == '__main__':
    demo_schema = {'users': [{'name': 'id', 'type': 'INT'}], 'orders': [{'name': 'id', 'type': 'INT'}]}
    print(format_schema_with_descriptions(demo_schema))

Caching the Schema

Database schemas rarely change. Fetching INFORMATION_SCHEMA on every query adds latency and load. Cache the formatted schema string and invalidate it on schema change events or on a time-based TTL.

import time

class SchemaCache:
    def __init__(self, ttl_seconds=300):
        self._cache = None
        self._timestamp = 0
        self.ttl = ttl_seconds

    def get(self, conn):
        now = time.time()
        if self._cache is None or (now - self._timestamp) > self.ttl:
            print('Refreshing schema cache...')
            schema_dict = build_schema_dict(conn)
            fk_info = get_foreign_keys(conn)
            self._cache = format_schema_for_prompt(schema_dict, fk_info=fk_info)
            self._timestamp = now
        return self._cache

schema_cache = SchemaCache(ttl_seconds=300)

Full Schema Injection Flow

Combining all techniques: cache the compressed schema, inject it into the system prompt, and use selective table filtering for large databases.

def build_sql_agent_prompt(question, conn, large_db=False):
    if large_db:
        schema = compressed_schema(question, conn)
        schema_text = format_schema_with_descriptions(schema)
    else:
        schema_text = schema_cache.get(conn)

    system = f'''You are a PostgreSQL expert.
Return ONLY a valid SELECT query based on this schema:

{schema_text}

Rules:
- Use only SELECT statements
- Use table aliases for clarity
- Limit results to 100 rows unless asked for all
'''
    return system

Knowledge Check

When should you use selective table injection instead of injecting the full schema?

Recap: Schema Understanding and Injection

Effective schema injection is the foundation of reliable NL-to-SQL agents. Extract structure from INFORMATION_SCHEMA, include primary and foreign key relationships, and format as compact text for the LLM.

For large databases: cache the schema, strip audit columns, and use selective injection to only send tables relevant to each question. Table descriptions in natural language further improve query quality.

Frequently asked questions

Is the “Schema Understanding and Injection” lesson free?

Yes — the full text of “Schema Understanding and Injection” is free to read here on the web, and the AI Agents course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the AI Agents course, upgrade to CoddyKit PRO.

What will I learn in “Schema Understanding and Injection”?

Extracting and formatting DB schema for LLM context: tables, columns, relations. You practise AI Agents with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start AI Agents?

No prior experience is required. AI Agents on CoddyKit is structured for beginners through advanced learners; this is — lesson 2 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Schema Understanding and Injection” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this AI Agents lesson?

Yes. Every AI Agents lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

← Back to AI Agents