Checkable Shelf for python3

safeshelve.py

#coding: utf8
# safeshelve
#
# Copyright (c) 2015, multiSnow <infinity.blick.winkel@gmail.com>
#
# Permission to use, copy, modify, and/or distribute this software for
# any purpose with or without fee is hereby granted, provided that the
# above copyright notice and this permission notice appear in all
# copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
# WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
# AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
# DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
# PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
# TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.

from collections.abc import MutableMapping
from hashlib import sha512
from shelve import Shelf
from threading import RLock

# maybe use more efficiently sql driver (pgsql...)
import sqlite3

# split something from sql driver
dbopen=sqlite3.connect
DBOpError=sqlite3.OperationalError
Connection=sqlite3.Connection
keycode='utf8'
# e.g. '%s' for pgsql
placeholder='?'
# True if sql driver can run in threading in python
sqlthreadsafe=False

class CheckError(Exception):
    def __init__(self,key):
        self.key=key
    def __str__(self):
        return 'data checksum error: {}'.format(self.key)

class DoubleKeyError(Exception):
    def __init__(self,key):
        self.key=key
    def __str__(self):
        return 'double key in database: {}'.format(self.key)

class KeyTypeError(Exception):
    def __init__(self,warg):
        self.warg=warg
    def __str__(self):
        return 'wrong type: {}'.format(type(self.warg).__name__)

# return sha512 of data, or None for empty data
def _digest(data):
    return sha512(data).hexdigest() if data else None

def _cmd(s):
    return s.format(ph=placeholder)

class _DLock:
    def __init__(self,nolock=sqlthreadsafe):
        self.lock=None if nolock else RLock()

    def __enter__(self):
        if self.lock:
            self.lock.acquire()
    def __exit__(self,t,v,tb):
        if self.lock:
            self.lock.release()

class SimpleConnection(Connection):
    def __init__(self,*args,**kwargs):
        super(SimpleConnection,self).__init__(*args,**kwargs)
        self.lock=_DLock()
        self._cursor=self.cursor()
        try:
            self._exc('create table _main (key blob,hashsum text,item blob)')
        except DBOpError:
            # _main table already in database
            # TODO: a batter way to check existing table?
            pass
        self.datacache={}
        self.indexcache=list(rkey for rkey, in self._exc('select key from _main'))

    def _exc(self,*args,**kwargs):
        with self.lock:
            return self._cursor.execute(*args,**kwargs)

    def _rkey(self,key):
        # force using bytes
        if type(key) is bytes:
            rkey=key
        elif type(key) is str:
            rkey=key.encode(keycode)
        else:
            raise KeyTypeError(key)
        if self.indexcache.count(rkey)>1:
            raise DoubleKeyError(key)
        return rkey

    def __len__(self):
        return len(self.indexcache)

    def __iter__(self):
        yield from self.indexcache

    def keys(self):
        yield from iter(self)

    def checkbad(self,key):
        # check data of key using sha512, return True if data is damaged
        rkey=self._rkey(key)
        for hashsum,item in self._exc(_cmd('select hashsum,item from _main where key={ph}'),
                                      (rkey,)):
            if _digest(item)==hashsum:
                self.datacache[rkey]=item
                return
            else:
                return True
        raise KeyError(key)

    def checkall(self):
        # check all of data using sha512, raise CheckError of first damaged key of data
        for rkey in self.indexcache:
            if self.checkbad(rkey):
                raise CheckError(rkey)

    def __setitem__(self,key,item):
        rkey=self._rkey(key)
        if rkey in self.indexcache:
            self._exc(_cmd('update _main set hashsum={ph},item={ph} where key={ph}'),
                         (_digest(item),item,rkey))
        else:
            self._exc(_cmd('insert into _main values ({ph},{ph},{ph})'),
                      (rkey,_digest(item),item))
            self.indexcache.append(rkey)

    def __getitem__(self,key):
        rkey=self._rkey(key)
        if rkey in self.datacache:
            return self.datacache[rkey]
        for hashsum,item in self._exc(_cmd('select hashsum,item from _main where key={ph}'),
                                      (rkey,)):
            if _digest(item)==hashsum:
                self.datacache[rkey]=item
                return item
            raise CheckError(key)
        raise KeyError(key)

    def __delitem__(self,key):
        rkey=self._rkey(key)
        if rkey in self.indexcache:
            self._exc(_cmd('delete from _main where key={ph}'),
                      (rkey,))
            del self.indexcache[self.indexcache.index(rkey)]
            if rkey in self.datacache:
                del self.datacache[rkey]
        else:
            raise KeyError(key)

    def clear(self):
        self._exc('delete from _main')
        self.datacache.clear()
        self.indexcache.clear()

    def popitem(self):
        try:
            rkey=self.indexcache[0]
        except IndexError:
            raise KeyError
        item=self[rkey]
        del self[rkey]
        return rkey,item

    def update(self,*args,**kwargs):
        cache=dict()
        cache.update(*args,**kwargs)
        for key,item in cache.items():
            self[self._rkey(key)]=item

    def vacuum(self):
        # do vacuum
        self._exc('vacuum')

def sqlopen(filename=None,check_same_thread=True):
    # create database, in memory if not specify the filename
    return dbopen(filename if filename else ':memory:',
                  factory=SimpleConnection,check_same_thread=check_same_thread)

class SafeShelf(MutableMapping):
    # safely Shelf, almost the same with DbfilenameShelf
    # startcheck for checkall at class init
    # endclean for vaccum at class close
    def __init__(self,filename,protocol=3,writeback=False,check_same_thread=True,
                 startcheck=False,endclean=False):
        super(SafeShelf,self).__init__()
        self.db=sqlopen(filename,check_same_thread=check_same_thread)
        self.shelf=Shelf(self.db,protocol=protocol,writeback=writeback)
        self.endclean=endclean
        if startcheck:
            self.db.checkall()

    def __getitem__(self,key):return self.shelf[key]
    def __setitem__(self,key,item):self.shelf[key]=item
    def __delitem__(self,key):del self.shelf[key]
    def __iter__(self):yield from iter(self.shelf)
    def __contains__(self,key):return key in self.shelf
    def __len__(self):return len(self.shelf)

    def get(self,key,default=None):return self.shelf.get(key,default=default)
    def clear(self):return self.db.clear()
    def checkbad(self,key):return self.db.checkbad(key)
    def checkall(self):return self.db.checkall()
    def vacuum(self):return self.db.vacuum()

    def sync(self):
        self.shelf.sync()
        self.db.commit()

    def close(self):
        self.sync()
        if self.endclean:
            self.vacuum()
        self.shelf.close()

    def __enter__(self):return self
    def __exit__(self,exc_type,exc_value,traceback):self.close()