import psycopg2
def get_schema(conn):
"""Retourne le schéma sous forme de CREATE TABLE simplifiés."""
query = """
SELECT table_name, column_name, data_type
FROM information_schema.columns
WHERE table_schema = 'public'
ORDER BY table_name, ordinal_position;
"""
tables = {}
with conn.cursor() as cur:
cur.execute(query)
for table, col, dtype in cur.fetchall():
tables.setdefault(table, []).append(f"{col} {dtype}")
lines = []
for table, cols in tables.items():
cols_str = ", ".join(cols)
lines.append(f"TABLE {table} ({cols_str});")
return "\n".join(lines)
import re
import ollama
import psycopg2
import sqlparse
MODEL = "qwen3-coder:30b"
SYSTEM_PROMPT = """Tu es un expert PostgreSQL. Génère UNE seule requête SQL
qui répond à la question de l'utilisateur, en respectant ces règles :
- Uniquement des requêtes SELECT (jamais INSERT/UPDATE/DELETE/DROP).
- Utilise exactement les noms de tables et colonnes du schéma fourni.
- Ajoute toujours LIMIT 100 si la question ne précise pas de limite.
- Réponds UNIQUEMENT avec le SQL, sans explication ni balise Markdown.
Schéma de la base :
{schema}"""
def generate_sql(question, schema):
resp = ollama.chat(
model=MODEL,
messages=[
{"role": "system", "content": SYSTEM_PROMPT.format(schema=schema)},
{"role": "user", "content": question},
],
options={"temperature": 0}, # déterminisme : crucial pour du SQL
)
return clean_sql(resp["message"]["content"])
def clean_sql(raw):
# Retire les fences Markdown ```sql ... ``` si le modèle en ajoute
raw = re.sub(r"```(?:sql)?", "", raw).strip()
return raw.rstrip(";") + ";"
def is_read_only(sql):
statements = sqlparse.parse(sql)
if len(statements) != 1:
return False # une seule requête, pas d'empilement
stmt = statements[0]
if stmt.get_type() != "SELECT":
return False
forbidden = ("insert", "update", "delete", "drop",
"alter", "truncate", "grant", "create")
lowered = sql.lower()
return not any(kw in lowered for kw in forbidden)
def run_query(sql):
if not is_read_only(sql):
raise ValueError(f"Requête refusée (non lecture seule) : {sql}")
# Rôle 'readonly' : ne dispose QUE du privilège SELECT côté base
conn = psycopg2.connect(
dbname="analytics", user="readonly",
password="...", host="localhost",
)
with conn.cursor() as cur:
cur.execute("SET statement_timeout = '5s';") # anti-requête folle
cur.execute(sql)
cols = [d[0] for d in cur.description]
rows = cur.fetchall()
conn.close()
return cols, rows
if __name__ == "__main__":
from schema import get_schema
ro = psycopg2.connect(dbname="analytics", user="readonly",
password="...", host="localhost")
schema = get_schema(ro)
question = "Combien de commandes par mois en 2025 ?"
sql = generate_sql(question, schema)
print("SQL généré :", sql)
cols, rows = run_query(sql)
print(cols)
for r in rows:
print(r)
SET statement_timeout 可防止错误的查询（如在数百万行上产生笛卡尔积）导致数据库资源耗尽。
04
强制使用 LIMIT
在提示和代码侧均设置 LIMIT，以确保从不将整张表加载到内存中。
05
纠错循环
若执行返回SQL错误，请将错误信息返回给模型，并请求修正后的查询（最多尝试1或2次）。
PostgreSQL 只读角色
-- À exécuter une fois par un admin
CREATE ROLE readonly WITH LOGIN PASSWORD '...';
GRANT CONNECT ON DATABASE analytics TO readonly;
GRANT USAGE ON SCHEMA public TO readonly;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO readonly;
-- Les tables créées plus tard héritent aussi du SELECT
ALTER DEFAULT PRIVILEGES IN SCHEMA public
GRANT SELECT ON TABLES TO readonly;