initial commit

db/db.py

@@ -0,0 +1,95 @@
+# python
+import os
+import threading
+from tinydb import TinyDB
+
+DB_ENV = os.environ.get('DB_ENV', 'prod')
+base_dir = os.path.dirname(__file__)
+
+class LockedTable:
+    """
+    Thread-safe wrapper around a TinyDB table. All callable attribute access
+    is wrapped to acquire a reentrant lock while calling the underlying method.
+    Non-callable attributes are returned directly.
+    """
+    def __init__(self, table):
+        self._table = table
+        self._lock = threading.RLock()
+
+    def __getattr__(self, name):
+        # avoid proxying internal attrs
+        if name in ('_table', '_lock'):
+            return super().__getattribute__(name)
+
+        attr = getattr(self._table, name)
+
+        if callable(attr):
+            def locked_call(*args, **kwargs):
+                with self._lock:
+                    return attr(*args, **kwargs)
+            return locked_call
+        return attr
+
+    # convenience explicit methods (ensure these are class methods, not top-level)
+    def insert(self, *args, **kwargs):
+        with self._lock:
+            return self._table.insert(*args, **kwargs)
+
+    def insert_multiple(self, *args, **kwargs):
+        with self._lock:
+            return self._table.insert_multiple(*args, **kwargs)
+
+    def search(self, *args, **kwargs):
+        with self._lock:
+            return self._table.search(*args, **kwargs)
+
+    def get(self, *args, **kwargs):
+        with self._lock:
+            return self._table.get(*args, **kwargs)
+
+    def all(self, *args, **kwargs):
+        with self._lock:
+            return self._table.all(*args, **kwargs)
+
+    def remove(self, *args, **kwargs):
+        with self._lock:
+            return self._table.remove(*args, **kwargs)
+
+    def update(self, *args, **kwargs):
+        with self._lock:
+            return self._table.update(*args, **kwargs)
+
+    def truncate(self):
+        with self._lock:
+            return self._table.truncate()
+
+# Setup DB files next to this module
+
+if DB_ENV == 'test':
+    child_path = os.path.join(base_dir, 'test_children.json')
+    task_path = os.path.join(base_dir, 'test_tasks.json')
+    reward_path = os.path.join(base_dir, 'test_rewards.json')
+    image_path = os.path.join(base_dir, 'test_images.json')
+else:
+    child_path = os.path.join(base_dir, 'children.json')
+    task_path = os.path.join(base_dir, 'tasks.json')
+    reward_path = os.path.join(base_dir, 'rewards.json')
+    image_path = os.path.join(base_dir, 'images.json')
+
+# Use separate TinyDB instances/files for each collection
+_child_db = TinyDB(child_path, indent=2)
+_task_db = TinyDB(task_path, indent=2)
+_reward_db = TinyDB(reward_path, indent=2)
+_image_db = TinyDB(image_path, indent=2)
+
+# Expose table objects wrapped with locking
+child_db = LockedTable(_child_db)
+task_db = LockedTable(_task_db)
+reward_db = LockedTable(_reward_db)
+image_db = LockedTable(_image_db)
+
+if DB_ENV == 'test':
+    child_db.truncate()
+    task_db.truncate()
+    reward_db.truncate()
+    image_db.truncate()