UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-11-16 13:40:55 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

format

Browse Source
This commit is contained in:
perf3ct 2024-06-17 16:47:09 -07:00
parent 59c6ff727a
commit 81b2416923
No known key found for this signature in database
GPG Key ID: 569C4EEC436F5232
1 changed files with 21 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
backend/apps/webui/internal/wrappers.py
View File

@ -4,6 +4,7 @@ from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError
import logging import logging
from playhouse.db_url import connect, parse from playhouse.db_url import connect, parse
from playhouse.shortcuts import ReconnectMixin
from config import SRC_LOG_LEVELS from config import SRC_LOG_LEVELS
@ -13,6 +14,7 @@ log.setLevel(SRC_LOG_LEVELS["DB"])
db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None} db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None}
db_state = ContextVar("db_state", default=db_state_default.copy()) db_state = ContextVar("db_state", default=db_state_default.copy())
class PeeweeConnectionState(object): class PeeweeConnectionState(object):
def __init__(self, **kwargs): def __init__(self, **kwargs):
super().__setattr__("_state", db_state) super().__setattr__("_state", db_state)
@ -25,18 +27,21 @@ class PeeweeConnectionState(object):
value = self._state.get()[name] value = self._state.get()[name]
return value return value
class CustomReconnectMixin(ReconnectMixin): class CustomReconnectMixin(ReconnectMixin):
reconnect_errors = ( reconnect_errors = (
# psycopg2 # psycopg2
(OperationalError, 'termin'), (OperationalError, "termin"),
(InterfaceError, 'closed'), (InterfaceError, "closed"),
# peewee # peewee
(PeeWeeInterfaceError, 'closed'), (PeeWeeInterfaceError, "closed"),
) )
class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase): class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase):
pass pass
def register_connection(db_url): def register_connection(db_url):
db = connect(db_url) db = connect(db_url)
if isinstance(db, PostgresqlDatabase): if isinstance(db, PostgresqlDatabase):
@ -44,8 +49,18 @@ def register_connection(db_url):
db.autoconnect = True db.autoconnect = True
db.reuse_if_open = True db.reuse_if_open = True
log.info("Connected to PostgreSQL database") log.info("Connected to PostgreSQL database")
# Get the connection details
connection = parse(db_url) connection = parse(db_url)
db = ReconnectingPostgresqlDatabase(connection['database'], user=connection['user'], password=connection['password'],host=connection['host'], port=connection['port'])
# 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) db.connect(reuse_if_open=True)
elif isinstance(db, SqliteDatabase): elif isinstance(db, SqliteDatabase):
# Enable autoconnect for SQLite databases, managed by Peewee # Enable autoconnect for SQLite databases, managed by Peewee
@ -53,5 +68,5 @@ def register_connection(db_url):
db.reuse_if_open = True db.reuse_if_open = True
log.info("Connected to SQLite database") log.info("Connected to SQLite database")
else: else:
raise ValueError('Unsupported database connection') raise ValueError("Unsupported database connection")
return db return db