From 257256f9155653bde88c2ef3ee9f832c5ee48d76 Mon Sep 17 00:00:00 2001
From: sunghyunkang1111 <114709653+sunghyunkang1111@users.noreply.github.com>
Date: Wed, 26 Mar 2025 11:26:47 -0500
Subject: [PATCH] fix partition key missing not being able to load the document
 (#2085)

---
 .../Tabs/DocumentsTabV2/DocumentsTabV2.tsx    |  5 ++
 src/Utils/QueryUtils.test.ts                  | 78 ++++++++++++++++---
 src/Utils/QueryUtils.ts                       |  7 +-
 3 files changed, 76 insertions(+), 14 deletions(-)

diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
index c8004fe80..efd0bb996 100644
--- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
@@ -1037,6 +1037,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     );
 
     const selectedDocumentId = documentIds[clickedRowIndex as number];
+    const originalPartitionKeyValue = selectedDocumentId.partitionKeyValue;
     selectedDocumentId.partitionKeyValue = partitionKeyValueArray;
 
     onExecutionErrorChange(false);
@@ -1072,6 +1073,10 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
           setColumnDefinitionsFromDocument(documentContent);
         },
         (error) => {
+          // in case of any kind of failures of accidently changing partition key, restore the original
+          // so that when user navigates away from current document and comes back,
+          // it doesnt fail to load due to using the invalid partition keys
+          selectedDocumentId.partitionKeyValue = originalPartitionKeyValue;
           onExecutionErrorChange(true);
           const errorMessage = getErrorMessage(error);
           useDialog.getState().showOkModalDialog("Update document failed", errorMessage);
diff --git a/src/Utils/QueryUtils.test.ts b/src/Utils/QueryUtils.test.ts
index ab2b2a8b1..7da3d381c 100644
--- a/src/Utils/QueryUtils.test.ts
+++ b/src/Utils/QueryUtils.test.ts
@@ -35,6 +35,13 @@ describe("Query Utils", () => {
       version: 2,
     };
   };
+  const generatePartitionKeysForPaths = (paths: string[]): DataModels.PartitionKey => {
+    return {
+      paths: paths,
+      kind: "Hash",
+      version: 2,
+    };
+  };
 
   describe("buildDocumentsQueryPartitionProjections()", () => {
     it("should return empty string if partition key is undefined", () => {
@@ -89,6 +96,18 @@ describe("Query Utils", () => {
 
       expect(query).toContain("c.id");
     });
+
+    it("should always include {} for any missing partition keys", () => {
+      const query = QueryUtils.buildDocumentsQuery(
+        "",
+        ["a", "b", "c"],
+        generatePartitionKeysForPaths(["/a", "/b", "/c"]),
+        [],
+      );
+      expect(query).toContain('IIF(IS_DEFINED(c["a"]), c["a"], {})');
+      expect(query).toContain('IIF(IS_DEFINED(c["b"]), c["b"], {})');
+      expect(query).toContain('IIF(IS_DEFINED(c["c"]), c["c"], {})');
+    });
   });
 
   describe("queryPagesUntilContentPresent()", () => {
@@ -201,18 +220,6 @@ describe("Query Utils", () => {
       expect(expectedPartitionKeyValues).toContain(documentContent["Category"]);
     });
 
-    it("should extract no partition key values in the case nested partition key", () => {
-      const singlePartitionKeyDefinition: PartitionKeyDefinition = {
-        kind: PartitionKeyKind.Hash,
-        paths: ["/Location.type"],
-      };
-      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(
-        documentContent,
-        singlePartitionKeyDefinition,
-      );
-      expect(partitionKeyValues.length).toBe(0);
-    });
-
     it("should extract all partition key values for hierarchical and nested partition keys", () => {
       const mixedPartitionKeyDefinition: PartitionKeyDefinition = {
         kind: PartitionKeyKind.MultiHash,
@@ -225,5 +232,52 @@ describe("Query Utils", () => {
       expect(partitionKeyValues.length).toBe(2);
       expect(partitionKeyValues).toEqual(["United States", "Point"]);
     });
+
+    it("if any partition key is null or empty string, the partitionKeyValues shall match", () => {
+      const newDocumentContent = {
+        ...documentContent,
+        ...{
+          Country: null,
+          Location: {
+            type: "",
+            coordinates: [-121.49, 46.206],
+          },
+        },
+      };
+
+      const mixedPartitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.MultiHash,
+        paths: ["/Country", "/Location/type"],
+      };
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(
+        newDocumentContent,
+        mixedPartitionKeyDefinition,
+      );
+      expect(partitionKeyValues.length).toBe(2);
+      expect(partitionKeyValues).toEqual([null, ""]);
+    });
+
+    it("if any partition key doesn't exist, it should still set partitionkey value as {}", () => {
+      const newDocumentContent = {
+        ...documentContent,
+        ...{
+          Country: null,
+          Location: {
+            coordinates: [-121.49, 46.206],
+          },
+        },
+      };
+
+      const mixedPartitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.MultiHash,
+        paths: ["/Country", "/Location/type"],
+      };
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(
+        newDocumentContent,
+        mixedPartitionKeyDefinition,
+      );
+      expect(partitionKeyValues.length).toBe(2);
+      expect(partitionKeyValues).toEqual([null, {}]);
+    });
   });
 });
diff --git a/src/Utils/QueryUtils.ts b/src/Utils/QueryUtils.ts
index f0b39e4e2..6aa20bbce 100644
--- a/src/Utils/QueryUtils.ts
+++ b/src/Utils/QueryUtils.ts
@@ -61,8 +61,9 @@ export function buildDocumentsQueryPartitionProjections(
         projectedProperty += `[${projection}]`;
       }
     });
-
-    projections.push(`${collectionAlias}${projectedProperty}`);
+    const fullAccess = `${collectionAlias}${projectedProperty}`;
+    const wrappedProjection = `IIF(IS_DEFINED(${fullAccess}), ${fullAccess}, {})`;
+    projections.push(wrappedProjection);
   }
 
   return projections.join(",");
@@ -130,6 +131,8 @@ export const extractPartitionKeyValues = (
 
     if (value !== undefined) {
       partitionKeyValues.push(value);
+    } else {
+      partitionKeyValues.push({});
     }
   });