prazy1208/text2sql

1"""Apply scripts/create_system_schema_few_shot_examples.sql. Run from project root."""2 3import os4import sys5from pathlib import Path6 7PROJECT_ROOT = Path(__file__).resolve().parent.parent8if str(PROJECT_ROOT) not in sys.path:9    sys.path.insert(0, str(PROJECT_ROOT))10 11os.chdir(PROJECT_ROOT)12 13from dotenv import load_dotenv14from sqlalchemy import create_engine15 16load_dotenv()17 18 19def get_engine():20    database_url = os.getenv("DATABASE_URL")21    if database_url:22        return create_engine(database_url)23    host = os.getenv("DB_HOST", "localhost")24    port = os.getenv("DB_PORT", "5432")25    user = os.getenv("DB_USER", "postgres")26    password = os.getenv("DB_PASSWORD", "")27    dbname = os.getenv("DB_NAME", "text2sql_db")28    url = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"29    return create_engine(url)30 31 32def main():33    sql_file = PROJECT_ROOT / "scripts" / "create_system_schema_few_shot_examples.sql"34    if not sql_file.exists():35        print(f"SQL file not found: {sql_file}")36        sys.exit(1)37    sql = sql_file.read_text(encoding="utf-8")38    engine = get_engine()39    raw_conn = engine.raw_connection()40    try:41        cur = raw_conn.cursor()42        cur.execute(sql)43        raw_conn.commit()44    finally:45        raw_conn.close()46    print("Done. system_schema.few_shot_examples created.")47 48 49if __name__ == "__main__":50    main()51