mirror of
https://github.com/open-webui/open-webui
synced 2024-11-16 13:40:55 +00:00
borrow some of the previous PRs reconnection code
This commit is contained in:
parent
48e1356ed9
commit
59c6ff727a
@ -46,7 +46,7 @@ router = Router(
|
|||||||
)
|
)
|
||||||
router.run()
|
router.run()
|
||||||
try:
|
try:
|
||||||
DB.connect()
|
DB.connect(reuse_if_open=True)
|
||||||
except OperationalError as e:
|
except OperationalError as e:
|
||||||
log.info(f"Failed to connect to database again due to: {e}")
|
log.info(f"Failed to connect to database again due to: {e}")
|
||||||
pass
|
pass
|
@ -1,7 +1,9 @@
|
|||||||
from contextvars import ContextVar
|
from contextvars import ContextVar
|
||||||
from peewee import *
|
from peewee import *
|
||||||
|
from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from playhouse.db_url import connect
|
from playhouse.db_url import connect, parse
|
||||||
|
|
||||||
from config import SRC_LOG_LEVELS
|
from config import SRC_LOG_LEVELS
|
||||||
|
|
||||||
@ -23,15 +25,32 @@ class PeeweeConnectionState(object):
|
|||||||
value = self._state.get()[name]
|
value = self._state.get()[name]
|
||||||
return value
|
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):
|
def register_connection(db_url):
|
||||||
db = connect(db_url)
|
db = connect(db_url)
|
||||||
if isinstance(db, PostgresqlDatabase):
|
if isinstance(db, PostgresqlDatabase):
|
||||||
# Enable autoconnect for SQLite databases, managed by Peewee
|
# Enable autoconnect for SQLite databases, managed by Peewee
|
||||||
db.autoconnect = True
|
db.autoconnect = True
|
||||||
|
db.reuse_if_open = True
|
||||||
log.info("Connected to PostgreSQL database")
|
log.info("Connected to PostgreSQL database")
|
||||||
|
connection = parse(db_url)
|
||||||
|
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):
|
elif isinstance(db, SqliteDatabase):
|
||||||
# Enable autoconnect for SQLite databases, managed by Peewee
|
# Enable autoconnect for SQLite databases, managed by Peewee
|
||||||
db.autoconnect = True
|
db.autoconnect = 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')
|
||||||
|
Loading…
Reference in New Issue
Block a user