From 5a2c8c6b0450d825a527532d8095f839bfa03eb2 Mon Sep 17 00:00:00 2001
From: Jonathan Mortensen <56177725+jmo-qap@users.noreply.github.com>
Date: Fri, 4 Dec 2020 18:37:24 -0800
Subject: [PATCH] Handle forking somewhat more safely

---
 lib/pool.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/lib/pool.py b/lib/pool.py
index 30a29c33..87ad9055 100644
--- a/lib/pool.py
+++ b/lib/pool.py
@@ -185,3 +185,24 @@ class ThreadedConnectionPool(AbstractConnectionPool):
             self._closeall()
         finally:
             self._lock.release()
+
+class MultiProcessThreadedConnectionPool(ThreadedConnectionPool):
+
+    def __init__(self, *args, **kwargs):
+        self.pid = os.getpid()
+        self.args = args
+        self.kwargs = kwargs
+        super().__init__(*args, **kwargs)
+
+    def getconn(self, *args, **kwargs):
+        current_pid = os.getpid()
+        if current_pid != self.pid:
+            super().__init__(*self.args, **self.kwargs)
+            self.pid = current_pid
+        return super().getconn(*args, **kwargs)
+
+    def putconn(self, *args, **kwargs):
+        current_pid = os.getpid()
+        if current_pid != self.pid:
+            raise Exception("Returned a connection to the pool that was acquired pre-fork. Don't use pre-fork connections post-fork.")
+        return super().putconn(*args, **kwargs)