[didex.git] / didex / index.py

import struct, contextlib
from . import db, lib
from .db import bd

deadlock = bd.DBLockDeadlockError
notfound = bd.DBNotFoundError

class simpletype(object):
    def __init__(self, encode, decode):
        self.enc = encode
        self.dec = decode

    def encode(self, ob):
        return self.enc(ob)
    def decode(self, dat):
        return self.dec(dat)
    def compare(self, a, b):
        if a < b:
            return -1
        elif a > b:
            return 1
        else:
            return 0

    @classmethod
    def struct(cls, fmt):
        return cls(lambda ob: struct.pack(fmt, ob),
                   lambda dat: struct.unpack(fmt, dat)[0])

class maybe(object):
    def __init__(self, bk):
        self.bk = bk

    def encode(self, ob):
        if ob is None: return b""
        return b"\0" + self.bk.encode(ob)
    def decode(self, dat):
        if dat == b"": return None
        return self.bk.dec(dat[1:])
    def compare(self, a, b):
        if a is b is None:
            return 0
        elif a is None:
            return -1
        elif b is None:
            return 1
        else:
            return self.bk.compare(a[1:], b[1:])

t_int = simpletype.struct(">Q")

class index(object):
    def __init__(self, db, name, datatype):
        self.db = db
        self.nm = name
        self.typ = datatype

missing = object()

class ordered(index, lib.closable):
    def __init__(self, db, name, datatype, duplicates, create=True):
        super().__init__(db, name, datatype)
        self.dup = duplicates
        fl = bd.DB_THREAD | bd.DB_AUTO_COMMIT
        if create: fl |= bd.DB_CREATE
        def initdb(db):
            def compare(a, b):
                if a == b == "": return 0
                return self.typ.compare(self.typ.decode(a), self.typ.decode(b))
            db.set_flags(bd.DB_DUPSORT)
            db.set_bt_compare(compare)
        self.bk = db._opendb("i-" + name, bd.DB_BTREE, fl, initdb)
        self.bk.set_get_returns_none(False)

    def close(self):
        self.bk.close()

    class cursor(lib.closable):
        def __init__(self, idx, cur, item, stop):
            self.idx = idx
            self.cur = cur
            self.item = item
            self.stop = stop

        def close(self):
            if self.cur is not None:
                self.cur.close()

        def __iter__(self):
            return self

        def peek(self):
            if self.item is None:
                raise StopIteration()
            rk, rv = self.item
            rk = self.idx.typ.decode(rk)
            rv = struct.unpack(">Q", rv)[0]
            if self.stop(rk):
                self.item = None
                raise StopIteration()
            return rk, rv

        def __next__(self):
            rk, rv = self.peek()
            try:
                while True:
                    try:
                        self.item = self.cur.next()
                        break
                    except deadlock:
                        continue
            except notfound:
                self.item = None
            return rk, rv

        def skip(self, n=1):
            try:
                for i in range(n):
                    next(self)
            except StopIteration:
                return

    def get(self, *, match=missing, ge=missing, gt=missing, lt=missing, le=missing, all=False):
        while True:
            try:
                cur = self.bk.cursor()
                done = False
                try:
                    if match is not missing:
                        try:
                            k, v = cur.set(self.typ.encode(match))
                        except notfound:
                            return self.cursor(None, None, None, None)
                        else:
                            done = True
                            return self.cursor(self, cur, (k, v), lambda o: (self.typ.compare(o, match) != 0))
                    elif all:
                        try:
                            k, v = cur.first()
                        except notfound:
                            return self.cursor(None, None, None, None)
                        else:
                            done = True
                            return self.cursor(self, cur, (k, v), lambda o: False)
                    elif ge is not missing or gt is not missing or lt is not missing or le is not missing:
                        skip = False
                        try:
                            if ge is not missing:
                                k, v = cur.set_range(self.typ.encode(ge))
                            elif gt is not missing:
                                k, v = cur.set_range(self.typ.encode(gt))
                                skip = True
                            else:
                                k, v = cur.first()
                        except notfound:
                            return self.cursor(None, None, None, None)
                        if lt is not missing:
                            stop = lambda o: self.typ.compare(o, lt) >= 0
                        elif le is not missing:
                            stop = lambda o: self.typ.compare(o, le) > 0
                        else:
                            stop = lambda o: False
                        ret = self.cursor(self, cur, (k, v), stop)
                        if skip:
                            try:
                                while self.typ.compare(ret.peek()[0], gt) == 0:
                                    next(ret)
                            except StopIteration:
                                pass
                        done = True
                        return ret
                    else:
                        raise NameError("invalid get() specification")
                finally:
                    if not done:
                        cur.close()
            except deadlock:
                continue

    def put(self, key, id):
        while True:
            try:
                with db.txn(self.db.env.env) as tx:
                    obid = struct.pack(">Q", id)
                    if not self.db.ob.has_key(obid, txn=tx.tx):
                        raise ValueError("no such object in database: " + str(id))
                    try:
                        self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA)
                    except bd.DBKeyExistError:
                        return False
                    tx.commit()
                    return True
            except deadlock:
                continue

    def remove(self, key, id):
        while True:
            try:
                with db.txn(self.db.env.env) as tx:
                    obid = struct.pack(">Q", id)
                    if not self.db.ob.has_key(obid, txn=tx.tx):
                        raise ValueError("no such object in database: " + str(id))
                    cur = self.bk.cursor(txn=tx.tx)
                    try:
                        try:
                            cur.get_both(self.typ.encode(key), obid)
                        except notfound:
                            return False
                        cur.delete()
                    finally:
                        cur.close()
                    tx.commit()
                    return True
            except deadlock:
                continue
