From 58b841518a86ac3afc300af0eb499f626e88b1a9 Mon Sep 17 00:00:00 2001
From: Andrey Sobolev <haiodo@users.noreply.github.com>
Date: Thu, 14 Dec 2023 15:53:27 +0700
Subject: [PATCH] UBERF-4649: Fix query projection/cache issue (#4200)

Signed-off-by: Andrey Sobolev <haiodo@gmail.com>
---
 packages/query/src/index.ts | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/packages/query/src/index.ts b/packages/query/src/index.ts
index f13fbeaa9b..ff556bb2ce 100644
--- a/packages/query/src/index.ts
+++ b/packages/query/src/index.ts
@@ -510,7 +510,7 @@ export class LiveQuery extends TxProcessor implements Client {
   }
 
   protected async txUpdateDoc (tx: TxUpdateDoc<Doc>): Promise<TxResult> {
-    const docCache = new Map<Ref<Doc>, Doc>()
+    const docCache = new Map<string, Doc>()
     for (const queries of this.queries) {
       const isTx = this.client.getHierarchy().isDerived(queries[0], core.class.Tx)
       for (const q of queries[1]) {
@@ -525,7 +525,7 @@ export class LiveQuery extends TxProcessor implements Client {
     return {}
   }
 
-  private async handleDocUpdate (q: Query, tx: TxUpdateDoc<Doc>, docCache?: Map<Ref<Doc>, Doc>): Promise<void> {
+  private async handleDocUpdate (q: Query, tx: TxUpdateDoc<Doc>, docCache?: Map<string, Doc>): Promise<void> {
     if (q.result instanceof Promise) {
       q.result = await q.result
     }
@@ -658,7 +658,7 @@ export class LiveQuery extends TxProcessor implements Client {
   }
 
   // Check if query is partially matched.
-  private async matchQuery (q: Query, tx: TxUpdateDoc<Doc>, docCache?: Map<Ref<Doc>, Doc>): Promise<boolean> {
+  private async matchQuery (q: Query, tx: TxUpdateDoc<Doc>, docCache?: Map<string, Doc>): Promise<boolean> {
     if (!this.client.getHierarchy().isDerived(tx.objectClass, q._class)) {
       return false
     }
@@ -687,10 +687,15 @@ export class LiveQuery extends TxProcessor implements Client {
     }
 
     if (matched) {
-      const realDoc = docCache?.get(doc._id) ?? (await this.client.findOne(q._class, { _id: doc._id }, q.options))
+      const docIdKey = doc._id + JSON.stringify(q.options?.lookup)
+
+      const lookup = q.options?.lookup
+      const realDoc =
+        docCache?.get(docIdKey) ??
+        (await this.client.findOne(q._class, { _id: doc._id }, lookup !== undefined ? { lookup } : undefined))
 
       if (realDoc != null && docCache != null) {
-        docCache?.set(realDoc._id, realDoc)
+        docCache?.set(docIdKey, realDoc)
       }
       if (realDoc === undefined) return false
       const res = matchQuery([realDoc], q.query, q._class, this.client.getHierarchy())