42 lines
931 B
Python
42 lines
931 B
Python
import psycopg2
|
|
|
|
from flask import current_app, g
|
|
from pathlib import Path
|
|
from psycopg2.extras import RealDictCursor
|
|
|
|
|
|
def get_db():
|
|
if "db" not in g:
|
|
# TODO: Make this configurable and use a default that works
|
|
# on the pool computers.
|
|
g.db = psycopg2.connect(current_app.config["DB_CONNECTION_STRING"])
|
|
run_query_on(g.db, "add_views.sql", None)
|
|
|
|
return g.db
|
|
|
|
|
|
def close_db(e=None):
|
|
db = g.pop("db", None)
|
|
if db is not None:
|
|
db.close()
|
|
|
|
|
|
def init_app(app):
|
|
app.teardown_appcontext(close_db)
|
|
|
|
|
|
def run_query(query_name, params=None):
|
|
return run_query_on(get_db(), query_name, params)
|
|
|
|
|
|
def run_query_on(db, query_name, params):
|
|
query = (Path(__file__).parent / Path(query_name)).read_text()
|
|
|
|
cursor = db.cursor(cursor_factory=RealDictCursor)
|
|
try:
|
|
cursor.execute(query, params)
|
|
return cursor
|
|
except:
|
|
db.rollback()
|
|
raise
|