From bc8ebb921319123db099cb7bc934e375d31189ef Mon Sep 17 00:00:00 2001 From: Keith Williams Date: Wed, 11 Feb 2026 11:33:11 -0300 Subject: [PATCH] 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> --- .../bookings/get.handler.integration-test.ts | 316 ++++++++++++++++++ .../viewer/bookings/get.handler.test.ts | 94 +++++- .../routers/viewer/bookings/get.handler.ts | 9 +- 3 files changed, 403 insertions(+), 16 deletions(-) create mode 100644 packages/trpc/server/routers/viewer/bookings/get.handler.integration-test.ts diff --git a/packages/trpc/server/routers/viewer/bookings/get.handler.integration-test.ts b/packages/trpc/server/routers/viewer/bookings/get.handler.integration-test.ts new file mode 100644 index 0000000000..fc0cc4b14e --- /dev/null +++ b/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); + }); +}); diff --git a/packages/trpc/server/routers/viewer/bookings/get.handler.test.ts b/packages/trpc/server/routers/viewer/bookings/get.handler.test.ts index 28bd067fee..fe12d66efa 100644 --- a/packages/trpc/server/routers/viewer/bookings/get.handler.test.ts +++ b/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; + _mockQueryBuilder: mockQueryBuilder, + } as unknown as Kysely & { _mockQueryBuilder: typeof mockQueryBuilder }; }; - let mockKysely: Kysely; + let mockKysely: ReturnType; 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, 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, 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, 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, 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, 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, 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, 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, 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, + 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, + 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, + 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(); + }); + }); }); diff --git a/packages/trpc/server/routers/viewer/bookings/get.handler.ts b/packages/trpc/server/routers/viewer/bookings/get.handler.ts index 77c30f2c23..4c64b466bb 100644 --- a/packages/trpc/server/routers/viewer/bookings/get.handler.ts +++ b/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 );