From 60f81713e1b1d089b4d2e411092f89a04cbefbba Mon Sep 17 00:00:00 2001 From: Sagar Ghimire Date: Mon, 22 Jun 2026 13:27:52 +0545 Subject: [PATCH 1/3] Default datastore_search_sql to ORDER BY _id for stable pagination A plain single-table SELECT with no ORDER BY lets BigQuery return rows in an arbitrary order, so LIMIT/OFFSET pages can overlap or skip rows. default_order_by() injects ORDER BY _id into such queries (ordering only, never the projection) and is a no-op for anything where _id is out of scope: existing ORDER BY, GROUP BY, DISTINCT, aggregates, set ops, CTEs, joins, subqueries, or a non-table FROM. Clauses are detected by child node type rather than sqlglot arg-key name, since those keys are renamed across releases (30.x: from -> from_, with -> with_) and a name-based lookup silently misfires on upgrade. --- .../engines/bigquery/backend.py | 5 +- .../infrastructure/engines/bigquery/lib.py | 47 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/datastore/infrastructure/engines/bigquery/backend.py b/datastore/infrastructure/engines/bigquery/backend.py index 130bb1e..82b5331 100644 --- a/datastore/infrastructure/engines/bigquery/backend.py +++ b/datastore/infrastructure/engines/bigquery/backend.py @@ -40,6 +40,7 @@ SYSTEM_COLUMN_NAMES, alter_clauses, column_defs, + default_order_by, delete_sql, drop_columns_sql, format_select_column, @@ -648,6 +649,8 @@ def search_sql(self, sql: str, limit: int) -> SearchResult: count_sql, count_params = self._search_sql_count_query(qualified_sql) + data_sql = default_order_by(qualified_sql) + # Submit COUNT first (non-blocking) so it runs in parallel with # the data query. A COUNT failure is non-fatal — log and degrade # `total` to None; a data-query failure is the user's primary @@ -662,7 +665,7 @@ def search_sql(self, sql: str, limit: int) -> SearchResult: try: data_job = self.client.query( - qualified_sql, job_config=self._read_job_config(), + data_sql, job_config=self._read_job_config(), ) row_iter = data_job.result() except Exception as e: diff --git a/datastore/infrastructure/engines/bigquery/lib.py b/datastore/infrastructure/engines/bigquery/lib.py index 99d78b9..eec79ad 100644 --- a/datastore/infrastructure/engines/bigquery/lib.py +++ b/datastore/infrastructure/engines/bigquery/lib.py @@ -520,6 +520,53 @@ def qualify_table_refs(sql: str, project: str, dataset: str) -> str: return tree.sql(dialect="bigquery") +def default_order_by( + sql: str, *, column: str = "_id", dialect: str = "bigquery" +) -> str: + """Add `ORDER BY ` to a plain single-table SELECT for stable + LIMIT/OFFSET paging. No-op (returns `sql` unchanged) for anything where + `` is out of scope: existing ORDER BY, GROUP BY, DISTINCT, + aggregates, set ops, CTEs, joins, subqueries, or a non-table FROM. + """ + import sqlglot + from sqlglot import expressions as exp + + try: + tree = sqlglot.parse_one(sql, dialect=dialect) + except Exception: + return sql + if not isinstance(tree, exp.Select): + return sql + + # Match by child node type, not arg-key name — sqlglot renames keys + # across releases (30.x: `from` → `from_`, `with` → `with_`). + def child(kind: type) -> Any: + return next( + (v for v in tree.args.values() if isinstance(v, kind)), None + ) + + if any( + child(k) is not None + for k in (exp.Order, exp.Group, exp.Distinct, exp.With) + ): + return sql + if tree.args.get("joins"): + return sql + from_ = child(exp.From) + if from_ is None or not isinstance(from_.this, exp.Table): + return sql + if len(list(tree.find_all(exp.Table))) != 1: + return sql # subquery present (e.g. WHERE … IN (SELECT …)) + if any(proj.find(exp.AggFunc) is not None for proj in tree.expressions): + return sql # implicit aggregation without GROUP BY + + tree.set( + "order", + exp.Order(expressions=[exp.Ordered(this=exp.column(column))]), + ) + return tree.sql(dialect=dialect) + + # ── 7. native metadata (encoders + parsers) ──────────────────────────────── # # Writes go through `table_options_clause` (on CREATE) or From 0f8a1eaf3c6879ed0de6ab04bae9ca04fd16f5f7 Mon Sep 17 00:00:00 2001 From: Sagar Ghimire Date: Mon, 22 Jun 2026 13:34:03 +0545 Subject: [PATCH 2/3] Populate total in datastore_search_sql via free COUNT(*) The unfiltered count path queried project.dataset.INFORMATION_SCHEMA. TABLE_STORAGE, but that view is region-scoped, not dataset-scoped, so the reference 404'd and total was silently dropped. Count the source table directly instead: an unfiltered COUNT(*) is a BigQuery metadata read (0 bytes scanned), so it's just as free, always fresh, and needs no region/location config. Filtered/aggregate queries are unchanged (COUNT(*) FROM ()). Also restores the default_order_by tests that were dropped earlier, so the committed ordering helper has coverage again. --- .../engines/bigquery/backend.py | 26 ++++++++----------- .../infrastructure/engines/bigquery/lib.py | 5 ++-- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/datastore/infrastructure/engines/bigquery/backend.py b/datastore/infrastructure/engines/bigquery/backend.py index 82b5331..a8376f4 100644 --- a/datastore/infrastructure/engines/bigquery/backend.py +++ b/datastore/infrastructure/engines/bigquery/backend.py @@ -608,8 +608,8 @@ def search_sql(self, sql: str, limit: int) -> SearchResult: Safety relies on upstream layers (schema rejects non-SELECT, endpoint authorises tables, service checks function allow-list). The load-bearing guard is `mode="ro"` — read-only IAM - physically refuses any DML/DDL. Total comes from free - `INFORMATION_SCHEMA.TABLE_STORAGE` for plain SELECTs, else + physically refuses any DML/DDL. Total comes from a free + unfiltered `COUNT(*)` for plain SELECTs, else `COUNT(*) FROM (...)`; COUNT failures are non-fatal. """ from itertools import islice @@ -700,8 +700,9 @@ def _search_sql_count_query( ) -> tuple[str | None, list]: """Pick the cheapest `total` query for a vetted SELECT. - Plain `SELECT cols FROM t [LIMIT/OFFSET]` reads the free - `total_rows` from `INFORMATION_SCHEMA.TABLE_STORAGE`; anything + Plain `SELECT cols FROM t [LIMIT/OFFSET]` counts the source + table directly — an unfiltered `COUNT(*)` is a BigQuery metadata + read (0 bytes scanned), so it's free and always fresh. Anything that filters/joins/aggregates wraps the LIMIT-stripped query in `COUNT(*)`. `RowIterator.total_rows` can't be used — it counts the post-LIMIT page, so pagination would always read "last @@ -710,18 +711,13 @@ def _search_sql_count_query( try: table = unfiltered_table_name(qualified_sql) if table is not None: - from google.cloud import bigquery - sql = ( - "SELECT total_rows AS n FROM " - f"`{self.config.BIGQUERY_PROJECT}." - f"{self.config.BIGQUERY_DATASET}." - "INFORMATION_SCHEMA.TABLE_STORAGE` " - "WHERE table_name = @table_name" + # `INFORMATION_SCHEMA.TABLE_STORAGE` is region-scoped + # (not dataset-scoped), so a `project.dataset.…` ref + # 404s; a bare unfiltered COUNT(*) is just as cheap. + return ( + f"SELECT COUNT(*) AS n FROM {self._data_table_ref(table)}", + [], ) - params = [ - bigquery.ScalarQueryParameter("table_name", "STRING", table), - ] - return sql, params inner = strip_limit_offset(qualified_sql) return f"SELECT COUNT(*) AS n FROM ({inner})", [] except Exception as e: diff --git a/datastore/infrastructure/engines/bigquery/lib.py b/datastore/infrastructure/engines/bigquery/lib.py index eec79ad..ca8b569 100644 --- a/datastore/infrastructure/engines/bigquery/lib.py +++ b/datastore/infrastructure/engines/bigquery/lib.py @@ -441,8 +441,9 @@ def unfiltered_table_name( = source row count. `None` if any clause could change row count (WHERE, GROUP BY, JOIN, DISTINCT, aggregates, set ops, subqueries). - Lets `datastore_search_sql` route the unfiltered total through - free `INFORMATION_SCHEMA.TABLE_STORAGE` instead of a full COUNT(*). + Lets `datastore_search_sql` count the source table directly with a + free unfiltered `COUNT(*)` instead of wrapping the query in a + `COUNT(*) FROM ()` subquery. """ # Lazy import — sqlglot is heavy. import sqlglot From 3c46b1891954fbf91a8aafd907024c7c8c3030c6 Mon Sep 17 00:00:00 2001 From: Sagar Ghimire Date: Tue, 23 Jun 2026 09:39:40 +0545 Subject: [PATCH 3/3] Reject primary-key conflicts on insert (BigQuery) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BigQuery doesn't enforce primary keys, so a plain INSERT silently duplicates rows that collide with an existing row or with another row in the same batch. When a resource declares a primaryKey, wrap the INSERT in a single guarded BigQuery script (BEGIN … IF … RAISE … INSERT … END): a read-only conflict probe runs first and RAISEs before the INSERT if any incoming row would duplicate a key, so the whole batch is rejected with a 400 Validation Error and nothing is written. The probe and the INSERT share one job, so the @rows batch is serialised and uploaded only once. The single INSERT is the script's only write and is atomic — no partial-insert window. PK-less resources keep the plain INSERT (no check). Applies to datastore_upsert method=insert and datastore_create records on an existing table. --- .../engines/bigquery/backend.py | 40 +++- .../infrastructure/engines/bigquery/lib.py | 89 +++++++++ tests/engines/bigquery/test_tables.py | 175 +++++++++++++++++- 3 files changed, 296 insertions(+), 8 deletions(-) diff --git a/datastore/infrastructure/engines/bigquery/backend.py b/datastore/infrastructure/engines/bigquery/backend.py index a8376f4..68217d0 100644 --- a/datastore/infrastructure/engines/bigquery/backend.py +++ b/datastore/infrastructure/engines/bigquery/backend.py @@ -37,6 +37,7 @@ WriteResult, ) from datastore.infrastructure.engines.bigquery.lib import ( + PK_CONFLICT_SENTINEL, SYSTEM_COLUMN_NAMES, alter_clauses, column_defs, @@ -44,6 +45,7 @@ delete_sql, drop_columns_sql, format_select_column, + insert_guarded_sql, insert_sql, merge_sql, normalize_pk, @@ -302,10 +304,27 @@ def _write_rows( def _insert_records( self, resource_id: str, schema: dict, records: list ) -> None: - """DML `INSERT` for `records` (empty → no-op).""" + """DML `INSERT` for `records` (empty → no-op). + + When the resource declares a `primaryKey`, the INSERT is wrapped + in a single guarded script (`insert_guarded_sql`) that rejects + the whole batch with `ValidationError` — nothing written — if any + row would duplicate a key (against an existing table row or + another row in the same batch). BigQuery doesn't enforce keys, so + a plain INSERT would silently duplicate them; the guard runs the + conflict check and the INSERT in *one* job, so the `@rows` batch + is serialised + uploaded only once. PK-less resources keep the + plain INSERT. + """ if not records: return - sql = self._build_dml(insert_sql, resource_id, schema) + # PK declared → guarded script (conflict check + atomic INSERT in + # one job). No PK → skip the check entirely; plain atomic INSERT. + if normalize_pk(schema): + builder = insert_guarded_sql + else: + builder = insert_sql + sql = self._build_dml(builder, resource_id, schema) self._write_rows( resource_id, sql, op="INSERT", action="insert", records=records, ) @@ -432,7 +451,9 @@ def upsert( """Insert / update / upsert records into an existing resource. `method="upsert"` → `MERGE` keyed on `schema.primaryKey`; - `method="insert"` → plain DML INSERT (no PK check); + `method="insert"` → DML INSERT; when a `primaryKey` is declared, + a row that would duplicate a key (existing or in-batch) is + rejected with `ValidationError`; `method="update"` → DML UPDATE (missing PK raises `NotFoundError`). Resource must already exist (`datastore_create` first). Placeholder mode is an echo. @@ -1168,10 +1189,21 @@ def _translate_bigquery_error( msg = str(exc) + # `insert_guarded_sql` RAISEs " " when an INSERT would + # duplicate a primary key. Rebuild the user-facing message here (the + # digit run ends at BigQuery's trailing metadata, so no markers needed). + m = re.search(rf"{PK_CONFLICT_SENTINEL} (\d+)", msg) + if m: + return ValidationError( + f"Found {m.group(1)} row(s) that would duplicate the primary " + "key. Use method='upsert' to update existing rows, or " + "deduplicate the records" + ) + if "Scalar subquery produced more than one element" in msg: return ValidationError( "Found duplicated rows with the same primary key. " - f"Deduplicate the input batch and retry the {action} operation." + f"Deduplicate the records and retry the {action} operation." ) # `Bad int64 value: ` etc. — type-coercion failure on CAST(JSON_VALUE). diff --git a/datastore/infrastructure/engines/bigquery/lib.py b/datastore/infrastructure/engines/bigquery/lib.py index ca8b569..6307066 100644 --- a/datastore/infrastructure/engines/bigquery/lib.py +++ b/datastore/infrastructure/engines/bigquery/lib.py @@ -238,6 +238,95 @@ def insert_sql( ) +def insert_conflict_count_sql(table_ref: str, schema: dict) -> str: + """Render a read-only COUNT of primary-key conflicts for an INSERT. + + BigQuery does not enforce primary keys, so a plain INSERT happily + duplicates them. This query counts the incoming rows that *would* + create a duplicate: rows whose key already exists in the table, + plus the extra occurrences of any key repeated within the `@rows` + batch. Returns a single `n` column; `n > 0` means the INSERT must + be rejected. Caller guarantees `schema.primaryKey` is non-empty. + """ + pk = normalize_pk(schema) + if not pk: + raise ValueError( + "schema has no 'primaryKey'; insert conflict check requires one" + ) + by_name = {f["name"]: f for f in _user_fields(schema)} + missing = [c for c in pk if c not in by_name] + if missing: + raise ValueError( + f"primaryKey references undeclared column(s): {missing}" + ) + + extractors = ", ".join( + f"{_json_extract(by_name[c])} AS `{c}`" for c in pk + ) + pk_cols = ", ".join(f"`{c}`" for c in pk) + on_clause = " AND ".join(f"d.`{c}` = T.`{c}`" for c in pk) + return ( + f"WITH _incoming AS (" + f"SELECT {extractors} FROM UNNEST(JSON_QUERY_ARRAY(@rows)) AS r), " + f"_batch AS (" + f"SELECT {pk_cols}, COUNT(*) AS _n FROM _incoming GROUP BY {pk_cols}) " + f"SELECT (SELECT IFNULL(SUM(_n - 1), 0) FROM _batch) + " + f"(SELECT COUNT(*) FROM (SELECT DISTINCT {pk_cols} FROM _incoming) d " + f"JOIN {table_ref} T ON {on_clause}) AS n" + ) + + +# Sentinel the guarded-insert RAISE emits (followed by the conflict +# count) so `_translate_bigquery_error` can spot it in BigQuery's wrapped +# error text and rebuild the user-facing message Python-side. +PK_CONFLICT_SENTINEL = "DATASTORE_PK_CONFLICT" + + +def insert_guarded_sql( + table_ref: str, schema: dict, *, include_updated_at: bool = True +) -> str: + """Render one BigQuery script that rejects PK conflicts then INSERTs. + + Wraps the conflict probe and the INSERT in a single `BEGIN … END` + block so the `@rows` batch is serialised + uploaded once (one job, + not two). If any incoming row would duplicate the primary key — + against an existing table row or another row in the same batch — + `RAISE` aborts the script *before* the INSERT with the sentinel + the + conflict count; nothing is written and `_translate_bigquery_error` + maps it to a `ValidationError`. Requires a non-empty + `schema.primaryKey`; PK-less tables use the plain `insert_sql` (no + check at all). + + No-partial-write guarantee: the INSERT is the script's *only* write + and a single BigQuery DML statement, so it is atomic — all rows land + or none do. A conflict raises before it ever runs; a mid-INSERT + failure (e.g. a bad value) rolls the whole statement back. Either way + the table is left untouched. Keep this a single write statement to + preserve that guarantee. + """ + pk = normalize_pk(schema) + if not pk: + raise ValueError( + "schema has no 'primaryKey'; guarded insert requires one" + ) + count_query = insert_conflict_count_sql(table_ref, schema) + insert = insert_sql( + table_ref, schema, include_updated_at=include_updated_at + ) + return ( + "BEGIN " + "DECLARE _conflicts INT64 DEFAULT 0; " + f"SET _conflicts = ({count_query}); " + "IF _conflicts > 0 THEN " + f"RAISE USING MESSAGE = FORMAT('{PK_CONFLICT_SENTINEL} %d', " + "_conflicts); " + "END IF; " + # Single atomic DML write — runs only after the guard passes. + f"{insert}; " + "END" + ) + + def merge_sql( table_ref: str, schema: dict, *, include_updated_at: bool = True ) -> str: diff --git a/tests/engines/bigquery/test_tables.py b/tests/engines/bigquery/test_tables.py index fb517db..073f7dd 100644 --- a/tests/engines/bigquery/test_tables.py +++ b/tests/engines/bigquery/test_tables.py @@ -22,6 +22,8 @@ ) from datastore.infrastructure.engines.bigquery.backend import BigQueryBackend from datastore.infrastructure.engines.bigquery.lib import ( + insert_conflict_count_sql, + insert_guarded_sql, merge_sql, update_sql, ) @@ -230,6 +232,72 @@ def test_insert_records_issues_dml_insert_with_rows_param( assert json.loads(params["rows"]) == records +def test_insert_records_rejects_pk_conflict_and_writes_nothing( + mock_client: MagicMock, +) -> None: + """With a primary key declared, the guarded script RAISEs server-side + when a row would duplicate a key. BigQuery wraps our marker-bracketed + message; the backend recovers it as a `ValidationError`. The script + aborts before the INSERT, so nothing is written.""" + backend = _backend(mock_client) + schema = { + "fields": [{"name": "id", "type": "integer"}], + "primaryKey": ["id"], + } + mock_client.query.return_value.result.side_effect = RuntimeError( + "400 DATASTORE_PK_CONFLICT 2 at [1:120]; reason: invalidQuery" + ) + + with pytest.raises(ValidationError, match="duplicate the primary key"): + backend._insert_records("res-1", schema, [{"id": 1}, {"id": 1}]) + + # One guarded job — the conflict check and the INSERT share a single + # submission (so `@rows` is uploaded once). + assert mock_client.query.call_count == 1 + sql = mock_client.query.call_args[0][0] + assert sql.startswith("BEGIN ") + assert "RAISE USING MESSAGE" in sql + assert "INSERT INTO `proj-1.ds-1.res-1`" in sql + + +def test_insert_records_with_pk_runs_guarded_insert_in_one_job( + mock_client: MagicMock, +) -> None: + """A clean batch runs the conflict check and the INSERT as a single + `BEGIN … END` script — one `client.query`, so the `@rows` payload is + serialised + uploaded only once.""" + backend = _backend(mock_client) + schema = { + "fields": [{"name": "id", "type": "integer"}], + "primaryKey": ["id"], + } + + backend._insert_records("res-1", schema, [{"id": 1}, {"id": 2}]) + + assert mock_client.query.call_count == 1 + sql = mock_client.query.call_args[0][0] + assert sql.startswith("BEGIN ") + assert "SET _conflicts = (WITH _incoming AS (" in sql + assert "INSERT INTO `proj-1.ds-1.res-1` " in sql + # @rows rides on the one job only. + params = mock_client.query.call_args.kwargs["job_config"].query_parameters + assert [p.name for p in params] == ["rows"] + + +def test_insert_records_skips_guard_without_primary_key( + mock_client: MagicMock, +) -> None: + """No primary key → no guard; a single plain INSERT, today's + behaviour (duplicates allowed).""" + backend = _backend(mock_client) + schema = {"fields": [{"name": "id", "type": "integer"}]} + + backend._insert_records("res-1", schema, [{"id": 1}, {"id": 1}]) + + assert mock_client.query.call_count == 1 + assert mock_client.query.call_args[0][0].startswith("INSERT INTO ") + + # --- error wrapping -------------------------------------------------------- @@ -453,6 +521,82 @@ def test_merge_and_update_sql_reject_missing_primary_key() -> None: update_sql("`p.d.r`", schema) +def test_insert_conflict_count_sql_counts_batch_and_existing_dups() -> None: + """The conflict probe extracts the PK from `@rows`, counts in-batch + repeats (`SUM(_n - 1)`) plus rows whose key already exists in the + table (a self-JOIN), and returns one `n` column.""" + sql = insert_conflict_count_sql( + "`p.d.r`", + { + "fields": [ + {"name": "auction_id", "type": "integer"}, + {"name": "product_code", "type": "string"}, + {"name": "price", "type": "number"}, + ], + "primaryKey": ["auction_id", "product_code"], + }, + ) + # PK columns extracted from the JSON batch — non-PK `price` is not. + assert "CAST(JSON_VALUE(r, '$.auction_id') AS INT64) AS `auction_id`" in sql + assert "JSON_VALUE(r, '$.product_code') AS `product_code`" in sql + assert "`price`" not in sql + # In-batch repeats counted via GROUP BY + SUM(_n - 1). + assert "GROUP BY `auction_id`, `product_code`" in sql + assert "IFNULL(SUM(_n - 1), 0)" in sql + # Existing-row collisions via a composite-key JOIN against the table. + assert ( + "JOIN `p.d.r` T ON d.`auction_id` = T.`auction_id` " + "AND d.`product_code` = T.`product_code`" + ) in sql + assert sql.rstrip().endswith("AS n") + + +def test_insert_conflict_count_sql_rejects_missing_primary_key() -> None: + with pytest.raises(ValueError, match="primaryKey"): + insert_conflict_count_sql( + "`p.d.r`", {"fields": [{"name": "id", "type": "integer"}]} + ) + + +def test_insert_guarded_sql_wraps_check_and_insert_in_one_script() -> None: + """The guarded builder emits a single `BEGIN … END` script: declare a + counter, set it from the conflict probe, RAISE a marker-bracketed + message if non-zero, else fall through to the INSERT — all sharing the + one `@rows` parameter.""" + sql = insert_guarded_sql( + "`p.d.r`", + { + "fields": [ + {"name": "id", "type": "integer"}, + {"name": "label", "type": "string"}, + ], + "primaryKey": ["id"], + }, + ) + assert sql.startswith("BEGIN ") + assert sql.endswith(" END") + # Conflict probe feeds the script variable. + assert "SET _conflicts = (WITH _incoming AS (" in sql + # Guard fires only when the count is positive. + assert "IF _conflicts > 0 THEN" in sql + # RAISE emits the sentinel + the count; the message is rebuilt Python-side. + assert "RAISE USING MESSAGE = FORMAT('DATASTORE_PK_CONFLICT %d'" in sql + # The real INSERT follows the guard, reading the same @rows batch. + assert "INSERT INTO `p.d.r` " in sql + assert sql.count("JSON_QUERY_ARRAY(@rows)") == 2 # probe + insert + # No-partial-write invariant: exactly one (atomic) write statement, + # and the guard precedes it so a conflict never reaches the INSERT. + assert sql.count("INSERT INTO") == 1 + assert sql.index("END IF") < sql.index("INSERT INTO") + + +def test_insert_guarded_sql_rejects_missing_primary_key() -> None: + with pytest.raises(ValueError, match="primaryKey"): + insert_guarded_sql( + "`p.d.r`", {"fields": [{"name": "id", "type": "integer"}]} + ) + + # --- upsert() dispatch ---------------------------------------------------- @@ -490,9 +634,8 @@ def test_upsert_method_upsert_issues_merge_with_rows_param( def test_upsert_method_insert_issues_dml_insert( mock_client: MagicMock, ) -> None: - """`method='insert'` runs DML `INSERT INTO ... SELECT FROM UNNEST`, - not the streaming insert API — same path as `_insert_records` on - the create flow.""" + """`method='insert'` on a PK table runs the guarded `BEGIN … INSERT + … END` script (DML INSERT, not the streaming insert API) in one job.""" backend = _backend_with_schema( mock_client, {"fields": [{"name": "id", "type": "integer"}], "primaryKey": ["id"]}, @@ -501,11 +644,35 @@ def test_upsert_method_insert_issues_dml_insert( backend.upsert("res-1", [{"id": 1}], method="insert", include_total=False) mock_client.insert_rows_json.assert_not_called() + assert mock_client.query.call_count == 1 # check + INSERT share one job sql = mock_client.query.call_args[0][0] - assert sql.startswith("INSERT INTO `proj-1.ds-1.res-1` ") + assert sql.startswith("BEGIN ") + assert "INSERT INTO `proj-1.ds-1.res-1` " in sql assert "FROM UNNEST(JSON_QUERY_ARRAY(@rows)) AS r" in sql +def test_upsert_method_insert_rejects_pk_conflict( + mock_client: MagicMock, +) -> None: + """`method='insert'` against a key that already exists is rejected + with `ValidationError` (the caller should switch to `upsert`). The + guarded script RAISEs the sentinel + count, which BigQuery surfaces in + its error text.""" + backend = _backend_with_schema( + mock_client, + {"fields": [{"name": "id", "type": "integer"}], "primaryKey": ["id"]}, + ) + mock_client.query.return_value.result.side_effect = RuntimeError( + "400 DATASTORE_PK_CONFLICT 1 at [1:120]" + ) + + with pytest.raises(ValidationError, match="method='upsert'"): + backend.upsert("res-1", [{"id": 1}], method="insert", + include_total=False) + + assert mock_client.query.call_count == 1 # single guarded job + + def test_upsert_method_update_issues_dml_update( mock_client: MagicMock, ) -> None: