Skip to content

fix(graph): make concerns/reflection writes actually atomic - #69

Open
pradeepmouli wants to merge 2 commits into
intuit:mainfrom
pradeepmouli:fix/graphstore-transaction-helper
Open

fix(graph): make concerns/reflection writes actually atomic#69
pradeepmouli wants to merge 2 commits into
intuit:mainfrom
pradeepmouli:fix/graphstore-transaction-helper

Conversation

@pradeepmouli

Copy link
Copy Markdown
Contributor

Fixes #67 (Bug 2 of 2). Depends on #68 (this branch is stacked on it) -- the diff will shrink to just this PR's changes once #68 merges.

What

write_concerns and reflection's equivalent write_resolves_to issued BEGIN TRANSACTION/COMMIT through raw_query, believing that made their delete-then-recreate loop atomic. It never did:

  • KuzuBackend::raw_query opens a fresh connection per call, so transaction-control statements issued through it can never span more than the one statement they're attached to.
  • Neo4jBackend::raw_query no-ops the same statements for a different reason (Neo4j transactions are driver-level, not Cypher).

Both backends' raw_query already document this and deliberately no-op those statements rather than error -- but neither caller knew that, so every individual DETACH DELETE/CREATE in the loop auto-committed independently. A crash mid-loop today already means the old concerns/RESOLVES_TO edges are gone and only some of the new ones landed -- live data-loss exposure, not hypothetical.

Fix

  • GraphStore::transaction<T>(&self, f: impl FnOnce(&Connection) -> Result<T>) -> Result<T> -- opens one connection, issues a real BEGIN TRANSACTION, runs f against it, commits on Ok, rolls back on Err.
  • Two new GraphBackend trait methods, replace_concerns and replace_resolves_to, following the same "backend owns the transaction" design as the existing write_calls_service_edges. New Concern/ResolvesToEdge structs live in backend.rs (mirroring CallsServiceEdge) rather than reusing the analysis-pass types (ConcernMatch/ReflectionSite), keeping those out of the storage layer.
  • KuzuBackend implements both via the new store.transaction().
  • Neo4jBackend implements both as a single chained Cypher statement (delete + UNWIND-based recreate) -- matches write_calls_service_edges's existing precedent of "one auto-committed statement, no driver-level transaction needed" for that backend.
  • write_calls_service_edges itself also migrates onto store.transaction(), removing its hand-rolled BEGIN/loop/COMMIT/ROLLBACK boilerplate (same behavior, no longer duplicated).
  • The free write_concerns/write_resolves_to functions in concerns/mod.rs/reflection/mod.rs are deleted; their one call site each now builds the new backend-layer structs and calls the trait method directly.

Testing

