waxsql 1.0.0__py3-none-any.whl

waxsql/gen/data/rows.py

@@ -0,0 +1,236 @@
+"""Per-table row materialization and FK resolution.
+
+The walk is one topological pass over the FK DAG. For each table, we
+materialize all rows in memory as Python tuples, capture the PK column
+values into a per-table ID store, and hand the tuples to the emitter.
+Children resolve FK columns by sampling from the parent's ID list,
+which is guaranteed populated by the topological order.
+
+Self-referential FKs (a column in table T referencing T.id) are handled
+as a special case: row with pk=N can reference any row in [1..N] including
+itself, because PostgreSQL checks FK constraints per-row by default (NOT
+DEFERRABLE). Nullable self-FK columns honour null_fraction (the null roll
+fires above, before the FK branch, and controls whether the column is NULL);
+NOT NULL self-FK columns always sample from [1..pk].
+
+In-memory materialization is fine at demo scales (default 100 rows ×
+fanout 5 across a few levels of depth is comfortably under a megabyte).
+If volume ever becomes a real constraint, a streaming variant can be
+added behind the same public API.
+"""
+from __future__ import annotations
+
+import random
+from collections.abc import Mapping
+
+from waxsql.gen.data.columns import strategy_for
+from waxsql.schema import Schema, Table
+
+
+def topological_order(schema: Schema) -> list[Table]:
+    """Return the schema's tables in FK-topological order: every table
+    appears before any table that has an FK referencing it.
+
+    Within each topological rank, tables are sorted alphabetically by
+    name. This eliminates dict-insertion-order dependency and keeps the
+    walk byte-identical across Python builds.
+
+    Raises ValueError if a cycle is detected (the current generator does
+    not produce cycles, but we assert loudly rather than loop forever).
+
+    Algorithm: Kahn-style level-set walk. Each iteration extracts the
+    set of tables whose remaining parents are empty (i.e. all their FK
+    targets have already been placed), sorts that set alphabetically,
+    and appends them. A "no ready tables but remaining is non-empty"
+    state means a cycle and raises.
+    """
+    by_name = {t.name: t for t in schema.tables}
+    # Collect UNIQUE external parents per table: FKs that point to a
+    # *different* table. Self-referential FKs (`t → t`) are intentionally
+    # excluded because they impose no ordering constraint — the table
+    # trivially appears before itself.
+    parents: dict[str, set[str]] = {
+        t.name: {fk.ref_table for fk in t.foreign_keys if fk.ref_table != t.name}
+        for t in schema.tables
+    }
+    remaining = set(by_name)
+    out: list[Table] = []
+    while remaining:
+        # sorted(...) is essential here, not cosmetic: `set` iteration
+        # order is not stable across Python builds (hash-randomized for
+        # strings since 3.3), and any set-driven choice would silently
+        # break the determinism contract.
+        ready = sorted(
+            n for n in remaining
+            if not (parents[n] & remaining)
+        )
+        if not ready:
+            # No table has all its parents already placed AND `remaining`
+            # is non-empty → cycle. The data generator can't currently
+            # handle cycles (would need deferred constraints + UPDATE
+            # patches); the CLI catches this ValueError and produces a
+            # clean usage message pointing at the cycle.
+            raise ValueError(
+                f"FK cycle in schema: remaining {sorted(remaining)!r}"
+            )
+        for n in ready:
+            out.append(by_name[n])
+            remaining.discard(n)
+    return out
+
+
+def depth_of(table: Table, schema: Schema) -> int:
+    """The longest FK chain from any root to `table`.
+
+    Roots have depth 0. A table directly referencing a root has depth 1.
+    When a table has multiple FK parents, depth is 1 + max(parent depths)
+    — the longest chain, not the shortest.
+
+    Depth feeds the row-count formula: `rows * fanout ** depth`. Choosing
+    longest-chain (not shortest) means a leaf with even one deep ancestor
+    gets the deep row count — appropriate, since the leaf has to spread
+    across the parent's larger ID pool to avoid clumping.
+
+    The internal `memo` dict is local to this call, so memoization
+    avoids re-walking shared ancestry WITHIN a single invocation but
+    is NOT shared across calls. `rows_for_table` calls `depth_of`
+    once per table in the schema, so the total work for a full
+    data-generation pass is O(N · D) where N is the number of
+    tables and D the maximum chain length — fine at the schema
+    generator's N=12 ceiling. Promote to a shared memo threaded
+    through `topological_order`/`rows_for_table` only if N grows
+    materially. The `stack` parameter threads the current path for
+    cycle detection; a cycle raises ValueError rather than
+    recursing to Python's stack limit.
+    """
+    by_name = {t.name: t for t in schema.tables}
+    memo: dict[str, int] = {}
+
+    def visit(name: str, stack: tuple[str, ...] = ()) -> int:
+        if name in stack:
+            # Cycle detected on the recursion path itself — different
+            # from the topological_order cycle check, which is global.
+            # Tuple-based stack is cheap (small N) and keeps the path
+            # for the error message.
+            raise ValueError(f"FK cycle through {name}: {stack}")
+        if name in memo:
+            return memo[name]
+        t = by_name[name]
+        # Only consider FKs to *other* tables. Self-referential FKs
+        # (`t → t`) would recurse forever; they carry no depth information
+        # because the row materializer handles self-references separately
+        # (sample from [1..pk] in the same row generation).
+        external_parents = [fk.ref_table for fk in t.foreign_keys if fk.ref_table != name]
+        d = 0 if not external_parents else 1 + max(
+            visit(p, stack + (name,)) for p in external_parents
+        )
+        memo[name] = d
+        return d
+
+    return visit(table.name)
+
+
+# PK column name convention is fixed by the schema generator: every table
+# has an `id BIGINT NOT NULL` PK at the first column position. We don't
+# search for it dynamically; if that invariant ever changes, this module
+# will produce nonsense for that table and tests will fail loudly, which
+# is the desired behavior.
+_PK_COLUMN_NAME = "id"
+
+
+def rows_for_table(table: Table, schema: Schema, *, base: int, fanout: int) -> int:
+    """Row count for `table`: `base * fanout ** depth`. No ceiling; the
+    user is responsible for not asking for combinations that produce
+    absurd output. Depth is 0 for root tables, so they always get exactly
+    `base` rows regardless of fanout.
+
+    Why the power law: it lets child tables outnumber their parents
+    geometrically, which is what real one-to-many shapes look like
+    (orders >> customers, line_items >> orders). With base=100, fanout=5,
+    depth=3: 12500 rows — large enough to exercise EXPLAIN's choices,
+    small enough to fit comfortably in memory and load in seconds.
+    """
+    return base * fanout ** depth_of(table, schema)
+
+
+def generate_row(
+    *,
+    table: Table,
+    pk: int,
+    rng: random.Random,
+    id_store: Mapping[str, list[int]],
+    null_fraction: float,
+) -> tuple[object, ...]:
+    """Generate one row for `table` as a Python tuple.
+
+    PK column gets `pk`; FK columns sample from `id_store[ref_table]`;
+    other columns dispatch through `strategy_for`. Nullable columns —
+    including nullable FK columns — roll for NULL first; only columns
+    that survive the null roll reach FK or strategy dispatch.
+
+    Self-referential FKs bypass `id_store` entirely: PostgreSQL's default
+    NOT DEFERRABLE constraint checking allows row pk=N to reference any
+    row in [1..N] (itself included, since the row exists by check time).
+    Both nullable and non-nullable self-FK columns sample from [1..pk] once
+    they reach the FK branch — the null roll above already gave nullable
+    columns their chance to become NULL, so null_fraction is honoured
+    correctly instead of forcing self-FK nullables to always be NULL.
+    """
+    # Build a column-name → ref_table lookup once per row. With most
+    # schemas having ≤ a handful of FKs per table this is cheap.
+    # Dict order doesn't matter here — we lookup by name, not iterate.
+    fk_by_col: dict[str, str] = {}
+    for fk in table.foreign_keys:
+        for child_col, _parent_col in zip(fk.columns, fk.ref_columns, strict=True):
+            fk_by_col[child_col] = fk.ref_table
+
+    values: list[object] = []
+    # Column ORDER matters: every consumed rng call advances the shared
+    # RNG state, so a swap here changes downstream output. The iteration
+    # order follows `table.columns`, which is the schema generator's
+    # declared order — same source of truth as the DDL emitter.
+    for col in table.columns:
+        if col.name == _PK_COLUMN_NAME:
+            # PK is supplied by the caller (sequential 1..n). PK never
+            # rolls for NULL and never goes through a strategy — bypass.
+            values.append(pk)
+            continue
+        if col.nullable and rng.random() < null_fraction:
+            # Null roll consumes one rng tick — this consumption happens
+            # for EVERY nullable column regardless of FK status, which
+            # keeps the FK and non-FK paths byte-stable when a column's
+            # type changes between FK and non-FK in future refactors.
+            values.append(None)
+            continue
+        if col.name in fk_by_col:
+            ref_table = fk_by_col[col.name]
+            if ref_table == table.name:
+                # Self-FK: PG enforces FKs per-row (default NOT DEFERRABLE),
+                # so row pk=N can reference any row in [1..N] including itself.
+                # Without this branch we'd try to read id_store[table.name]
+                # which is empty (populated only after this loop in data.py).
+                # For nullable columns, the null roll has already fired above
+                # — any column that reaches this branch survived that roll
+                # and should get a real value, not unconditional NULL.
+                values.append(rng.randint(1, pk))
+                continue
+            parent_ids = id_store[ref_table]
+            # If the parent table has zero rows, this FK column must be NULL.
+            # A NOT NULL FK into an empty parent table is a generator error —
+            # the topological walk should have materialized parents first.
+            # This branch ONLY fires when the user passes --rows=0 (or the
+            # parent's depth-adjusted count rounds to 0, which can't happen
+            # with the current `base * fanout**depth` formula and base > 0).
+            if not parent_ids:
+                if not col.nullable:
+                    raise ValueError(
+                        f"NOT NULL FK {table.name}.{col.name} references "
+                        f"empty parent {ref_table}"
+                    )
+                values.append(None)
+                continue
+            values.append(rng.choice(parent_ids))
+            continue
+        strat = strategy_for(col)
+        values.append(strat(rng, col))
+    return tuple(values)

