Team Ai
Apppublic

prazy1208/text2sql

sourceHugging Faceupdated 5mo agoView on Hugging Face
0likes
run_full_schema_refresh.py115 linesDownload Raw Back to scripts
1"""2Run full domain schema refresh in one command.3 4Pipeline order:51) Apply scripts/create_domain_schemas.sql62) Ensure domain table_relationships tables exist73) Extract + upsert FK relationships into domain table_relationships84) Build relationship embedding metadata JSON95) Build table/column metadata JSON + FAISS indexes10 11Run from project root: python scripts/run_full_schema_refresh.py12"""13 14from __future__ import annotations15 16import os17import subprocess18import sys19import time20from pathlib import Path21 22from dotenv import load_dotenv23from sqlalchemy import create_engine24 25PROJECT_ROOT = Path(__file__).resolve().parent.parent26if str(PROJECT_ROOT) not in sys.path:27    sys.path.insert(0, str(PROJECT_ROOT))28 29os.chdir(PROJECT_ROOT)30load_dotenv()31 32 33def get_engine():34    """Build SQLAlchemy engine from .env (DATABASE_URL or DB_* variables)."""35    database_url = os.getenv("DATABASE_URL")36    if database_url:37        return create_engine(database_url)38    host = os.getenv("DB_HOST", "localhost")39    port = os.getenv("DB_PORT", "5432")40    user = os.getenv("DB_USER", "postgres")41    password = os.getenv("DB_PASSWORD", "")42    dbname = os.getenv("DB_NAME", "text2sql_db")43    url = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"44    return create_engine(url)45 46 47def apply_domain_schema_sql() -> None:48    """Execute scripts/create_domain_schemas.sql (multi-statement SQL)."""49    sql_file = PROJECT_ROOT / "scripts" / "create_domain_schemas.sql"50    if not sql_file.exists():51        raise FileNotFoundError(f"SQL file not found: {sql_file}")52 53    sql = sql_file.read_text(encoding="utf-8")54    engine = get_engine()55    raw_conn = engine.raw_connection()56    try:57        cur = raw_conn.cursor()58        cur.execute(sql)59        raw_conn.commit()60    finally:61        raw_conn.close()62 63 64def run_python_script(relative_path: str) -> None:65    """Run a Python script and raise if it fails."""66    script_path = PROJECT_ROOT / relative_path67    if not script_path.exists():68        raise FileNotFoundError(f"Script not found: {script_path}")69    cmd = [sys.executable, str(script_path)]70    subprocess.run(cmd, cwd=str(PROJECT_ROOT), check=True)71 72 73def run_step(step_name: str, fn) -> None:74    start = time.perf_counter()75    print(f"[START] {step_name}")76    fn()77    elapsed = time.perf_counter() - start78    print(f"[OK] {step_name} ({elapsed:.2f}s)")79 80 81def main() -> None:82    print("Starting full schema refresh pipeline...")83    steps: list[tuple[str, callable]] = [84        ("Apply domain schema SQL", apply_domain_schema_sql),85        (86            "Ensure table_relationships tables",87            lambda: run_python_script("scripts/run_create_domain_schema_table_relationships.py"),88        ),89        (90            "Extract and load relationships",91            lambda: run_python_script("scripts/extract_and_load_relationships.py"),92        ),93        (94            "Build relationship embeddings metadata",95            lambda: run_python_script("build_relationship_embeddings.py"),96        ),97        (98            "Build metadata and FAISS indexes",99            lambda: run_python_script("build_vector_store.py"),100        ),101    ]102 103    try:104        for step_name, fn in steps:105            run_step(step_name, fn)106    except Exception as exc:107        print(f"[FAILED] Pipeline stopped: {exc}")108        raise109 110    print("Full schema refresh completed successfully.")111 112 113if __name__ == "__main__":114    main()115