feat: page entities with StableBTreeMap.range instead of probing every id

README.md

@@ -11,6 +11,7 @@ A lightweight key-value database with entity relationships and audit logging cap
 ## Features
 
 - **Persistent Storage**: Works with StableBTreeMap stable structure for persistent storage on your canister's stable memory so your data persists automatically across canister upgrades.
+- **Range-based paging**: `load_some`, `find`, `instances` and `rebuild_field_index` read a page of entities with a single ordered `range` call on the stable map (Basilisk `StableBTreeMap.range`), so cost tracks the rows returned rather than the highest id ever issued, and deleted ids are free. Backends without `range` (older CDK builds, plain dict test doubles) transparently fall back to one read per id.
 - **Entity-Relational Database**: Create, read and write entities with OneToOne, OneToMany, ManyToOne, and ManyToMany relationships.
 - **Schema Versioning & Upgrade Safety**: Automatic schema introspection, compatibility checking, and auto-migration for safe changes. Breaking changes without a `migrate()` method are rejected, preventing data corruption on canister upgrades.
 - **Entity Hooks**: Intercept and control entity lifecycle events (create, modify, delete) with `on_event` hooks.

ic_python_db/db_engine.py

@@ -9,10 +9,41 @@
 
 from ic_python_logging import get_logger
 
-from .storage import MemoryStorage, Storage
+from .storage import MemoryStorage, Storage, supports_range
 
 logger = get_logger(__name__)
 
+_DIGITS = frozenset("0123456789")
+
+
+def _is_sequential_id(id_str: str) -> bool:
+    """True for canonical sequential ids: ASCII digits, no leading zero."""
+    return bool(id_str) and id_str[0] != "0" and all(c in _DIGITS for c in id_str)
+
+
+def _next_sequential_id(id_str: str, digits: int) -> Optional[int]:
+    """Smallest ``digits``-digit sequential id that sorts after ``id_str``.
+
+    ``id_str`` is a key suffix returned by a range over the ``digits``-length
+    bucket, so it starts with a digit but may be a custom id such as
+    ``"1x"`` or ``"1-a"``. Returns None when no such id exists in the bucket.
+    """
+    i = 0
+    while i < len(id_str) and id_str[i] in _DIGITS:
+        i += 1
+    numeric_prefix, rest = id_str[:i], id_str[i:]
+    if not rest:
+        candidate = str(int(numeric_prefix) + 1)
+    elif rest[0] > "9":
+        # "1x" > every "1<digit>": jump to the next prefix, zero-filled
+        candidate = str(int(numeric_prefix) + 1) + "0" * len(rest)
+    else:
+        # "1-a" < "10": the zero-filled prefix is the next candidate
+        candidate = numeric_prefix + "0" * len(rest)
+    if len(candidate) > digits:
+        return None
+    return max(int(candidate), 10 ** (digits - 1))
+
 
 class Database:
     """Main database class providing high-level operations"""
@@ -181,6 +212,102 @@ def load(self, type_name: str, id: str) -> Optional[dict]:
             return json.loads(data)
         return None
 
