Fix $lookup mixins (#1085)

Signed-off-by: Andrey Sobolev <haiodo@gmail.com>

packages/core/src/__tests__/memdb.test.ts

@@ -281,4 +281,35 @@ describe('memdb', () => {
     const reverse = await client.findAll(spaces[0]._class, { _id: spaces[0]._id }, { lookup: { _id: { comments: test.class.TestComment } } })
     expect((reverse[0].$lookup as any).comments).toHaveLength(2)
   })
+
+  it('mixin lookups', async () => {
+    const { model } = await createModel()
+
+    const client = new TxOperations(model, core.account.System)
+    const spaces = await client.findAll(core.class.Space, {})
+    expect(spaces).toHaveLength(2)
+
+    const task = await client.createDoc(test.class.Task, spaces[0]._id, {
+      name: 'TSK1',
+      number: 1,
+      state: 0
+    })
+
+    await client.createMixin(task, test.class.Task, spaces[0]._id, test.mixin.TaskMixinTodos, {
+      todos: 0
+    })
+
+    await client.addCollection(test.class.TestMixinTodo, spaces[0]._id, task, test.mixin.TaskMixinTodos, 'todos', {
+      text: 'qwe'
+    })
+    await client.addCollection(test.class.TestMixinTodo, spaces[0]._id, task, test.mixin.TaskMixinTodos, 'todos', {
+      text: 'qwe2'
+    })
+
+    const results = await client.findAll(test.class.TestMixinTodo, {}, { lookup: { attachedTo: test.mixin.TaskMixinTodos } })
+    expect(results.length).toEqual(2)
+    const attached = results[0].$lookup?.attachedTo
+    expect(attached).toBeDefined()
+    expect(Hierarchy.mixinOrClass(attached as Doc)).toEqual(test.mixin.TaskMixinTodos)
+  })
 })

packages/core/src/__tests__/minmodel.ts

@@ -51,6 +51,14 @@ export interface Task extends Doc, WithState {
   name: string
 }
 
