Capstone Project 1: Building a Self-Healing SQL Agent

In this hands-on capstone project, you will build an Autonomous Text-to-SQL Agent capable of inspecting database schemas, generating SQL queries, validating syntax using dry-run tests, and automatically self-correcting syntax errors before returning verified results to the user.

1. System Architecture

flowchart TD
    UserQ["User Business Question:<br/>'Which 5 enterprise customers spent the most in Q3?'"] --> Agent["SQL Agent Brain"]

    Agent --> CheckSchema["Action: get_schema('fct_orders', 'dim_customers')"]
    CheckSchema --> SchemaObs["Observation: Table DDLs & Column Names"]

    SchemaObs --> GenSQL["Thought: Draft SQL query with JOIN and GROUP BY"]
    GenSQL --> DryRun["Action: dry_run_explain(sql)"]

    DryRun --> Evaluation{Did EXPLAIN succeed?}
    Evaluation -->|Error: Column 'cust_tier' doesn't exist| SelfCorrect["Self-Correction Loop:<br/>Inspect error, rewrite query with 'customer_tier'"]
    SelfCorrect --> DryRun

    Evaluation -->|Success| Exec["Action: execute_query(sql)"]
    Exec --> Out["Final Formatted Table & Business Explanation"]

    style UserQ fill:#334155,color:#fff
    style Agent fill:#2563eb,color:#fff
    style DryRun fill:#d97706,color:#fff
    style SelfCorrect fill:#b91c1c,color:#fff
    style Exec fill:#059669,color:#fff
    style Out fill:#10b981,color:#fff

2. Complete Runnable Implementation

import re
from typing import Dict, Any

# Mock Database Engine
class MockWarehouseEngine:
    def __init__(self):
        self.tables = {
            "dim_customers": {"customer_id": "INT", "name": "VARCHAR", "customer_tier": "VARCHAR"},
            "fct_orders": {"order_id": "INT", "customer_id": "INT", "amount": "NUMERIC", "order_date": "DATE"}
        }

    def get_schema(self, table_name: str) -> str:
        if table_name in self.tables:
            cols = ", ".join([f"{col} {dtype}" for col, dtype in self.tables[table_name].items()])
            return f"Table {table_name}: ({cols})"
        return f"Error: Table '{table_name}' does not exist."

    def validate_sql(self, sql: str) -> Dict[str, Any]:
        # Simulated SQL compiler catching a deliberate schema bug
        if "cust_tier" in sql:
            return {
                "valid": False, 
                "error": "SQLCompilationError: Column 'cust_tier' does not exist in table 'dim_customers'. Did you mean 'customer_tier'?"
            }
        return {"valid": True, "error": None}

    def execute_sql(self, sql: str) -> str:
        return "[('Acme Corp', '$120,000'), ('GlobalTech', '$95,000'), ('Starlight Inc', '$82,000')]"

class SelfHealingSQLAgent:
    def __init__(self, db: MockWarehouseEngine):
        self.db = db

    def run(self, user_question: str) -> str:
        print(f"=== Starting SQL Agent for Goal: '{user_question}' ===")

        # Step 1: Discover schemas
        print("
[Step 1: Inspecting Schema]")
        schema_info = self.db.get_schema("dim_customers") + " | " + self.db.get_schema("fct_orders")
        print(f"Discovered: {schema_info}")

        # Step 2: First attempt (with deliberate typo to demonstrate self-healing)
        attempt_1_sql = (
            "SELECT c.name, SUM(o.amount) FROM fct_orders o "
            "JOIN dim_customers c ON o.customer_id = c.customer_id "
            "WHERE c.cust_tier = 'ENTERPRISE' GROUP BY 1 ORDER BY 2 DESC LIMIT 3;"
        )
        print(f"
[Step 2: Drafted SQL Query]:
{attempt_1_sql}")

        # Step 3: Dry-run validation
        print("
[Step 3: Running Validation / EXPLAIN]")
        val_result = self.db.validate_sql(attempt_1_sql)

        if not val_result["valid"]:
            print(f"⚠️ Validation Failed: {val_result['error']}")
            print("
🔄 [Self-Correction Loop Triggered]")
            # Autonomous correction: replace 'cust_tier' with 'customer_tier'
            corrected_sql = attempt_1_sql.replace("cust_tier", "customer_tier")
            print(f"Corrected Query:
{corrected_sql}")

            # Step 4: Re-validate
            val_retry = self.db.validate_sql(corrected_sql)
            if val_retry["valid"]:
                print("✅ Validation Succeeded!")
                data = self.db.execute_sql(corrected_sql)
                return f"Final Answer (Validated SQL Execution):
{data}"
            else:
                return f"Failed after self-correction: {val_retry['error']}"
        else:
            data = self.db.execute_sql(attempt_1_sql)
            return f"Executed Successfully:
{data}"

if __name__ == "__main__":
    db = MockWarehouseEngine()
    agent = SelfHealingSQLAgent(db)
    final_output = agent.run("Show top spending Enterprise tier customers")
    print(f"
{final_output}")

Key Takeaways

• Dry-run validation (EXPLAIN) catches bugs before executing against production data.
• The Self-Correction Loop turns syntax errors into learning feedback for the next generation turn.
• Strict schema grounding prevents hallucinated column names.