diff --git a/datastore/infrastructure/engines/bigquery/backend.py b/datastore/infrastructure/engines/bigquery/backend.py index 130bb1e..68217d0 100644 --- a/datastore/infrastructure/engines/bigquery/backend.py +++ b/datastore/infrastructure/engines/bigquery/backend.py @@ -37,12 +37,15 @@ WriteResult, ) from datastore.infrastructure.engines.bigquery.lib import ( + PK_CONFLICT_SENTINEL, SYSTEM_COLUMN_NAMES, alter_clauses, column_defs, + default_order_by, delete_sql, drop_columns_sql, format_select_column, + insert_guarded_sql, insert_sql, merge_sql, normalize_pk, @@ -301,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, ) @@ -431,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. @@ -607,8 +629,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 @@ -648,6 +670,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 +686,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: @@ -697,8 +721,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 @@ -707,18 +732,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: @@ -1169,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 99d78b9..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: @@ -441,8 +530,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 @@ -520,6 +610,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 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: