tiramisu/tiramisu/storage/sqlite3/storage.py

# -*- coding: utf-8 -*-
"default plugin for cache: set it in a simple dictionary"
# Copyright (C) 2013 Team tiramisu (see AUTHORS for all contributors)
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# ____________________________________________________________

from pickle import dumps, loads
from os import unlink
from os.path import basename, splitext, join
import sqlite3
from glob import glob


class Setting(object):
    extension = 'db'
    dir_database = '/tmp'


setting = Setting()


def _gen_filename(name):
    return join(setting.dir_database, '{0}.{1}'.format(name,
                                                       setting.extension))


def list_sessions():
    names = []
    for filename in glob(_gen_filename('*')):
        names.append(basename(splitext(filename)[0]))
    return names


def delete_session(session_id):
    unlink(_gen_filename(session_id))


class Storage(object):
    __slots__ = ('_conn', '_cursor', 'is_persistent', '_session_id')
    storage = 'sqlite3'

    def __init__(self, session_id, is_persistent):
        self.is_persistent = is_persistent
        self._session_id = session_id
        self._conn = sqlite3.connect(_gen_filename(self._session_id))
        self._conn.text_factory = str
        self._cursor = self._conn.cursor()

    def execute(self, sql, params=None, commit=True):
        if params is None:
            params = tuple()
        self._cursor.execute(sql, params)
        if commit:
            self._conn.commit()

    def select(self, sql, params=None, only_one=True):
        self.execute(sql, params=params, commit=False)
        if only_one:
            return self._cursor.fetchone()
        else:
            return self._cursor.fetchall()

    def __del__(self):
        self._cursor.close()
        self._conn.close()
        if not self.is_persistent:
            delete_session(self._session_id)


class Cache(object):
    __slots__ = ('storage',)
    key_is_path = True

    def __init__(self, cache_type, storage):
        self.storage = storage
        cache_table = 'CREATE TABLE IF NOT EXISTS cache_{0}(path '.format(
            cache_type)
        cache_table += 'text primary key, value text, time real)'
        self.storage.execute(cache_table)

    # value
    def _sqlite_decode_path(self, path):
        if path == '_none':
            return None
        else:
            return path

    def _sqlite_encode_path(self, path):
        if path is None:
            return '_none'
        else:
            return path

    def _sqlite_decode(self, value):
        return loads(value)

    def _sqlite_encode(self, value):
        if isinstance(value, list):
            value = list(value)
        return dumps(value)

    def setcache(self, cache_type, path, val, time):
        convert_value = self._sqlite_encode(val)
        path = self._sqlite_encode_path(path)
        self.storage.execute("DELETE FROM cache_{0} WHERE path = ?".format(
            cache_type), (path,), False)
        self.storage.execute("INSERT INTO cache_{0}(path, value, time) "
                             "VALUES (?, ?, ?)".format(cache_type),
                             (path, convert_value, time))

    def getcache(self, cache_type, path, exp):
        path = self._sqlite_encode_path(path)
        cached = self.storage.select("SELECT value FROM cache_{0} WHERE "
                                     "path = ? AND time >= ?".format(
                                         cache_type), (path, exp))
        if cached is None:
            return False, None
        else:
            return True, self._sqlite_decode(cached[0])

    def hascache(self, cache_type, path):
        path = self._sqlite_encode_path(path)
        return self.storage.select("SELECT value FROM cache_{0} WHERE "
                                   "path = ?".format(cache_type),
                                   (path,)) is not None

    def reset_expired_cache(self, cache_type, exp):
        self.storage.execute("DELETE FROM cache_{0} WHERE time < ?".format(
            cache_type), (exp,))

    def reset_all_cache(self, cache_type):
        self.storage.execute("DELETE FROM cache_{0}".format(cache_type))

    def get_cached(self, cache_type, context):
        """return all values in a dictionary
        example: {'path1': ('value1', 'time1'), 'path2': ('value2', 'time2')}
        """
        ret = {}
        for path, value, time in self.storage.select("SELECT * FROM cache_{0}"
                                                     "".format(cache_type),
                                                     only_one=False):
            path = self._sqlite_decode_path(path)
            value = self._sqlite_decode(value)
            ret[path] = (value, time)
        return ret