File size: 1,834 Bytes
0733fd6 825a805 a467728 825a805 0733fd6 a467728 0733fd6 a467728 0733fd6 a467728 0733fd6 825a805 0733fd6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
# database_module/db.py
import os
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
# Configure your database URL (env var or fallback to SQLite file)
DATABASE_URL = os.getenv(
"DATABASE_URL",
"sqlite:///./drift_detector.sqlite3"
)
# Create engine and session factory
# For SQLite, disable same-thread check
connect_args = {"check_same_thread": False} if DATABASE_URL.startswith("sqlite") else {}
engine = create_engine(
DATABASE_URL,
connect_args=connect_args,
pool_pre_ping=True
)
SessionLocal = sessionmaker(
autocommit=False,
autoflush=False,
bind=engine
)
Base = declarative_base()
def apply_migrations():
"""
Apply any necessary migrations to existing tables.
"""
with engine.connect() as conn:
# Check if the models table exists and has the required columns
inspector = inspect(engine)
if "models" in inspector.get_table_names():
columns = [col['name'] for col in inspector.get_columns('models')]
# Add capabilities column if missing
if "capabilities" not in columns:
conn.execute(text("ALTER TABLE models ADD COLUMN capabilities TEXT"))
conn.commit()
print("Migration: Added capabilities column to models table")
# Add updated column if missing
if "updated" not in columns:
conn.execute(text("ALTER TABLE models ADD COLUMN updated DATETIME"))
conn.commit()
print("Migration: Added updated column to models table")
def init_db():
"""
Create tables if they don't exist.
Call this once at application startup.
"""
Base.metadata.create_all(bind=engine)
apply_migrations() |