enh: handle peewee migration

backend/apps/webui/internal/db.py

@@ -2,6 +2,10 @@ import os
 import logging
 import json
 from contextlib import contextmanager
+
+from peewee_migrate import Router
+from apps.webui.internal.wrappers import register_connection
+
 from typing import Optional, Any
 from typing_extensions import Self
 
@@ -46,6 +50,35 @@ if os.path.exists(f"{DATA_DIR}/ollama.db"):
 else:
     pass
 
+
+# Workaround to handle the peewee migration
+# This is required to ensure the peewee migration is handled before the alembic migration
+def handle_peewee_migration():
+    try:
+        db = register_connection(DATABASE_URL)
+        migrate_dir = BACKEND_DIR / "apps" / "webui" / "internal" / "migrations"
+        router = Router(db, logger=log, migrate_dir=migrate_dir)
+        router.run()
+        db.close()
+
+        # check if db connection has been closed
+
+    except Exception as e:
+        log.error(f"Failed to initialize the database connection: {e}")
+        raise
+
+    finally:
+        # Properly closing the database connection
+        if db and not db.is_closed():
+            db.close()
+
+        # Assert if db connection has been closed
+        assert db.is_closed(), "Database connection is still open."
+
+
+handle_peewee_migration()
+
+
 SQLALCHEMY_DATABASE_URL = DATABASE_URL
 if "sqlite" in SQLALCHEMY_DATABASE_URL:
     engine = create_engine(
@@ -62,9 +95,6 @@ Base = declarative_base()
 Session = scoped_session(SessionLocal)
 
 
-from contextlib import contextmanager
-
-
 # Dependency
 def get_session():
     db = SessionLocal()