Reject primary-key conflicts on insert (BigQuery)

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 "<sentinel> <count>" 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: <v>` etc. — type-coercion failure on CAST(JSON_VALUE).

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 (<inner>)` 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 <column>` to a plain single-table SELECT for stable
+    LIMIT/OFFSET paging. No-op (returns `sql` unchanged) for anything where
+    `<column>` 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