prazy1208/text2sql

1"""2Backfill app_schema.sessions: titles from first user message; optional client_id.3 4From project root:5  python scripts/run_backfill_session_metadata.py6  python scripts/run_backfill_session_metadata.py --client-id "<uuid from text2sql_client_id>"7 8Or set BACKFILL_CLIENT_ID in .env (used when --client-id is omitted).9"""10 11import argparse12import os13import sys14import uuid15from pathlib import Path16 17PROJECT_ROOT = Path(__file__).resolve().parent.parent18if str(PROJECT_ROOT) not in sys.path:19    sys.path.insert(0, str(PROJECT_ROOT))20 21os.chdir(PROJECT_ROOT)22 23from dotenv import load_dotenv24from sqlalchemy import create_engine, text25 26load_dotenv()27 28from backend.config import APP_SCHEMA29 30 31def get_engine():32    database_url = os.getenv("DATABASE_URL")33    if database_url:34        return create_engine(database_url)35    host = os.getenv("DB_HOST", "localhost")36    port = os.getenv("DB_PORT", "5432")37    user = os.getenv("DB_USER", "postgres")38    password = os.getenv("DB_PASSWORD", "")39    dbname = os.getenv("DB_NAME", "text2sql_db")40    url = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"41    return create_engine(url)42 43 44def run_title_backfill(engine) -> None:45    sql_file = PROJECT_ROOT / "scripts" / "backfill_session_metadata.sql"46    if not sql_file.exists():47        print(f"SQL file not found: {sql_file}")48        sys.exit(1)49    sql = sql_file.read_text(encoding="utf-8")50    raw_conn = engine.raw_connection()51    try:52        cur = raw_conn.cursor()53        cur.execute(sql)54        raw_conn.commit()55    finally:56        raw_conn.close()57 58 59def run_client_id_backfill(engine, client_id: str) -> int:60    """Assign client_id to sessions that have messages but NULL client_id. Returns rowcount."""61    cid = uuid.UUID(str(client_id).strip())62    with engine.begin() as conn:63        result = conn.execute(64            text(f"""65                UPDATE {APP_SCHEMA}.sessions s66                SET client_id = :cid,67                    updated_at = CURRENT_TIMESTAMP68                WHERE s.client_id IS NULL69                  AND EXISTS (70                      SELECT 1 FROM {APP_SCHEMA}.chat_messages m71                      WHERE m.session_id = s.session_id72                  )73            """),74            {"cid": cid},75        )76        return result.rowcount or 077 78 79def main() -> None:80    parser = argparse.ArgumentParser(description="Backfill session title and optional client_id")81    parser.add_argument(82        "--client-id",83        metavar="UUID",84        help="Browser client id (localStorage text2sql_client_id). Also reads BACKFILL_CLIENT_ID env.",85    )86    args = parser.parse_args()87    client_id = (args.client_id or os.getenv("BACKFILL_CLIENT_ID") or "").strip() or None88 89    engine = get_engine()90    run_title_backfill(engine)91    print("Done: title backfill (from backfill_session_metadata.sql).")92 93    if client_id:94        try:95            n = run_client_id_backfill(engine, client_id)96            print(f"Done: client_id backfill — updated {n} session(s) with NULL client_id (that have messages).")97        except ValueError as e:98            print(f"Invalid --client-id / BACKFILL_CLIENT_ID: {e}", file=sys.stderr)99            sys.exit(1)100    else:101        print(102            "Skipped client_id backfill (no --client-id or BACKFILL_CLIENT_ID). "103            "GET /sessions only lists rows where client_id matches your browser; "104            "run again with: python scripts/run_backfill_session_metadata.py --client-id \"<uuid>\""105        )106 107 108if __name__ == "__main__":109    main()110