Pk missing fix (#2094)

* fix partition key missing not being able to load the document

* Implement E2E tests for documents with different partitionkeys

* Implement E2E tests for documents with different partitionkeys

* Implement E2E tests for documents with different partitionkeys

* Updated snapshot

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, {}]);
+    });
   });
 });

src/Utils/QueryUtils.ts

@@ -47,6 +47,7 @@ export function buildDocumentsQueryPartitionProjections(
   for (const index in partitionKey.paths) {
     // TODO: Handle "/" in partition key definitions
     const projectedProperties: string[] = partitionKey.paths[index].split("/").slice(1);
+    const isSystemPartitionKey: boolean = partitionKey.systemKey || false;
     let projectedProperty = "";
 
     projectedProperties.forEach((property: string) => {
@@ -61,8 +62,13 @@ export function buildDocumentsQueryPartitionProjections(
         projectedProperty += `[${projection}]`;
       }
     });
-
-    projections.push(`${collectionAlias}${projectedProperty}`);
+    const fullAccess = `${collectionAlias}${projectedProperty}`;
+    if (!isSystemPartitionKey) {
+      const wrappedProjection = `IIF(IS_DEFINED(${fullAccess}), ${fullAccess}, {})`;
+      projections.push(wrappedProjection);
+    } else {
+      projections.push(fullAccess);
+    }
   }
 
   return projections.join(",");
@@ -118,7 +124,7 @@ export const extractPartitionKeyValues = (
   documentContent: any,
   partitionKeyDefinition: PartitionKeyDefinition,
 ): PartitionKey[] => {
-  if (!partitionKeyDefinition.paths || partitionKeyDefinition.paths.length === 0) {
+  if (!partitionKeyDefinition.paths || partitionKeyDefinition.paths.length === 0 || partitionKeyDefinition.systemKey) {
     return undefined;
   }
 
@@ -130,6 +136,8 @@ export const extractPartitionKeyValues = (
 
     if (value !== undefined) {
       partitionKeyValues.push(value);
+    } else {
+      partitionKeyValues.push({});
     }
   });