refactor(electron): sqlite db data workflow (remove symlink & fs watcher) (#2491)

apps/electron/layers/main/src/db/base-db-adapter.ts

@@ -0,0 +1,152 @@
+import assert from 'assert';
+import type { Database } from 'better-sqlite3';
+import sqlite from 'better-sqlite3';
+
+import { logger } from '../logger';
+
+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
+)`,
+];
+
+interface UpdateRow {
+  id: number;
+  data: Buffer;
+  timestamp: string;
+}
+
+interface BlobRow {
+  key: string;
+  data: Buffer;
+  timestamp: string;
+}
+
+/**
+ * A base class for SQLite DB adapter that provides basic methods around updates & blobs
+ */
+export abstract class BaseSQLiteAdapter {
+  db: Database | null = null;
+  abstract role: string;
+
+  constructor(public path: string) {}
+
+  ensureTables() {
+    assert(this.db, 'db is not connected');
+    this.db.exec(schemas.join(';'));
+  }
+
+  // todo: what if SQLite DB wrapper later is not sync?
+  connect() {
+    if (this.db) {
+      return;
+    }
+    logger.log(`[SQLiteAdapter][${this.role}] open db`, this.path);
+    const db = (this.db = sqlite(this.path));
+    this.ensureTables();
+    return db;
+  }
+
+  destroy() {
+    this.db?.close();
+    this.db = null;
+  }
+
+  addBlob(key: string, data: Uint8Array) {
+    try {
+      assert(this.db, 'db is not connected');
+      const statement = this.db.prepare(
+        'INSERT INTO blobs (key, data) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET data = ?'
+      );
+      statement.run(key, data, data);
+      return key;
+    } catch (error) {
+      logger.error('addBlob', error);
+    }
+  }
+
+  getBlob(key: string) {
+    try {
+      assert(this.db, 'db is not connected');
+      const statement = this.db.prepare('SELECT data FROM blobs WHERE key = ?');
+      const row = statement.get(key) as BlobRow;
+      if (!row) {
+        return null;
+      }
+      return row.data;
+    } catch (error) {
+      logger.error('getBlob', error);
+      return null;
+    }
+  }
+
+  deleteBlob(key: string) {
+    try {
+      assert(this.db, 'db is not connected');
+      const statement = this.db.prepare('DELETE FROM blobs WHERE key = ?');
+      statement.run(key);
+    } catch (error) {
+      logger.error('deleteBlob', error);
+    }
+  }
+
+  getBlobKeys() {
+    try {
+      assert(this.db, 'db is not connected');
+      const statement = this.db.prepare('SELECT key FROM blobs');
+      const rows = statement.all() as BlobRow[];
+      return rows.map(row => row.key);
+    } catch (error) {
+      logger.error('getBlobKeys', error);
+      return [];
+    }
+  }
+
+  getUpdates() {
+    try {
+      assert(this.db, 'db is not connected');
+      const statement = this.db.prepare('SELECT * FROM updates');
+      const rows = statement.all() as UpdateRow[];
+      return rows;
+    } catch (error) {
+      logger.error('getUpdates', error);
+      return [];
+    }
+  }
+
+  // add a single update to SQLite
+  addUpdateToSQLite(updates: Uint8Array[]) {
+    // batch write instead write per key stroke?
+    try {
+      assert(this.db, 'db is not connected');
+      const start = performance.now();
+      const statement = this.db.prepare(
+        'INSERT INTO updates (data) VALUES (?)'
+      );
+      const insertMany = this.db.transaction(updates => {
+        for (const d of updates) {
+          statement.run(d);
+        }
+      });
+
+      insertMany(updates);
+
+      logger.debug(
+        `[SQLiteAdapter][${this.role}] addUpdateToSQLite`,
+        'length:',
+        updates.length,
+        performance.now() - start,
+        'ms'
+      );
+    } catch (error) {
+      logger.error('addUpdateToSQLite', error);
+    }
+  }
+}

apps/electron/layers/main/src/db/ensure-db.ts

@@ -0,0 +1,94 @@
+import { app } from 'electron';
+import {
+  defer,
+  firstValueFrom,
+  from,
+  fromEvent,
+  interval,
+  Observable,
+} from 'rxjs';
+import {
+  distinctUntilChanged,
+  filter,
+  ignoreElements,
+  last,
+  map,
+  shareReplay,
+  startWith,
+  switchMap,
+  takeUntil,
+  tap,
+} from 'rxjs/operators';
+
+import { appContext } from '../context';
+import { logger } from '../logger';
+import { getWorkspaceMeta$ } from '../workspace';
+import { SecondaryWorkspaceSQLiteDB } from './secondary-db';
+import type { WorkspaceSQLiteDB } from './workspace-db-adapter';
+import { openWorkspaceDatabase } from './workspace-db-adapter';
+
+const db$Map = new Map<string, Observable<WorkspaceSQLiteDB>>();
+
+const beforeQuit$ = defer(() => fromEvent(app, 'before-quit'));
+
+function getWorkspaceDB$(id: string) {
+  if (!db$Map.has(id)) {
+    db$Map.set(
+      id,
+      from(openWorkspaceDatabase(appContext, id)).pipe(
+        switchMap(db => {
+          return startPollingSecondaryDB(db).pipe(
+            ignoreElements(),
+            startWith(db),
+            takeUntil(beforeQuit$),
+            tap({
+              complete: () => {
+                logger.info('[ensureSQLiteDB] close db connection');
+                db.destroy();
+                db$Map.delete(id);
+              },
+            })
+          );
+        }),
+        shareReplay(1)
+      )
+    );
+  }
+  return db$Map.get(id)!;
+}
+
+// fixme: this function has issue on registering multiple times...
+function startPollingSecondaryDB(db: WorkspaceSQLiteDB) {
+  const meta$ = getWorkspaceMeta$(db.workspaceId);
+  const secondaryDB$ = meta$.pipe(
+    map(meta => meta?.secondaryDBPath),
+    distinctUntilChanged(),
+    filter((p): p is string => !!p),
+    switchMap(path => {
+      const secondaryDB = new SecondaryWorkspaceSQLiteDB(path, db);
+      return new Observable<SecondaryWorkspaceSQLiteDB>(observer => {
+        observer.next(secondaryDB);
+        return () => {
+          secondaryDB.destroy();
+        };
+      });
+    })
+  );
+
+  // pull every 30 seconds
+  const poll$ = interval(30000).pipe(
+    switchMap(() => secondaryDB$),
+    tap({
+      next: secondaryDB => {
+        secondaryDB.pull();
+      },
+    }),
+    shareReplay(1)
+  );
+
+  return poll$.pipe(takeUntil(db.update$.pipe(last())), shareReplay(1));
+}
+
+export function ensureSQLiteDB(id: string) {
+  return firstValueFrom(getWorkspaceDB$(id));
+}

apps/electron/layers/main/src/db/helper.ts

@@ -0,0 +1,38 @@
+import type { Database } from 'better-sqlite3';
+import sqlite from 'better-sqlite3';
+
+import { logger } from '../logger';
+
+export function isValidateDB(db: Database) {
+  // check if db has two tables, one for updates and one for blobs
+  const statement = db.prepare(
+    `SELECT name FROM sqlite_schema WHERE type='table'`
+  );
+  const rows = statement.all() as { name: string }[];
+  const tableNames = rows.map(row => row.name);
+  if (!tableNames.includes('updates') || !tableNames.includes('blobs')) {
+    return false;
+  }
+}
+
+export function isValidDBFile(path: string) {
+  let db: Database | null = null;
+  try {
+    db = sqlite(path);
+    // check if db has two tables, one for updates and one for blobs
+    const statement = db.prepare(
+      `SELECT name FROM sqlite_schema WHERE type='table'`
+    );
+    const rows = statement.all() as { name: string }[];
+    const tableNames = rows.map(row => row.name);
+    if (!tableNames.includes('updates') || !tableNames.includes('blobs')) {
+      return false;
+    }
+    return true;
+  } catch (error) {
+    logger.error('isValidDBFile', error);
+    return false;
+  } finally {
+    db?.close();
+  }
+}

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

@@ -0,0 +1,48 @@
+import { appContext } from '../context';
+import type { MainEventListener, NamespaceHandlers } from '../type';
+import { ensureSQLiteDB } from './ensure-db';
+import { dbSubjects } from './subjects';
+
+export * from './ensure-db';
+export * from './subjects';
+
+export const dbHandlers = {
+  getDocAsUpdates: async (_, id: string) => {
+    const workspaceDB = await ensureSQLiteDB(id);
+    return workspaceDB.getDocAsUpdates();
+  },
+  applyDocUpdate: async (_, id: string, update: Uint8Array) => {
+    const workspaceDB = await ensureSQLiteDB(id);
+    return workspaceDB.applyUpdate(update);
+  },
+  addBlob: async (_, workspaceId: string, key: string, data: Uint8Array) => {
+    const workspaceDB = await ensureSQLiteDB(workspaceId);
+    return workspaceDB.addBlob(key, data);
+  },
+  getBlob: async (_, workspaceId: string, key: string) => {
+    const workspaceDB = await ensureSQLiteDB(workspaceId);
+    return workspaceDB.getBlob(key);
+  },
+  deleteBlob: async (_, workspaceId: string, key: string) => {
+    const workspaceDB = await ensureSQLiteDB(workspaceId);
+    return workspaceDB.deleteBlob(key);
+  },
+  getBlobKeys: async (_, workspaceId: string) => {
+    const workspaceDB = await ensureSQLiteDB(workspaceId);
+    return workspaceDB.getBlobKeys();
+  },
+  getDefaultStorageLocation: async () => {
+    return appContext.appDataPath;
+  },
+} satisfies NamespaceHandlers;
+
+export const dbEvents = {
+  onExternalUpdate: (
+    fn: (update: { workspaceId: string; update: Uint8Array }) => void
+  ) => {
+    const sub = dbSubjects.externalUpdate.subscribe(fn);
+    return () => {
+      sub.unsubscribe();
+    };
+  },
+} satisfies Record<string, MainEventListener>;

apps/electron/layers/main/src/db/secondary-db.ts

@@ -0,0 +1,191 @@
+import { debounce } from 'lodash-es';
+import * as Y from 'yjs';
+
+import type { AppContext } from '../context';
+import { logger } from '../logger';
+import type { YOrigin } from '../type';
+import { getWorkspaceMeta } from '../workspace';
+import { BaseSQLiteAdapter } from './base-db-adapter';
+import type { WorkspaceSQLiteDB } from './workspace-db-adapter';
+
+const FLUSH_WAIT_TIME = 5000;
+const FLUSH_MAX_WAIT_TIME = 10000;
+
+export class SecondaryWorkspaceSQLiteDB extends BaseSQLiteAdapter {
+  role = 'secondary';
+  yDoc = new Y.Doc();
+  firstConnected = false;
+
+  updateQueue: Uint8Array[] = [];
+
+  unsubscribers = new Set<() => void>();
+
+  constructor(
+    public override path: string,
+    public upstream: WorkspaceSQLiteDB
+  ) {
+    super(path);
+    this.setupAndListen();
+  }
+
+  close() {
+    this.db?.close();
+    this.db = null;
+  }
+
+  override destroy() {
+    this.unsubscribers.forEach(unsub => unsub());
+    this.db?.close();
+    this.yDoc.destroy();
+    this.close();
+  }
+
+  get workspaceId() {
+    return this.upstream.workspaceId;
+  }
+
+  // do not update db immediately, instead, push to a queue
+  // and flush the queue in a future time
+  addUpdateToUpdateQueue(update: Uint8Array) {
+    this.updateQueue.push(update);
+    this.debouncedFlush();
+  }
+
+  flushUpdateQueue() {
+    logger.debug(
+      'flushUpdateQueue',
+      this.workspaceId,
+      'queue',
+      this.updateQueue.length
+    );
+    const updates = [...this.updateQueue];
+    this.updateQueue = [];
+    this.connect();
+    this.addUpdateToSQLite(updates);
+    this.close();
+  }
+
+  // flush after 5s, but will not wait for more than 10s
+  debouncedFlush = debounce(this.flushUpdateQueue, FLUSH_WAIT_TIME, {
+    maxWait: FLUSH_MAX_WAIT_TIME,
+  });
+
+  runCounter = 0;
+
+  // wrap the fn with connect and close
+  // it only works for sync functions
+  run = (fn: () => void) => {
+    try {
+      if (this.runCounter === 0) {
+        this.connect();
+      }
+      this.runCounter++;
+      fn();
+    } catch (err) {
+      logger.error(err);
+    } finally {
+      this.runCounter--;
+      if (this.runCounter === 0) {
+        this.close();
+      }
+    }
+  };
+
+  setupAndListen() {
+    if (this.firstConnected) {
+      return;
+    }
+    this.firstConnected = true;
+
+    const onUpstreamUpdate = (update: Uint8Array, origin: YOrigin) => {
+      if (origin === 'renderer') {
+        // update to upstream yDoc should be replicated to self yDoc
+        this.applyUpdate(update, 'upstream');
+      }
+    };
+
+    const onSelfUpdate = (update: Uint8Array, origin: YOrigin) => {
+      // for self update from upstream, we need to push it to external DB
+      if (origin === 'upstream') {
+        this.addUpdateToUpdateQueue(update);
+      }
+
+      if (origin === 'self') {
+        this.upstream.applyUpdate(update, 'external');
+      }
+    };
+
+    // listen to upstream update
+    this.upstream.yDoc.on('update', onUpstreamUpdate);
+    this.yDoc.on('update', onSelfUpdate);
+
+    this.unsubscribers.add(() => {
+      this.upstream.yDoc.off('update', onUpstreamUpdate);
+      this.yDoc.off('update', onSelfUpdate);
+    });
+
+    this.run(() => {
+      // apply all updates from upstream
+      const upstreamUpdate = this.upstream.getDocAsUpdates();
+      // to initialize the yDoc, we need to apply all updates from the db
+      this.applyUpdate(upstreamUpdate, 'upstream');
+
+      this.pull();
+    });
+  }
+
+  applyUpdate = (data: Uint8Array, origin: YOrigin = 'upstream') => {
+    Y.applyUpdate(this.yDoc, data, origin);
+  };
+
+  // TODO: have a better solution to handle blobs
+  syncBlobs() {
+    this.run(() => {
+      // pull blobs
+      const blobsKeys = this.getBlobKeys();
+      const upstreamBlobsKeys = this.upstream.getBlobKeys();
+      // put every missing blob to upstream
+      for (const key of blobsKeys) {
+        if (!upstreamBlobsKeys.includes(key)) {
+          const blob = this.getBlob(key);
+          if (blob) {
+            this.upstream.addBlob(key, blob);
+            logger.debug('syncBlobs', this.workspaceId, key);
+          }
+        }
+      }
+    });
+  }
+
+  /**
+   * pull from external DB file and apply to embedded yDoc
+   * workflow:
+   * - connect to external db
+   * - get updates
+   * - apply updates to local yDoc
+   * - get blobs and put new blobs to upstream
+   * - disconnect
+   */
+  pull() {
+    this.run(() => {
+      // TODO: no need to get all updates, just get the latest ones (using a cursor, etc)?
+      const updates = this.getUpdates().map(update => update.data);
+      Y.transact(this.yDoc, () => {
+        updates.forEach(update => {
+          this.applyUpdate(update, 'self');
+        });
+      });
+      logger.debug('pull external updates', this.path, updates.length);
+
+      this.syncBlobs();
+    });
+  }
+}
+
+export async function getSecondaryWorkspaceDBPath(
+  context: AppContext,
+  workspaceId: string
+) {
+  const meta = await getWorkspaceMeta(context, workspaceId);
+  return meta?.secondaryDBPath;
+}

apps/electron/layers/main/src/db/subjects.ts

@@ -0,0 +1,7 @@
+import { Subject } from 'rxjs';
+
+export const dbSubjects = {
+  // emit workspace id when the db file is missing
+  fileMissing: new Subject<string>(),
+  externalUpdate: new Subject<{ workspaceId: string; update: Uint8Array }>(),
+};

apps/electron/layers/main/src/db/workspace-db-adapter.ts

@@ -0,0 +1,105 @@
+import { Subject } from 'rxjs';
+import * as Y from 'yjs';
+
+import type { AppContext } from '../context';
+import { logger } from '../logger';
+import type { YOrigin } from '../type';
+import { getWorkspaceMeta } from '../workspace';
+import { BaseSQLiteAdapter } from './base-db-adapter';
+import { dbSubjects } from './subjects';
+
+export class WorkspaceSQLiteDB extends BaseSQLiteAdapter {
+  role = 'primary';
+  yDoc = new Y.Doc();
+  firstConnected = false;
+
+  update$ = new Subject<void>();
+
+  constructor(public override path: string, public workspaceId: string) {
+    super(path);
+  }
+
+  override destroy() {
+    this.db?.close();
+    this.yDoc.destroy();
+
+    // when db is closed, we can safely remove it from ensure-db list
+    this.update$.complete();
+  }
+
+  getWorkspaceName = () => {
+    return this.yDoc.getMap('space:meta').get('name') as string;
+  };
+
+  override connect() {
+    const db = super.connect();
+
+    if (!this.firstConnected) {
+      this.yDoc.on('update', (update: Uint8Array, origin: YOrigin) => {
+        if (origin !== 'self') {
+          this.addUpdateToSQLite([update]);
+        }
+        if (origin === 'external') {
+          logger.debug('external update', this.workspaceId);
+          dbSubjects.externalUpdate.next({
+            workspaceId: this.workspaceId,
+            update,
+          });
+        }
+      });
+    }
+
+    const updates = this.getUpdates();
+    // to initialize the yDoc, we need to apply all updates from the db
+    Y.transact(this.yDoc, () => {
+      updates.forEach(update => {
+        this.applyUpdate(update.data, 'self');
+      });
+    });
+
+    this.firstConnected = true;
+    this.update$.next();
+
+    return db;
+  }
+
+  getDocAsUpdates = () => {
+    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, origin: YOrigin = 'renderer') => {
+    // 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
+    Y.applyUpdate(this.yDoc, data, origin);
+  };
+
+  override addBlob(key: string, value: Uint8Array) {
+    const res = super.addBlob(key, value);
+    this.update$.next();
+    return res;
+  }
+
+  override deleteBlob(key: string) {
+    super.deleteBlob(key);
+    this.update$.next();
+  }
+
+  override addUpdateToSQLite(data: Uint8Array[]) {
+    super.addUpdateToSQLite(data);
+    this.update$.next();
+  }
+}
+
+export async function openWorkspaceDatabase(
+  context: AppContext,
+  workspaceId: string
+) {
+  const meta = await getWorkspaceMeta(context, workspaceId);
+  const db = new WorkspaceSQLiteDB(meta.mainDBPath, workspaceId);
+  await db.connect();
+  return db;
+}