feat(nbstore): add sqlite implementation (#8811)

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,54 @@
+import { share } from '../../connection';
+import { type DocClock, 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);
+  }
+
+  override getDocTimestamp(docId: string): Promise<DocClock | null> {
+    return this.db.getDocTimestamp(docId);
+  }
+
+  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,53 @@
+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 getPeerRemoteClocks(peer: string) {
+    return this.db.getPeerRemoteClocks(peer);
+  }
+
+  override async getPeerRemoteClock(peer: string, docId: string) {
+    return this.db.getPeerRemoteClock(peer, docId);
+  }
+
+  override async setPeerRemoteClock(peer: string, clock: DocClock) {
+    await this.db.setPeerRemoteClock(peer, clock);
+  }
+
+  override async getPeerPulledRemoteClocks(peer: string) {
+    return this.db.getPeerPulledRemoteClocks(peer);
+  }
+
+  override async getPeerPulledRemoteClock(peer: string, docId: string) {
+    return this.db.getPeerPulledRemoteClock(peer, docId);
+  }
+
+  override async setPeerPulledRemoteClock(peer: string, clock: DocClock) {
+    await this.db.setPeerPulledRemoteClock(peer, clock);
+  }
+
+  override async getPeerPushedClocks(peer: string) {
+    return this.db.getPeerPushedClocks(peer);
+  }
+
+  override async getPeerPushedClock(peer: string, docId: string) {
+    return this.db.getPeerPushedClock(peer, docId);
+  }
+
+  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,67 @@
+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 {};
+  }
+
+  override async getDocTimestamp() {
+    return null;
+  }
+
+  protected override async setDocSnapshot(): Promise<boolean> {
+    return false;
+  }
+
+  protected override async getDocUpdates(): Promise<DocRecord[]> {
+    return [];
+  }
+
+  protected override async markUpdatesMerged(): Promise<number> {
+    return 0;
+  }
+}

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

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

packages/common/nbstore/src/op/index.ts

@@ -3,6 +3,8 @@ import { OpClient } from '@toeverything/infra/op';
 import type { Storage } from '../storage';
 import type { SpaceStorageOps } from './ops';
 
