Initial import.

[didex.git] / didex / db.py

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

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

    def commit(self):
        self.done = True
        self.tx.commit(0)

    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

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

    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

    def add(self, ob):
        while True:
            try:
                with txn(self.env.env) as tx:
                    seq = self._nextseq(tx)
                    self.ob.put(struct.pack(">Q", seq), ob, txn=tx.tx, flags=bd.DB_NOOVERWRITE)
                    tx.commit()
                    return seq
            except deadlock:
                continue

    def replace(self, id, ob):
        while True:
            try:
                with txn(self.env.env) as 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)
                    tx.commit()
                    return
            except deadlock:
                continue

    def get(self, id):
        while True:
            try:
                return self.ob[struct.pack(">Q", id)]
            except KeyError:
                raise KeyError(id) from None
            except deadlock:
                continue

    def remove(self, id):
        while True:
            try:
                with txn(self.env.env) as tx:
                    self.ob.delete(struct.pack(">Q", id), txn=tx.tx)
                    tx.commit()
                    return
            except deadlock:
                continue