174 lines
4.5 KiB
Python
174 lines
4.5 KiB
Python
# 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
|