feat(nbstore): add sqlite implementation

.github/workflows/build-test.yml

@@ -379,8 +379,8 @@ jobs:
           name: affine
           fail_ci_if_error: false
 
-  server-native-test:
-    name: Run server native tests
+  rust-test:
+    name: Run native tests
     runs-on: ubuntu-latest
     env:
       RUSTFLAGS: -D warnings
@@ -772,7 +772,7 @@ jobs:
       - build-server-native
       - build-electron-renderer
       - server-test
-      - server-native-test
+      - rust-test
       - copilot-api-test
       - copilot-e2e-test
       - server-e2e-test

Cargo.toml

@@ -40,3 +40,6 @@ codegen-units = 1
 lto           = true
 opt-level     = 3
 strip         = "symbols"
+
+[target.'cfg(target_os = "linux")']
+rustflags = ["-C", "link-args=-Wl,--warn-unresolved-symbols"]

packages/common/nbstore/package.json

@@ -9,7 +9,9 @@
     "./op": "./src/op/index.ts",
     "./idb": "./src/impls/idb/index.ts",
     "./idb/v1": "./src/impls/idb/v1/index.ts",
-    "./cloud": "./src/impls/cloud/index.ts"
+    "./cloud": "./src/impls/cloud/index.ts",
+    "./sqlite": "./src/impls/sqlite/index.ts",
+    "./sqlite/v1": "./src/impls/sqlite/v1/index.ts"
   },
   "dependencies": {
     "@toeverything/infra": "workspace:*",
@@ -19,11 +21,14 @@
     "yjs": "patch:yjs@npm%3A13.6.18#~/.yarn/patches/yjs-npm-13.6.18-ad0d5f7c43.patch"
   },
   "devDependencies": {
+    "@affine/electron-api": "workspace:*",
     "@affine/graphql": "workspace:*",
     "idb": "^8.0.0",
-    "socket.io-client": "^4.7.5"
+    "socket.io-client": "^4.7.5",
+    "vitest": "2.1.4"
   },
   "peerDependencies": {
+    "@affine/electron-api": "workspace:*",
     "@affine/graphql": "workspace:*",
     "idb": "^8.0.0",
     "socket.io-client": "^4.7.5"

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

@@ -0,0 +1,33 @@
+import { share } from '../../connection';
+import { type BlobRecord, BlobStorage } from '../../storage';
+import { NativeDBConnection } from './db';
+
+export class SqliteBlobStorage extends BlobStorage {
+  override connection = share(
+    new NativeDBConnection(this.peer, this.spaceType, this.spaceId)
+  );
+
+  get db() {
+    return this.connection.apis;
+  }
+
+  override async get(key: string) {
+    return this.db.getBlob(key);
+  }
+
+  override async set(blob: BlobRecord) {
+    await this.db.setBlob(blob);
+  }
+
+  override async delete(key: string, permanently: boolean) {
+    await this.db.deleteBlob(key, permanently);
+  }
+
+  override async release() {
+    await this.db.releaseBlobs();
+  }
+
+  override async list() {
+    return this.db.listBlobs();
+  }
+}

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

@@ -0,0 +1,83 @@
+import { apis, events } from '@affine/electron-api';
+
+import { Connection, type ConnectionStatus } from '../../connection';
+import { type SpaceType, universalId } from '../../storage';
+
+type NativeDBApis = NonNullable<typeof apis>['nbstore'] extends infer APIs
+  ? {
+      [K in keyof APIs]: APIs[K] extends (...args: any[]) => any
+        ? Parameters<APIs[K]> extends [string, ...infer Rest]
+          ? (...args: Rest) => ReturnType<APIs[K]>
+          : never
+        : never;
+    }
+  : never;
+
+export class NativeDBConnection extends Connection<void> {
+  readonly apis: NativeDBApis;
+
+  constructor(
+    private readonly peer: string,
+    private readonly type: SpaceType,
+    private readonly id: string
+  ) {
+    super();
+    if (!apis) {
+      throw new Error('Not in electron context.');
+    }
+
+    this.apis = this.bindApis(apis.nbstore);
+    this.listenToConnectionEvents();
+  }
+
+  override get shareId(): string {
+    return `sqlite:${this.peer}:${this.type}:${this.id}`;
+  }
+
+  bindApis(originalApis: NonNullable<typeof apis>['nbstore']): NativeDBApis {
+    const id = universalId({
+      peer: this.peer,
+      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 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;
+  }
+
+  override async doConnect() {
+    await this.apis.connect();
+  }
+
+  override async doDisconnect() {
+    await this.apis.close();
+  }
+
+  private listenToConnectionEvents() {
+    events?.nbstore.onConnectionStatusChanged(
+      ({ peer, spaceType, spaceId, status, error }) => {
+        if (
+          peer === this.peer &&
+          spaceType === this.type &&
+          spaceId === this.id
+        ) {
+          this.setStatus(status as ConnectionStatus, error);
+        }
+      }
+    );
+  }
+}

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

@@ -0,0 +1,50 @@
+import { share } from '../../connection';
+import { DocStorage, type DocUpdate } from '../../storage';
+import { NativeDBConnection } from './db';
+
+export class SqliteDocStorage extends DocStorage {
+  override connection = share(
+    new NativeDBConnection(this.peer, this.spaceType, this.spaceId)
+  );
+
+  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);
+  }
+
+  override async deleteDoc(docId: string) {
+    return this.db.deleteDoc(docId);
+  }
+
+  override async getDocTimestamps(after?: Date) {
+    return this.db.getDocTimestamps(after ? new Date(after) : undefined);
+  }
+
+  protected override async getDocSnapshot() {
+    // handled in db
+    // see electron/src/helper/nbstore/doc.ts
+    return null;
+  }
+
+  protected override async setDocSnapshot(): Promise<boolean> {
+    // handled in db
+    return true;
+  }
+
+  protected override async getDocUpdates() {
+    // handled in db
+    return [];
+  }
+
+  protected override markUpdatesMerged() {
+    // handled in db
+    return Promise.resolve(0);
+  }
+}

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

@@ -0,0 +1,3 @@
+export * from './blob';
+export * from './doc';
+export * from './sync';

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

@@ -0,0 +1,33 @@
+import { share } from '../../connection';
+import { type DocClock, SyncStorage } from '../../storage';
+import { NativeDBConnection } from './db';
+
+export class SqliteSyncStorage extends SyncStorage {
+  override connection = share(
+    new NativeDBConnection(this.peer, this.spaceType, this.spaceId)
+  );
+
+  get db() {
+    return this.connection.apis;
+  }
+
+  override async getPeerClocks(peer: string) {
+    return this.db.getPeerClocks(peer);
+  }
+
+  override async setPeerClock(peer: string, clock: DocClock) {
+    await this.db.setPeerClock(peer, clock);
+  }
+
+  override async getPeerPushedClocks(peer: string) {
+    return this.db.getPeerPushedClocks(peer);
+  }
+
+  override async setPeerPushedClock(peer: string, clock: DocClock) {
+    await this.db.setPeerPushedClock(peer, clock);
+  }
+
+  override async clearClocks() {
+    await this.db.clearClocks();
+  }
+}

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

@@ -0,0 +1,62 @@
+import { apis } from '@affine/electron-api';
+
+import { DummyConnection, share } from '../../../connection';
+import { BlobStorage } from '../../../storage';
+
+/**
+ * @deprecated readonly
+ */
+export class SqliteV1BlobStorage extends BlobStorage {
+  override connection = share(new DummyConnection());
+
+  get db() {
+    if (!apis) {
+      throw new Error('Not in electron context.');
+    }
+
+    return apis.db;
+  }
+
+  override async get(key: string) {
+    const data: Uint8Array | null = await this.db.getBlob(
+      this.spaceType,
+      this.spaceId,
+      key
+    );
+
+    if (!data) {
+      return null;
+    }
+
+    return {
+      key,
+      data,
+      mime: '',
+      createdAt: new Date(),
+    };
+  }
+
+  override async delete(key: string, permanently: boolean) {
+    if (permanently) {
+      await this.db.deleteBlob(this.spaceType, this.spaceId, key);
+    }
+  }
+
+  override async list() {
+    const keys = await this.db.getBlobKeys(this.spaceType, this.spaceId);
+
+    return keys.map(key => ({
+      key,
+      mime: '',
+      size: 0,
+      createdAt: new Date(),
+    }));
+  }
+
+  override async set() {
+    // no more writes
+  }
+  override async release() {
+    // no more writes
+  }
+}

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

@@ -0,0 +1,63 @@
+import { apis } from '@affine/electron-api';
+
+import { DummyConnection, share } from '../../../connection';
+import { type DocRecord, DocStorage, type DocUpdate } from '../../../storage';
+
+/**
+ * @deprecated readonly
+ */
+export class SqliteV1DocStorage extends DocStorage {
+  override connection = share(new DummyConnection());
+
+  get db() {
+    if (!apis) {
+      throw new Error('Not in electron context.');
+    }
+
+    return apis.db;
+  }
+
+  override async pushDocUpdate(update: DocUpdate) {
+    // no more writes
+
+    return { docId: update.docId, timestamp: new Date() };
+  }
+
+  override async getDoc(docId: string) {
+    const bin = await this.db.getDocAsUpdates(
+      this.spaceType,
+      this.spaceId,
+      docId
+    );
+
+    return {
+      docId,
+      bin,
+      timestamp: new Date(),
+    };
+  }
+
+  override async deleteDoc(docId: string) {
+    await this.db.deleteDoc(this.spaceType, this.spaceId, docId);
+  }
+
+  protected override async getDocSnapshot() {
+    return null;
+  }
+
+  override async getDocTimestamps() {
+    return {};
+  }
+
+  protected override async setDocSnapshot(): Promise<boolean> {
+    return false;
+  }
+
+  protected override async getDocUpdates(): Promise<DocRecord[]> {
+    return [];
+  }
+
+  protected override async markUpdatesMerged(): Promise<number> {
+    return 0;
+  }
+}