model2data 1.0.0__tar.gz → 1.2.0__tar.gz

{model2data-1.0.0/model2data.egg-info → model2data-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model2data
-Version: 1.0.0
+Version: 1.2.0
 Summary: Generate analytics-ready datasets from DBML models
 Author: JB Analytica
 License-Expression: MIT

{model2data-1.0.0 → model2data-1.2.0}/README.md

@@ -5,6 +5,8 @@
 [![codecov](https://codecov.io/gh/JB-Analytica/model2data/branch/main/graph/badge.svg)](https://codecov.io/gh/JB-Analytica/model2data)
 [![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
 
+Built and maintained by [JB Analytica](https://www.jbanalytica.com/) — data platform architecture and analytics engineering.
+
 **Turn a data model into a running analytics stack in one command.**
 
 Give `model2data` a [DBML](https://dbml.dbdiagram.io/docs/) schema — hand-written or exported
@@ -13,17 +15,17 @@ from an existing database — and it generates realistic, relationship-preservin
 DuckDB or Postgres profile. No sample data to hunt down, no dbt boilerplate to hand-write, no
 production data to risk exposing.
 
+A working analytics stack — real (synthetic) data, tested dbt models, queryable in
+DuckDB — from a schema file, in seconds:
+
+![model2data generating a project and running it with dbt](https://raw.githubusercontent.com/JB-Analytica/model2data/main/assets/demo.gif)
+
 ```bash
 pip install model2data
 model2data --file examples/ecommerce.dbml --rows 200 --seed 42
 cd dbt_ecommerce && dbt build
 ```
 
-That's a working analytics stack — real (synthetic) data, tested dbt models, queryable in
-DuckDB — from a schema file, in seconds:
-
-![model2data generating a project and running it with dbt](https://raw.githubusercontent.com/JB-Analytica/model2data/main/assets/demo.gif)
-
 ---
 
 ## Why this exists
@@ -37,7 +39,8 @@ access required.
 - **Privacy-safe.** Nothing but a schema definition goes in; nothing but synthetic data comes out.
 - **Realistic, not random.** Column names are matched against ~35 common patterns — `email`,
   `first_name`, `city`, `phone`, `company`, ... — so a column called `email` gets real-looking
-  emails, not `Lorem ipsum` text.
+  emails, not `Lorem ipsum` text. Type a column with any Faker provider (`billing_country state`,
+  `sku ean13`) to pick its generator outright when the name is wrong for the data.
 - **Relationship-preserving.** Foreign keys resolve to real parent rows; tables are generated in
   dependency order.
 - **Deterministic.** Pass `--seed` and the same schema always produces the same data — safe to
@@ -95,8 +98,9 @@ flowchart LR
 
 1. **Parse.** Reads tables, columns, types, and `Ref` relationships from a DBML file.
 2. **Generate.** Produces synthetic values per column — typed generation for known SQL types
-   (int, date, timestamp, ...), name-aware inference for everything else (`email`, `phone`,
-   `city`, ...), foreign keys resolved against already-generated parent rows.
+   (int, date, timestamp, ...), then a Faker provider named as the type (`sku ean13`), then
+   name-aware inference for everything else (`email`, `phone`, `city`, ...), foreign keys
+   resolved against already-generated parent rows.
 3. **Scaffold.** Writes a complete dbt project around that data: CSV seeds, staging models that
    `ref()` those seeds, `not_null`/`unique`/`relationships` tests, `accepted_values` tests for
    DBML `Enum`-typed columns, singular SQL tests for composite primary/unique keys, table and
@@ -126,6 +130,14 @@ model2data --file examples/ecommerce.dbml --rows 200 --seed 42
 
 This creates a `dbt_ecommerce/` folder with your data and dbt setup.
 
+Real schemas are rarely uniform. `--rows-for` sizes individual tables, so a handful of customers
+can sit behind a large orders table the way they would in the warehouse you're modelling:
+
+```bash
+model2data --file examples/ecommerce.dbml --rows 200 --seed 42 \
+  --rows-for customers=50 --rows-for order_items=5000
+```
+
 Run dbt to load, transform, and test the data:
 
 ```bash
@@ -251,10 +263,11 @@ If you're pinned to an older dbt-core, use model2data 0.5.x, which supported dow
 
 ## Project status
 
-As of `1.0.0`, model2data is considered **feature-complete for its intended use case**: turning a
-DBML schema into realistic synthetic data and a runnable dbt project, reliably. There's no active
-roadmap of new capabilities planned — the focus from here is maintenance: bug fixes, keeping pace
-with new dbt-core releases, and reviewing community contributions.
+As of `1.0.0`, model2data is stable and feature-complete for its intended use case: turning a
+DBML schema into realistic synthetic data and a runnable dbt project, reliably. It is actively
+maintained — CI runs a real `dbt build` against both the oldest supported dbt-core and the newest
+release on every push, so compatibility is proven rather than assumed. New capabilities are not
+the focus; correctness, dbt-core compatibility and community contributions are.
 
 Ideas that came up during development but were deliberately left out of scope, in case anyone
 wants to pick them up as a contribution:

{model2data-1.0.0 → model2data-1.2.0}/model2data/cli.py

@@ -27,6 +27,53 @@ from model2data.utils import normalize_identifier
 
 SUPPORTED_ADAPTERS = ("duckdb", "postgres")
 
+
+def _parse_row_overrides(
+    raw: Optional[list[str]],
+    tables: dict,
+) -> dict[str, int]:
+    """Turn repeated `--rows-for TABLE=N` values into a {table: rows} mapping.
+
+    Fails loudly rather than silently ignoring a typo: naming a table that
+    isn't in the schema almost always means a misspelling, and quietly
+    generating the default row count for it would be discovered only by
+    counting rows in the output.
+    """
+    # `main` is also called directly as a plain function (see tests/), which
+    # bypasses Typer and leaves this parameter holding its `OptionInfo` default
+    # rather than None. Anything that isn't an actual list means "not supplied".
+    if not isinstance(raw, (list, tuple)):
+        return {}
+
+    overrides: dict[str, int] = {}
+    for item in raw:
+        table_name, separator, count = item.partition("=")
+        table_name = table_name.strip()
+        if not separator or not table_name:
+            raise typer.BadParameter(f"Expected TABLE=N, got {item!r}.", param_hint="--rows-for")
+
+        try:
+            rows = int(count)
+        except ValueError:
+            raise typer.BadParameter(
+                f"Row count for {table_name!r} must be a whole number, got {count!r}.",
+                param_hint="--rows-for",
+            ) from None
+        if rows < 1:
+            raise typer.BadParameter(
+                f"Row count for {table_name!r} must be at least 1, got {rows}.",
+                param_hint="--rows-for",
+            )
+        if table_name not in tables:
+            known = ", ".join(sorted(tables)) or "none"
+            raise typer.BadParameter(
+                f"No table named {table_name!r} in this schema. Tables: {known}.",
+                param_hint="--rows-for",
+            )
+        overrides[table_name] = rows
+    return overrides
+
+
 app = typer.Typer(
     help=(
         "model2data: Generate analytics-ready datasets from DBML models.\n\n"
@@ -59,6 +106,17 @@ def main(
         min=10,
         help="Number of rows to generate per table.",
     ),
+    # noqa: B008 is only needed here (not on the other options) because a
+    # repeatable option must be annotated with a mutable `list` type.
+    rows_for: Optional[list[str]] = typer.Option(  # noqa: B008
+        None,
+        "--rows-for",
+        metavar="TABLE=N",
+        help=(
+            "Row count for one table, overriding --rows. Repeatable, e.g.\n"
+            "--rows-for customers=200 --rows-for orders=5000."
+        ),
+    ),
     seed: Optional[int] = typer.Option(
         None,
         "--seed",
@@ -124,6 +182,11 @@ def main(
         typer.echo("❌ No tables found in the provided DBML file.")
         raise typer.Exit(1)
 
+    # Validated before anything touches the filesystem: a typo'd table name here
+    # should not leave a half-scaffolded project behind for the next run to trip
+    # over with a confusing "destination already exists".
+    row_overrides = _parse_row_overrides(rows_for, tables)
+
     project_name = normalize_identifier(name or file.stem)
     dest = Path.cwd() / f"dbt_{project_name}"
     profile_name = f"{project_name}_profile"
@@ -150,6 +213,7 @@ def main(
         refs=refs,
         base_rows=rows,
         seed=seed,
+        row_overrides=row_overrides,
     )
 
     # -------------------------

{model2data-1.0.0 → model2data-1.2.0}/model2data/generate/core.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import random
 from collections import defaultdict, deque
+from collections.abc import Mapping
 from typing import Optional
 
 import pandas as pd
@@ -64,10 +65,19 @@ def generate_data_from_dbml(
     refs: list[dict],
     base_rows: int = 100,
     seed: Optional[int] = None,
+    row_overrides: Optional[Mapping[str, int]] = None,
 ) -> dict[str, pd.DataFrame]:
     """
     Generate synthetic datasets from parsed DBML definitions.
 
+    `base_rows` is the row count for every table; `row_overrides` sets it per
+    table, keyed by DBML table name. Real schemas are rarely uniform -- a
+    handful of dimension rows against a fact table two orders of magnitude
+    larger is the normal shape, and generating 100 of each makes joins and
+    aggregates behave nothing like the warehouse being modelled. Names not
+    present in `row_overrides` fall back to `base_rows`; unknown names are
+    ignored.
+
     This function is deterministic if a seed is provided.
     It performs no filesystem I/O and returns pandas DataFrames.
     """
@@ -93,7 +103,7 @@ def generate_data_from_dbml(
 
     for table_name in ordered_tables:
         table_def = tables[table_name]
-        row_count = _determine_row_count(table_def.name, base_rows)
+        row_count = _determine_row_count(table_def.name, base_rows, row_overrides)
 
         data: dict[str, list] = {}
 
@@ -341,10 +351,16 @@ def _resolve_self_referencing_fks(
     return df
 
 
-def _determine_row_count(table_name: str, base_rows: int) -> int:
-    """
-    Return the base number of rows for all tables.
-    """
+def _determine_row_count(
+    table_name: str,
+    base_rows: int,
+    row_overrides: Optional[Mapping[str, int]] = None,
+) -> int:
+    """Return the row count for one table: its override, else `base_rows`."""
+    if row_overrides:
+        override = row_overrides.get(table_name)
+        if override is not None:
+            return override
     return base_rows
 
 

{model2data-1.0.0 → model2data-1.2.0}/model2data/generate/faker.py

@@ -150,6 +150,33 @@ def _infer_by_name(column_name: str) -> Optional[Callable[[], object]]:
     return None
 
 
+# Type names that are ordinary SQL types first and Faker providers only by
+# coincidence. Someone typing `email text` means the SQL type, so these must
+# not count as a deliberate choice of provider -- otherwise the commonest
+# declaration in any schema would stop generating emails.
+_SQL_TYPES_SHADOWING_A_PROVIDER = frozenset({"text", "json", "jsonb", "xml", "binary", "year"})
+
+
+def _infer_by_type(base_type: str) -> Optional[Callable[[], object]]:
+    """The provider a column's declared type names, if it names one deliberately.
+
+    `sku ean13` and `home_state state` are the user saying which generator they
+    want, in the only place DBML gives them to say it. That has to outrank the
+    guess made from the column's name: a name pattern is inferred, a type is
+    declared, and `first_name email` silently generating first names -- the
+    type having no effect whatsoever -- is the single most confusing thing this
+    module did.
+    """
+    if base_type in _SQL_TYPES_SHADOWING_A_PROVIDER:
+        return None
+    try:
+        fake.format(base_type)
+    except (AttributeError, TypeError):
+        # Not a provider, or one that needs arguments: nothing was declared.
+        return None
+    return lambda: fake.format(base_type)
+
+
 # ---------------------------------------------------------
 # Public API
 # ---------------------------------------------------------
@@ -278,16 +305,16 @@ def generate_column_values(
         values = [_random_datetime().isoformat(sep=" ") for _ in range(row_count)]
 
     # -----------------------------------------------------
-    # Untyped / generic string columns: infer intent from the
-    # column name first (email, city, phone...), then fall back
-    # to a literal Faker provider name, then to a generic value.
+    # Untyped / generic string columns: honour a type that names
+    # a Faker provider (`sku ean13`), then infer intent from the
+    # column name (email, city, phone...), then a generic value.
     # -----------------------------------------------------
     else:
-        name_generator = _infer_by_name(column.name)
-        if name_generator is not None:
-            values = [name_generator() for _ in range(row_count)]
+        generator = _infer_by_type(base_type) or _infer_by_name(column.name)
+        if generator is not None:
+            values = [generator() for _ in range(row_count)]
             values = (
-                _deduplicate(values, name_generator, column_name=unique_label)
+                _deduplicate(values, generator, column_name=unique_label)
                 if ensure_unique
                 else values
             )

{model2data-1.0.0 → model2data-1.2.0/model2data.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model2data
-Version: 1.0.0
+Version: 1.2.0
 Summary: Generate analytics-ready datasets from DBML models
 Author: JB Analytica
 License-Expression: MIT

{model2data-1.0.0 → model2data-1.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "model2data"
-version = "1.0.0"
+version = "1.2.0"
 description = "Generate analytics-ready datasets from DBML models"
 readme = "README_PYPI.md"
 requires-python = ">=3.10"

{model2data-1.0.0 → model2data-1.2.0}/tests/test_cli.py

@@ -777,3 +777,58 @@ def test_cli_summary_is_quiet_when_uniqueness_is_satisfiable(tmp_path):
     assert result.exit_code == 0
     assert "duplicate row(s)" not in result.stdout
     assert "duplicate value(s)" not in result.stdout
+
+
+ROWS_FOR_SCHEMA = """
+Table users {
+    id int [pk]
+    email email [unique]
+}
+Table orders {
+    id int [pk]
+    user_id int [not null]
+}
+Ref: orders.user_id > users.id
+"""
+
+
+def _run_rows_for(tmp_path, *extra_args):
+    dbml_file = tmp_path / "shop.dbml"
+    dbml_file.write_text(ROWS_FOR_SCHEMA)
+    cwd = os.getcwd()
+    os.chdir(tmp_path)
+    try:
+        return runner.invoke(app, ["--file", str(dbml_file), "--rows", "50", *extra_args])
+    finally:
+        os.chdir(cwd)
+
+
+def test_cli_rows_for_sets_row_counts_per_table(tmp_path):
+    result = _run_rows_for(tmp_path, "--rows-for", "orders=120")
+    assert result.exit_code == 0, result.output
+
+    seeds = tmp_path / "dbt_shop" / "seeds" / "raw"
+    assert len(pd.read_csv(seeds / "orders.csv")) == 120
+    assert len(pd.read_csv(seeds / "users.csv")) == 50
+
+
+def test_cli_rows_for_is_repeatable(tmp_path):
+    result = _run_rows_for(tmp_path, "--rows-for", "users=10", "--rows-for", "orders=200")
+    assert result.exit_code == 0, result.output
+
+    seeds = tmp_path / "dbt_shop" / "seeds" / "raw"
+    assert len(pd.read_csv(seeds / "users.csv")) == 10
+    assert len(pd.read_csv(seeds / "orders.csv")) == 200
+
+
+def test_cli_rows_for_rejects_an_unknown_table(tmp_path):
+    """A typo'd table name would otherwise be discovered by counting rows."""
+    result = _run_rows_for(tmp_path, "--rows-for", "userz=10")
+    assert result.exit_code != 0
+    assert "No table named 'userz'" in result.output
+
+
+def test_cli_rows_for_rejects_malformed_values(tmp_path):
+    assert "Expected TABLE=N" in _run_rows_for(tmp_path, "--rows-for", "users").output
+    assert "whole number" in _run_rows_for(tmp_path, "--rows-for", "users=lots").output
+    assert "at least 1" in _run_rows_for(tmp_path, "--rows-for", "users=0").output

{model2data-1.0.0 → model2data-1.2.0}/tests/test_faker_name_inference.py

@@ -82,6 +82,73 @@ class TestNameInference:
         assert result == ["x", "x", "x"]
 
 
+class TestTypeBeatsName:
+    """A declared Faker provider type outranks the guess made from the name.
+
+    DBML has one place to say which generator a column should use -- its type --
+    and until this held, a recognised column name silently overruled it:
+    `home_state state` generated countries and `first_name email` generated
+    first names, with the declared type having no effect at all.
+    """
+
+    def test_a_declared_provider_overrules_the_name_pattern(self):
+        values = generate_column_values(
+            ColumnDef(name="first_name", data_type="email", settings={"not null"}),
+            row_count=20,
+        )
+        assert all(EMAIL_RE.match(v) for v in values)
+
+    def test_a_column_can_be_typed_against_its_own_name(self):
+        # The case the feature exists for: a state column that isn't called one.
+        # Without this the `country` in its name won and it generated countries.
+        from faker.providers.address.en_US import Provider as UsAddress
+
+        values = generate_column_values(
+            ColumnDef(name="billing_country", data_type="state", settings={"not null"}),
+            row_count=40,
+        )
+        assert set(values) <= set(UsAddress.states)
+
+    def test_a_plain_sql_type_does_not_count_as_a_declaration(self):
+        # `text` and `json` are Faker providers by coincidence; a column typed
+        # that way means the SQL type, and its name must still be inferred.
+        for sql_type in ("text", "json", "varchar", "char(64)"):
+            values = generate_column_values(
+                ColumnDef(name="email", data_type=sql_type, settings={"not null"}),
+                row_count=10,
+            )
+            assert all(EMAIL_RE.match(v) for v in values), sql_type
+
+    def test_an_unknown_type_still_falls_back_to_the_name(self):
+        reset_stats()
+        values = generate_column_values(
+            ColumnDef(name="city", data_type="weird_custom_type", settings={"not null"}),
+            row_count=10,
+        )
+        assert all(isinstance(value, str) and value for value in values)
+        assert get_unmapped_columns() == []
+
+    def test_a_declared_provider_is_still_deduplicated_when_unique(self):
+        values = generate_column_values(
+            ColumnDef(name="first_name", data_type="email", settings={"not null", "unique"}),
+            row_count=50,
+            ensure_unique=True,
+        )
+        assert len(values) == len(set(values)) == 50
+
+    def test_structured_types_are_untouched_by_this(self):
+        # `date` and `boolean` are providers too, but they are decided long
+        # before either inference runs, and must keep their real Python types.
+        dates = generate_column_values(
+            ColumnDef(name="signup_date", data_type="date", settings={"not null"}), row_count=5
+        )
+        assert all(hasattr(value, "year") for value in dates)
+        flags = generate_column_values(
+            ColumnDef(name="is_active", data_type="boolean", settings={"not null"}), row_count=5
+        )
+        assert set(flags) <= {True, False}
+
+
 class TestEnumGeneration:
     def test_enum_values_only_ever_come_from_the_enum_set(self):
         allowed = {"active", "inactive", "pending"}

{model2data-1.0.0 → model2data-1.2.0}/tests/test_generation.py

@@ -670,7 +670,7 @@ def test_enum_column_whose_name_contains_int_does_not_crash_dtype_coercion():
     assert set(df["maintenance_type"]) <= {"oil_change", "repair", "cleaning"}
 
 
-def test_composite_key_dedup_retry_preserves_fk_validity(monkeypatch):
+def test_composite_key_dedup_retry_preserves_fk_validity():
     # Regression test: a join/bridge table's composite key is almost always
     # built from FK columns (posts/tags -> post_tags is the canonical
     # example). When two parent tables are small relative to the bridge
@@ -679,7 +679,6 @@ def test_composite_key_dedup_retry_preserves_fk_validity(monkeypatch):
     # type-based generator instead of resampling from the real parent id
     # pool, silently producing post_id/tag_id values that referenced no
     # real parent row at all.
-    import model2data.generate.core as core
 
     tables = {
         "posts": TableDef(name="posts", columns=[ColumnDef("id", "int", {"pk"})]),
@@ -710,12 +709,13 @@ def test_composite_key_dedup_retry_preserves_fk_validity(monkeypatch):
 
     # Force posts/tags to 5 rows each (25 possible combos) while post_tags
     # gets 200 rows, guaranteeing heavy dedup-retry activity.
-    def small_parents_large_bridge(table_name, base_rows):
-        return 5 if table_name in ("posts", "tags") else base_rows
-
-    monkeypatch.setattr(core, "_determine_row_count", small_parents_large_bridge)
-
-    data = generate_data_from_dbml(tables, refs, base_rows=200, seed=123)
+    data = generate_data_from_dbml(
+        tables,
+        refs,
+        base_rows=200,
+        seed=123,
+        row_overrides={"posts": 5, "tags": 5},
+    )
 
     post_ids = set(data["posts"]["id"].tolist())
     tag_ids = set(data["tags"]["id"].tolist())
@@ -885,3 +885,76 @@ def test_timestamp_generation_is_reproducible_across_runs():
     assert list(first["created_at"]) == list(second["created_at"])
     # No sub-second component at all, so a CSV round-trip stays stable too.
     assert all(ts.endswith(":00") or "." not in ts for ts in map(str, first["created_at"]))
+
+
+def _users_and_orders() -> tuple[dict, list[dict]]:
+    tables = {
+        "users": TableDef(
+            name="users",
+            columns=[ColumnDef("id", "int", {"pk"}), ColumnDef("email", "email", {"unique"})],
+        ),
+        "orders": TableDef(
+            name="orders",
+            columns=[
+                ColumnDef("id", "int", {"pk"}),
+                ColumnDef("user_id", "int", {"not null"}),
+                ColumnDef("total", "numeric"),
+            ],
+        ),
+    }
+    refs = [
+        {
+            "source_table": "orders",
+            "source_column": "user_id",
+            "target_table": "users",
+            "target_column": "id",
+        }
+    ]
+    return tables, refs
+
+
+def test_row_overrides_set_row_counts_per_table():
+    tables, refs = _users_and_orders()
+
+    data = generate_data_from_dbml(
+        tables, refs, base_rows=50, seed=1, row_overrides={"orders": 300}
+    )
+
+    assert len(data["orders"]) == 300
+    assert len(data["users"]) == 50, "tables without an override keep base_rows"
+
+
+def test_row_overrides_default_to_base_rows_and_ignore_unknown_tables():
+    tables, refs = _users_and_orders()
+
+    data = generate_data_from_dbml(
+        tables, refs, base_rows=20, seed=1, row_overrides={"nonexistent": 999}
+    )
+
+    assert {name: len(df) for name, df in data.items()} == {"users": 20, "orders": 20}
+
+
+def test_foreign_keys_stay_valid_when_the_child_is_much_larger_than_its_parent():
+    """The point of per-table counts: a small dimension against a large fact."""
+    tables, refs = _users_and_orders()
+
+    data = generate_data_from_dbml(
+        tables, refs, base_rows=100, seed=42, row_overrides={"users": 5, "orders": 2000}
+    )
+
+    assert len(data["users"]) == 5
+    assert len(data["orders"]) == 2000
+    assert set(data["orders"]["user_id"]) <= set(data["users"]["id"])
+    # A 5-row parent still has to have unique keys for the FK to mean anything.
+    assert data["users"]["id"].is_unique
+
+
+def test_row_overrides_are_deterministic_with_a_seed():
+    tables, refs = _users_and_orders()
+    overrides = {"users": 25, "orders": 175}
+
+    first = generate_data_from_dbml(tables, refs, base_rows=50, seed=9, row_overrides=overrides)
+    second = generate_data_from_dbml(tables, refs, base_rows=50, seed=9, row_overrides=overrides)
+
+    assert first["users"].equals(second["users"])
+    assert first["orders"].equals(second["orders"])

{model2data-1.0.0 → model2data-1.2.0}/tests/test_release_stress.py

@@ -26,7 +26,6 @@ from pathlib import Path
 
 import pytest
 
-import model2data.generate.core as core
 from model2data.generate.core import generate_data_from_dbml
 from model2data.generate.relationships import build_fk_lookup, classify_refs
 from model2data.parse.dbml import ColumnDef, TableDef, parse_dbml
@@ -187,7 +186,7 @@ def _bridge_schema():
 )
 @pytest.mark.parametrize("seed", list(range(1, 16)))  # 15 seeds
 def test_composite_key_dedup_extreme_value_space_stress(
-    seed, parent_rows, bridge_rows, uniqueness_achievable, monkeypatch
+    seed, parent_rows, bridge_rows, uniqueness_achievable
 ):
     # Extends the original bug-1 regression test
     # (test_composite_key_dedup_retry_preserves_fk_validity in
@@ -199,12 +198,13 @@ def test_composite_key_dedup_extreme_value_space_stress(
     # surface at all.
     tables, refs = _bridge_schema()
 
-    def sized_row_counts(table_name, base_rows):
-        return parent_rows if table_name in ("posts", "tags") else base_rows
-
-    monkeypatch.setattr(core, "_determine_row_count", sized_row_counts)
-
-    data = generate_data_from_dbml(tables, refs, base_rows=bridge_rows, seed=seed)
+    data = generate_data_from_dbml(
+        tables,
+        refs,
+        base_rows=bridge_rows,
+        seed=seed,
+        row_overrides={"posts": parent_rows, "tags": parent_rows},
+    )
 
     post_ids = set(data["posts"]["id"].tolist())
     tag_ids = set(data["tags"]["id"].tolist())