Permalink
Switch branches/tags
Nothing to show
Find file
Fetching contributors…
Cannot retrieve contributors at this time
146 lines (113 sloc) 4.39 KB
# -*- coding: utf-8 -*-
"""
flask_moresql
=============
A thin layer of glue between PostgreSQL stored procedures and Flask web
applications.
:copyright: (C) 2012 by Emanuele Rocca.
:license: BSD, see LICENSE for more details.
"""
import re
import urllib
import datetime
import simplejson
import psycopg2
import psycopg2.extras
from flask import request, make_response
def parse_rfc1738_args(name):
"""Parse the given database URI and return a dictionary of database
connection values.
Eg: parse_rfc1738_args('postgres://user:pass@host:5432/dbname')
-> {
'username': 'user',
'password': 'pass',
'host': 'host',
'port': '5432',
'database': 'dbname'
}
"""
if name is None:
raise RuntimeError("MORESQL_DATABASE_URI needs to be specified")
pattern = re.compile(r'''
(?P<name>[\w\+]+)://
(?:
(?P<username>[^:/]*)
(?::(?P<password>[^/]*))?
@)?
(?:
(?P<host>[^/:]*)
(?::(?P<port>[^/]*))?
)?
(?:/(?P<database>.*))?
''', re.X)
match = pattern.match(name)
if match is None:
raise RuntimeError("Wrong DB URI string '%s'" % name)
components = match.groupdict()
if components['name'] != 'postgres':
raise RuntimeError("MoreSQL only supports postgres databases")
if components['password'] is not None:
components['password'] = urllib.unquote_plus(components['password'])
del components['name']
return components
def _convert_http_value(value):
"""Try to parse the given JSON value. Return raw value on failure."""
try:
return simplejson.loads(value)
except simplejson.JSONDecodeError:
return value
def _get_procedure_arguments(fields, values):
"""Return a list of parameters to be passed to the stored procedure."""
if fields is None:
# The stored procedure has been called without parameters
return []
if values:
# Use user-supplied values
return [ values.get(field) for field in fields ]
# Use HTTP request values
return [ _convert_http_value(request.values.get(field))
for field in fields if request.values.get(field) is not None ]
def dthandler(obj):
if isinstance(obj, datetime.datetime):
return obj.isoformat()
class MoreSQL(object):
"""Used to connect to a given PostgreSQL database.
To use MoreSQL you need to first create a Flask application, and then bind
a MoreSQL instance to it::
app = Flask(__name__)
db = MoreSQL(app)
"""
def __init__(self, app):
self.app = app
creds = parse_rfc1738_args(app.config.get('MORESQL_DATABASE_URI'))
self.connection = psycopg2.connect(user=creds['username'],
password=creds['password'],
dbname=creds['database'],
host=creds['host'],
port=creds['port'])
self.cursor = self.connection.cursor(
cursor_factory=psycopg2.extras.RealDictCursor)
def execute(self, procname, fields=None, values=None):
"""Execute the given stored procedure. Return results as a JSON
HTTP response.
:param procname: the stored procedure name
:param fields: a list of dictionary fields used as parameters of the
stored procedure. The procedure will be called with no
arguments if omitted
:param values: an optional dictionary of values from which the
parameters should be taken. If omitted, default to the
values passed via HTTP
"""
procargs = _get_procedure_arguments(fields, values)
result = self.cursor.callproc(procname, procargs)
self.connection.commit()
if procargs == result:
# The parameters have not been modified. There was no in/out
# parameter, we have to fetch returned results from the cursor.
result = self.cursor.fetchall()
if len(result) == 1:
if len(result[0]) == 1:
result = result[0].values()[0]
else:
result = result[0]
return make_response(simplejson.dumps(result, default=dthandler))