fix: add DB migration to add workspace (#3115)

apps/electron/src/helper/db/__tests__/migration.spec.ts

@@ -0,0 +1,69 @@
+import path from 'node:path';
+
+import { SqliteConnection } from '@affine/native';
+import { afterEach, describe, expect, it, vi } from 'vitest';
+import * as Y from 'yjs';
+
+import { removeWithRetry } from '../../../../tests/utils';
+import { copyToTemp, migrateToSubdocAndReplaceDatabase } from '../migration';
+
+const tmpDir = path.join(__dirname, 'tmp');
+const testDBFilePath = path.resolve(__dirname, 'old-db.affine');
+
+const appDataPath = path.join(tmpDir, 'app-data');
+
+vi.mock('../../main-rpc', () => ({
+  mainRPC: {
+    getPath: async () => appDataPath,
+  },
+}));
+
+afterEach(async () => {
+  await removeWithRetry(tmpDir);
+});
+
+describe('migrateToSubdocAndReplaceDatabase', () => {
+  it('should migrate and replace the database', async () => {
+    const copiedDbFilePath = await copyToTemp(testDBFilePath);
+    await migrateToSubdocAndReplaceDatabase(copiedDbFilePath);
+
+    const db = new SqliteConnection(copiedDbFilePath);
+    await db.connect();
+
+    // check if db has two rows, one for root doc and one for subdoc
+    const rows = await db.getAllUpdates();
+    expect(rows.length).toBe(2);
+
+    const rootUpdate = rows.find(row => row.docId === undefined)!.data;
+    const subdocUpdate = rows.find(row => row.docId !== undefined)!.data;
+
+    expect(rootUpdate).toBeDefined();
+    expect(subdocUpdate).toBeDefined();
+
+    // apply updates
+    const rootDoc = new Y.Doc();
+    Y.applyUpdate(rootDoc, rootUpdate);
+
+    // check if root doc has one subdoc
+    expect(rootDoc.subdocs.size).toBe(1);
+
+    // populates subdoc
+    Y.applyUpdate(rootDoc.subdocs.values().next().value, subdocUpdate);
+
+    // check if root doc's meta is correct
+    const meta = rootDoc.getMap('meta').toJSON();
+    expect(meta.workspaceVersion).toBe(1);
+    expect(meta.name).toBe('hiw');
+    expect(meta.pages.length).toBe(1);
+    const pageMeta = meta.pages[0];
+    expect(pageMeta.title).toBe('Welcome to AFFiNEd');
+
+    // get the subdoc through id
+    const subDoc = rootDoc
+      .getMap('spaces')
+      .get(`space:${pageMeta.id}`) as Y.Doc;
+    expect(subDoc).toEqual(rootDoc.subdocs.values().next().value);
+
+    await db.close();
+  });
+});

apps/electron/src/helper/db/migration.ts

@@ -0,0 +1,55 @@
+import { resolve } from 'node:path';
+
+import { migrateToSubdoc } from '@affine/env/blocksuite';
+import { SqliteConnection } from '@affine/native';
+import fs from 'fs-extra';
+import { nanoid } from 'nanoid';
+import * as Y from 'yjs';
+
+import { mainRPC } from '../main-rpc';
+
+export const migrateToSubdocAndReplaceDatabase = async (path: string) => {
+  const db = new SqliteConnection(path);
+  await db.connect();
+
+  const rows = await db.getAllUpdates();
+  const originalDoc = new Y.Doc();
+
+  // 1. apply all updates to the root doc
+  rows.forEach(row => {
+    Y.applyUpdate(originalDoc, row.data);
+  });
+
+  // 2. migrate using migrateToSubdoc
+  const migratedDoc = migrateToSubdoc(originalDoc);
+
+  // 3. replace db rows with the migrated doc
+  await replaceRows(db, migratedDoc, true);
+
+  // 4. close db
+  await db.close();
+};
+
+export const copyToTemp = async (path: string) => {
+  const tmpDirPath = resolve(await mainRPC.getPath('sessionData'), 'tmp');
+  const tmpFilePath = resolve(tmpDirPath, nanoid());
+  await fs.ensureDir(tmpDirPath);
+  await fs.copyFile(path, tmpFilePath);
+  return tmpFilePath;
+};
+
+async function replaceRows(
+  db: SqliteConnection,
+  doc: Y.Doc,
+  isRoot: boolean
+): Promise<void> {
+  const migratedUpdates = Y.encodeStateAsUpdate(doc);
+  const docId = isRoot ? undefined : doc.guid;
+  const rows = [{ data: migratedUpdates, docId: docId }];
+  await db.replaceUpdates(docId, rows);
+  await Promise.all(
+    [...doc.subdocs].map(async subdoc => {
+      await replaceRows(db, subdoc, false);
+    })
+  );
+}

apps/electron/src/helper/dialog/dialog.ts

@@ -1,9 +1,11 @@
 import path from 'node:path';
 
+import { ValidationResult } from '@affine/native';
 import fs from 'fs-extra';
 import { nanoid } from 'nanoid';
 
 import { ensureSQLiteDB } from '../db/ensure-db';
+import { copyToTemp, migrateToSubdocAndReplaceDatabase } from '../db/migration';
 import type { WorkspaceSQLiteDB } from '../db/workspace-db-adapter';
 import { logger } from '../logger';
 import { mainRPC } from '../main-rpc';
@@ -55,6 +57,7 @@ const ErrorMessages = [
   'DB_FILE_ALREADY_LOADED',
   'DB_FILE_PATH_INVALID',
   'DB_FILE_INVALID',
+  'DB_FILE_MIGRATION_FAILED',
   'FILE_ALREADY_EXISTS',
   'UNKNOWN_ERROR',
 ] as const;
@@ -191,27 +194,42 @@ export async function loadDBFile(): Promise<LoadDBFileResult> {
         ],
         message: 'Load Workspace from a AFFiNE file',
       }));
