feat: add index for snapshots (#8163)

packages/backend/server/migrations/20240908130725_add_workspace_level_snapshot_index/migration.sql

@@ -0,0 +1,12 @@
+/*
+  Warnings:
+
+  - The primary key for the `snapshots` table will be changed. If it partially fails, the table could be left without primary key constraint.
+
+*/
+-- AlterTable
+ALTER TABLE "snapshots" DROP CONSTRAINT "snapshots_pkey",
+ADD CONSTRAINT "snapshots_pkey" PRIMARY KEY ("workspace_id", "guid");
+
+-- CreateIndex
+CREATE INDEX "snapshots_workspace_id_updated_at_idx" ON "snapshots"("workspace_id", "updated_at");

packages/backend/server/schema.prisma

@@ -262,7 +262,8 @@ model Snapshot {
   // we need to clear all hanging updates and snapshots before enable the foreign key on workspaceId
   // workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
 
-  @@id([id, workspaceId])
+  @@id([workspaceId, id])
+  @@index([workspaceId, updatedAt])
   @@map("snapshots")
 }
 

packages/backend/server/src/core/doc/adapters/workspace.ts

@@ -143,8 +143,6 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
   }
 
   async getSpaceDocTimestamps(workspaceId: string, after?: number) {
-    // TODO(@forehalo): do we need a [Clock] table to store the last seen time of each doc?
-    // SLOW if query DB in large workspace by `updatedAt: { gt: after }`
     const snapshots = await this.db.snapshot.findMany({
       select: {
         id: true,
@@ -152,6 +150,13 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
       },
       where: {
         workspaceId,
+        ...(after
+          ? {
+              updatedAt: {
+                gt: new Date(after),
+              },
+            }
+          : {}),
       },
     });
 
@@ -176,9 +181,7 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
     const result: Record<string, number> = {};
 
     snapshots.forEach(s => {
-      if (!after || s.updatedAt.getTime() > after) {
-        result[s.id] = s.updatedAt.getTime();
-      }
+      result[s.id] = s.updatedAt.getTime();
     });
 
     updates.forEach(u => {
@@ -378,7 +381,7 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
   protected async getDocSnapshot(workspaceId: string, docId: string) {
     const snapshot = await this.db.snapshot.findUnique({
       where: {
-        id_workspaceId: {
+        workspaceId_id: {
           workspaceId,
           id: docId,
         },
@@ -414,7 +417,7 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
     //
     //  ii. Prisma doesn't support `upsert` with additional `where` condition along side unique constraint.
     //      In our case, we need to manually check the `updatedAt` to avoid overriding the newer snapshot.
-    //      where: { id_workspaceId: {}, updatedAt: { lt: updatedAt } }
+    //      where: { workspaceId_id: {}, updatedAt: { lt: updatedAt } }
     //                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     try {
       const result: { updatedAt: Date }[] = await this.db.$queryRaw`
@@ -432,7 +435,7 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
       //     seq: true,
       //   },
       //   where: {
-      //     id_workspaceId: {
+      //     workspaceId_id: {
       //       workspaceId,
       //       id: guid,
       //     },
@@ -571,7 +574,7 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
           seq: true,
         },
         where: {
-          id_workspaceId: {
+          workspaceId_id: {
             workspaceId,
             id: guid,
           },
@@ -594,7 +597,7 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
             seq: true,
           },
           where: {
-            id_workspaceId: {
+            workspaceId_id: {
               workspaceId,
               id: guid,
             },

packages/backend/server/src/data/migrations/1698398506533-guid.ts

@@ -78,7 +78,7 @@ export class Guid1698398506533 {
               }),
               db.snapshot.update({
                 where: {
-                  id_workspaceId: {
+                  workspaceId_id: {
                     id: docId.guid,
                     workspaceId: doc.workspaceId,
                   },
@@ -93,7 +93,7 @@ export class Guid1698398506533 {
             // just modify the id the required one
             await db.snapshot.update({
               where: {
-                id_workspaceId: {
+                workspaceId_id: {
                   id: doc.id,
                   workspaceId: doc.workspaceId,
                 },

packages/backend/server/tests/doc/workspace.spec.ts

@@ -79,7 +79,7 @@ test('should have sequential update number', async t => {
   // fake the seq num is about to overflow
   await db.snapshot.update({
     where: {
-      id_workspaceId: {
+      workspaceId_id: {
         id: '2',
         workspaceId: '2',
       },
@@ -273,7 +273,7 @@ test('should not update snapshot if doc is outdated', async t => {
   // fake the snapshot is a lot newer
   await db.snapshot.update({
     where: {
-      id_workspaceId: {
+      workspaceId_id: {
         workspaceId: '2',
         id: '1',
       },