"...git@developer.sourcefind.cn:chenpangpang/open-webui.git" did not exist on "43e08c6afa00a536fc50bd81a667bc935e1d5fae"
Unverified Commit 59c6ff72 authored by perf3ct's avatar perf3ct
Browse files

borrow some of the previous PRs reconnection code

parent 48e1356e
...@@ -46,7 +46,7 @@ router = Router( ...@@ -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
\ No newline at end of file
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): ...@@ -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')
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment