8481c0f6d6a8fb7ff4c6fac828fc7586483e9c3a
[didex.git] / didex / store.py
1 import threading, pickle
2 from . import db, index, cache
3 from .db import txnfun
4
5 __all__ = ["environment", "datastore", "autostore"]
6
7 class environment(object):
8     def __init__(self, *, path=None, getpath=None, recover=False):
9         if path is not None:
10             self.path = path
11             self.getpath = None
12         else:
13             self.path = None
14             self.getpath = getpath
15         self.recover = recover
16         self.lk = threading.Lock()
17         self.bk = None
18
19     def __call__(self):
20         with self.lk:
21             if self.bk is None:
22                 if self.path is None:
23                     self.path = self.getpath()
24                 self.bk = db.environment(self.path, recover=self.recover)
25             return self.bk
26
27     def close(self):
28         with self.lk:
29             if self.bk is not None:
30                 self.bk.close()
31                 self.bk = None
32
33 class storedesc(object):
34     pass
35
36 def storedescs(obj):
37     t = type(obj)
38     ret = getattr(t, "__didex_attr", None)
39     if ret is None:
40         ret = []
41         for nm, val in t.__dict__.items():
42             if isinstance(val, storedesc):
43                 ret.append((nm, val))
44         t.__didex_attr = ret
45     return ret
46
47 class datastore(object):
48     def __init__(self, name, *, env=None, path=".", ncache=None):
49         self.name = name
50         self.lk = threading.Lock()
51         if env:
52             self.env = env
53         else:
54             self.env = environment(path=path)
55         self._db = None
56         if ncache is None:
57             ncache = cache.cache()
58         self.cache = ncache
59         self.cache.load = self._load
60
61     def db(self):
62         with self.lk:
63             if self._db is None:
64                 self._db = self.env().db(self.name)
65             return self._db
66
67     def _load(self, id):
68         try:
69             return pickle.loads(self.db().get(id))
70         except:
71             raise KeyError(id, "could not unpickle data")
72
73     def _encode(self, obj):
74         return pickle.dumps(obj)
75
76     def get(self, id, *, load=True):
77         return self.cache.get(id, load=load)
78
79     @txnfun(lambda self: self.db().env.env)
80     def register(self, obj, *, tx):
81         id = self.db().add(self._encode(obj), tx=tx)
82         for nm, attr in storedescs(obj):
83             attr.register(id, obj, tx)
84         self.cache.put(id, obj)
85         return id
86
87     @txnfun(lambda self: self.db().env.env)
88     def unregister(self, id, *, vfy=None, tx):
89         obj = self.get(id)
90         if vfy is not None and obj is not vfy:
91             raise RuntimeError("object identity crisis: " + str(vfy) + " is not cached object " + obj)
92         for nm, attr in storedescs(obj):
93             attr.unregister(id, obj, tx)
94         self.db().remove(id, tx=tx)
95         self.cache.remove(id)
96
97     @txnfun(lambda self: self.db().env.env)
98     def update(self, id, *, vfy=None, tx):
99         obj = self.get(id, load=False)
100         if vfy is not None and obj is not vfy:
101             raise RuntimeError("object identity crisis: " + str(vfy) + " is not cached object " + obj)
102         for nm, attr, in storedescs(obj):
103             attr.update(id, obj, tx)
104         self.db().replace(id, self._encode(obj), tx=tx)
105
106 class autotype(type):
107     def __call__(self, *args, **kwargs):
108         new = super().__call__(*args, **kwargs)
109         new.id = self.store.register(new)
110         return new
111
112 class autostore(object, metaclass=autotype):
113     def __init__(self):
114         self.id = None
115
116     def save(self):
117         self.store.update(self.id, vfy=self)
118
119     def remove(self):
120         self.store.unregister(self.id, vfy=self)
121         self.id = None