feat: store local data to local db (#2037)

apps/electron/layers/main/src/app-state/index.ts

@@ -1,69 +0,0 @@
-import * as os from 'node:os';
-import path from 'node:path';
-
-import { app, shell } from 'electron';
-import { BrowserWindow, ipcMain, nativeTheme } from 'electron';
-import fs from 'fs-extra';
-import { parse } from 'url';
-
-import { isMacOS } from '../../../utils';
-import { getExchangeTokenParams, oauthEndpoint } from './google-auth';
-
-const AFFINE_ROOT = path.join(os.homedir(), '.affine');
-
-fs.ensureDirSync(AFFINE_ROOT);
-
-const logger = console;
-
-export const registerHandlers = () => {
-  ipcMain.handle('ui:theme-change', async (_, theme) => {
-    nativeTheme.themeSource = theme;
-    logger.info('theme change', theme);
-  });
-
-  ipcMain.handle('ui:sidebar-visibility-change', async (_, visible) => {
-    // todo
-    // detect if os is macos
-    if (isMacOS()) {
-      const windows = BrowserWindow.getAllWindows();
-      windows.forEach(w => {
-        // hide window buttons when sidebar is not visible
-        w.setWindowButtonVisibility(visible);
-      });
-      logger.info('sidebar visibility change', visible);
-    }
-  });
-
-  ipcMain.handle('ui:get-google-oauth-code', async () => {
-    logger.info('starting google sign in ...');
-    shell.openExternal(oauthEndpoint);
-
-    return new Promise((resolve, reject) => {
-      const handleOpenUrl = async (_: any, url: string) => {
-        const mainWindow = BrowserWindow.getAllWindows().find(
-          w => !w.isDestroyed()
-        );
-        const urlObj = parse(url.replace('??', '?'), true);
-        if (!mainWindow || !url.startsWith('affine://auth-callback')) return;
-        const code = urlObj.query['code'] as string;
-        if (!code) return;
-
-        logger.info('google sign in code received from callback', code);
-
-        app.removeListener('open-url', handleOpenUrl);
-        resolve(getExchangeTokenParams(code));
-      };
-
-      app.on('open-url', handleOpenUrl);
-
-      setTimeout(() => {
-        reject(new Error('Timed out'));
-        app.removeListener('open-url', handleOpenUrl);
-      }, 30000);
-    });
-  });
-
-  ipcMain.handle('main:env-update', async (_, env, value) => {
-    process.env[env] = value;
-  });
-};

apps/electron/layers/main/src/context.ts

@@ -0,0 +1,9 @@
+import { app } from 'electron';
+import path from 'path';
+
+export const appContext = {
+  appName: app.name,
+  appDataPath: path.join(app.getPath('appData'), app.name),
+};
+
+export type AppContext = typeof appContext;

apps/electron/layers/main/src/data/export.ts

@@ -0,0 +1,34 @@
+import fs from 'fs-extra';
+
+import { logger } from '../../../logger';
+import type { WorkspaceDatabase } from './sqlite';
+
+/**
+ * Start a backup of the database to the given destination.
+ */
+export async function exportDatabase(db: WorkspaceDatabase, dest: string) {
+  await fs.copyFile(db.path, dest);
+  logger.log('export: ', dest);
+}
+
+// export async function startBackup(db: WorkspaceDatabase, dest: string) {
+//   let timeout: NodeJS.Timeout | null;
+//   async function backup() {
+//     await fs.copyFile(db.path, dest);
+//     logger.log('backup: ', dest);
+//   }
+
+//   backup();
+
+//   const _db = await db.sqliteDB$;
+
+//   _db.on('change', () => {
+//     if (timeout) {
+//       clearTimeout(timeout);
+//     }
+//     timeout = setTimeout(async () => {
+//       await backup();
+//       timeout = null;
+//     }, 1000);
+//   });
+// }

apps/electron/layers/main/src/data/sqlite.ts

@@ -0,0 +1,216 @@
+import path from 'node:path';
+
+import fs from 'fs-extra';
+import type { Database } from 'sqlite3';
+import sqlite3Default from 'sqlite3';
+import * as Y from 'yjs';
+
+import { logger } from '../../../logger';
+import type { AppContext } from '../context';
+
+const sqlite3 = sqlite3Default.verbose();
+
+const schemas = [
+  `CREATE TABLE IF NOT EXISTS "updates" (
+  id INTEGER PRIMARY KEY AUTOINCREMENT,
+  data BLOB NOT NULL,
+  timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
+)`,
+  `CREATE TABLE IF NOT EXISTS "blobs" (
+  key TEXT PRIMARY KEY NOT NULL,
+  data BLOB NOT NULL,
+  timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
+)`,
+];
+
+sqlite3.verbose();
+
+interface UpdateRow {
+  id: number;
+  data: Buffer;
+  timestamp: string;
+}
+
+interface BlobRow {
+  key: string;
+  data: Buffer;
+  timestamp: string;
+}
+
+export class WorkspaceDatabase {
+  sqliteDB$: Promise<Database>;
+  ydoc = new Y.Doc();
+  _db: Database | null = null;
+
+  ready: Promise<Uint8Array>;
+
+  constructor(public path: string) {
+    this.sqliteDB$ = this.reconnectDB();
+    logger.log('open db', path);
+
+    this.ydoc.on('update', update => {
+      this.addUpdateToSQLite(update);
+    });
+
+    this.ready = (async () => {
+      const updates = await this.getUpdates();
+      updates.forEach(update => {
+        Y.applyUpdate(this.ydoc, update.data);
+      });
+      return this.getEncodedDocUpdates();
+    })();
+  }
+
+  // release resources
+  destroy = () => {
+    this._db?.close();
+    this.ydoc.destroy();
+  };
+
+  reconnectDB = async () => {
+    logger.log('open db', this.path);
+    if (this._db) {
+      const _db = this._db;
+      await new Promise<void>(res =>
+        _db.close(() => {
+          res();
+        })
+      );
+    }
+    return (this.sqliteDB$ = new Promise(res => {
+      // use cached version?
+      const db = new sqlite3.Database(this.path, error => {
+        if (error) {
+          logger.error('open db error', error);
+        }
+      });
+      this._db = db;
+
+      db.exec(schemas.join(';'), () => {
+        res(db);
+      });
+    }));
+  };
+
+  getEncodedDocUpdates = () => {
+    return Y.encodeStateAsUpdate(this.ydoc);
+  };
+
+  // non-blocking and use yDoc to validate the update
+  // after that, the update is added to the db
+  applyUpdate = (data: Uint8Array) => {
+    Y.applyUpdate(this.ydoc, data);
+    // todo: trim the updates when the number of records is too large
+    // 1. store the current ydoc state in the db
+    // 2. then delete the old updates
+    // yjs-idb will always trim the db for the first time after DB is loaded
+  };
+
+  addBlob = async (key: string, data: Uint8Array) => {
+    const db = await this.sqliteDB$;
+    return new Promise<void>((resolve, reject) => {
+      db.run(
+        'INSERT INTO blobs (key, data) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET data = ?',
+        [key, data, data],
+        err => {
+          if (err) {
+            logger.error('addBlob', err);
+            reject(err);
+          } else {
+            resolve();
+          }
+        }
+      );
+    });
+  };
+
+  getBlob = async (key: string) => {
+    const db = await this.sqliteDB$;
+    return new Promise<Uint8Array | null>((resolve, reject) => {
+      db.get<BlobRow>(
+        'SELECT data FROM blobs WHERE key = ?',
+        [key],
+        (err, row) => {
+          if (err) {
+            logger.error('getBlob', err);
+            reject(err);
+          } else if (!row) {
+            logger.error('getBlob', 'not found');
+            resolve(null);
+          } else {
+            resolve(row.data);
+          }
+        }
+      );
+    });
+  };
+
+  deleteBlob = async (key: string) => {
+    const db = await this.sqliteDB$;
+    return new Promise<void>((resolve, reject) => {
+      db.run('DELETE FROM blobs WHERE key = ?', [key], err => {
+        if (err) {
+          logger.error('deleteBlob', err);
+          reject(err);
+        } else {
+          resolve();
+        }
+      });
+    });
+  };
+
+  getPersistentBlobKeys = async () => {
+    const db = await this.sqliteDB$;
+    return new Promise<string[]>((resolve, reject) => {
+      db.all<BlobRow>('SELECT key FROM blobs', (err, rows) => {
+        if (err) {
+          logger.error('getPersistentBlobKeys', err);
+          reject(err);
+        } else {
+          resolve(rows.map(row => row.key));
+        }
+      });
+    });
+  };
+
+  private getUpdates = async () => {
+    const db = await this.sqliteDB$;
+    return new Promise<{ id: number; data: any }[]>((resolve, reject) => {
+      // do we need to order by id?
+      db.all<UpdateRow>('SELECT * FROM updates', (err, rows) => {
+        if (err) {
+          logger.error('getUpdates', err);
+          reject(err);
+        } else {
+          resolve(rows);
+        }
+      });
+    });
+  };
+
+  private addUpdateToSQLite = async (data: Uint8Array) => {
+    const db = await this.sqliteDB$;
+    return new Promise<void>((resolve, reject) => {
+      db.run('INSERT INTO updates (data) VALUES (?)', [data], err => {
+        if (err) {
+          logger.error('addUpdateToSQLite', err);
+          reject(err);
+        } else {
+          resolve();
+        }
+      });
+    });
+  };
+}
+
+export async function openWorkspaceDatabase(
+  context: AppContext,
+  workspaceId: string
+) {
+  const basePath = path.join(context.appDataPath, 'workspaces', workspaceId);
+  // hmmm.... blocking api but it should be fine, right?
+  await fs.ensureDir(basePath);
+  const dbPath = path.join(basePath, 'storage.db');
+
+  return new WorkspaceDatabase(dbPath);
+}

apps/electron/layers/main/src/data/workspace.ts

@@ -0,0 +1,20 @@
+import path from 'node:path';
+
+import fs from 'fs-extra';
+
+import type { AppContext } from '../context';
+
+export async function listWorkspaces(context: AppContext) {
+  const basePath = path.join(context.appDataPath, 'workspaces');
+  return fs.readdir(basePath);
+}
+
+export async function deleteWorkspace(context: AppContext, id: string) {
+  const basePath = path.join(context.appDataPath, 'workspaces', id);
+  const movedPath = path.join(
+    context.appDataPath,
+    'delete-workspaces',
+    `${id}`
+  );
+  return fs.move(basePath, movedPath);
+}

apps/electron/layers/main/src/handlers.ts

@@ -0,0 +1,184 @@
+import {
+  app,
+  BrowserWindow,
+  dialog,
+  ipcMain,
+  nativeTheme,
+  shell,
+} from 'electron';
+import { parse } from 'url';
+
+import { logger } from '../../logger';
+import { isMacOS } from '../../utils';
+import { appContext } from './context';
+import { exportDatabase } from './data/export';
+import type { WorkspaceDatabase } from './data/sqlite';
+import { openWorkspaceDatabase } from './data/sqlite';
+import { deleteWorkspace, listWorkspaces } from './data/workspace';
+import { getExchangeTokenParams, oauthEndpoint } from './google-auth';
+
+let currentWorkspaceId = '';
+
+const dbMapping = new Map<string, WorkspaceDatabase>();
+
+async function ensureWorkspaceDB(id: string) {
+  let workspaceDB = dbMapping.get(id);
+  if (!workspaceDB) {
+    // hmm... potential race condition?
+    workspaceDB = await openWorkspaceDatabase(appContext, id);
+    dbMapping.set(id, workspaceDB);
+  }
+  await workspaceDB.ready;
+  return workspaceDB;
+}
+
+function registerWorkspaceHandlers() {
+  ipcMain.handle('workspace:list', async _ => {
+    logger.info('list workspaces');
+    return listWorkspaces(appContext);
+  });
+
+  ipcMain.handle('workspace:delete', async (_, id) => {
+    logger.info('delete workspace', id);
+    return deleteWorkspace(appContext, id);
+  });
+}
+
+function registerUIHandlers() {
+  ipcMain.handle('ui:theme-change', async (_, theme) => {
+    nativeTheme.themeSource = theme;
+    logger.info('theme change', theme);
+  });
+
+  ipcMain.handle('ui:sidebar-visibility-change', async (_, visible) => {
+    // todo
+    // detect if os is macos
+    if (isMacOS()) {
+      const windows = BrowserWindow.getAllWindows();
+      windows.forEach(w => {
+        // hide window buttons when sidebar is not visible
+        w.setWindowButtonVisibility(visible);
+      });
+      logger.info('sidebar visibility change', visible);
+    }
+  });
+
+  ipcMain.handle('ui:workspace-change', async (_, workspaceId) => {
+    logger.info('workspace change', workspaceId);
+    currentWorkspaceId = workspaceId;
+  });
+
+  // @deprecated
+  ipcMain.handle('ui:get-google-oauth-code', async () => {
+    logger.info('starting google sign in ...');
+    shell.openExternal(oauthEndpoint);
+
+    return new Promise((resolve, reject) => {
+      const handleOpenUrl = async (_: any, url: string) => {
+        const mainWindow = BrowserWindow.getAllWindows().find(
+          w => !w.isDestroyed()
+        );
+        const urlObj = parse(url.replace('??', '?'), true);
+        if (!mainWindow || !url.startsWith('affine://auth-callback')) return;
+        const code = urlObj.query['code'] as string;
+        if (!code) return;
+
+        logger.info('google sign in code received from callback', code);
+
+        app.removeListener('open-url', handleOpenUrl);
+        resolve(getExchangeTokenParams(code));
+      };
+
+      app.on('open-url', handleOpenUrl);
+
+      setTimeout(() => {
+        reject(new Error('Timed out'));
+        app.removeListener('open-url', handleOpenUrl);
+      }, 30000);
+    });
+  });
+
+  ipcMain.handle('main:env-update', async (_, env, value) => {
+    process.env[env] = value;
+  });
+}
+
+function registerDBHandlers() {
+  app.on('activate', () => {
+    for (const [_, workspaceDB] of dbMapping) {
+      workspaceDB.reconnectDB();
+    }
+  });
+
+  ipcMain.handle('db:get-doc', async (_, id) => {
+    logger.log('main: get doc', id);
+    const workspaceDB = await ensureWorkspaceDB(id);
+    return workspaceDB.getEncodedDocUpdates();
+  });
+
+  ipcMain.handle('db:apply-doc-update', async (_, id, update) => {
+    logger.log('main: apply doc update', id);
+    const workspaceDB = await ensureWorkspaceDB(id);
+    return workspaceDB.applyUpdate(update);
+  });
+
+  ipcMain.handle('db:add-blob', async (_, workspaceId, key, data) => {
+    logger.log('main: add blob', workspaceId, key);
+    const workspaceDB = await ensureWorkspaceDB(workspaceId);
+    return workspaceDB.addBlob(key, data);
+  });
+
+  ipcMain.handle('db:get-blob', async (_, workspaceId, key) => {
+    logger.log('main: get blob', workspaceId, key);
+    const workspaceDB = await ensureWorkspaceDB(workspaceId);
+    return workspaceDB.getBlob(key);
+  });
+
+  ipcMain.handle('db:get-persisted-blobs', async (_, workspaceId) => {
+    logger.log('main: get persisted blob keys', workspaceId);
+    const workspaceDB = await ensureWorkspaceDB(workspaceId);
+    return workspaceDB.getPersistentBlobKeys();
+  });
+
+  ipcMain.handle('db:delete-blob', async (_, workspaceId, key) => {
+    logger.log('main: delete blob', workspaceId, key);
+    const workspaceDB = await ensureWorkspaceDB(workspaceId);
+    return workspaceDB.deleteBlob(key);
+  });
+
+  ipcMain.handle('ui:open-db-folder', async _ => {
+    const workspaceDB = await ensureWorkspaceDB(currentWorkspaceId);
+    logger.log('main: open db folder', workspaceDB.path);
+    shell.showItemInFolder(workspaceDB.path);
+  });
+
+  ipcMain.handle('ui:open-load-db-file-dialog', async () => {
+    // todo
+  });
+
+  ipcMain.handle('ui:open-save-db-file-dialog', async () => {
+    logger.log('main: open save db file dialog', currentWorkspaceId);
+    const workspaceDB = await ensureWorkspaceDB(currentWorkspaceId);
+    const ret = await dialog.showSaveDialog({
+      properties: ['showOverwriteConfirmation'],
+      title: 'Save Workspace',
+      buttonLabel: 'Save',
+      defaultPath: currentWorkspaceId + '.db',
+      message: 'Save Workspace as SQLite Database',
+    });
+    const filePath = ret.filePath;
+    if (ret.canceled || !filePath) {
+      return null;
+    }
+
+    await exportDatabase(workspaceDB, filePath);
+    shell.showItemInFolder(filePath);
+    return filePath;
+  });
+}
+
+export const registerHandlers = () => {
+  registerWorkspaceHandlers();
+  registerUIHandlers();
+  registerDBHandlers();
+};

apps/electron/layers/main/src/index.ts

@@ -3,7 +3,8 @@ import './security-restrictions';
 import { app } from 'electron';
 import path from 'path';
 
-import { registerHandlers } from './app-state';
+import { logger } from '../../logger';
+import { registerHandlers } from './handlers';
 import { restoreOrCreateWindow } from './main-window';
 import { registerProtocol } from './protocol';
 
@@ -22,6 +23,7 @@ if (process.defaultApp) {
  */
 const isSingleInstance = app.requestSingleInstanceLock();
 if (!isSingleInstance) {
+  logger.info('Another instance is running, exiting...');
   app.quit();
   process.exit(0);
 }

apps/electron/layers/main/src/main-window.ts

@@ -2,6 +2,7 @@ import { BrowserWindow, nativeTheme } from 'electron';
 import electronWindowState from 'electron-window-state';
 import { join } from 'path';
 
+import { logger } from '../../logger';
 import { isMacOS } from '../../utils';
 
 const IS_DEV = process.env.NODE_ENV === 'development';
@@ -87,7 +88,7 @@ export async function restoreOrCreateWindow() {
     browserWindow.restore();
   }
 
-  browserWindow.focus();
+  logger.info('Create main window');
 
   return browserWindow;
 }