+    def max_id(self, type_name: str) -> int:
+        """Highest sequential id ever issued for ``type_name`` (0 if none)."""
+        raw = self.load("_system", f"{type_name}_id")
+        return int(raw) if raw else 0
+
+    def load_page(
+        self,
+        type_name: str,
+        from_id: int,
+        count: int,
+        max_id: Optional[int] = None,
+    ) -> List[Tuple[str, dict]]:
+        """Persisted entities of ``type_name`` with sequential id >= ``from_id``.
+
+        Returns up to ``count`` ``(id, data)`` pairs in ascending numeric id
+        order, skipping ids that were deleted. Only canonical sequential ids
+        (``"1"``, ``"42"``; no leading zeros, no custom string ids) are
+        returned, matching what probing ``1..max_id`` would find.
+
+        When the storage backend supports ``range`` (Basilisk's
+        ``StableBTreeMap``, ``MemoryStorage``) a page costs one range call per
+        digit-count bucket touched (usually one), independent of how many ids
+        were deleted. Otherwise it falls back to one ``get`` per id.
+        """
+        if from_id < 1:
+            raise ValueError("from_id must be at least 1")
+        if count < 1:
+            raise ValueError("count must be at least 1")
+        if max_id is None:
+            max_id = self.max_id(type_name)
+        if max_id < from_id:
+            return []
+        if supports_range(self._db_storage):
+            return self._load_page_range(type_name, from_id, count, max_id)
+        return self._load_page_probe(type_name, from_id, count, max_id)
+
+    def _load_page_probe(
+        self, type_name: str, from_id: int, count: int, max_id: int
+    ) -> List[Tuple[str, dict]]:
+        results: List[Tuple[str, dict]] = []
+        entity_id = from_id
+        while len(results) < count and entity_id <= max_id:
+            raw = self._db_storage.get(f"{type_name}@{entity_id}")
+            if raw:
+                results.append((str(entity_id), json.loads(raw)))
+            entity_id += 1
+        return results
+
+    def _load_page_range(
+        self, type_name: str, from_id: int, count: int, max_id: int
+    ) -> List[Tuple[str, dict]]:
+        # Keys are "Type@<id>" with unpadded decimal ids. The backend orders
+        # str keys by byte length first, then bytewise (see storage_sort_key),
+        # so all ids with the same digit count form one contiguous, numerically
+        # ordered run. Walk one run ("bucket") per digit count, from the digit
+        # count of from_id up to that of max_id.
+        prefix = f"{type_name}@"
+        results: List[Tuple[str, dict]] = []
+        digits = len(str(from_id))
+        max_digits = len(str(max_id))
+        while digits <= max_digits and len(results) < count:
+            lo_id = max(from_id, 10 ** (digits - 1))
+            if digits == max_digits:
+                # Exclusive upper bound right after max_id (same digit count),
+                # or the whole bucket if max_id is the largest d-digit number.
+                hi_id = max_id + 1
+                end = (
+                    f"{prefix}{hi_id}"
+                    if len(str(hi_id)) == digits
+                    else f"{prefix}{':' * digits}"
+                )
+            else:
+                # ':' is the character after '9': every same-length key of
+                # the form "Type@<digits>" sorts before "Type@:::".
+                end = f"{prefix}{':' * digits}"
+
+            start = f"{prefix}{lo_id}"
+            plen = len(prefix)
+            while len(results) < count:
+                need = count - len(results)
+                page = self._db_storage.range(start, end, need)
+                for key, raw in page:
+                    id_str = key[plen:]
+                    if _is_sequential_id(id_str):
+                        results.append((id_str, json.loads(raw)))
+                if len(page) < need:
+                    break  # bucket exhausted
+                # Page was full but may have contained skipped custom ids:
+                # continue from the next possible sequential id in this bucket.
+                next_id = _next_sequential_id(page[-1][0][plen:], digits)
+                if next_id is None:
+                    break
+                start = f"{prefix}{next_id}"
+            digits += 1
+        return results
+
     def delete(self, type_name: str, entity_id: str) -> None:
         """Delete the data associated with the key
 

ic_python_db/entity.py

@@ -176,8 +176,11 @@ class User(Entity):
             next_id = str(int(current_id) + 1)
             self._id = next_id
             db.save("_system", f"{type_name}_id", self._id)
-        else:
-            # Update max_id if custom ID is higher than current max
+        elif not self._loaded:
+            # Update max_id if custom ID is higher than current max.
+            # Skipped when hydrating persisted state: that id already went
+            # through this check when it was first created, so re-reading the
+            # counter would only add one storage read per loaded entity.
             db = self.db()
             type_name = self._type
             current_max_id = db.load("_system", f"{type_name}_id")
@@ -418,6 +421,16 @@ def load(
         if not data:
             return None
 
+        return cls._hydrate(entity_id, data)
+
+    @classmethod
+    def _hydrate(cls: Type[T], entity_id: str, data: dict) -> T:
+        """Build an instance from its persisted dict, migrating if needed.
+
+        Shared by ``load`` (single ``get``) and ``load_some`` (range page).
+        Callers must have checked the identity registry first.
+        """
+        type_name = cls.get_full_type_name()
         stored_version = data.get("__version__", 1)
         current_version = cls.__version__
 
@@ -582,19 +595,17 @@ def rebuild_field_index(
 
         db = cls.db()
         type_name = cls.get_full_type_name()
-        end = min(from_id + batch - 1, max_id)
-        for entity_id in range(from_id, end + 1):
-            try:
-                entity = cls.load(str(entity_id))
-            except (ValueError, AttributeError):
-                continue
-            if entity is None:
-                continue
+        # One range read per batch of *live* entities (deleted ids cost nothing)
+        entities = cls.load_some(from_id, batch)
+        for entity in entities:
             value = getattr(entity, field, None)
             if value is not None:
                 db.field_index_add(type_name, field, str(value), entity._id)
 
-        return end + 1 if end < max_id else None
+        if len(entities) < batch:
+            return None
+        next_from_id = int(entities[-1]._id) + 1
+        return next_from_id if next_from_id <= max_id else None
 
     @classmethod
     def instances(cls: Type[T]) -> List[T]:
@@ -676,11 +687,7 @@ def max_id(cls: Type[T]) -> int:
         Returns:
             int: Maximum entity ID
         """
