jon/jon/db/__init__.py

40 lines
830 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:
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