From 773022cfaee18c6b9b2604840dae1f3039e63ab4 Mon Sep 17 00:00:00 2001
From: Denis Bykhov <80476319+BykhovDenis@users.noreply.github.com>
Date: Fri, 18 Mar 2022 10:49:56 +0600
Subject: [PATCH] LiveQuery lookup fix (#1161)

Signed-off-by: Denis Bykhov <80476319+BykhovDenis@users.noreply.github.com>
---
 packages/query/src/index.ts | 57 +++++++++++++++++++++++++++++++++----
 1 file changed, 52 insertions(+), 5 deletions(-)

diff --git a/packages/query/src/index.ts b/packages/query/src/index.ts
index 7cd699401f..da83b669c7 100644
--- a/packages/query/src/index.ts
+++ b/packages/query/src/index.ts
@@ -247,7 +247,7 @@ export class LiveQuery extends TxProcessor implements Client {
         const updatedDoc = q.result[pos]
         if (updatedDoc.modifiedOn > tx.modifiedOn) return
         if (updatedDoc.modifiedOn === tx.modifiedOn) {
-          const current = await this.findOne(q._class, { _id: updatedDoc._id })
+          const current = await this.findOne(q._class, { _id: updatedDoc._id }, q.options)
           if (current !== undefined) {
             q.result[pos] = current
           } else {
@@ -614,7 +614,21 @@ export class LiveQuery extends TxProcessor implements Client {
         if (q.options !== undefined) {
           const lookup = (q.options.lookup as any)?.[key]
           if (lookup !== undefined) {
-            ;(updatedDoc.$lookup as any)[key] = await this.client.findOne(lookup, { _id: ops[key] })
+            const lookupClass = getLookupClass(lookup)
+            const nestedLookup = getNestedLookup(lookup)
+            if (Array.isArray(ops[key])) {
+              ;(updatedDoc.$lookup as any)[key] = await this.client.findAll(
+                lookupClass,
+                { _id: { $in: ops[key] } },
+                { lookup: nestedLookup }
+              )
+            } else {
+              ;(updatedDoc.$lookup as any)[key] = await this.client.findOne(
+                lookupClass,
+                { _id: ops[key] },
+                { lookup: nestedLookup }
+              )
+            }
           }
         }
       } else {
@@ -624,7 +638,24 @@ export class LiveQuery extends TxProcessor implements Client {
             if (q.options !== undefined) {
               const lookup = (q.options.lookup as any)?.[pkey]
               if (lookup !== undefined) {
-                ;(updatedDoc.$lookup as any)[pkey].push(await this.client.findOne(lookup, { _id: (pops as any)[pkey] as Ref<Doc> }))
+                const lookupClass = getLookupClass(lookup)
+                const nestedLookup = getNestedLookup(lookup)
+                const pp = updatedDoc.$lookup as any
+                if (pp[pkey] === undefined) {
+                  pp[pkey] = []
+                }
+                if (Array.isArray((pops as any)[pkey])) {
+                  const pushData = await this.client.findAll(
+                    lookupClass,
+                    { _id: { $in: (pops as any)[pkey] } },
+                    { lookup: nestedLookup }
+                  )
+                  pp[pkey].push(...pushData)
+                } else {
+                  pp[pkey].push(
+                    await this.client.findOne(lookupClass, { _id: (pops as any)[pkey] }, { lookup: nestedLookup })
+                  )
+                }
               }
             }
           }
@@ -634,8 +665,16 @@ export class LiveQuery extends TxProcessor implements Client {
             if (q.options !== undefined) {
               const lookup = (q.options.lookup as any)?.[pkey]
               if (lookup !== undefined) {
-                const pid = (pops as any)[pkey] as Ref<Doc>
-                ;(updatedDoc.$lookup as any)[pkey] = ((updatedDoc.$lookup as any)[pkey]).filter((it: Doc) => it._id !== pid)
+                const pid = (pops as any)[pkey]
+                const pp = updatedDoc.$lookup as any
+                if (pp[pkey] === undefined) {
+                  pp[pkey] = []
+                }
+                if (Array.isArray(pid)) {
+                  pp[pkey] = pp[pkey].filter((it: Doc) => !pid.includes(it._id))
+                } else {
+                  pp[pkey] = pp[pkey].filter((it: Doc) => it._id !== pid)
+                }
               }
             }
           }
@@ -680,3 +719,11 @@ export class LiveQuery extends TxProcessor implements Client {
     }
   }
 }
+
+function getNestedLookup (lookup: Ref<Class<Doc>> | [Ref<Class<Doc>>, Lookup<Doc>]): Lookup<Doc> | undefined {
+  return Array.isArray(lookup) ? lookup[1] : undefined
+}
+
+function getLookupClass (lookup: Ref<Class<Doc>> | [Ref<Class<Doc>>, Lookup<Doc>]): Ref<Class<Doc>> {
+  return Array.isArray(lookup) ? lookup[0] : lookup
+}