From 73be0eecf98f9ade5ae2e627dbe1f58e16766b54 Mon Sep 17 00:00:00 2001
From: Brendan Howell <brendan@howell-ersatz.com>
Date: Mon, 26 Feb 2018 18:17:49 +0100
Subject: [PATCH] implements new KeyValStore for convenient, fast databases for
 all bureaus

---
 screenless/bureau/bureau.py | 71 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 70 insertions(+), 1 deletion(-)

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