Level 0: Code generation only (AI never connects to data)
Level 1: Metadata access (schemas, row counts, column stats)
Level 2: Sampled data access (AI sees representative samples, never full tables)
Level 3: Guarded query execution (read-only, row limits, cost caps, human approval)

# Level 0: AI generates dbt model SQL based on requirements
# No data connection needed — works in Tier 3 (air-gapped)

prompt = """
Generate a dbt staging model for Stripe charges with these columns:
- charge_id (string, primary key)
- customer_id (string, foreign key)
- amount_cents (integer)
- currency (string, 3-letter ISO)
- status (string: succeeded, pending, failed)
- created_at (timestamp)

Include:
- Rename from source column conventions (camelCase → snake_case)
- Cast amount_cents to decimal dollars
- Filter to only succeeded charges
- Add unique + not_null tests on charge_id
"""

# AI generates complete dbt model + schema.yml — no data access required

# Level 1: AI reads INFORMATION_SCHEMA to understand actual table structure
# Connection is read-only, limited to metadata views only

METADATA_QUERY = """
SELECT
    table_schema,
    table_name,
    column_name,
    data_type,
    is_nullable,
    character_maximum_length
FROM information_schema.columns
WHERE table_schema IN ('RAW', 'STAGING')
ORDER BY table_schema, table_name, ordinal_position
"""

# AI uses real schema info to generate accurate SQL
# Never sees actual row data — only metadata

# Level 2: AI reads a small sample to understand data patterns
# Enforced via LIMIT and TABLESAMPLE

SAMPLE_QUERY = """
SELECT *
FROM raw.orders
TABLESAMPLE (10 ROWS)  -- Or: LIMIT 10
"""

# AI sees 10 real rows to understand:
# - Actual value formats (dates, IDs, enums)
# - Null patterns
# - Data quality issues
# Then generates more accurate transforms

# Level 3: AI runs NL-to-SQL with guardrails
# Only for Tier 1 (cloud-native) environments with explicit approval

GUARDRAILS = {
    "max_rows": 1000,           # Row limit on all queries
    "max_cost_credits": 1.0,    # Snowflake credit cap
    "timeout_seconds": 30,      # Query timeout
    "allowed_schemas": ["STAGING", "MARTS"],  # Whitelist
    "blocked_tables": ["PII_*", "AUDIT_*"],   # Blacklist
    "read_only": True,          # No DDL/DML
    "require_where_clause": True,  # Prevent full table scans
}

from mcp.server import Server
from mcp.types import Tool, TextContent
import os

# ── Credential boundary ──────────────────────────────────────────────
# Configure: export SNOWFLAKE_ACCOUNT, SNOWFLAKE_USER, SNOWFLAKE_ROLE
# Configure: export SNOWFLAKE_PRIVATE_KEY_PATH
# The MCP server connects to the warehouse — credentials must be scoped.
# See: shared-references/data-engineering/security-compliance-patterns.md
# ─────────────────────────────────────────────────────────────────────

app = Server("warehouse-mcp")

@app.tool()
async def list_schemas() -> list[TextContent]:
    """List available database schemas."""
    schemas = execute_query(
        "SELECT schema_name FROM information_schema.schemata ORDER BY schema_name"
    )
    return [TextContent(type="text", text=format_results(schemas))]

@app.tool()
async def describe_table(schema_name: str, table_name: str) -> list[TextContent]:
    """Describe columns and types for a specific table."""
    # Validate input to prevent SQL injection
    if not schema_name.isalnum() or not table_name.isalnum():
        return [TextContent(type="text", text="Error: Invalid schema or table name")]

    columns = execute_query(f"""
        SELECT column_name, data_type, is_nullable, comment
        FROM information_schema.columns
        WHERE table_schema = '{schema_name.upper()}'
          AND table_name = '{table_name.upper()}'
        ORDER BY ordinal_position
    """)
    return [TextContent(type="text", text=format_results(columns))]

@app.tool()
async def query_sample(
    schema_name: str,
    table_name: str,
    limit: int = 10,
) -> list[TextContent]:
    """Return a small sample from a table (max 100 rows)."""
    # Enforce guardrails
    limit = min(limit, 100)

    if not schema_name.isalnum() or not table_name.isalnum():
        return [TextContent(type="text", text="Error: Invalid schema or table name")]

    rows = execute_query(
        f"SELECT * FROM {schema_name}.{table_name} LIMIT {limit}"
    )
    return [TextContent(type="text", text=format_results(rows))]

