django-orchestra/orchestra/utils/db.py

50 lines
1.5 KiB
Python
Raw Normal View History

2015-05-03 17:44:46 +00:00
import ast
from django import db
def close_connection(execute):
""" Threads have their own connection pool, closing it when finishing """
def wrapper(*args, **kwargs):
try:
log = execute(*args, **kwargs)
except Exception as e:
pass
else:
wrapper.log = log
finally:
db.connection.close()
return wrapper
def get_settings(settings_file):
""" get db settings from settings.py file without importing """
settings = {}
with open(settings_file, 'r') as handler:
body = ast.parse(handler.read()).body
for var in body:
targets = getattr(var, 'targets', None)
if targets and targets[0].id == 'DATABASES':
keys = var.value.values[0].keys
values = var.value.values[0].values
for key, value in zip(keys, values):
if key.s == 'ENGINE':
2015-05-03 18:33:07 +00:00
if value.s not in ('django.db.backends.sqlite3', 'django.db.backends.postgresql_psycopg2'):
2015-05-03 17:44:46 +00:00
raise ValueError("%s engine not supported." % value)
settings[key.s] = getattr(value, 's', None)
return settings
def get_connection(settings):
import psycopg2
conn = psycopg2.connect("dbname='{NAME}' user='{USER}' host='{HOST}' password='{PASSWORD}'".format(**settings))
return conn
def run_query(conn, query):
cur = conn.cursor()
cur.execute(query)
result = cur.fetchall()
cur.close()
return result