| 1 | import struct, contextlib, math |
| 2 | from . import db, lib |
| 3 | from .db import bd, txnfun |
| 4 | |
| 5 | __all__ = ["maybe", "t_int", "t_uint", "t_float", "t_str", "ordered"] |
| 6 | |
| 7 | deadlock = bd.DBLockDeadlockError |
| 8 | notfound = bd.DBNotFoundError |
| 9 | |
| 10 | class simpletype(object): |
| 11 | def __init__(self, encode, decode): |
| 12 | self.enc = encode |
| 13 | self.dec = decode |
| 14 | |
| 15 | def encode(self, ob): |
| 16 | return self.enc(ob) |
| 17 | def decode(self, dat): |
| 18 | return self.dec(dat) |
| 19 | def compare(self, a, b): |
| 20 | if a < b: |
| 21 | return -1 |
| 22 | elif a > b: |
| 23 | return 1 |
| 24 | else: |
| 25 | return 0 |
| 26 | |
| 27 | @classmethod |
| 28 | def struct(cls, fmt): |
| 29 | return cls(lambda ob: struct.pack(fmt, ob), |
| 30 | lambda dat: struct.unpack(fmt, dat)[0]) |
| 31 | |
| 32 | class maybe(object): |
| 33 | def __init__(self, bk): |
| 34 | self.bk = bk |
| 35 | |
| 36 | def encode(self, ob): |
| 37 | if ob is None: return b"" |
| 38 | return b"\0" + self.bk.encode(ob) |
| 39 | def decode(self, dat): |
| 40 | if dat == b"": return None |
| 41 | return self.bk.dec(dat[1:]) |
| 42 | def compare(self, a, b): |
| 43 | if a is b is None: |
| 44 | return 0 |
| 45 | elif a is None: |
| 46 | return -1 |
| 47 | elif b is None: |
| 48 | return 1 |
| 49 | else: |
| 50 | return self.bk.compare(a[1:], b[1:]) |
| 51 | |
| 52 | class compound(object): |
| 53 | def __init__(self, *parts): |
| 54 | self.parts = parts |
| 55 | |
| 56 | def encode(self, obs): |
| 57 | if len(obs) != len(self.parts): |
| 58 | raise ValueError("invalid length of compound data: " + str(len(obs)) + ", rather than " + len(self.parts)) |
| 59 | buf = bytearray() |
| 60 | for ob, part in zip(obs, self.parts): |
| 61 | dat = part.encode(ob) |
| 62 | if len(dat) < 128: |
| 63 | buf.append(0x80 | len(dat)) |
| 64 | buf.extend(dat) |
| 65 | else: |
| 66 | buf.extend(struct.pack(">i", len(dat))) |
| 67 | buf.extend(dat) |
| 68 | return bytes(buf) |
| 69 | def decode(self, dat): |
| 70 | ret = [] |
| 71 | off = 0 |
| 72 | for part in self.parts: |
| 73 | if dat[off] & 0x80: |
| 74 | ln = dat[off] & 0x7f |
| 75 | off += 1 |
| 76 | else: |
| 77 | ln = struct.unpack(">i", dat[off:off + 4])[0] |
| 78 | off += 4 |
| 79 | ret.append(part.decode(dat[off:off + len])) |
| 80 | off += len |
| 81 | return tuple(ret) |
| 82 | def compare(self, al, bl): |
| 83 | if (len(al) != len(self.parts)) or (len(bl) != len(self.parts)): |
| 84 | raise ValueError("invalid length of compound data: " + str(len(al)) + ", " + str(len(bl)) + ", rather than " + len(self.parts)) |
| 85 | for a, b, part in zip(al, bl, self.parts): |
| 86 | c = part.compare(a, b) |
| 87 | if c != 0: |
| 88 | return c |
| 89 | return 0 |
| 90 | |
| 91 | def floatcmp(a, b): |
| 92 | if math.isnan(a) and math.isnan(b): |
| 93 | return 0 |
| 94 | elif math.isnan(a): |
| 95 | return -1 |
| 96 | elif math.isnan(b): |
| 97 | return 1 |
| 98 | elif a < b: |
| 99 | return -1 |
| 100 | elif a > b: |
| 101 | return 1 |
| 102 | else: |
| 103 | return 0 |
| 104 | |
| 105 | t_int = simpletype.struct(">q") |
| 106 | t_uint = simpletype.struct(">Q") |
| 107 | t_float = simpletype.struct(">d") |
| 108 | t_float.compare = floatcmp |
| 109 | t_str = simpletype((lambda ob: ob.encode("utf-8")), (lambda dat: dat.decode("utf-8"))) |
| 110 | |
| 111 | class index(object): |
| 112 | def __init__(self, db, name, datatype): |
| 113 | self.db = db |
| 114 | self.nm = name |
| 115 | self.typ = datatype |
| 116 | |
| 117 | missing = object() |
| 118 | |
| 119 | class ordered(index, lib.closable): |
| 120 | def __init__(self, db, name, datatype, create=True): |
| 121 | super().__init__(db, name, datatype) |
| 122 | fl = bd.DB_THREAD | bd.DB_AUTO_COMMIT |
| 123 | if create: fl |= bd.DB_CREATE |
| 124 | def initdb(db): |
| 125 | def compare(a, b): |
| 126 | if a == b == "": return 0 |
| 127 | return self.typ.compare(self.typ.decode(a), self.typ.decode(b)) |
| 128 | db.set_flags(bd.DB_DUPSORT) |
| 129 | db.set_bt_compare(compare) |
| 130 | self.bk = db._opendb("i-" + name, bd.DB_BTREE, fl, initdb) |
| 131 | self.bk.set_get_returns_none(False) |
| 132 | |
| 133 | def close(self): |
| 134 | self.bk.close() |
| 135 | |
| 136 | class cursor(lib.closable): |
| 137 | def __init__(self, idx, cur, item, stop): |
| 138 | self.idx = idx |
| 139 | self.cur = cur |
| 140 | self.item = item |
| 141 | self.stop = stop |
| 142 | |
| 143 | def close(self): |
| 144 | if self.cur is not None: |
| 145 | self.cur.close() |
| 146 | |
| 147 | def __iter__(self): |
| 148 | return self |
| 149 | |
| 150 | def peek(self): |
| 151 | if self.item is None: |
| 152 | raise StopIteration() |
| 153 | rk, rv = self.item |
| 154 | rk = self.idx.typ.decode(rk) |
| 155 | rv = struct.unpack(">Q", rv)[0] |
| 156 | if self.stop(rk): |
| 157 | self.item = None |
| 158 | raise StopIteration() |
| 159 | return rk, rv |
| 160 | |
| 161 | def __next__(self): |
| 162 | rk, rv = self.peek() |
| 163 | try: |
| 164 | while True: |
| 165 | try: |
| 166 | self.item = self.cur.next() |
| 167 | break |
| 168 | except deadlock: |
| 169 | continue |
| 170 | except notfound: |
| 171 | self.item = None |
| 172 | return rk, rv |
| 173 | |
| 174 | def skip(self, n=1): |
| 175 | try: |
| 176 | for i in range(n): |
| 177 | next(self) |
| 178 | except StopIteration: |
| 179 | return |
| 180 | |
| 181 | def get(self, *, match=missing, ge=missing, gt=missing, lt=missing, le=missing, all=False): |
| 182 | while True: |
| 183 | try: |
| 184 | cur = self.bk.cursor() |
| 185 | done = False |
| 186 | try: |
| 187 | if match is not missing: |
| 188 | try: |
| 189 | k, v = cur.set(self.typ.encode(match)) |
| 190 | except notfound: |
| 191 | return self.cursor(None, None, None, None) |
| 192 | else: |
| 193 | done = True |
| 194 | return self.cursor(self, cur, (k, v), lambda o: (self.typ.compare(o, match) != 0)) |
| 195 | elif all: |
| 196 | try: |
| 197 | k, v = cur.first() |
| 198 | except notfound: |
| 199 | return self.cursor(None, None, None, None) |
| 200 | else: |
| 201 | done = True |
| 202 | return self.cursor(self, cur, (k, v), lambda o: False) |
| 203 | elif ge is not missing or gt is not missing or lt is not missing or le is not missing: |
| 204 | skip = False |
| 205 | try: |
| 206 | if ge is not missing: |
| 207 | k, v = cur.set_range(self.typ.encode(ge)) |
| 208 | elif gt is not missing: |
| 209 | k, v = cur.set_range(self.typ.encode(gt)) |
| 210 | skip = True |
| 211 | else: |
| 212 | k, v = cur.first() |
| 213 | except notfound: |
| 214 | return self.cursor(None, None, None, None) |
| 215 | if lt is not missing: |
| 216 | stop = lambda o: self.typ.compare(o, lt) >= 0 |
| 217 | elif le is not missing: |
| 218 | stop = lambda o: self.typ.compare(o, le) > 0 |
| 219 | else: |
| 220 | stop = lambda o: False |
| 221 | ret = self.cursor(self, cur, (k, v), stop) |
| 222 | if skip: |
| 223 | try: |
| 224 | while self.typ.compare(ret.peek()[0], gt) == 0: |
| 225 | next(ret) |
| 226 | except StopIteration: |
| 227 | pass |
| 228 | done = True |
| 229 | return ret |
| 230 | else: |
| 231 | raise NameError("invalid get() specification") |
| 232 | finally: |
| 233 | if not done: |
| 234 | cur.close() |
| 235 | except deadlock: |
| 236 | continue |
| 237 | |
| 238 | @txnfun(lambda self: self.db.env.env) |
| 239 | def put(self, key, id, *, tx): |
| 240 | obid = struct.pack(">Q", id) |
| 241 | if not self.db.ob.has_key(obid, txn=tx.tx): |
| 242 | raise ValueError("no such object in database: " + str(id)) |
| 243 | try: |
| 244 | self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA) |
| 245 | except bd.DBKeyExistError: |
| 246 | return False |
| 247 | return True |
| 248 | |
| 249 | @txnfun(lambda self: self.db.env.env) |
| 250 | def remove(self, key, id, *, tx): |
| 251 | obid = struct.pack(">Q", id) |
| 252 | if not self.db.ob.has_key(obid, txn=tx.tx): |
| 253 | raise ValueError("no such object in database: " + str(id)) |
| 254 | cur = self.bk.cursor(txn=tx.tx) |
| 255 | try: |
| 256 | try: |
| 257 | cur.get_both(self.typ.encode(key), obid) |
| 258 | except notfound: |
| 259 | return False |
| 260 | cur.delete() |
| 261 | finally: |
| 262 | cur.close() |
| 263 | return True |