refac

backend/apps/webui/internal/db.py

@@ -53,9 +53,12 @@ else:
 
 # Workaround to handle the peewee migration
 # This is required to ensure the peewee migration is handled before the alembic migration
-def handle_peewee_migration():
+def handle_peewee_migration(DATABASE_URL):
     try:
-        db = register_connection(DATABASE_URL)
+        # Replace the postgresql:// with postgres:// and %40 with @ in the DATABASE_URL
+        db = register_connection(
+            DATABASE_URL.replace("postgresql://", "postgres://").replace("%40", "@")
+        )
         migrate_dir = BACKEND_DIR / "apps" / "webui" / "internal" / "migrations"
         router = Router(db, logger=log, migrate_dir=migrate_dir)
         router.run()
@@ -76,11 +79,10 @@ def handle_peewee_migration():
         assert db.is_closed(), "Database connection is still open."
 
 
-handle_peewee_migration()
+handle_peewee_migration(DATABASE_URL)
 
 
 SQLALCHEMY_DATABASE_URL = DATABASE_URL
-
 if "sqlite" in SQLALCHEMY_DATABASE_URL:
     engine = create_engine(
         SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}

backend/apps/webui/models/chats.py

@@ -248,9 +248,7 @@ class ChatTable:
     def get_chat_list_by_chat_ids(
         self, chat_ids: List[str], skip: int = 0, limit: int = 50
     ) -> List[ChatModel]:
-
         with get_db() as db:
-
             all_chats = (
                 db.query(Chat)
                 .filter(Chat.id.in_(chat_ids))

backend/migrations/env.py

@@ -41,17 +41,11 @@ target_metadata = Auth.metadata
 # ... etc.
 
 DB_URL = DATABASE_URL
-# Replace the postgres:// with postgresql://
-if "postgres://" in DB_URL:
-    DB_URL = DB_URL.replace("postgres://", "postgresql://")
 
 if DB_URL:
     config.set_main_option("sqlalchemy.url", DB_URL)
 
 
-print("DB_URL", DB_URL)
-
-
 def run_migrations_offline() -> None:
     """Run migrations in 'offline' mode.