-        type_name = cls.get_full_type_name()
-        db = cls.db()
-        max_id_key = f"{type_name}_id"
-        max_id = db.load("_system", max_id_key)
-        return int(max_id) if max_id else 0
+        return cls.db().max_id(cls.get_full_type_name())
 
     @classmethod
     def load_some(
@@ -698,29 +705,42 @@ def load_some(
             List[T]: List of entities for the requested page
 
         Raises:
-            ValueError: If page or page_size is less than 1
+            ValueError: If from_id or count is less than 1
+
+        Cost: with a range-capable storage backend (Basilisk's StableBTreeMap,
+        MemoryStorage) this is one ordered range read per page regardless of
+        how many ids in between were deleted; otherwise one read per id.
         """
-        logger.info(f"Loading entities from {from_id} to {from_id + count}")
+        logger.debug(f"Loading up to {count} entities from id {from_id}")
 
         if from_id < 1:
             raise ValueError("from_id must be at least 1")
         if count < 1:
             raise ValueError("count must be at least 1")
 
-        # Return the slice of entities for the requested page
-        ret = []
+        db = cls.db()
+        type_name = cls.get_full_type_name()
 
-        while len(ret) < count and from_id <= cls.max_id():
-            logger.info(f"Loading entity {from_id}")
-            try:
-                entity = cls.load(str(from_id))
-                if entity:
+        ret: List[T] = []
+        max_id = db.max_id(type_name)
+        while len(ret) < count:
+            need = count - len(ret)
+            page = db.load_page(type_name, from_id, need, max_id=max_id)
+            for entity_id, data in page:
+                try:
+                    # Identity map first, so callers keep getting the same instance
+                    entity = db.get_entity(type_name, entity_id)
+                    if entity is None:
+                        entity = cls._hydrate(entity_id, data)
                     ret.append(entity)
-            except (ValueError, AttributeError) as e:
-                # Skip entities with broken/dangling relation references
-                # (full fix: issue #4 — lazy relation resolution)
-                logger.warning(f"Skipping {cls.__name__}@{from_id}: {e}")
-            from_id += 1
+                except (ValueError, AttributeError) as e:
+                    # Skip entities with broken/dangling relation references
+                    # (full fix: issue #4 — lazy relation resolution)
+                    logger.warning(f"Skipping {cls.__name__}@{entity_id}: {e}")
+            if len(page) < need:
+                break  # no more persisted entities of this type
+            # Page was full but something was skipped: continue after it
+            from_id = int(page[-1][0]) + 1
 
         return ret
 

ic_python_db/storage.py

@@ -3,11 +3,39 @@
 """
 
 from abc import ABC, abstractmethod
-from typing import Dict, Iterator, Optional, Tuple
+from typing import Dict, Iterator, List, Optional, Tuple
+
+
+def storage_sort_key(key: str) -> Tuple[int, bytes]:
+    """Sort key reproducing the on-chain order of ``str`` keys.
+
+    Basilisk's ``StableBTreeMap`` stores a ``str`` key as
+    ``tag + 4-byte big-endian length + utf-8``, so the B-tree orders keys by
+    byte length first and bytewise second: ``"T@9" < "T@10" < "T@2x"``.
+    ``MemoryStorage.range`` uses this so local tests see exactly the order a
+    canister will.
+    """
+    raw = key.encode("utf-8")
+    return (len(raw), raw)
+
+
+def supports_range(storage) -> bool:
+    """True if ``storage`` offers ``range(start, end, limit)``.
+
+    Basilisk's ``StableBTreeMap`` (>= the version that added ``range``) and
+    ``MemoryStorage`` do; older CDK builds and simple dict-backed test
+    doubles don't, and callers fall back to key probing.
+    """
+    return callable(getattr(storage, "range", None))
 
 
 class Storage(ABC):
-    """Abstract base class for storage backends"""
+    """Abstract base class for storage backends
+
+    ``range`` is optional: backends that can walk keys in order should
+    implement it (see ``storage_sort_key`` for the order that is expected);
+    callers must check ``supports_range`` and fall back to ``get`` probing.
+    """
 
     @abstractmethod
     def insert(self, key: str, value: str) -> None:
@@ -67,3 +95,24 @@ def __contains__(self, key: str) -> bool:
     def keys(self) -> Iterator[str]:
         """Return all keys in storage"""
         return iter(self._data.keys())
+
+    def range(
+        self, start: str, end: Optional[str] = None, limit: int = 1000
+    ) -> List[Tuple[str, str]]:
+        """Ordered page of ``(key, value)`` with ``start <= key < end``.
+
+        Same contract as Basilisk's ``StableBTreeMap.range``: half-open,
+        ``end=None`` means unbounded, keys ordered by ``storage_sort_key``.
+        O(n log n) here, which is fine for an in-memory test backend.
+        """
+        if limit < 0:
+            raise ValueError("limit must be >= 0")
+        lo = storage_sort_key(start)
+        hi = None if end is None else storage_sort_key(end)
+        selected = [
+            k
+            for k in self._data
+            if lo <= storage_sort_key(k) and (hi is None or storage_sort_key(k) < hi)
+        ]
+        selected.sort(key=storage_sort_key)
+        return [(k, self._data[k]) for k in selected[:limit]]

tests/entrypoint.sh

@@ -14,7 +14,11 @@ TEMPLATE_PATH="${TEMPLATE_DIR}/cpython_canister_template.wasm"
 if [ ! -f "$TEMPLATE_PATH" ]; then
     echo "Downloading CPython canister template..."
     mkdir -p "$TEMPLATE_DIR"
-    curl -fL https://github.com/smart-social-contracts/basilisk/releases/download/cpython-wasm-3.13.0/cpython_canister_template.wasm \
+    # Prefer the template attached to the installed Basilisk release (this is
+    # where basilisk.find_template_wasm looks); fall back to the rolling CI asset.
+    curl -fL "https://github.com/smart-social-contracts/basilisk/releases/download/v${BASILISK_VERSION}/cpython_canister_template.wasm" \
+         -o "$TEMPLATE_PATH" || \
+    curl -fL https://github.com/smart-social-contracts/basilisk/releases/download/cpython-wasm-3.13.0-ic1/cpython_canister_template.wasm \
          -o "$TEMPLATE_PATH"
     echo "Template downloaded: $(du -sh "$TEMPLATE_PATH" | cut -f1)"
 fi

tests/entrypoint_benchmark.sh

@@ -13,7 +13,11 @@ TEMPLATE_PATH="${TEMPLATE_DIR}/cpython_canister_template.wasm"
 if [ ! -f "$TEMPLATE_PATH" ]; then
     echo "Downloading CPython canister template..."
     mkdir -p "$TEMPLATE_DIR"
-    curl -fL https://github.com/smart-social-contracts/basilisk/releases/download/cpython-wasm-3.13.0/cpython_canister_template.wasm \
+    # Prefer the template attached to the installed Basilisk release (this is
+    # where basilisk.find_template_wasm looks); fall back to the rolling CI asset.
+    curl -fL "https://github.com/smart-social-contracts/basilisk/releases/download/v${BASILISK_VERSION}/cpython_canister_template.wasm" \
+         -o "$TEMPLATE_PATH" || \
+    curl -fL https://github.com/smart-social-contracts/basilisk/releases/download/cpython-wasm-3.13.0-ic1/cpython_canister_template.wasm \
          -o "$TEMPLATE_PATH"
     echo "Template downloaded: $(du -sh "$TEMPLATE_PATH" | cut -f1)"
 fi

tests/entrypoint_stress.sh

@@ -14,7 +14,11 @@ TEMPLATE_PATH="${TEMPLATE_DIR}/cpython_canister_template.wasm"
 if [ ! -f "$TEMPLATE_PATH" ]; then
     echo "Downloading CPython canister template..."
     mkdir -p "$TEMPLATE_DIR"
-    curl -fL https://github.com/smart-social-contracts/basilisk/releases/download/cpython-wasm-3.13.0/cpython_canister_template.wasm \
+    # Prefer the template attached to the installed Basilisk release (this is
+    # where basilisk.find_template_wasm looks); fall back to the rolling CI asset.
+    curl -fL "https://github.com/smart-social-contracts/basilisk/releases/download/v${BASILISK_VERSION}/cpython_canister_template.wasm" \
+         -o "$TEMPLATE_PATH" || \
+    curl -fL https://github.com/smart-social-contracts/basilisk/releases/download/cpython-wasm-3.13.0-ic1/cpython_canister_template.wasm \
          -o "$TEMPLATE_PATH"
     echo "Template downloaded: $(du -sh "$TEMPLATE_PATH" | cut -f1)"
 fi