Eliminate N+1 subgraph hydration queries with batched entity/relation fetches#4
Conversation
Agent-Logs-Url: https://github.com/vndee/memex/sessions/1ec419ab-55f2-4063-8471-c0e8891d649c Co-authored-by: vndee <28271488+vndee@users.noreply.github.com>
Agent-Logs-Url: https://github.com/vndee/memex/sessions/1ec419ab-55f2-4063-8471-c0e8891d649c Co-authored-by: vndee <28271488+vndee@users.noreply.github.com>
Agent-Logs-Url: https://github.com/vndee/memex/sessions/1ec419ab-55f2-4063-8471-c0e8891d649c Co-authored-by: vndee <28271488+vndee@users.noreply.github.com>
Agent-Logs-Url: https://github.com/vndee/memex/sessions/1ec419ab-55f2-4063-8471-c0e8891d649c Co-authored-by: vndee <28271488+vndee@users.noreply.github.com>
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
|
Important Review skippedBot user detected. To trigger a single review, invoke the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR implements batch fetching for entities and relations, replacing individual per-item lookups with bulk retrieval methods. It adds Changes
Possibly related issues
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/storage/storage_test.go (1)
296-354: Good test coverage for batch entity fetching.The tests correctly validate KB scoping (entities from other KBs excluded), partial ID matches, and large batch handling (905 IDs exercises the chunking logic).
Consider adding an edge case test for empty
idsslice to verify the early return behavior:💡 Optional test for empty input
func TestGetEntitiesByIDs_EmptyInput(t *testing.T) { store := newTestStore(t) createTestKB(t, store, "kb1") entities, err := store.GetEntitiesByIDs(context.Background(), "kb1", []string{}) if err != nil { t.Fatal("get entities by ids:", err) } if len(entities) != 0 { t.Fatalf("got %d entities, want 0", len(entities)) } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/storage/storage_test.go` around lines 296 - 354, Add a unit test to cover the empty-ids edge case: create a test named TestGetEntitiesByIDs_EmptyInput that uses newTestStore and createTestKB, calls store.GetEntitiesByIDs(ctx, "kb1", []string{}) and asserts no error and that the returned map length is 0; this verifies the GetEntitiesByIDs early-return behavior and prevents regressions in the chunking/early-exit logic exercised by TestGetEntitiesByIDs and TestGetEntitiesByIDs_LargeBatch.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@internal/storage/storage_test.go`:
- Around line 296-354: Add a unit test to cover the empty-ids edge case: create
a test named TestGetEntitiesByIDs_EmptyInput that uses newTestStore and
createTestKB, calls store.GetEntitiesByIDs(ctx, "kb1", []string{}) and asserts
no error and that the returned map length is 0; this verifies the
GetEntitiesByIDs early-return behavior and prevents regressions in the
chunking/early-exit logic exercised by TestGetEntitiesByIDs and
TestGetEntitiesByIDs_LargeBatch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c5ad4916-3f28-429e-bb2d-bf98517ad45d
📒 Files selected for processing (8)
internal/search/graph.gointernal/search/graph_test.gointernal/server/helpers.gointernal/server/helpers_test.gointernal/storage/entities.gointernal/storage/relations.gointernal/storage/sqlite.gointernal/storage/storage_test.go
There was a problem hiding this comment.
Pull request overview
This PR removes N+1 lookup behavior during graph/subgraph hydration by adding batch-fetch primitives to the storage layer and updating hydration paths to populate results from in-memory maps.
Changes:
- Added
GetEntitiesByIDs/GetRelationsByIDstostorage.Store, with SQLite implementations using chunkedIN (...)queries. - Updated
HydrateSubgraphandhydrateEntityResultsto hydrate via batched lookups rather than per-item store calls. - Added tests for batch semantics (scoping, missing IDs, large batches) and for ensuring hydration uses batch methods.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| internal/storage/storage_test.go | Adds tests for entity/relation batch lookup semantics and large-batch chunking. |
| internal/storage/sqlite.go | Extends storage.Store with batch lookup methods. |
| internal/storage/entities.go | Implements GetEntitiesByIDs with chunked IN queries returning an ID-keyed map. |
| internal/storage/relations.go | Implements GetRelationsByIDs with chunked IN queries returning an ID-keyed map. |
| internal/server/helpers.go | Updates HydrateSubgraph to batch-fetch nodes/edges once and hydrate from maps. |
| internal/server/helpers_test.go | Adds a regression test asserting batch methods are used by HydrateSubgraph. |
| internal/search/graph.go | Updates hydrateEntityResults to batch-fetch entity metadata and skip missing IDs via map lookup. |
| internal/search/graph_test.go | Adds a regression test asserting batch method usage in hydrateEntityResults. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| } | ||
| entities[e.ID] = e | ||
| } | ||
| if err := rows.Err(); err != nil { |
There was a problem hiding this comment.
rows.Err() can return a non-nil error after iteration; in that branch the code returns without closing rows, which can leak resources. Close rows before returning (or defer rows.Close() immediately after QueryContext and then drop the explicit Close() calls).
| if err := rows.Err(); err != nil { | |
| if err := rows.Err(); err != nil { | |
| if closeErr := rows.Close(); closeErr != nil { | |
| return nil, fmt.Errorf("iterate entities by ids: %w (close rows: %v)", err, closeErr) | |
| } |
| rows, err := s.db.QueryContext(ctx, query, args...) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("query relations by ids: %w", err) | ||
| } | ||
| for rows.Next() { | ||
| r, err := scanRelation(rows) | ||
| if err != nil { | ||
| rows.Close() | ||
| return nil, err | ||
| } | ||
| rels[r.ID] = r | ||
| } | ||
| if err := rows.Err(); err != nil { | ||
| return nil, fmt.Errorf("iterate relations by ids: %w", err) | ||
| } | ||
| if err := rows.Close(); err != nil { | ||
| return nil, fmt.Errorf("close relations by ids rows: %w", err) | ||
| } |
There was a problem hiding this comment.
If rows.Err() is non-nil, the function returns without closing rows, which can leak the underlying statement/connection resources. Ensure rows is closed on all paths (e.g., defer rows.Close() right after QueryContext, or explicitly rows.Close() before returning on the rows.Err() branch).
| rows, err := s.db.QueryContext(ctx, query, args...) | |
| if err != nil { | |
| return nil, fmt.Errorf("query relations by ids: %w", err) | |
| } | |
| for rows.Next() { | |
| r, err := scanRelation(rows) | |
| if err != nil { | |
| rows.Close() | |
| return nil, err | |
| } | |
| rels[r.ID] = r | |
| } | |
| if err := rows.Err(); err != nil { | |
| return nil, fmt.Errorf("iterate relations by ids: %w", err) | |
| } | |
| if err := rows.Close(); err != nil { | |
| return nil, fmt.Errorf("close relations by ids rows: %w", err) | |
| } | |
| if err := func() error { | |
| rows, err := s.db.QueryContext(ctx, query, args...) | |
| if err != nil { | |
| return fmt.Errorf("query relations by ids: %w", err) | |
| } | |
| defer rows.Close() | |
| for rows.Next() { | |
| r, err := scanRelation(rows) | |
| if err != nil { | |
| return err | |
| } | |
| rels[r.ID] = r | |
| } | |
| if err := rows.Err(); err != nil { | |
| return fmt.Errorf("iterate relations by ids: %w", err) | |
| } | |
| return nil | |
| }(); err != nil { | |
| return nil, err | |
| } |
| entitiesByID, err := store.GetEntitiesByIDs(ctx, kbID, ids) | ||
| if err != nil { | ||
| slog.Warn("graph entity batch lookup failed", "count", len(ids), "error", err) | ||
| return nil, fmt.Errorf("get entities by ids: %w", err) | ||
| } |
There was a problem hiding this comment.
hydrateEntityResults now batch-fetches using the full ids slice even though it stops building results once limit is reached. For large traversals with small limits this can cause unnecessary DB work and chunked queries. Consider fetching only as many IDs as needed (e.g., progressively batch/paginate until limit results are filled, or cap the initial batch to a reasonable bound like min(len(ids), limit) with a fallback to fetch more if many IDs are missing).
|
@copilot check the reviews and comments |
Agent-Logs-Url: https://github.com/vndee/memex/sessions/c7178a26-748e-4cfb-a92a-ba8c423be949 Co-authored-by: vndee <28271488+vndee@users.noreply.github.com>
Agent-Logs-Url: https://github.com/vndee/memex/sessions/c7178a26-748e-4cfb-a92a-ba8c423be949 Co-authored-by: vndee <28271488+vndee@users.noreply.github.com>
Addressed the review items in commit |
HydrateSubgraphand graph-result hydration were doing per-node/per-edge lookups (GetEntity/GetRelation) in loops, creating N+1 query behavior on larger traversals. This change switches both paths to batched ID fetches and in-memory map hydration.Storage API: add batch lookup primitives
storage.Store:GetEntitiesByIDs(ctx, kbID, ids) (map[string]*domain.Entity, error)GetRelationsByIDs(ctx, kbID, ids) (map[string]*domain.Relation, error)SQLite implementation: single-query ID batching (chunked)
SQLiteStorewithWHERE kb_id = ? AND id IN (...).rowsare closed on all return paths.Hydration paths: remove per-item store calls
internal/server/helpers.go(HydrateSubgraph):internal/search/graph.go(hydrateEntityResults):Focused coverage for new behavior
Summary by CodeRabbit