New file tests/concerns_reflection_atomic_writes.rs, 7 tests:

  • replace_concerns_rolls_back_atomically_on_mid_batch_failure -- forces a mid-batch failure (two Concerns with the same id, colliding on Kùzu's primary key) and asserts a pre-existing concern survives untouched, proving the whole batch rolls back together. This fails against the pre-fix code (the old per-row let _ = swallowed the error entirely and always returned Ok, let alone rolled back).
  • Direct commit/rollback tests for GraphStore::transaction().
  • Basic correctness tests for both new trait methods (creates edges, replaces rather than accumulates, empty-input still clears old data).

All pass under both default and --features remote (neo4j). cargo fmt --all -- --check and cargo clippy --all-targets -- -D warnings both clean under both feature sets. Pre-commit hook's perf suite (write_lock_perf, groups_watch_perf, index_perf) all pass.

Note: the Neo4j-side Cypher hasn't been run against a live instance -- matches this codebase's existing precedent (the reference write_calls_service_edges Neo4j test is #[ignore]d for the same reason).

GitHub Copilot added 2 commits August 25, 2026 17:03
…ites

copy_edges_with_bad_record_retry took a borrowed Connection and reused it
across every retry attempt and the UNWIND fallback. None of these bulk
loads are wrapped in an explicit transaction, so sharing the connection
bought no atomicity -- it only created exposure to whatever internal state
a caught COPY failure leaves behind. Observed in production: a Symbol-table
COPY's bad-PK-drop-and-retry cycle left the connection wedged such that the
very next COPY on it (a different table, CALLS) failed immediately with
Kuzu's internal "Invalid transaction type to rollback." and fell back to
the slower per-row UNWIND path.

Fix: copy_edges_with_bad_record_retry now takes &GraphStore instead of
&Connection and asks for a fresh connection on every retry-loop iteration
and before the UNWIND fallback -- no "is this a retry" bookkeeping needed,
since GraphStore::connection() already mints a fresh, cheap Connection on
every call. The inline Symbol-node COPY-with-retry block in
import_scip_index (a near-duplicate of the same pattern for nodes instead
of edges) gets the same treatment.

Threading &GraphStore down to the one caller of this helper that didn't
already have it (resolve_with_map) also let resolve_inherits drop its now-
entirely-unused &Connection parameter.

Self-healing today via the UNWIND fallback (byte-for-byte identical output,
per store_bench::test_parquet_quality), so no data-loss exposure -- but
real and reproducible.

Also silences one pre-existing, unrelated clippy::chunks_exact_to_as_chunks
lint in embed/mod.rs (newer clippy than this branch's baseline; the
workspace-wide pre-commit hook blocks on it otherwise) -- no behavior
change, matches clippy's own suggested suppression.
Fixes intuit#67 (Bug 2 of 2) -- a live data-loss bug, not just
tech debt. write_concerns and reflection's equivalent write_resolves_to
issued BEGIN TRANSACTION/COMMIT through raw_query, believing that made the
delete-then-recreate loop atomic. It never did: KuzuBackend::raw_query
opens a fresh connection per call (so transaction-control statements can't
span more than the one statement they're attached to), and Neo4jBackend's
raw_query no-ops them too (Neo4j transactions are driver-level, not
Cypher). Both backends' raw_query already documented this and no-op those
statements rather than error -- but neither caller knew that, so every
individual DETACH DELETE/CREATE in the loop auto-committed independently.
A crash mid-loop today already means the old concerns/RESOLVES_TO edges
are gone and only some of the new ones landed.

Fix: add GraphStore::transaction<T>(&self, f) -- opens one connection,
issues a real BEGIN TRANSACTION, runs f against it, commits on Ok, rolls
back on Err. Add two new GraphBackend trait methods, replace_concerns and
replace_resolves_to (same "backend owns the transaction" design as the
existing write_calls_service_edges), with Concern/ResolvesToEdge structs
in backend.rs mirroring CallsServiceEdge -- keeps analysis-pass types
(ConcernMatch/ReflectionSite) out of the storage layer. KuzuBackend
implements both via the new store.transaction(). Neo4jBackend implements
both as a single chained Cypher statement (delete + UNWIND-recreate),
matching write_calls_service_edges's existing "one auto-committed
statement, no driver-level transaction needed" precedent for that backend.

write_calls_service_edges itself also migrates onto store.transaction(),
removing its hand-rolled BEGIN/loop/COMMIT/ROLLBACK boilerplate (same
behavior, no longer duplicated).

concerns/mod.rs and reflection/mod.rs's free write_concerns/
write_resolves_to functions are deleted; their one call site each now
converts to the new backend-layer structs and calls the trait method
directly.

New test: replace_concerns_rolls_back_atomically_on_mid_batch_failure
forces a mid-batch failure (two Concerns with the same id, colliding on
Kùzu's primary key) and asserts a pre-existing concern survives untouched
-- proving the whole batch rolls back together. This fails against the
pre-fix code (the old per-row `let _ =` swallowed the error entirely and
always returned Ok, let alone rolled back). Plus direct commit/rollback
tests for GraphStore::transaction() and basic correctness tests for both
new trait methods (7 tests total, new file
tests/concerns_reflection_atomic_writes.rs).

Builds and passes fmt/clippy/tests under both default and --features
remote (neo4j) -- the Neo4j-side Cypher hasn't been run against a live
instance (matches this codebase's existing precedent: the reference
write_calls_service_edges Neo4j test is #[ignore]'d for the same reason).
pradeepmouli pushed a commit to pradeepmouli/infigraph that referenced this pull request Aug 28, 2026
Port of upstream intuit#69 (tracks #119).

write_concerns and reflection's equivalent write_resolves_to issued BEGIN
TRANSACTION/COMMIT through raw_query, believing that made their
delete-then-recreate loop atomic. It never did: KuzuBackend::raw_query
opens a fresh connection per call (so transaction-control statements can't
span more than the one statement they're attached to), and Neo4jBackend's
raw_query no-ops them too (Neo4j transactions are driver-level, not
Cypher). Both backends' raw_query already documented this and no-op those
statements rather than error -- but neither caller knew that, so every
individual DETACH DELETE/CREATE in the loop auto-committed independently.
A crash mid-loop today already means the old concerns/RESOLVES_TO edges
are gone and only some of the new ones landed.

Fix: add GraphStore::transaction<T>(&self, f) -- opens one connection,
issues a real BEGIN TRANSACTION, runs f against it, commits on Ok, rolls
back on Err. Add two new GraphBackend trait methods, replace_concerns and
replace_resolves_to (same "backend owns the transaction" design as the
existing write_calls_service_edges), with Concern/ResolvesToEdge structs
in backend.rs mirroring CallsServiceEdge -- keeps analysis-pass types
(ConcernMatch/ReflectionSite) out of the storage layer.

Three backend implementations (this fork has one more than upstream):
- KuzuBackend: via the new store.transaction().
- Neo4jBackend: a single chained Cypher statement (delete +
  UNWIND-recreate), matching write_calls_service_edges's existing "one
  auto-committed statement, no driver-level transaction needed" precedent.
- DaemonKuzuBackend (fork-only): two new WriteRequest variants,
  ReplaceConcerns/ReplaceResolvesTo, riding inline in the request envelope
  (Concern/ResolvesToEdge are small serde-serializable payloads, same
  convention as UpsertDependencies/StoreConfigBindings -- no bespoke Arrow
  sibling file needed, unlike WriteCallsServiceEdges's genuinely-tabular
  edge lists). Daemon-side dispatch added to serve_one_request, mirroring
  UpsertDependencies's handling exactly.

write_calls_service_edges itself also migrates onto store.transaction(),
removing its hand-rolled BEGIN/loop/COMMIT/ROLLBACK boilerplate.

concerns/mod.rs and reflection/mod.rs's free write_concerns/
write_resolves_to functions are deleted; their one call site each now
converts to the new backend-layer structs and calls the trait method
directly.

New test (ported as-is from upstream): concerns_reflection_atomic_writes.rs,
7 tests including replace_concerns_rolls_back_atomically_on_mid_batch_failure,
which forces a mid-batch failure (two Concerns with the same id, colliding
on Kùzu's primary key) and asserts a pre-existing concern survives
untouched -- proving the whole batch rolls back together. This fails
against the pre-fix code (the old per-row `let _ =` swallowed the error
entirely and always returned Ok, let alone rolled back).

Verified: cargo fmt --all -- --check and cargo clippy --all-targets
-- -D warnings clean under both default and --features remote. Full
workspace (cargo check --workspace --all-targets) compiles clean --
DaemonKuzuBackend is the one GraphBackend implementer with no upstream
equivalent, and needed the daemon_protocol.rs wiring above to satisfy the
trait. Targeted test suites (concerns_reflection_atomic_writes,
calls_service_edges, concerns::, reflection::, daemon_protocol::) all
pass. Confirmed via git stash that the 16 daemon::drain::/tests::init_*
failures seen under `cargo test --lib` pre-exist this change (same
failures reproduce identically on the prior commit) -- matches this
codebase's documented embedded-DB resource-contention flakiness on this
dev machine, not a regression.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Connection/transaction lifecycle bugs in KuzuBackend: wedged COPY retries + silently non-atomic writes

1 participant