Add db connection pooling

app.py

@@ -204,3 +204,7 @@ if __name__ == '__main__':
     # Bind to PORT if defined, otherwise default to 5000.
     port = int(os.environ.get('PORT', 5000))
     app.run(host='127.0.0.1', port=port)
+
+@app.teardown_appcontext
+def teardown_db(exception):
+    db.close_conn()

db.py

@@ -1,51 +1,61 @@
 import json
 import os
 import psycopg2
+from psycopg2 import pool
 from psycopg2.extras import RealDictCursor
 from urllib.parse import urlparse
 from flask import g
 
 class DataBase():
-    def __init__(self, app=None):
+    def __init__(self):
+        self.pool = None
+
+    def init_app(self, app):
         db_url = urlparse(os.environ['DATABASE_URL'])
-        # if db_url is null then throw error
         if not db_url:
             raise Exception("No DATABASE_URL environment variable set")
         
-    def getDB(self):
-        db = getattr(g, 'database', None)
-        if db is None:
-            db_url = urlparse(os.environ['DATABASE_URL'])
-            g.database = psycopg2.connect(
+        self.pool = psycopg2.pool.SimpleConnectionPool(
+            1, 20,
             database=db_url.path[1:],
             user=db_url.username,
             password=db_url.password,
             host=db_url.hostname,
             port=db_url.port
         )
-            db = g.database
-        return db
         
-    def close_connection(exception):
-        db = getattr(g, 'database', None)
-        if db is not None:
-            db.close()
+        app.teardown_appcontext(self.close_conn)
+
+    def get_conn(self):
+        if 'db_conn' not in g:
+            g.db_conn = self.pool.getconn()
+        return g.db_conn
+
+    def close_conn(self, e=None):
+        db_conn = g.pop('db_conn', None)
+        if db_conn is not None:
+            self.pool.putconn(db_conn)
+
+    def close_all_connections(self):
+        if self.pool:
+            self.pool.closeall()
 
     def execute(self, query, args=(), one=False, commit=False):
-        conn = self.getDB()
+        conn = self.get_conn()
         cur = conn.cursor(cursor_factory=RealDictCursor)
+        try:
             cur.execute(query, args)
             rv = None
             if cur.description is not None:
                 rv = cur.fetchall()
             if commit:
-            try:
                 conn.commit()
-            except:
-                conn.rollback()
-        cur.close()
-
             return (rv[0] if rv else None) if one else rv
+        except Exception as e:
+            conn.rollback()
+            raise e
+        finally:
+            cur.close()
     
     def get_http_functions_for_user(self, user_id):
         http_functions = self.execute(

extensions.py

@@ -16,6 +16,7 @@ environment = Environment(
 
 def init_app(app):
     htmx.init_app(app)
+    db.init_app(app)
     
     # Add all Flask's default Jinja2 globals and filters
     environment.globals.update(