Add database connection pooling

This commit is contained in:
Peter Stockings
2026-02-04 00:03:03 +11:00
parent a6a71f3139
commit b4121eada7

29
db.py
View File

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