perf: Avoid membership joins in search (#13269)

* perf: Avoid membership joins in search

* feedback

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", () => {

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<RankedDocument[]>;
+      })) as unknown as RankedDocument[];
-
-      const countQuery = Document.unscoped().count({
-        // @ts-expect-error Types are incorrect for count
-        replacements: findOptions.replacements,
-        where,
-      }) as unknown as Promise<number>;
-      const [results, count] = await Promise.all([resultsQuery, countQuery]);
 
       // Final query to get associated document data
-      const documents = await Document.findAll({
+      const [documents, count] = await Promise.all([
-        where: {
+        Document.findAll({
-          id: map(results, "id"),
+          where: {
-          teamId: team.id,
+            id: map(results, "id"),
-        },
+            teamId: team.id,
-        include: [
-          {
-            model: Collection,
-            as: "collection",
           },
-        ],
+          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<number>;
-
       // 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
+        PostgresSearchProvider.countResults({
-          ? Promise.resolve(results.length)
+          results,
-          : countQuery,
+          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<Document>;
+  }): Promise<number> {
+    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<number>;
+  }
+
   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(
+      if (membershipDocumentIds.length) {
-        { "$memberships.id$": { [Op.ne]: null } },
+        where[Op.or].push({ id: membershipDocumentIds });
-        { "$groupMemberships.id$": { [Op.ne]: null } }
+      }
-      );
 
       // 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 }]
+                : []),
             ],
           },
         ],

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();

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<string[]> {
+    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<Document> & { includeDrafts?: boolean }