diff --git a/apps/api/src/lib/db-context.ts b/apps/api/src/lib/db-context.ts index 0e16fc8..eac6f7c 100644 --- a/apps/api/src/lib/db-context.ts +++ b/apps/api/src/lib/db-context.ts @@ -25,7 +25,7 @@ function getPrismaInstance(): PrismaClient { * * Note: SET LOCAL must be used within a transaction to ensure it's scoped * correctly with connection pooling. This is a low-level function - prefer - * using withUserContext or withUserTransaction for most use cases. + * using withUserContext or withWorkspaceContext for most use cases. * * @param userId - The UUID of the current user * @param client - Prisma client (required - must be a transaction client) @@ -42,6 +42,53 @@ export async function setCurrentUser(userId: string, client: PrismaClient): Prom await client.$executeRaw`SET LOCAL app.current_user_id = ${userId}`; } +/** + * Sets the current workspace ID for RLS policies within a transaction context. + * Must be called before executing any workspace-scoped queries. + * + * @param workspaceId - The UUID of the current workspace + * @param client - Prisma client (required - must be a transaction client) + * + * @example + * ```typescript + * await prisma.$transaction(async (tx) => { + * await setCurrentWorkspace(workspaceId, tx); + * const tasks = await tx.task.findMany(); // Filtered by workspace via RLS + * }); + * ``` + */ +export async function setCurrentWorkspace( + workspaceId: string, + client: PrismaClient +): Promise { + await client.$executeRaw`SET LOCAL app.current_workspace_id = ${workspaceId}`; +} + +/** + * Sets both user and workspace context for RLS policies. + * This is the recommended way to set context for workspace-scoped operations. + * + * @param userId - The UUID of the current user + * @param workspaceId - The UUID of the current workspace + * @param client - Prisma client (required - must be a transaction client) + * + * @example + * ```typescript + * await prisma.$transaction(async (tx) => { + * await setWorkspaceContext(userId, workspaceId, tx); + * const tasks = await tx.task.findMany(); // Filtered by workspace via RLS + * }); + * ``` + */ +export async function setWorkspaceContext( + userId: string, + workspaceId: string, + client: PrismaClient +): Promise { + await client.$executeRaw`SET LOCAL app.current_user_id = ${userId}`; + await client.$executeRaw`SET LOCAL app.current_workspace_id = ${workspaceId}`; +} + /** * Clears the current user context within a transaction. * Use this to reset the session or when switching users. @@ -55,6 +102,19 @@ export async function clearCurrentUser(client: PrismaClient): Promise { await client.$executeRaw`SET LOCAL app.current_user_id = NULL`; } +/** + * Clears both user and workspace context within a transaction. + * + * Note: SET LOCAL is automatically cleared at transaction end, + * so explicit clearing is typically unnecessary. + * + * @param client - Prisma client (required - must be a transaction client) + */ +export async function clearWorkspaceContext(client: PrismaClient): Promise { + await client.$executeRaw`SET LOCAL app.current_user_id = NULL`; + await client.$executeRaw`SET LOCAL app.current_workspace_id = NULL`; +} + /** * Executes a function with the current user context set within a transaction. * Automatically sets the user context and ensures it's properly scoped. @@ -121,6 +181,36 @@ export async function withUserTransaction( }); } +/** + * Executes a function with workspace context (user + workspace) set within a transaction. + * This is the recommended way for workspace-scoped operations. + * + * @param userId - The UUID of the current user + * @param workspaceId - The UUID of the current workspace + * @param fn - The function to execute with context (receives transaction client) + * @returns The result of the function + * + * @example + * ```typescript + * const tasks = await withWorkspaceContext(userId, workspaceId, async (tx) => { + * return tx.task.findMany({ + * where: { status: 'IN_PROGRESS' } + * }); + * }); + * ``` + */ +export async function withWorkspaceContext( + userId: string, + workspaceId: string, + fn: (tx: PrismaClient) => Promise +): Promise { + const prismaClient = getPrismaInstance(); + return prismaClient.$transaction(async (tx) => { + await setWorkspaceContext(userId, workspaceId, tx as PrismaClient); + return fn(tx as PrismaClient); + }); +} + /** * Higher-order function that wraps a handler with user context. * Useful for API routes and tRPC procedures. @@ -270,9 +360,13 @@ export function createAuthMiddleware(client: PrismaClient) { export default { setCurrentUser, + setCurrentWorkspace, + setWorkspaceContext, clearCurrentUser, + clearWorkspaceContext, withUserContext, withUserTransaction, + withWorkspaceContext, withAuth, verifyWorkspaceAccess, getUserWorkspaces, diff --git a/apps/api/src/prisma/prisma.service.spec.ts b/apps/api/src/prisma/prisma.service.spec.ts index c8d956c..25d841c 100644 --- a/apps/api/src/prisma/prisma.service.spec.ts +++ b/apps/api/src/prisma/prisma.service.spec.ts @@ -119,4 +119,86 @@ describe("PrismaService", () => { }); }); }); + + describe("setWorkspaceContext", () => { + it("should set workspace context variables in transaction", async () => { + const userId = "user-123"; + const workspaceId = "workspace-456"; + const executeRawSpy = vi.spyOn(service, "$executeRaw").mockResolvedValue(0); + + await service.$transaction(async (tx) => { + await service.setWorkspaceContext(userId, workspaceId, tx); + }); + + expect(executeRawSpy).toHaveBeenCalledTimes(2); + // Check that both session variables were set + expect(executeRawSpy).toHaveBeenNthCalledWith(1, expect.anything()); + expect(executeRawSpy).toHaveBeenNthCalledWith(2, expect.anything()); + }); + + it("should work when called outside transaction using default client", async () => { + const userId = "user-123"; + const workspaceId = "workspace-456"; + const executeRawSpy = vi.spyOn(service, "$executeRaw").mockResolvedValue(0); + + await service.setWorkspaceContext(userId, workspaceId); + + expect(executeRawSpy).toHaveBeenCalledTimes(2); + }); + }); + + describe("withWorkspaceContext", () => { + it("should execute function with workspace context set", async () => { + const userId = "user-123"; + const workspaceId = "workspace-456"; + const executeRawSpy = vi.spyOn(service, "$executeRaw").mockResolvedValue(0); + + const result = await service.withWorkspaceContext(userId, workspaceId, async () => { + return "test-result"; + }); + + expect(result).toBe("test-result"); + expect(executeRawSpy).toHaveBeenCalledTimes(2); + }); + + it("should pass transaction client to callback", async () => { + const userId = "user-123"; + const workspaceId = "workspace-456"; + vi.spyOn(service, "$executeRaw").mockResolvedValue(0); + + let receivedClient: unknown = null; + await service.withWorkspaceContext(userId, workspaceId, async (tx) => { + receivedClient = tx; + return null; + }); + + expect(receivedClient).toBeDefined(); + expect(receivedClient).toHaveProperty("$executeRaw"); + }); + + it("should handle errors from callback", async () => { + const userId = "user-123"; + const workspaceId = "workspace-456"; + vi.spyOn(service, "$executeRaw").mockResolvedValue(0); + + const error = new Error("Callback error"); + await expect( + service.withWorkspaceContext(userId, workspaceId, async () => { + throw error; + }) + ).rejects.toThrow(error); + }); + }); + + describe("clearWorkspaceContext", () => { + it("should clear workspace context variables", async () => { + const executeRawSpy = vi.spyOn(service, "$executeRaw").mockResolvedValue(0); + + await service.$transaction(async (tx) => { + await service.clearWorkspaceContext(tx); + }); + + expect(executeRawSpy).toHaveBeenCalledTimes(2); + }); + }); }); diff --git a/apps/api/src/prisma/prisma.service.ts b/apps/api/src/prisma/prisma.service.ts index 0fc7310..6f33293 100644 --- a/apps/api/src/prisma/prisma.service.ts +++ b/apps/api/src/prisma/prisma.service.ts @@ -79,4 +79,72 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul return { connected: false }; } } + + /** + * Sets workspace context for Row-Level Security (RLS) + * Sets both user_id and workspace_id session variables for PostgreSQL RLS policies + * + * IMPORTANT: Must be called within a transaction or use the default client + * Session variables are transaction-scoped (SET LOCAL) for connection pool safety + * + * @param userId - The ID of the authenticated user + * @param workspaceId - The ID of the workspace context + * @param client - Optional Prisma client (uses 'this' if not provided) + * + * @example + * ```typescript + * await prisma.$transaction(async (tx) => { + * await prisma.setWorkspaceContext(userId, workspaceId, tx); + * const tasks = await tx.task.findMany(); // Filtered by RLS + * }); + * ``` + */ + async setWorkspaceContext( + userId: string, + workspaceId: string, + client: PrismaClient = this + ): Promise { + await client.$executeRaw`SET LOCAL app.current_user_id = ${userId}`; + await client.$executeRaw`SET LOCAL app.current_workspace_id = ${workspaceId}`; + } + + /** + * Clears workspace context session variables + * Typically not needed as SET LOCAL is automatically cleared at transaction end + * + * @param client - Optional Prisma client (uses 'this' if not provided) + */ + async clearWorkspaceContext(client: PrismaClient = this): Promise { + await client.$executeRaw`SET LOCAL app.current_user_id = NULL`; + await client.$executeRaw`SET LOCAL app.current_workspace_id = NULL`; + } + + /** + * Executes a function with workspace context set within a transaction + * Automatically sets the context and ensures proper scoping + * + * @param userId - The ID of the authenticated user + * @param workspaceId - The ID of the workspace context + * @param fn - Function to execute with context (receives transaction client) + * @returns The result of the function + * + * @example + * ```typescript + * const tasks = await prisma.withWorkspaceContext(userId, workspaceId, async (tx) => { + * return tx.task.findMany({ + * where: { status: 'IN_PROGRESS' } + * }); + * }); + * ``` + */ + async withWorkspaceContext( + userId: string, + workspaceId: string, + fn: (tx: PrismaClient) => Promise + ): Promise { + return this.$transaction(async (tx) => { + await this.setWorkspaceContext(userId, workspaceId, tx as PrismaClient); + return fn(tx as PrismaClient); + }); + } } diff --git a/docs/scratchpads/195-rls-context-helpers.md b/docs/scratchpads/195-rls-context-helpers.md new file mode 100644 index 0000000..828814b --- /dev/null +++ b/docs/scratchpads/195-rls-context-helpers.md @@ -0,0 +1,97 @@ +# Issue #195: Implement RLS context helpers consistently across all services + +## Objective + +Implement consistent RLS (Row-Level Security) context helpers across all services to ensure proper workspace isolation through PostgreSQL session variables. + +## Current State Analysis + +Need to examine: + +1. Existing db-context.ts helpers +2. How services currently handle workspace filtering +3. Whether RLS policies are defined in Prisma schema +4. Best approach for setting PostgreSQL session variables + +## Approach + +**Use Prisma Extension with PostgreSQL Session Variables:** + +- Create a Prisma extension that sets session variables before queries +- Session variables: `app.current_workspace_id` and `app.current_user_id` +- Apply to all workspace-scoped operations +- This works with connection pooling (uses `SET LOCAL` which is transaction-scoped) + +## Implementation Plan + +- [x] Examine existing db-context.ts +- [x] Examine current service implementations +- [x] Write tests for RLS context setting (11 tests passing, 5 need actual DB) +- [x] Implement `setWorkspaceContext()` method in PrismaService +- [x] Create helper methods for workspace-scoped queries +- [x] Update db-context.ts with workspace context functions +- [x] Export new helper functions +- [ ] Document RLS usage patterns +- [ ] Add example of service using RLS context + +## Changes Made + +### PrismaService + +Added three new methods: + +1. `setWorkspaceContext(userId, workspaceId, client?)` - Sets session variables +2. `clearWorkspaceContext(client?)` - Clears session variables +3. `withWorkspaceContext(userId, workspaceId, fn)` - Transaction wrapper + +### db-context.ts + +Added new functions: + +1. `setCurrentWorkspace(workspaceId, client)` - Set workspace ID +2. `setWorkspaceContext(userId, workspaceId, client)` - Set both user and workspace +3. `clearWorkspaceContext(client)` - Clear both variables +4. `withWorkspaceContext(userId, workspaceId, fn)` - High-level transaction wrapper + +All functions use `SET LOCAL` for transaction-scoped variables (connection pool safe). + +## Usage Pattern + +```typescript +// In a service +const tasks = await this.prisma.withWorkspaceContext(userId, workspaceId, async (tx) => { + return tx.task.findMany({ + where: { status: "IN_PROGRESS" }, + }); +}); +``` + +Or using db-context helpers: + +```typescript +import { withWorkspaceContext } from "../lib/db-context"; + +const tasks = await withWorkspaceContext(userId, workspaceId, async (tx) => { + return tx.task.findMany(); +}); +``` + +## Testing Strategy + +### Unit Tests + +- PrismaService sets context correctly +- Context is cleared after transaction +- Multiple concurrent requests don't interfere + +### Integration Tests + +- Workspace isolation is enforced +- Cross-workspace queries are blocked +- RLS policies work with context variables + +## Notes + +- Must use transaction-scoped `SET LOCAL` (not session-level `SET`) +- Connection pooling compatible +- Should work with or without actual RLS policies (defense in depth)