schema_reader.py

# schema_reader.py
import psycopg2
from config import db_config as default_db_config

def get_schema_hint(db_config=None):
    # 如果没有提供配置，使用默认配置
    if db_config is None:
        db_config = default_db_config
        
    conn = psycopg2.connect(**db_config)
    cur = conn.cursor()
    hint = ""
    try:
        cur.execute("""
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = 'public'
        """)
        tables = cur.fetchall()
        for table in tables:
            table_name = table[0]
            cur.execute(f"""
                SELECT column_name, data_type
                FROM information_schema.columns
                WHERE table_name = '{table_name}'
            """)
            columns = cur.fetchall()
            columns_str = ", ".join([f"{col[0]} ({col[1]})" for col in columns])
            hint += f"表 {table_name}：{columns_str}\n"
    except Exception as e:
        hint = f"获取schema失败：{str(e)}"
    finally:
        cur.close()
        conn.close()
    return hint
    return hint