fix: import workspace may only show default preload page (#2685)

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

@@ -3,14 +3,15 @@ import type { Y as YType } from '@blocksuite/store';
 import { uuidv4, Workspace } from '@blocksuite/store';
 import { beforeEach, describe, expect, test, vi } from 'vitest';
 
-import type { SQLiteProvider } from '../../type';
-import { createSQLiteProvider } from '../index';
+import type { SQLiteDBDownloadProvider, SQLiteProvider } from '../../type';
+import { createSQLiteDBDownloadProvider, createSQLiteProvider } from '../index';
 
 const Y = Workspace.Y;
 
 let id: string;
 let workspace: Workspace;
 let provider: SQLiteProvider;
+let downloadProvider: SQLiteDBDownloadProvider;
 
 let offlineYdoc: YType.Doc;
 
@@ -60,12 +61,13 @@ beforeEach(() => {
   });
   workspace.register(AffineSchemas).register(__unstableSchemas);
   provider = createSQLiteProvider(workspace);
+  downloadProvider = createSQLiteDBDownloadProvider(workspace);
   offlineYdoc = new Y.Doc();
   offlineYdoc.getText('text').insert(0, 'sqlite-hello');
 });
 
-describe('SQLite provider', () => {
-  test('connect', async () => {
+describe('SQLite download provider', () => {
+  test('sync updates', async () => {
     // on connect, the updates from sqlite should be sync'ed to the existing ydoc
     // and ydoc should be sync'ed back to sqlite
     // Workspace.Y.applyUpdate(workspace.doc);
@@ -73,7 +75,8 @@ describe('SQLite provider', () => {
 
     expect(offlineYdoc.getText('text').toString()).toBe('sqlite-hello');
 
-    await provider.connect();
+    downloadProvider.sync();
+    await downloadProvider.whenReady;
 
     // depending on the nature of the sync, the data can be sync'ed in either direction
     const options = ['mem-hellosqlite-hello', 'sqlite-hellomem-hello'];
@@ -83,10 +86,10 @@ describe('SQLite provider', () => {
     expect(synced.length).toBe(1);
     expect(workspace.doc.getText('text').toString()).toBe(synced[0]);
 
-    workspace.doc.getText('text').insert(0, 'world');
+    // workspace.doc.getText('text').insert(0, 'world');
 
-    // check if the data are sync'ed
-    expect(offlineYdoc.getText('text').toString()).toBe('world' + synced[0]);
+    // // check if the data are sync'ed
+    // expect(offlineYdoc.getText('text').toString()).toBe('world' + synced[0]);
   });
 
   test('blobs will be synced to sqlite on connect', async () => {
@@ -98,14 +101,15 @@ describe('SQLite provider', () => {
       return blob;
     });
 
-    await provider.connect();
+    downloadProvider.sync();
+    await downloadProvider.whenReady;
     await new Promise(resolve => setTimeout(resolve, 100));
 
     expect(mockedAddBlob).toBeCalledWith(id, 'blob1', bin);
   });
 
   test('on db update', async () => {
-    await provider.connect();
+    provider.connect();
 
     offlineYdoc.getText('text').insert(0, 'sqlite-world');
 
@@ -114,8 +118,8 @@ describe('SQLite provider', () => {
       update: Y.encodeStateAsUpdate(offlineYdoc),
     });
 
-    // not yet updated
-    expect(workspace.doc.getText('text').toString()).toBe('sqlite-hello');
+    // not yet updated (because the workspace id is different)
+    expect(workspace.doc.getText('text').toString()).toBe('');
 
     triggerDBUpdate?.({
       workspaceId: id,

packages/workspace/src/providers/index.ts

@@ -17,6 +17,7 @@ import type {
   LocalIndexedDBBackgroundProvider,
   LocalIndexedDBDownloadProvider,
   Provider,
+  SQLiteDBDownloadProvider,
   SQLiteProvider,
 } from '../type';
 import { CallbackSet } from '../utils';
@@ -156,10 +157,11 @@ const createIndexedDBDownloadProvider = (
   };
 };
 
+const sqliteOrigin = Symbol('sqlite-provider-origin');
+
 const createSQLiteProvider = (
   blockSuiteWorkspace: BlockSuiteWorkspace
 ): SQLiteProvider => {
-  const sqliteOrigin = Symbol('sqlite-provider-origin');
   const apis = window.apis!;
   const events = window.events!;
   // make sure it is being used in Electron with APIs
@@ -173,9 +175,87 @@ const createSQLiteProvider = (
     apis.db.applyDocUpdate(blockSuiteWorkspace.id, update);
   }
 
+  let unsubscribe = () => {};
+  let connected = false;
+
+  const callbacks = new CallbackSet();
+
+  const connect = () => {
+    logger.info('connecting sqlite provider', blockSuiteWorkspace.id);
+    blockSuiteWorkspace.doc.on('update', handleUpdate);
+    unsubscribe = events.db.onExternalUpdate(({ update, workspaceId }) => {
+      if (workspaceId === blockSuiteWorkspace.id) {
+        Y.applyUpdate(blockSuiteWorkspace.doc, update, sqliteOrigin);
+      }
+    });
+    connected = true;
+    logger.info('connecting sqlite done', blockSuiteWorkspace.id);
+  };
+
+  const cleanup = () => {
+    logger.info('disconnecting sqlite provider', blockSuiteWorkspace.id);
+    unsubscribe();
+    blockSuiteWorkspace.doc.off('update', handleUpdate);
+    connected = false;
+  };
+
+  return {
+    flavour: 'sqlite',
+    background: true,
+    callbacks,
+    get connected(): boolean {
+      return connected;
+    },
+    cleanup,
+    connect,
+    disconnect: cleanup,
+  };
+};
+
+const createSQLiteDBDownloadProvider = (
+  blockSuiteWorkspace: BlockSuiteWorkspace
+): SQLiteDBDownloadProvider => {
+  const apis = window.apis!;
+  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', blockSuiteWorkspace.id);
+    const updates = await apis.db.getDocAsUpdates(blockSuiteWorkspace.id);
+
+    if (disconnected) {
+      return;
+    }
+
+    if (updates) {
+      Y.applyUpdate(blockSuiteWorkspace.doc, updates, sqliteOrigin);
+    }
+
+    const diff = Y.encodeStateAsUpdate(blockSuiteWorkspace.doc, updates);
+
+    // also apply updates to sqlite
+    apis.db.applyDocUpdate(blockSuiteWorkspace.id, diff);
+
+    const bs = blockSuiteWorkspace.blobs;
+
+    if (bs && !disconnected) {
+      await syncBlobIntoSQLite(bs);
+    }
+  }
+
   async function syncBlobIntoSQLite(bs: BlobManager) {
     const persistedKeys = await apis.db.getBlobKeys(blockSuiteWorkspace.id);
 
+    if (disconnected) {
+      return;
+    }
+
     const allKeys = await bs.list().catch(() => []);
     const keysToPersist = allKeys.filter(k => !persistedKeys.includes(k));
 
@@ -187,6 +267,11 @@ const createSQLiteProvider = (
           logger.warn('blob not found for', k);
           return;
         }
+
+        if (disconnected) {
+          return;
+        }
+
         return window.apis?.db.addBlob(
           blockSuiteWorkspace.id,
           k,
@@ -196,61 +281,23 @@ const createSQLiteProvider = (
     );
   }
 
-  async function syncUpdates() {
-    logger.info('syncing updates from sqlite', blockSuiteWorkspace.id);
-    const updates = await apis.db.getDocAsUpdates(blockSuiteWorkspace.id);
-
-    if (updates) {
-      Y.applyUpdate(blockSuiteWorkspace.doc, updates, sqliteOrigin);
-    }
-
-    const mergeUpdates = Y.encodeStateAsUpdate(blockSuiteWorkspace.doc);
-
-    // also apply updates to sqlite
-    apis.db.applyDocUpdate(blockSuiteWorkspace.id, mergeUpdates);
-
-    const bs = blockSuiteWorkspace.blobs;
-
-    if (bs) {
-      // this can be non-blocking
-      syncBlobIntoSQLite(bs);
-    }
-  }
-
-  let unsubscribe = () => {};
-  let connected = false;
-  const callbacks = new CallbackSet();
-
   return {
-    flavour: 'sqlite',
-    background: true,
-    callbacks,
-    get connected(): boolean {
-      return connected;
+    flavour: 'sqlite-download',
+    necessary: true,
+    get whenReady() {
+      return promise;
     },
     cleanup: () => {
-      throw new Error('Method not implemented.');
+      disconnected = true;
     },
-    connect: async () => {
-      logger.info('connecting sqlite provider', blockSuiteWorkspace.id);
-      await syncUpdates();
-      connected = true;
-
-      blockSuiteWorkspace.doc.on('update', handleUpdate);
-
-      unsubscribe = events.db.onExternalUpdate(({ update, workspaceId }) => {
-        if (workspaceId === blockSuiteWorkspace.id) {
-          Y.applyUpdate(blockSuiteWorkspace.doc, update, sqliteOrigin);
-        }
-      });
-
-      // blockSuiteWorkspace.doc.on('destroy', ...);
-      logger.info('connecting sqlite done', blockSuiteWorkspace.id);
-    },
-    disconnect: () => {
-      unsubscribe();
-      blockSuiteWorkspace.doc.off('update', handleUpdate);
-      connected = false;
+    sync: async () => {
+      logger.info('connect indexeddb provider', blockSuiteWorkspace.id);
+      try {
+        await syncUpdates();
+        _resolve();
+      } catch (error) {
+        _reject(error);
+      }
     },
   };
 };
@@ -261,21 +308,30 @@ export {
   createBroadCastChannelProvider,
   createIndexedDBBackgroundProvider,
   createIndexedDBDownloadProvider,
+  createSQLiteDBDownloadProvider,
   createSQLiteProvider,
 };
 
 export const createLocalProviders = (
   blockSuiteWorkspace: BlockSuiteWorkspace
 ): Provider[] => {
-  return (
-    [
-      config.enableBroadCastChannelProvider &&
-        createBroadCastChannelProvider(blockSuiteWorkspace),
-      createIndexedDBBackgroundProvider(blockSuiteWorkspace),
-      createIndexedDBDownloadProvider(blockSuiteWorkspace),
-      environment.isDesktop && createSQLiteProvider(blockSuiteWorkspace),
-    ] as any[]
-  ).filter(v => Boolean(v));
+  const providers = [
+    createIndexedDBBackgroundProvider(blockSuiteWorkspace),
+    createIndexedDBDownloadProvider(blockSuiteWorkspace),
+  ] as Provider[];
+
+  if (config.enableBroadCastChannelProvider) {
+    providers.push(createBroadCastChannelProvider(blockSuiteWorkspace));
+  }
+
+  if (environment.isDesktop) {
+    providers.push(
+      createSQLiteProvider(blockSuiteWorkspace),
+      createSQLiteDBDownloadProvider(blockSuiteWorkspace)
+    );
+  }
+
+  return providers;
 };
 
 export const createAffineProviders = (

packages/workspace/src/type.ts

@@ -78,12 +78,16 @@ export interface LocalIndexedDBBackgroundProvider extends BackgroundProvider {
   flavour: 'local-indexeddb-background';
 }
 
+export interface LocalIndexedDBDownloadProvider extends NecessaryProvider {
+  flavour: 'local-indexeddb';
+}
+
 export interface SQLiteProvider extends BackgroundProvider {
   flavour: 'sqlite';
 }
 
-export interface LocalIndexedDBDownloadProvider extends NecessaryProvider {
-  flavour: 'local-indexeddb';
+export interface SQLiteDBDownloadProvider extends NecessaryProvider {
+  flavour: 'sqlite-download';
 }
 
 export interface AffineWebSocketProvider extends BackgroundProvider {