perf: use UNION ALL with DISTINCT for bookings query optimization (#27841)

* perf: use UNION ALL with DISTINCT for bookings query optimization

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

* test: add unionAll and distinct to kysely test mocks

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

* test: add integration tests for getBookings query

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

* test: add totalCount deduplication test for multi-branch union scenario

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

* test: fix totalCount assertion to be exact instead of gte

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

* Apply suggestion from @keithwillcode

* perf: restore COUNT(DISTINCT id) for totalCount query

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

* Add tests: packages/trpc/server/routers/viewer/bookings/get.handler.integration-test.ts

Generated by Paragon from proposal for PR #27841

* Add tests: packages/trpc/server/routers/viewer/bookings/get.handler.test.ts

Generated by Paragon from proposal for PR #27841

---------

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

packages/trpc/server/routers/viewer/bookings/get.handler.integration-test.ts

@@ -0,0 +1,316 @@
+import prisma from "@calcom/prisma";
+import kysely from "@calcom/kysely";
+import type { Booking, EventType, Team, User } from "@calcom/prisma/client";
+import { BookingStatus, MembershipRole } from "@calcom/prisma/enums";
+import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import { getBookings } from "./get.handler";
+
+let user1: User;
+let user2: User;
+let team1: Team;
+let eventType1: EventType;
+let booking1: Booking;
+let booking2: Booking;
+let booking3: Booking;
+let booking4: Booking;
+let teamEventType: EventType;
+
+const timestamp = Date.now();
+
+describe("getBookings - integration", () => {
+  beforeAll(async () => {
+    user1 = await prisma.user.create({
+      data: {
+        username: `getbookings-user1-${timestamp}`,
+        email: `getbookings-user1-${timestamp}@example.com`,
+        name: "GetBookings User 1",
+      },
+    });
+
+    user2 = await prisma.user.create({
+      data: {
+        username: `getbookings-user2-${timestamp}`,
+        email: `getbookings-user2-${timestamp}@example.com`,
+        name: "GetBookings User 2",
+      },
+    });
+
+    team1 = await prisma.team.create({
+      data: {
+        name: `GetBookings Team ${timestamp}`,
+        slug: `getbookings-team-${timestamp}`,
+        members: {
+          create: {
+            userId: user1.id,
+            role: MembershipRole.ADMIN,
+            accepted: true,
+          },
+        },
+      },
+    });
+
+    await prisma.membership.create({
+      data: {
+        userId: user2.id,
+        teamId: team1.id,
+        role: MembershipRole.MEMBER,
+        accepted: true,
+      },
+    });
+
+    eventType1 = await prisma.eventType.create({
+      data: {
+        title: `GetBookings Event ${timestamp}`,
+        slug: `getbookings-event-${timestamp}`,
+        length: 30,
+        userId: user1.id,
+      },
+    });
+
+    const futureDate = new Date(Date.now() + 7 * 24 * 60 * 60 * 1000);
+    const futureDate2 = new Date(Date.now() + 8 * 24 * 60 * 60 * 1000);
+    const futureDate3 = new Date(Date.now() + 9 * 24 * 60 * 60 * 1000);
+
+    booking1 = await prisma.booking.create({
+      data: {
+        uid: `getbookings-booking1-${timestamp}`,
+        title: "Booking 1",
+        startTime: futureDate,
+        endTime: new Date(futureDate.getTime() + 30 * 60 * 1000),
+        userId: user1.id,
+        eventTypeId: eventType1.id,
+        status: BookingStatus.ACCEPTED,
+        attendees: {
+          create: {
+            email: user2.email,
+            name: user2.name ?? "User 2",
+            timeZone: "UTC",
+          },
+        },
+      },
+    });
+
+    booking2 = await prisma.booking.create({
+      data: {
+        uid: `getbookings-booking2-${timestamp}`,
+        title: "Booking 2",
+        startTime: futureDate2,
+        endTime: new Date(futureDate2.getTime() + 30 * 60 * 1000),
+        userId: user1.id,
+        eventTypeId: eventType1.id,
+        status: BookingStatus.ACCEPTED,
+        attendees: {
+          create: {
+            email: `getbookings-external-${timestamp}@example.com`,
+            name: "External Attendee",
+            timeZone: "UTC",
+          },
+        },
+      },
+    });
+
+    booking3 = await prisma.booking.create({
+      data: {
+        uid: `getbookings-booking3-${timestamp}`,
+        title: "Booking 3",
+        startTime: futureDate3,
+        endTime: new Date(futureDate3.getTime() + 30 * 60 * 1000),
+        userId: user2.id,
+        eventTypeId: eventType1.id,
+        status: BookingStatus.ACCEPTED,
+        attendees: {
+          create: {
+            email: user1.email,
+            name: user1.name ?? "User 1",
+            timeZone: "UTC",
+          },
+        },
+      },
+    });
+
+    teamEventType = await prisma.eventType.create({
+      data: {
+        title: `GetBookings Team Event ${timestamp}`,
+        slug: `getbookings-team-event-${timestamp}`,
+        length: 30,
+        teamId: team1.id,
+      },
+    });
+
+    const futureDate4 = new Date(Date.now() + 10 * 24 * 60 * 60 * 1000);
+    booking4 = await prisma.booking.create({
+      data: {
+        uid: `getbookings-booking4-${timestamp}`,
+        title: "Team Booking - multi branch",
+        startTime: futureDate4,
+        endTime: new Date(futureDate4.getTime() + 30 * 60 * 1000),
+        userId: user1.id,
+        eventTypeId: teamEventType.id,
+        status: BookingStatus.ACCEPTED,
+        attendees: {
+          create: {
+            email: user1.email,
+            name: user1.name ?? "User 1",
+            timeZone: "UTC",
+          },
+        },
+      },
+    });
+  });
+
+  afterAll(async () => {
+    try {
+      const bookingIds = [booking1?.id, booking2?.id, booking3?.id, booking4?.id].filter(Boolean);
+      if (bookingIds.length > 0) {
+        await prisma.attendee.deleteMany({ where: { bookingId: { in: bookingIds } } });
+        await prisma.booking.deleteMany({ where: { id: { in: bookingIds } } });
+      }
+      const eventTypeIds = [eventType1?.id, teamEventType?.id].filter(Boolean);
+      if (eventTypeIds.length > 0) {
+        await prisma.eventType.deleteMany({ where: { id: { in: eventTypeIds } } });
+      }
+      const teamIds = [team1?.id].filter(Boolean);
+      if (teamIds.length > 0) {
+        await prisma.team.deleteMany({ where: { id: { in: teamIds } } });
+      }
+      const userIds = [user1?.id, user2?.id].filter(Boolean);
+      if (userIds.length > 0) {
+        await prisma.user.deleteMany({ where: { id: { in: userIds } } });
+      }
+    } catch (error) {
+      console.warn("Test cleanup failed:", error);
+    }
+  });
+
+  it("should return bookings created by the user", async () => {
+    const result = await getBookings({
+      user: { id: user1.id, email: user1.email, orgId: null },
+      prisma,
+      kysely,
+      bookingListingByStatus: ["upcoming"],
+      filters: {},
+      take: 50,
+      skip: 0,
+    });
+
+    const bookingIds = result.bookings.map((b) => b.id);
+    expect(bookingIds).toContain(booking1.id);
+    expect(bookingIds).toContain(booking2.id);
+  });
+
+  it("should return bookings where user is an attendee", async () => {
+    const result = await getBookings({
+      user: { id: user1.id, email: user1.email, orgId: null },
+      prisma,
+      kysely,
+      bookingListingByStatus: ["upcoming"],
+      filters: {},
+      take: 50,
+      skip: 0,
+    });
+
+    const bookingIds = result.bookings.map((b) => b.id);
+    expect(bookingIds).toContain(booking3.id);
+  });
+
+  it("should not return duplicate bookings when user is both creator and attendee matches", async () => {
+    const result = await getBookings({
+      user: { id: user1.id, email: user1.email, orgId: null },
+      prisma,
+      kysely,
+      bookingListingByStatus: ["upcoming"],
+      filters: {},
+      take: 50,
+      skip: 0,
+    });
+
+    const testBookingIds = [booking1.id, booking2.id, booking3.id];
+    const relevantBookingIds = result.bookings
+      .map((b) => b.id)
+      .filter((id) => testBookingIds.includes(id));
+
+    const uniqueIds = new Set(relevantBookingIds);
+    expect(relevantBookingIds.length).toBe(uniqueIds.size);
+  });
+
+  it("should return correct totalCount without duplicates", async () => {
+    const resultUser1 = await getBookings({
+      user: { id: user1.id, email: user1.email, orgId: null },
+      prisma,
+      kysely,
+      bookingListingByStatus: ["upcoming"],
+      filters: {},
+      take: 50,
+      skip: 0,
+    });
+
+    const expectedBookingIds = [booking1.id, booking2.id, booking3.id, booking4.id];
+    const returnedIds = resultUser1.bookings.map((b) => b.id);
+    for (const id of expectedBookingIds) {
+      expect(returnedIds).toContain(id);
+    }
+    expect(resultUser1.totalCount).toBe(resultUser1.bookings.length);
+  });
+
+  it("should count booking4 exactly once in totalCount even though it matches multiple union branches", async () => {
+    const result = await getBookings({
+      user: { id: user1.id, email: user1.email, orgId: null },
+      prisma,
+      kysely,
+      bookingListingByStatus: ["upcoming"],
+      filters: {},
+      take: 50,
+      skip: 0,
+    });
+
+    const booking4Occurrences = result.bookings.filter((b) => b.id === booking4.id);
+    expect(booking4Occurrences).toHaveLength(1);
+    expect(result.totalCount).toBe(result.bookings.length);
+  });
+
+  it("should respect pagination with correct ordering", async () => {
+    const page1 = await getBookings({
+      user: { id: user1.id, email: user1.email, orgId: null },
+      prisma,
+      kysely,
+      bookingListingByStatus: ["upcoming"],
+      filters: {},
+      take: 1,
+      skip: 0,
+    });
+
+    const page2 = await getBookings({
+      user: { id: user1.id, email: user1.email, orgId: null },
+      prisma,
+      kysely,
+      bookingListingByStatus: ["upcoming"],
+      filters: {},
+      take: 1,
+      skip: 1,
+    });
+
+    expect(page1.bookings.length).toBe(1);
+    expect(page2.bookings.length).toBe(1);
+    expect(page1.bookings[0].id).not.toBe(page2.bookings[0].id);
+
+    const date1 = new Date(page1.bookings[0].startTime);
+    const date2 = new Date(page2.bookings[0].startTime);
+    expect(date1.getTime()).toBeLessThanOrEqual(date2.getTime());
+  });
+
+  it("should return bookings for user2 including those where user2 is an attendee", async () => {
+    const result = await getBookings({
+      user: { id: user2.id, email: user2.email, orgId: null },
+      prisma,
+      kysely,
+      bookingListingByStatus: ["upcoming"],
+      filters: {},
+      take: 50,
+      skip: 0,
+    });
+
+    const bookingIds = result.bookings.map((b) => b.id);
+    expect(bookingIds).toContain(booking3.id);
+    expect(bookingIds).toContain(booking1.id);
+  });
+});

packages/trpc/server/routers/viewer/bookings/get.handler.test.ts

@@ -137,6 +137,8 @@ describe("getBookings - PBAC Permission Checks", () => {
       where: vi.fn(() => mockQueryBuilder),
       innerJoin: vi.fn(() => mockQueryBuilder),
       union: vi.fn(() => mockQueryBuilder),
+      unionAll: vi.fn(() => mockQueryBuilder),
+      distinct: vi.fn(() => mockQueryBuilder),
       as: vi.fn(() => mockQueryBuilder),
       $if: vi.fn(() => mockQueryBuilder),
       orderBy: vi.fn(() => mockQueryBuilder),
@@ -150,10 +152,11 @@ describe("getBookings - PBAC Permission Checks", () => {
     return {
       selectFrom: vi.fn(() => mockQueryBuilder),
       executeQuery: vi.fn().mockResolvedValue({ rows: [] }),
-    } as unknown as Kysely<DB>;
+      _mockQueryBuilder: mockQueryBuilder,
+    } as unknown as Kysely<DB> & { _mockQueryBuilder: typeof mockQueryBuilder };
   };
 
-  let mockKysely: Kysely<DB>;
+  let mockKysely: ReturnType<typeof createMockKysely>;
 
   const mockGetTeamIdsWithPermission = vi.fn();
 
@@ -178,7 +181,7 @@ describe("getBookings - PBAC Permission Checks", () => {
       await getBookings({
         user: mockUser,
         prisma: mockPrisma,
-        kysely: mockKysely,
+        kysely: mockKysely as unknown as Kysely<DB>,
         bookingListingByStatus: ["upcoming"],
         filters: {
           userIds: [2],
@@ -210,6 +213,7 @@ describe("getBookings - PBAC Permission Checks", () => {
         where: vi.fn(() => mockQueryBuilder),
         innerJoin: vi.fn(() => mockQueryBuilder),
         union: vi.fn(() => mockQueryBuilder),
+        unionAll: vi.fn(() => mockQueryBuilder),
         as: vi.fn(() => ({
           select: vi.fn(() => ({
             selectAll: vi.fn(() => ({
@@ -226,14 +230,14 @@ describe("getBookings - PBAC Permission Checks", () => {
         })),
       };
 
-      mockKysely.selectFrom = vi.fn(() => mockQueryBuilder as any);
-      mockKysely.executeQuery = vi.fn().mockResolvedValue({ rows: [] });
+      (mockKysely as any).selectFrom = vi.fn(() => mockQueryBuilder as any);
+      (mockKysely as any).executeQuery = vi.fn().mockResolvedValue({ rows: [] });
 
       await expect(
         getBookings({
           user: mockUser,
           prisma: mockPrisma,
-          kysely: mockKysely,
+          kysely: mockKysely as unknown as Kysely<DB>,
           bookingListingByStatus: ["upcoming"],
           filters: {
             userIds: [4], // User 4 is not in team 1
@@ -267,7 +271,7 @@ describe("getBookings - PBAC Permission Checks", () => {
         getBookings({
           user: mockUser,
           prisma: mockPrisma,
-          kysely: mockKysely,
+          kysely: mockKysely as unknown as Kysely<DB>,
           bookingListingByStatus: ["upcoming"],
           filters: {
             userIds: [1], // Own user ID
@@ -289,7 +293,7 @@ describe("getBookings - PBAC Permission Checks", () => {
         getBookings({
           user: mockUser,
           prisma: mockPrisma,
-          kysely: mockKysely,
+          kysely: mockKysely as unknown as Kysely<DB>,
           bookingListingByStatus: ["upcoming"],
           filters: {
             userIds: [2],
@@ -323,7 +327,7 @@ describe("getBookings - PBAC Permission Checks", () => {
         getBookings({
           user: mockUser,
           prisma: mockPrisma,
-          kysely: mockKysely,
+          kysely: mockKysely as unknown as Kysely<DB>,
           bookingListingByStatus: ["upcoming"],
           filters: {
             userIds: [2, 3],
@@ -344,7 +348,7 @@ describe("getBookings - PBAC Permission Checks", () => {
       await getBookings({
         user: mockUser,
         prisma: mockPrisma,
-        kysely: mockKysely,
+        kysely: mockKysely as unknown as Kysely<DB>,
         bookingListingByStatus: ["upcoming"],
         filters: {},
         take: 10,
@@ -354,7 +358,7 @@ describe("getBookings - PBAC Permission Checks", () => {
       // PERFORMANCE: With the optimization, we no longer call $queryRaw to fetch event type IDs.
       // Instead, we use subqueries directly in the SQL to avoid materializing large arrays.
       // The query is built using Kysely and executed via executeQuery.
-      expect(mockKysely.executeQuery).toHaveBeenCalled();
+      expect((mockKysely as any).executeQuery).toHaveBeenCalled();
     });
   });
 
@@ -371,7 +375,7 @@ describe("getBookings - PBAC Permission Checks", () => {
       await getBookings({
         user: mockUser,
         prisma: mockPrisma,
-        kysely: mockKysely,
+        kysely: mockKysely as unknown as Kysely<DB>,
         bookingListingByStatus: ["upcoming"],
         filters: {
           userIds: [2], // When userIds filter is provided, we need to validate permissions
@@ -394,7 +398,7 @@ describe("getBookings - PBAC Permission Checks", () => {
       await getBookings({
         user: mockUser,
         prisma: mockPrisma,
-        kysely: mockKysely,
+        kysely: mockKysely as unknown as Kysely<DB>,
         bookingListingByStatus: ["upcoming"],
         filters: {}, // No userIds filter
         take: 10,
@@ -407,4 +411,68 @@ describe("getBookings - PBAC Permission Checks", () => {
       expect(mockPrisma.user.findMany).not.toHaveBeenCalled();
     });
   });
+
+  describe("UNION ALL with DISTINCT query optimization", () => {
+    it("should use unionAll instead of union for combining booking queries", async () => {
+      mockGetTeamIdsWithPermission.mockResolvedValue([1]);
+      mockPrisma.user.findMany = vi.fn().mockResolvedValue([]);
+      mockPrisma.booking.groupBy = vi.fn().mockResolvedValue([]);
+
+      await getBookings({
+        user: mockUser,
+        prisma: mockPrisma,
+        kysely: mockKysely as unknown as Kysely<DB>,
+        bookingListingByStatus: ["upcoming"],
+        filters: {},
+        take: 10,
+        skip: 0,
+      });
+
+      // Verify unionAll is used (UNION ALL) instead of union (UNION)
+      // This is the performance optimization: UNION ALL avoids an implicit DISTINCT
+      // at each union step, deferring deduplication to the outer SELECT DISTINCT
+      expect(mockKysely._mockQueryBuilder.unionAll).toHaveBeenCalled();
+    });
+
+    it("should apply DISTINCT on the outer select from union subquery", async () => {
+      mockGetTeamIdsWithPermission.mockResolvedValue([1]);
+      mockPrisma.user.findMany = vi.fn().mockResolvedValue([]);
+      mockPrisma.booking.groupBy = vi.fn().mockResolvedValue([]);
+
+      await getBookings({
+        user: mockUser,
+        prisma: mockPrisma,
+        kysely: mockKysely as unknown as Kysely<DB>,
+        bookingListingByStatus: ["upcoming"],
+        filters: {},
+        take: 10,
+        skip: 0,
+      });
+
+      // Verify distinct() is called on the outer query to deduplicate
+      // results from UNION ALL
+      expect(mockKysely._mockQueryBuilder.distinct).toHaveBeenCalled();
+    });
+
+    it("should use count with distinct for totalCount calculation", async () => {
+      mockGetTeamIdsWithPermission.mockResolvedValue([1]);
+      mockPrisma.user.findMany = vi.fn().mockResolvedValue([]);
+      mockPrisma.booking.groupBy = vi.fn().mockResolvedValue([]);
+
+      await getBookings({
+        user: mockUser,
+        prisma: mockPrisma,
+        kysely: mockKysely as unknown as Kysely<DB>,
+        bookingListingByStatus: ["upcoming"],
+        filters: {},
+        take: 10,
+        skip: 0,
+      });
+
+      // The count query uses fn.count("union_subquery.id").distinct()
+      // instead of fn.countAll() to ensure duplicates from UNION ALL
+      // are not counted multiple times
+      expect(mockKysely._mockQueryBuilder.executeTakeFirst).toHaveBeenCalled();
+    });
+  });
 });

packages/trpc/server/routers/viewer/bookings/get.handler.ts

@@ -132,7 +132,9 @@ export async function getBookings({
     getAttendeeEmailsFromUserIdsFilter(prisma, user.email, filters?.userIds),
     getEventTypeIdsFromEventTypeIdsFilter(prisma, filters?.eventTypeIds),
     // Only fetch accessible user IDs when we need to validate the userIds filter
-    needsUserIdsValidation ? getUserIdsFromTeamIds(prisma, teamIdsWithBookingPermission) : Promise.resolve([]),
+    needsUserIdsValidation
+      ? getUserIdsFromTeamIds(prisma, teamIdsWithBookingPermission)
+      : Promise.resolve([]),
   ]);
 
   const bookingQueries: { query: BookingsUnionQuery; tables: (keyof DB)[] }[] = [];
@@ -403,13 +405,14 @@ export async function getBookings({
   });
 
   const queryUnion = queriesWithFilters.reduce((acc, query) => {
-    return acc.union(query);
+    return acc.unionAll(query);
   });
 
   const orderBy = getOrderBy(bookingListingByStatus, sort);
 
   const getBookingsUnionCompiled = kysely
     .selectFrom(queryUnion.as("union_subquery"))
+    .distinct()
     .selectAll("union_subquery")
     .$if(Boolean(filters?.afterUpdatedDate), (eb) =>
       eb.where("union_subquery.updatedAt", ">=", dayjs.utc(filters.afterUpdatedDate).toDate())
@@ -436,7 +439,7 @@ export async function getBookings({
     (
       await kysely
         .selectFrom(queryUnion.as("union_subquery"))
-        .select(({ fn }) => fn.countAll().as("bookingCount"))
+        .select(({ fn }) => fn.count("union_subquery.id").distinct().as("bookingCount"))
         .executeTakeFirst()
     )?.bookingCount ?? 0
   );