DB용 SQL 어시스턴트

프로젝트 목표

아키텍처

Step 1: Schema Tools

import psycopg

conn = psycopg.connect(DATABASE_URL)

def list_tables():
    with conn.cursor() as cur:
        cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema='public'")
        return [r[0] for r in cur.fetchall()]

def describe_table(name):
    with conn.cursor() as cur:
        cur.execute('''
            SELECT column_name, data_type FROM information_schema.columns
            WHERE table_name = %s
        ''', (name,))
        return cur.fetchall()

2단계: SAFE run_sql 도구

import re

def run_sql(query: str, limit: int = 100):
    q = query.strip().rstrip(';').lower()
    if not q.startswith('select'):
        return {'error': 'Only SELECT statements are allowed.'}
    if any(bad in q for bad in [' drop ', ' delete ', ' update ', ' insert ', ' alter ', ' truncate ']):
        return {'error': 'Statement contains a disallowed keyword.'}
    with conn.cursor() as cur:
        cur.execute(f'SET statement_timeout = 5000')   # 5 seconds
        cur.execute(f'SELECT * FROM ({query}) sub LIMIT {limit}')
        cols = [c.name for c in cur.description]
        rows = cur.fetchall()
        return {'columns': cols, 'rows': rows}

Tool Definitions

tools = [
    {'type': 'function', 'function': {'name': 'list_tables', 'description': 'List tables in the database', 'parameters': {'type': 'object', 'properties': {}}}},
    {'type': 'function', 'function': {'name': 'describe_table', 'description': 'Get columns of a table', 'parameters': {'type': 'object', 'properties': {'name': {'type': 'string'}}, 'required': ['name']}}},
    {'type': 'function', 'function': {'name': 'run_sql', 'description': 'Execute a SELECT query (read-only, max 100 rows, 5s timeout)', 'parameters': {'type': 'object', 'properties': {'query': {'type': 'string'}}, 'required': ['query']}}}
]
import json
print(json.dumps(tools, indent=2))

System Prompt

system = '''
You are a SQL analyst assistant for a Postgres database.
First use list_tables and describe_table to learn the schema.
Then write a single SELECT query to answer the user.
Never modify data.
After receiving results, explain them in plain language.
'''
print(system.strip())

시스템 프롬프트에 스키마 전달하기

schema = ''
for t in list_tables():
    cols = describe_table(t)
    schema += f'{t}: {cols}\n'
system = system + f'\nSchema:\n{schema}'

읽기 전용 데이터베이스 사용자

PII 마스킹

PII_COLS = {'email', 'phone'}
for row in rows:
    for i, col in enumerate(cols):
        if col in PII_COLS:
            row[i] = '[REDACTED]'

query 비용 추정

대화 예시

차트 출력

우아하게 실패 처리하기

모든 query 감사하기

SELECT로 제한하는 이유

정리