Problem 72 · Object & System Design

KV Store with Nested Transactions

MediumClassic DSA & systems

01Problem

Implement a KV store supporting set, get, delete, begin, rollback, and commit. Transactions nest: rollback undoes only the innermost; commit applies all open transactions. Rollback/commit with no open transaction returns False.

02Key insight

Maintain a stack of write-sets layered over the base map; reads check from innermost layer outward. A tombstone sentinel records deletes. Rollback pops a layer; commit flattens all layers into base.

03Solution

solution.pypython51 lines

from typing import Any

_TOMBSTONE = object()


class TransactionalKV:
    def __init__(self) -> None:
        self.base: dict[str, Any] = {}
        self.layers: list[dict[str, Any]] = []

    def _write(self, key: str, value: Any) -> None:
        if self.layers:
            self.layers[-1][key] = value
        elif value is _TOMBSTONE:
            self.base.pop(key, None)
        else:
            self.base[key] = value

    def set(self, key: str, value: Any) -> None:
        self._write(key, value)

    def delete(self, key: str) -> None:
        self._write(key, _TOMBSTONE)

    def get(self, key: str, default: Any = None) -> Any:
        for layer in reversed(self.layers):
            if key in layer:
                v = layer[key]
                return default if v is _TOMBSTONE else v
        return self.base.get(key, default)

    def begin(self) -> None:
        self.layers.append({})

    def rollback(self) -> bool:
        if not self.layers:
            return False
        self.layers.pop()
        return True

    def commit(self) -> bool:
        if not self.layers:
            return False
        for layer in self.layers:  # oldest -> newest so inner writes win
            for k, v in layer.items():
                if v is _TOMBSTONE:
                    self.base.pop(k, None)
                else:
                    self.base[k] = v
        self.layers.clear()
        return True

Try it first.

TimeO(1) write; O(depth) read; O(total writes) commit

SpaceO(writes in open txns)

tests.py · 4 passing

tests.pypython9 lines

kv = TransactionalKV()
kv.set("a", 1)
kv.begin(); kv.set("a", 2)
kv.begin(); kv.delete("a")
assert kv.get("a") is None
assert kv.rollback() and kv.get("a") == 2
assert kv.rollback() and kv.get("a") == 1 and not kv.rollback()
kv.begin(); kv.set("b", 5); kv.begin(); kv.delete("a")
assert kv.commit() and kv.get("a") is None and kv.get("b") == 5 and not kv.commit()

04Staff-level follow-up

Reads are O(depth) and some variants need COUNT(value). Alternative: undo-log per layer (apply writes to base, store prior values) gives O(1) reads and makes value-count indexes easy to maintain.