Add database connection pooling

db.py

@@ -1,10 +1,11 @@
 import os
 import psycopg2
+from psycopg2 import pool
 from psycopg2.extras import RealDictCursor
 from datetime import datetime
 from dateutil.relativedelta import relativedelta
 from urllib.parse import urlparse
-from flask import g
+from flask import g, current_app
 from features.exercises import Exercises
 from features.people_graphs import PeopleGraphs
 from features.person_overview import PersonOverview
@@ -16,6 +17,8 @@ from utils import get_exercise_graph_model
 
 
 class DataBase():
+    _pool = None
+
     def __init__(self, app=None):
         self.stats = Stats(self.execute)
         self.exercises = Exercises(self.execute)
@@ -25,29 +28,29 @@ class DataBase():
         self.schema = Schema(self.execute)
         self.activityRequest = Activity(self.execute)
 
-        db_url = urlparse(os.environ['DATABASE_URL'])
-        # if db_url is null then throw error
-        if not db_url:
+        if not os.environ.get('DATABASE_URL'):
             raise Exception("No DATABASE_URL environment variable set")
 
-    def getDB(self):
-        db = getattr(g, 'database', None)
-        if db is None:
+        if DataBase._pool is None:
             db_url = urlparse(os.environ['DATABASE_URL'])
-            g.database = psycopg2.connect(
+            DataBase._pool = pool.ThreadedConnectionPool(
+                1, 20, # minconn, maxconn
                 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)
+    def getDB(self):
+        if 'database' not in g:
+            g.database = self._pool.getconn()
+        return g.database
+
+    def close_connection(self, exception=None):
+        db = g.pop('database', None)
         if db is not None:
-            db.close()
+            self._pool.putconn(db)
 
     def execute(self, query, args=(), one=False, commit=False):
         conn = self.getDB()