feat(core): Add statement timeout parameter.#3544
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a configurable statement timeout for database connections. This enhancement is designed to protect the database from potentially expensive, long-running operations, particularly in anticipation of upcoming substring search functionality in List RPCs. The implementation adds the parameter to the database configuration struct and ensures it is properly handled during connection pool initialization. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. A query runs, it takes too long, The database feels something's wrong. We set a limit, time to cease, And bring the system back to peace. Footnotes
|
📝 WalkthroughWalkthroughAdds an optional DB statement_timeout config field (StatementTimeoutSeconds), emits it in logs, applies it to pgx runtime parameters as "Ns" when >0, updates three YAML examples with a commented setting, and extends tests to validate runtime parameter behavior for unset and 30s cases. ChangesDatabase Statement Timeout Configuration
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 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 docstrings
🧪 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a statement_timeout configuration option for database connections, updating the configuration files, the database configuration struct, logging, connection initialization, and corresponding unit tests. The reviewer feedback suggests changing the parameter type from a string to an integer representing seconds (statement_timeout_seconds) to maintain consistency with other timeout settings and prevent PostgreSQL from misinterpreting bare integers as milliseconds. This change requires updating the configuration files, struct fields, logger, connection builder, and tests accordingly.
|
Actionable comments posted: 0 |
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@service/pkg/db/db.go`:
- Around line 255-257: The code currently sets
parsed.ConnConfig.RuntimeParams["statement_timeout"] using
c.StatementTimeoutSeconds which applies a pool-wide timeout and will also affect
migrations/admin work (see migrationInit which casts c.Pgx to *pgxpool.Pool and
uses stdlib.OpenDBFromPool). Remove the RuntimeParams assignment and instead
apply the timeout only where needed: use context.WithTimeout around RPC
handlers/queries that require a per-request timeout (or run "SET LOCAL
statement_timeout = 'Ns'" at the start of the specific transaction) in the
List-query code paths; keep references to c.StatementTimeoutSeconds,
parsed.ConnConfig.RuntimeParams, migrationInit, c.Pgx and stdlib.OpenDBFromPool
to locate the related code.
🪄 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: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 3e6ea1cf-a481-40f1-a240-49b167802ede
📒 Files selected for processing (5)
opentdf-core-mode.yamlopentdf-dev.yamlopentdf-example.yamlservice/pkg/db/db.goservice/pkg/db/db_test.go
| if c.StatementTimeoutSeconds > 0 { | ||
| parsed.ConnConfig.RuntimeParams["statement_timeout"] = fmt.Sprintf("%ds", c.StatementTimeoutSeconds) | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# How are migrations run? Do they reuse the pool/SQLDB built from buildConfig?
rg -nP --type=go -C4 '\b(RunMigrations|runMigrations|Migrate|migrat)' service/pkg/db
# Find the migration runner entrypoint and what conn/pool it uses
ast-grep --pattern 'func ($_ $_) RunMigrations($$$) { $$$ }'Repository: opentdf/platform
Length of output: 10075
statement_timeout is pool-wide and will also affect migrations/admin work
The statement_timeout startup runtime param applies to every statement executed via this client’s shared pgxpool.Pool:
migrationInitcastsc.Pgxto*pgxpool.Pooland creates the migrations DB handle withstdlib.OpenDBFromPool(pool), so Goose runs migrations on connections from the same pool that carriesRuntimeParams["statement_timeout"].- This means a pool-wide timeout could fail slow startup migrations/backfills, not just the intended List RPC queries.
"%ds"formatting is correct for PostgreSQLstatement_timeout(unit-suffixed seconds like"30s").
if c.StatementTimeoutSeconds > 0 {
parsed.ConnConfig.RuntimeParams["statement_timeout"] = fmt.Sprintf("%ds", c.StatementTimeoutSeconds)
}🤖 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 `@service/pkg/db/db.go` around lines 255 - 257, The code currently sets
parsed.ConnConfig.RuntimeParams["statement_timeout"] using
c.StatementTimeoutSeconds which applies a pool-wide timeout and will also affect
migrations/admin work (see migrationInit which casts c.Pgx to *pgxpool.Pool and
uses stdlib.OpenDBFromPool). Remove the RuntimeParams assignment and instead
apply the timeout only where needed: use context.WithTimeout around RPC
handlers/queries that require a per-request timeout (or run "SET LOCAL
statement_timeout = 'Ns'" at the start of the specific transaction) in the
List-query code paths; keep references to c.StatementTimeoutSeconds,
parsed.ConnConfig.RuntimeParams, migrationInit, c.Pgx and stdlib.OpenDBFromPool
to locate the related code.
There was a problem hiding this comment.
This is a good point. Should we override this client setting while running migrations? @c-r33d
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
There was a problem hiding this comment.
Yes, I actually was rethinking how I was going to do this because of the coderabbit comment. The problem is that we run migrations during startup and that uses the pool that where we have already specified. I will check if we can override this with a command-level SET.
|
2 participants
Note
Adding a query timeout will be useful for protecting the database from
long running searches when substring search is added to the List RPCs.
1.) Add the ability to specify a statement timeout for our queries.
2.) By default we do not specify a timeout, but allow one to be set if needed.
Summary by CodeRabbit