waxsql/gen/data/strategies.py

@@ -0,0 +1,299 @@
+"""Per-PgType value strategies. Each strategy maps (rng, Column) → object;
+the emitter is responsible for formatting that object for COPY.
+
+The split between strategies and the emitter is deliberate: strategies
+return native Python values (Decimal, datetime, UUID, dict, list) and
+have no knowledge of tab encoding or NULL sentinels. That keeps the
+strategy registry trivially testable and lets the emitter own all of
+PostgreSQL's COPY text-format escape rules in one place.
+
+Determinism invariants (load-bearing across this whole module):
+  * Only the injected `rng` is allowed as a source of randomness.
+    No `uuid.uuid4()`, no `random.gauss()` from the global module,
+    no `datetime.now()`, no environment lookups.
+  * `_EPOCH` is a fixed date constant — wall-clock input would mean
+    the same seed produces different data tomorrow, breaking the
+    "reproduce a bug from a seed years later" guarantee.
+  * Dict iteration order is stable in Python 3.7+, so `_TYPE_STRATEGIES`
+    is fine; but anything that iterates a `set` for rng decisions must
+    `sorted()` first (none currently do).
+"""
+from __future__ import annotations
+
+import datetime as _dt
+import random
+import uuid
+from dataclasses import dataclass
+from decimal import Decimal
+from collections.abc import Callable
+
+from waxsql.schema import Column
+from waxsql.types import PgType
+
+
+# Hand-curated short list of simple English words. Used as the default
+# text/varchar value source: each cell gets one of these. Not a name
+# dictionary, not Lorem Ipsum — just enough variety that SELECT * FROM t
+# doesn't look like line noise. Grow as taste dictates.
+WORDLIST: tuple[str, ...] = (
+    "alpha", "amber", "anchor", "apple", "arrow", "atlas", "azure",
+    "badger", "basin", "beacon", "berry", "birch", "blossom", "boulder",
+    "breeze", "bridge", "bronze", "buffalo", "cabin", "candle", "canyon",
+    "cedar", "cherry", "cinder", "clover", "cobalt", "comet", "copper",
+    "coral", "cottage", "crater", "crescent", "crimson", "crystal",
+    "dahlia", "dawn", "delta", "diamond", "dolphin", "dove", "dusk",
+    "ember", "emerald", "falcon", "feather", "fern", "festival", "fjord",
+    "forest", "fossil", "frost", "galaxy", "garnet", "gentian", "geyser",
+    "glacier", "glade", "granite", "harbor", "harvest", "haven", "hazel",
+    "heron", "hickory", "horizon", "indigo", "iris", "island", "ivory",
+    "jasper", "juniper", "kestrel", "lagoon", "lantern", "lavender",
+    "library", "linden", "lotus", "magnet", "maple", "marble", "marigold",
+    "marsh", "meadow", "mercury", "midnight", "mineral", "mint", "mirage",
+    "morning", "mosaic", "mountain", "nectar", "nimbus", "nocturne",
+    "oasis", "obsidian", "ocean", "olive", "opal", "orchid", "otter",
+    "panther", "parsley", "peach", "pebble", "pelican", "petal", "phlox",
+    "pine", "plateau", "poppy", "prairie", "prism", "quartz", "quail",
+    "quill", "rainbow", "raven", "redwood", "river", "robin", "rose",
+    "ruby", "saffron", "sage", "salmon", "sapphire", "scarlet", "sequoia",
+    "shadow", "shoreline", "silver", "slate", "solstice", "sparrow",
+    "spruce", "starling", "stone", "summit", "sunset", "swallow", "tangerine",
+    "thicket", "thistle", "thunder", "tide", "topaz", "tulip", "turquoise",
+    "twilight", "umber", "valley", "velvet", "verdant", "violet", "walnut",
+    "waterfall", "willow", "winter", "wisteria", "yarrow", "yew", "zenith",
+    "zephyr", "zircon",
+)
+
+
+def pick_word(rng: random.Random) -> str:
+    """Pick a deterministic word from `WORDLIST`. Uses `rng.choice` rather
+    than `rng.randint`+index so adding/removing wordlist entries shifts
+    output predictably rather than scrambling it.
+    """
+    return rng.choice(WORDLIST)
+
+
+# ---------------------------------------------------------------------------
+# Per-type strategy functions
+# ---------------------------------------------------------------------------
+
+# A strategy is a pure function: (rng, column) -> Python value.
+# Native types are returned; the emitter formats them for COPY.
+Strategy = Callable[[random.Random, Column], object]
+
+
+def _int4(rng: random.Random, col: Column) -> int:
+    # PG int4 is signed 32-bit: -2147483648..2147483647. The -1 floor
+    # avoids the minimum-int corner that some downstream string
+    # formatters mishandle; the loss of one value is irrelevant.
+    return rng.randint(-(2**31) + 1, (2**31) - 1)
+
+
+def _int8(rng: random.Random, col: Column) -> int:
+    # PG int8 is signed 64-bit but we deliberately bound this tighter
+    # than the full range. Full-range int8 values aren't interesting
+    # for demo data and produce visually ugly output.
+    # Trade-off: planner-cost estimates that depend on value distribution
+    # see a slightly narrower range; this hasn't bitten anything yet.
+    return rng.randint(-(2**62), (2**62))
+
+
+def _text(rng: random.Random, col: Column) -> str:
+    return pick_word(rng)
+
+
+def _varchar(rng: random.Random, col: Column) -> str:
+    word = pick_word(rng)
+    # Respect typmod when present; varchar(N) rejects values longer than N.
+    # When typmod is absent (`varchar` with no length), cap at 32 — generous
+    # enough for any single wordlist entry, modest enough to keep COPY
+    # output readable. Note: pick_word ALWAYS fires (consumes one rng tick)
+    # even if the cap would truncate to empty — keeps the rng stream
+    # independent of typmod, so adding a length to a column doesn't
+    # cascade into downstream byte-shifts.
+    cap = col.type.typmod[0] if col.type.typmod else 32
+    return word[:cap]
+
+
+def _bool(rng: random.Random, col: Column) -> bool:
+    return rng.choice((True, False))
+
+
+def _uuid(rng: random.Random, col: Column) -> uuid.UUID:
+    # rng.getrandbits(128) keeps determinism within our RNG; uuid.uuid4()
+    # would reach for os.urandom and break that. The resulting UUID won't
+    # have the version-4 bit pattern set — PG doesn't care, the uuid column
+    # accepts any 128-bit value, and the determinism contract trumps RFC
+    # 4122 cosmetic correctness.
+    return uuid.UUID(int=rng.getrandbits(128))
+
+
+def _float8(rng: random.Random, col: Column) -> float:
+    # Bounded so output stays human-readable. Avoid `random.gauss()` —
+    # we want a flat distribution for COPY, not a bell curve.
+    return rng.uniform(-1_000_000.0, 1_000_000.0)
+
+
+def _numeric(rng: random.Random, col: Column) -> Decimal:
+    # numeric(precision, scale): `precision` total digits, `scale` after
+    # the decimal point. Magnitude < 10^(precision-scale); fractional
+    # digits = scale. Without typmod, fall back to a sensible default.
+    # The default (10, 4) is arbitrary but matches a common business-data
+    # shape (six integer digits, four fractional) and keeps output narrow.
+    if col.type.typmod and len(col.type.typmod) >= 2:
+        precision, scale = col.type.typmod[0], col.type.typmod[1]
+    elif col.type.typmod and len(col.type.typmod) == 1:
+        # PG allows `numeric(P)` (scale implicit 0). Mirror that.
+        precision, scale = col.type.typmod[0], 0
+    else:
+        precision, scale = 10, 4
+    integer_digits = precision - scale
+    upper = 10**integer_digits - 1
+    # Pick a raw integer in the value space [−upper·10^scale, upper·10^scale],
+    # then divide by 10^scale. Doing integer arithmetic first keeps every
+    # produced Decimal exactly representable (no float rounding intrusion).
+    raw = rng.randint(-upper * 10**scale, upper * 10**scale)
+    return Decimal(raw) / (Decimal(10) ** scale)
+
+
+# Fixed reference epoch — NOT `datetime.now()`. Determinism contract:
+# same seed must produce same output years from now, which means no
+# wall-clock input anywhere in the generator.
+_EPOCH = _dt.date(2025, 1, 1)
+_WINDOW_DAYS = 5 * 365  # ±5 years
+
+
+def _date(rng: random.Random, col: Column) -> _dt.date:
+    days = rng.randint(-_WINDOW_DAYS, _WINDOW_DAYS)
+    return _EPOCH + _dt.timedelta(days=days)
+
+
+def _timestamptz(rng: random.Random, col: Column) -> _dt.datetime:
+    # Three rng calls in fixed order: days, seconds, microseconds.
+    # Changing the order would shift every downstream value — these
+    # are part of the determinism contract.
+    days = rng.randint(-_WINDOW_DAYS, _WINDOW_DAYS)
+    seconds = rng.randint(0, 86_399)
+    micros = rng.randint(0, 999_999)
+    # tz: stick with UTC. timestamptz stores UTC internally regardless
+    # of the input tz; emitting UTC keeps COPY output canonical and
+    # avoids planner statistics being affected by client TZ settings.
+    return _dt.datetime(
+        _EPOCH.year, _EPOCH.month, _EPOCH.day, tzinfo=_dt.timezone.utc,
+    ) + _dt.timedelta(days=days, seconds=seconds, microseconds=micros)
+
+
+def _interval(rng: random.Random, col: Column) -> _dt.timedelta:
+    # Bounded interval: at most a few months. PG intervals can encode
+    # year/month parts that timedelta cannot; the emitter formats as
+    # ISO-8601-like and the server accepts it cleanly.
+    days = rng.randint(0, 120)
+    seconds = rng.randint(0, 86_399)
+    return _dt.timedelta(days=days, seconds=seconds)
+
+
+def _jsonb(rng: random.Random, col: Column) -> dict:
+    """Shallow random object: 1-4 string keys, scalar values.
+
+    Deliberately not nested. The spec allows for richer JSON later, but
+    for parse/plan validation and demo readability, shallow is enough.
+
+    Key collisions across the loop are tolerated: when two iterations
+    pick the same word, the later iteration overwrites the earlier.
+    That's why output dicts may have fewer than `n_keys` entries; this
+    is intentional and keeps the rng-call count fixed at `n_keys` keys
+    regardless of collisions (every iteration consumes the same ticks).
+    """
+    n_keys = rng.randint(1, 4)
+    out: dict[str, object] = {}
+    for _ in range(n_keys):
+        key = pick_word(rng)
+        # Cap the dict at n_keys; collisions just overwrite, which is fine.
+        # Each branch below consumes EXACTLY ONE rng call (after the kind
+        # roll), so the total rng consumption per _jsonb is deterministic:
+        # 1 (n_keys) + n_keys × (1 key + 1 kind + 0-or-1 value).
+        kind = rng.randint(0, 4)
+        if kind == 0:
+            out[key] = pick_word(rng)
+        elif kind == 1:
+            out[key] = rng.randint(-1000, 1000)
+        elif kind == 2:
+            out[key] = round(rng.uniform(-1000.0, 1000.0), 3)
+        elif kind == 3:
+            out[key] = rng.choice((True, False))
+        else:
+            # kind == 4: literal JSON null. No rng call. The asymmetry
+            # (4 of 5 branches consume an rng tick, 1 doesn't) is
+            # intentional — null is a real JSON value, not an error path,
+            # and balancing rng consumption would force a dummy call.
+            out[key] = None
+    return out
+
+
+@dataclass(frozen=True)
+class _ColumnAdapter:
+    """Lightweight Column stand-in used when array strategies recurse on
+    the element type. Avoids depending on the full Column constructor
+    surface in case it grows constraints we don't care about here.
+    """
+    name: str
+    type: PgType
+    nullable: bool
+
+
+def _array(element_strategy: Strategy, element_type: PgType) -> Strategy:
+    """Build a strategy that returns a 0-5 element list. The factory
+    closes over the element strategy so call-time work is minimal —
+    we pay the `strategy_for_type` lookup once at factory time, not
+    once per row.
+
+    The closure pattern matters because `strategy_for_type` is called
+    every time a row generates an array column; without memoization via
+    the factory, recursive array types would re-resolve their element
+    strategies on every row.
+    """
+    def gen(rng: random.Random, col: Column) -> list:
+        n = rng.randint(0, 5)
+        # Element nullability is intentionally suppressed; NULL injection
+        # is an outer-row concern, not an element-level one. A PG array
+        # can contain NULL elements, but we don't generate them today —
+        # would require differently-quoted output ({NULL} not {""}) and
+        # the emit layer doesn't currently model that distinction.
+        elem_col = _ColumnAdapter(name=col.name, type=element_type, nullable=False)
+        return [element_strategy(rng, elem_col) for _ in range(n)]  # type: ignore[arg-type]
+    return gen
+
+
+# Type-name → strategy lookup. Keyed on `PgType.name` (matches pg_type.typname).
+# Adding a new scalar PgType requires adding an entry here AND making sure
+# the emit module's `encode_value` knows how to format the strategy's
+# return type. The two registries together define what the data generator
+# can produce.
+_TYPE_STRATEGIES: dict[str, Strategy] = {
+    "int4": _int4,
+    "int8": _int8,
+    "text": _text,
+    "varchar": _varchar,
+    "bool": _bool,
+    "uuid": _uuid,
+    "float8": _float8,
+    "numeric": _numeric,
+    "date": _date,
+    "timestamptz": _timestamptz,
+    "interval": _interval,
+    "jsonb": _jsonb,
+}
+
+
+def strategy_for_type(t: PgType) -> Strategy:
+    """Return the strategy for type `t`, or raise KeyError if no scalar strategy
+    is registered for the element type.
+
+    Array types are handled directly: we recurse on the element type and
+    wrap the result in `_array`. The element strategy is looked up once at
+    factory time (not per row), keeping data generation cheap. Recursion
+    bottoms out at scalar types, which are in `_TYPE_STRATEGIES`.
+    """
+    if t.is_array():
+        assert t.element is not None  # guaranteed by is_array()
+        return _array(strategy_for_type(t.element), t.element)
+    return _TYPE_STRATEGIES[t.name]