Commit	Line	Data
a95055e8 FT	1	import struct, contextlib
	2	from . import db, lib
	3	from .db import bd
	4
	5	deadlock = bd.DBLockDeadlockError
	6	notfound = bd.DBNotFoundError
	7
	8	class simpletype(object):
	9	def __init__(self, encode, decode):
	10	self.enc = encode
	11	self.dec = decode
	12
	13	def encode(self, ob):
	14	return self.enc(ob)
	15	def decode(self, dat):
	16	return self.dec(dat)
	17	def compare(self, a, b):
	18	if a < b:
	19	return -1
	20	elif a > b:
	21	return 1
	22	else:
	23	return 0
	24
	25	@classmethod
	26	def struct(cls, fmt):
	27	return cls(lambda ob: struct.pack(fmt, ob),
	28	lambda dat: struct.unpack(fmt, dat)[0])
	29
	30	class maybe(object):
	31	def __init__(self, bk):
	32	self.bk = bk
	33
	34	def encode(self, ob):
	35	if ob is None: return b""
	36	return b"\0" + self.bk.encode(ob)
	37	def decode(self, dat):
	38	if dat == b"": return None
	39	return self.bk.dec(dat[1:])
	40	def compare(self, a, b):
	41	if a is b is None:
	42	return 0
	43	elif a is None:
	44	return -1
	45	elif b is None:
	46	return 1
	47	else:
	48	return self.bk.compare(a[1:], b[1:])
	49
	50	t_int = simpletype.struct(">Q")
	51
	52	class index(object):
	53	def __init__(self, db, name, datatype):
	54	self.db = db
	55	self.nm = name
	56	self.typ = datatype
	57
	58	missing = object()
	59
	60	class ordered(index, lib.closable):
	61	def __init__(self, db, name, datatype, duplicates, create=True):
	62	super().__init__(db, name, datatype)
	63	self.dup = duplicates
	64	fl = bd.DB_THREAD \| bd.DB_AUTO_COMMIT
