flask_snooze.py

# coding: utf-8

"""
Snooze: a backend-agnostic REST API provider for Flask.

e.g.

    from flask import app, Blueprint
    from flask.ext.sqlalchemy import SQLAlchemy
    from flask.ext.snooze import Snooze, SqlAlchemyEndpoint
    from my_model import sqlalchemy_db, Book

    api = Blueprint('api_v1', __name__)
    apimgr = Snooze(api)
    apimgr.add(SqlAlchemyEndpoint(sqlalchemy_db, Book, ['author', 'title']))

    app.register_blueprint(api, url_prefix='/api_v1')
"""

from flask import request, make_response
import re

try:
    import simplejson as json
except ImportError:
    import json


class NotFoundError(Exception):

    """
    Resource not found.
    """

    def __init__(self, cls, path):
        super(NotFoundError, self).__init__()

        self.cls = cls
        self.path = path
        self.message = 'No %(cls)s exists with an ID of %(path)s' % dict(
            cls=cls.__name__,
            path=path
        )


def error_dict(etype, message, **kwargs):
    d = dict(type=etype, message=message)
    if kwargs:
        d['detail'] = kwargs
    return d


class CoerceToDictEncoder(json.JSONEncoder):

    """
    A fairly naive encoder that will try to convert unrecognised types to dict.
    The idea being that objects can be made iterable quite easily as a bridge
    to being converted to JSON.
    """

    def default(self, obj):
        if obj is None or type(obj) in (
                dict,
                list, tuple,
                str, unicode,
                int, long, float,
                bool):
            return json.JSONEncoder.default(self, obj)

        return dict(obj)


def wrap_verb_call(call, endpoint, data_in, data_out):
    """
    Construct a callback that will wrap a given HTTP Verb call, passing a path.
    """
    def f(path=None):
        data = data_in(request.data) if request.data != '' else dict()
        assert isinstance(data, dict), "Data must be a dict"
        try:
            res = call(endpoint, path, data)
            try:
                # NB. error_data used because Flask stringifies stuff we put
                #     into res.data, which isn't good for us
                res.data = data_out(res.error_data)
            except AttributeError:
                try:
                    res.data = data_out(res.data)
                except AttributeError:
                    res = data_out(res)
        except NotFoundError, e:
            res = make_response()
            res.status = '404'
            res.data = data_out(error_dict(**{
                'etype': type(e).__name__,
                'message': e.message,
                'class': e.cls.__name__,
                'path': e.path
            }))
        except:
            import sys
            from traceback import extract_tb
            exc_type, exc_value, exc_traceback = sys.exc_info()
            res = data_out(error_dict(exc_type.__name__,
                           exc_value.message,
                           traceback=extract_tb(exc_traceback))), '500'
        return res
    return f


def response_redirect(endpoint, o, code):
    r = make_response()
    r.headers['Location'] = '%(path)s%(id)s' % dict(
        path=re.sub('[^/]*$', '', request.path),
        id=getattr(o, endpoint.id_key)
    )
    r.status = str(code)
    return r


