restricting access by context (#796)

* restricting access by context

* restricted context updates

---------

Co-authored-by: Wojtek <wskozlowski@dbzero.io>

Author: wskozlowski

dbzero/dbzero/dbzero.pyi

@@ -30,6 +30,7 @@ def open(prefix_name: str, open_mode: str = "rw", **kwargs: Any) -> None:
         
         * autocommit (bool) to disable automatic commits for this prefix
         * restricted (bool) to restrict memo reflection access for this prefix
+        * restricted_context (ContextVar) to dynamically restrict memo reflection access for this prefix
         * slab_size (int) for memory slab allocation size in bytes
         * meta_io_step_size (int) for metadata I/O operation chunk size
         * lock_flags (dict) to change locking behavior when opening the prefix for read-write
@@ -55,6 +56,14 @@ def open(prefix_name: str, open_mode: str = "rw", **kwargs: Any) -> None:
     """
     ...
 
+def set_restricted(*, restricted: Optional[bool] = None, restricted_context: Any = None,
+                   prefix: Optional[str] = None) -> None:
+    """Increase memo restricted-access mode globally or for an open prefix.
+
+    The restriction level can only increase: unrestricted -> restricted_context -> restricted.
+    """
+    ...
+
 def close(prefix_name: Optional[str] = None) -> None:
     """Gracefully shut down dbzero, persisting changes and releasing resources.
 
@@ -1545,7 +1554,7 @@ def get_config() -> Dict[str, Any]:
         
         * autocommit (bool): True if autocommit enabled. Defaults to True.
         * autocommit_interval (int): Milliseconds between commits. Defaults to 250.
-        * restricted (bool): True if future opened prefixes default to restricted mode.
+        * restricted (bool): True if future opened prefixes default to static restricted mode.
 
     Raises
     ------

dbzero/dbzero/initialization.py

@@ -26,6 +26,7 @@ def init(dbzero_root: str, **kwargs: Any) -> None:
         
         Configure global dbzero behavior:
         * restricted (bool, default False) to restrict memo reflection access for future opened prefixes
+        * restricted_context (ContextVar) to dynamically restrict memo reflection access
         * autocommit (bool, default True) to enable automatic commits
         * autocommit_interval (int, default 367) for commit interval in milliseconds
         * cache_size (int, default 2 GiB) for main object cache size in bytes
@@ -56,6 +57,9 @@ def init(dbzero_root: str, **kwargs: Any) -> None:
     if "restricted" in kwargs:
         init_kwargs["restricted"] = kwargs["restricted"]
 
+    if "restricted_context" in kwargs:
+        init_kwargs["restricted_context"] = kwargs["restricted_context"]
+
     rpc = kwargs.get("rpc")
     if rpc is not None and not isinstance(rpc, Mapping):
         raise TypeError("rpc must be a mapping")
@@ -71,4 +75,6 @@ def init(dbzero_root: str, **kwargs: Any) -> None:
         open_kwargs = {}
         if "restricted" in kwargs:
             open_kwargs["restricted"] = kwargs["restricted"]
+        if "restricted_context" in kwargs:
+            open_kwargs["restricted_context"] = kwargs["restricted_context"]
         dbzero_open(kwargs["prefix"], open_mode=open_mode, **open_kwargs)

dbzero/setup.py

@@ -10,7 +10,7 @@
 
 setup(
     name='dbzero',
-    version='0.4.1',
+    version='0.4.2',
     description='DBZero community edition',
     packages=['dbzero'],
     python_requires='>=3.9',

pyproject.toml

@@ -4,7 +4,7 @@ requires = ['meson-python']
 
 [project]
 name = 'dbzero'
-version = '0.4.1'
+version = '0.4.2'
 description = 'A state management system for Python 3.x that unifies your applications business logic, data persistence, and caching into a single, efficient layer.'
 readme = 'README.md'
 requires-python = '>=3.9'

python_tests/test_restricted_memo_access.py

@@ -7,6 +7,7 @@
 import subprocess
 import sys
 import textwrap
+from contextvars import ContextVar
 
 import pytest
 import dbzero as db0
@@ -200,3 +201,208 @@ def test_global_restricted_applies_to_auto_opened_prefix(db0_restricted_root):
     assert db0.get_config()["restricted"] is True
     assert db0.get_prefix_stats("auto-open-prefix")["restricted"] is True
     _assert_restricted(fetched, writable=False)
+
+
+def test_global_restricted_context_controls_existing_objects(db0_restricted_root):
+    restricted_context = ContextVar("restricted_context", default=False)
+    db0.init(DB0_DIR, restricted_context=restricted_context)
+    assert db0.get_config()["restricted"] is False
+    db0.open("dynamic-global-prefix")
+    assert db0.get_prefix_stats("dynamic-global-prefix")["restricted"] is False
+    obj = RestrictedMemoAccessData(123)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+    assert obj.get_value.__self__ is obj
+
+    token = restricted_context.set(True)
+    try:
+        _assert_restricted(obj)
+    finally:
+        restricted_context.reset(token)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+    assert obj.get_value.__self__ is obj
+
+
+def test_prefix_restricted_context_overrides_global_context(db0_restricted_root):
+    global_context = ContextVar("global_context", default=True)
+    prefix_context = ContextVar("prefix_context", default=False)
+    db0.init(DB0_DIR, restricted_context=global_context)
+
+    db0.open("global-dynamic-prefix")
+    global_obj = RestrictedMemoAccessData(123)
+
+    db0.open("prefix-dynamic-prefix", restricted_context=prefix_context)
+    prefix_obj = RestrictedMemoAccessData(123)
+
+    _assert_restricted(global_obj)
+    assert prefix_obj.__class__ is RestrictedMemoAccessData
+    assert prefix_obj.get_value.__self__ is prefix_obj
+
+    token = prefix_context.set(True)
+    try:
+        _assert_restricted(prefix_obj)
+    finally:
+        prefix_context.reset(token)
+
+
+def test_prefix_restricted_context_none_cannot_replace_existing_context(db0_restricted_root):
+    global_context = ContextVar("global_context", default=True)
+    prefix_context = ContextVar("prefix_context", default=False)
+    db0.init(DB0_DIR, restricted_context=global_context)
+    db0.open("prefix-reset-dynamic-prefix", restricted_context=prefix_context)
+    obj = RestrictedMemoAccessData(123)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+    assert obj.get_value.__self__ is obj
+
+    with pytest.raises(RuntimeError):
+        db0.open("prefix-reset-dynamic-prefix", restricted_context=None)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+
+
+def test_static_restricted_without_restricted_context_remains_restricted(db0_restricted_root):
+    db0.init(DB0_DIR)
+    db0.open("static-restricted-prefix", restricted=True)
+    assert db0.get_prefix_stats("static-restricted-prefix")["restricted"] is True
+
+    _assert_restricted(RestrictedMemoAccessData(123))
+
+
+def test_static_restricted_object_does_not_resolve_later_restricted_context(db0_restricted_root):
+    class ExplodingContext:
+        def get(self):
+            raise AssertionError("restricted_context should not be resolved")
+
+    db0.init(DB0_DIR)
+    db0.open("static-no-context-prefix", restricted=True)
+    obj = RestrictedMemoAccessData(123)
+
+    with pytest.raises(RuntimeError):
+        db0.open("static-no-context-prefix", restricted_context=ExplodingContext())
+
+    _assert_restricted(obj)
+
+
+def test_unset_restricted_context_is_treated_as_false(db0_restricted_root):
+    restricted_context = ContextVar("restricted_context")
+    db0.init(DB0_DIR, restricted_context=restricted_context)
+    db0.open("unset-dynamic-prefix")
+    obj = RestrictedMemoAccessData(123)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+    assert obj.get_value.__self__ is obj
+
+
+def test_restricted_context_uses_python_truthiness(db0_restricted_root):
+    restricted_context = ContextVar("restricted_context", default=[])
+    db0.init(DB0_DIR, restricted_context=restricted_context)
+    db0.open("truthy-dynamic-prefix")
+    obj = RestrictedMemoAccessData(123)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+
+    token = restricted_context.set([object()])
+    try:
+        _assert_restricted(obj)
+    finally:
+        restricted_context.reset(token)
+
+
+def test_init_prefix_shortcut_forwards_restricted_context(db0_restricted_root):
+    restricted_context = ContextVar("restricted_context", default=True)
+    db0.init(DB0_DIR, prefix="shortcut-dynamic-prefix", restricted_context=restricted_context)
+    assert db0.get_prefix_stats("shortcut-dynamic-prefix")["restricted"] is False
+
+    _assert_restricted(RestrictedMemoAccessData(123))
+
+
+def test_set_restricted_context_upgrades_existing_and_future_prefixes(db0_restricted_root):
+    restricted_context = ContextVar("restricted_context", default=False)
+    db0.init(DB0_DIR)
+    db0.open("dynamic-upgrade-prefix")
+    obj = RestrictedMemoAccessData(123)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+    db0.set_restricted(restricted_context=restricted_context)
+    assert db0.get_config()["restricted"] is False
+
+    assert obj.__class__ is RestrictedMemoAccessData
+    token = restricted_context.set(True)
+    try:
+        _assert_restricted(obj)
+        db0.open("future-dynamic-upgrade-prefix")
+        _assert_restricted(RestrictedMemoAccessData(123))
+    finally:
+        restricted_context.reset(token)
+
+
+def test_set_restricted_static_upgrades_dynamic_context(db0_restricted_root):
+    restricted_context = ContextVar("restricted_context", default=False)
+    db0.init(DB0_DIR)
+    db0.set_restricted(restricted_context=restricted_context)
+    db0.open("dynamic-to-static-prefix")
+    obj = RestrictedMemoAccessData(123)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+    db0.set_restricted(restricted=True)
+
+    _assert_restricted(obj)
+    db0.open("future-static-upgrade-prefix")
+    assert db0.get_prefix_stats("future-static-upgrade-prefix")["restricted"] is True
+    _assert_restricted(RestrictedMemoAccessData(123))
+
+
+def test_set_restricted_rejects_context_after_static_restricted(db0_restricted_root):
+    restricted_context = ContextVar("restricted_context", default=False)
+    db0.init(DB0_DIR, restricted=True)
+    db0.open("static-terminal-prefix")
+    obj = RestrictedMemoAccessData(123)
+
+    with pytest.raises(RuntimeError):
+        db0.set_restricted(restricted_context=restricted_context)
+
+    _assert_restricted(obj)
+
+
+def test_set_restricted_rejects_context_replacement(db0_restricted_root):
+    first_context = ContextVar("first_context", default=False)
+    second_context = ContextVar("second_context", default=True)
+    db0.init(DB0_DIR)
+    db0.set_restricted(restricted_context=first_context)
+    db0.open("context-replacement-prefix")
+    obj = RestrictedMemoAccessData(123)
+
+    with pytest.raises(RuntimeError):
+        db0.set_restricted(restricted_context=second_context)
+
+    assert obj.__class__ is RestrictedMemoAccessData
+    token = first_context.set(True)
+    try:
+        _assert_restricted(obj)
+    finally:
+        first_context.reset(token)
+
+
+def test_set_restricted_rejects_false(db0_restricted_root):
+    db0.init(DB0_DIR)
+
+    with pytest.raises(RuntimeError):
+        db0.set_restricted(restricted=False)
+
+
+def test_set_restricted_prefix_only_upgrades_target_prefix(db0_restricted_root):
+    restricted_context = ContextVar("restricted_context", default=True)
+    db0.init(DB0_DIR)
+
+    db0.open("target-prefix")
+    target = RestrictedMemoAccessData(123)
+    db0.open("other-prefix")
+    other = RestrictedMemoAccessData(123)
+
+    db0.set_restricted(prefix="target-prefix", restricted_context=restricted_context)
+
+    _assert_restricted(target)
+    assert other.__class__ is RestrictedMemoAccessData
+    assert other.get_value.__self__ is other

src/dbzero/bindings/python/Memo.cpp

@@ -71,17 +71,16 @@ namespace db0::python
         MemoImplT *memo_obj = reinterpret_cast<MemoImplT*>(py_type->tp_alloc(py_type, 0));
 
         // if type cannot be retrieved due to access mode then defer this operation (fallback)
-        auto restricted = fixture->isRestricted();
         if (type) {
             // prepare a new dbzero instance of a known db0 class
-            memo_obj->makeNew(type, restricted);
+            memo_obj->makeNew(type);
         } else {
             auto type_initializer = [py_type](db0::swine_ptr<Fixture> &fixture) {
                 auto &class_factory = fixture->get<db0::object_model::ClassFactory>();
                 return class_factory.getOrCreateType(py_type);
             };
             // prepare a new db0 instance of a known db0 class
-            memo_obj->makeNew(std::move(type_initializer), restricted);
+            memo_obj->makeNew(std::move(type_initializer));
         }
 
         return memo_obj;
@@ -192,18 +191,17 @@ namespace db0::python
         // find py type associated dbzero class with the ClassFactory
         auto type = class_factory.tryGetOrCreateType(py_type);
         MemoObject *memo_obj = nullptr;
-        auto restricted = fixture->isRestricted();
         if (type) {
             memo_obj = reinterpret_cast<MemoObject*>(py_type->tp_alloc(py_type, 0));
-            memo_obj->makeNew(type, restricted);
+            memo_obj->makeNew(type);
         } else {
             // if type cannot be retrieved due to access mode then deferr this operation (fallback)
             auto type_initializer = [py_type](db0::swine_ptr<Fixture> &fixture) {
                 auto &class_factory = fixture->get<db0::object_model::ClassFactory>();
                 return class_factory.getOrCreateType(py_type);
             };
             memo_obj = reinterpret_cast<MemoObject*>(py_type->tp_alloc(py_type, 0));
-            memo_obj->makeNew(std::move(type_initializer), restricted);
+            memo_obj->makeNew(std::move(type_initializer));
         }
 
         return memo_obj;
@@ -260,7 +258,18 @@ namespace db0::python
             return false;
         }
 
-        return memo_obj->ext().isRestricted();
+        auto fixture = memo_obj->ext().tryGetFixture();
+        if (!fixture) {
+            return false;
+        }
+
+        if (fixture->isRestricted()) {
+            return true;
+        }
+        if (!fixture->isRestrictedCtx()) {
+            return false;
+        }
+        return resolveRestrictedCtx(*fixture);
     }
 
     template <typename MemoImplT>
@@ -404,7 +413,18 @@ namespace db0::python
             return false;
         }
 
-        return memo_obj->ext().isRestricted() || isRestrictedMemoContextActive();
+        auto fixture = memo_obj->ext().tryGetFixture();
+        if (!fixture) {
+            return false;
+        }
+
+        if (fixture->isRestricted()) {
+            return true;
+        }
+        if (!fixture->isRestrictedCtx()) {
+            return false;
+        }
+        return resolveRestrictedCtx(*fixture);
     }
 
     template <typename MemoImplT>