feat: granularised db connection usage in superposition crate

[ayushjain17]

Problem

DB Connection from pool is being held up for the entirety of the request processing duration

Solution

Granularise getting and usage of DB Connection from the pool limited to the query execution only

Summary by CodeRabbit

• Refactor
  • Updated database connection handling across multiple handler modules to use a centralized pool-based query execution pattern.
  • Simplified database operation abstraction with unified transaction handling across the service.

[semanticdiff-com]

Review changes with  

Changed Files

File	Status
crates/experimentation_platform/src/api/experiments/handlers.rs	81% smaller
crates/superposition/src/resolve/handlers.rs	73% smaller
crates/service_utils/src/middlewares/workspace_context.rs	52% smaller
crates/service_utils/src/helpers.rs	39% smaller
crates/superposition/src/organisation/handlers.rs	15% smaller
crates/service_utils/src/middlewares/auth_n/helpers.rs	1% smaller
crates/service_utils/src/db.rs	0% smaller
crates/superposition/src/webhooks/handlers.rs	0% smaller
crates/superposition/src/webhooks/helper.rs	0% smaller
crates/superposition/src/workspace/handlers.rs	0% smaller

[coderabbitai]

Walkthrough

This PR migrates database access patterns across multiple handlers and utilities from per-request DbConnection parameters to a centralized connection pool approach, introducing a new run_query macro in service_utils to manage transactions, connection acquisition, and error handling uniformly.

Changes

Cohort / File(s)	Summary
Database Macro Infrastructure crates/service_utils/src/db.rs	Added new run_query macro that abstracts connection pool acquisition, prepared statement cache disabling, transaction execution, and error mapping to unified AppError handling.
Experiments Handler crates/experimentation_platform/src/api/experiments/handlers.rs	Converted get_applicable_variants_helper from async to synchronous function; updated call site to use direct invocation instead of .await?.
Service Utils & Middleware crates/service_utils/src/helpers.rs, crates/service_utils/src/middlewares/auth_n/helpers.rs, crates/service_utils/src/middlewares/workspace_context.rs	Refactored database access to use connection pools and run_query macro; removed manual prepared statement cache management and direct connection extraction; simplified workspace context retrieval.
Organisation Handlers crates/superposition/src/organisation/handlers.rs	Migrated all handlers (create, update, get, list) from accepting DbConnection parameters to using Data with run_query-wrapped database operations.
Resolve Handler crates/superposition/src/resolve/handlers.rs	Replaced direct connection usage in resolve_with_exp_handler with run_query abstraction for all database interactions including max_created_at, config generation, variant fetching, and audit header operations.
Webhooks Handlers & Helpers crates/superposition/src/webhooks/handlers.rs, crates/superposition/src/webhooks/helper.rs	Migrated all webhook handlers and helper functions from direct PooledConnection parameters to connection pool references; refactored validate_events and fetch_webhook to use run_query pattern; removed explicit per-request connection handling.
Workspace Handlers crates/superposition/src/workspace/handlers.rs	Migrated all handlers (get, create, update, list, migrate_schema, rotate_encryption_key) from DbConnection to Data with run_query pattern; updated imports to include DBConnection and run_query.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• feat: AuthZ setup #543: Modifies the same experiments handlers module and get_applicable_variants handler call path with annotation changes alongside the async-to-sync conversion.
• feat: secrets support for functions/webhooks #773: Related migration effort moving handlers and helpers from per-request DbConnection to pooled run_query pattern with AppState threading through handler signatures.

Suggested reviewers

• knutties
• mahatoankitkumar
• Datron

Poem

🐰 A pool of connections flows so free,
From handlers deep through the service tree,
No more per-request, the macros declare—
Transactions managed with utmost care!
The refactoring hops along with delight! ✨

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately reflects the main objective of the pull request: granularizing database connection usage in the superposition crate to limit connection pool acquisition to only query execution periods.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch feat/granular/connection/1

---

No actionable comments were generated in the recent review. 🎉

🧹 Recent nitpick comments

crates/service_utils/src/db.rs (3)

15-18: Disabling the prepared statement cache on every connection checkout is a performance concern.

set_prepared_statement_cache_size(..., CacheSize::Disabled) is called each time a connection is obtained from the pool. Prepared statements are a significant Diesel/PostgreSQL optimization for repeated queries. Was this intentionally added to work around a specific issue (e.g., schema switching), or is it carried over from the previous per-request pattern? If it's needed for schema-name dynamic queries, consider documenting the rationale inline.

---

19-21: Every query is unconditionally wrapped in a transaction, even single read-only statements.

While PostgreSQL implicitly wraps single statements, the explicit BEGIN/COMMIT round-trips add slight overhead. For single-statement reads (e.g., SELECT in get_workspace, fetch_webhook, list_handler), this is unnecessary. Consider offering a non-transactional variant (e.g., run_query_no_tx!) or letting callers opt in to transactions only when needed.

---

10-27: The macro relies on external crate paths being available at the call site.

The macro uses diesel::Connection, superposition_types::result::AppError, and anyhow::anyhow! without importing them—this works because they're fully qualified. However, the Into::into call on Line 20 (map_err(Into::into)) depends on the caller's error type implementing Into<AppError>. This is fine for Diesel errors but consider adding a brief doc comment on the macro explaining the expected error conversion constraint.

