Added quote escaping for partition key extraction (#2403)

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -789,7 +789,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
   );
   let partitionKeyProperties = useMemo(() => {
     return partitionKeyPropertyHeaders?.map((partitionKeyPropertyHeader) =>
-      partitionKeyPropertyHeader.replace(/[/]+/g, ".").substring(1).replace(/[']+/g, ""),
+      partitionKeyPropertyHeader.replace(/[/]+/g, ".").substring(1).replace(/[']+/g, "").replace(/["]+/g, ""),
     );
   }, [partitionKeyPropertyHeaders]);
 
@@ -1470,7 +1470,11 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
                 const partitionKey = _partitionKey || (_collection && _collection.partitionKey);
                 const partitionKeyPropertyHeaders = _collection?.partitionKeyPropertyHeaders || partitionKey?.paths;
                 const partitionKeyProperties = partitionKeyPropertyHeaders?.map((partitionKeyPropertyHeader) =>
-                  partitionKeyPropertyHeader.replace(/[/]+/g, ".").substring(1).replace(/[']+/g, ""),
+                  partitionKeyPropertyHeader
+                    .replace(/[/]+/g, ".")
+                    .substring(1)
+                    .replace(/[']+/g, "")
+                    .replace(/["]+/g, ""),
                 );
 
                 return newDocumentId(rawDocument, partitionKeyProperties, partitionKeyValue);

src/Utils/QueryUtils.test.ts

@@ -4,7 +4,12 @@ import * as sinon from "sinon";
 import * as DataModels from "../Contracts/DataModels";
 import * as ViewModels from "../Contracts/ViewModels";
 import * as QueryUtils from "./QueryUtils";
-import { defaultQueryFields, extractPartitionKeyValues, getValueForPath } from "./QueryUtils";
+import {
+  defaultQueryFields,
+  extractPartitionKeyValues,
+  getValueForPath,
+  stripDoubleQuotesFromSegment,
+} from "./QueryUtils";
 
 const documentContent = {
   "Volcano Name": "Adams",
@@ -279,5 +284,97 @@ describe("Query Utils", () => {
       expect(partitionKeyValues.length).toBe(2);
       expect(partitionKeyValues).toEqual([null, {}]);
     });
+
+    it("should extract partition key value when path has enclosing double quotes", () => {
+      const docWithSpecialKey = {
+        id: "test-id",
+        "partition-key": "some-value",
+      };
+
+      const partitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.Hash,
+        paths: ['/"partition-key"'],
+      };
+
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(docWithSpecialKey, partitionKeyDefinition);
+      expect(partitionKeyValues.length).toBe(1);
+      expect(partitionKeyValues[0]).toEqual("some-value");
+    });
+
+    it("should extract nested partition key value when path segments have enclosing double quotes", () => {
+      const docWithSpecialKey = {
+        id: "test-id",
+        "my-field": {
+          "sub-field": 42,
+        },
+      };
+
+      const partitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.Hash,
+        paths: ['/"my-field"/"sub-field"'],
+      };
+
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(docWithSpecialKey, partitionKeyDefinition);
+      expect(partitionKeyValues.length).toBe(1);
+      expect(partitionKeyValues[0]).toEqual(42);
+    });
+
+    it("should return {} for missing double-quoted partition key", () => {
+      const docWithSpecialKey = {
+        id: "test-id",
+      };
+
+      const partitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.Hash,
+        paths: ['/"partition-key"'],
+      };
+
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(docWithSpecialKey, partitionKeyDefinition);
+      expect(partitionKeyValues.length).toBe(1);
+      expect(partitionKeyValues[0]).toEqual({});
+    });
+
+    it("should handle multi-hash with mixed quoted and unquoted paths", () => {
+      const doc = {
+        id: "test-id",
+        Country: "Japan",
+        "partition-key": "hello",
+      };
+
+      const partitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.MultiHash,
+        paths: ["/Country", '/"partition-key"'],
+      };
+
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(doc, partitionKeyDefinition);
+      expect(partitionKeyValues.length).toBe(2);
+      expect(partitionKeyValues).toEqual(["Japan", "hello"]);
+    });
+  });
+
+  describe("stripDoubleQuotesFromSegment", () => {
+    it("should strip enclosing double quotes", () => {
+      expect(stripDoubleQuotesFromSegment('"partition-key"')).toBe("partition-key");
+    });
+
+    it("should not strip if only opening quote", () => {
+      expect(stripDoubleQuotesFromSegment('"partition-key')).toBe('"partition-key');
+    });
+
+    it("should not strip if only closing quote", () => {
+      expect(stripDoubleQuotesFromSegment('partition-key"')).toBe('partition-key"');
+    });
+
+    it("should return empty string when stripping quotes from empty quoted string", () => {
+      expect(stripDoubleQuotesFromSegment('""')).toBe("");
+    });
+
+    it("should not modify unquoted segments", () => {
+      expect(stripDoubleQuotesFromSegment("Country")).toBe("Country");
+    });
+
+    it("should not strip single quotes", () => {
+      expect(stripDoubleQuotesFromSegment("'partition-key'")).toBe("'partition-key'");
+    });
   });
 });

src/Utils/QueryUtils.ts

@@ -116,6 +116,17 @@ export const queryPagesUntilContentPresent = async (
   return await doRequest(firstItemIndex);
 };
 
+/**
+ * Strips enclosing double quotes from a partition key path segment.
+ * e.g., '"partition-key"' -> 'partition-key'
+ */
+export const stripDoubleQuotesFromSegment = (segment: string): string => {
+  if (segment.length >= 2 && segment.charAt(0) === '"' && segment.charAt(segment.length - 1) === '"') {
+    return segment.slice(1, -1);
+  }
+  return segment;
+};
+
 /* eslint-disable  @typescript-eslint/no-explicit-any */
 export const getValueForPath = (content: any, pathSegments: string[]): any => {
   if (pathSegments.length === 0) {
@@ -146,7 +157,7 @@ export const extractPartitionKeyValues = (
   const partitionKeyValues: PartitionKey[] = [];
 
   partitionKeyDefinition.paths.forEach((partitionKeyPath: string) => {
-    const pathSegments: string[] = partitionKeyPath.substring(1).split("/");
+    const pathSegments: string[] = partitionKeyPath.substring(1).split("/").map(stripDoubleQuotesFromSegment);
     const value = getValueForPath(documentContent, pathSegments);
 
     if (value !== undefined) {

test/sql/testCases.ts

@@ -249,4 +249,27 @@ export const documentTestCases: DocumentTestCase[] = [
       },
     ],
   },
+  {
+    name: "Single Double-Quoted Partition Key",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "doubleQuotedPartitionKey",
+    documents: [
+      {
+        documentId: "doubleQuotedPartitionKey",
+        partitionKeys: [{ key: "/partition-key", value: "doubleQuotedValue" }],
+      },
+      {
+        documentId: "doubleQuotedPartitionKey_empty_string",
+        partitionKeys: [{ key: "/partition-key", value: "" }],
+      },
+      {
+        documentId: "doubleQuotedPartitionKey_null",
+        partitionKeys: [{ key: "/partition-key", value: null }],
+      },
+      {
+        documentId: "doubleQuotedPartitionKey_missing",
+        partitionKeys: [],
+      },
+    ],
+  },
 ];

test/testData.ts

@@ -251,7 +251,14 @@ export const setPartitionKeys = (partitionKeys: PartitionKey[]) => {
   partitionKeys.forEach((partitionKey) => {
     const { key: keyPath, value: keyValue } = partitionKey;
     const cleanPath = keyPath.startsWith("/") ? keyPath.slice(1) : keyPath;
-    const keys = cleanPath.split("/");
+    const keys = cleanPath.split("/").map((segment) => {
+      // Strip enclosing double quotes from partition key path segments
+      // e.g., '"partition-key"' -> 'partition-key'
+      if (segment.length >= 2 && segment.charAt(0) === '"' && segment.charAt(segment.length - 1) === '"') {
+        return segment.slice(1, -1);
+      }
+      return segment;
+    });
     let current = result;
 
     keys.forEach((key, index) => {