SaltmModule to store key/value in sqlite

kvsqlite.py

# -*- coding: utf-8 -*-
from __future__ import absolute_import

# Import python libs
import logging
import codecs

# Import third party libs
try:
    import msgpack
    import sqlite3
    IMPORTS_OK = True
except ImportError:
    IMPORTS_OK = False

log = logging.getLogger(__name__)


def __virtual__():
    '''
    Only load this module if sqlite3 is installed
    '''
    if IMPORTS_OK:
        return True
    else:
        return False


def _connect(bucket):
    '''
    Returns a sqlite3 cursor.
    '''
    sqlitedb = __salt__['config.option']('saltdb.sqlitedb', '/srv/saltdb/db')
    db = sqlite3.connect(sqlitedb)
    db.text_factory = str
    cursor = db.cursor()
    try:
        q = 'CREATE TABLE {0} (key string not null unique, value text)'.format(
                _escape(bucket))
        cursor.execute(q)
    except sqlite3.OperationalError:
        pass
    q = 'CREATE UNIQUE INDEX IF NOT EXISTS {0}_key_idx ON {0} (key)'.format(
            _escape(bucket))
    cursor.execute(q)
    return db, cursor


def _escape(s, errors="strict"):
    encodable = s.encode("utf-8", errors).decode("utf-8")
    nul_index = encodable.find("\x00")
    if nul_index >= 0:
        error = UnicodeEncodeError("NUL-terminated utf-8", encodable,
                                   nul_index, nul_index + 1, "NUL not allowed")
        error_handler = codecs.lookup_error(errors)
        replacement, _ = error_handler(error)
        encodable = encodable.replace("\x00", replacement)
    return encodable.replace("\"", "\"\"")


def update(key, value, bucket='salt'):
    '''
    Update or set the value for a key

    CLI Example:

    .. code-block:: bash

        salt '*' kvsqlite.update key value [bucket=bucket_name]
    '''
    db, cursor = _connect(bucket)
    value = msgpack.packs(value)
    cursor.execute('INSERT OR REPLACE INTO {0} VALUES (?, ?)'.format(
        _escape(bucket)), (key, value))
    db.commit()
    cursor.close()
    db.close()
    return True


def getval(key, default=None, bucket='salt'):
    db, cursor = _connect(bucket)
    query = 'SELECT value FROM {0} WHERE key=?'.format(_escape(bucket))
    res = cursor.execute(query, (key,)).fetchone()

    if res:
        ret = msgpack.loads(res[0])
    else:
        ret = None
    cursor.close()
    db.close()
    return ret