diff --git a/db.py b/db.py index 66ae1f1..06b5252 100644 --- a/db.py +++ b/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()