implements new KeyValStore for convenient, fast databases for all bureaus

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