diff --git a/screenless/bureau/bureau.py b/screenless/bureau/bureau.py index 9ab10b2..77c2c4f 100644 --- a/screenless/bureau/bureau.py +++ b/screenless/bureau/bureau.py @@ -4,12 +4,13 @@ import inspect import json import logging import os.path +import random +import string import subprocess import sys import tempfile import textwrap import threading -import time import lmdb import PIL @@ -70,6 +71,66 @@ class LogPrinter(logging.Handler): prn.cut() +class KeyValStore(object): + """ + A KeyValStore is a simple wrapper for LMDB flat file storage. It's very + fast and simple for large databases with small (less than 4kb) entries. + If you need something larger try the filesystem. If you need more structure + or indexes try sqlite. Keys and values MUST BE UNICODE STRINGS! + """ + def __init__(self, env, name): + self.env = env + self.db = env.open_db(name.encode()) + + def store(self, key, val): + """ + Store a key-val pair. + Returns True on success. + """ + with self.env.begin(write=True, db=self.db) as txn: + ret = txn.put(key.encode(), val.encode()) + return ret + + def store_and_get_shortcode(self, val): + """ + Find an un-used shortcode and use it as a key to store the value given. + Note, each db is limited to about a billion keys so don't go too crazy. + returns a 5-char shortcode string. + """ + def _shortcode(): + # returns a random 5-char string + return ''.join(random.choice(string.ascii_letters + string.digits) + for _ in range(5)) + + # we only have about a billion so make sure we don't collide keys + with self.env.begin(write=True, db=self.db) as txn: + res = "not None" + while res is not None: + tmpcode = _shortcode() + res = txn.get(tmpcode.encode()) + txn.put(tmpcode.encode(), val.encode()) + + return tmpcode + + def get(self, key): + """ + Look up a value. + Returns value as a unicode string or None if nonexistent. + """ + with self.env.begin(db=self.db) as txn: + res = txn.get(key.encode()) + return res.decode("utf-8") + + def delete(self, key): + """ + Delete a key-val pair. + Returns True on success. + """ + with self.env.begin(write=True, db=self.db) as txn: + ret = txn.delete(key.encode) + return ret + + class Bureau(object): """ Bureau is a base class that implements standard methods for inter-bureau communication, IO, registration and some convenient stuff @@ -139,6 +200,14 @@ class Bureau(object): self.log.error("CRASH TRACE: {0}".format(str(value)), exc_info=(typ, value, tb)) sys.__excepthook__(typ, value, tb) + def open_db(self, name): + """ + Loads and if not yet existing, creates, an LMDB database + returns a KeyValStore object + """ + db = KeyValStore(self.dbenv, name) + return db + def load_config(self): """ load (or reload) config data from file