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

use autoconnect and stop using the mixin

Browse Source
This commit is contained in:
perf3ct 2024-06-17 07:50:47 -07:00
parent 10fa887eab
commit 981866eb93
No known key found for this signature in database
GPG Key ID: 569C4EEC436F5232
1 changed files with 19 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,8 +1,8 @@
from contextvars import ContextVar
from peewee import *
from playhouse.db_url import connect
from playhouse.pool import PooledPostgresqlDatabase
from playhouse.shortcuts import ReconnectMixin
from playhouse.pool import PooledPostgresqlExtDatabase
from playhouse.pool import PooledSqliteDatabase
db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None}
db_state = ContextVar("db_state", default=db_state_default.copy())
@ -19,26 +19,26 @@ class PeeweeConnectionState(object):
value = self._state.get()[name]
return value
class ReconnectingPostgresqlDatabase(ReconnectMixin, PostgresqlDatabase):
pass
class ReconnectingPooledPostgresqlDatabase(ReconnectMixin, PooledPostgresqlDatabase):
pass
class ReconnectingSqliteDatabase(ReconnectMixin, SqliteDatabase):
pass
def register_connection(db_url):
# Connect using the playhouse.db_url module, which supports multiple
# database types, then wrap the connection in a ReconnectMixin to handle dropped connections
db = connect(db_url)
if isinstance(db, PostgresqlDatabase):
db = ReconnectingPostgresqlDatabase(db.database, **db.connect_params)
elif isinstance(db, PooledPostgresqlDatabase):
db = ReconnectingPooledPostgresqlDatabase(db.database, **db.connect_params)
db = PooledPostgresqlExtDatabase(
db.database,
max_connections=8,
stale_timeout=300,
timeout=None,
autoconnect=True,
**db.connect_params
)
elif isinstance(db, SqliteDatabase):
db = ReconnectingSqliteDatabase(db.database, **db.connect_params)
db = PooledSqliteDatabase(
db.database,
max_connections=8,
stale_timeout=300,
timeout=None,
autoconnect=True,
**db.connect_params
)
else:
raise ValueError('Unsupported database connection')
return db