waxsql 1.0.0__py3-none-any.whl

waxsql/context.py

@@ -0,0 +1,255 @@
+"""Generation context — the bag of state threaded through every
+generator call.
+
+A `GenContext` holds:
+
+  * `rng` — the random.Random instance. Shared across the whole
+    generation run; advancing it in one call affects what later calls
+    see, which is exactly what makes the (seed → output) mapping
+    deterministic.
+
+  * `scope` — the current binding scope. Determines which columns are
+    visible at expression generation time.
+
+  * `schema`, `catalog` — read-only references to the world. Passed
+    through so generators can look up tables, FKs, functions,
+    operators without re-importing globals.
+
+  * `config` — the ComplexityConfig. Tells generators how aggressive
+    to be (how many FROM items, how deep, which features are unlocked).
+
+  * `depth_remaining` — the depth budget. Decremented on each
+    recursive descent; at zero, generators are required to pick leaf
+    productions only. This is what guarantees termination.
+
+  * Expression-context flags (`in_aggregate`, `allow_aggregates`,
+    `allow_window`, `allow_subquery`) — track what kind of expression
+    we're inside, gating which production candidates the expression
+    generator considers at each call site.
+
+`GenContext` is frozen. To "modify" it (descend into a sub-expression,
+swap scope for a subquery, flip an aggregate flag), use `replace(...)`
+or one of the helper methods like `descend()`. Forcing all changes
+through replace makes "use this for one descent, then revert" the
+default control-flow shape, which prevents whole categories of
+"forgot to reset the flag" bugs.
+
+Why this is threaded explicitly through every call instead of being
+held in a module-level / thread-local / context-var slot:
+
+  * Testability. A test can build a GenContext directly with a
+    hand-chosen rng/scope/config and call a generator function in
+    isolation. Thread-local state would force tests through a
+    setup-fixture-then-call shape that's harder to read and easier
+    to leak between tests.
+
+  * Parallelism. Multiple generation runs (e.g. parallel fuzzing
+    workers in one process) can interleave without contending on
+    a shared context slot. Each worker holds its own GenContext.
+
+  * Reproducibility. The state that controls "what gets generated"
+    is visible in the function signature. Reading a generator and
+    tracing what's in scope doesn't require knowing about an
+    ambient module-level variable that some other call mutated.
+
+The cost is verbosity — `ctx` shows up in nearly every signature.
+That's a price worth paying for the property that "given this ctx,
+this function does X" is true regardless of where it's called from.
+"""
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass, field, replace
+
+from .catalog import Catalog
+from .config import ComplexityConfig
+from .schema import Schema
+from .scope import Scope
+
+
+@dataclass
+class Counter:
+    """Mutable monotonic counter used as a cell on the frozen GenContext.
+
+    Replaces the bare `list[int]`-of-one idiom previously used for alias
+    and CTE numbering. The point is type-visibility: `alias_counter:
+    Counter` says "shared mutable counter" outright, and an accidental
+    reset via `dataclasses.replace(ctx, alias_counter=Counter())` reads
+    as obviously suspicious in a way `alias_counter=[1]` never did — that
+    silent-reset footgun is exactly what this guards against. Determinism
+    depends on the produced value sequence being identical across runs,
+    so the only operation is a monotonic advance.
+    """
+
+    value: int = 1
+
+    def take(self, n: int = 1) -> int:
+        """Return the current value, then advance by `n` (default 1).
+
+        The returned value names an alias or CTE; because `replace()`
+        preserves this same Counter instance (not a copy), the
+        post-advance state is immediately visible to sibling generators
+        that share the context — which is what keeps sibling subqueries
+        from picking colliding alias/CTE numbers.
+        """
+        current = self.value
+        self.value += n
+        return current
+
+
+@dataclass(frozen=True)
+class GenContext:
+    """Generation state passed through the recursive generator calls.
+
+    Frozen — see module docstring. The contained `rng` and `scope`
+    objects are themselves mutable; that's intentional. The frozen
+    wrapper protects against re-binding the field, not against the
+    pointed-to object's natural mutation.
+    """
+
+    rng: random.Random
+    scope: Scope
+    schema: Schema
+    catalog: Catalog
+    config: ComplexityConfig
+    depth_remaining: int
+
+    # Expression-context flags. Defaults reflect "outermost SELECT
+    # body": no aggregate context active, but aggregates are allowed
+    # (in the SELECT list and HAVING; the SELECT generator flips the
+    # `allow_aggregates` flag off when descending into WHERE).
+    in_aggregate: bool = False
+    allow_aggregates: bool = True
+    allow_window: bool = False  # True in SELECT-list expr generation
+    in_window: bool = False     # True inside a window FuncCall's args
+    allow_subquery: bool = True
+    # True iff the current ctx is a body of a correlated subquery
+    # (or LATERAL derived table) — anything where the inner can see
+    # outer columns via the parent-chain walk in scope. When True
+    # AND in_aggregate=True, gen_expr restricts column-ref candidates
+    # to the LOCAL scope only — outer-column refs inside aggregates
+    # trigger PG's implicit-grouping inference (PARSE-tier 42803).
+    in_correlated_subquery: bool = False
+    # Whether a WITH clause may be generated at this gen_select call.
+    # Top-level queries default to True; descend_subquery flips it
+    # off so subqueries / derived tables / CTE bodies don't generate
+    # their own nested WITHs (PG allows nested WITH but milestone 5
+    # keeps WITH top-level only — defer nested to a later milestone).
+    allow_with: bool = True
+
+    # Subquery-nesting budget. 0 (the default) means "no subqueries
+    # allowed", so code that builds a GenContext without thinking
+    # about subqueries gets the safe behavior. The public
+    # generate_query sets this from cfg.max_subquery_depth.
+    #
+    # Counted separately from depth_remaining because subquery
+    # nesting is a different rationing dimension — a deep
+    # `a + b + c` expression shouldn't share a budget with
+    # `(SELECT ... WHERE x IN (SELECT ...))`.
+    subquery_depth_remaining: int = 0
+
+    # Mutable monotonic counter for FROM-clause alias generation.
+    # Shared across all GenContexts in a single query (replace()
+    # preserves the same Counter instance), so sibling subqueries
+    # don't pick colliding aliases. See the Counter docstring for why
+    # this is a dedicated type rather than a bare list-of-one cell.
+    #
+    # Generator-internal: not part of the public API. Tests that
+    # build GenContext directly can ignore it.
+    alias_counter: Counter = field(default_factory=Counter)
+
+    # Companion monotonic counter for CTE-name generation. Separate
+    # from alias_counter because CTE names live in a different
+    # namespace (lookup_cte vs lookup_alias) — combining them would
+    # interleave `cte1, cte5, ...` with `t2, t3, t4, t6, ...` for
+    # no semantic benefit. Same shared-mutable-cell mechanism so
+    # nested subqueries' CTE names don't collide either.
+    cte_counter: Counter = field(default_factory=Counter)
+
+    def descend(self, cost: int = 1) -> "GenContext":
+        """Return a copy with `depth_remaining` decremented.
+
+        Every recursive expression-generator call must descend by at
+        least 1; that's what enforces termination. Higher-cost
+        productions can pass `cost > 1` to penalize expensive shapes
+        (e.g. function calls more than column refs); current callers
+        all use the default cost of 1, but the parameter is there for
+        when biasing depth budget by production cost becomes useful.
+
+        Note this only decrements the expression-depth budget. Use
+        `descend_subquery` when entering a SELECT body — that's a
+        different rationing dimension and inherits-or-resets a
+        different set of flags.
+        """
+        return replace(self, depth_remaining=self.depth_remaining - cost)
+
+    def at_leaf(self) -> bool:
+        """True iff the depth budget is exhausted; generators must
+        pick leaf productions only.
+
+        Centralizing this check makes the budget rule one line in
+        the generator instead of an inline `<= 0` test scattered
+        across files.
+        """
+        return self.depth_remaining <= 0
+
+    def at_subquery_leaf(self) -> bool:
+        """True iff no further subqueries can be opened. Used by
+        gen_expr to gate subquery candidates the same way at_leaf
+        gates recursive expression productions."""
+        return self.subquery_depth_remaining <= 0
+
+    def descend_subquery(self, *, correlated: bool) -> "GenContext":
+        """Return a child context for entering a subquery body.
+
+        The child context gets:
+
+          * a fresh child scope (correlated/uncorrelated per arg),
+            so inner FROM tables don't pollute the outer scope and
+            outer columns are visible only when correlated;
+          * decremented `subquery_depth_remaining`;
+          * fresh `depth_remaining` — the subquery's expression tree
+            starts with a full budget, not whatever's left of the
+            outer expression's;
+          * `allow_aggregates=True`, `in_aggregate=False` — the
+            subquery is its own expression context, so the outer
+            "no aggregates here" constraint doesn't propagate.
+
+        The shared rng / schema / catalog / config remain shared.
+
+        IMPORTANT — this method is the single chokepoint for entering
+        a child query body. Every expression-context flag must be
+        explicitly considered here: inheriting an outer flag through
+        replace()'s "preserve unspecified fields" default is exactly
+        the bug pattern this central method exists to prevent (e.g.
+        an inherited allow_window=True would leak window functions
+        into a subquery's WHERE clause). When adding a new flag to
+        GenContext, decide here whether it propagates or resets
+        before doing anything else.
+        """
+        return replace(
+            self,
+            scope=self.scope.push_subquery(correlated=correlated),
+            subquery_depth_remaining=self.subquery_depth_remaining - 1,
+            depth_remaining=self.config.max_expr_depth,
+            allow_aggregates=True,
+            in_aggregate=False,
+            # No nested WITH inside a subquery — milestone 5 keeps
+            # WITH top-level only.
+            allow_with=False,
+            # Window flags reset too. allow_window inherited as True
+            # from a SELECT-list-context parent would leak windows
+            # into the SUBQUERY's WHERE/HAVING/ON clauses (PARSE-tier
+            # error 42P20 "window functions are not allowed in WHERE").
+            # The subquery's gen_select will independently flip
+            # allow_window=True for ITS own SELECT-list expressions.
+            allow_window=False,
+            in_window=False,
+            # Track whether THIS ctx is a correlated-subquery body.
+            # Used by gen_expr's aggregate-arg column-ref candidate
+            # restriction (see in_correlated_subquery field doc).
+            in_correlated_subquery=correlated,
+        )
+
+
+__all__ = ["GenContext"]

