import os import psycopg2 from psycopg2.extras import RealDictCursor from urllib.parse import urlparse from flask import g class DataBase(): def __init__(self, app): db_url = urlparse(os.environ['DATABASE_URL']) # if db_url is null then throw error if not db_url: raise Exception("No DATABASE_URL environment variable set") def getDB(self): db = getattr(g, 'database', None) if db is None: db_url = urlparse(os.environ['DATABASE_URL']) g.database = psycopg2.connect( 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) if db is not None: db.close() def execute(self, query, args=(), one=False, commit=False): conn = self.getDB() cur = conn.cursor(cursor_factory=RealDictCursor) cur.execute(query, args) rv = None if cur.description is not None: rv = cur.fetchall() if commit: try: conn.commit() except: conn.rollback() cur.close() return (rv[0] if rv else None) if one else rv def get_http_functions(self): http_functions = self.execute( 'SELECT id, NAME, script_content, invoked_count FROM http_functions', []) return http_functions def get_http_function(self, name): http_function = self.execute( 'SELECT id, NAME, script_content, invoked_count FROM http_functions WHERE NAME=%s', [name], one=True) return http_function def create_new_http_function(self, name, script_content): self.execute( 'INSERT INTO http_functions (NAME, script_content) VALUES (%s, %s)', [name, script_content], commit=True)