@app.tool()
async def run_read_only_query(sql: str) -> list[TextContent]:
    """
    Execute a read-only SQL query with guardrails.

    Tier 1 only. Query is validated before execution:
    - Must be SELECT (no DDL/DML)
    - Must include LIMIT clause (max 1000)
    - Must not reference blocked schemas/tables
    - Has 30-second timeout
    """
    validation = validate_query(sql)
    if not validation.is_valid:
        return [TextContent(type="text", text=f"Query rejected: {validation.reason}")]

    try:
        results = execute_query(sql, timeout=30)
        return [TextContent(type="text", text=format_results(results))]
    except Exception as e:
        return [TextContent(type="text", text=f"Query error: {str(e)}")]

import anthropic
import os

# ── Credential boundary ──────────────────────────────────────────────
# Configure: export ANTHROPIC_API_KEY="sk-ant-xxx"
# Configure: export SNOWFLAKE_ACCOUNT (for metadata queries)
# ─────────────────────────────────────────────────────────────────────

def build_schema_context(tables: list[str]) -> str:
    """Build schema context string for NL-to-SQL prompt."""
    context_parts = []
    for table in tables:
        columns = get_table_columns(table)  # From INFORMATION_SCHEMA
        col_descriptions = []
        for col in columns:
            desc = f"  - {col['name']} ({col['type']})"
            if col.get("comment"):
                desc += f" -- {col['comment']}"
            col_descriptions.append(desc)

        context_parts.append(
            f"Table: {table}\n"
            f"Columns:\n" + "\n".join(col_descriptions)
        )

    return "\n\n".join(context_parts)

def generate_sql(question: str, schema_context: str) -> str:
    """Generate SQL from natural language question."""
    client = anthropic.Anthropic()

    response = client.messages.create(
        model="claude-sonnet-4-5-20250929",
        max_tokens=1024,
        system="""You are a SQL expert. Generate a single SQL query that answers the user's question.

Rules:
- Use only the tables and columns provided in the schema context
- Always include a LIMIT clause (default 100, max 1000)
- Use explicit column names (never SELECT *)
- Add comments explaining complex logic
- Return ONLY the SQL query, no explanation""",
        messages=[{
            "role": "user",
            "content": f"""Schema context:
{schema_context}

Question: {question}

Generate SQL:"""
        }],
    )

    return response.content[0].text.strip()

import sqlglot

def validate_generated_sql(
    sql: str,
    allowed_schemas: list[str],
    max_limit: int = 1000,
) -> tuple[bool, str]:
    """Validate AI-generated SQL meets guardrails."""
    try:
        parsed = sqlglot.parse_one(sql)
    except sqlglot.errors.ParseError as e:
        return False, f"Invalid SQL syntax: {e}"

    # Must be a SELECT statement
    if parsed.key != "select":
        return False, "Only SELECT queries are allowed"

    # Must have LIMIT clause
    limit_node = parsed.find(sqlglot.exp.Limit)
    if not limit_node:
        return False, "Query must include a LIMIT clause"

    limit_value = int(limit_node.expression.this)
    if limit_value > max_limit:
        return False, f"LIMIT {limit_value} exceeds maximum of {max_limit}"

    # Check table references against allowlist
    for table in parsed.find_all(sqlglot.exp.Table):
        schema = (table.db or "").upper()
        if schema and schema not in [s.upper() for s in allowed_schemas]:
            return False, f"Schema '{schema}' is not in allowed list"

    # No subqueries that bypass limits
    subqueries = list(parsed.find_all(sqlglot.exp.Subquery))
    if len(subqueries) > 3:
        return False, "Too many subqueries (max 3)"

    return True, "Valid"

import anthropic
from pydantic import BaseModel
from typing import Literal
import json

class ClassificationResult(BaseModel):
    id: str
    category: Literal["bug", "feature", "question", "docs", "other"]
    confidence: float

