Added convenience get() wrappers.

[didex.git] / didex / db.py

import time, threading, struct
from . import lib
from bsddb3 import db as bd

__all__ = ["environment", "database"]

deadlock = bd.DBLockDeadlockError

class environment(lib.closable):
    def __init__(self, path, *, create=True, recover=False, mode=0o666):
        self.env = bd.DBEnv()
        self.env.set_lk_detect(bd.DB_LOCK_RANDOM)
        fl = bd.DB_THREAD | bd.DB_INIT_MPOOL | bd.DB_INIT_LOCK | bd.DB_INIT_LOG | bd.DB_INIT_TXN
        if recover:
            fl |= bd.DB_RECOVER
        if create:
            fl |= bd.DB_CREATE
        self.env.open(path, fl, mode)
        self.lastckp = self.lastarch = time.time()
        self.lock = threading.Lock()
        self.dbs = {}

    def close(self):
        env = self.env
        if env is None:
            return
        env.close()
        self.env = None

    def maint(self):
        now = time.time()
        try:
            if now - self.lastckp > 60:
                self.env.txn_checkpoint(1024)
                self.lastckp = now
            if now - self.lastarch > 3600:
                self.env.log_archive(bd.DB_ARCH_REMOVE)
                self.lastarch = now
        except deadlock:
            pass

    def db(self, name, create=True, mode=0o666):
        with self.lock:
            if name not in self.dbs:
                self.dbs[name] = database(self, name, create, mode)
            return self.dbs[name]

    def __del__(self):
        self.close()
    def __enter__(self):
        return self
    def __exit__(self, *excinfo):
        self.close()
        return False

def opendb(env, fnm, dnm, typ, fl, mode):
    ret = bd.DB(env)
    while True:
        try:
            self.main.open(fnm, dnm, typ, fl, mode)
        except deadlock:
            continue
        return ret

class txn(object):
    def __init__(self, env, flags=bd.DB_TXN_WRITE_NOSYNC):
        self.tx = env.txn_begin(None, flags)
        self.done = False
        self.pcommit = set()

    def commit(self):
        self.done = True
        self.tx.commit(0)
        def run1(list):
            if len(list) > 0:
                try:
                    list[0]()
                finally:
                    run1(list[1:])
        run1(list(self.pcommit))

    def abort(self):
        self.done = True
        self.tx.abort()

    def __enter__(self):
        return self

    def __exit__(self, etype, exc, tb):
        if not self.done:
            self.abort()
        return False

    def postcommit(self, fun):
        self.pcommit.add(fun)

def dloopfun(fun):
    def wrapper(self, *args, **kwargs):
        while True:
            try:
                return fun(self, *args, **kwargs)
            except deadlock:
                continue
    return wrapper

def txnfun(envfun):
    def fxf(fun):
        def wrapper(self, *args, tx=None, **kwargs):
            if tx is None:
                while True:
                    try:
                        with txn(envfun(self)) as ltx:
                            ret = fun(self, *args, tx=ltx, **kwargs)
                            ltx.commit()
                            return ret
                    except deadlock:
                        continue
            else:
                return fun(self, *args, tx=tx, **kwargs)
        return wrapper
    return fxf

class database(object):
    def __init__(self, env, name, create, mode):
        self.env = env
        self.mode = mode
        self.fnm = name
        fl = bd.DB_THREAD | bd.DB_AUTO_COMMIT
        if create:
            fl |= bd.DB_CREATE
        self.cf = self._opendb("cf", bd.DB_HASH, fl)
        self.ob = self._opendb("ob", bd.DB_HASH, fl)

    def _opendb(self, dnm, typ, fl, init=None):
        ret = bd.DB(self.env.env)
        if init: init(ret)
        while True:
            try:
                ret.open(self.fnm, dnm, typ, fl, self.mode)
            except deadlock:
                continue
            return ret

    @txnfun(lambda self: self.env.env)
    def _nextseq(self, *, tx):
        if self.cf.has_key(b"seq", txn=tx.tx):
            seq = struct.unpack(">Q", self.cf.get(b"seq", txn=tx.tx))[0]
        else:
            seq = 1
        self.cf.put(b"seq", struct.pack(">Q", seq + 1), txn=tx.tx)
        return seq

    @txnfun(lambda self: self.env.env)
    def add(self, ob, *, tx):
        seq = self._nextseq(tx=tx)
        self.ob.put(struct.pack(">Q", seq), ob, txn=tx.tx, flags=bd.DB_NOOVERWRITE)
        return seq

    @txnfun(lambda self: self.env.env)
    def replace(self, id, ob, *, tx):
        key = struct.pack(">Q", id)
        if not self.ob.has_key(key, txn=tx.tx):
            raise KeyError(id)
        self.ob.put(key, ob, txn=tx.tx)

    @txnfun(lambda self: self.env.env)
    def get(self, id, *, tx):
        ret = self.ob.get(struct.pack(">Q", id), None)
        if ret is None:
            raise KeyError(id)
        return ret

    @txnfun(lambda self: self.env.env)
    def remove(self, id, *, tx):
        key = struct.pack(">Q", id)
        if not self.ob.has_key(key, txn=tx.tx):
            raise KeyError(id)
        self.ob.delete(key, txn=tx.tx)