feat: add fallbackRoles parameter to getTeamIdsWithPermissions method (#24042)

* feat: add fallbackRoles parameter to getTeamIdsWithPermissions method

- Add fallbackRoles parameter to method signature in interface and implementation
- Implement second query for teams without PBAC where user has fallback roles
- Combine and deduplicate results from both PBAC-enabled and fallback role teams
- Supports fallback to role-based permissions when PBAC is disabled
- Fix linting issue in getUserMemberships method by replacing include with select

Co-Authored-By: eunjae@cal.com <hey@eunjae.dev>

* fix usages

* revert some change

* fix unit tests

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

packages/features/pbac/domain/repositories/IPermissionRepository.ts

@@ -1,3 +1,5 @@
+import type { MembershipRole } from "@calcom/prisma/enums";
+
 import type { TeamPermissions } from "../models/Permission";
 import type { PermissionString, Resource, CrudAction, CustomAction } from "../types/permission-registry";
 
@@ -62,10 +64,18 @@ export interface IPermissionRepository {
   /**
    * Gets all team IDs where the user has a specific permission
    */
-  getTeamIdsWithPermission(userId: number, permission: PermissionString): Promise<number[]>;
+  getTeamIdsWithPermission(params: {
+    userId: number;
+    permission: PermissionString;
+    fallbackRoles: MembershipRole[];
+  }): Promise<number[]>;
 
   /**
    * Gets all team IDs where the user has all of the specified permissions
    */
-  getTeamIdsWithPermissions(userId: number, permissions: PermissionString[]): Promise<number[]>;
+  getTeamIdsWithPermissions(params: {
+    userId: number;
+    permissions: PermissionString[];
+    fallbackRoles: MembershipRole[];
+  }): Promise<number[]>;
 }

packages/features/pbac/infrastructure/repositories/PermissionRepository.ts

@@ -1,5 +1,6 @@
 import db from "@calcom/prisma";
 import type { PrismaClient as PrismaClientWithExtensions } from "@calcom/prisma";
+import type { MembershipRole } from "@calcom/prisma/enums";
 
 import { PermissionMapper } from "../../domain/mappers/PermissionMapper";
 import type { TeamPermissions } from "../../domain/models/Permission";
@@ -12,6 +13,7 @@ import {
 } from "../../domain/types/permission-registry";
 
 export class PermissionRepository implements IPermissionRepository {
+  private readonly PBAC_FEATURE_FLAG = "pbac" as const;
   private client: PrismaClientWithExtensions;
 
   constructor(client: PrismaClientWithExtensions = db) {
@@ -210,11 +212,27 @@ export class PermissionRepository implements IPermissionRepository {
     return permissions.map((p) => p.action as CrudAction | CustomAction);
   }
 
-  async getTeamIdsWithPermission(userId: number, permission: PermissionString): Promise<number[]> {
-    return this.getTeamIdsWithPermissions(userId, [permission]);
+  async getTeamIdsWithPermission({
+    userId,
+    permission,
+    fallbackRoles,
+  }: {
+    userId: number;
+    permission: PermissionString;
+    fallbackRoles: MembershipRole[];
+  }): Promise<number[]> {
+    return this.getTeamIdsWithPermissions({ userId, permissions: [permission], fallbackRoles });
   }
 
-  async getTeamIdsWithPermissions(userId: number, permissions: PermissionString[]): Promise<number[]> {
+  async getTeamIdsWithPermissions({
+    userId,
+    permissions,
+    fallbackRoles,
+  }: {
+    userId: number;
+    permissions: PermissionString[];
+    fallbackRoles: MembershipRole[];
+  }): Promise<number[]> {
     // Validate that permissions array is not empty to prevent privilege escalation
     if (permissions.length === 0) {
       return [];
@@ -225,7 +243,7 @@ export class PermissionRepository implements IPermissionRepository {
       return { resource, action };
     });
 
-    const teamsWithPermission = await this.client.$queryRaw<{ teamId: number }[]>`
+    const teamsWithPermissionPromise = this.client.$queryRaw<{ teamId: number }[]>`
       SELECT DISTINCT m."teamId"
       FROM "Membership" m
       INNER JOIN "Role" r ON m."customRoleId" = r.id
@@ -249,6 +267,26 @@ export class PermissionRepository implements IPermissionRepository {
         ) = ${permissions.length}
     `;
 
-    return teamsWithPermission.map((team) => team.teamId);
+    const teamsWithFallbackRolesPromise = this.client.$queryRaw<{ teamId: number }[]>`
+      SELECT DISTINCT m."teamId"
+      FROM "Membership" m
+      INNER JOIN "Team" t ON m."teamId" = t.id
+      LEFT JOIN "TeamFeatures" f ON t.id = f."teamId" AND f."featureId" = ${this.PBAC_FEATURE_FLAG}
+      WHERE m."userId" = ${userId}
+        AND m."accepted" = true
+        AND m."role"::text = ANY(${fallbackRoles})
+        AND f."teamId" IS NULL
+    `;
+
+    const [teamsWithPermission, teamsWithFallbackRoles] = await Promise.all([
+      teamsWithPermissionPromise,
+      teamsWithFallbackRolesPromise,
+    ]);
+
+    const pbacTeamIds = teamsWithPermission.map((team) => team.teamId);
+    const fallbackTeamIds = teamsWithFallbackRoles.map((team) => team.teamId);
+
+    const allTeamIds = Array.from(new Set([...pbacTeamIds, ...fallbackTeamIds]));
+    return allTeamIds;
   }
 }

packages/features/pbac/services/__tests__/permission-check.service.test.ts

@@ -327,10 +327,18 @@ describe("PermissionCheckService", () => {
       const expectedTeamIds = [1, 2, 3];
       mockRepository.getTeamIdsWithPermission.mockResolvedValueOnce(expectedTeamIds);
 
-      const result = await service.getTeamIdsWithPermission(1, "eventType.read");
+      const result = await service.getTeamIdsWithPermission({
+        userId: 1,
+        permission: "eventType.read",
+        fallbackRoles: [],
+      });
 
       expect(result).toEqual(expectedTeamIds);
-      expect(mockRepository.getTeamIdsWithPermission).toHaveBeenCalledWith(1, "eventType.read");
+      expect(mockRepository.getTeamIdsWithPermission).toHaveBeenCalledWith({
+        userId: 1,
+        permission: "eventType.read",
+        fallbackRoles: [],
+      });
     });
 
     it("should return empty array when permission validation fails", async () => {
@@ -339,7 +347,11 @@ describe("PermissionCheckService", () => {
         error: "Invalid permissions",
       });
 
-      const result = await service.getTeamIdsWithPermission(1, "eventType.read");
+      const result = await service.getTeamIdsWithPermission({
+        userId: 1,
+        permission: "eventType.read",
+        fallbackRoles: [],
+      });
 
       expect(result).toEqual([]);
       expect(mockRepository.getTeamIdsWithPermission).not.toHaveBeenCalled();
@@ -348,10 +360,18 @@ describe("PermissionCheckService", () => {
     it("should return empty array and log error when repository throws", async () => {
       mockRepository.getTeamIdsWithPermission.mockRejectedValueOnce(new Error("Database error"));
 
-      const result = await service.getTeamIdsWithPermission(1, "eventType.read");
+      const result = await service.getTeamIdsWithPermission({
+        userId: 1,
+        permission: "eventType.read",
+        fallbackRoles: [],
+      });
 
       expect(result).toEqual([]);
-      expect(mockRepository.getTeamIdsWithPermission).toHaveBeenCalledWith(1, "eventType.read");
+      expect(mockRepository.getTeamIdsWithPermission).toHaveBeenCalledWith({
+        userId: 1,
+        permission: "eventType.read",
+        fallbackRoles: [],
+      });
     });
   });
 
@@ -361,10 +381,14 @@ describe("PermissionCheckService", () => {
       const permissions: PermissionString[] = ["eventType.read", "eventType.create"];
       mockRepository.getTeamIdsWithPermissions.mockResolvedValueOnce(expectedTeamIds);
 
-      const result = await service.getTeamIdsWithPermissions(1, permissions);
+      const result = await service.getTeamIdsWithPermissions({ userId: 1, permissions, fallbackRoles: [] });
 
       expect(result).toEqual(expectedTeamIds);
-      expect(mockRepository.getTeamIdsWithPermissions).toHaveBeenCalledWith(1, permissions);
+      expect(mockRepository.getTeamIdsWithPermissions).toHaveBeenCalledWith({
+        userId: 1,
+        permissions,
+        fallbackRoles: [],
+      });
     });
 
     it("should return empty array when permissions validation fails", async () => {
@@ -373,7 +397,11 @@ describe("PermissionCheckService", () => {
         error: "Invalid permissions",
       });
 
-      const result = await service.getTeamIdsWithPermissions(1, ["eventType.read"] as PermissionString[]);
+      const result = await service.getTeamIdsWithPermissions({
+        userId: 1,
+        permissions: ["eventType.read"] as PermissionString[],
+        fallbackRoles: [],
+      });
 
       expect(result).toEqual([]);
       expect(mockRepository.getTeamIdsWithPermissions).not.toHaveBeenCalled();
@@ -382,20 +410,32 @@ describe("PermissionCheckService", () => {
     it("should return empty array when permissions array is empty", async () => {
       mockRepository.getTeamIdsWithPermissions.mockResolvedValueOnce([]);
 
-      const result = await service.getTeamIdsWithPermissions(1, []);
+      const result = await service.getTeamIdsWithPermissions({
+        userId: 1,
+        permissions: [],
+        fallbackRoles: [],
+      });
 
       expect(result).toEqual([]);
-      expect(mockRepository.getTeamIdsWithPermissions).toHaveBeenCalledWith(1, []);
+      expect(mockRepository.getTeamIdsWithPermissions).toHaveBeenCalledWith({
+        userId: 1,
+        permissions: [],
+        fallbackRoles: [],
+      });
     });
 
     it("should return empty array and log error when repository throws", async () => {
       const permissions: PermissionString[] = ["eventType.read", "eventType.create"];
       mockRepository.getTeamIdsWithPermissions.mockRejectedValueOnce(new Error("Database error"));
 
-      const result = await service.getTeamIdsWithPermissions(1, permissions);
+      const result = await service.getTeamIdsWithPermissions({ userId: 1, permissions, fallbackRoles: [] });
 
       expect(result).toEqual([]);
-      expect(mockRepository.getTeamIdsWithPermissions).toHaveBeenCalledWith(1, permissions);
+      expect(mockRepository.getTeamIdsWithPermissions).toHaveBeenCalledWith({
+        userId: 1,
+        permissions,
+        fallbackRoles: [],
+      });
     });
   });
 

packages/features/pbac/services/permission-check.service.ts

@@ -255,7 +255,15 @@ export class PermissionCheckService {
   /**
    * Gets all team IDs where the user has a specific permission
    */
-  async getTeamIdsWithPermission(userId: number, permission: PermissionString): Promise<number[]> {
+  async getTeamIdsWithPermission({
+    userId,
+    permission,
+    fallbackRoles,
+  }: {
+    userId: number;
+    permission: PermissionString;
+    fallbackRoles: MembershipRole[];
+  }): Promise<number[]> {
     try {
       const validationResult = this.permissionService.validatePermission(permission);
       if (!validationResult.isValid) {
@@ -263,7 +271,7 @@ export class PermissionCheckService {
         return [];
       }
 
-      return await this.repository.getTeamIdsWithPermission(userId, permission);
+      return await this.repository.getTeamIdsWithPermission({ userId, permission, fallbackRoles });
     } catch (error) {
       this.logger.error(error);
       return [];
@@ -273,7 +281,15 @@ export class PermissionCheckService {
   /**
    * Gets all team IDs where the user has all of the specified permissions
    */
-  async getTeamIdsWithPermissions(userId: number, permissions: PermissionString[]): Promise<number[]> {
+  async getTeamIdsWithPermissions({
+    userId,
+    permissions,
+    fallbackRoles,
+  }: {
+    userId: number;
+    permissions: PermissionString[];
+    fallbackRoles: MembershipRole[];
+  }): Promise<number[]> {
     try {
       const validationResult = this.permissionService.validatePermissions(permissions);
       if (!validationResult.isValid) {
@@ -281,7 +297,7 @@ export class PermissionCheckService {
         return [];
       }
 
-      return await this.repository.getTeamIdsWithPermissions(userId, permissions);
+      return await this.repository.getTeamIdsWithPermissions({ userId, permissions, fallbackRoles });
     } catch (error) {
       this.logger.error(error);
       return [];

packages/trpc/server/routers/viewer/teams/listSimpleMembers.handler.ts

@@ -4,6 +4,7 @@
  */
 import { PermissionCheckService } from "@calcom/features/pbac/services/permission-check.service";
 import type { PrismaClient } from "@calcom/prisma";
+import { MembershipRole } from "@calcom/prisma/enums";
 import type { TrpcSessionUser } from "@calcom/trpc/server/types";
 
 type ListSimpleMembersOptions = {
@@ -23,7 +24,11 @@ export const listSimpleMembers = async ({ ctx }: ListSimpleMembersOptions) => {
   }
 
   const permissionCheckService = new PermissionCheckService();
-  const teamsToQuery = await permissionCheckService.getTeamIdsWithPermission(ctx.user.id, "team.listMembers");
+  const teamsToQuery = await permissionCheckService.getTeamIdsWithPermission({
+    userId: ctx.user.id,
+    permission: "team.listMembers",
+    fallbackRoles: [MembershipRole.OWNER, MembershipRole.ADMIN],
+  });
 
   if (!teamsToQuery.length) {
     return [];

packages/trpc/server/routers/viewer/teams/removeMember/PBACRemoveMemberService.ts

@@ -2,6 +2,7 @@ import * as teamQueries from "@calcom/features/ee/teams/lib/queries";
 import { PermissionMapper } from "@calcom/features/pbac/domain/mappers/PermissionMapper";
 import { Resource, CustomAction } from "@calcom/features/pbac/domain/types/permission-registry";
 import { PermissionCheckService } from "@calcom/features/pbac/services/permission-check.service";
+import { MembershipRole } from "@calcom/prisma/enums";
 
 import { TRPCError } from "@trpc/server";
 
@@ -20,10 +21,11 @@ export class PBACRemoveMemberService extends BaseRemoveMemberService {
       action: CustomAction.Remove,
     });
 
-    const teamsWithPermission = await this.permissionService.getTeamIdsWithPermission(
+    const teamsWithPermission = await this.permissionService.getTeamIdsWithPermission({
       userId,
-      removePermission
-    );
+      permission: removePermission,
+      fallbackRoles: [MembershipRole.OWNER, MembershipRole.ADMIN],
+    });
 
     // Convert to Set for O(1) lookup
     const teamsWithPermissionSet = new Set(teamsWithPermission);

packages/trpc/server/routers/viewer/teams/removeMember/__tests__/PBACRemoveMemberService.test.ts

@@ -68,10 +68,11 @@ describe("PBACRemoveMemberService", () => {
           action: CustomAction.Remove,
         });
 
-        expect(mockPermissionCheckService.getTeamIdsWithPermission).toHaveBeenCalledWith(
+        expect(mockPermissionCheckService.getTeamIdsWithPermission).toHaveBeenCalledWith({
           userId,
-          removePermission
-        );
+          permission: removePermission,
+          fallbackRoles: [MembershipRole.OWNER, MembershipRole.ADMIN],
+        });
       });
 
       it("should check organization.remove permission for org context", async () => {