feat: sqlite subdocument (#2816)

Co-authored-by: Alex Yang <himself65@outlook.com>

packages/workspace/src/affine/__tests__/sync.spec.ts

@@ -140,7 +140,7 @@ describe('ydoc sync', () => {
 
       const pageId = uuidv4();
       const page1 = workspace1.createPage({ id: pageId });
-      await page1.waitForLoaded()
+      await page1.waitForLoaded();
       const pageBlockId = page1.addBlock('affine:page', {
         title: new page1.Text(''),
       });
@@ -153,7 +153,7 @@ describe('ydoc sync', () => {
         workspace1.doc.getMap(`space:${pageId}`).toJSON()
       );
       const page2 = workspace2.getPage(pageId) as Page;
-      await page2.waitForLoaded()
+      await page2.waitForLoaded();
       page1.updateBlock(
         page1.getBlockById(paragraphId) as ParagraphBlockModel,
         {

packages/workspace/src/providers/__tests__/sqlite-provider.spec.ts

@@ -7,7 +7,10 @@ import type { Y as YType } from '@blocksuite/store';
 import { uuidv4, Workspace } from '@blocksuite/store';
 import { beforeEach, describe, expect, test, vi } from 'vitest';
 
-import { createSQLiteDBDownloadProvider, createSQLiteProvider } from '../index';
+import {
+  createSQLiteDBDownloadProvider,
+  createSQLiteProvider,
+} from '../sqlite-providers';
 
 const Y = Workspace.Y;
 
@@ -148,15 +151,21 @@ describe('SQLite download provider', () => {
   test('disconnect handlers', async () => {
     const offHandler = vi.fn();
     let handleUpdate = () => {};
-    workspace.doc.on = (_: string, fn: () => void) => {
-      handleUpdate = fn;
+    let handleSubdocs = () => {};
+    workspace.doc.on = (event: string, fn: () => void) => {
+      if (event === 'update') {
+        handleUpdate = fn;
+      } else if (event === 'subdocs') {
+        handleSubdocs = fn;
+      }
     };
     workspace.doc.off = offHandler;
-    await provider.connect();
+    provider.connect();
 
     provider.disconnect();
 
     expect(triggerDBUpdate).toBe(null);
     expect(offHandler).toBeCalledWith('update', handleUpdate);
+    expect(offHandler).toBeCalledWith('subdocs', handleSubdocs);
   });
 });

packages/workspace/src/providers/index.ts

@@ -3,8 +3,6 @@ import type {
   AffineWebSocketProvider,
   LocalIndexedDBBackgroundProvider,
   LocalIndexedDBDownloadProvider,
-  SQLiteDBDownloadProvider,
-  SQLiteProvider,
 } from '@affine/env/workspace';
 import type { Disposable, DocProviderCreator } from '@blocksuite/store';
 import { assertExists, Workspace } from '@blocksuite/store';
@@ -21,6 +19,10 @@ import { getLoginStorage, storageChangeSlot } from '../affine/login';
 import { CallbackSet } from '../utils';
 import { createAffineDownloadProvider } from './affine-download';
 import { localProviderLogger as logger } from './logger';
+import {
+  createSQLiteDBDownloadProvider,
+  createSQLiteProvider,
+} from './sqlite-providers';
 
 const Y = Workspace.Y;
 
@@ -151,151 +153,6 @@ const createIndexedDBDownloadProvider: DocProviderCreator = (
   };
 };
 
-const sqliteOrigin = Symbol('sqlite-provider-origin');
-
-const createSQLiteProvider: DocProviderCreator = (id, doc): SQLiteProvider => {
-  const { apis, events } = window;
-  // make sure it is being used in Electron with APIs
-  assertExists(apis);
-  assertExists(events);
-
-  function handleUpdate(update: Uint8Array, origin: unknown) {
-    if (origin === sqliteOrigin) {
-      return;
-    }
-    apis.db.applyDocUpdate(id, update).catch(err => {
-      console.error(err);
-    });
-  }
-
-  let unsubscribe = () => {};
-  let connected = false;
-
-  const connect = () => {
-    logger.info('connecting sqlite provider', id);
-    doc.on('update', handleUpdate);
-    unsubscribe = events.db.onExternalUpdate(
-      ({
-        update,
-        workspaceId,
-      }: {
-        workspaceId: string;
-        update: Uint8Array;
-      }) => {
-        if (workspaceId === id) {
-          Y.applyUpdate(doc, update, sqliteOrigin);
-        }
-      }
-    );
-    connected = true;
-    logger.info('connecting sqlite done', id);
-  };
-
-  const cleanup = () => {
-    logger.info('disconnecting sqlite provider', id);
-    unsubscribe();
-    doc.off('update', handleUpdate);
-    connected = false;
-  };
-
-  return {
-    flavour: 'sqlite',
-    passive: true,
-    get connected(): boolean {
-      return connected;
-    },
-    cleanup,
-    connect,
-    disconnect: cleanup,
-  };
-};
-
-const createSQLiteDBDownloadProvider: DocProviderCreator = (
-  id,
-  doc
-): SQLiteDBDownloadProvider => {
-  const { apis } = window;
-  let disconnected = false;
-
-  let _resolve: () => void;
-  let _reject: (error: unknown) => void;
-  const promise = new Promise<void>((resolve, reject) => {
-    _resolve = resolve;
-    _reject = reject;
-  });
-
-  async function syncUpdates() {
-    logger.info('syncing updates from sqlite', id);
-    const updates = await apis.db.getDocAsUpdates(id);
-
-    if (disconnected) {
-      return;
-    }
-
-    if (updates) {
-      Y.applyUpdate(doc, updates, sqliteOrigin);
-    }
-
-    const diff = Y.encodeStateAsUpdate(doc, updates);
-
-    // also apply updates to sqlite
-    await apis.db.applyDocUpdate(id, diff);
-  }
-
-  // fixme(pengx17): should n't sync blob in doc provider
-  // async function _syncBlobIntoSQLite(bs: BlobManager) {
-  //   const persistedKeys = await apis.db.getBlobKeys(id);
-  //
-  //   if (disconnected) {
-  //     return;
-  //   }
-  //
-  //   const allKeys = await bs.list().catch(() => []);
-  //   const keysToPersist = allKeys.filter(k => !persistedKeys.includes(k));
-  //
-  //   logger.info('persisting blobs', keysToPersist, 'to sqlite');
-  //   return Promise.all(
-  //     keysToPersist.map(async k => {
-  //       const blob = await bs.get(k);
-  //       if (!blob) {
-  //         logger.warn('blob not found for', k);
-  //         return;
-  //       }
-  //
-  //       if (disconnected) {
-  //         return;
-  //       }
-  //
-  //       return apis?.db.addBlob(
-  //         id,
-  //         k,
-  //         new Uint8Array(await blob.arrayBuffer())
-  //       );
-  //     })
-  //   );
-  // }
-
-  return {
-    flavour: 'sqlite-download',
-    active: true,
-    get whenReady() {
-      return promise;
-    },
-    cleanup: () => {
-      disconnected = true;
-    },
-    sync: async () => {
-      logger.info('connect indexeddb provider', id);
-      try {
-        await syncUpdates();
-        _resolve();
-      } catch (error) {
-        _reject(error);
-      }
-    },
-  };
-};
-
 export {
   createAffineDownloadProvider,
   createAffineWebSocketProvider,

packages/workspace/src/providers/sqlite-providers.ts

@@ -0,0 +1,212 @@
+import type {
+  SQLiteDBDownloadProvider,
+  SQLiteProvider,
+} from '@affine/env/workspace';
+import type { DocProviderCreator } from '@blocksuite/store';
+import {
+  assertExists,
+  Workspace as BlockSuiteWorkspace,
+} from '@blocksuite/store';
+import type { Doc } from 'yjs';
+
+import { localProviderLogger as logger } from './logger';
+
+const Y = BlockSuiteWorkspace.Y;
+
+const sqliteOrigin = Symbol('sqlite-provider-origin');
+
+type SubDocsEvent = {
+  added: Set<Doc>;
+  removed: Set<Doc>;
+  loaded: Set<Doc>;
+};
+
+/**
+ * A provider that is responsible for syncing updates the workspace with the local SQLite database.
+ */
+export const createSQLiteProvider: DocProviderCreator = (
+  id,
+  rootDoc
+): SQLiteProvider => {
+  const { apis, events } = window;
+  // make sure it is being used in Electron with APIs
+  assertExists(apis);
+  assertExists(events);
+
+  const updateHandlerMap = new WeakMap<
+    Doc,
+    (update: Uint8Array, origin: unknown) => void
+  >();
+  const subDocsHandlerMap = new WeakMap<Doc, (event: SubDocsEvent) => void>();
+
+  const createOrHandleUpdate = (doc: Doc) => {
+    if (updateHandlerMap.has(doc)) {
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      return updateHandlerMap.get(doc)!;
+    }
+
+    function handleUpdate(update: Uint8Array, origin: unknown) {
+      if (origin === sqliteOrigin) {
+        return;
+      }
+      const subdocId = doc.guid === id ? undefined : doc.guid;
+      apis.db.applyDocUpdate(id, update, subdocId).catch(err => {
+        logger.error(err);
+      });
+    }
+    updateHandlerMap.set(doc, handleUpdate);
+    return handleUpdate;
+  };
+
+  const createOrGetHandleSubDocs = (doc: Doc) => {
+    if (subDocsHandlerMap.has(doc)) {
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      return subDocsHandlerMap.get(doc)!;
+    }
+    function handleSubdocs(event: SubDocsEvent) {
+      event.removed.forEach(doc => {
+        untrackDoc(doc);
+      });
+      event.loaded.forEach(doc => {
+        trackDoc(doc);
+      });
+    }
+    subDocsHandlerMap.set(doc, handleSubdocs);
+    return handleSubdocs;
+  };
+
+  function trackDoc(doc: Doc) {
+    doc.on('update', createOrHandleUpdate(doc));
+    doc.on('subdocs', createOrGetHandleSubDocs(doc));
+    doc.subdocs.forEach(doc => {
+      trackDoc(doc);
+    });
+  }
+
+  function untrackDoc(doc: Doc) {
+    doc.subdocs.forEach(doc => {
+      untrackDoc(doc);
+    });
+    doc.off('update', createOrHandleUpdate(doc));
+    doc.off('subdocs', createOrGetHandleSubDocs(doc));
+  }
+
+  let unsubscribe = () => {};
+  let connected = false;
+
+  const connect = () => {
+    logger.info('connecting sqlite provider', id);
+    trackDoc(rootDoc);
+
+    unsubscribe = events.db.onExternalUpdate(
+      ({
+        update,
+        workspaceId,
+        docId,
+      }: {
+        workspaceId: string;
+        update: Uint8Array;
+        docId?: string;
+      }) => {
+        if (workspaceId === id) {
+          if (docId) {
+            for (const doc of rootDoc.subdocs) {
+              if (doc.guid === docId) {
+                Y.applyUpdate(doc, update, sqliteOrigin);
+                return;
+              }
+            }
+          } else {
+            Y.applyUpdate(rootDoc, update, sqliteOrigin);
+          }
+        }
+      }
+    );
+    connected = true;
+    logger.info('connecting sqlite done', id);
+  };
+
+  const cleanup = () => {
+    logger.info('disconnecting sqlite provider', id);
+    unsubscribe();
+    untrackDoc(rootDoc);
+    connected = false;
+  };
+
+  return {
+    flavour: 'sqlite',
+    passive: true,
+    get connected(): boolean {
+      return connected;
+    },
+    cleanup,
+    connect,
+    disconnect: cleanup,
+  };
+};
+
+/**
+ * A provider that is responsible for DOWNLOADING updates from the local SQLite database.
+ */
+export const createSQLiteDBDownloadProvider: DocProviderCreator = (
+  id,
+  rootDoc
+): SQLiteDBDownloadProvider => {
+  const { apis } = window;
+  let disconnected = false;
+
+  let _resolve: () => void;
+  let _reject: (error: unknown) => void;
+  const promise = new Promise<void>((resolve, reject) => {
+    _resolve = resolve;
+    _reject = reject;
+  });
+
+  async function syncUpdates(doc: Doc) {
+    logger.info('syncing updates from sqlite', id);
+    const subdocId = doc.guid === id ? undefined : doc.guid;
+    const updates = await apis.db.getDocAsUpdates(id, subdocId);
+
+    if (disconnected) {
+      return false;
+    }
+
+    if (updates) {
+      Y.applyUpdate(doc, updates, sqliteOrigin);
+    }
+
+    const diff = Y.encodeStateAsUpdate(doc, updates);
+
+    // also apply updates to sqlite
+    await apis.db.applyDocUpdate(id, diff, subdocId);
+
+    return true;
+  }
+
+  async function syncAllUpdates(doc: Doc) {
+    if (await syncUpdates(doc)) {
+      const subdocs = Array.from(doc.subdocs).filter(d => d.shouldLoad);
+      await Promise.all(subdocs.map(syncAllUpdates));
+    }
+  }
+
+  return {
+    flavour: 'sqlite-download',
+    active: true,
+    get whenReady() {
+      return promise;
+    },
+    cleanup: () => {
+      disconnected = true;
+    },
+    sync: async () => {
+      logger.info('connect indexeddb provider', id);
+      try {
+        await syncAllUpdates(rootDoc);
+        _resolve();
+      } catch (error) {
+        _reject(error);
+      }
+    },
+  };
+};