waxsql/data.py

@@ -0,0 +1,99 @@
+"""Public API for the data generator.
+
+`generate_data(schema, *, seed, rows, fanout, null_fraction)` walks the
+schema's FK DAG topologically and produces a string containing one
+`COPY ... FROM STDIN; ...; \\.` block per table, in dependency order.
+
+Determinism contract: same (schema, seed, rows, fanout, null_fraction)
+produces byte-identical output across Python versions. The schema and
+data generators each construct their own `random.Random(seed)`, so
+they share no RNG state even when given the same seed.
+
+Role in the system: this is the thin orchestration layer. The heavy
+lifting (topological walk, per-row dispatch, COPY text formatting,
+strategy registries) lives under `waxsql.gen.data.*`; this module
+just stitches them together and owns the public function signature.
+Keeping the public surface here means callers can `from waxsql import
+generate_data` without ever touching the internal package layout.
+"""
+from __future__ import annotations
+
+import random
+
+from waxsql.gen.data.emit import emit_copy_block
+from waxsql.gen.data.rows import (
+    generate_row,
+    rows_for_table,
+    topological_order,
+)
+from waxsql.schema import Schema
+
+
+def generate_data(
+    schema: Schema,
+    *,
+    seed: int,
+    rows: int = 100,
+    fanout: int = 5,
+    null_fraction: float = 0.05,
+) -> str:
+    """Emit COPY blocks for every table in `schema`, in FK-topological order.
+
+    `rows` is the base row count; tables deeper in the FK DAG get
+    `rows * fanout ** depth`. `null_fraction` is the per-nullable-column
+    probability of emitting NULL. Same arguments + same schema → byte-
+    identical output.
+
+    `id_store` is populated as each table is finished — critically, the
+    current table's own IDs are added AFTER all its rows are materialized.
+    This lets the self-FK branch in `generate_row` use `rng.randint(1, pk)`
+    for forward-safe references without consulting `id_store` for the
+    current table.
+
+    Raises ValueError if the schema contains an FK cycle. Cycle handling
+    (deferred constraints + UPDATE-patches) is a known follow-up; today
+    the CLI catches this and surfaces a clean usage error.
+    """
+    # One RNG seeded once and threaded through every row. The data
+    # generator does NOT share RNG state with the schema generator —
+    # each constructs its own `random.Random(seed)` — so passing the
+    # same seed to both produces independent deterministic streams.
+    rng = random.Random(seed)
+    # Per-table list of PK values produced so far. Read by child tables
+    # to resolve their FK columns; written ONCE per table, after that
+    # table's rows are fully materialized (see ordering note below).
+    id_store: dict[str, list[int]] = {}
+    blocks: list[str] = []
+    for table in topological_order(schema):
+        n = rows_for_table(table, schema, base=rows, fanout=fanout)
+        table_rows: list[tuple[object, ...]] = []
+        ids: list[int] = []
+        for pk in range(1, n + 1):
+            # PKs are sequential 1..n. The schema generator emits
+            # `id BIGINT NOT NULL` (no sequence), so the data generator
+            # owns the PK numbering — no gaps, no collisions.
+            row = generate_row(
+                table=table,
+                pk=pk,
+                rng=rng,
+                id_store=id_store,
+                null_fraction=null_fraction,
+            )
+            ids.append(pk)
+            table_rows.append(row)
+        # Populate id_store only after all rows for this table are done.
+        # Children (tables with an FK to this table) will be visited later
+        # in topological order and can safely read from id_store then.
+        # This ordering is load-bearing: a child must not be able to see
+        # a parent's IDs mid-materialization (which could only happen
+        # under a buggy reordering), and self-FK resolution within the
+        # inner loop intentionally uses `rng.randint(1, pk)` instead of
+        # consulting id_store, so the current table is absent from the
+        # store throughout its own row loop.
+        id_store[table.name] = ids
+        columns = tuple(c.name for c in table.columns)
+        blocks.append(emit_copy_block(table.name, columns, table_rows))
+    # Blocks are joined with a single "\n" so callers can split/process
+    # them; each block already ends with "\.\n" so adjacent blocks are
+    # separated by a blank line in the rendered output.
+    return "\n".join(blocks)

