fix peewee and playhouse connections to retry

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