SQL-query's genereren en valideren

import sqlparse

def is_select_only(sql):
    parsed = sqlparse.parse(sql)
    if not parsed:
        return False
    for statement in parsed:
        stmt_type = statement.get_type()
        if stmt_type != 'SELECT':
            print(f'Blocked statement type: {stmt_type}')
            return False
    return True

# Test
print(is_select_only('SELECT * FROM users'))  # True
print(is_select_only('DROP TABLE users'))      # False — Blocked

DANGEROUS_KEYWORDS = [
    'INSERT', 'UPDATE', 'DELETE', 'DROP', 'CREATE',
    'ALTER', 'TRUNCATE', 'GRANT', 'REVOKE', 'EXEC',
    'EXECUTE', 'CALL', 'MERGE'
]

def passes_blocklist(sql):
    sql_upper = sql.upper()
    for keyword in DANGEROUS_KEYWORDS:
        # Check as whole word to avoid false positives like 'CREATED_AT'
        import re
        if re.search(r'\b' + keyword + r'\b', sql_upper):
            raise ValueError(f'Blocked keyword detected: {keyword}')
    return True

def validate_sql(sql):
    if not is_select_only(sql):
        raise ValueError('Only SELECT statements are allowed')
    passes_blocklist(sql)
    return True

import sqlparse
from sqlparse.sql import IdentifierList, Identifier
from sqlparse.tokens import Keyword, DML

def extract_table_names(sql):
    parsed = sqlparse.parse(sql)[0]
    tables = []
    from_seen = False
    for token in parsed.tokens:
        if token.ttype is DML and token.value.upper() == 'SELECT':
            continue
        if token.ttype is Keyword and token.value.upper() in ('FROM', 'JOIN'):
            from_seen = True
            continue
        if from_seen:
            if isinstance(token, Identifier):
                tables.append(token.get_name())
            elif isinstance(token, IdentifierList):
                for item in token.get_identifiers():
                    tables.append(item.get_name())
            from_seen = False
    return tables

print(extract_table_names('SELECT u.name FROM users u JOIN orders o ON u.id = o.user_id'))
# ['users', 'orders']

def validate_tables_exist(sql, known_tables):
    used_tables = extract_table_names(sql)
    invalid = [t for t in used_tables if t and t not in known_tables]
    if invalid:
        raise ValueError(
            f'Query references non-existent tables: {invalid}. '
            f'Available tables: {list(known_tables)[:10]}...'
        )
    return True

# Usage
known = set(build_schema_dict(conn).keys())
try:
    validate_tables_exist(generated_sql, known)
except ValueError as e:
    # Send error back to LLM for correction
    corrected_sql = llm_fix_sql(generated_sql, str(e))
    print('Corrected SQL:', corrected_sql)

import sqlite3

conn = sqlite3.connect(':memory:')
conn.execute('CREATE TABLE orders (status TEXT, user_id INTEGER)')
conn.execute("INSERT INTO orders VALUES ('pending', 42)")

def safe_execute(conn, sql_template, params=()):
    """Execute with parameterized values."""
    cur = conn.cursor()
    cur.execute(sql_template, params)  # driver handles escaping
    columns = [d[0] for d in cur.description]
    rows = cur.fetchmany(200)
    return {'columns': columns, 'rows': rows}

sql = 'SELECT * FROM orders WHERE status = ? AND user_id = ?'
result = safe_execute(conn, sql, params=('pending', 42))
print(result)

def check_explain_plan(conn, sql):
    explain_sql = f'EXPLAIN {sql}'
    with conn.cursor() as cur:
        cur.execute(explain_sql)
        plan = '\n'.join(row[0] for row in cur.fetchall())

    # Check for sequential scans on large tables
    if 'Seq Scan' in plan:
        print('WARNING: Query involves a sequential scan')
        print(plan)
        return {'safe': False, 'plan': plan, 'warning': 'Sequential scan detected'}

    return {'safe': True, 'plan': plan}

# Use before executing
plan_result = check_explain_plan(conn, generated_sql)
if not plan_result['safe']:
    print(f'Optimization hint: {plan_result["warning"]}')

import re

MAX_ROWS = 500

def enforce_row_limit(sql, max_rows=MAX_ROWS):
    sql_upper = sql.upper().rstrip().rstrip(';')

    # Check if LIMIT already present
    if re.search(r'\bLIMIT\b', sql_upper):
        # Extract current limit and enforce maximum
        match = re.search(r'LIMIT\s+(\d+)', sql_upper)
        if match:
            current = int(match.group(1))
            if current > max_rows:
                sql = re.sub(r'LIMIT\s+\d+', f'LIMIT {max_rows}', sql, flags=re.IGNORECASE)
    else:
        sql = sql.rstrip(';') + f' LIMIT {max_rows}'

    return sql

print(enforce_row_limit('SELECT * FROM users'))
# SELECT * FROM users LIMIT 500

import re

CODE_FENCE = chr(96) * 3  # three backticks, built at runtime to avoid template issues

def extract_sql(llm_response):
    # Remove markdown code blocks ('''sql ... ''' or ''' ... ''')
    pattern = CODE_FENCE + r'(?:sql)?\s*([\s\S]+?)' + CODE_FENCE
    match = re.search(pattern, llm_response, re.IGNORECASE)
    if match:
        return match.group(1).strip()

    # If no code block, look for SELECT statement
    match = re.search(r'(SELECT\s+[\s\S]+?;)', llm_response, re.IGNORECASE)
    if match:
        return match.group(1).strip()

    # Fallback: strip common preamble phrases
    cleaned = re.sub(r'^(Here is|The SQL query is|Query:)[^\n]*\n', '',
                     llm_response, flags=re.IGNORECASE).strip()
    return cleaned

if __name__ == '__main__':
    demo_response = 'Here is the SQL query:\n' + CODE_FENCE + 'sql\nSELECT * FROM users;\n' + CODE_FENCE
    print(extract_sql(demo_response))

def validate_and_prepare_sql(llm_output, known_tables, max_rows=500):
    # Step 1: extract raw SQL
    sql = extract_sql(llm_output)
    if not sql:
        raise ValueError('No SQL found in LLM response')

    # Step 2: type check
    if not is_select_only(sql):
        raise ValueError('Only SELECT queries allowed')

    # Step 3: keyword blocklist
    passes_blocklist(sql)

    # Step 4: table existence check
    validate_tables_exist(sql, known_tables)

    # Step 5: row limit
    sql = enforce_row_limit(sql, max_rows)

    return sql

# Full flow
try:
    safe_sql = validate_and_prepare_sql(llm_output, known_tables)
    result = safe_execute(conn, safe_sql)
except ValueError as e:
    corrected = llm_fix_sql(llm_output, str(e))
    safe_sql = validate_and_prepare_sql(corrected, known_tables)
    result = safe_execute(conn, safe_sql)

# Create read-only user in PostgreSQL:
# CREATE USER nl_to_sql_reader WITH PASSWORD 'secure_password';
# GRANT CONNECT ON DATABASE yourdb TO nl_to_sql_reader;
# GRANT USAGE ON SCHEMA public TO nl_to_sql_reader;
# GRANT SELECT ON ALL TABLES IN SCHEMA public TO nl_to_sql_reader;

import os
import psycopg2

def get_readonly_connection():
    return psycopg2.connect(
        host=os.getenv('DB_HOST'),
        database=os.getenv('DB_NAME'),
        user='nl_to_sql_reader',       # read-only account
        password=os.getenv('DB_READER_PASS')
    )