-    const filePath = ret.filePaths?.[0];
-    if (ret.canceled || !filePath) {
+    let originalPath = ret.filePaths?.[0];
+    if (ret.canceled || !originalPath) {
       logger.info('loadDBFile canceled');
       return { canceled: true };
     }
 
     // the imported file should not be in app data dir
-    if (filePath.startsWith(await getWorkspacesBasePath())) {
+    if (originalPath.startsWith(await getWorkspacesBasePath())) {
       logger.warn('loadDBFile: db file in app data dir');
       return { error: 'DB_FILE_PATH_INVALID' };
     }
 
-    if (await dbFileAlreadyLoaded(filePath)) {
+    if (await dbFileAlreadyLoaded(originalPath)) {
       logger.warn('loadDBFile: db file already loaded');
       return { error: 'DB_FILE_ALREADY_LOADED' };
     }
 
     const { SqliteConnection } = await import('@affine/native');
 
-    if (!(await SqliteConnection.validate(filePath))) {
-      // TODO: report invalid db file error?
+    const validationResult = await SqliteConnection.validate(originalPath);
+
+    if (validationResult === ValidationResult.MissingDocIdColumn) {
+      try {
+        const tmpDBPath = await copyToTemp(originalPath);
+        await migrateToSubdocAndReplaceDatabase(tmpDBPath);
+        originalPath = tmpDBPath;
+      } catch (error) {
+        logger.warn(`loadDBFile, migration failed: ${originalPath}`, error);
+        return { error: 'DB_FILE_MIGRATION_FAILED' };
+      }
+    }
+
+    if (
+      validationResult !== ValidationResult.MissingDocIdColumn &&
+      validationResult !== ValidationResult.Valid
+    ) {
       return { error: 'DB_FILE_INVALID' }; // invalid db file
     }
 
@@ -220,9 +238,8 @@ export async function loadDBFile(): Promise<LoadDBFileResult> {
     const internalFilePath = await getWorkspaceDBPath(workspaceId);
 
     await fs.ensureDir(await getWorkspacesBasePath());
-
-    await fs.copy(filePath, internalFilePath);
-    logger.info(`loadDBFile, copy: ${filePath} -> ${internalFilePath}`);
+    await fs.copy(originalPath, internalFilePath);
+    logger.info(`loadDBFile, copy: ${originalPath} -> ${internalFilePath}`);
 
     await storeWorkspaceMeta(workspaceId, {
       id: workspaceId,