UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-01-18 16:51:17 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-07-08 12:42:52 -07:00
parent 68d775e1ab
commit d3ef3a7494
3 changed files with 6 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
backend/apps/webui/internal/db.py
View File

@ -53,9 +53,12 @@ else:
# Workaround to handle the peewee migration # Workaround to handle the peewee migration
# This is required to ensure the peewee migration is handled before the alembic 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: 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" migrate_dir = BACKEND_DIR / "apps" / "webui" / "internal" / "migrations"
router = Router(db, logger=log, migrate_dir=migrate_dir) router = Router(db, logger=log, migrate_dir=migrate_dir)
router.run() router.run()
@ -76,11 +79,10 @@ def handle_peewee_migration():
assert db.is_closed(), "Database connection is still open." assert db.is_closed(), "Database connection is still open."
handle_peewee_migration() handle_peewee_migration(DATABASE_URL)
SQLALCHEMY_DATABASE_URL = DATABASE_URL SQLALCHEMY_DATABASE_URL = DATABASE_URL
if "sqlite" in SQLALCHEMY_DATABASE_URL: if "sqlite" in SQLALCHEMY_DATABASE_URL:
engine = create_engine( engine = create_engine(
SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False} SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}

2
backend/apps/webui/models/chats.py
View File

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

6
backend/migrations/env.py
View File

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