From 454a02bc53486263c4461c9f7526d11d6f87a3b2 Mon Sep 17 00:00:00 2001
From: asier-isayas <asier.isayas@gmail.com>
Date: Mon, 16 Mar 2026 07:21:41 -0700
Subject: [PATCH] Handle partition key path with whitespace (#2423)

* Handle partition key path with white space

* add whitespace test case for pkey

* remove trailing whitespace

---------

Co-authored-by: Asier Isayas <aisayas@microsoft.com>
---
 src/Utils/QueryUtils.test.ts | 19 +++++++++++++++++++
 src/Utils/QueryUtils.ts      |  6 +++++-
 test/sql/testCases.ts        | 23 +++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/Utils/QueryUtils.test.ts b/src/Utils/QueryUtils.test.ts
index e1dfd6bf7..b8d072e28 100644
--- a/src/Utils/QueryUtils.test.ts
+++ b/src/Utils/QueryUtils.test.ts
@@ -350,6 +350,25 @@ describe("Query Utils", () => {
       expect(partitionKeyValues.length).toBe(2);
       expect(partitionKeyValues).toEqual(["Japan", "hello"]);
     });
+
+    it("should extract partition key value when path has surrounding whitespace", () => {
+      const doc = {
+        id: "test-id",
+        Country: "United States",
+        Location: {
+          type: "Point",
+        },
+      };
+
+      const partitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.MultiHash,
+        paths: ["/ Country ", "/ Location / type "],
+      };
+
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(doc, partitionKeyDefinition);
+      expect(partitionKeyValues.length).toBe(2);
+      expect(partitionKeyValues).toEqual(["United States", "Point"]);
+    });
   });
 
   describe("stripDoubleQuotesFromSegment", () => {
diff --git a/src/Utils/QueryUtils.ts b/src/Utils/QueryUtils.ts
index 2146f678a..d5a93a660 100644
--- a/src/Utils/QueryUtils.ts
+++ b/src/Utils/QueryUtils.ts
@@ -157,7 +157,11 @@ export const extractPartitionKeyValues = (
   const partitionKeyValues: PartitionKey[] = [];
 
   partitionKeyDefinition.paths.forEach((partitionKeyPath: string) => {
-    const pathSegments: string[] = partitionKeyPath.substring(1).split("/").map(stripDoubleQuotesFromSegment);
+    const pathSegments: string[] = partitionKeyPath
+      .substring(1)
+      .split("/")
+      .map(stripDoubleQuotesFromSegment)
+      .map((s) => s.trim());
     const value = getValueForPath(documentContent, pathSegments);
 
     if (value !== undefined) {
diff --git a/test/sql/testCases.ts b/test/sql/testCases.ts
index 2a31cb235..872f427b4 100644
--- a/test/sql/testCases.ts
+++ b/test/sql/testCases.ts
@@ -272,4 +272,27 @@ export const documentTestCases: DocumentTestCase[] = [
       },
     ],
   },
+  {
+    name: "Single Partition Key With Whitespace",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "whitespacePartitionKey",
+    documents: [
+      {
+        documentId: "whitespacePartitionKey",
+        partitionKeys: [{ key: "/ partitionKey", value: "whitespaceValue" }],
+      },
+      {
+        documentId: "whitespacePartitionKey_empty_string",
+        partitionKeys: [{ key: "/ partitionKey", value: "" }],
+      },
+      {
+        documentId: "whitespacePartitionKey_null",
+        partitionKeys: [{ key: "/ partitionKey", value: null }],
+      },
+      {
+        documentId: "whitespacePartitionKey_missing",
+        partitionKeys: [],
+      },
+    ],
+  },
 ];