feat(nbstore): improve nbstore (#9512)

packages/common/nbstore/src/impls/sqlite/blob.ts

@@ -1,11 +1,15 @@
 import { share } from '../../connection';
 import { type BlobRecord, BlobStorageBase } from '../../storage';
-import { NativeDBConnection } from './db';
+import { NativeDBConnection, type SqliteNativeDBOptions } from './db';
 
 export class SqliteBlobStorage extends BlobStorageBase {
-  override connection = share(
-    new NativeDBConnection(this.peer, this.spaceType, this.spaceId)
-  );
+  static readonly identifier = 'SqliteBlobStorage';
+
+  override connection = share(new NativeDBConnection(this.options));
+
+  constructor(private readonly options: SqliteNativeDBOptions) {
+    super();
+  }
 
   get db() {
     return this.connection.apis;

packages/common/nbstore/src/impls/sqlite/db.ts

@@ -1,9 +1,81 @@
-import { apis } from '@affine/electron-api';
-
 import { AutoReconnectConnection } from '../../connection';
-import { type SpaceType, universalId } from '../../storage';
+import type {
+  BlobRecord,
+  DocClock,
+  DocRecord,
+  ListedBlobRecord,
+} from '../../storage';
+import { type SpaceType, universalId } from '../../utils/universal-id';
 
-type NativeDBApis = NonNullable<typeof apis>['nbstore'] extends infer APIs
+export interface SqliteNativeDBOptions {
+  readonly flavour: string;
+  readonly type: SpaceType;
+  readonly id: string;
+}
+
+export type NativeDBApis = {
+  connect(id: string): Promise<void>;
+  disconnect(id: string): Promise<void>;
+  pushUpdate(id: string, docId: string, update: Uint8Array): Promise<Date>;
+  getDocSnapshot(id: string, docId: string): Promise<DocRecord | null>;
+  setDocSnapshot(id: string, snapshot: DocRecord): Promise<boolean>;
+  getDocUpdates(id: string, docId: string): Promise<DocRecord[]>;
+  markUpdatesMerged(
+    id: string,
+    docId: string,
+    updates: Date[]
+  ): Promise<number>;
+  deleteDoc(id: string, docId: string): Promise<void>;
+  getDocClocks(
+    id: string,
+    after?: Date | undefined | null
+  ): Promise<DocClock[]>;
+  getDocClock(id: string, docId: string): Promise<DocClock | null>;
+  getBlob(id: string, key: string): Promise<BlobRecord | null>;
+  setBlob(id: string, blob: BlobRecord): Promise<void>;
+  deleteBlob(id: string, key: string, permanently: boolean): Promise<void>;
+  releaseBlobs(id: string): Promise<void>;
+  listBlobs(id: string): Promise<ListedBlobRecord[]>;
+  getPeerRemoteClocks(id: string, peer: string): Promise<DocClock[]>;
+  getPeerRemoteClock(
+    id: string,
+    peer: string,
+    docId: string
+  ): Promise<DocClock>;
+  setPeerRemoteClock(
+    id: string,
+    peer: string,
+    docId: string,
+    clock: Date
+  ): Promise<void>;
+  getPeerPulledRemoteClocks(id: string, peer: string): Promise<DocClock[]>;
+  getPeerPulledRemoteClock(
+    id: string,
+    peer: string,
+    docId: string
+  ): Promise<DocClock>;
+  setPeerPulledRemoteClock(
+    id: string,
+    peer: string,
+    docId: string,
+    clock: Date
+  ): Promise<void>;
+  getPeerPushedClocks(id: string, peer: string): Promise<DocClock[]>;
+  getPeerPushedClock(
+    id: string,
+    peer: string,
+    docId: string
+  ): Promise<DocClock>;
+  setPeerPushedClock(
+    id: string,
+    peer: string,
+    docId: string,
+    clock: Date
+  ): Promise<void>;
+  clearClocks(id: string): Promise<void>;
+};
+
+type NativeDBApisWrapper = NativeDBApis extends infer APIs
   ? {
       [K in keyof APIs]: APIs[K] extends (...args: any[]) => any
         ? Parameters<APIs[K]> extends [string, ...infer Rest]
@@ -13,49 +85,56 @@ type NativeDBApis = NonNullable<typeof apis>['nbstore'] extends infer APIs
     }
   : never;
 
-export class NativeDBConnection extends AutoReconnectConnection<void> {
-  readonly apis: NativeDBApis;
+let apis: NativeDBApis | null = null;
 
-  constructor(
-    private readonly peer: string,
-    private readonly type: SpaceType,
-    private readonly id: string
-  ) {
+export function bindNativeDBApis(a: NativeDBApis) {
+  apis = a;
+}
+
+export class NativeDBConnection extends AutoReconnectConnection<void> {
+  readonly apis: NativeDBApisWrapper;
+
+  readonly flavour = this.options.flavour;
+  readonly type = this.options.type;
+  readonly id = this.options.id;
+
+  constructor(private readonly options: SqliteNativeDBOptions) {
     super();
+
     if (!apis) {
-      throw new Error('Not in electron context.');
+      throw new Error('Not in native context.');
     }
 
-    this.apis = this.bindApis(apis.nbstore);
+    this.apis = this.warpApis(apis);
   }
 
   override get shareId(): string {
-    return `sqlite:${this.peer}:${this.type}:${this.id}`;
+    return `sqlite:${this.flavour}:${this.type}:${this.id}`;
   }
 
-  bindApis(originalApis: NonNullable<typeof apis>['nbstore']): NativeDBApis {
+  warpApis(originalApis: NativeDBApis): NativeDBApisWrapper {
     const id = universalId({
-      peer: this.peer,
+      peer: this.flavour,
       type: this.type,
       id: this.id,
     });
-    return new Proxy(originalApis, {
-      get: (target, key: keyof NativeDBApis) => {
-        const v = target[key];
-        if (typeof v !== 'function') {
-          return v;
-        }
+    return new Proxy(
+      {},
+      {
+        get: (_target, key: keyof NativeDBApisWrapper) => {
+          const v = originalApis[key];
 
-        return async (...args: any[]) => {
-          return v.call(
-            originalApis,
-            id,
-            // @ts-expect-error I don't know why it complains ts(2556)
-            ...args
-          );
-        };
-      },
-    }) as unknown as NativeDBApis;
+          return async (...args: any[]) => {
+            return v.call(
+              originalApis,
+              id,
+              // @ts-expect-error I don't know why it complains ts(2556)
+              ...args
+            );
+          };
+        },
+      }
+    ) as unknown as NativeDBApisWrapper;
   }
 
   override async doConnect() {
@@ -63,7 +142,7 @@ export class NativeDBConnection extends AutoReconnectConnection<void> {
   }
 
   override doDisconnect() {
-    this.apis.close().catch(err => {
+    this.apis.disconnect().catch(err => {
       console.error('NativeDBConnection close failed', err);
     });
   }

packages/common/nbstore/src/impls/sqlite/doc.ts

@@ -1,54 +1,82 @@
 import { share } from '../../connection';
-import { type DocClock, DocStorageBase, type DocUpdate } from '../../storage';
-import { NativeDBConnection } from './db';
+import {
+  type DocClocks,
+  type DocRecord,
+  DocStorageBase,
+  type DocUpdate,
+} from '../../storage';
+import { NativeDBConnection, type SqliteNativeDBOptions } from './db';
 
-export class SqliteDocStorage extends DocStorageBase {
-  override connection = share(
-    new NativeDBConnection(this.peer, this.spaceType, this.spaceId)
-  );
+export class SqliteDocStorage extends DocStorageBase<SqliteNativeDBOptions> {
+  static readonly identifier = 'SqliteDocStorage';
+  override connection = share(new NativeDBConnection(this.options));
 
   get db() {
     return this.connection.apis;
   }
 
-  override async getDoc(docId: string) {
-    return this.db.getDoc(docId);
-  }
-
   override async pushDocUpdate(update: DocUpdate) {
-    return this.db.pushDocUpdate(update);
+    const timestamp = await this.db.pushUpdate(update.docId, update.bin);
+
+    this.emit(
+      'update',
+      {
+        docId: update.docId,
+        bin: update.bin,
+        timestamp,
+        editor: update.editor,
+      },
+      origin
+    );
+
+    return { docId: update.docId, timestamp };
   }
 
   override async deleteDoc(docId: string) {
-    return this.db.deleteDoc(docId);
+    await this.db.deleteDoc(docId);
   }
 
   override async getDocTimestamps(after?: Date) {
-    return this.db.getDocTimestamps(after ? new Date(after) : undefined);
+    const clocks = await this.db.getDocClocks(after);
+
+    return clocks.reduce((ret, cur) => {
+      ret[cur.docId] = cur.timestamp;
+      return ret;
+    }, {} as DocClocks);
   }
 
-  override getDocTimestamp(docId: string): Promise<DocClock | null> {
-    return this.db.getDocTimestamp(docId);
+  override async getDocTimestamp(docId: string) {
+    return this.db.getDocClock(docId);
   }
 
-  protected override async getDocSnapshot() {
-    // handled in db
-    // see electron/src/helper/nbstore/doc.ts
-    return null;
+  protected override async getDocSnapshot(docId: string) {
+    const snapshot = await this.db.getDocSnapshot(docId);
+
+    if (!snapshot) {
+      return null;
+    }
+
+    return snapshot;
   }
 
-  protected override async setDocSnapshot(): Promise<boolean> {
-    // handled in db
-    return true;
+  protected override async setDocSnapshot(
+    snapshot: DocRecord
+  ): Promise<boolean> {
+    return this.db.setDocSnapshot({
+      docId: snapshot.docId,
+      bin: snapshot.bin,
+      timestamp: snapshot.timestamp,
+    });
   }
 
-  protected override async getDocUpdates() {
-    // handled in db
-    return [];
+  protected override async getDocUpdates(docId: string) {
+    return this.db.getDocUpdates(docId);
   }
 
-  protected override markUpdatesMerged() {
-    // handled in db
-    return Promise.resolve(0);
+  protected override markUpdatesMerged(docId: string, updates: DocRecord[]) {
+    return this.db.markUpdatesMerged(
+      docId,
+      updates.map(update => update.timestamp)
+    );
   }
 }

packages/common/nbstore/src/impls/sqlite/index.ts

@@ -1,3 +1,16 @@
+import type { StorageConstructor } from '..';
+import { SqliteBlobStorage } from './blob';
+import { SqliteDocStorage } from './doc';
+import { SqliteSyncStorage } from './sync';
+
 export * from './blob';
+export { bindNativeDBApis, type NativeDBApis } from './db';
 export * from './doc';
 export * from './sync';
+export * from './v1';
+
+export const sqliteStorages = [
+  SqliteDocStorage,
+  SqliteBlobStorage,
+  SqliteSyncStorage,
+] satisfies StorageConstructor[];

packages/common/nbstore/src/impls/sqlite/sync.ts

@@ -1,18 +1,26 @@
 import { share } from '../../connection';
-import { BasicSyncStorage, type DocClock } from '../../storage';
-import { NativeDBConnection } from './db';
+import { type DocClock, SyncStorageBase } from '../../storage';
+import { NativeDBConnection, type SqliteNativeDBOptions } from './db';
 
-export class SqliteSyncStorage extends BasicSyncStorage {
-  override connection = share(
-    new NativeDBConnection(this.peer, this.spaceType, this.spaceId)
-  );
+export class SqliteSyncStorage extends SyncStorageBase {
+  static readonly identifier = 'SqliteSyncStorage';
+
+  override connection = share(new NativeDBConnection(this.options));
+
+  constructor(private readonly options: SqliteNativeDBOptions) {
+    super();
+  }
 
   get db() {
     return this.connection.apis;
   }
 
   override async getPeerRemoteClocks(peer: string) {
-    return this.db.getPeerRemoteClocks(peer);
+    return this.db
+      .getPeerRemoteClocks(peer)
+      .then(clocks =>
+        Object.fromEntries(clocks.map(clock => [clock.docId, clock.timestamp]))
+      );
   }
 
   override async getPeerRemoteClock(peer: string, docId: string) {
@@ -20,11 +28,15 @@ export class SqliteSyncStorage extends BasicSyncStorage {
   }
 
   override async setPeerRemoteClock(peer: string, clock: DocClock) {
-    await this.db.setPeerRemoteClock(peer, clock);
+    await this.db.setPeerRemoteClock(peer, clock.docId, clock.timestamp);
   }
 
   override async getPeerPulledRemoteClocks(peer: string) {
-    return this.db.getPeerPulledRemoteClocks(peer);
+    return this.db
+      .getPeerPulledRemoteClocks(peer)
+      .then(clocks =>
+        Object.fromEntries(clocks.map(clock => [clock.docId, clock.timestamp]))
+      );
   }
 
   override async getPeerPulledRemoteClock(peer: string, docId: string) {
@@ -32,11 +44,15 @@ export class SqliteSyncStorage extends BasicSyncStorage {
   }
 
   override async setPeerPulledRemoteClock(peer: string, clock: DocClock) {
-    await this.db.setPeerPulledRemoteClock(peer, clock);
+    await this.db.setPeerPulledRemoteClock(peer, clock.docId, clock.timestamp);
   }
 
   override async getPeerPushedClocks(peer: string) {
-    return this.db.getPeerPushedClocks(peer);
+    return this.db
+      .getPeerPushedClocks(peer)
+      .then(clocks =>
+        Object.fromEntries(clocks.map(clock => [clock.docId, clock.timestamp]))
+      );
   }
 
   override async getPeerPushedClock(peer: string, docId: string) {
@@ -44,7 +60,7 @@ export class SqliteSyncStorage extends BasicSyncStorage {
   }
 
   override async setPeerPushedClock(peer: string, clock: DocClock) {
-    await this.db.setPeerPushedClock(peer, clock);
+    await this.db.setPeerPushedClock(peer, clock.docId, clock.timestamp);
   }
 
   override async clearClocks() {

packages/common/nbstore/src/impls/sqlite/v1/blob.ts

@@ -1,7 +1,7 @@
-import { apis } from '@affine/electron-api';
-
 import { DummyConnection } from '../../../connection';
 import { BlobStorageBase } from '../../../storage';
+import type { SpaceType } from '../../../utils/universal-id';
+import { apis } from './db';
 
 /**
  * @deprecated readonly
@@ -9,18 +9,22 @@ import { BlobStorageBase } from '../../../storage';
 export class SqliteV1BlobStorage extends BlobStorageBase {
   override connection = new DummyConnection();
 
-  get db() {
+  constructor(private readonly options: { type: SpaceType; id: string }) {
+    super();
+  }
+
+  private get db() {
     if (!apis) {
       throw new Error('Not in electron context.');
     }
 
-    return apis.db;
+    return apis;
   }
 
   override async get(key: string) {
     const data: Uint8Array | null = await this.db.getBlob(
-      this.spaceType,
-      this.spaceId,
+      this.options.type,
+      this.options.id,
       key
     );
 
@@ -38,12 +42,12 @@ export class SqliteV1BlobStorage extends BlobStorageBase {
 
   override async delete(key: string, permanently: boolean) {
     if (permanently) {
-      await this.db.deleteBlob(this.spaceType, this.spaceId, key);
+      await this.db.deleteBlob(this.options.type, this.options.id, key);
     }
   }
 
   override async list() {
-    const keys = await this.db.getBlobKeys(this.spaceType, this.spaceId);
+    const keys = await this.db.getBlobKeys(this.options.type, this.options.id);
 
     return keys.map(key => ({
       key,

packages/common/nbstore/src/impls/sqlite/v1/db.ts

@@ -0,0 +1,26 @@
+import type { SpaceType } from '../../../utils/universal-id';
+
+interface NativeDBV1Apis {
+  getBlob: (
+    spaceType: SpaceType,
+    workspaceId: string,
+    key: string
+  ) => Promise<Buffer | null>;
+  deleteBlob: (
+    spaceType: SpaceType,
+    workspaceId: string,
+    key: string
+  ) => Promise<void>;
+  getBlobKeys: (spaceType: SpaceType, workspaceId: string) => Promise<string[]>;
+  getDocAsUpdates: (
+    spaceType: SpaceType,
+    workspaceId: string,
+    subdocId: string
+  ) => Promise<Uint8Array>;
+}
+
+export let apis: NativeDBV1Apis | null = null;
+
+export function bindNativeDBV1Apis(a: NativeDBV1Apis) {
+  apis = a;
+}

packages/common/nbstore/src/impls/sqlite/v1/doc.ts

@@ -1,24 +1,27 @@
-import { apis } from '@affine/electron-api';
-
 import { DummyConnection } from '../../../connection';
 import {
   type DocRecord,
   DocStorageBase,
   type DocUpdate,
 } from '../../../storage';
+import type { SpaceType } from '../../../utils/universal-id';
+import { apis } from './db';
 
 /**
  * @deprecated readonly
  */
-export class SqliteV1DocStorage extends DocStorageBase {
+export class SqliteV1DocStorage extends DocStorageBase<{
+  type: SpaceType;
+  id: string;
+}> {
   override connection = new DummyConnection();
 
-  get db() {
+  private get db() {
     if (!apis) {
       throw new Error('Not in electron context.');
     }
 
-    return apis.db;
+    return apis;
   }
 
   override async pushDocUpdate(update: DocUpdate) {
@@ -29,8 +32,8 @@ export class SqliteV1DocStorage extends DocStorageBase {
 
   override async getDoc(docId: string) {
     const bin = await this.db.getDocAsUpdates(
-      this.spaceType,
-      this.spaceId,
+      this.options.type,
+      this.options.id,
       docId
     );
 
@@ -41,8 +44,8 @@ export class SqliteV1DocStorage extends DocStorageBase {
     };
   }
 
-  override async deleteDoc(docId: string) {
-    await this.db.deleteDoc(this.spaceType, this.spaceId, docId);
+  override async deleteDoc() {
+    return;
   }
 
   protected override async getDocSnapshot() {

packages/common/nbstore/src/impls/sqlite/v1/index.ts

@@ -1,2 +1,3 @@
 export * from './blob';
+export { bindNativeDBV1Apis } from './db';
 export * from './doc';