sushy-tools/sushy_tools/emulator/memoize.py

# Copyright 2018 Red Hat, Inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

try:
    from collections import abc as collections

except ImportError:
    import collections

import contextlib
from functools import wraps
import os
import pickle
import sqlite3
import tempfile

import tenacity

# Python 3.8
MutableMapping = getattr(collections, 'abc', collections).MutableMapping


def memoize(permanent_cache=None):
    """Cache the return value of the decorated method.

    :param permanent_cache: a `dict` like object to use as a cache.
        If not given, the `._cache` attribute would be added to
        the object of the decorated method pointing to a newly
        created `dict`.
    :return: decorated function
    """

    def decorator(method):

        @wraps(method)
        def wrapped(self, *args, **kwargs):
            if permanent_cache is None:
                try:
                    cache = self._cache

                except AttributeError:
                    cache = self._cache = {}

            else:
                cache = permanent_cache

            method_cache = cache.setdefault(method, {})

            key = frozenset(args), frozenset(kwargs)

            try:
                return method_cache[key]

            except KeyError:
                rv = method(self, *args, **kwargs)
                method_cache[key] = rv
                return rv

        return wrapped

    return decorator


_retry = tenacity.retry(
    retry=tenacity.retry_if_exception_type(sqlite3.OperationalError),
    wait=tenacity.wait_exponential(min=0.1, max=2, multiplier=1),
    stop=tenacity.stop_after_delay(5),
    reraise=True)


class PersistentDict(MutableMapping):
    DBPATH = os.path.join(tempfile.gettempdir(), 'sushy-emulator')

    _dbpath = None

    def make_permanent(self, dbpath, dbfile):
        dbpath = dbpath or self.DBPATH
        os.makedirs(dbpath, exist_ok=True)

        self._dbpath = os.path.join(dbpath, dbfile) + '.sqlite'

        with self.connection() as cursor:
            cursor.execute(
                'create table if not exists cache '
                '(key blob primary key not null, value blob not null)'
            )

    @staticmethod
    def encode(obj):
        return pickle.dumps(obj)

    @staticmethod
    def decode(blob):
        return pickle.loads(blob)

    @contextlib.contextmanager
    def connection(self):
        if not self._dbpath:
            raise TypeError('Dict is not yet persistent')

        with sqlite3.connect(self._dbpath) as connection:
            yield connection.cursor()

    @_retry
    def __getitem__(self, key):
        key = self.encode(key)

        with self.connection() as cursor:
            cursor.execute(
                'select value from cache where key=?',
                (key,)
            )
            value = cursor.fetchone()

        if value is None:
            raise KeyError(key)

        return self.decode(value[0])

    @_retry
    def __setitem__(self, key, value):
        key = self.encode(key)
        value = self.encode(value)

        with self.connection() as cursor:
            cursor.execute(
                'insert or replace into cache values (?, ?)',
                (key, value)
            )

    @_retry
    def __delitem__(self, key):
        key = self.encode(key)

        with self.connection() as cursor:
            cursor.execute(
                'delete from cache where key=?',
                (key,)
            )
            if not cursor.rowcount:
                raise KeyError(key)

    @_retry
    def __iter__(self):
        with self.connection() as cursor:
            cursor.execute(
                'select key from cache'
            )
            records = cursor.fetchall()

        for r in records:
            yield self.decode(r[0])

    @_retry
    def __len__(self):
        with self.connection() as cursor:
            cursor.execute(
                'select count(*) from cache'
            )
            count = cursor.fetchone()[0]
            return count