feat(sdk/go): add context.Context to Store interface methods#476
Merged
Conversation
The Store interface methods accepted no context, forcing the PostgreSQL adapter to use context.Background() for every pgx call. A stalled connection could block indefinitely while holding the store mutex with no way for callers to cancel. Add context.Context as the first parameter to Unit, All, Insert, Update, Delete, Query, and Stats. Close is unchanged — it does no cancellable I/O. All three implementations (memory, SQLite, PostgreSQL), the Client, the conformance suite, and tests are updated in lockstep. The SQLite store now uses context-aware database/sql methods (QueryContext, ExecContext, BeginTx, PrepareContext). Fixes #468
postgres.New used context.Background() for pool creation, server ping, and schema migration — the same indefinite-blocking problem the Store interface change addressed. Thread caller-provided context through New, ensureSchema, and stampWriter so callers can set timeouts on connection setup.
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
WalkthroughGo store methods now accept ChangesGo store context propagation
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the Go SDK’s cq.Store SPI to accept context.Context as the first parameter for all cancellable operations, and propagates that change through the SDK client, store implementations, conformance suite, and tests. This addresses the prior risk of indefinite blocking (notably in the PostgreSQL adapter) by allowing callers to enforce timeouts/cancellation.
Changes:
- Add
context.Contextas the first argument to allStoreinterface methods exceptClose. - Update SQLite and PostgreSQL stores to use context-aware DB operations (including
database/sql*Contextmethods andpgxcalls). - Update the Go SDK client, conformance suite, and store tests to pass contexts through.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| sdk/go/store.go | Adds context.Context to all Store methods (except Close) to enable cancellation/deadlines. |
| sdk/go/client.go | Threads caller context through client → store calls for local operations. |
| sdk/go/store_memory.go | Updates in-memory store to satisfy new interface (context ignored). |
| sdk/go/store_memory_test.go | Updates in-memory store tests to pass context.Background(). |
| sdk/go/store_sqlite.go | Updates SQLite store methods and helpers to use database/sql context-aware APIs. |
| sdk/go/store_sqlite_test.go | Updates SQLite store tests to pass contexts. |
| sdk/go/stores/postgres/postgres.go | Updates PostgreSQL store API, propagates ctx into pool creation/ping/migrations and all queries. |
| sdk/go/stores/postgres/postgres_test.go | Updates PostgreSQL constructor tests to pass context. |
| sdk/go/storetest/storetest.go | Updates the conformance suite to pass contexts to all store calls. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
sdk/go/client.go (1)
185-192: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winApply
operationContextinDrainableCount.Unlike the other client entry points, this method never normalises the caller context before Line 192. That means
c.timeoutis ignored forStore.All(ctx), and a nilctxstill panics at Line 187. Now thatAllcan perform cancellable backing-store I/O, this path can still block indefinitely.Suggested fix
func (c *Client) DrainableCount(ctx context.Context) (int, error) { - select { - case <-ctx.Done(): - return 0, ctx.Err() - default: - } + ctx, cancel := c.operationContext(ctx) + defer cancel() + + if err := ctx.Err(); err != nil { + return 0, err + } units, err := c.store.All(ctx) if err != nil { return 0, fmt.Errorf("reading local units: %w", err) }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sdk/go/client.go` around lines 185 - 192, The DrainableCount method is using the caller context directly, so it can ignore c.timeout and still panic on a nil ctx before calling Store.All. Update Client.DrainableCount to normalize the input with operationContext, matching the other client entry points, and then pass that derived context into c.store.All so timeout and cancellation are consistently applied.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sdk/go/store_memory.go`:
- Line 31: The in-memory store methods on inMemoryStore still ignore ctx, so
cancellation is not honored consistently with the SQLite and PostgreSQL
backends. Update the affected methods, especially All, Query, and Stats, to
perform a pre-flight context check before blocking on s.mu and add loop-level
ctx checks while iterating the store; apply the same pattern to the other
inMemoryStore methods referenced in the diff so Client.Query and DrainableCount
stop behaving differently by backend.
In `@sdk/go/stores/postgres/postgres_test.go`:
- Around line 41-43: Add a regression case in TestNew that calls postgres.New
with an already-cancelled context instead of only context.Background(), so the
test verifies pool creation, ping, and schema setup all respect cancellation.
Use the existing postgres.New test table or add a small separate subtest keyed
off tc.connString, and assert the returned error reflects context cancellation.
---
Outside diff comments:
In `@sdk/go/client.go`:
- Around line 185-192: The DrainableCount method is using the caller context
directly, so it can ignore c.timeout and still panic on a nil ctx before calling
Store.All. Update Client.DrainableCount to normalize the input with
operationContext, matching the other client entry points, and then pass that
derived context into c.store.All so timeout and cancellation are consistently
applied.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: a4f8de4b-2b67-494e-a2da-e86073dc27ad
📒 Files selected for processing (9)
sdk/go/client.gosdk/go/store.gosdk/go/store_memory.gosdk/go/store_memory_test.gosdk/go/store_sqlite.gosdk/go/store_sqlite_test.gosdk/go/stores/postgres/postgres.gosdk/go/stores/postgres/postgres_test.gosdk/go/storetest/storetest.go
The deferred rollback in Insert and Update used the caller's ctx. When that ctx is cancelled or times out mid-operation, the rollback cannot run, and pgx churns the pooled connection instead of releasing it cleanly. Extract a rollback helper that discards the transaction with a short background-derived context, independent of caller cancellation but still bounded so it cannot block indefinitely.
Guard against future drift back to context.Background() inside New by asserting that an already-cancelled context aborts connection setup. Asserts errors.Is(err, context.Canceled) so it holds regardless of which setup stage observes the cancellation.
The in-memory store discarded ctx, so a cancelled call could still acquire the mutex and walk the store while the SQLite and Postgres backends would return early. Add a pre-flight select/default check to each context-taking method so cancellation behaves consistently across backends. The store does no blocking I/O, so per-iteration checks are unnecessary.
operationContext skipped the client timeout entirely when the caller's context already carried a deadline, making c.timeout a default-when-unset rather than a ceiling. Always derive from the caller's context with WithTimeout, which expires at the sooner of the two — so c.timeout caps every operation without ever extending a shorter caller deadline. Also drop the nil-context guard: passing a nil context is a caller bug and panicking is the idiomatic Go outcome.
Convert the pre-flight cancellation guards across the client entry points from ctx.Err() to a uniform select/default form, and route DrainableCount through operationContext like every other entry point so it honors the client timeout instead of using the raw caller context. The mid-loop check in Drain stays on ctx.Err() because a select there would need a labeled break.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
sdk/go/client.go (1)
169-177: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUse a non-cancellable cleanup context for the post-push delete.
After the remote
proposesucceeds, Line 176 still uses the caller context for local cleanup. If the deadline fires in that gap, the unit is already persisted remotely but remains in the local store, so the next drain can re-push the same item. The delete needs a short-lived background-derived context so cancellation cannot strand a successfully drained unit locally.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sdk/go/client.go` around lines 169 - 177, The post-push cleanup in the drain flow still uses the caller’s context for the local delete, which can leave a successfully proposed unit stranded if the request is canceled before deletion. Update the cleanup path in c.remote.propose handling within c.store.Delete to use a short-lived background-derived context for the local delete, so cancellation or deadline expiry cannot block the cleanup after a successful remote push.
♻️ Duplicate comments (1)
sdk/go/store_memory.go (1)
31-50: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe in-memory scans still ignore cancellation after the pre-flight check.
These new
selectblocks only catch contexts that are already cancelled. OnceAll,Query, orStatsenter theirs.orderloops, they never re-checkctx, so a cancelledDrainableCount,Query, orStatuscall can still walk the whole store while holdings.mu. Add loop-time cancellation checks before cloning/ranking each item so the in-memory backend matches the cancellability promised by the other stores.Also applies to: 126-166, 170-220
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sdk/go/store_memory.go` around lines 31 - 50, The in-memory store methods All, Query, and Stats only check ctx before entering their loops, so cancellation is ignored once iteration starts. Update the implementations in inMemoryStore to re-check ctx during the s.order traversal, before cloning/ranking each item and while holding s.mu, and return ctx.Err() immediately when cancelled. Apply the same pattern across All, Query, and Stats so DrainableCount, Query, and Status stop promptly instead of scanning the whole store.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sdk/go/client.go`:
- Around line 511-515: The public client methods are now assuming a non-nil
context, but operationContext still can panic if ctx is nil. Add the nil guard
at the exported method boundary or restore the context.Background() fallback
before calling operationContext, and keep the timeout behavior unchanged for
non-nil callers.
In `@sdk/go/stores/postgres/postgres_test.go`:
- Around line 50-59: The postgres.New examples are stale and still show the old
context-free API. Update the usages in postgres.New-related docs and package
comments to the new postgres.New(ctx, connString) signature, matching the
cancellation-aware form exercised by TestNewRespectsCancelledContext and the New
function in postgres.go. Keep the examples consistent across README.md files and
the postgres package comment so they demonstrate passing a context as the first
argument.
---
Outside diff comments:
In `@sdk/go/client.go`:
- Around line 169-177: The post-push cleanup in the drain flow still uses the
caller’s context for the local delete, which can leave a successfully proposed
unit stranded if the request is canceled before deletion. Update the cleanup
path in c.remote.propose handling within c.store.Delete to use a short-lived
background-derived context for the local delete, so cancellation or deadline
expiry cannot block the cleanup after a successful remote push.
---
Duplicate comments:
In `@sdk/go/store_memory.go`:
- Around line 31-50: The in-memory store methods All, Query, and Stats only
check ctx before entering their loops, so cancellation is ignored once iteration
starts. Update the implementations in inMemoryStore to re-check ctx during the
s.order traversal, before cloning/ranking each item and while holding s.mu, and
return ctx.Err() immediately when cancelled. Apply the same pattern across All,
Query, and Stats so DrainableCount, Query, and Status stop promptly instead of
scanning the whole store.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: bfcbe834-1430-45f5-bafe-9cbce56d219b
📒 Files selected for processing (5)
sdk/go/client.gosdk/go/options.gosdk/go/store_memory.gosdk/go/stores/postgres/postgres.gosdk/go/stores/postgres/postgres_test.go
The README usage snippets still showed the context-free postgres.New(connString) form. Update both to the new postgres.New(ctx, connString) signature so the docs match the API.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
context.Contextas the first parameter to all Store interface methods(Unit, All, Insert, Update, Delete, Query, Stats) so callers can set
timeouts and cancellation. Close is unchanged — it does no cancellable I/O.
Client, the conformance suite, and tests.
ExecContext, BeginTx, PrepareContext).
indefinite-blocking problem via context.Background() for pool creation,
ping, and schema migration.
Fixes #468
Test plan
make test-sdk-go— all tests passmake test-cli— all tests passmake lint-sdk-go— 0 issuesmake lint-cli— 0 issuesSummary by CodeRabbit
contextfor cancellation and deadlines, including prompt early cancellation.