def classify_batch(
    records: list[dict],
    batch_size: int = 50,
) -> list[ClassificationResult]:
    """Classify records in batches for cost efficiency."""
    client = anthropic.Anthropic()
    results = []

    for i in range(0, len(records), batch_size):
        batch = records[i:i + batch_size]

        # Format batch for single LLM call
        batch_text = "\n".join(
            f"[{r['id']}] {r['title']}: {r['description'][:200]}"
            for r in batch
        )

        response = client.messages.create(
            model="claude-haiku-4-5-20251001",  # Use cheapest model for classification
            max_tokens=2048,
            system="""Classify each item into exactly one category: bug, feature, question, docs, other.
Return JSON array: [{"id": "...", "category": "...", "confidence": 0.0-1.0}]""",
            messages=[{"role": "user", "content": batch_text}],
        )

        batch_results = json.loads(response.content[0].text)
        results.extend(
            ClassificationResult.model_validate(r) for r in batch_results
        )

    return results

import polars as pl

def enrich_with_llm(
    df: pl.DataFrame,
    text_column: str,
    output_column: str,
    prompt_template: str,
    cache_path: str = "llm_cache.parquet",
) -> pl.DataFrame:
    """
    LLM enrichment with caching to minimize API costs.

    1. Check cache for previously enriched records
    2. Only call LLM for new/uncached records
    3. Batch LLM calls for efficiency
    4. Save results to cache
    """
    # Load cache
    try:
        cache = pl.read_parquet(cache_path)
        cached_keys = set(cache[text_column].to_list())
    except FileNotFoundError:
        cache = pl.DataFrame()
        cached_keys = set()

    # Filter to uncached records
    uncached = df.filter(~pl.col(text_column).is_in(list(cached_keys)))

    if uncached.height == 0:
        # All cached — just join
        return df.join(cache.select(text_column, output_column), on=text_column, how="left")

    # Call LLM for uncached records (batched)
    new_results = classify_batch(uncached.to_dicts())

    # Update cache
    new_cache_rows = pl.DataFrame([
        {text_column: r.id, output_column: r.category}
        for r in new_results
    ])
    updated_cache = pl.concat([cache, new_cache_rows]) if cache.height > 0 else new_cache_rows
    updated_cache.write_parquet(cache_path)

    # Join enrichment back
    return df.join(updated_cache.select(text_column, output_column), on=text_column, how="left")

import anthropic
import numpy as np
from typing import Optional

# ── Credential boundary ──────────────────────────────────────────────
# Configure: export ANTHROPIC_API_KEY="sk-ant-xxx" (or use Voyage/OpenAI)
# Configure: export VECTOR_DB_URL="..." (if using external vector store)
# ─────────────────────────────────────────────────────────────────────

def embed_table_metadata(tables: list[dict]) -> list[dict]:
    """Create embeddings for table metadata (schema + descriptions)."""
    client = anthropic.Anthropic()
    results = []

    for table in tables:
        # Build rich text representation of table
        text = (
            f"Table: {table['schema']}.{table['name']}\n"
            f"Description: {table.get('description', 'No description')}\n"
            f"Columns: {', '.join(c['name'] for c in table['columns'])}\n"
            f"Column details: {'; '.join(f'{c[\"name\"]} ({c[\"type\"]}): {c.get(\"description\", \"\")}' for c in table['columns'])}"
        )

        # Generate embedding (use your preferred provider)
        embedding = generate_embedding(text)

        results.append({
            "table_fqn": f"{table['schema']}.{table['name']}",
            "text": text,
            "embedding": embedding,
        })

    return results

def search_tables(query: str, table_embeddings: list[dict], top_k: int = 5) -> list[dict]:
    """Find tables most relevant to a natural language query."""
    query_embedding = generate_embedding(query)

    # Cosine similarity
    scored = []
    for te in table_embeddings:
        similarity = np.dot(query_embedding, te["embedding"]) / (
            np.linalg.norm(query_embedding) * np.linalg.norm(te["embedding"])
        )
        scored.append({**te, "similarity": similarity})

    return sorted(scored, key=lambda x: x["similarity"], reverse=True)[:top_k]

-- Snowflake: Create audit table for AI-generated queries
CREATE TABLE IF NOT EXISTS audit.ai_query_log (
    query_id STRING DEFAULT UUID_STRING(),
    agent_id STRING,           -- Which AI agent/MCP server
    prompt_hash STRING,        -- SHA-256 of the originating prompt
    generated_sql STRING,      -- The SQL that was generated
    execution_status STRING,   -- 'approved', 'rejected', 'executed', 'failed'
    result_row_count INTEGER,
    execution_time_ms INTEGER,
    warehouse_credits_used FLOAT,
    reviewed_by STRING,        -- Human reviewer (Tier 2)
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP()
);