crates/service_utils/src/middlewares/auth_n/helpers.rs (1)

21-40: Inconsistent with the run_query! pattern used elsewhere in this PR.

This function manually calls db_pool.get() without using run_query!. While the different error type (&'static str vs AppError) makes direct run_query! usage awkward, the prepared statement cache is no longer disabled here (it was removed per the summary). If disabling the cache was important for correctness with schema-scoped queries, note that this function queries the superposition schema directly so it may be unaffected. Still worth confirming this is intentional.

crates/superposition/src/organisation/handlers.rs (1)

124-142: Count and data queries use separate run_query! calls — minor consistency concern.

Lines 124–126 fetch total_items in one connection/transaction, and Line 142 fetches the paginated data in a separate connection/transaction. Between these two calls, rows could be inserted or deleted, leading to a stale total_items relative to the returned data. This is a pre-existing pattern common in paginated APIs and unlikely to cause real issues, but if exact consistency matters, both queries could be combined into a single run_query! block.

Combine count and data fetch into a single `run_query!` block

-    let total_items = run_query!(state.db_pool, |conn| organisations::table
-        .count()
-        .get_result(conn))?;
-
-    // Set up pagination
-    let limit = filters.count.unwrap_or(10);
-    let mut builder = organisations::table
-        .into_boxed()
-        .order(organisations::created_at.desc())
-        .limit(limit);
-
-    // Apply offset if page is specified
-    if let Some(page) = filters.page {
-        let offset = (page - 1) * limit;
-        builder = builder.offset(offset);
-    }
-
-    // Get paginated results
-    let data = run_query!(state.db_pool, |conn| builder.load(conn))?;
+    let limit = filters.count.unwrap_or(10);
+    let page = filters.page;
+
+    let (total_items, data) = run_query!(state.db_pool, |conn| {
+        let total_items: i64 = organisations::table
+            .count()
+            .get_result(conn)?;
+
+        let mut builder = organisations::table
+            .into_boxed()
+            .order(organisations::created_at.desc())
+            .limit(limit);
+
+        if let Some(page) = page {
+            let offset = (page - 1) * limit;
+            builder = builder.offset(offset);
+        }
+
+        let data: Vec<Organisation> = builder.load(conn)?;
+        Ok((total_items, data))
+    })?;

crates/superposition/src/workspace/handlers.rs (1)

198-215: Unnecessary transaction for a single UPDATE statement.

The update_handler wraps a single diesel::update call in what appears to be a transactional run_query! (via the |transaction_conn: &mut DBConnection| annotation), but there's only one DB operation here. Using the simpler |conn| form (as in get_handler) would avoid the transaction overhead.

♻️ Suggested simplification

     let updated_workspace =
-        run_query!(state.db_pool, |transaction_conn: &mut DBConnection| {
-            let updated_workspace = diesel::update(workspaces::table)
+        run_query!(state.db_pool, |conn| {
+            diesel::update(workspaces::table)
                 .filter(workspaces::organisation_id.eq(&org_id.0))
                 .filter(workspaces::workspace_name.eq(workspace_name))
                 .set((
                     request,
                     workspaces::last_modified_by.eq(user.email),
                     workspaces::last_modified_at.eq(timestamp),
                 ))
-                .get_result::<Workspace>(transaction_conn)
+                .get_result::<Workspace>(conn)
                 .map_err(|err| {
                     log::error!("failed to update workspace with error: {}", err);
                     err
-                })?;
-
-            Ok::<Workspace, superposition::AppError>(updated_workspace)
+                })
         })?;

crates/superposition/src/resolve/handlers.rs (2)

46-51: Five separate pool checkouts on a hot resolve path — consider consolidating.

resolve_with_exp_handler acquires a connection from the pool five times per request (get_max_created_at, generate_config_from_version, get_applicable_variants_helper, resolve, add_audit_id_to_header). The TODO comments acknowledge this, but for a latency-sensitive resolve endpoint this is a meaningful overhead increase compared to the previous single-connection approach. Each checkout involves pool contention, potential waits, and handshake cost.

While the TODOs indicate this is interim, consider grouping at least the tightly coupled calls (e.g., generate_config_from_version + resolve, or the final resolve + add_audit_id_to_header) into a single run_query! block now to reduce checkout count on this critical path.

Also applies to: 66-73, 90-101

---

105-109: Wrapping a non-query side-effect in run_query! solely for connection access.

add_audit_id_to_header mutates the HTTP response builder and happens to need a DB connection. Wrapping it in run_query! with an artificial Ok::<(), superposition::AppError>(()) return feels like a workaround. If the macro is designed for query execution (possibly with transaction semantics), using it for a header-mutation side-effect is a mismatch. Consider whether add_audit_id_to_header could be refactored to accept a pool reference directly instead.

crates/superposition/src/webhooks/handlers.rs (1)

34-82: Inconsistent DB access pattern between validate_change_reason and validate_events.

validate_change_reason (line 45) is wrapped in run_query! and receives a raw connection, while validate_events (line 52) receives &state.db_pool directly. This inconsistency suggests the two helpers have divergent signatures — one expects a connection, the other a pool. For maintainability, consider aligning them to the same pattern (both accepting a pool, or both called via run_query!).

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}