feat(infra): improve orm (#7475)

packages/common/infra/src/orm/core/__tests__/entity.spec.ts

@@ -6,7 +6,6 @@ import {
   type DBSchemaBuilder,
   f,
   MemoryORMAdapter,
-  type ORMClient,
   Table,
 } from '../';
 
@@ -18,12 +17,14 @@ const TEST_SCHEMA = {
   },
 } satisfies DBSchemaBuilder;
 
+const ORMClient = createORMClient(TEST_SCHEMA);
+
 type Context = {
-  client: ORMClient<typeof TEST_SCHEMA>;
+  client: InstanceType<typeof ORMClient>;
 };
 
 beforeEach<Context>(async t => {
-  t.client = createORMClient(TEST_SCHEMA, MemoryORMAdapter);
+  t.client = new ORMClient(new MemoryORMAdapter());
 });
 
 const test = t as TestAPI<Context>;
@@ -94,7 +95,7 @@ describe('ORM entity CRUD', () => {
     });
 
     // old tag should not be updated
-    expect(tag.name).not.toBe(tag2.name);
+    expect(tag.name).not.toBe(tag2!.name);
   });
 
   test('should be able to delete entity', async t => {

packages/common/infra/src/orm/core/__tests__/hook.spec.ts

@@ -7,7 +7,6 @@ import {
   type Entity,
   f,
   MemoryORMAdapter,
-  type ORMClient,
 } from '../';
 
 const TEST_SCHEMA = {
@@ -23,15 +22,10 @@ const TEST_SCHEMA = {
   },
 } satisfies DBSchemaBuilder;
 
-type Context = {
+const ORMClient = createORMClient(TEST_SCHEMA);
-  client: ORMClient<typeof TEST_SCHEMA>;
-};
-
-beforeEach<Context>(async t => {
-  t.client = createORMClient(TEST_SCHEMA, MemoryORMAdapter);
 
 // define the hooks
-  t.client.defineHook('tags', 'migrate field `color` to field `colors`', {
+ORMClient.defineHook('tags', 'migrate field `color` to field `colors`', {
   deserialize(data) {
     if (!data.colors && data.color) {
       data.colors = [data.color];
@@ -40,6 +34,13 @@ beforeEach<Context>(async t => {
     return data;
   },
 });
+
+type Context = {
+  client: InstanceType<typeof ORMClient>;
+};
+
+beforeEach<Context>(async t => {
+  t.client = new ORMClient(new MemoryORMAdapter());
 });
 
 const test = t as TestAPI<Context>;
@@ -65,7 +66,7 @@ describe('ORM hook mixin', () => {
     });
 
     const tag2 = client.tags.get(tag.id);
-    expect(tag2.colors).toStrictEqual(['red']);
+    expect(tag2!.colors).toStrictEqual(['red']);
   });
 
   test('update entity', t => {
@@ -77,7 +78,7 @@ describe('ORM hook mixin', () => {
     });
 
     const tag2 = client.tags.update(tag.id, { color: 'blue' });
-    expect(tag2.colors).toStrictEqual(['blue']);
+    expect(tag2!.colors).toStrictEqual(['blue']);
   });
 
   test('subscribe entity', t => {

packages/common/infra/src/orm/core/__tests__/schema.spec.ts

@@ -1,21 +1,12 @@
 import { nanoid } from 'nanoid';
 import { describe, expect, test } from 'vitest';
 
-import {
+import { createORMClient, f, MemoryORMAdapter } from '../';
-  createORMClient,
-  type DBSchemaBuilder,
-  f,
-  MemoryORMAdapter,
-} from '../';
-
-function createClient<Schema extends DBSchemaBuilder>(schema: Schema) {
-  return createORMClient(schema, MemoryORMAdapter);
-}
 
 describe('Schema validations', () => {
   test('primary key must be set', () => {
     expect(() =>
-      createClient({
+      createORMClient({
         tags: {
           id: f.string(),
           name: f.string(),
@@ -28,7 +19,7 @@ describe('Schema validations', () => {
 
   test('primary key must be unique', () => {
     expect(() =>
-      createClient({
+      createORMClient({
         tags: {
           id: f.string().primaryKey(),
           name: f.string().primaryKey(),
@@ -41,7 +32,7 @@ describe('Schema validations', () => {
 
   test('primary key should not be optional without default value', () => {
     expect(() =>
-      createClient({
+      createORMClient({
         tags: {
           id: f.string().primaryKey().optional(),
           name: f.string(),
@@ -54,7 +45,7 @@ describe('Schema validations', () => {
 
   test('primary key can be optional with default value', async () => {
     expect(() =>
-      createClient({
+      createORMClient({
         tags: {
           id: f.string().primaryKey().optional().default(nanoid),
           name: f.string(),
@@ -65,14 +56,16 @@ describe('Schema validations', () => {
 });
 
 describe('Entity validations', () => {
-  function createTagsClient() {
+  const Client = createORMClient({
-    return createClient({
     tags: {
       id: f.string().primaryKey().default(nanoid),
       name: f.string(),
       color: f.string(),
     },
   });
+
+  function createTagsClient() {
+    return new Client(new MemoryORMAdapter());
   }
 
   test('should not update primary key', () => {
@@ -123,13 +116,15 @@ describe('Entity validations', () => {
 
   test('should be able to assign `null` to json field', () => {
     expect(() => {
-      const client = createClient({
+      const Client = createORMClient({
         tags: {
           id: f.string().primaryKey().default(nanoid),
           info: f.json(),
         },
       });
 
+      const client = new Client(new MemoryORMAdapter());
+
       const tag = client.tags.create({ info: null });
 
       expect(tag.info).toBe(null);

packages/common/infra/src/orm/core/__tests__/sync.spec.ts

@@ -13,13 +13,7 @@ import { Doc } from 'yjs';
 import { DocEngine } from '../../../sync';
 import { MiniSyncServer } from '../../../sync/doc/__tests__/utils';
 import { MemoryStorage } from '../../../sync/doc/storage';
-import {
+import { createORMClient, type DBSchemaBuilder, f, YjsDBAdapter } from '../';
-  createORMClient,
-  type DBSchemaBuilder,
-  f,
-  type ORMClient,
-  YjsDBAdapter,
-} from '../';
 
 const TEST_SCHEMA = {
   tags: {
@@ -30,14 +24,16 @@ const TEST_SCHEMA = {
   },
 } satisfies DBSchemaBuilder;
 
+const ORMClient = createORMClient(TEST_SCHEMA);
+
 type Context = {
   server: MiniSyncServer;
   user1: {
-    client: ORMClient<typeof TEST_SCHEMA>;
+    client: InstanceType<typeof ORMClient>;
     engine: DocEngine;
   };
   user2: {
-    client: ORMClient<typeof TEST_SCHEMA>;
+    client: InstanceType<typeof ORMClient>;
     engine: DocEngine;
   };
 };
@@ -48,17 +44,10 @@ function createEngine(server: MiniSyncServer) {
 
 async function createClient(server: MiniSyncServer, clientId: number) {
   const engine = createEngine(server);
-  const client = createORMClient(TEST_SCHEMA, YjsDBAdapter, {
+  const Client = createORMClient(TEST_SCHEMA);
-    getDoc(guid: string) {
-      const doc = new Doc({ guid });
-      doc.clientID = clientId;
-      engine.addDoc(doc);
-      return doc;
-    },
-  });
 
   // define the hooks
-  client.defineHook('tags', 'migrate field `color` to field `colors`', {
+  Client.defineHook('tags', 'migrate field `color` to field `colors`', {
     deserialize(data) {
       if (!data.colors && data.color) {
         data.colors = [data.color];
@@ -68,6 +57,17 @@ async function createClient(server: MiniSyncServer, clientId: number) {
     },
   });
 
+  const client = new Client(
+    new YjsDBAdapter(TEST_SCHEMA, {
+      getDoc(guid: string) {
+        const doc = new Doc({ guid });
+        doc.clientID = clientId;
+        engine.addDoc(doc);
+        return doc;
+      },
+    })
+  );
+
   return {
     engine,
     client,

packages/common/infra/src/orm/core/__tests__/yjs.spec.ts

@@ -8,17 +8,25 @@ import {
   type DocProvider,
   type Entity,
   f,
-  type ORMClient,
   Table,
   YjsDBAdapter,
 } from '../';
 
+function incremental() {
+  let i = 0;
+  return () => i++;
+}
+
 const TEST_SCHEMA = {
   tags: {
     id: f.string().primaryKey().default(nanoid),
     name: f.string(),
     color: f.string(),
   },
+  users: {
+    id: f.number().primaryKey().default(incremental()),
+    name: f.string(),
+  },
 } satisfies DBSchemaBuilder;
 
 const docProvider: DocProvider = {
@@ -27,12 +35,13 @@ const docProvider: DocProvider = {
   },
 };
 
+const Client = createORMClient(TEST_SCHEMA);
 type Context = {
-  client: ORMClient<typeof TEST_SCHEMA>;
+  client: InstanceType<typeof Client>;
 };
 
 beforeEach<Context>(async t => {
-  t.client = createORMClient(TEST_SCHEMA, YjsDBAdapter, docProvider);
+  t.client = new Client(new YjsDBAdapter(TEST_SCHEMA, docProvider));
 });
 
 const test = t as TestAPI<Context>;
@@ -55,6 +64,13 @@ describe('ORM entity CRUD', () => {
     expect(tag.id).toBeDefined();
     expect(tag.name).toBe('test');
     expect(tag.color).toBe('red');
+
+    const user = client.users.create({
+      name: 'user1',
+    });
+
+    expect(typeof user.id).toBe('number');
+    expect(user.name).toBe('user1');
   });
 
   test('should be able to read entity', t => {
@@ -67,6 +83,12 @@ describe('ORM entity CRUD', () => {
 
     const tag2 = client.tags.get(tag.id);
     expect(tag2).toEqual(tag);
+
+    const user = client.users.create({
+      name: 'user1',
+    });
+    const user2 = client.users.get(user.id);
+    expect(user2).toEqual(user);
   });
 
   test('should be able to update entity', t => {
@@ -89,7 +111,7 @@ describe('ORM entity CRUD', () => {
     });
 
     // old tag should not be updated
-    expect(tag.name).not.toBe(tag2.name);
+    expect(tag.name).not.toBe(tag2!.name);
   });
 
   test('should be able to delete entity', t => {
@@ -149,6 +171,7 @@ describe('ORM entity CRUD', () => {
     const { client } = t;
 
     let tag: Entity<(typeof TEST_SCHEMA)['tags']> | null = null;
+
     const subscription1 = client.tags.get$('test').subscribe(data => {
       tag = data;
     });
@@ -210,15 +233,73 @@ describe('ORM entity CRUD', () => {
     subscription.unsubscribe();
   });
 
-  test('can not use reserved keyword as field name', () => {
+  test('should be able to subscribe to filtered entity changes', t => {
-    const schema = {
+    const { client } = t;
-      tags: {
-        $$KEY: f.string().primaryKey().default(nanoid),
-      },
-    };
 
-    expect(() => createORMClient(schema, YjsDBAdapter, docProvider)).toThrow(
+    let entities: any[] = [];
-      "[Table(tags)]: Field '$$KEY' is reserved keyword and can't be used"
+    const subscription = client.tags.find$({ name: 'test' }).subscribe(data => {
+      entities = data;
+    });
+
+    const tag1 = client.tags.create({
+      id: '1',
+      name: 'test',
+      color: 'red',
+    });
+
+    expect(entities).toStrictEqual([tag1]);
+
+    const tag2 = client.tags.create({
+      id: '2',
+      name: 'test',
+      color: 'blue',
+    });
+
+    expect(entities).toStrictEqual([tag1, tag2]);
+
+    subscription.unsubscribe();
+  });
+
+  test('should be able to subscription to any entity changes', t => {
+    const { client } = t;
+
+    let entities: any[] = [];
+    const subscription = client.tags.find$({}).subscribe(data => {
+      entities = data;
+    });
+
+    const tag1 = client.tags.create({
+      id: '1',
+      name: 'tag1',
+      color: 'red',
+    });
+
+    expect(entities).toStrictEqual([tag1]);
+
+    const tag2 = client.tags.create({
+      id: '2',
+      name: 'tag2',
+      color: 'blue',
+    });
+
+    expect(entities).toStrictEqual([tag1, tag2]);
+
+    subscription.unsubscribe();
+  });
+
+  test('can not use reserved keyword as field name', () => {
+    expect(
+      () =>
+        new YjsDBAdapter(
+          {
+            tags: {
+              $$DELETED: f.string().primaryKey().default(nanoid),
+            },
+          },
+          docProvider
+        )
+    ).toThrow(
+      "[Table(tags)]: Field '$$DELETED' is reserved keyword and can't be used"
     );
   });
 });

packages/common/infra/src/orm/core/adapters/memory/table.ts

@@ -1,19 +1,36 @@
-import { merge } from 'lodash-es';
+import { merge, pick } from 'lodash-es';
 
 import { HookAdapter } from '../mixins';
-import type { Key, TableAdapter, TableOptions } from '../types';
+import type {
+  DeleteQuery,
+  FindQuery,
+  InsertQuery,
+  ObserveQuery,
+  Select,
+  TableAdapter,
+  TableAdapterOptions,
+  UpdateQuery,
+  WhereCondition,
+} from '../types';
 
 @HookAdapter()
 export class MemoryTableAdapter implements TableAdapter {
-  data = new Map<Key, any>();
+  private readonly data = new Map<string, any>();
-  subscriptions = new Map<Key, Array<(data: any) => void>>();
+  private keyField = 'key';
+  private readonly subscriptions = new Set<(key: string, data: any) => void>();
 
   constructor(private readonly tableName: string) {}
 
-  setup(_opts: TableOptions) {}
+  setup(opts: TableAdapterOptions) {
+    this.keyField = opts.keyField;
+  }
+
   dispose() {}
 
-  create(key: Key, data: any) {
+  insert(query: InsertQuery) {
+    const { data, select } = query;
+    const key = String(data[this.keyField]);
+
     if (this.data.has(key)) {
       throw new Error(
         `Record with key ${key} already exists in table ${this.tableName}`
@@ -22,79 +39,125 @@ export class MemoryTableAdapter implements TableAdapter {
 
     this.data.set(key, data);
     this.dispatch(key, data);
-    this.dispatch('$$KEYS', this.keys());
+    return this.value(data, select);
-    return data;
   }
 
-  get(key: Key) {
+  find(query: FindQuery) {
-    return this.data.get(key) || null;
+    const { where, select } = query;
+    const result = [];
+
+    for (const record of this.iterate(where)) {
+      result.push(this.value(record, select));
     }
 
-  subscribe(key: Key, callback: (data: any) => void): () => void {
+    return result;
-    const sKey = key.toString();
-    let subs = this.subscriptions.get(sKey.toString());
-
-    if (!subs) {
-      subs = [];
-      this.subscriptions.set(sKey, subs);
   }
 
-    subs.push(callback);
+  observe(query: ObserveQuery): () => void {
-    callback(this.data.get(key) || null);
+    const { where, select, callback } = query;
+
+    let listeningOnAll = false;
+    const obKeys = new Set<string>();
+    const results = [];
+
+    if (!where) {
+      listeningOnAll = true;
+    } else if ('byKey' in where) {
+      obKeys.add(where.byKey.toString());
+    }
+
+    for (const record of this.iterate(where)) {
+      const key = String(record[this.keyField]);
+      if (!listeningOnAll) {
+        obKeys.add(key);
+      }
+      results.push(this.value(record, select));
+    }
+
+    callback(results);
+
+    const ob = (key: string, data: any) => {
+      if (
+        listeningOnAll ||
+        obKeys.has(key) ||
+        (where && this.match(data, where))
+      ) {
+        callback(this.find({ where, select }));
+        return;
+      }
+    };
+
+    this.subscriptions.add(ob);
 
     return () => {
-      this.subscriptions.set(
+      this.subscriptions.delete(ob);
-        sKey,
-        subs.filter(s => s !== callback)
-      );
     };
   }
 
-  keys(): Key[] {
+  update(query: UpdateQuery) {
-    return Array.from(this.data.keys());
+    const { where, data, select } = query;
-  }
+    const result = [];
-
-  subscribeKeys(callback: (keys: Key[]) => void): () => void {
-    const sKey = `$$KEYS`;
-    let subs = this.subscriptions.get(sKey);
-
-    if (!subs) {
-      subs = [];
-      this.subscriptions.set(sKey, subs);
-    }
-    subs.push(callback);
-    callback(this.keys());
-
-    return () => {
-      this.subscriptions.set(
-        sKey,
-        subs.filter(s => s !== callback)
-      );
-    };
-  }
-
-  update(key: Key, data: any) {
-    let record = this.data.get(key);
-
-    if (!record) {
-      throw new Error(
-        `Record with key ${key} does not exist in table ${this.tableName}`
-      );
-    }
 
+    for (let record of this.iterate(where)) {
       record = merge({}, record, data);
+      const key = String(record[this.keyField]);
       this.data.set(key, record);
       this.dispatch(key, record);
+      result.push(this.value(this.value(record, select)));
+    }
+
+    return result;
+  }
+
+  delete(query: DeleteQuery) {
+    const { where } = query;
+
+    for (const record of this.iterate(where)) {
+      const key = String(record[this.keyField]);
+      this.data.delete(key);
+      this.dispatch(key, null);
+    }
+  }
+
+  toObject(record: any): Record<string, any> {
     return record;
   }
 
-  delete(key: Key) {
+  value(data: any, select: Select = '*') {
-    this.data.delete(key);
+    if (select === 'key') {
-    this.dispatch(key, null);
+      return data[this.keyField];
-    this.dispatch('$$KEYS', this.keys());
     }
 
-  dispatch(key: Key, data: any) {
+    if (select === '*') {
-    this.subscriptions.get(key)?.forEach(callback => callback(data));
+      return this.toObject(data);
+    }
+
+    return pick(this.toObject(data), select);
+  }
+
+  private *iterate(where: WhereCondition = []) {
+    if (Array.isArray(where)) {
+      for (const value of this.data.values()) {
+        if (this.match(value, where)) {
+          yield value;
+        }
+      }
+    } else {
+      const key = where.byKey;
+      const record = this.data.get(key.toString());
+      if (record) {
+        yield record;
+      }
+    }
+  }
+
+  private match(record: any, where: WhereCondition) {
+    return Array.isArray(where)
+      ? where.every(c => record[c.field] === c.value)
+      : where.byKey === record[this.keyField];
+  }
+
+  private dispatch(key: string, data: any) {
+    this.subscriptions.forEach(callback => callback(key, data));
   }
 }

packages/common/infra/src/orm/core/adapters/mixins/hook.ts

@@ -1,6 +1,5 @@
-import type { Key, TableAdapter, TableOptions } from '../types';
+import type { TableAdapter, TableAdapterOptions } from '../types';
-
+declare module '../../types' {
-declare module '../types' {
   interface TableOptions {
     hooks?: Hook<unknown>[];
   }
@@ -15,12 +14,17 @@ export interface TableAdapterWithHook<T = unknown> extends Hook<T> {}
 export function HookAdapter(): ClassDecorator {
   // @ts-expect-error allow
   return (Class: { new (...args: any[]): TableAdapter }) => {
-    return class TableAdapterImpl
+    return class TableAdapterExtensions
       extends Class
       implements TableAdapterWithHook
     {
       hooks: Hook<unknown>[] = [];
 
+      override setup(opts: TableAdapterOptions): void {
+        super.setup(opts);
+        this.hooks = opts.hooks ?? [];
+      }
+
       deserialize(data: unknown) {
         if (!this.hooks.length) {
           return data;
@@ -32,28 +36,8 @@ export function HookAdapter(): ClassDecorator {
         );
       }
 
-      override setup(opts: TableOptions) {
+      override toObject(data: any): Record<string, any> {
-        this.hooks = opts.hooks || [];
+        return this.deserialize(super.toObject(data));
-        super.setup(opts);
-      }
-
-      override create(key: Key, data: any) {
-        return this.deserialize(super.create(key, data));
-      }
-
-      override get(key: Key) {
-        return this.deserialize(super.get(key));
-      }
-
-      override update(key: Key, data: any) {
-        return this.deserialize(super.update(key, data));
-      }
-
-      override subscribe(
-        key: Key,
-        callback: (data: unknown) => void
-      ): () => void {
-        return super.subscribe(key, data => callback(this.deserialize(data)));
       }
     };
   };

packages/common/infra/src/orm/core/adapters/types.ts

@@ -1,23 +1,66 @@
-import type { TableSchemaBuilder } from '../schema';
+import type { Key, TableOptions } from '../types';
 
-export interface Key {
+export interface TableAdapterOptions extends TableOptions {
-  toString(): string;
+  keyField: string;
 }
 
-export interface TableOptions {
+type WhereEqCondition = {
-  schema: TableSchemaBuilder;
+  field: string;
-}
+  value: any;
+};
 
-export interface TableAdapter<K extends Key = any, T = unknown> {
+type WhereByKeyCondition = {
-  setup(opts: TableOptions): void;
+  byKey: Key;
+};
+
+// currently only support eq condition
+// TODO(@forehalo): on the way [gt, gte, lt, lte, in, notIn, like, notLike, isNull, isNotNull, And, Or]
+export type WhereCondition = WhereEqCondition[] | WhereByKeyCondition;
+export type Select = '*' | 'key' | string[];
+
+export type InsertQuery = {
+  data: any;
+  select?: Select;
+};
+
+export type DeleteQuery = {
+  where?: WhereCondition;
+};
+
+export type UpdateQuery = {
+  where?: WhereCondition;
+  data: any;
+  select?: Select;
+};
+
+export type FindQuery = {
+  where?: WhereCondition;
+  select?: Select;
+};
+
+export type ObserveQuery = {
+  where?: WhereCondition;
+  select?: Select;
+  callback: (data: any[]) => void;
+};
+
+export type Query =
+  | InsertQuery
+  | DeleteQuery
+  | UpdateQuery
+  | FindQuery
+  | ObserveQuery;
+
+export interface TableAdapter {
+  setup(opts: TableAdapterOptions): void;
   dispose(): void;
-  create(key: K, data: Partial<T>): T;
+
-  get(key: K): T;
+  toObject(record: any): Record<string, any>;
-  subscribe(key: K, callback: (data: T) => void): () => void;
+  insert(query: InsertQuery): any;
-  keys(): K[];
+  update(query: UpdateQuery): any[];
-  subscribeKeys(callback: (keys: K[]) => void): () => void;
+  delete(query: DeleteQuery): void;
-  update(key: K, data: Partial<T>): T;
+  find(query: FindQuery): any[];
-  delete(key: K): void;
+  observe(query: ObserveQuery): () => void;
 }
 
 export interface DBAdapter {

packages/common/infra/src/orm/core/adapters/yjs/table.ts

@@ -1,9 +1,24 @@
-import { omit } from 'lodash-es';
+import { pick } from 'lodash-es';
-import type { Doc, Map as YMap, Transaction, YMapEvent } from 'yjs';
+import {
+  type AbstractType,
+  type Doc,
+  Map as YMap,
+  type Transaction,
+} from 'yjs';
 
 import { validators } from '../../validators';
 import { HookAdapter } from '../mixins';
-import type { Key, TableAdapter, TableOptions } from '../types';
+import type {
+  DeleteQuery,
+  FindQuery,
+  InsertQuery,
+  ObserveQuery,
+  Select,
+  TableAdapter,
+  TableAdapterOptions,
+  UpdateQuery,
+  WhereCondition,
+} from '../types';
 
 /**
  * Yjs Adapter for AFFiNE ORM
@@ -22,33 +37,29 @@ import type { Key, TableAdapter, TableOptions } from '../types';
 @HookAdapter()
 export class YjsTableAdapter implements TableAdapter {
   private readonly deleteFlagKey = '$$DELETED';
-  private readonly keyFlagKey = '$$KEY';
+  private keyField: string = 'key';
-  private readonly hiddenFields = [this.deleteFlagKey, this.keyFlagKey];
+  private fields: string[] = [];
 
   private readonly origin = 'YjsTableAdapter';
 
-  keysCache: Set<Key> | null = null;
-  cacheStaled = true;
-
   constructor(
     private readonly tableName: string,
     private readonly doc: Doc
   ) {}
 
-  setup(_opts: TableOptions): void {
+  setup(opts: TableAdapterOptions): void {
-    this.doc.on('update', (_, origin) => {
+    this.keyField = opts.keyField;
-      if (origin !== this.origin) {
+    this.fields = Object.keys(opts.schema);
-        this.markCacheStaled();
-      }
-    });
   }
 
   dispose() {
     this.doc.destroy();
   }
 
-  create(key: Key, data: any) {
+  insert(query: InsertQuery) {
+    const { data, select } = query;
     validators.validateYjsEntityData(this.tableName, data);
+    const key = data[this.keyField];
     const record = this.doc.getMap(key.toString());
 
     this.doc.transact(() => {
@@ -56,139 +67,174 @@ export class YjsTableAdapter implements TableAdapter {
         record.set(key, data[key]);
       }
 
-      this.keyBy(record, key);
+      record.delete(this.deleteFlagKey);
-      record.set(this.deleteFlagKey, false);
     }, this.origin);
 
-    this.markCacheStaled();
+    return this.value(record, select);
-    return this.value(record);
   }
 
-  update(key: Key, data: any) {
+  update(query: UpdateQuery) {
+    const { data, select, where } = query;
     validators.validateYjsEntityData(this.tableName, data);
-    const record = this.record(key);
-
-    if (this.isDeleted(record)) {
-      return;
-    }
 
+    const results: any[] = [];
     this.doc.transact(() => {
+      for (const record of this.iterate(where)) {
+        results.push(this.value(record, select));
         for (const key in data) {
-        record.set(key, data[key]);
+          this.setField(record, key, data[key]);
+        }
       }
     }, this.origin);
 
-    return this.value(record);
+    return results;
   }
 
-  get(key: Key) {
+  find(query: FindQuery) {
-    const record = this.record(key);
+    const { where, select } = query;
-    return this.value(record);
+    const records: any[] = [];
+    for (const record of this.iterate(where)) {
+      records.push(this.value(record, select));
     }
 
-  subscribe(key: Key, callback: (data: any) => void) {
+    return records;
-    const record: YMap<any> = this.record(key);
-    // init callback
-    callback(this.value(record));
-
-    const ob = (event: YMapEvent<any>) => {
-      callback(this.value(event.target));
-    };
-    record.observe(ob);
-
-    return () => {
-      record.unobserve(ob);
-    };
   }
 
-  keys() {
+  observe(query: ObserveQuery) {
-    const keysCache = this.buildKeysCache();
+    const { where, select, callback } = query;
-    return Array.from(keysCache);
+
+    let listeningOnAll = false;
+    const obKeys = new Set<any>();
+    const results = [];
+
+    if (!where) {
+      listeningOnAll = true;
+    } else if ('byKey' in where) {
+      obKeys.add(where.byKey.toString());
     }
 
-  subscribeKeys(callback: (keys: Key[]) => void) {
+    for (const record of this.iterate(where)) {
-    const keysCache = this.buildKeysCache();
+      if (!listeningOnAll) {
-    // init callback
+        obKeys.add(this.keyof(record));
-    callback(Array.from(keysCache));
+      }
+      results.push(this.value(record, select));
+    }
+
+    callback(results);
 
     const ob = (tx: Transaction) => {
-      const keysCache = this.buildKeysCache();
+      for (const [ty] of tx.changed) {
-
+        const record = ty as unknown as AbstractType<any>;
-      for (const [type] of tx.changed) {
+        if (
-        const data = type as unknown as YMap<any>;
+          listeningOnAll ||
-        const key = this.keyof(data);
+          obKeys.has(this.keyof(record)) ||
-        if (this.isDeleted(data)) {
+          (where && this.match(record, where))
-          keysCache.delete(key);
+        ) {
-        } else {
+          callback(this.find({ where, select }));
-          keysCache.add(key);
+          return;
         }
       }
-
-      callback(Array.from(keysCache));
     };
 
     this.doc.on('afterTransaction', ob);
-
     return () => {
       this.doc.off('afterTransaction', ob);
     };
   }
 
-  delete(key: Key) {
+  delete(query: DeleteQuery) {
-    const record = this.record(key);
+    const { where } = query;
 
     this.doc.transact(() => {
-      for (const key of record.keys()) {
+      for (const record of this.iterate(where)) {
-        if (!this.hiddenFields.includes(key)) {
+        this.deleteTy(record);
-          record.delete(key);
       }
-      }
-      record.set(this.deleteFlagKey, true);
     }, this.origin);
-    this.markCacheStaled();
   }
 
-  private isDeleted(record: YMap<any>) {
+  toObject(ty: AbstractType<any>): Record<string, any> {
-    return record.get(this.deleteFlagKey) === true;
+    return YMap.prototype.toJSON.call(ty);
   }
 
-  private record(key: Key) {
+  private recordByKey(key: string): AbstractType<any> | null {
-    return this.doc.getMap(key.toString());
+    // detect if the record is there otherwise yjs will create an empty Map.
+    if (this.doc.share.has(key)) {
+      return this.doc.getMap(key);
     }
 
-  private value(record: YMap<any>) {
-    if (this.isDeleted(record) || !record.size) {
     return null;
   }
 
-    return omit(record.toJSON(), this.hiddenFields);
+  private *iterate(where: WhereCondition = []) {
+    // fast pass for key lookup without iterating the whole table
+    if ('byKey' in where) {
+      const record = this.recordByKey(where.byKey.toString());
+      if (record) {
+        yield record;
+      }
+    } else if (Array.isArray(where)) {
+      for (const map of this.doc.share.values()) {
+        if (this.match(map, where)) {
+          yield map;
+        }
+      }
+    }
   }
 
-  private buildKeysCache() {
+  private value(record: AbstractType<any>, select: Select = '*') {
-    if (!this.keysCache || this.cacheStaled) {
+    if (this.isDeleted(record) || this.isEmpty(record)) {
-      this.keysCache = new Set();
+      return null;
-
-      for (const key of this.doc.share.keys()) {
-        const record = this.doc.getMap(key);
-        if (!this.isDeleted(record)) {
-          this.keysCache.add(this.keyof(record));
-        }
-      }
-      this.cacheStaled = false;
     }
 
-    return this.keysCache;
+    let selectedFields: string[];
+    if (select === 'key') {
+      return this.keyof(record);
+    } else if (select === '*') {
+      selectedFields = this.fields;
+    } else {
+      selectedFields = select;
     }
 
-  private markCacheStaled() {
+    return pick(this.toObject(record), selectedFields);
-    this.cacheStaled = true;
   }
 
-  private keyof(record: YMap<any>) {
+  private match(record: AbstractType<any>, where: WhereCondition) {
-    return record.get(this.keyFlagKey);
+    return (
+      !this.isDeleted(record) &&
+      (Array.isArray(where)
+        ? where.every(c => this.field(record, c.field) === c.value)
+        : where.byKey === this.keyof(record))
+    );
   }
 
-  private keyBy(record: YMap<any>, key: Key) {
+  private isDeleted(record: AbstractType<any>) {
-    record.set(this.keyFlagKey, key);
+    return (
+      this.field(record, this.deleteFlagKey) === true || this.isEmpty(record)
+    );
+  }
+
+  private keyof(record: AbstractType<any>) {
+    return this.field(record, this.keyField);
+  }
+
+  private field(ty: AbstractType<any>, field: string) {
+    return YMap.prototype.get.call(ty, field);
+  }
+
+  private setField(ty: AbstractType<any>, field: string, value: any) {
+    YMap.prototype.set.call(ty, field, value);
+  }
+
+  private isEmpty(ty: AbstractType<any>) {
+    return ty._map.size === 0;
+  }
+
+  private deleteTy(ty: AbstractType<any>) {
+    this.fields.forEach(field => {
+      if (field !== this.keyField) {
+        YMap.prototype.delete.call(ty, field);
+      }
+    });
+    YMap.prototype.set.call(ty, this.deleteFlagKey, true);
   }
 }

packages/common/infra/src/orm/core/client.ts

@@ -1,10 +1,10 @@
 import { type DBAdapter, type Hook } from './adapters';
 import type { DBSchemaBuilder } from './schema';
-import { Table, type TableMap } from './table';
+import { type CreateEntityInput, Table, type TableMap } from './table';
 import { validators } from './validators';
 
-class RawORMClient {
+export class ORMClient {
-  hooksMap: Map<string, Hook<any>[]> = new Map();
+  static hooksMap: Map<string, Hook<any>[]> = new Map();
   private readonly tables = new Map<string, Table<any>>();
   constructor(
     protected readonly db: DBSchemaBuilder,
@@ -17,7 +17,7 @@ class RawORMClient {
           if (!table) {
             table = new Table(this.adapter, tableName, {
               schema: tableSchema,
-              hooks: this.hooksMap.get(tableName),
+              hooks: ORMClient.hooksMap.get(tableName),
             });
             this.tables.set(tableName, table);
           }
@@ -27,7 +27,7 @@ class RawORMClient {
     });
   }
 
-  defineHook(tableName: string, _desc: string, hook: Hook<any>) {
+  static defineHook(tableName: string, _desc: string, hook: Hook<any>) {
     let hooks = this.hooksMap.get(tableName);
     if (!hooks) {
       hooks = [];
@@ -38,28 +38,28 @@ class RawORMClient {
   }
 }
 
-export function createORMClient<
+export function createORMClient<Schema extends DBSchemaBuilder>(
-  const Schema extends DBSchemaBuilder,
+  db: Schema
-  AdapterConstructor extends new (...args: any[]) => DBAdapter,
+): ORMClientWithTablesClass<Schema> {
-  AdapterConstructorParams extends
-    any[] = ConstructorParameters<AdapterConstructor> extends [
-    DBSchemaBuilder,
-    ...infer Args,
-  ]
-    ? Args
-    : never,
->(
-  db: Schema,
-  adapter: AdapterConstructor,
-  ...args: AdapterConstructorParams
-): ORMClient<Schema> {
   Object.entries(db).forEach(([tableName, schema]) => {
     validators.validateTableSchema(tableName, schema);
   });
 
-  return new RawORMClient(db, new adapter(db, ...args)) as TableMap<Schema> &
+  class ORMClientWithTables extends ORMClient {
-    RawORMClient;
+    constructor(adapter: DBAdapter) {
+      super(db, adapter);
+    }
   }
 
-export type ORMClient<Schema extends DBSchemaBuilder> = RawORMClient &
+  return ORMClientWithTables as any;
-  TableMap<Schema>;
+}
+
+export type ORMClientWithTablesClass<Schema extends DBSchemaBuilder> = {
+  new (adapter: DBAdapter): TableMap<Schema> & ORMClient;
+
+  defineHook<TableName extends keyof Schema>(
+    tableName: TableName,
+    desc: string,
+    hook: Hook<CreateEntityInput<Schema[TableName]>>
+  ): void;
+};

packages/common/infra/src/orm/core/table.ts

@@ -1,13 +1,14 @@
 import { isUndefined, omitBy } from 'lodash-es';
 import { Observable, shareReplay } from 'rxjs';
 
-import type { DBAdapter, Key, TableAdapter, TableOptions } from './adapters';
+import type { DBAdapter, TableAdapter } from './adapters';
 import type {
   DBSchemaBuilder,
   FieldSchemaBuilder,
   TableSchema,
   TableSchemaBuilder,
 } from './schema';
+import type { Key, TableOptions } from './types';
 import { validators } from './validators';
 
 type Pretty<T> = T extends any
@@ -74,10 +75,16 @@ export type UpdateEntityInput<T extends TableSchemaBuilder> = Pretty<{
     : never;
 }>;
 
+export type FindEntityInput<T extends TableSchemaBuilder> = Pretty<{
+  [key in keyof T]?: T[key] extends FieldSchemaBuilder<infer Type>
+    ? Type
+    : never;
+}>;
+
 export class Table<T extends TableSchemaBuilder> {
   readonly schema: TableSchema;
   readonly keyField: string = '';
-  private readonly adapter: TableAdapter<PrimaryKeyFieldType<T>, Entity<T>>;
+  private readonly adapter: TableAdapter;
 
   private readonly subscribedKeys: Map<Key, Observable<any>> = new Map();
 
@@ -87,7 +94,6 @@ export class Table<T extends TableSchemaBuilder> {
     private readonly opts: TableOptions
   ) {
     this.adapter = db.table(name) as any;
-    this.adapter.setup(opts);
     this.schema = Object.entries(this.opts.schema).reduce(
       (acc, [fieldName, fieldBuilder]) => {
         acc[fieldName] = fieldBuilder.schema;
@@ -99,6 +105,7 @@ export class Table<T extends TableSchemaBuilder> {
       },
       {} as TableSchema
     );
+    this.adapter.setup({ ...opts, keyField: this.keyField });
   }
 
   create(input: CreateEntityInput<T>): Entity<T> {
@@ -123,16 +130,35 @@ export class Table<T extends TableSchemaBuilder> {
 
     validators.validateCreateEntityData(this, data);
 
-    return this.adapter.create(data[this.keyField], data);
+    return this.adapter.insert({
+      data: data,
+    });
   }
 
-  update(key: PrimaryKeyFieldType<T>, input: UpdateEntityInput<T>): Entity<T> {
+  update(
+    key: PrimaryKeyFieldType<T>,
+    input: UpdateEntityInput<T>
+  ): Entity<T> | null {
     validators.validateUpdateEntityData(this, input);
-    return this.adapter.update(key, omitBy(input, isUndefined) as any);
+
+    const [record] = this.adapter.update({
+      where: {
+        byKey: key,
+      },
+      data: input,
+    });
+
+    return record || null;
   }
 
-  get(key: PrimaryKeyFieldType<T>): Entity<T> {
+  get(key: PrimaryKeyFieldType<T>): Entity<T> | null {
-    return this.adapter.get(key);
+    const [record] = this.adapter.find({
+      where: {
+        byKey: key,
+      },
+    });
+
+    return record || null;
   }
 
   get$(key: PrimaryKeyFieldType<T>): Observable<Entity<T>> {
@@ -140,8 +166,13 @@ export class Table<T extends TableSchemaBuilder> {
 
     if (!ob$) {
       ob$ = new Observable<Entity<T>>(subscriber => {
-        const unsubscribe = this.adapter.subscribe(key, data => {
+        const unsubscribe = this.adapter.observe({
-          subscriber.next(data);
+          where: {
+            byKey: key,
+          },
+          callback: ([data]) => {
+            subscriber.next(data || null);
+          },
         });
 
         return () => {
@@ -161,8 +192,35 @@ export class Table<T extends TableSchemaBuilder> {
     return ob$;
   }
 
+  find(where: FindEntityInput<T>): Entity<T>[] {
+    return this.adapter.find({
+      where: Object.entries(where).map(([field, value]) => ({
+        field,
+        value,
+      })),
+    });
+  }
+
+  find$(where: FindEntityInput<T>): Observable<Entity<T>[]> {
+    return new Observable<Entity<T>[]>(subscriber => {
+      const unsubscribe = this.adapter.observe({
+        where: Object.entries(where).map(([field, value]) => ({
+          field,
+          value,
+        })),
+        callback: data => {
+          subscriber.next(data);
+        },
+      });
+
+      return unsubscribe;
+    });
+  }
+
   keys(): PrimaryKeyFieldType<T>[] {
-    return this.adapter.keys();
+    return this.adapter.find({
+      select: 'key',
+    });
   }
 
   keys$(): Observable<PrimaryKeyFieldType<T>[]> {
@@ -170,8 +228,11 @@ export class Table<T extends TableSchemaBuilder> {
 
     if (!ob$) {
       ob$ = new Observable<PrimaryKeyFieldType<T>[]>(subscriber => {
-        const unsubscribe = this.adapter.subscribeKeys(keys => {
+        const unsubscribe = this.adapter.observe({
+          select: 'key',
+          callback: (keys: PrimaryKeyFieldType<T>[]) => {
             subscriber.next(keys);
+          },
         });
 
         return () => {
@@ -192,7 +253,11 @@ export class Table<T extends TableSchemaBuilder> {
   }
 
   delete(key: PrimaryKeyFieldType<T>) {
-    return this.adapter.delete(key);
+    this.adapter.delete({
+      where: {
+        byKey: key,
+      },
+    });
   }
 }
 

packages/common/infra/src/orm/core/types.ts

@@ -0,0 +1,9 @@
+import type { TableSchemaBuilder } from './schema';
+
+export interface Key {
+  toString(): string;
+}
+
+export interface TableOptions {
+  schema: TableSchemaBuilder;
+}

packages/common/infra/src/orm/core/validators/yjs.ts

@@ -1,6 +1,6 @@
 import type { TableSchemaValidator } from './types';
 
-const PRESERVED_FIELDS = ['$$KEY', '$$DELETED'];
+const PRESERVED_FIELDS = ['$$DELETED'];
 
 interface DataValidator {
   validate(tableName: string, data: any): void;