65	if create: fl \|= bd.DB_CREATE
66	def initdb(db):
67	def compare(a, b):
68	if a == b == "": return 0
69	return self.typ.compare(self.typ.decode(a), self.typ.decode(b))
70	db.set_flags(bd.DB_DUPSORT)
71	db.set_bt_compare(compare)
72	self.bk = db._opendb("i-" + name, bd.DB_BTREE, fl, initdb)
73	self.bk.set_get_returns_none(False)
74
75	def close(self):
76	self.bk.close()
77
78	class cursor(lib.closable):
79	def __init__(self, idx, cur, item, stop):
80	self.idx = idx
81	self.cur = cur
82	self.item = item
83	self.stop = stop
84
85	def close(self):
86	if self.cur is not None:
87	self.cur.close()
88
89	def __iter__(self):
90	return self
91
92	def peek(self):
93	if self.item is None:
94	raise StopIteration()
95	rk, rv = self.item
96	rk = self.idx.typ.decode(rk)
97	rv = struct.unpack(">Q", rv)[0]
98	if self.stop(rk):
99	self.item = None
100	raise StopIteration()
101	return rk, rv
102
103	def __next__(self):
104	rk, rv = self.peek()
105	try:
106	while True:
107	try:
108	self.item = self.cur.next()
109	break
110	except deadlock:
111	continue
112	except notfound:
113	self.item = None
114	return rk, rv
115
116	def skip(self, n=1):
117	try:
118	for i in range(n):
119	next(self)
120	except StopIteration:
121	return
122
123	def get(self, *, match=missing, ge=missing, gt=missing, lt=missing, le=missing, all=False):
124	while True:
125	try:
126	cur = self.bk.cursor()
127	done = False
128	try:
129	if match is not missing:
130	try:
131	k, v = cur.set(self.typ.encode(match))
132	except notfound:
133	return self.cursor(None, None, None, None)
134	else:
135	done = True
136	return self.cursor(self, cur, (k, v), lambda o: (self.typ.compare(o, match) != 0))
137	elif all:
138	try:
139	k, v = cur.first()
140	except notfound:
141	return self.cursor(None, None, None, None)
142	else:
143	done = True
144	return self.cursor(self, cur, (k, v), lambda o: False)
145	elif ge is not missing or gt is not missing or lt is not missing or le is not missing:
146	skip = False
147	try:
148	if ge is not missing:
149	k, v = cur.set_range(self.typ.encode(ge))
150	elif gt is not missing:
151	k, v = cur.set_range(self.typ.encode(gt))
152	skip = True
153	else:
154	k, v = cur.first()
155	except notfound:
156	return self.cursor(None, None, None, None)
157	if lt is not missing:
158	stop = lambda o: self.typ.compare(o, lt) >= 0
159	elif le is not missing:
160	stop = lambda o: self.typ.compare(o, le) > 0
161	else:
162	stop = lambda o: False
163	ret = self.cursor(self, cur, (k, v), stop)
164	if skip:
165	try:
166	while self.typ.compare(ret.peek()[0], gt) == 0:
167	next(ret)
168	except StopIteration:
169	pass
170	done = True
171	return ret
172	else:
173	raise NameError("invalid get() specification")
174	finally:
175	if not done:
176	cur.close()
177	except deadlock:
178	continue
179
180	def put(self, key, id):
181	while True:
182	try:
183	with db.txn(self.db.env.env) as tx:
184	obid = struct.pack(">Q", id)
185	if not self.db.ob.has_key(obid, txn=tx.tx):
186	raise ValueError("no such object in database: " + str(id))
187	try:
188	self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA)
189	except bd.DBKeyExistError:
190	return False
191	tx.commit()
192	return True
193	except deadlock:
194	continue
195
196	def remove(self, key, id):
197	while True:
198	try:
199	with db.txn(self.db.env.env) as tx:
200	obid = struct.pack(">Q", id)
201	if not self.db.ob.has_key(obid, txn=tx.tx):
202	raise ValueError("no such object in database: " + str(id))
203	cur = self.bk.cursor(txn=tx.tx)
204	try:
205	try:
206	cur.get_both(self.typ.encode(key), obid)
207	except notfound:
208	return False
209	cur.delete()
210	finally:
211	cur.close()
212	tx.commit()
213	return True
214	except deadlock:
215	continue