class Snooze(object):

    """
    The API context manager,
    The api level means:
        every verb takes in and gives out data in the same ways
    """

    def __init__(self, app, hooks=None):
        self._app = app
        hooks = dict() if hooks is None else hooks
        self._hook_data_in = hooks.get('data_in', json.loads)
        self._hook_data_out = hooks.get('data_out', CoerceToDictEncoder().encode)
        self._routes = {}

    def add(self, endpoint, name=None, methods=(
            'OPTIONS', 'POST', 'GET', 'PUT', 'PATCH', 'DELETE')):
        """
        Add an endpoint for a class, the name defaults to a lowercase version
        of the class name but can be overriden.

        Methods can be specified, note that HEAD is automatically generated by
        Flask to execute the GET method without returning a body.
        """
        obj_name = endpoint.cls.__name__.lower() if name is None else name
        methods = [m.upper() for m in methods]

        for verb in 'OPTIONS', 'POST', 'GET', 'PUT', 'PATCH', 'DELETE':
            if verb not in methods:
                continue

            l = wrap_verb_call(call=getattr(self, '_%s' % verb.lower()),
                               endpoint=endpoint,
                               data_in=self._hook_data_in,
                               data_out=self._hook_data_out)

            self._register(obj_name=obj_name,
                           verb=verb,
                           func=l)

    #
    # Verbs
    #

    def _options(self, endpoint, path, data):
        """HTTP Verb endpoint"""
        return self._routes

    def _post(self, endpoint, path, data):
        """HTTP Verb endpoint"""
        o = endpoint.create(path)
        if data is not None:
            self._fill(endpoint, o, data)

        return response_redirect(endpoint, o, 201)

    def _get(self, endpoint, path, data):
        """HTTP Verb endpoint"""
        return endpoint.read(path)

    def _put(self, endpoint, path, data):
        """HTTP Verb endpoint"""
        created = False
        try:
            o = endpoint.read(path)
        except NotFoundError:
            o = endpoint.create(path)
            created = True

        self._fill(endpoint, o, data)

        if created:
            return response_redirect(endpoint, o, 201)

    def _patch(self, endpoint, path, data):
        """HTTP Verb endpoint"""
        o = endpoint.read(path)
        self._update(endpoint, o, data)

    def _delete(self, endpoint, path, data):
        """HTTP Verb endpoint"""
        endpoint.delete(path)

    #
    # Tools
    #

    def _update(self, endpoint, o, data):
        for k in data:
            assert k in endpoint.writeable_keys, \
                "Cannot update key %s, valid keys for update: %s" % \
                    (k, ', '.join(endpoint.writeable_keys))
            setattr(o, k, data[k])
        endpoint.finalize(o)

    def _fill(self, endpoint, o, data):
        items_set = set(endpoint.writeable_keys)
        keys_set = set(data.keys())
        assert items_set == keys_set, \
            "The provided keys (%s) do not match the expected items (%s)" % \
                (', '.join(keys_set), ', '.join(items_set))

        self._update(endpoint, o, data)

    def _register(self, obj_name, verb, func):
        func.provide_automatic_options = False

        route = '/%s/<path:path>' % obj_name
        self._app.route(route,
                        methods=(verb,),
                        endpoint="%s:%s" % (verb, route))(func)

        self._reg_options(verb, route)

        if verb in ('OPTIONS', 'GET', 'POST'):
            route = '/%s/' % obj_name
            self._app.route(route,
                            methods=(verb,),
                            endpoint="%s:%s" % (verb, route),
                            defaults={'path': None})(func)

            self._reg_options(verb, route)

    def _reg_options(self, verb, route):
        verbs = self._routes.get(route, [])
        verbs.append(verb)
        if verb == 'GET':
            # Flask adds 'HEAD' for GET
            verbs.append('HEAD')
        self._routes[route] = verbs


class Endpoint(object):

    """
    Base Endpoint object.
    """

    def __init__(self, cls, id_key, writeable_keys):
        """
        cls:            Class of object being represented by this endpoint
        id_key:         Identifying key of an object
        writeable_keys: A list of keys that may be written to on an object
        """
        self.cls = cls
        self.id_key = id_key
        self.writeable_keys = writeable_keys

    def create(self, path=None):
        """Create a new object"""
        raise NotImplementedError()

    def read(self, path):
        """Load an existing object"""
        raise NotImplementedError()

    def finalize(self, obj):
        """Save an object (if required)"""
        raise NotImplementedError()

    def delete(self, path):
        """Delete the data for the provided ID"""
        raise NotImplementedError()


#
# SQLAlchemy Land
#


def row2dict(row):
    """
    Convert a SQLAlchemy row/object to a dict, found on:
    http://stackoverflow.com/questions/
        1958219/convert-sqlalchemy-row-object-to-python-dict
    """
    d = {}
    for col_name in row.__table__.columns.keys():
        d[col_name] = getattr(row, col_name)

    return d


class SqlAlchemyEndpoint(Endpoint):

    def __init__(self, db, cls, items):
        from sqlalchemy.orm import class_mapper
        self.db = db
        self.pk = class_mapper(cls).primary_key[0]
        super(SqlAlchemyEndpoint, self).__init__(cls, self.pk.name, items)

    def create(self, path=None):
        o = self.cls()
        if path is not None:
            setattr(o, self.id_key, path)
        return o

    def read(self, path):
        if path == None:
            return [pk[0] for pk in \
                self.db.session.query(self.pk).all()]

        try:
            return self.cls.query.filter(self.pk == path).all()[0]
        except IndexError:
            raise NotFoundError(self.cls, path)

    def finalize(self, obj):
        self.db.session.add(obj)
        self.db.session.commit()

    def delete(self, path):
        o = self.read(path)
        self.db.session.delete(o)