mirror of
https://github.com/open-webui/open-webui
synced 2025-01-18 00:30:51 +00:00
use autoconnect and stop using the mixin
This commit is contained in:
parent
10fa887eab
commit
981866eb93
@ -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
|
||||
return db
|
Loading…
Reference in New Issue
Block a user