You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

352 lines
9.7 KiB

4 years ago
  1. #!/usr/bin/env python
  2. """ PickleShare - a small 'shelve' like datastore with concurrency support
  3. Like shelve, a PickleShareDB object acts like a normal dictionary. Unlike
  4. shelve, many processes can access the database simultaneously. Changing a
  5. value in database is immediately visible to other processes accessing the
  6. same database.
  7. Concurrency is possible because the values are stored in separate files. Hence
  8. the "database" is a directory where *all* files are governed by PickleShare.
  9. Example usage::
  10. from pickleshare import *
  11. db = PickleShareDB('~/testpickleshare')
  12. db.clear()
  13. print "Should be empty:",db.items()
  14. db['hello'] = 15
  15. db['aku ankka'] = [1,2,313]
  16. db['paths/are/ok/key'] = [1,(5,46)]
  17. print db.keys()
  18. del db['aku ankka']
  19. This module is certainly not ZODB, but can be used for low-load
  20. (non-mission-critical) situations where tiny code size trumps the
  21. advanced features of a "real" object database.
  22. Installation guide: pip install pickleshare
  23. Author: Ville Vainio <vivainio@gmail.com>
  24. License: MIT open source license.
  25. """
  26. from __future__ import print_function
  27. __version__ = "0.7.5"
  28. try:
  29. from pathlib import Path
  30. except ImportError:
  31. # Python 2 backport
  32. from pathlib2 import Path
  33. import os,stat,time
  34. try:
  35. import collections.abc as collections_abc
  36. except ImportError:
  37. import collections as collections_abc
  38. try:
  39. import cPickle as pickle
  40. except ImportError:
  41. import pickle
  42. import errno
  43. import sys
  44. if sys.version_info[0] >= 3:
  45. string_types = (str,)
  46. else:
  47. string_types = (str, unicode)
  48. def gethashfile(key):
  49. return ("%02x" % abs(hash(key) % 256))[-2:]
  50. _sentinel = object()
  51. class PickleShareDB(collections_abc.MutableMapping):
  52. """ The main 'connection' object for PickleShare database """
  53. def __init__(self,root):
  54. """ Return a db object that will manage the specied directory"""
  55. if not isinstance(root, string_types):
  56. root = str(root)
  57. root = os.path.abspath(os.path.expanduser(root))
  58. self.root = Path(root)
  59. if not self.root.is_dir():
  60. # catching the exception is necessary if multiple processes are concurrently trying to create a folder
  61. # exists_ok keyword argument of mkdir does the same but only from Python 3.5
  62. try:
  63. self.root.mkdir(parents=True)
  64. except OSError as e:
  65. if e.errno != errno.EEXIST:
  66. raise
  67. # cache has { 'key' : (obj, orig_mod_time) }
  68. self.cache = {}
  69. def __getitem__(self,key):
  70. """ db['key'] reading """
  71. fil = self.root / key
  72. try:
  73. mtime = (fil.stat()[stat.ST_MTIME])
  74. except OSError:
  75. raise KeyError(key)
  76. if fil in self.cache and mtime == self.cache[fil][1]:
  77. return self.cache[fil][0]
  78. try:
  79. # The cached item has expired, need to read
  80. with fil.open("rb") as f:
  81. obj = pickle.loads(f.read())
  82. except:
  83. raise KeyError(key)
  84. self.cache[fil] = (obj,mtime)
  85. return obj
  86. def __setitem__(self,key,value):
  87. """ db['key'] = 5 """
  88. fil = self.root / key
  89. parent = fil.parent
  90. if parent and not parent.is_dir():
  91. parent.mkdir(parents=True)
  92. # We specify protocol 2, so that we can mostly go between Python 2
  93. # and Python 3. We can upgrade to protocol 3 when Python 2 is obsolete.
  94. with fil.open('wb') as f:
  95. pickle.dump(value, f, protocol=2)
  96. try:
  97. self.cache[fil] = (value, fil.stat().st_mtime)
  98. except OSError as e:
  99. if e.errno != errno.ENOENT:
  100. raise
  101. def hset(self, hashroot, key, value):
  102. """ hashed set """
  103. hroot = self.root / hashroot
  104. if not hroot.is_dir():
  105. hroot.mkdir()
  106. hfile = hroot / gethashfile(key)
  107. d = self.get(hfile, {})
  108. d.update( {key : value})
  109. self[hfile] = d
  110. def hget(self, hashroot, key, default = _sentinel, fast_only = True):
  111. """ hashed get """
  112. hroot = self.root / hashroot
  113. hfile = hroot / gethashfile(key)
  114. d = self.get(hfile, _sentinel )
  115. #print "got dict",d,"from",hfile
  116. if d is _sentinel:
  117. if fast_only:
  118. if default is _sentinel:
  119. raise KeyError(key)
  120. return default
  121. # slow mode ok, works even after hcompress()
  122. d = self.hdict(hashroot)
  123. return d.get(key, default)
  124. def hdict(self, hashroot):
  125. """ Get all data contained in hashed category 'hashroot' as dict """
  126. hfiles = self.keys(hashroot + "/*")
  127. hfiles.sort()
  128. last = len(hfiles) and hfiles[-1] or ''
  129. if last.endswith('xx'):
  130. # print "using xx"
  131. hfiles = [last] + hfiles[:-1]
  132. all = {}
  133. for f in hfiles:
  134. # print "using",f
  135. try:
  136. all.update(self[f])
  137. except KeyError:
  138. print("Corrupt",f,"deleted - hset is not threadsafe!")
  139. del self[f]
  140. self.uncache(f)
  141. return all
  142. def hcompress(self, hashroot):
  143. """ Compress category 'hashroot', so hset is fast again
  144. hget will fail if fast_only is True for compressed items (that were
  145. hset before hcompress).
  146. """
  147. hfiles = self.keys(hashroot + "/*")
  148. all = {}
  149. for f in hfiles:
  150. # print "using",f
  151. all.update(self[f])
  152. self.uncache(f)
  153. self[hashroot + '/xx'] = all
  154. for f in hfiles:
  155. p = self.root / f
  156. if p.name == 'xx':
  157. continue
  158. p.unlink()
  159. def __delitem__(self,key):
  160. """ del db["key"] """
  161. fil = self.root / key
  162. self.cache.pop(fil,None)
  163. try:
  164. fil.unlink()
  165. except OSError:
  166. # notfound and permission denied are ok - we
  167. # lost, the other process wins the conflict
  168. pass
  169. def _normalized(self, p):
  170. """ Make a key suitable for user's eyes """
  171. return str(p.relative_to(self.root)).replace('\\','/')
  172. def keys(self, globpat = None):
  173. """ All keys in DB, or all keys matching a glob"""
  174. if globpat is None:
  175. files = self.root.rglob('*')
  176. else:
  177. files = self.root.glob(globpat)
  178. return [self._normalized(p) for p in files if p.is_file()]
  179. def __iter__(self):
  180. return iter(self.keys())
  181. def __len__(self):
  182. return len(self.keys())
  183. def uncache(self,*items):
  184. """ Removes all, or specified items from cache
  185. Use this after reading a large amount of large objects
  186. to free up memory, when you won't be needing the objects
  187. for a while.
  188. """
  189. if not items:
  190. self.cache = {}
  191. for it in items:
  192. self.cache.pop(it,None)
  193. def waitget(self,key, maxwaittime = 60 ):
  194. """ Wait (poll) for a key to get a value
  195. Will wait for `maxwaittime` seconds before raising a KeyError.
  196. The call exits normally if the `key` field in db gets a value
  197. within the timeout period.
  198. Use this for synchronizing different processes or for ensuring
  199. that an unfortunately timed "db['key'] = newvalue" operation
  200. in another process (which causes all 'get' operation to cause a
  201. KeyError for the duration of pickling) won't screw up your program
  202. logic.
  203. """
  204. wtimes = [0.2] * 3 + [0.5] * 2 + [1]
  205. tries = 0
  206. waited = 0
  207. while 1:
  208. try:
  209. val = self[key]
  210. return val
  211. except KeyError:
  212. pass
  213. if waited > maxwaittime:
  214. raise KeyError(key)
  215. time.sleep(wtimes[tries])
  216. waited+=wtimes[tries]
  217. if tries < len(wtimes) -1:
  218. tries+=1
  219. def getlink(self,folder):
  220. """ Get a convenient link for accessing items """
  221. return PickleShareLink(self, folder)
  222. def __repr__(self):
  223. return "PickleShareDB('%s')" % self.root
  224. class PickleShareLink:
  225. """ A shortdand for accessing nested PickleShare data conveniently.
  226. Created through PickleShareDB.getlink(), example::
  227. lnk = db.getlink('myobjects/test')
  228. lnk.foo = 2
  229. lnk.bar = lnk.foo + 5
  230. """
  231. def __init__(self, db, keydir ):
  232. self.__dict__.update(locals())
  233. def __getattr__(self,key):
  234. return self.__dict__['db'][self.__dict__['keydir']+'/' + key]
  235. def __setattr__(self,key,val):
  236. self.db[self.keydir+'/' + key] = val
  237. def __repr__(self):
  238. db = self.__dict__['db']
  239. keys = db.keys( self.__dict__['keydir'] +"/*")
  240. return "<PickleShareLink '%s': %s>" % (
  241. self.__dict__['keydir'],
  242. ";".join([Path(k).basename() for k in keys]))
  243. def main():
  244. import textwrap
  245. usage = textwrap.dedent("""\
  246. pickleshare - manage PickleShare databases
  247. Usage:
  248. pickleshare dump /path/to/db > dump.txt
  249. pickleshare load /path/to/db < dump.txt
  250. pickleshare test /path/to/db
  251. """)
  252. DB = PickleShareDB
  253. import sys
  254. if len(sys.argv) < 2:
  255. print(usage)
  256. return
  257. cmd = sys.argv[1]
  258. args = sys.argv[2:]
  259. if cmd == 'dump':
  260. if not args: args= ['.']
  261. db = DB(args[0])
  262. import pprint
  263. pprint.pprint(db.items())
  264. elif cmd == 'load':
  265. cont = sys.stdin.read()
  266. db = DB(args[0])
  267. data = eval(cont)
  268. db.clear()
  269. for k,v in db.items():
  270. db[k] = v
  271. elif cmd == 'testwait':
  272. db = DB(args[0])
  273. db.clear()
  274. print(db.waitget('250'))
  275. elif cmd == 'test':
  276. test()
  277. stress()
  278. if __name__== "__main__":
  279. main()