From 33c8d51e089c6f06beb19c345a738f7fef78eeab Mon Sep 17 00:00:00 2001 From: Tom Moor Date: Sun, 2 Aug 2026 18:29:16 -0400 Subject: [PATCH] perf: Avoid membership joins in search (#13269) * perf: Avoid membership joins in search * feedback --- .../server/PostgresSearchProvider.test.ts | 60 ++++++ .../server/PostgresSearchProvider.ts | 199 +++++++----------- server/models/Document.test.ts | 131 ++++++++++++ server/models/Document.ts | 54 +++++ 4 files changed, 324 insertions(+), 120 deletions(-) diff --git a/plugins/search-postgres/server/PostgresSearchProvider.test.ts b/plugins/search-postgres/server/PostgresSearchProvider.test.ts index c65fa2d968..e500617a9e 100644 --- a/plugins/search-postgres/server/PostgresSearchProvider.test.ts +++ b/plugins/search-postgres/server/PostgresSearchProvider.test.ts @@ -763,6 +763,66 @@ describe("PostgresSearchProvider", () => { expect(results[0].ranking).toBeTruthy(); expect(results[0].document?.id).toBe(document.id); }); + + it("should return no results for a user with no collection or document access", async () => { + const team = await buildTeam(); + const user = await buildUser({ teamId: team.id }); + const otherUser = await buildUser({ teamId: team.id }); + const privateCollection = await buildCollection({ + teamId: team.id, + userId: otherUser.id, + permission: null, + }); + await buildDocument({ + teamId: team.id, + userId: otherUser.id, + collectionId: privateCollection.id, + title: "test", + }); + + const { results, total } = await provider.searchForUser(user, { + query: "test", + }); + + expect(results.length).toBe(0); + expect(total).toBe(0); + }); + + it("should include drafts shared with the user through a group", async () => { + const team = await buildTeam(); + const user = await buildUser({ teamId: team.id }); + const otherUser = await buildUser({ teamId: team.id }); + const draft = await buildDraftDocument({ + teamId: team.id, + userId: otherUser.id, + createdById: otherUser.id, + collectionId: null, + title: "group draft test", + }); + + const group = await buildGroup({ + teamId: team.id, + }); + await group.$add("user", user, { + through: { + createdById: otherUser.id, + }, + }); + await GroupMembership.create({ + createdById: otherUser.id, + groupId: group.id, + documentId: draft.id, + permission: DocumentPermission.Read, + }); + + const { results } = await provider.searchForUser(user, { + query: "group draft", + statusFilter: [StatusFilter.Draft], + }); + + expect(results.length).toBe(1); + expect(results[0].document?.id).toBe(draft.id); + }); }); describe("#searchTitlesForUser", () => { diff --git a/plugins/search-postgres/server/PostgresSearchProvider.ts b/plugins/search-postgres/server/PostgresSearchProvider.ts index 86f7a75130..b8e0df61d1 100644 --- a/plugins/search-postgres/server/PostgresSearchProvider.ts +++ b/plugins/search-postgres/server/PostgresSearchProvider.ts @@ -230,33 +230,35 @@ export default class PostgresSearchProvider extends BaseSearchProvider { }); try { - const resultsQuery = Document.unscoped().findAll({ + const results = (await Document.unscoped().findAll({ ...findOptions, where, limit, offset, - }) as unknown as Promise; - - const countQuery = Document.unscoped().count({ - // @ts-expect-error Types are incorrect for count - replacements: findOptions.replacements, - where, - }) as unknown as Promise; - const [results, count] = await Promise.all([resultsQuery, countQuery]); + })) as unknown as RankedDocument[]; // Final query to get associated document data - const documents = await Document.findAll({ - where: { - id: map(results, "id"), - teamId: team.id, - }, - include: [ - { - model: Collection, - as: "collection", + const [documents, count] = await Promise.all([ + Document.findAll({ + where: { + id: map(results, "id"), + teamId: team.id, }, - ], - }); + include: [ + { + model: Collection, + as: "collection", + }, + ], + }), + PostgresSearchProvider.countResults({ + results, + limit, + offset, + replacements: findOptions.replacements, + where, + }), + ]); return PostgresSearchProvider.buildResponse({ query, @@ -288,59 +290,16 @@ export default class PostgresSearchProvider extends BaseSearchProvider { }); } - const include = [ - { - association: "memberships", - where: { - userId: user.id, - }, - required: false, - separate: false, - }, - { - association: "groupMemberships", - required: false, - separate: false, - include: [ - { - association: "group", - required: true, - include: [ - { - association: "groupUsers", - required: true, - where: { - userId: user.id, - }, - }, - ], - }, - ], - }, - { - model: User, - as: "createdBy", - paranoid: false, - }, - { - model: User, - as: "updatedBy", - paranoid: false, - }, - ]; - return Document.withMembershipScope(user.id, { includeDrafts: true, }).findAll({ where, - subQuery: false, order: [ [ options.sort ?? SortFilter.UpdatedAt, options.direction ?? DirectionFilter.DESC, ], ], - include, offset, limit, }); @@ -389,55 +348,14 @@ export default class PostgresSearchProvider extends BaseSearchProvider { direction: options.direction, }); - const include = [ - { - association: "memberships", - where: { - userId: user.id, - }, - required: false, - separate: false, - }, - { - association: "groupMemberships", - required: false, - separate: false, - include: [ - { - association: "group", - required: true, - include: [ - { - association: "groupUsers", - required: true, - where: { - userId: user.id, - }, - }, - ], - }, - ], - }, - ]; - try { const results = (await Document.unscoped().findAll({ ...findOptions, - subQuery: false, - include, where, limit, offset, })) as unknown as RankedDocument[]; - const countQuery = Document.unscoped().count({ - // @ts-expect-error Types are incorrect for count - subQuery: false, - include, - replacements: findOptions.replacements, - where, - }) as unknown as Promise; - // Final query to get associated document data const [documents, count] = await Promise.all([ Document.withMembershipScope(user.id, { includeDrafts: true }).findAll({ @@ -446,9 +364,13 @@ export default class PostgresSearchProvider extends BaseSearchProvider { id: map(results, "id"), }, }), - results.length < limit && offset === 0 - ? Promise.resolve(results.length) - : countQuery, + PostgresSearchProvider.countResults({ + results, + limit, + offset, + replacements: findOptions.replacements, + where, + }), ]); return PostgresSearchProvider.buildResponse({ @@ -511,6 +433,35 @@ export default class PostgresSearchProvider extends BaseSearchProvider { // PostgreSQL metadata lives in the same row as the document } + /** + * Returns the total number of documents matching the search, avoiding a + * second query over the search conditions when the requested page was not + * filled and the total can be inferred. + */ + private static countResults({ + results, + limit, + offset, + replacements, + where, + }: { + results: RankedDocument[]; + limit: number; + offset: number; + replacements?: BindOrReplacements; + where: WhereOptions; + }): Promise { + if (results.length < limit && (offset === 0 || results.length > 0)) { + return Promise.resolve(offset + results.length); + } + + return Document.unscoped().count({ + // @ts-expect-error Types are incorrect for count + replacements, + where, + }) as unknown as Promise; + } + private static buildFindOptions({ query, sort, @@ -632,11 +583,23 @@ export default class PostgresSearchProvider extends BaseSearchProvider { ], }; + // Resolve the collections and individual documents accessible to the + // model upfront so the search query needs no membership joins. If + // collectionId is passed as an option it is assumed that the + // authorization has already been done in the router. + const [membershipDocumentIds, collectionIds] = await Promise.all([ + model instanceof User + ? Document.membershipDocumentIds(model.id) + : Promise.resolve([]), + options.collectionId + ? Promise.resolve([options.collectionId]) + : model.collectionIds(), + ]); + if (model instanceof User) { - where[Op.or].push( - { "$memberships.id$": { [Op.ne]: null } }, - { "$groupMemberships.id$": { [Op.ne]: null } } - ); + if (membershipDocumentIds.length) { + where[Op.or].push({ id: membershipDocumentIds }); + } // Allow users to see their own drafts that have no collection, where no // membership or collection access applies. Drafts in collections remain @@ -651,13 +614,6 @@ export default class PostgresSearchProvider extends BaseSearchProvider { } } - // Ensure we're filtering by the users accessible collections. If - // collectionId is passed as an option it is assumed that the authorization - // has already been done in the router - const collectionIds = options.collectionId - ? [options.collectionId] - : await model.collectionIds(); - if (options.collectionId) { where[Op.and].push({ collectionId: options.collectionId }); } @@ -707,7 +663,8 @@ export default class PostgresSearchProvider extends BaseSearchProvider { if ( options.statusFilter?.includes(StatusFilter.Draft) && - // Only ever include draft results for the user's own documents + // Only include draft results for the user's own documents, or those + // explicitly shared with them model instanceof User ) { statusQuery.push({ @@ -721,7 +678,9 @@ export default class PostgresSearchProvider extends BaseSearchProvider { }, [Op.or]: [ { createdById: model.id }, - { "$memberships.id$": { [Op.ne]: null } }, + ...(membershipDocumentIds.length + ? [{ id: membershipDocumentIds }] + : []), ], }, ], diff --git a/server/models/Document.test.ts b/server/models/Document.test.ts index 825ba9cdbd..b58c07e2ac 100644 --- a/server/models/Document.test.ts +++ b/server/models/Document.test.ts @@ -9,11 +9,13 @@ import { buildCollection, buildComment, buildResolvedComment, + buildGroup, buildTeam, buildUser, buildGuestUser, } from "@server/test/factories"; import { withAPIContext } from "@server/test/support"; +import GroupMembership from "./GroupMembership"; import UserMembership from "./UserMembership"; beforeEach(() => { @@ -257,6 +259,135 @@ describe("#findAllChildDocumentIds", () => { }); }); +describe("#membershipDocumentIds", () => { + it("should return empty array when the user has no document memberships", async () => { + const user = await buildUser(); + const ids = await Document.membershipDocumentIds(user.id); + expect(ids).toEqual([]); + }); + + it("should return documents shared directly with the user", async () => { + const team = await buildTeam(); + const user = await buildUser({ teamId: team.id }); + const otherUser = await buildUser({ teamId: team.id }); + const document = await buildDocument({ + teamId: team.id, + userId: otherUser.id, + }); + await UserMembership.create({ + createdById: otherUser.id, + documentId: document.id, + userId: user.id, + permission: DocumentPermission.Read, + }); + + const ids = await Document.membershipDocumentIds(user.id); + expect(ids).toEqual([document.id]); + }); + + it("should return documents shared with the user through a group", async () => { + const team = await buildTeam(); + const user = await buildUser({ teamId: team.id }); + const otherUser = await buildUser({ teamId: team.id }); + const document = await buildDocument({ + teamId: team.id, + userId: otherUser.id, + }); + const group = await buildGroup({ teamId: team.id }); + await group.$add("user", user, { through: { createdById: otherUser.id } }); + await GroupMembership.create({ + createdById: otherUser.id, + groupId: group.id, + documentId: document.id, + permission: DocumentPermission.Read, + }); + + const ids = await Document.membershipDocumentIds(user.id); + expect(ids).toEqual([document.id]); + }); + + it("should deduplicate documents shared both directly and through a group", async () => { + const team = await buildTeam(); + const user = await buildUser({ teamId: team.id }); + const otherUser = await buildUser({ teamId: team.id }); + const document = await buildDocument({ + teamId: team.id, + userId: otherUser.id, + }); + await UserMembership.create({ + createdById: otherUser.id, + documentId: document.id, + userId: user.id, + permission: DocumentPermission.Read, + }); + const group = await buildGroup({ teamId: team.id }); + await group.$add("user", user, { through: { createdById: otherUser.id } }); + await GroupMembership.create({ + createdById: otherUser.id, + groupId: group.id, + documentId: document.id, + permission: DocumentPermission.Read, + }); + + const ids = await Document.membershipDocumentIds(user.id); + expect(ids).toEqual([document.id]); + }); + + it("should not return memberships of other users or collections", async () => { + const team = await buildTeam(); + const user = await buildUser({ teamId: team.id }); + const otherUser = await buildUser({ teamId: team.id }); + const collection = await buildCollection({ + teamId: team.id, + userId: otherUser.id, + }); + const document = await buildDocument({ + teamId: team.id, + userId: otherUser.id, + collectionId: collection.id, + }); + // document membership for another user + await UserMembership.create({ + createdById: otherUser.id, + documentId: document.id, + userId: otherUser.id, + permission: DocumentPermission.Read, + }); + // collection-level membership for this user + await UserMembership.create({ + createdById: otherUser.id, + collectionId: collection.id, + userId: user.id, + permission: CollectionPermission.Read, + }); + + const ids = await Document.membershipDocumentIds(user.id); + expect(ids).toEqual([]); + }); + + it("should not return documents from deleted group memberships", async () => { + const team = await buildTeam(); + const user = await buildUser({ teamId: team.id }); + const otherUser = await buildUser({ teamId: team.id }); + const document = await buildDocument({ + teamId: team.id, + userId: otherUser.id, + }); + const group = await buildGroup({ teamId: team.id }); + await group.$add("user", user, { through: { createdById: otherUser.id } }); + const membership = await GroupMembership.create({ + createdById: otherUser.id, + groupId: group.id, + documentId: document.id, + permission: DocumentPermission.Read, + }); + await membership.destroy(); + + const ids = await Document.membershipDocumentIds(user.id); + expect(ids).toEqual([]); + }); +}); + describe("#findByPk", () => { test("should return document when urlId is correct", async () => { const document = await buildDocument(); diff --git a/server/models/Document.ts b/server/models/Document.ts index 15b9f40859..5c27d51ae4 100644 --- a/server/models/Document.ts +++ b/server/models/Document.ts @@ -761,6 +761,60 @@ class Document extends ArchivableModel< return uniq(membershipUserIds); } + /** + * Returns an array of unique document IDs that the user is a member of, + * either via direct membership or through a group membership. + * + * @param userId The user ID to find document memberships for. + * @returns A promise resolving to an array of document IDs. + */ + static async membershipDocumentIds(userId: string): Promise { + const [memberships, groupMemberships] = await Promise.all([ + UserMembership.findAll({ + attributes: ["documentId"], + where: { + userId, + documentId: { + [Op.ne]: null, + }, + }, + }), + GroupMembership.findAll({ + attributes: ["documentId"], + where: { + documentId: { + [Op.ne]: null, + }, + }, + include: [ + { + model: Group, + as: "group", + attributes: [], + required: true, + include: [ + { + model: GroupUser, + as: "groupUsers", + attributes: [], + required: true, + where: { + userId, + }, + }, + ], + }, + ], + }), + ]); + + return uniq( + [...memberships, ...groupMemberships] + .map((membership) => membership.documentId) + .filter((id): id is string => !isNil(id)) + ); + } + static withMembershipScope( userId: string, options?: FindOptions & { includeDrafts?: boolean }