+export { SpaceStorageConsumer } from './consumer';
+
 export class SpaceStorageClient extends OpClient<SpaceStorageOps> {
   /**
    * Adding a storage implementation to the backend.

packages/common/nbstore/src/storage/__tests__/__snapshots__/storage.spec.ts.snap

@@ -0,0 +1,41 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`parseUniversalId > should parse universal id > @peer(@name);@type(userspace);@id(@id); 1`] = `
+{
+  "id": "@id",
+  "peer": "@name",
+  "type": "userspace",
+}
+`;
+
+exports[`parseUniversalId > should parse universal id > @peer(@peer(name);@type(userspace);@id(@id); 1`] = `
+{
+  "id": "@id",
+  "peer": "@peer(name",
+  "type": "userspace",
+}
+`;
+
+exports[`parseUniversalId > should parse universal id > @peer(123);@type(userspace);@id(456); 1`] = `
+{
+  "id": "456",
+  "peer": "123",
+  "type": "userspace",
+}
+`;
+
+exports[`parseUniversalId > should parse universal id > @peer(123);@type(workspace);@id(456); 1`] = `
+{
+  "id": "456",
+  "peer": "123",
+  "type": "workspace",
+}
+`;
+
+exports[`parseUniversalId > should parse universal id > @peer(https://app.affine.pro);@type(userspace);@id(hello:world); 1`] = `
+{
+  "id": "hello:world",
+  "peer": "https://app.affine.pro",
+  "type": "userspace",
+}
+`;

packages/common/nbstore/src/storage/__tests__/storage.spec.ts

@@ -0,0 +1,36 @@
+import { describe, expect, it } from 'vitest';
+
+import { parseUniversalId, universalId } from '../storage';
+
+describe('parseUniversalId', () => {
+  it('should generate universal id', () => {
+    expect(universalId({ peer: '123', type: 'workspace', id: '456' })).toEqual(
+      '@peer(123);@type(workspace);@id(456);'
+    );
+  });
+
+  it('should parse universal id', () => {
+    const testcases = [
+      '@peer(123);@type(userspace);@id(456);',
+      '@peer(123);@type(workspace);@id(456);',
+      '@peer(https://app.affine.pro);@type(userspace);@id(hello:world);',
+      '@peer(@name);@type(userspace);@id(@id);',
+      '@peer(@peer(name);@type(userspace);@id(@id);',
+    ];
+
+    testcases.forEach(id => {
+      expect(parseUniversalId(id)).toMatchSnapshot(id);
+    });
+  });
+
+  it('should throw invalid universal id', () => {
+    const testcases = [
+      '@peer(123);@type(anyspace);@id(456);', // invalid space type
+      '@peer(@peer(name););@type(userspace);@id(@id);', // invalid peer
+    ];
+
+    testcases.forEach(id => {
+      expect(() => parseUniversalId(id)).toThrow();
+    });
+  });
+});

packages/common/nbstore/src/storage/index.ts

@@ -6,6 +6,8 @@ import type { BlobStorage } from './blob';
 import type { DocStorage } from './doc';
 import type { SyncStorage } from './sync';
 
+type Storages = DocStorage | BlobStorage | SyncStorage;
+
 export class SpaceStorage {
   protected readonly storages: Map<StorageType, Storage> = new Map();
   private readonly event = new EventEmitter2();
@@ -17,24 +19,20 @@ export class SpaceStorage {
     );
   }
 
-  tryGet(type: 'blob'): BlobStorage | undefined;
-  tryGet(type: 'sync'): SyncStorage | undefined;
-  tryGet(type: 'doc'): DocStorage | undefined;
-  tryGet(type: StorageType) {
-    return this.storages.get(type);
+  tryGet<T extends StorageType>(
+    type: T
+  ): Extract<Storages, { storageType: T }> | undefined {
+    return this.storages.get(type) as Extract<Storages, { storageType: T }>;
   }
 
-  get(type: 'blob'): BlobStorage;
-  get(type: 'sync'): SyncStorage;
-  get(type: 'doc'): DocStorage;
-  get(type: StorageType) {
-    const storage = this.storages.get(type);
+  get<T extends StorageType>(type: T): Extract<Storages, { storageType: T }> {
+    const storage = this.tryGet(type);
 
     if (!storage) {
       throw new Error(`Storage ${type} not registered.`);
     }
 
-    return storage;
+    return storage as Extract<Storages, { storageType: T }>;
   }
 
   async connect() {

packages/common/nbstore/src/storage/storage.ts

@@ -9,6 +9,77 @@ export interface StorageOptions {
   id: string;
 }
 
+export function universalId({ peer, type, id }: StorageOptions) {
+  return `@peer(${peer});@type(${type});@id(${id});`;
+}
+
+export function isValidUniversalId(opts: Record<string, string>): boolean {
+  const requiredKeys: Array<keyof StorageOptions> = [
+    'peer',
+    'type',
+    'id',
+  ] as const;
+
+  for (const key of requiredKeys) {
+    if (!opts[key]) {
+      return false;
+    }
+  }
+
+  return opts.type === 'userspace' || opts.type === 'workspace';
+}
+
+export function parseUniversalId(id: string) {
+  const result: Record<string, string> = {};
+  let key = '';
+  let value = '';
+  let isInValue = false;
+
+  let i = -1;
+
+  while (++i < id.length) {
+    const ch = id[i];
+    const nextCh = id[i + 1];
+
+    // when we are in value string, we only care about ch and next char to be [')', ';'] to end the id part
+    if (isInValue) {
+      if (ch === ')' && nextCh === ';') {
+        result[key] = value;
+        key = '';
+        value = '';
+        isInValue = false;
+        i++;
+        continue;
+      }
+
+      value += ch;
+      continue;
+    }
+
+    if (ch === '@') {
+      const keyEnd = id.indexOf('(', i);
+      // we find '@' but no '(' in lookahead or '(' is immediately after '@', invalid id
+      if (keyEnd === -1 || keyEnd === i + 1) {
+        break;
+      }
+
+      key = id.slice(i + 1, keyEnd);
+      i = keyEnd;
+      isInValue = true;
+    } else {
+      break;
+    }
+  }
+
+  if (!isValidUniversalId(result)) {
+    throw new Error(
+      `Invalid universal storage id: ${id}. It should be in format of @peer(\${peer});@type(\${type});@id(\${id});`
+    );
+  }
+
+  return result as any;
+}
+
 export abstract class Storage<Opts extends StorageOptions = StorageOptions> {
   abstract readonly storageType: StorageType;
   abstract readonly connection: Connection;
@@ -25,6 +96,10 @@ export abstract class Storage<Opts extends StorageOptions = StorageOptions> {
     return this.options.id;
   }
 
+  get universalId() {
+    return universalId(this.options);
+  }
+
   constructor(public readonly options: Opts) {}
 
   async connect() {