From f92776e64625b3be4c29ee18f8f277e7e79ea5a9 Mon Sep 17 00:00:00 2001
From: simon987 <me@simon987.net>
Date: Fri, 22 May 2020 18:56:35 -0400
Subject: [PATCH] add volatile state

---
 hexlib/db.py               | 36 ++++++++++++++++++++++++++++++++++
 setup.py                   |  4 ++--
 test/test_VolatileState.py | 40 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 2 deletions(-)
 create mode 100644 test/test_VolatileState.py

diff --git a/hexlib/db.py b/hexlib/db.py
index 6384330..77bb672 100644
--- a/hexlib/db.py
+++ b/hexlib/db.py
@@ -1,5 +1,7 @@
 import base64
 import sqlite3
+import redis
+import ujson as json
 
 
 class PersistentState:
@@ -15,6 +17,40 @@ class PersistentState:
         return Table(self, table)
 
 
+class VolatileState:
+    """Quick and dirty volatile dict-like redis wrapper"""
+    def __init__(self, prefix, ttl=3600, **redis_args):
+        self.rdb = redis.Redis(**redis_args)
+        self.prefix = prefix
+        self.ttl = 3600
+
+    def __getitem__(self, table):
+        return RedisTable(self, table)
+
+
+class RedisTable:
+    def __init__(self, state, table):
+        self._state = state
+        self._table = table
+
+    def __setitem__(self, key, value):
+        self._state.rdb.set(self._state.prefix + self._table + ":" + str(key), json.dumps(value), ex=self._state.ttl)
+
+    def __getitem__(self, key):
+        val = self._state.rdb.get(self._state.prefix + self._table + ":" + str(key))
+        if val:
+            return json.loads(val)
+        return None
+
+    def __delitem__(self, key):
+        self._state.rdb.delete(self._state.prefix + self._table + ":" + str(key))
+
+    def __iter__(self):
+        for key in self._state.rdb.scan_iter(self._state.prefix + self._table + "*"):
+            val = self._state.rdb.get(key)
+            yield json.loads(val) if val else None
+
+
 class Table:
     def __init__(self, state, table):
         self._state = state
diff --git a/setup.py b/setup.py
index 29c03a2..59e6f72 100644
--- a/setup.py
+++ b/setup.py
@@ -2,12 +2,12 @@ from setuptools import setup
 
 setup(
     name="hexlib",
-    version="1.7",
+    version="1.8",
     description="Misc utility methods",
     author="simon987",
     author_email="me@simon987.net",
     packages=["hexlib"],
     install_requires=[
-        "ImageHash", "influxdb", "siphash", "python-dateutil"
+        "ImageHash", "influxdb", "siphash", "python-dateutil", "redis", "ujson"
     ]
 )
diff --git a/test/test_VolatileState.py b/test/test_VolatileState.py
new file mode 100644
index 0000000..399d550
--- /dev/null
+++ b/test/test_VolatileState.py
@@ -0,0 +1,40 @@
+from unittest import TestCase
+from hexlib.db import VolatileState
+
+
+class TestVolatileState(TestCase):
+
+    def test_get_set(self):
+        s = VolatileState(prefix="test1")
+        val = {
+            "field1": 1,
+            "arr1": [1,2,3]
+        }
+
+        s["a"]["1"] = val
+
+        self.assertDictEqual(val, s["a"]["1"])
+
+    def test_iter(self):
+        s = VolatileState(prefix="test2")
+
+        s["b"]["1"] = 1
+        s["b"]["2"] = 2
+        s["b"]["3"] = 3
+        s["b"]["4"] = 4
+
+        self.assertEqual(sum(s["b"]), 10)
+
+    def test_int_key(self):
+        s = VolatileState(prefix="test2")
+        s["b"][1] = 1
+        del s["b"][1]
+
+    def test_delete(self):
+        s = VolatileState(prefix="test3")
+
+        s["c"]["1"] = 1
+        self.assertIsNotNone(s["c"]["1"])
+        del s["c"]["1"]
+        self.assertIsNone(s["c"]["1"])
+