Design an in-memory database with transactions

01The design, in the six-step framework

1. Requirements and use cases

A machine-coding staple at the SDE II and III level: an in-memory database small enough to write in an hour but shaped like a real one, with transactions that can be rolled back.

• A client creates a table with named columns and a primary key.
• A client inserts, updates and deletes rows by primary key, and selects rows by a column value.
• A client begins a transaction, makes changes, and either commits them all or rolls back so nothing is visible.
• Reads inside an open transaction see that transaction's own uncommitted writes.
• A client can declare a secondary index on a column so selects on it avoid a full scan.

Out of scope: SQL parsing, joins, persistence, and multi-client isolation levels beyond "one open transaction at a time per session". Non-functional constraints: point reads and writes in O(1); rollback cost proportional to the number of changes in the transaction, not to the table size; and the storage layer must be swappable without touching the transaction logic, which is the named extension.

2. Core entities and relationships

• Database: the facade; owns tables and the current transaction; every client call enters here.
• Table: rows keyed by primary key, a schema, and zero or more indexes; knows how to apply and reverse one row change.
• Index: a map from a column value to the set of primary keys holding it; kept in step with the rows by the table.
• Transaction: an ordered log of Change objects; commit discards the log, rollback replays it backwards.
• Change: one reversible row mutation holding the previous row (or none) and the new row (or none).

The database composes tables and the transaction; the table composes its rows and its indexes. The transaction aggregates changes that reference tables by name, so a change can undo itself by asking the database for the table. The key relationship decision: a Change holds the before-image of the row, not a diff, because restoring a whole row is simpler to reason about and rows are small. This is the layering of a tiny database: a facade on top, a transaction layer in the middle, storage at the bottom.

3. Interfaces and key methods

from typing import Any, Dict, List, Optional

Row = Dict[str, Any]


class Database:
    def create_table(self, name: str, primary_key: str, columns: List[str]) -> None:
        """Fails if the table exists or primary_key is not among columns."""

    def create_index(self, table: str, column: str) -> None:
        """Builds the index from existing rows; later writes maintain it."""

    def insert(self, table: str, row: Row) -> None:
        """Fails on a duplicate primary key; recorded in the open transaction if any."""

    def update(self, table: str, pk: Any, changes: Row) -> None:
        """Merges changes into the existing row; fails if the row is absent."""

    def delete(self, table: str, pk: Any) -> None: ...
    def select(self, table: str, column: str, value: Any) -> List[Row]:
        """Uses the index on column if one exists, else scans."""

    def begin(self) -> None: ...
    def commit(self) -> None: ...
    def rollback(self) -> None: ...

Invariants: every row in a table satisfies the schema and has a unique primary key; every index reflects exactly the current rows; the transaction log, replayed backwards, restores the tables to their state at begin; and outside a transaction every write is its own implicit transaction, so autocommit and explicit transactions use one code path.

4. Design patterns and the choices behind them

Command objects with before-images for rollback

The force is undo. Each write produces a Command-like Change that knows how to reverse itself. Storing the before-image is a memento of one row; replaying the log in reverse order restores everything, including index entries, because undo goes through the same Table.put and Table.remove that maintain indexes. The rejected alternative is copying every table at begin: O(size) per transaction, hopeless for a large table and a tiny transaction.

Facade over tables

The force is a small public surface over several cooperating classes. The Database facade is the only object clients hold; it routes every call to a table, records the change in the transaction, and hides that tables and indexes exist as objects. Exposing Table directly is rejected because a client could then mutate rows behind the transaction's back.

Indexes maintained by the table, not by callers

The force is consistency. If callers had to update indexes, one forgotten call would corrupt a select forever. The table's two mutation methods update every index, so the invariant is enforced in one place. Adding a unique-index rule later is a check in the same place; that is the extension point you can name when asked.

5. Core flow and code

Flow for begin(); update("users", 7, {"city": "Pune"}); rollback():

1. begin creates an empty Transaction.
2. update asks the users table for row 7, builds the merged row, calls Table.put which returns the old row and fixes the indexes, and records a Change(before=old, after=new).
3. rollback walks the log backwards; the single change calls Table.put(before), which restores the old row and the index entries; the transaction is dropped.

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set

Row = Dict[str, Any]


class Index:
    def __init__(self, column: str) -> None:
        self.column = column
        self.buckets: Dict[Any, Set[Any]] = {}

    def add(self, pk: Any, row: Row) -> None:
        self.buckets.setdefault(row.get(self.column), set()).add(pk)

    def drop(self, pk: Any, row: Row) -> None:
        bucket = self.buckets.get(row.get(self.column))
        if bucket:
            bucket.discard(pk)

    def lookup(self, value: Any) -> Set[Any]:
        return set(self.buckets.get(value, ()))


