From dfbc12594702bf601c4c391cfe080f8e0844d01e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=91=D0=B5=D0=BA=D0=BB=D0=B5=D0=BC=D0=B8=D1=88=D0=B5?= =?UTF-8?q?=D0=B2=20=D0=9F=D0=B5=D1=82=D1=80=20=D0=90=D0=BB=D0=B5=D0=BA?= =?UTF-8?q?=D1=81=D0=B5=D0=B5=D0=B2=D0=B8=D1=87?= Date: Thu, 30 May 2024 18:55:58 +0700 Subject: [PATCH 01/10] Reconnect to postgresql & mysql external databases when getting disconnected --- backend/apps/webui/internal/db.py | 9 ++++ backend/apps/webui/internal/wrappers.py | 59 +++++++++++++++++++++++++ 2 files changed, 68 insertions(+) create mode 100644 backend/apps/webui/internal/wrappers.py diff --git a/backend/apps/webui/internal/db.py b/backend/apps/webui/internal/db.py index 0e7b1f95d..dda58a4e1 100644 --- a/backend/apps/webui/internal/db.py +++ b/backend/apps/webui/internal/db.py @@ -7,6 +7,12 @@ from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR import os import logging +from peewee_migrate import Router +from playhouse.db_url import connect + +from apps.webui.internal.wrappers import PeeweeConnectionState, register_peewee_databases +from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL + log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["DB"]) @@ -20,6 +26,8 @@ class JSONField(TextField): return json.loads(value) +register_peewee_databases() + # Check if the file exists if os.path.exists(f"{DATA_DIR}/ollama.db"): # Rename the file @@ -29,6 +37,7 @@ else: pass DB = connect(DATABASE_URL) +DB._state = PeeweeConnectionState() log.info(f"Connected to a {DB.__class__.__name__} database.") router = Router( DB, diff --git a/backend/apps/webui/internal/wrappers.py b/backend/apps/webui/internal/wrappers.py new file mode 100644 index 000000000..406599b5a --- /dev/null +++ b/backend/apps/webui/internal/wrappers.py @@ -0,0 +1,59 @@ +from contextvars import ContextVar + +from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError, MySQLDatabase, _ConnectionState +from playhouse.db_url import register_database +from playhouse.pool import PooledPostgresqlDatabase, PooledMySQLDatabase +from playhouse.shortcuts import ReconnectMixin +from psycopg2 import OperationalError +from psycopg2.errors import InterfaceError + + +db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None} +db_state = ContextVar("db_state", default=db_state_default.copy()) + + +class PeeweeConnectionState(_ConnectionState): + def __init__(self, **kwargs): + super().__setattr__("_state", db_state) + super().__init__(**kwargs) + + def __setattr__(self, name, value): + self._state.get()[name] = value + + def __getattr__(self, name): + return self._state.get()[name] + + +class CustomReconnectMixin(ReconnectMixin): + reconnect_errors = ( + # default ReconnectMixin exceptions (MySQL specific) + *ReconnectMixin.reconnect_errors, + # psycopg2 + (OperationalError, 'termin'), + (InterfaceError, 'closed'), + # peewee + (PeeWeeInterfaceError, 'closed'), + ) + + +class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase): + pass + + +class ReconnectingPooledPostgresqlDatabase(CustomReconnectMixin, PooledPostgresqlDatabase): + pass + + +class ReconnectingMySQLDatabase(CustomReconnectMixin, MySQLDatabase): + pass + + +class ReconnectingPooledMySQLDatabase(CustomReconnectMixin, PooledMySQLDatabase): + pass + + +def register_peewee_databases(): + register_database(MySQLDatabase, 'mysql') + register_database(PooledMySQLDatabase, 'mysql+pool') + register_database(ReconnectingPostgresqlDatabase, 'postgres', 'postgresql') + register_database(ReconnectingPooledPostgresqlDatabase, 'postgres+pool', 'postgresql+pool') From e1fa453edae689c5eea6af0b23e3dc1e158b302f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=91=D0=B5=D0=BA=D0=BB=D0=B5=D0=BC=D0=B8=D1=88=D0=B5?= =?UTF-8?q?=D0=B2=20=D0=9F=D0=B5=D1=82=D1=80=20=D0=90=D0=BB=D0=B5=D0=BA?= =?UTF-8?q?=D1=81=D0=B5=D0=B5=D0=B2=D0=B8=D1=87?= Date: Thu, 30 May 2024 20:20:09 +0700 Subject: [PATCH 02/10] Test reconnection to postgres in gh actions --- .github/workflows/integration-test.yml | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/.github/workflows/integration-test.yml b/.github/workflows/integration-test.yml index 2426aff27..e64f93bc1 100644 --- a/.github/workflows/integration-test.yml +++ b/.github/workflows/integration-test.yml @@ -170,6 +170,26 @@ jobs: echo "Server has stopped" exit 1 fi + + # Check that service will reconnect to postgres when connection will be closed + status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/api/tags) + if [[ "$status_code" -ne 200 ]] ; then + echo "Server has failed before postgres reconnect check" + exit 1 + fi + + echo "Terminating all connections to postgres..." + python -c "import os, psycopg2 as pg2; \ + conn = pg2.connect(dsn=os.environ['DATABASE_URL'].replace('+pool', '')); \ + cur = conn.cursor(); \ + cur.execute('SELECT pg_terminate_backend(psa.pid) FROM pg_stat_activity psa WHERE datname = current_database() AND pid <> pg_backend_pid();')" + + status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/api/tags) + if [[ "$status_code" -ne 200 ]] ; then + echo "Server has not reconnected to postgres after connection was closed: returned status $status_code" + exit 1 + fi + # - name: Test backend with MySQL # if: success() || steps.sqlite.conclusion == 'failure' || steps.postgres.conclusion == 'failure' From e59e1f5049c1b3d9b5ff7f61d7f6051cd2950b22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=91=D0=B5=D0=BA=D0=BB=D0=B5=D0=BC=D0=B8=D1=88=D0=B5?= =?UTF-8?q?=D0=B2=20=D0=9F=D0=B5=D1=82=D1=80=20=D0=90=D0=BB=D0=B5=D0=BA?= =?UTF-8?q?=D1=81=D0=B5=D0=B5=D0=B2=D0=B8=D1=87?= Date: Thu, 30 May 2024 20:44:13 +0700 Subject: [PATCH 03/10] Fix rebase artifacts --- backend/apps/webui/internal/db.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/backend/apps/webui/internal/db.py b/backend/apps/webui/internal/db.py index dda58a4e1..7420bd019 100644 --- a/backend/apps/webui/internal/db.py +++ b/backend/apps/webui/internal/db.py @@ -1,17 +1,13 @@ +import os +import logging import json from peewee import * -from peewee_migrate import Router -from playhouse.db_url import connect -from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR -import os -import logging - from peewee_migrate import Router from playhouse.db_url import connect from apps.webui.internal.wrappers import PeeweeConnectionState, register_peewee_databases -from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL +from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["DB"]) From ad32a2ef3cc508d6540b16a762dec31b773818c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=91=D0=B5=D0=BA=D0=BB=D0=B5=D0=BC=D0=B8=D1=88=D0=B5?= =?UTF-8?q?=D0=B2=20=D0=9F=D0=B5=D1=82=D1=80=20=D0=90=D0=BB=D0=B5=D0=BA?= =?UTF-8?q?=D1=81=D0=B5=D0=B5=D0=B2=D0=B8=D1=87?= Date: Fri, 31 May 2024 13:26:23 +0700 Subject: [PATCH 04/10] Drop mysql restarts --- backend/apps/webui/internal/wrappers.py | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/backend/apps/webui/internal/wrappers.py b/backend/apps/webui/internal/wrappers.py index 406599b5a..53869b916 100644 --- a/backend/apps/webui/internal/wrappers.py +++ b/backend/apps/webui/internal/wrappers.py @@ -1,8 +1,8 @@ from contextvars import ContextVar -from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError, MySQLDatabase, _ConnectionState +from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError, _ConnectionState from playhouse.db_url import register_database -from playhouse.pool import PooledPostgresqlDatabase, PooledMySQLDatabase +from playhouse.pool import PooledPostgresqlDatabase from playhouse.shortcuts import ReconnectMixin from psycopg2 import OperationalError from psycopg2.errors import InterfaceError @@ -26,7 +26,7 @@ class PeeweeConnectionState(_ConnectionState): class CustomReconnectMixin(ReconnectMixin): reconnect_errors = ( - # default ReconnectMixin exceptions (MySQL specific) + # default ReconnectMixin exceptions *ReconnectMixin.reconnect_errors, # psycopg2 (OperationalError, 'termin'), @@ -44,16 +44,6 @@ class ReconnectingPooledPostgresqlDatabase(CustomReconnectMixin, PooledPostgresq pass -class ReconnectingMySQLDatabase(CustomReconnectMixin, MySQLDatabase): - pass - - -class ReconnectingPooledMySQLDatabase(CustomReconnectMixin, PooledMySQLDatabase): - pass - - def register_peewee_databases(): - register_database(MySQLDatabase, 'mysql') - register_database(PooledMySQLDatabase, 'mysql+pool') register_database(ReconnectingPostgresqlDatabase, 'postgres', 'postgresql') register_database(ReconnectingPooledPostgresqlDatabase, 'postgres+pool', 'postgresql+pool') From 10fa887eab491f1fcfb175b9c458d202e2198732 Mon Sep 17 00:00:00 2001 From: perf3ct Date: Sun, 16 Jun 2024 15:25:48 -0700 Subject: [PATCH 05/10] fix peewee and playhouse connections to retry --- backend/apps/webui/internal/db.py | 28 +++++++++----- backend/apps/webui/internal/wrappers.py | 51 +++++++++++-------------- 2 files changed, 41 insertions(+), 38 deletions(-) diff --git a/backend/apps/webui/internal/db.py b/backend/apps/webui/internal/db.py index 7420bd019..a4b86ae1f 100644 --- a/backend/apps/webui/internal/db.py +++ b/backend/apps/webui/internal/db.py @@ -4,15 +4,13 @@ import json from peewee import * from peewee_migrate import Router -from playhouse.db_url import connect -from apps.webui.internal.wrappers import PeeweeConnectionState, register_peewee_databases +from apps.webui.internal.wrappers import register_connection from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["DB"]) - class JSONField(TextField): def db_value(self, value): return json.dumps(value) @@ -21,9 +19,6 @@ class JSONField(TextField): if value is not None: return json.loads(value) - -register_peewee_databases() - # Check if the file exists if os.path.exists(f"{DATA_DIR}/ollama.db"): # Rename the file @@ -32,13 +27,26 @@ if os.path.exists(f"{DATA_DIR}/ollama.db"): else: pass -DB = connect(DATABASE_URL) -DB._state = PeeweeConnectionState() -log.info(f"Connected to a {DB.__class__.__name__} database.") + +# The `register_connection` function encapsulates the logic for setting up +# the database connection based on the connection string, while `connect` +# is a Peewee-specific method to manage the connection state and avoid errors +# when a connection is already open. +try: + DB = register_connection(DATABASE_URL) + log.info(f"Connected to a {DB.__class__.__name__} database.") +except Exception as e: + log.error(f"Failed to initialize the database connection: {e}") + raise + router = Router( DB, migrate_dir=BACKEND_DIR / "apps" / "webui" / "internal" / "migrations", logger=log, ) router.run() -DB.connect(reuse_if_open=True) +try: + DB.connect() +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 53869b916..11c91034f 100644 --- a/backend/apps/webui/internal/wrappers.py +++ b/backend/apps/webui/internal/wrappers.py @@ -1,18 +1,13 @@ from contextvars import ContextVar - -from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError, _ConnectionState -from playhouse.db_url import register_database +from peewee import * +from playhouse.db_url import connect from playhouse.pool import PooledPostgresqlDatabase from playhouse.shortcuts import ReconnectMixin -from psycopg2 import OperationalError -from psycopg2.errors import InterfaceError - db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None} db_state = ContextVar("db_state", default=db_state_default.copy()) - -class PeeweeConnectionState(_ConnectionState): +class PeeweeConnectionState(object): def __init__(self, **kwargs): super().__setattr__("_state", db_state) super().__init__(**kwargs) @@ -21,29 +16,29 @@ class PeeweeConnectionState(_ConnectionState): self._state.get()[name] = value def __getattr__(self, name): - return self._state.get()[name] + value = self._state.get()[name] + return value +class ReconnectingPostgresqlDatabase(ReconnectMixin, PostgresqlDatabase): + pass -class CustomReconnectMixin(ReconnectMixin): - reconnect_errors = ( - # default ReconnectMixin exceptions - *ReconnectMixin.reconnect_errors, - # psycopg2 - (OperationalError, 'termin'), - (InterfaceError, 'closed'), - # peewee - (PeeWeeInterfaceError, 'closed'), - ) +class ReconnectingPooledPostgresqlDatabase(ReconnectMixin, PooledPostgresqlDatabase): + pass - -class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase): +class ReconnectingSqliteDatabase(ReconnectMixin, SqliteDatabase): pass -class ReconnectingPooledPostgresqlDatabase(CustomReconnectMixin, PooledPostgresqlDatabase): - pass - - -def register_peewee_databases(): - register_database(ReconnectingPostgresqlDatabase, 'postgres', 'postgresql') - register_database(ReconnectingPooledPostgresqlDatabase, 'postgres+pool', 'postgresql+pool') +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) + elif isinstance(db, SqliteDatabase): + db = ReconnectingSqliteDatabase(db.database, **db.connect_params) + else: + raise ValueError('Unsupported database connection') + return db From 981866eb930b4ef26bd93b89313c70983951ee75 Mon Sep 17 00:00:00 2001 From: perf3ct Date: Mon, 17 Jun 2024 07:50:47 -0700 Subject: [PATCH 06/10] use autoconnect and stop using the mixin --- backend/apps/webui/internal/wrappers.py | 38 ++++++++++++------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/backend/apps/webui/internal/wrappers.py b/backend/apps/webui/internal/wrappers.py index 11c91034f..4285efd21 100644 --- a/backend/apps/webui/internal/wrappers.py +++ b/backend/apps/webui/internal/wrappers.py @@ -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 \ No newline at end of file From 5c655f298b6b692dc936eae0e81f39bd174e54b6 Mon Sep 17 00:00:00 2001 From: perf3ct Date: Mon, 17 Jun 2024 09:56:31 -0700 Subject: [PATCH 07/10] stop even using pooled DBs in peewee --- backend/apps/webui/internal/wrappers.py | 20 ++++---------------- 1 file changed, 4 insertions(+), 16 deletions(-) diff --git a/backend/apps/webui/internal/wrappers.py b/backend/apps/webui/internal/wrappers.py index 4285efd21..e90e9860f 100644 --- a/backend/apps/webui/internal/wrappers.py +++ b/backend/apps/webui/internal/wrappers.py @@ -22,23 +22,11 @@ class PeeweeConnectionState(object): def register_connection(db_url): db = connect(db_url) if isinstance(db, PostgresqlDatabase): - db = PooledPostgresqlExtDatabase( - db.database, - max_connections=8, - stale_timeout=300, - timeout=None, - autoconnect=True, - **db.connect_params - ) + # Directly use PostgresqlDatabase without pooling + db.autoconnect = True elif isinstance(db, SqliteDatabase): - db = PooledSqliteDatabase( - db.database, - max_connections=8, - stale_timeout=300, - timeout=None, - autoconnect=True, - **db.connect_params - ) + # Directly use SqliteDatabase without pooling + db.autoconnect = True else: raise ValueError('Unsupported database connection') return db \ No newline at end of file From 48e1356ed9cc54b2ca3852a2349b6f74871dc5f9 Mon Sep 17 00:00:00 2001 From: perf3ct Date: Mon, 17 Jun 2024 10:34:19 -0700 Subject: [PATCH 08/10] add logging for user upon db connection --- backend/apps/webui/internal/wrappers.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/backend/apps/webui/internal/wrappers.py b/backend/apps/webui/internal/wrappers.py index e90e9860f..b4da0ce25 100644 --- a/backend/apps/webui/internal/wrappers.py +++ b/backend/apps/webui/internal/wrappers.py @@ -1,8 +1,12 @@ from contextvars import ContextVar from peewee import * +import logging from playhouse.db_url import connect -from playhouse.pool import PooledPostgresqlExtDatabase -from playhouse.pool import PooledSqliteDatabase + +from config import SRC_LOG_LEVELS + +log = logging.getLogger(__name__) +log.setLevel(SRC_LOG_LEVELS["DB"]) db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None} db_state = ContextVar("db_state", default=db_state_default.copy()) @@ -22,11 +26,13 @@ class PeeweeConnectionState(object): def register_connection(db_url): db = connect(db_url) if isinstance(db, PostgresqlDatabase): - # Directly use PostgresqlDatabase without pooling + # Enable autoconnect for SQLite databases, managed by Peewee db.autoconnect = True + log.info("Connected to PostgreSQL database") elif isinstance(db, SqliteDatabase): - # Directly use SqliteDatabase without pooling + # Enable autoconnect for SQLite databases, managed by Peewee db.autoconnect = True + log.info("Connected to SQLite database") else: raise ValueError('Unsupported database connection') return db \ No newline at end of file From 59c6ff727af459ca274170afd05e8a49eb5d611c Mon Sep 17 00:00:00 2001 From: perf3ct Date: Mon, 17 Jun 2024 16:44:20 -0700 Subject: [PATCH 09/10] borrow some of the previous PRs reconnection code --- backend/apps/webui/internal/db.py | 2 +- backend/apps/webui/internal/wrappers.py | 21 ++++++++++++++++++++- 2 files changed, 21 insertions(+), 2 deletions(-) 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') From 81b2416923137e2e4eef820a7b85c8930154e1b8 Mon Sep 17 00:00:00 2001 From: perf3ct Date: Mon, 17 Jun 2024 16:47:09 -0700 Subject: [PATCH 10/10] format --- backend/apps/webui/internal/wrappers.py | 27 +++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/backend/apps/webui/internal/wrappers.py b/backend/apps/webui/internal/wrappers.py index d1cdecd29..2b5551ce2 100644 --- a/backend/apps/webui/internal/wrappers.py +++ b/backend/apps/webui/internal/wrappers.py @@ -4,6 +4,7 @@ from peewee import PostgresqlDatabase, InterfaceError as PeeWeeInterfaceError import logging from playhouse.db_url import connect, parse +from playhouse.shortcuts import ReconnectMixin from config import SRC_LOG_LEVELS @@ -13,6 +14,7 @@ log.setLevel(SRC_LOG_LEVELS["DB"]) db_state_default = {"closed": None, "conn": None, "ctx": None, "transactions": None} db_state = ContextVar("db_state", default=db_state_default.copy()) + class PeeweeConnectionState(object): def __init__(self, **kwargs): super().__setattr__("_state", db_state) @@ -25,18 +27,21 @@ class PeeweeConnectionState(object): value = self._state.get()[name] return value + class CustomReconnectMixin(ReconnectMixin): reconnect_errors = ( # psycopg2 - (OperationalError, 'termin'), - (InterfaceError, 'closed'), + (OperationalError, "termin"), + (InterfaceError, "closed"), # peewee - (PeeWeeInterfaceError, 'closed'), + (PeeWeeInterfaceError, "closed"), ) + class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase): pass + def register_connection(db_url): db = connect(db_url) if isinstance(db, PostgresqlDatabase): @@ -44,8 +49,18 @@ def register_connection(db_url): db.autoconnect = True db.reuse_if_open = True log.info("Connected to PostgreSQL database") + + # Get the connection details connection = parse(db_url) - db = ReconnectingPostgresqlDatabase(connection['database'], user=connection['user'], password=connection['password'],host=connection['host'], port=connection['port']) + + # Use our custom database class that supports reconnection + 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 @@ -53,5 +68,5 @@ def register_connection(db_url): db.reuse_if_open = True log.info("Connected to SQLite database") else: - raise ValueError('Unsupported database connection') - return db \ No newline at end of file + raise ValueError("Unsupported database connection") + return db