feat(sqlalchemy): Replace peewee with sqlalchemy

backend/apps/webui/internal/db.py

@@ -1,18 +1,34 @@
 import os
 import logging
 import json
+from typing import Optional, Any
+from typing_extensions import Self
 
-from peewee import *
-from peewee_migrate import Router
+from sqlalchemy import create_engine, types, Dialect
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.sql.type_api import _T
 
-from apps.webui.internal.wrappers import register_connection
 from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["DB"])
 
 
-class JSONField(TextField):
+class JSONField(types.TypeDecorator):
+    impl = types.Text
+    cache_ok = True
+
+    def process_bind_param(self, value: Optional[_T], dialect: Dialect) -> Any:
+        return json.dumps(value)
+
+    def process_result_value(self, value: Optional[_T], dialect: Dialect) -> Any:
+        if value is not None:
+            return json.loads(value)
+
+    def copy(self, **kw: Any) -> Self:
+        return JSONField(self.impl.length)
+
     def db_value(self, value):
         return json.dumps(value)
 
@@ -29,26 +45,24 @@ if os.path.exists(f"{DATA_DIR}/ollama.db"):
 else:
     pass
 
+SQLALCHEMY_DATABASE_URL = DATABASE_URL
+if "sqlite" in SQLALCHEMY_DATABASE_URL:
+    engine = create_engine(
+        SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
+    )
+else:
+    engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+Base = declarative_base()
 
-# The `register_connection` function encapsulates the logic for setting up
-# the database connection based on the connection string, while `connect`
-# is a Peewee-specific method to manage the connection state and avoid errors
-# when a connection is already open.
-try:
-    DB = register_connection(DATABASE_URL)
-    log.info(f"Connected to a {DB.__class__.__name__} database.")
-except Exception as e:
-    log.error(f"Failed to initialize the database connection: {e}")
-    raise
 
-router = Router(
-    DB,
-    migrate_dir=BACKEND_DIR / "apps" / "webui" / "internal" / "migrations",
-    logger=log,
-)
-router.run()
-try:
-    DB.connect(reuse_if_open=True)
-except OperationalError as e:
-    log.info(f"Failed to connect to database again due to: {e}")
-    pass
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+        db.commit()
+    except Exception as e:
+        db.rollback()
+        raise e
+    finally:
+        db.close()

backend/apps/webui/internal/wrappers.py

@@ -1,72 +0,0 @@
-from contextvars import ContextVar
-from peewee import *
-from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError
-
-import logging
-from playhouse.db_url import connect, parse
-from playhouse.shortcuts import ReconnectMixin
-
-from config import SRC_LOG_LEVELS
-
-log = logging.getLogger(__name__)
-log.setLevel(SRC_LOG_LEVELS["DB"])
-
-db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None}
-db_state = ContextVar("db_state", default=db_state_default.copy())
-
-
-class PeeweeConnectionState(object):
-    def __init__(self, **kwargs):
-        super().__setattr__("_state", db_state)
-        super().__init__(**kwargs)
-
-    def __setattr__(self, name, value):
-        self._state.get()[name] = value
-
-    def __getattr__(self, name):
-        value = self._state.get()[name]
-        return value
-
-
-class CustomReconnectMixin(ReconnectMixin):
-    reconnect_errors = (
-        # psycopg2
-        (OperationalError, "termin"),
-        (InterfaceError, "closed"),
-        # peewee
-        (PeeWeeInterfaceError, "closed"),
-    )
-
-
-class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase):
-    pass
-
-
-def register_connection(db_url):
-    db = connect(db_url)
-    if isinstance(db, PostgresqlDatabase):
-        # Enable autoconnect for SQLite databases, managed by Peewee
-        db.autoconnect = True
-        db.reuse_if_open = True
-        log.info("Connected to PostgreSQL database")
-
-        # Get the connection details
-        connection = parse(db_url)
-
-        # Use our custom database class that supports reconnection
-        db = ReconnectingPostgresqlDatabase(
-            connection["database"],
-            user=connection["user"],
-            password=connection["password"],
-            host=connection["host"],
-            port=connection["port"],
-        )
-        db.connect(reuse_if_open=True)
-    elif isinstance(db, SqliteDatabase):
-        # Enable autoconnect for SQLite databases, managed by Peewee
-        db.autoconnect = True
-        db.reuse_if_open = True
-        log.info("Connected to SQLite database")
-    else:
-        raise ValueError("Unsupported database connection")
-    return db