waxsql/gen/__init__.py

@@ -0,0 +1,51 @@
+"""Query-generation entry points.
+
+Role: the public surface of the `gen/` subpackage. Re-exports a flat
+set of generator functions so callers (validators, tests, the CLI)
+can `from waxsql.gen import gen_select` without touching the internal
+file layout. The submodules themselves still import each other
+directly — the re-exports here are for *external* convenience, not
+intra-package wiring.
+
+The package exposes a flat set of generator functions, one per
+SQL surface. The two foundational ones are:
+
+  * `expr.gen_expr(ctx, target_type)` — typed expression generator.
+    The recursive core; every other generator that needs to emit an
+    expression goes through this.
+
+  * `select.gen_select(ctx)` — SELECT-statement generator. Picks a
+    FROM clause, populates the scope, then asks `gen_expr` for each
+    select-list / WHERE / ORDER BY expression.
+
+The rest (`subquery`, `cte`, `setop`, `window`) are full peers, each
+emitting a different SQL surface but all called from `gen_select` (or
+recursively from each other) for their part of the tree.
+
+CONTRACT: these generators MUTATE `ctx` as they run. The rng advances
+deterministically (intentional — that's the determinism guarantee),
+but `ctx.scope`, `ctx.alias_counter`, and `ctx.cte_counter`
+also get mutated as tables are added to the FROM clause and aliases
+are minted. Callers must NOT assume `ctx` is unchanged across a
+generator call. The `descend_subquery` chokepoint on GenContext is
+where children get a fresh scope so cross-arm mutations don't
+collide; production code always goes through it.
+"""
+from .cte import gen_cte_def, gen_recursive_cte_def
+from .expr import gen_expr, gen_literal
+from .select import gen_select
+from .setop import gen_setop
+from .subquery import (
+    gen_derived_table, gen_exists_subquery, gen_in_subquery,
+    gen_scalar_subquery,
+)
+from .window import gen_window_spec
+
+__all__ = [
+    "gen_expr", "gen_literal", "gen_select",
+    "gen_scalar_subquery", "gen_exists_subquery", "gen_in_subquery",
+    "gen_derived_table",
+    "gen_cte_def", "gen_recursive_cte_def",
+    "gen_window_spec",
+    "gen_setop",
+]