stack/apps/api/src/lib/db-context.ts

/**
 * Database Context Utilities for Row-Level Security (RLS)
 *
 * This module provides utilities for setting the current user context
 * in the database, enabling Row-Level Security policies to automatically
 * filter queries to only the data the user is authorized to access.
 *
 * @see docs/design/multi-tenant-rls.md for full documentation
 */

import { PrismaClient } from "@prisma/client";

// Global prisma instance for standalone usage
// Note: In NestJS controllers/services, inject PrismaService instead
let prisma: PrismaClient | null = null;

function getPrismaInstance(): PrismaClient {
  prisma ??= new PrismaClient();
  return prisma;
}

/**
 * Sets the current user ID for RLS policies within a transaction context.
 * Must be called before executing any queries that rely on RLS.
 *
 * 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 withWorkspaceContext for most use cases.
 *
 * @param userId - The UUID of the current user
 * @param client - Prisma client (required - must be a transaction client)
 *
 * @example
 * ```typescript
 * await prisma.$transaction(async (tx) => {
 *   await setCurrentUser(userId, tx);
 *   const tasks = await tx.task.findMany(); // Automatically filtered by RLS
 * });
 * ```
 */
export async function setCurrentUser(userId: string, client: PrismaClient): Promise<void> {
  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<void> {
  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<void> {
  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.
 *
 * 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 clearCurrentUser(client: PrismaClient): Promise<void> {
  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<void> {
  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.
 *
 * @param userId - The UUID of the current user
 * @param fn - The function to execute with user context (receives transaction client)
 * @returns The result of the function
 *
 * @example
 * ```typescript
 * const tasks = await withUserContext(userId, async (tx) => {
 *   return tx.task.findMany({
 *     where: { workspaceId }
 *   });
 * });
 * ```
 */
export async function withUserContext<T>(
  userId: string,
  fn: (tx: PrismaClient) => Promise<T>
): Promise<T> {
  const prismaClient = getPrismaInstance();
  return prismaClient.$transaction(async (tx) => {
    await setCurrentUser(userId, tx as PrismaClient);
    return fn(tx as PrismaClient);
  });
}

/**
 * Executes a function within a transaction with the current user context set.
 * Useful for operations that need atomicity and RLS.
 *
 * @param userId - The UUID of the current user
 * @param fn - The function to execute with transaction and user context
 * @returns The result of the function
 *
 * @example
 * ```typescript
 * const workspace = await withUserTransaction(userId, async (tx) => {
 *   const workspace = await tx.workspace.create({
 *     data: { name: 'New Workspace', ownerId: userId }
 *   });
 *
 *   await tx.workspaceMember.create({
 *     data: {
 *       workspaceId: workspace.id,
 *       userId,
 *       role: 'OWNER'
 *     }
 *   });
 *
 *   return workspace;
 * });
 * ```
 */
export async function withUserTransaction<T>(
  userId: string,
  fn: (tx: PrismaClient) => Promise<T>
): Promise<T> {
  const prismaClient = getPrismaInstance();
  return prismaClient.$transaction(async (tx) => {
    await setCurrentUser(userId, tx as PrismaClient);
    return fn(tx as PrismaClient);
  });
}

/**
 * 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<T>(
  userId: string,
  workspaceId: string,
  fn: (tx: PrismaClient) => Promise<T>
): Promise<T> {
  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.
 *
 * @param handler - The handler function that requires user context
 * @returns A new function that sets user context before calling the handler
 *
 * @example
 * ```typescript
 * // In a tRPC procedure
 * export const getTasks = withAuth(async ({ ctx, input }) => {
 *   return prisma.task.findMany({
 *     where: { workspaceId: input.workspaceId }
 *   });
 * });
 * ```
 */
export function withAuth<TArgs extends { ctx: { userId: string } }, TResult>(
  handler: (args: TArgs) => Promise<TResult>
) {
  return async (args: TArgs): Promise<TResult> => {
    return withUserContext(args.ctx.userId, () => handler(args));
  };
}

/**
 * Verifies that a user has access to a specific workspace.
 * This is an additional application-level check on top of RLS.
 *
 * @param userId - The UUID of the user
 * @param workspaceId - The UUID of the workspace
 * @returns True if the user is a member of the workspace
 *
 * @example
 * ```typescript
 * if (!await verifyWorkspaceAccess(userId, workspaceId)) {
 *   throw new Error('Access denied');
 * }
 * ```
 */
export async function verifyWorkspaceAccess(userId: string, workspaceId: string): Promise<boolean> {
  return withUserContext(userId, async (tx) => {
    const member = await tx.workspaceMember.findUnique({
      where: {
        workspaceId_userId: {
          workspaceId,
          userId,
        },
      },
    });
    return member !== null;
  });
}

/**
 * Gets all workspaces accessible by a user.
 * Uses RLS to automatically filter to authorized workspaces.
 *
 * @param userId - The UUID of the user
 * @returns Array of workspaces the user can access
 *
 * @example
 * ```typescript
 * const workspaces = await getUserWorkspaces(userId);
 * ```
 */
export async function getUserWorkspaces(userId: string) {
  return withUserContext(userId, async (tx) => {
    return tx.workspace.findMany({
      include: {
        members: {
          where: { userId },
          select: { role: true },
        },
      },
    });
  });
}

/**
 * Type guard to check if a user has admin access to a workspace.
 *
 * @param userId - The UUID of the user
 * @param workspaceId - The UUID of the workspace
 * @returns True if the user is an OWNER or ADMIN
 */
export async function isWorkspaceAdmin(userId: string, workspaceId: string): Promise<boolean> {
  return withUserContext(userId, async (tx) => {
    const member = await tx.workspaceMember.findUnique({
      where: {
        workspaceId_userId: {
          workspaceId,
          userId,
        },
      },
    });
    return member?.role === "OWNER" || member?.role === "ADMIN";
  });
}

/**
 * Executes a query without RLS restrictions.
 * ⚠️ USE WITH EXTREME CAUTION - Only for system-level operations!
 *
 * @param fn - The function to execute without RLS
 * @returns The result of the function
 *
 * @example
 * ```typescript
 * // Only use for system operations like migrations or admin cleanup
 * const allTasks = await withoutRLS(async () => {
 *   return prisma.task.findMany();
 * });
 * ```
 */
export async function withoutRLS<T>(fn: (client: PrismaClient) => Promise<T>): Promise<T> {
  const prismaClient = getPrismaInstance();
  return prismaClient.$transaction(async (tx) => {
    await clearCurrentUser(tx as PrismaClient);
    return fn(tx as PrismaClient);
  });
}

/**
 * Middleware factory for tRPC that automatically sets user context.
 *
 * @example
 * ```typescript
 * const authMiddleware = createAuthMiddleware();
 *
 * const protectedProcedure = publicProcedure.use(authMiddleware);
 * ```
 */
export function createAuthMiddleware(client: PrismaClient) {
  return async function authMiddleware(opts: {
    ctx: { userId?: string };
    next: () => Promise<unknown>;
  }): Promise<unknown> {
    const { userId } = opts.ctx;
    if (!userId) {
      throw new Error("User not authenticated");
    }

    // SEC-API-27: SET LOCAL must be called inside a transaction boundary.
    // Without a transaction, SET LOCAL behaves as a session-level SET,
    // which can leak RLS context to other requests via connection pooling.
    return client.$transaction(async (tx) => {
      await setCurrentUser(userId, tx as PrismaClient);
      return opts.next();
    });
  };
}

export default {
  setCurrentUser,
  setCurrentWorkspace,
  setWorkspaceContext,
  clearCurrentUser,
  clearWorkspaceContext,
  withUserContext,
  withUserTransaction,
  withWorkspaceContext,
  withAuth,
  verifyWorkspaceAccess,
  getUserWorkspaces,
  isWorkspaceAdmin,
  withoutRLS,
  createAuthMiddleware,
};