Avoid postgres timeout errors in watcher

VIDA-NYU · Aug 21, 2023 · e98f85d · e98f85d
1 parent 3d9bdfa
commit e98f85d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 12 deletions.
diff --git a/reproserver/database.py b/reproserver/database.py
@@ -341,7 +341,15 @@ def connect(url=None, *, create=False):
             host=os.environ['POSTGRES_HOST'],
             database=os.environ['POSTGRES_DB'],
         )
-        engine = create_engine(url, connect_args={'connect_timeout': 10})
+        engine = create_engine(
+            url,
+            connect_args={
+                'connect_timeout': 10,
+                'keepalives_idle': 30,
+                'keepalives_interval': 10,
+                'keepalives_count': 3,
+            },
+        )
     else:
         engine = create_engine(url)
 

diff --git a/reproserver/run/k8s.py b/reproserver/run/k8s.py
@@ -370,15 +370,15 @@ async def _full_sync(self, api, DBSession):
                         raise
 
         # Handle runs with missing pods
-        db = DBSession()
-        runs = (
-            db.query(database.Run)
-            .filter(database.Run.done == None)  # noqa: E711
-            .all()
-        )
-        for run in runs:
-            if run.id not in self.running:
-                await self.connector.run_failed(run.id, "Pod deleted")
+        with DBSession() as db:
+            runs = (
+                db.query(database.Run)
+                .filter(database.Run.done == None)  # noqa: E711
+                .all()
+            )
+            for run in runs:
+                if run.id not in self.running:
+                    await self.connector.run_failed(run.id, "Pod deleted")
 
         logger.info(
             "Full sync complete, %d pods: %s",
@@ -425,8 +425,8 @@ async def _handle_watch_event(self, api, DBSession, event):
             return
 
         # Get run
-        db = DBSession()
-        run = db.query(database.Run).get(run_id)
+        with DBSession() as db:
+            run = db.query(database.Run).get(run_id)
         if run is None:
             logger.warning("Event in pod for unknown run %d", run_id)
             return