Initial test

sdss · Feb 2, 2024 · 7163fa5 · 7163fa5
1 parent 4e168e4
commit 7163fa5
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/python/sdssdb/connection.py b/python/sdssdb/connection.py
@@ -21,7 +21,8 @@
 from sqlalchemy.orm import scoped_session, sessionmaker
 
 import peewee
-from peewee import OperationalError, PostgresqlDatabase
+from peewee import OperationalError
+from playhouse.pool import PooledPostgresqlExtDatabase
 from playhouse.postgres_ext import ArrayField
 from playhouse.reflection import Introspector, UnknownField
 
@@ -357,7 +358,7 @@ def post_connect(self):
         pass
 
 
-class PeeweeDatabaseConnection(DatabaseConnection, PostgresqlDatabase):
+class PeeweeDatabaseConnection(DatabaseConnection, PooledPostgresqlExtDatabase):
     """Peewee database connection implementation.
 
     Attributes
@@ -376,8 +377,12 @@ def __init__(self, *args, **kwargs):
         self._metadata = {}
 
         autorollback = kwargs.pop('autorollback', True)
+        max_connections = kwargs.pop('max_connections', 10)
+        stale_timeout = kwargs.pop('stale_timeout', None)
 
-        PostgresqlDatabase.__init__(self, None, autorollback=autorollback)
+        PooledPostgresqlExtDatabase.__init__(self, None, autorollback=autorollback,
+                                             stale_timeout=stale_timeout,
+                                             max_connections=max_connections)
         DatabaseConnection.__init__(self, *args, **kwargs)
 
     @property
@@ -408,16 +413,16 @@ def _conn(self, dbname, silent_on_fail=False, **params):
             except pgpasslib.FileNotFound:
                 params['password'] = None
 
-        PostgresqlDatabase.init(self, dbname, **params)
+        PooledPostgresqlExtDatabase.init(self, dbname, **params)
         self._metadata = {}
 
         try:
-            PostgresqlDatabase.connect(self)
+            PooledPostgresqlExtDatabase.connect(self)
             self.dbname = dbname
         except OperationalError as ee:
             if not silent_on_fail:
                 log.warning(f'failed to connect to database {self.database!r}: {ee}')
-            PostgresqlDatabase.init(self, None)
+            PooledPostgresqlExtDatabase.init(self, None)
 
         if self.is_connection_usable() and self.auto_reflect:
             with self.atomic():