diff --git a/backend/apps/webui/internal/db.py b/backend/apps/webui/internal/db.py index a4b86ae1f..b61eb012d 100644 --- a/backend/apps/webui/internal/db.py +++ b/backend/apps/webui/internal/db.py @@ -46,7 +46,7 @@ router = Router( ) router.run() try: - DB.connect() + DB.connect(reuse_if_open=True) except OperationalError as e: log.info(f"Failed to connect to database again due to: {e}") pass \ No newline at end of file diff --git a/backend/apps/webui/internal/wrappers.py b/backend/apps/webui/internal/wrappers.py index b4da0ce25..d1cdecd29 100644 --- a/backend/apps/webui/internal/wrappers.py +++ b/backend/apps/webui/internal/wrappers.py @@ -1,7 +1,9 @@ from contextvars import ContextVar from peewee import * +from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError + import logging -from playhouse.db_url import connect +from playhouse.db_url import connect, parse from config import SRC_LOG_LEVELS @@ -23,15 +25,32 @@ class PeeweeConnectionState(object): value = self._state.get()[name] 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): db = connect(db_url) if isinstance(db, PostgresqlDatabase): # Enable autoconnect for SQLite databases, managed by Peewee db.autoconnect = True + db.reuse_if_open = True 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): # Enable autoconnect for SQLite databases, managed by Peewee db.autoconnect = True + db.reuse_if_open = True log.info("Connected to SQLite database") else: raise ValueError('Unsupported database connection')