from sqlalchemy import create_engine, inspect, text from sqlalchemy.orm import sessionmaker, DeclarativeBase from app.config import get_settings settings = get_settings() engine = create_engine( settings.database_url, connect_args={"check_same_thread": False}, ) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) class Base(DeclarativeBase): pass def get_db(): db = SessionLocal() try: yield db finally: db.close() def init_db(): """Create all tables.""" # import models so SQLAlchemy registers them from app.models import admin_user, upstream, snapshot, webhook_config, notification_log, custom_page, website, revoked_token # noqa: F401 Base.metadata.create_all(bind=engine) _migrate_custom_pages() def _migrate_custom_pages(): """Apply small SQLite-safe migrations for deployments without Alembic.""" inspector = inspect(engine) if "custom_pages" not in inspector.get_table_names(): return columns = {col["name"] for col in inspector.get_columns("custom_pages")} with engine.begin() as conn: if "access_mode" not in columns: conn.execute(text("ALTER TABLE custom_pages ADD COLUMN access_mode VARCHAR(32) NOT NULL DEFAULT 'direct'")) conn.execute(text("UPDATE custom_pages SET access_mode = CASE WHEN use_proxy = 1 THEN 'proxy' ELSE 'direct' END")) if "login_username" not in columns: conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_username VARCHAR(255)")) if "login_password" not in columns: conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_password TEXT")) if "login_username_selector" not in columns: conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_username_selector VARCHAR(512)")) if "login_password_selector" not in columns: conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_password_selector VARCHAR(512)")) if "login_submit_selector" not in columns: conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_submit_selector VARCHAR(512)")) if "login_autofill_enabled" not in columns: conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_autofill_enabled BOOLEAN NOT NULL DEFAULT 0")) if "login_autofill_backfilled_at" not in columns: conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_autofill_backfilled_at DATETIME")) conn.execute( text( "UPDATE custom_pages " "SET login_autofill_enabled = 1, login_autofill_backfilled_at = CURRENT_TIMESTAMP " "WHERE login_autofill_enabled = 0 " "AND NULLIF(TRIM(login_username), '') IS NOT NULL " "AND NULLIF(TRIM(login_password), '') IS NOT NULL" ) )