From 1ad0a6fb18c3cb804ead25b1c34ea2eee3e3542c Mon Sep 17 00:00:00 2001 From: Luccas Gomes Date: Wed, 23 Sep 2026 07:16:30 -0300 Subject: [PATCH] Fix non-sequential _id gaps on mixed upsert batches merge_sql() numbered every incoming row with a single ROW_NUMBER() OVER() spanning the whole batch, but only unmatched (new) rows ever consumed that number for their _id (MAX(_id) + _rn). A batch mixing already-present unchanged rows with new rows left gaps: unchanged rows' _rn slots were never inserted anywhere. Partition the window on whether the row's primary key is found in the target table (via a self-join probe) so _rn is numbered 1..k within just the unmatched rows, independent of their position among matched rows in the same batch. Co-Authored-By: Claude Sonnet 5 --- .../infrastructure/engines/bigquery/lib.py | 14 +++++-- tests/engines/bigquery/test_tables.py | 39 ++++++++++++++++++- 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/datastore/infrastructure/engines/bigquery/lib.py b/datastore/infrastructure/engines/bigquery/lib.py index 0f1c45b..b1fe8dd 100644 --- a/datastore/infrastructure/engines/bigquery/lib.py +++ b/datastore/infrastructure/engines/bigquery/lib.py @@ -291,7 +291,11 @@ def merge_sql(table_ref: str, schema: dict, *, include_updated_at: bool = True) Matched rows update only when a non-PK column differs (so `_updated_at` advances only on real changes). Unmatched rows - insert with `_id` continuing from `MAX(_id) + _rn`. + insert with `_id` continuing from `MAX(_id) + _rn`, where `_rn` + is numbered only among the unmatched rows (via a self-join + against the target on the primary key, partitioned on match + status) so already-present rows in the batch don't consume `_id` + values and leave gaps in the sequence for the rows that do insert. """ fields = _user_fields(schema) pk = normalize_pk(schema) @@ -306,10 +310,14 @@ def merge_sql(table_ref: str, schema: dict, *, include_updated_at: bool = True) insert_cols = ", ".join(f"`{f['name']}`" for f in fields) insert_vals = ", ".join(f"S.`{f['name']}`" for f in fields) + pk_probe_on = " AND ".join(f"S0.`{n}` = X.`{n}`" for n in pk) parts = [ f"MERGE {table_ref} T", - f"USING (SELECT {using_cols}, ROW_NUMBER() OVER () AS _rn " - f"FROM UNNEST(JSON_QUERY_ARRAY(@rows)) AS r) S", + f"USING (SELECT S0.*, " + f"ROW_NUMBER() OVER (PARTITION BY X.`{pk[0]}` IS NULL) AS _rn " + f"FROM (SELECT {using_cols} " + f"FROM UNNEST(JSON_QUERY_ARRAY(@rows)) AS r) AS S0 " + f"LEFT JOIN {table_ref} AS X ON {pk_probe_on}) S", f"ON {on_clause}", ] if non_pk: diff --git a/tests/engines/bigquery/test_tables.py b/tests/engines/bigquery/test_tables.py index 9505f83..9b62bce 100644 --- a/tests/engines/bigquery/test_tables.py +++ b/tests/engines/bigquery/test_tables.py @@ -518,7 +518,13 @@ def test_merge_sql_renders_typed_extractors_on_match_update_no_match_insert() -> assert "JSON_VALUE(r, '$.label') AS `label`" in sql assert "PARSE_JSON(JSON_QUERY(r, '$.meta')) AS `meta`" in sql # USING attaches ROW_NUMBER() as _rn for auto-`_id` on NOT MATCHED. - assert "ROW_NUMBER() OVER () AS _rn" in sql + # It's numbered per match-status partition (via a self-join probe + # against the target on the PK) so unmatched rows are numbered + # 1..k among themselves, not by position in the whole batch — a + # batch of unmatched rows preceded by matched ones doesn't leave + # gaps in the `_id` sequence. + assert "LEFT JOIN `p.d.r` AS X ON S0.`id` = X.`id`" in sql + assert "ROW_NUMBER() OVER (PARTITION BY X.`id` IS NULL) AS _rn" in sql assert "ON T.`id` = S.`id`" in sql # WHEN MATCHED only fires when some non-PK column actually differs # — `_updated_at` advances on real changes, not on no-op upserts. @@ -558,6 +564,37 @@ def test_update_sql_renders_dml_update_keyed_on_primary_key() -> None: assert "MERGE" not in sql # plain DML, not MERGE +def test_merge_sql_numbers_unmatched_rows_within_their_own_partition() -> None: + """Regression test: `_rn` must be independent per match-status + partition so `_id = MAX(_id) + _rn` stays gap-free for unmatched + rows even when a batch mixes them with already-present rows. + + Before this fix, `_rn` was a single `ROW_NUMBER() OVER ()` over + the whole batch, so unmatched rows inherited their position among + *all* incoming rows (matched no-ops included) rather than their + position among just the unmatched ones — leaving permanent gaps + in `_id` whenever a batch mixed unchanged rows with new ones. + """ + sql = merge_sql( + "`p.d.r`", + { + "fields": [ + {"name": "id", "type": "integer"}, + {"name": "label", "type": "string"}, + ], + "primaryKey": ["id"], + }, + ) + # _rn comes from a window partitioned on whether the row's PK was + # found in a self-join against the target table — not a bare + # ROW_NUMBER() OVER () spanning the whole incoming batch. + assert "ROW_NUMBER() OVER () AS _rn" not in sql + assert "PARTITION BY X.`id` IS NULL) AS _rn" in sql + assert "LEFT JOIN `p.d.r` AS X ON S0.`id` = X.`id`" in sql + # `_id` assignment for NOT MATCHED rows is unchanged: MAX + _rn. + assert "(SELECT IFNULL(MAX(`_id`), 0) FROM `p.d.r`) + S._rn" in sql + + def test_merge_and_update_sql_reject_missing_primary_key() -> None: schema = {"fields": [{"name": "id", "type": "integer"}]} with pytest.raises(ValueError, match="primaryKey"):