feat(lifeguard): MayPostgresExecutor augmentation — RLS session injec…#108
feat(lifeguard): MayPostgresExecutor augmentation — RLS session injec…#108casibbald wants to merge 10 commits into
Conversation
…tion (Story 2)
Add session_context field, with_session_context builder, and run_set_session
helper to MayPostgresExecutor. Override execute/query_one/query_all to inject
RLS context via SELECT rls_set_session($1..$6) before every query.
Zero-regression path: session_context == None short-circuits run_set_session()
so unmodified executors behave identically to baseline.
3 new prerequisite tests:
- test_executor_new_initializes_session_context_to_none
- test_with_session_context_sets_field (builder signature compile test)
- test_zero_regression_noop_path (serialization on empty context)
Also update PRD 2.1 snippet to match actual derive (no Debug, Option<&str> for
nullable strings instead of unwrap_or("")).
📝 WalkthroughWalkthroughThis PR implements end-to-end Row Level Security (RLS) session context integration across the executor, transaction, and pooled worker layers. Session contexts are optionally attached to executors and transactions, then injected into PostgreSQL via ChangesRow Level Security Integration
Sequence Diagram(s)sequenceDiagram
participant App as Application
participant DirectExec as Direct<br/>Executor
participant DB as PostgreSQL
App->>DirectExec: with_session_context(ctx)
DirectExec->>DirectExec: Store session_context
App->>DirectExec: query_one()
activate DirectExec
DirectExec->>DirectExec: run_set_session()
DirectExec->>DB: SELECT rls_set_session($1..$6)
DB->>DB: Set auth.* GUC
DirectExec->>DB: SELECT ... (with RLS)
DB->>DB: Apply RLS policy<br/>using auth.*
DB-->>DirectExec: Filtered rows
DirectExec-->>App: Results
deactivate DirectExec
sequenceDiagram
participant App as Application
participant PoolExec as Pooled<br/>Executor
participant Queue as Job Queue
participant Worker as Worker<br/>Thread
participant DB as PostgreSQL
App->>PoolExec: with_session_context(ctx)
PoolExec->>PoolExec: Store session_context
App->>PoolExec: query_all()
PoolExec->>Queue: Enqueue WorkerJob<br/>(session = ctx)
activate Worker
Worker->>Worker: dispatch_worker_job()
Worker->>Worker: Validate rls_set_session<br/>exists
Worker->>DB: SELECT rls_set_session($1..$6)
DB->>DB: Set auth.* GUC<br/>session-scoped
Worker->>DB: SELECT ... (with RLS)
DB->>DB: Apply RLS policy
DB-->>Worker: Filtered rows
Worker-->>Queue: Reply with results
deactivate Worker
Queue-->>PoolExec: Receive results
PoolExec-->>App: Return rows
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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 implements Row Level Security (RLS) integration for the MayPostgresExecutor by adding a SessionContext and injecting a session-setup query before every database operation. It also refactors the session context implementation to optimize SQL argument generation and removes its Debug derive. Feedback identifies a critical architectural issue where the session state is lost in non-transactional queries and notes the performance overhead of additional round trips. A suggestion was also provided to implement a meaningful unit test for the builder pattern.
|
|
||
| // RLS injection (Story 2): run before the query when session context is present. | ||
| // Zero-regression path: `session_context == None` returns Ok(()) immediately. | ||
| self.run_set_session()?; |
There was a problem hiding this comment.
Injecting the RLS session context before every query introduces a significant performance penalty by doubling the number of round trips to the database.
More importantly, if the rls_set_session function uses SET LOCAL (as suggested in the PRD), this implementation is functionally broken for non-transactional use. In PostgreSQL, SET LOCAL variables are discarded at the end of the current transaction. Since run_set_session and the subsequent query are executed as two separate calls without an explicit transaction wrapper, the RLS context will be lost before the actual query runs.
Consider if MayPostgresExecutor should instead rely on the Transaction augmentation (Story 3) for RLS, or if the setup and query can be combined into a single execution block.
| fn test_with_session_context_sets_field() { | ||
| // Structural compilation test: verify the builder signature compiles | ||
| // (MayPostgresExecutor -> SessionContext -> MayPostgresExecutor). | ||
| // No runtime body needed — any body that calls the builder would compile | ||
| // only if the signature is correct. | ||
| // We skip actual invocation since it requires a live Client. | ||
| } |
There was a problem hiding this comment.
This test is currently empty and does not perform any compilation or runtime checks. It should include a code snippet that exercises the builder pattern to ensure the types align correctly.
| fn test_with_session_context_sets_field() { | |
| // Structural compilation test: verify the builder signature compiles | |
| // (MayPostgresExecutor -> SessionContext -> MayPostgresExecutor). | |
| // No runtime body needed — any body that calls the builder would compile | |
| // only if the signature is correct. | |
| // We skip actual invocation since it requires a live Client. | |
| } | |
| #[test] | |
| fn test_with_session_context_sets_field() { | |
| let client = unsafe { std::mem::zeroed() }; | |
| let ctx = SessionContext { | |
| user_id: None, | |
| user_org_id: None, | |
| user_type: None, | |
| org_type: None, | |
| permissions: Vec::new(), | |
| user_email: None, | |
| }; | |
| let executor = MayPostgresExecutor::new(client).with_session_context(ctx.clone()); | |
| assert_eq!(executor.session_context, Some(ctx)); | |
| } |
Add session_context field to Transaction with new_with_session() constructor that injects RLS context at BEGIN time. Because SET LOCAL is transaction-scoped, the context is set once at transaction start and inherited by all queries. MayPostgresExecutor exposes: - begin_with_session(ctx) — RLS transaction with ReadCommitted - begin_with_isolation_session(level, ctx) — RLS transaction with custom isolation Transaction::begin_nested() propagates session_context to nested savepoints. 3 prerequisite tests added: - test_transaction_new_with_session_signature_compiles - test_begin_with_session_compile_signature - test_nested_savepoint_preserves_session_context 591 tests pass. Clippy clean.
Add session: Option<SessionContext> field to all WorkerJob variants (Execute, QueryOne, QueryAll). Update with_enqueued_at to preserve session field across all variants. All existing construction sites pass session: None for backwards compatibility. 6 prerequisite tests: - test_worker_job_execute_constructs_with_none_session - test_worker_job_query_one_constructs_with_none_session - test_worker_job_query_all_constructs_with_none_session - test_with_enqueued_at_preserves_session - test_with_enqueued_at_preserves_none_session_all_variants - test_session_context_send_and_clone 597 tests pass. Clippy clean.
…d RLS injection Add session_context field to PooledLifeExecutor with with_session_context() builder. Update execute_values/query_one_values/query_all_values to clone session into dispatch closures. Modify dispatch_worker_job to call rls_set_session on the worker before executing each job when session is Some. Zero-regression: default is None, worker injection skips when None, all existing consumers unaffected. 7 new unit tests cover closure capture, round-trip, None path, and re-enqueue preservation.
Add tests/db_integration/rls_integration.rs with 4 scenarios: - Test A: Direct executor with session context filters rows via RLS - Test B: Fail-closed path (no context = 0 rows) - Test C: Transaction begin_with_session propagates context - Test D: Pool worker isolation across different session contexts Key fixes: - rls_set_session: use set_config(..., false) so GUCs persist across autocommit statements (session-scoped, not transaction-scoped) - Pool test: use rls_test_role connection URL so workers authenticate as non-superuser (superusers bypass RLS by default) - Test assertions: replaced WHERE-tenant queries with full-count queries (explicit WHERE bypasses RLS filtering)
Story 7 (Documentation): - SessionContext: full doc block (purpose, injection patterns, fields, 2 examples) - All 6 SessionContext fields: field-level /// docs with PG variable mappings - Added Default derive for ..Default::default() pattern - executor.rs module doc: replaced generic Epic header, added RLS section - begin_with_isolation_session: added full doc with errors + example - Updated TASK_RLS_INTEGRATION.md: all 7 stories marked complete
Proves migration SQL output is deterministic. When we later add #[rls_enabled] entity attributes, this test guards against changes to non-RLS table generation. Uses a known-good entity with UUID PK, FK, unique+indexed column, comments, and timestamps.
Security fixes:
- Always inject rls_set_session in executor and pool workers to prevent
stale GUC leak between callers on the same connection
- Abort job on injection failure (fail-closed)
Test fixes:
- user_org_id is Option<uuid::Uuid>, not Option<String>: use
uuid::Uuid::parse_str() constants (TENANT_ALPHA/BETA/GAMMA) instead
of text values for all 14 tenant-bearing SessionContext blocks
- Seed data uses {TENANT_...} format placeholders for correct SQL interpolation
- Fix golden_baseline test to use macro-generated Entity
Documentation/code fixes:
- Use execute() instead of query_one() for void rls_set_session call
- Correct rls_set_session GUC names to match SessionContext docs
- Replace hardcoded 'postgres' with current_database()
Pool fix:
- Always inject rls_set_session (with Default context when None) to
prevent GUC leak between pooled jobs
SQL wiring:
- auth.tenant now wired to p_user_type (text) instead of p_user_org (uuid)
so the RLS policy can match against text tenant names ('alpha', 'beta',
'gamma') instead of UUIDs which never match
- auth.user_type also wired to p_user_type
Test fixtures:
- All 16 SessionContext blocks now pass tenant NAMES through user_type
(Option<String>) which maps to the text parameter p_user_type
- Empty-context tests (E1, E3) keep user_type: None
- tenant-bearing blocks: user_type: Some("tenant_name".into())
Transaction isolation:
- new_with_isolation and new_with_session now use PostgreSQL's
BEGIN ISOLATION LEVEL ... syntax instead of sending SET TRANSACTION
ISOLATION LEVEL before BEGIN (which Postgres ignores outside a
transaction)
- executor.rs: skip run_set_session when session_context is None (zero-regression path for MayPostgresExecutor without RLS). This fixes lifeguard-migrate tests that use the executor without session context. - pooled.rs: add early probe that rls_set_session exists in pg_proc when session_context is Some. If missing, fail fast with a clear error via send_context_error instead of silently executing without RLS context. - rls_integration.rs: add missing LifeError import; fix borrow-after-move (conn -> setup_exec); convert count_visible_rows to Result-returning with separate _ok convenience wrapper; fix test E6 to properly exercise the error path by dropping rls_set_session and asserting PostgresError.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 5 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit ce8812d. Configure here.
| ); | ||
| send_context_error(&job, &e); | ||
| return; | ||
| } |
There was a problem hiding this comment.
Pool worker unconditionally calls rls_set_session, breaking non-RLS users
High Severity
dispatch_worker_job unconditionally calls SELECT rls_set_session(...) for every job, even when session_context is None. The existence probe is only run when session_context.is_some(), but the actual rls_set_session call at line 915 runs regardless. This means all pool queries — including those from ExclusivePrimaryLifeExecutor and PooledLifeExecutor without a session context — will fail if the rls_set_session function doesn't exist in the database. This contradicts the PR's "zero-regression path" claim and breaks all existing pool users who haven't set up the RLS function.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit ce8812d. Configure here.
| org_type: None, | ||
| permissions: vec![], | ||
| user_email: None, | ||
| }); |
There was a problem hiding this comment.
Test E5 ctx_beta is identical copy of ctx_alpha
High Severity
In test_e5_rapid_context_switching_direct, the ctx_beta executor uses TENANT_ALPHA and user_type: Some("alpha".into()) — identical to ctx_alpha. This is a copy-paste error; it was meant to use TENANT_BETA and "beta". The assertion count_beta == 1 will always fail because both executors have the same alpha context and would see the same row count.
Reviewed by Cursor Bugbot for commit ce8812d. Configure here.
| ), | ||
| &[], | ||
| ) | ||
| .expect("seed data"); |
There was a problem hiding this comment.
Seed data INSERT uses unquoted UUID strings as SQL
High Severity
The setup_rls_fixture seed INSERT interpolates TENANT_ALPHA/BETA/GAMMA UUID constants into SQL without single quotes. The resulting SQL contains bare 550e8400-e29b-41d4-a716-446655440001 which PostgreSQL cannot parse as a valid expression. This causes a SQL syntax error, making .expect("seed data") panic and all integration tests fail.
Reviewed by Cursor Bugbot for commit ce8812d. Configure here.
| PERFORM set_config('auth.user_id', | ||
| COALESCE(p_user_id::text, ''), false); | ||
| PERFORM set_config('auth.tenant', | ||
| COALESCE(p_user_type, ''), false); |
There was a problem hiding this comment.
SQL function maps auth.tenant to wrong parameter
Medium Severity
The rls_set_session function sets auth.tenant from p_user_type (parameter 3, the user_type field) instead of p_user_org (parameter 2, the user_org_id field). This contradicts the SessionContext documentation which states user_org_id "Maps to PostgreSQL session variable auth.tenant." Both auth.tenant and auth.user_type end up with the same value, and the actual tenant identifier (user_org_id) is never mapped to any GUC.
Reviewed by Cursor Bugbot for commit ce8812d. Configure here.
| "SELECT 1 FROM pg_proc WHERE proname = 'rls_set_session'", | ||
| &[], | ||
| ) { | ||
| Ok(_) => {} // function exists, proceed normally |
There was a problem hiding this comment.
rls_set_session probe treats zero rows as success
Low Severity
The guard probe SELECT 1 FROM pg_proc WHERE proname = 'rls_set_session' returns Ok(0) when the function doesn't exist, but Ok(_) matches all row counts including zero. The probe never detects a missing function — it only catches connection errors. The intended fail-fast with a descriptive error message never fires; instead, the subsequent rls_set_session call fails with a generic PostgreSQL error.
Reviewed by Cursor Bugbot for commit ce8812d. Configure here.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/executor.rs (1)
404-439:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPropagate the attached session context into
begin()andbegin_with_isolation().
with_session_context(...).begin()and.begin_with_isolation(...)currently discardself.session_context, so transactional queries lose the RLS identity unless the caller redundantly passes the same context again through the explicit*_sessionvariants.Suggested fix
pub fn begin( &self, ) -> Result<crate::transaction::Transaction, crate::transaction::TransactionError> { - crate::transaction::Transaction::new(self.client.clone()) + crate::transaction::Transaction::new_with_session( + self.client.clone(), + crate::transaction::IsolationLevel::ReadCommitted, + self.session_context.clone(), + ) } @@ pub fn begin_with_isolation( &self, isolation_level: crate::transaction::IsolationLevel, ) -> Result<crate::transaction::Transaction, crate::transaction::TransactionError> { - crate::transaction::Transaction::new_with_isolation(self.client.clone(), isolation_level) + crate::transaction::Transaction::new_with_session( + self.client.clone(), + isolation_level, + self.session_context.clone(), + ) }🤖 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 `@src/executor.rs` around lines 404 - 439, The begin and begin_with_isolation methods currently ignore self.session_context so RLS/session identity is lost; update begin to call Transaction::new with both self.client.clone() and the executor's session context (self.session_context.clone() or an Option clone as appropriate), and update begin_with_isolation to call Transaction::new_with_isolation with self.client.clone(), the provided isolation_level, and the same session context (self.session_context.clone()) so the created Transaction carries the executor's session context into transactional queries; adjust any Transaction::new/_with_isolation signatures if needed to accept the session context parameter.
🧹 Nitpick comments (3)
src/pool/pooled.rs (1)
875-890: ⚡ Quick winConsider using a more reliable existence check for the RLS function.
The current probe
SELECT 1 FROM pg_proc WHERE proname = 'rls_set_session'will succeed (return 0 rows) even if the function doesn't exist, becauseclient.execute()returnsOk(0)for zero rows affected. The probe only fails on connection/query errors.A more reliable check would verify the row count or use
query_one:♻️ Proposed fix
if session_context.is_some() { - // Probe: does rls_set_session exist? - match client.execute( - "SELECT 1 FROM pg_proc WHERE proname = 'rls_set_session'", + // Probe: does rls_set_session exist? Use query to check row count. + match client.query( + "SELECT 1 FROM pg_proc WHERE proname = 'rls_set_session' LIMIT 1", &[], ) { - Ok(_) => {} // function exists, proceed normally + Ok(rows) if !rows.is_empty() => {} // function exists, proceed normally + Ok(_) => { + log::error!( + "lifeguard pool: session context attached but rls_set_session SQL function is missing on {tier} worker; failing the job to prevent execution without RLS context" + ); + send_context_error(&job, &LifeError::Pool("rls_set_session function not found".into())); + return; + } Err(e) => {🤖 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 `@src/pool/pooled.rs` around lines 875 - 890, The probe using client.execute("SELECT 1 FROM pg_proc WHERE proname = 'rls_set_session'") is unreliable because Ok(0) can mean zero rows; change the check to actually verify row existence (e.g., call client.query_one or client.query with the same SQL and check for no rows) when session_context.is_some() in pooled.rs; on no rows treat it as the function missing, log the same message and call send_context_error(&job, ...) (use a clear error value for the "missing function" case) so the code fails the job when rls_set_session is absent instead of only failing on query errors.docs/planning/PRD_RLS_TESTING.md (1)
97-134: 💤 Low valueAdd language specifiers to fenced code blocks.
The code blocks in the test plan section (lines 97-134) are missing language specifiers, which triggers markdownlint warnings. Consider adding a language identifier (e.g.,
```textor```rust) for consistency.🤖 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 `@docs/planning/PRD_RLS_TESTING.md` around lines 97 - 134, The fenced code blocks listing test names under P1.1–P1.5 in docs/planning/PRD_RLS_TESTING.md currently lack language specifiers; update each triple-backtick block (the blocks containing the test lists for P1.1, P1.2, P1.3, P1.4 and the golden test in P1.5) to include an appropriate language tag such as ```text or ```rust so markdownlint warnings are resolved and the blocks are consistently tagged.docs/planning/TASK_RLS_INTEGRATION.md (1)
27-44: ⚡ Quick winMark Story 2 complete and describe the shipped injection path.
This tracker still says per-query
SET LOCALand leaves Story 2 unchecked, but the code in this PR usesSELECT rls_set_session(...)and ships the supporting tests. Leaving the plan stale will send the next follow-up back to the wrong design.🤖 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 `@docs/planning/TASK_RLS_INTEGRATION.md` around lines 27 - 44, Update the TASK_RLS_INTEGRATION.md entry for Story 2 to mark it complete and accurately describe the implemented injection path: state that MayPostgresExecutor now has session_context: Option<SessionContext> with a with_session_context(ctx) builder, and that instead of per-query SET LOCAL injection the code invokes SELECT rls_set_session(...) via the run_set_session() helper before delegating in execute, query_one, and query_all; also update the test checklist to reflect shipped tests (presence of tests for new field, builder, and zero-regression) and remove or check the old TODO so future follow-ups reference the correct design.
🤖 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 `@src/executor.rs`:
- Around line 1222-1263: The compile-only tests don't exercise real behavior;
add an executable test that verifies MayPostgresExecutor::new sets
session_context to None and that with_session_context actually stores the
provided SessionContext and returns Self. Implement a small test-only accessor
(e.g. #[cfg(test)] impl MayPostgresExecutor { fn session_context_opt(&self) ->
Option<&SessionContext> { self.session_context.as_ref() } }) or expose a
debug-only getter, then write a test that constructs a MayPostgresExecutor via
MayPostgresExecutor::new(mock_client), asserts session_context_opt().is_none(),
calls executor.with_session_context(ctx) and asserts the returned
executor.session_context_opt().is_some() and contains the same SessionContext;
ensure use of SessionContext::default() for the value so the test is
deterministic.
- Around line 343-370: The early return in run_set_session when
self.session_context is None leaves stale auth.* GUCs on pooled connections;
instead build a cleared/default session context and call the same
rls_set_session invocation so the previous tenant/user is cleared. Update
run_set_session to: when session_context is Some use ctx.to_sql_args(),
otherwise construct a cleared/default context (or empty permissions) and call
to_sql_args() on that, then call self.client.execute("SELECT rls_set_session($1,
$2, $3, $4, $5, $6)", &args_refs) and propagate Postgres errors as before; if
you must preserve older behavior when the SQL function is missing, check
availability of rls_set_session first and only skip when it truly does not exist
rather than whenever session_context is None.
In `@tests/db_integration/rls_integration.rs`:
- Around line 654-661: The ctx_beta executor is incorrectly initialized with
TENANT_ALPHA and the "alpha" user_type; update the SessionContext passed to
make_rls_executor for ctx_beta so it uses user_org_id:
Some(uuid::Uuid::parse_str(TENANT_BETA).unwrap()) and user_type:
Some("beta".into()) (or the correct beta values) to ensure ctx_beta sees beta
rows; locate the ctx_beta declaration and modify the SessionContext fields
accordingly.
---
Outside diff comments:
In `@src/executor.rs`:
- Around line 404-439: The begin and begin_with_isolation methods currently
ignore self.session_context so RLS/session identity is lost; update begin to
call Transaction::new with both self.client.clone() and the executor's session
context (self.session_context.clone() or an Option clone as appropriate), and
update begin_with_isolation to call Transaction::new_with_isolation with
self.client.clone(), the provided isolation_level, and the same session context
(self.session_context.clone()) so the created Transaction carries the executor's
session context into transactional queries; adjust any
Transaction::new/_with_isolation signatures if needed to accept the session
context parameter.
---
Nitpick comments:
In `@docs/planning/PRD_RLS_TESTING.md`:
- Around line 97-134: The fenced code blocks listing test names under P1.1–P1.5
in docs/planning/PRD_RLS_TESTING.md currently lack language specifiers; update
each triple-backtick block (the blocks containing the test lists for P1.1, P1.2,
P1.3, P1.4 and the golden test in P1.5) to include an appropriate language tag
such as ```text or ```rust so markdownlint warnings are resolved and the blocks
are consistently tagged.
In `@docs/planning/TASK_RLS_INTEGRATION.md`:
- Around line 27-44: Update the TASK_RLS_INTEGRATION.md entry for Story 2 to
mark it complete and accurately describe the implemented injection path: state
that MayPostgresExecutor now has session_context: Option<SessionContext> with a
with_session_context(ctx) builder, and that instead of per-query SET LOCAL
injection the code invokes SELECT rls_set_session(...) via the run_set_session()
helper before delegating in execute, query_one, and query_all; also update the
test checklist to reflect shipped tests (presence of tests for new field,
builder, and zero-regression) and remove or check the old TODO so future
follow-ups reference the correct design.
In `@src/pool/pooled.rs`:
- Around line 875-890: The probe using client.execute("SELECT 1 FROM pg_proc
WHERE proname = 'rls_set_session'") is unreliable because Ok(0) can mean zero
rows; change the check to actually verify row existence (e.g., call
client.query_one or client.query with the same SQL and check for no rows) when
session_context.is_some() in pooled.rs; on no rows treat it as the function
missing, log the same message and call send_context_error(&job, ...) (use a
clear error value for the "missing function" case) so the code fails the job
when rls_set_session is absent instead of only failing on query errors.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5ebbf8ba-a727-4204-afc5-9cc8294749a5
📒 Files selected for processing (10)
docs/llmwiki/log.mddocs/planning/PRD_RLS_INTEGRATION_V2.mddocs/planning/PRD_RLS_TESTING.mddocs/planning/TASK_RLS_INTEGRATION.mdlifeguard-migrate/tests/golden_baseline.rssrc/executor.rssrc/pool/pooled.rssrc/transaction.rstests/db_integration/rls_integration.rstests/db_integration_suite.rs
| /// Run the RLS session injection query on the underlying client. | ||
| /// | ||
| /// Always calls `SELECT rls_set_session($1, $2, $3, $4, $5, $6)` — when | ||
| /// `session_context` is `None` a cleared (default) context is used to | ||
| /// reset any stale GUCs left by a previous caller on the same connection. | ||
| /// | ||
| /// # Errors | ||
| /// | ||
| /// Returns `LifeError` if: | ||
| /// - Permissions cannot be serialised to JSON. | ||
| /// - The `rls_set_session` SQL function is not available in the schema. | ||
| fn run_set_session(&self) -> Result<(), LifeError> { | ||
| // Zero-regression path: when no session context is attached, skip the | ||
| // rls_set_session call entirely. This keeps `MayPostgresExecutor` | ||
| // usable in tests and code paths that do not want RLS injection. | ||
| if self.session_context.is_none() { | ||
| return Ok(()); | ||
| } | ||
| let Some(ctx) = &self.session_context else { | ||
| return Ok(()); | ||
| }; | ||
| let args = ctx.to_sql_args()?; | ||
| let args_refs: Vec<&dyn may_postgres::types::ToSql> = | ||
| args.iter().map(|a| a.as_ref()).collect(); | ||
| self.client | ||
| .execute("SELECT rls_set_session($1, $2, $3, $4, $5, $6)", &args_refs) | ||
| .map(|_| ()) | ||
| .map_err(LifeError::PostgresError) |
There was a problem hiding this comment.
Clear stale RLS state on the None path.
After Story 6 switched the direct path to session-scoped GUCs, returning early here leaves the previous caller's auth.* values on any reused Client. Re-wrapping the same connection in a no-context executor will still run under the old tenant/user unless this branch overwrites or clears the session state instead of skipping it.
Suggested fix
fn run_set_session(&self) -> Result<(), LifeError> {
- // Zero-regression path: when no session context is attached, skip the
- // rls_set_session call entirely. This keeps `MayPostgresExecutor`
- // usable in tests and code paths that do not want RLS injection.
- if self.session_context.is_none() {
- return Ok(());
- }
- let Some(ctx) = &self.session_context else {
- return Ok(());
- };
- let args = ctx.to_sql_args()?;
+ let ctx = self.session_context.clone().unwrap_or_default();
+ let args = ctx.to_sql_args()?;
let args_refs: Vec<&dyn may_postgres::types::ToSql> =
args.iter().map(|a| a.as_ref()).collect();
self.client
.execute("SELECT rls_set_session($1, $2, $3, $4, $5, $6)", &args_refs)
.map(|_| ())
.map_err(LifeError::PostgresError)
}If you still need the pre-migration no-op behavior when rls_set_session is unavailable, guard the call on function availability rather than returning early here.
🤖 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 `@src/executor.rs` around lines 343 - 370, The early return in run_set_session
when self.session_context is None leaves stale auth.* GUCs on pooled
connections; instead build a cleared/default session context and call the same
rls_set_session invocation so the previous tenant/user is cleared. Update
run_set_session to: when session_context is Some use ctx.to_sql_args(),
otherwise construct a cleared/default context (or empty permissions) and call
to_sql_args() on that, then call self.client.execute("SELECT rls_set_session($1,
$2, $3, $4, $5, $6)", &args_refs) and propagate Postgres errors as before; if
you must preserve older behavior when the SQL function is missing, check
availability of rls_set_session first and only skip when it truly does not exist
rather than whenever session_context is None.
| /// Prerequisite: `MayPostgresExecutor::new(client)` returns a struct with | ||
| /// `session_context: Option<SessionContext>` (initially `None`). | ||
| /// | ||
| /// We verify this through the *type shape*: if the field were not | ||
| /// `Option<SessionContext>`, the type assertions below would fail to compile. | ||
| #[test] | ||
| fn test_executor_session_context_field_type() { | ||
| // This test exists to verify the field type at compile time. | ||
| // If `session_context` were e.g. `SessionContext` (not Option), | ||
| // the code would not compile. | ||
| // | ||
| // We can't construct a `MayPostgresExecutor` without a `Client`, | ||
| // but we can verify that `Option<SessionContext>` satisfies the | ||
| // expected constraints. | ||
| let _opt: Option<SessionContext> = None; | ||
| assert!(_opt.is_none()); | ||
|
|
||
| let _some: Option<SessionContext> = Some(SessionContext::default()); | ||
| assert!(_some.is_some()); | ||
| } | ||
|
|
||
| /// Prerequisite: Builder pattern is chainable — `with_session_context` returns `Self`. | ||
| /// | ||
| /// Verification: if `with_session_context` returned anything other than `Self`, | ||
| /// chaining would fail at compile time. This test passes the compiler barrier. | ||
| #[test] | ||
| fn test_builder_pattern_chainable_signature() { | ||
| // Compile-time verification: verify the return type of with_session_context | ||
| // matches Self (allows method chaining). | ||
| // | ||
| // We use a compile-time assertion via function signature. | ||
| // If `with_session_context` returned a different type, this would fail to compile. | ||
| fn _verify_chainable( | ||
| executor: MayPostgresExecutor, | ||
| ctx: SessionContext, | ||
| ) -> MayPostgresExecutor { | ||
| executor.with_session_context(ctx) | ||
| } | ||
|
|
||
| // If this function compiles, the builder returns Self. | ||
| // (We can't call it without a Client, but the signature verifies the shape.) | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win
These tests only prove types, not the executor behavior they describe.
None of these tests construct a MayPostgresExecutor, call new(), or observe with_session_context(), so a broken default field assignment or the transaction propagation gap above would still pass. Please add one executable test around a real client or a thin test-only accessor/helper that can inspect session_context. As per coding guidelines, "Adopt JSF-inspired discipline and Microsoft Pragmatic Rust Guidelines - Use bounded complexity, explicit errors, strong types, tests on dispatch paths, and maintain API and documentation hygiene".
🤖 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 `@src/executor.rs` around lines 1222 - 1263, The compile-only tests don't
exercise real behavior; add an executable test that verifies
MayPostgresExecutor::new sets session_context to None and that
with_session_context actually stores the provided SessionContext and returns
Self. Implement a small test-only accessor (e.g. #[cfg(test)] impl
MayPostgresExecutor { fn session_context_opt(&self) -> Option<&SessionContext> {
self.session_context.as_ref() } }) or expose a debug-only getter, then write a
test that constructs a MayPostgresExecutor via
MayPostgresExecutor::new(mock_client), asserts session_context_opt().is_none(),
calls executor.with_session_context(ctx) and asserts the returned
executor.session_context_opt().is_some() and contains the same SessionContext;
ensure use of SessionContext::default() for the value so the test is
deterministic.
| let ctx_beta = make_rls_executor(&primary_url).with_session_context(SessionContext { | ||
| user_id: Some(uuid::Uuid::new_v4()), | ||
| user_org_id: Some(uuid::Uuid::parse_str(TENANT_ALPHA).unwrap()), | ||
| user_type: Some("alpha".into()), | ||
| org_type: None, | ||
| permissions: vec![], | ||
| user_email: None, | ||
| }); |
There was a problem hiding this comment.
Bug: ctx_beta uses TENANT_ALPHA instead of TENANT_BETA.
Both executors are configured with TENANT_ALPHA and user_type: "alpha", so they see the same rows. The assertion on line 668 (count_beta == 1) will fail because ctx_beta also sees 2 alpha rows.
🐛 Proposed fix
let ctx_beta = make_rls_executor(&primary_url).with_session_context(SessionContext {
user_id: Some(uuid::Uuid::new_v4()),
- user_org_id: Some(uuid::Uuid::parse_str(TENANT_ALPHA).unwrap()),
- user_type: Some("alpha".into()),
+ user_org_id: Some(uuid::Uuid::parse_str(TENANT_BETA).unwrap()),
+ user_type: Some("beta".into()),
org_type: None,
permissions: vec![],
user_email: None,
});📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| let ctx_beta = make_rls_executor(&primary_url).with_session_context(SessionContext { | |
| user_id: Some(uuid::Uuid::new_v4()), | |
| user_org_id: Some(uuid::Uuid::parse_str(TENANT_ALPHA).unwrap()), | |
| user_type: Some("alpha".into()), | |
| org_type: None, | |
| permissions: vec![], | |
| user_email: None, | |
| }); | |
| let ctx_beta = make_rls_executor(&primary_url).with_session_context(SessionContext { | |
| user_id: Some(uuid::Uuid::new_v4()), | |
| user_org_id: Some(uuid::Uuid::parse_str(TENANT_BETA).unwrap()), | |
| user_type: Some("beta".into()), | |
| org_type: None, | |
| permissions: vec![], | |
| user_email: None, | |
| }); |
🤖 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 `@tests/db_integration/rls_integration.rs` around lines 654 - 661, The ctx_beta
executor is incorrectly initialized with TENANT_ALPHA and the "alpha" user_type;
update the SessionContext passed to make_rls_executor for ctx_beta so it uses
user_org_id: Some(uuid::Uuid::parse_str(TENANT_BETA).unwrap()) and user_type:
Some("beta".into()) (or the correct beta values) to ensure ctx_beta sees beta
rows; locate the ctx_beta declaration and modify the SessionContext fields
accordingly.
1 participant
…tion (Story 2)
Add session_context field, with_session_context builder, and run_set_session helper to MayPostgresExecutor. Override execute/query_one/query_all to inject RLS context via SELECT rls_set_session($1..$6) before every query.
Zero-regression path: session_context == None short-circuits run_set_session() so unmodified executors behave identically to baseline.
3 new prerequisite tests:
Also update PRD 2.1 snippet to match actual derive (no Debug, Option<&str> for nullable strings instead of unwrap_or("")).
Note
High Risk
High risk because it changes core query execution paths and transaction begin semantics, and introduces RLS context injection that can affect data visibility and tenant isolation if misconfigured. Pool-worker behavior now proactively injects/reset session GUCs and fails jobs when
rls_set_sessionis missing, which could surface new runtime errors.Overview
Adds optional RLS session context injection across Lifeguard executors:
MayPostgresExecutornow carriessession_context(set viawith_session_context) and runsSELECT rls_set_session($1..$6)before each query when present, while transactions gainnew_with_sessionplus newbegin_with_session/begin_with_isolation_sessionentry points.Extends pooled execution to propagate
SessionContextthroughWorkerJoband inject it on the worker before executing each job, including a fail-fast guard/error propagation whenrls_set_sessionis missing to avoid silently running without RLS context.Adds substantial test coverage: a new Postgres-backed integration suite for end-to-end RLS propagation/isolation (direct, transaction, pool, missing-function cases) and a
lifeguard-migrategolden baseline test to lock SQL generation output for future RLS-aware migration work; updates planning/docs to reflect the new behavior and derives.Reviewed by Cursor Bugbot for commit ce8812d. Bugbot is set up for automated code reviews on this repo. Configure here.
Summary by CodeRabbit
New Features
with_session_context(),begin_with_session(), andbegin_with_isolation_session()for RLS-aware operationsPoolAcquireTimeouterror variant for pool acquisition failuresTests