+export interface TaskMixinTodos extends Task {
+  todos: number
+}
+
+export interface TaskMixinTodo extends AttachedDoc {
+  text: string
+}
+
 export interface TaskCheckItem extends AttachedDoc, WithState {
   name: string
   complete: boolean
@@ -58,12 +66,14 @@ export interface TaskCheckItem extends AttachedDoc, WithState {
 
 export const test = plugin('test' as Plugin, {
   mixin: {
-    TestMixin: '' as Ref<Mixin<TestMixin>>
+    TestMixin: '' as Ref<Mixin<TestMixin>>,
+    TaskMixinTodos: '' as Ref<Mixin<TaskMixinTodos>>
   },
   class: {
     Task: '' as Ref<Class<Task>>,
     TaskCheckItem: '' as Ref<Class<TaskCheckItem>>,
-    TestComment: '' as Ref<Class<AttachedComment>>
+    TestComment: '' as Ref<Class<AttachedComment>>,
+    TestMixinTodo: '' as Ref<Mixin<TaskMixinTodo>>
   },
   interface: {
     WithState: '' as Ref<Interface<WithState>>,
@@ -102,6 +112,9 @@ export function genMinModel (): TxCUD<Doc>[] {
   txes.push(createClass(test.class.Task, { label: 'Task' as IntlString, extends: core.class.Doc, implements: [test.interface.DummyWithState], kind: ClassifierKind.CLASS }))
   txes.push(createClass(test.class.TaskCheckItem, { label: 'Task' as IntlString, extends: core.class.AttachedDoc, implements: [test.interface.WithState], kind: ClassifierKind.CLASS }))
 
+  txes.push(createClass(test.mixin.TaskMixinTodos, { label: 'TaskMixinTodos' as IntlString, extends: test.class.Task, kind: ClassifierKind.MIXIN }))
+  txes.push(createClass(test.class.TestMixinTodo, { label: 'TestMixinTodo' as IntlString, extends: core.class.AttachedDoc, kind: ClassifierKind.CLASS }))
+
   txes.push(
     createDoc(core.class.Space, {
       name: 'Sp1',

packages/core/src/client.ts

@@ -13,16 +13,14 @@
 // limitations under the License.
 //
 
-import type { Doc, Ref, Class } from './classes'
-import type { Tx } from './tx'
-import type { Storage, DocumentQuery, FindOptions, FindResult, WithLookup, TxResult } from './storage'
-
-import { SortingOrder } from './storage'
+import type { Class, Doc, Ref } from './classes'
+import { DOMAIN_MODEL } from './classes'
+import core from './component'
 import { Hierarchy } from './hierarchy'
 import { ModelDb } from './memdb'
-import { DOMAIN_MODEL } from './classes'
-
-import core from './component'
+import type { DocumentQuery, FindOptions, FindResult, Storage, TxResult, WithLookup } from './storage'
+import { SortingOrder } from './storage'
+import type { Tx } from './tx'
 
 /**
  * @public
@@ -67,15 +65,16 @@ class ClientImpl implements Client {
     options?: FindOptions<T>
   ): Promise<FindResult<T>> {
     const domain = this.hierarchy.getDomain(_class)
-    const result = (domain === DOMAIN_MODEL)
+    let result = (domain === DOMAIN_MODEL)
       ? await this.model.findAll(_class, query, options)
       : await this.conn.findAll(_class, query, options)
 
     // In case of mixin we need to create mixin proxies.
-    const baseClass = this.hierarchy.getBaseClass(_class)
-    if (baseClass !== _class) {
-      return result.map(v => this.hierarchy.as(v, _class))
-    }
+
+    // Update mixins & lookups
+    result = result.map(v => {
+      return this.hierarchy.updateLookupMixin(_class, v, options)
+    })
     return result
   }
 

packages/core/src/hierarchy.ts

@@ -13,6 +13,7 @@
 // limitations under the License.
 //
 
+import { FindOptions, Lookup, ToClassRefT, WithLookup } from '.'
 import type { AnyAttribute, Class, Classifier, Doc, Domain, Interface, Mixin, Obj, Ref } from './classes'
 import { ClassifierKind } from './classes'
 import core from './component'
@@ -333,6 +334,44 @@ export class Hierarchy {
       }
     }
   }
+
+  updateLookupMixin<T extends Doc>(_class: Ref<Class<T>>, result: WithLookup<T>, options?: FindOptions<T>): WithLookup<T> {
+    const baseClass = this.getBaseClass(_class)
+    const vResult = baseClass !== _class ? this.as(result, _class) : result
+    const lookup = result.$lookup
+    if (lookup !== undefined) {
+      // We need to check if lookup type is mixin and cast to it if required.
+      const lu = options?.lookup as Lookup<Doc>
+      if (lu?._id !== undefined) {
+        for (const [k, v] of Object.entries(lu._id)) {
+          const _cl = getClass(v as ToClassRefT<T, keyof T>)
+          if (this.isMixin(_cl)) {
+            const mval = (lookup as any)[k]
+            if (mval !== undefined) {
+              if (Array.isArray(mval)) {
+                (lookup as any)[k] = mval.map(it => this.as(it, _cl))
+              } else {
+                (lookup as any)[k] = this.as(mval, _cl)
+              }
+            }
+          }
+        }
+      }
+      for (const [k, v] of Object.entries(lu)) {
+        if (k === '_id') {
+          continue
+        }
+        const _cl = getClass(v as ToClassRefT<T, keyof T>)
+        if (this.isMixin(_cl)) {
+          const mval = (lookup as any)[k]
+          if (mval !== undefined) {
+            (lookup as any)[k] = this.as(mval, _cl)
+          }
+        }
+      }
+    }
+    return vResult
+  }
 }
 
 function addNew<T> (val: Set<T>, value: T): boolean {
@@ -348,3 +387,10 @@ function addIf<T> (array: T[], value: T): void {
     array.push(value)
   }
 }
+
+function getClass<T extends Doc> (vvv: ToClassRefT<T, keyof T>): Ref<Class<T>> {
+  if (Array.isArray(vvv)) {
+    return vvv[0]
+  }
+  return vvv
+}

packages/core/src/memdb.ts

@@ -147,7 +147,8 @@ export abstract class MemDb extends TxProcessor {
     if (options?.sort !== undefined) resultSort(result, options?.sort)
 
     result = result.slice(0, options?.limit)
-    return clone(result) as T[]
+    const tresult = clone(result) as T[]
+    return tresult.map(it => this.hierarchy.updateLookupMixin(_class, it, options))
   }
 
   addDoc (doc: Doc): void {

packages/core/src/storage.ts

@@ -45,11 +45,15 @@ export type DocumentQuery<T extends Doc> = {
   [key: string]: any
 }
 
+/**
+ * @public
+ */
+export type ToClassRefT<T extends object, P extends keyof T> = T[P] extends Ref<infer X> | null ? Ref<Class<X>> | [Ref<Class<X>>, Lookup<X>] : never
 /**
  * @public
  */
 export type ToClassRef<T extends object> = {
-  [P in keyof T]?: T[P] extends Ref<infer X> | null ? Ref<Class<X>> | [Ref<Class<X>>, Lookup<X>] : never
+  [P in keyof T]?: ToClassRefT<T, P>
 }
 
 /**