class Table:
    def __init__(self, name: str, primary_key: str, columns: List[str]) -> None:
        if primary_key not in columns:
            raise ValueError("primary key must be a column")
        self.name, self.pk, self.columns = name, primary_key, set(columns)
        self.rows: Dict[Any, Row] = {}
        self.indexes: Dict[str, Index] = {}

    def put(self, row: Optional[Row], pk: Any) -> Optional[Row]:
        """Set the row for pk (None deletes). Returns the previous row. Maintains indexes."""
        old = self.rows.pop(pk, None)
        if old is not None:
            for idx in self.indexes.values():
                idx.drop(pk, old)
        if row is not None:
            if set(row) - self.columns:
                raise ValueError(f"unknown columns in {self.name}")
            self.rows[pk] = row
            for idx in self.indexes.values():
                idx.add(pk, row)
        return old

    def find(self, column: str, value: Any) -> List[Row]:
        idx = self.indexes.get(column)
        if idx is not None:
            return [self.rows[pk] for pk in idx.lookup(value)]
        return [r for r in self.rows.values() if r.get(column) == value]


@dataclass
class Change:
    table: str
    pk: Any
    before: Optional[Row]

    def undo(self, table: Table) -> None:
        table.put(self.before, self.pk)


@dataclass
class Transaction:
    log: List[Change] = field(default_factory=list)


class Database:
    def __init__(self) -> None:
        self.tables: Dict[str, Table] = {}
        self.tx: Optional[Transaction] = None

    def create_table(self, name: str, primary_key: str, columns: List[str]) -> None:
        if name in self.tables:
            raise ValueError("table exists")
        self.tables[name] = Table(name, primary_key, columns)

    def create_index(self, table: str, column: str) -> None:
        t = self._table(table)
        idx = Index(column)
        for pk, row in t.rows.items():
            idx.add(pk, row)
        t.indexes[column] = idx

    def insert(self, table: str, row: Row) -> None:
        t = self._table(table)
        pk = row[t.pk]
        if pk in t.rows:
            raise KeyError("duplicate primary key")
        self._write(t, pk, dict(row))

    def update(self, table: str, pk: Any, changes: Row) -> None:
        t = self._table(table)
        if pk not in t.rows:
            raise KeyError("no such row")
        self._write(t, pk, {**t.rows[pk], **changes})

    def delete(self, table: str, pk: Any) -> None:
        self._write(self._table(table), pk, None)

    def select(self, table: str, column: str, value: Any) -> List[Row]:
        return [dict(r) for r in self._table(table).find(column, value)]

    def begin(self) -> None:
        if self.tx is not None:
            raise RuntimeError("transaction already open")
        self.tx = Transaction()

    def commit(self) -> None:
        self._require_tx()
        self.tx = None                       # changes are already applied

    def rollback(self) -> None:
        self._require_tx()
        for change in reversed(self.tx.log):  # type: ignore[union-attr]
            change.undo(self.tables[change.table])
        self.tx = None

    def _write(self, t: Table, pk: Any, row: Optional[Row]) -> None:
        before = t.put(row, pk)
        if self.tx is not None:
            self.tx.log.append(Change(t.name, pk, before))

    def _table(self, name: str) -> Table:
        if name not in self.tables:
            raise KeyError(f"no table {name}")
        return self.tables[name]

    def _require_tx(self) -> None:
        if self.tx is None:
            raise RuntimeError("no open transaction")


if __name__ == "__main__":
    db = Database()
    db.create_table("users", "id", ["id", "city"])
    db.create_index("users", "city")
    db.insert("users", {"id": 7, "city": "Delhi"})
    db.begin()
    db.update("users", 7, {"city": "Pune"})
    assert db.select("users", "city", "Pune")[0]["id"] == 7   # own write visible
    db.rollback()
    assert db.select("users", "city", "Pune") == [] and db.select("users", "city", "Delhi")

6. Extensions, concurrency and what interviewers probe

• Nested transactions. Make tx a stack; rollback undoes only the top log, commit appends the top log to the parent's so an outer rollback still undoes everything.
• Two sessions at once. Each session gets its own Transaction and sees the shared tables; the simple answer is one lock around every public method, the better answer is optimistic versioning per row with a check at commit, and you should say which isolation level you are offering.
• Swapping storage. Table's put and find are the whole storage contract; a disk-backed or sorted-map implementation behind the same two methods leaves Database and Transaction untouched.
• Unique secondary indexes and range queries. Uniqueness is a check in put; ranges need a sorted index, which is a different Index class behind the same lookup shape plus a between.
• What to test first. Rollback restores index entries, not just rows; the test in the code checks exactly that through select.
• Common mistake. Recording the change after mutating through a path that bypasses put, so the index drifts; every write must go through _write.

02Where this question was reported

• workat.tech · workat.tech machine coding (SDE II/III)
• Exponent / Aced · Aced/Exponent: 'Design an in-memory database' under API and low-level design
• LeetCode Discuss · LeetCode 981 Time Based Key-Value Store, Google-tagged (freq 32.6)