Spaces:
Sleeping
Sleeping
| from sqlalchemy import create_engine | |
| from sqlalchemy.engine import Engine | |
| from sqlalchemy.orm import sessionmaker | |
| from app.services.utility import UtilityClass | |
| import sqlite3 | |
| from pathlib import Path | |
| import os | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| # Path to SQLite DB file, configurable via environment variable or .env | |
| print("DB_PATH from env:", os.getenv("DB_PATH")) | |
| DB_PATH = os.getenv("DB_PATH") | |
| # DB_PATH = UtilityClass.download_sqlite_db_from_dropbox(os.getenv("DROPBOX_TOKEN"), os.getenv("DROPBOX_DB_PATH")) | |
| if not DB_PATH: | |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| DB_PATH = os.path.join(BASE_DIR, "Banking.db") | |
| DATABASE_URL = f"sqlite:///{DB_PATH}" | |
| def get_connection(): | |
| db_path = DB_PATH | |
| # Only allow connection if DB file exists | |
| if not os.path.isfile(db_path): | |
| raise FileNotFoundError(f"Database file not found at {db_path}. Set DB_PATH env variable or .env to the correct location.") | |
| try: | |
| # Remove timeout restriction to allow long-running queries | |
| # Complex AI-generated queries may take time and should not be interrupted | |
| conn = sqlite3.connect(db_path, check_same_thread=False) | |
| conn.row_factory = sqlite3.Row # access columns by name | |
| return conn | |
| except sqlite3.OperationalError as e: | |
| raise Exception(f"Error connecting to DB at {db_path}: {e}") | |
| # Create engine with no timeout restrictions for complex queries | |
| engine = create_engine( | |
| DATABASE_URL, | |
| connect_args={ | |
| "check_same_thread": False, # Needed for SQLite threading | |
| "timeout": 0 # No timeout - wait indefinitely for query completion | |
| }, | |
| pool_timeout=None, # No pool timeout | |
| pool_recycle=-1 # No connection recycling timeout | |
| ) | |
| # Session factory | |
| SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) | |
| def get_db_engine(): | |
| """Return SQLAlchemy engine (used by LangChain SQLDatabase).""" | |
| return engine | |
| def get_db(): | |
| """Provide DB session for queries.""" | |
| db = SessionLocal() | |
| try: | |
| yield db | |
| finally: | |
| db.close() |