Simple SQLite-backed key-value storage Rest API. Built with Flask & flask-restful.

sqlite-kv-restful.py

#!/usr/bin/env python3

"""
Requirements:

  pip install Flask==1.1.2 Flask-RESTful==0.3.8
"""

import os
import sqlite3
from hashlib import sha256
from time import time

from werkzeug.exceptions import BadRequest
from flask import Flask, g, request
from flask_restful import Api, Resource, abort

app = Flask(__name__)
api = Api(app)

PATH = os.path.dirname(os.path.abspath(__file__))
DATABASE = os.path.join(PATH, 'data.db')


CREATE_INDEX = ('CREATE TABLE IF NOT EXISTS %(index)s'
                '(key TEXT,'
                ' key_hash TEXT UNIQUE,'
                ' value BLOB,'
                ' modified REAL)')
INDEX_CACHE = {}


def get_db():
    db = getattr(g, '_database', None)
    if db is None:
        db = g._database = sqlite3.connect(DATABASE, isolation_level=None)
    return db


def query_db(query, args=()):
    cursor = get_db().execute(query, args)
    query_type = query.lower().strip().split()[0]
    result = None
    if query_type == 'select':
        result = cursor.fetchall()
    elif query_type in ['update', 'insert', 'replace', 'delete']:
        result = cursor.rowcount
        get_db().commit()
    cursor.close()
    return result


def fetch_one(query, args=()):
    result = query_db(query, args)
    if result:
        return result[0]


def ensure_index(name):
    global INDEX_CACHE
    if not index_exists(name):
        create_index(name)
    else:
        INDEX_CACHE[name] = index_hash(name)
    return INDEX_CACHE[name]


def h(value):
    if not isinstance(value, bytes):
        value = value.encode()
    return sha256(value).hexdigest()


def index_hash(name):
    return 'index_%s' % h(name)


def create_index(name):
    global INDEX_CACHE
    idx_hash = index_hash(name)
    query_db(CREATE_INDEX % {'index': idx_hash})
    INDEX_CACHE[name] = idx_hash


def index_exists(name):
    global INDEX_CACHE
    if name in INDEX_CACHE:
        return True
    return query_db('SELECT name FROM sqlite_master'
                    ' WHERE type=? AND name= ? COLLATE NOCASE',
                    ('table', index_hash(name)))


@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, '_database', None)
    if db is not None:
        db.close()


def index_store(index, key, value):
    key_hash = h(key)
    rows = query_db('INSERT OR REPLACE INTO %s(key, key_hash, value, modified)'
                    ' VALUES(?, ?, ?, ?)' % index,
                    (key, key_hash, value, time()))
    return key, rows


class Home(Resource):
    def get(self):
        return {'hello': 'world'}

    def post(self):
        return {'hello': 'world'}


class Index(Resource):
    def get(self, index, key):
        """
        http :5000/user/alice
        """
        index = ensure_index(index)
        row = fetch_one('SELECT * FROM %s WHERE key_hash = ?' % index,
                        (h(key),))
        if not row:
            abort(404)
        key, _, value, time = row
        value = value.decode()
        return {'key': key, 'value': value, 'modified': time}

    def delete(self, index, key):
        """
        http delete :5000/user/alice
        """
        index = ensure_index(index)
        rows = query_db('DELETE FROM %s WHERE key_hash = ?' % index,
                        (h(key),))
        if not rows:
            abort(404)
        return {'key': key, 'deleted': rows}

    def post(self, index, key):
        """
        http :5000/user/alice email=alice@test.com name=Alice
        curl -H "Content-Type: application/json" -X POST -d bad :5000/test/x
        """
        index = ensure_index(index)
        raw = request.get_data()
        try:
            # Validate JSON
            value = request.get_json()
        except BadRequest:
            abort(400)
        key, rows = index_store(index, key, raw)
        if not rows:
            abort(404)
        return {'key': key, 'updated': True}


api.add_resource(Home, '/')
api.add_resource(Index, '/<string:index>/<string:key>')

if __name__ == '__main__':
    app.run(debug=True)