Pk missing fix (#2107)

* 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

* Updated tests for MongoRU and add create/delete tests

* Fixing system partition key showing up in Data Explorer

playwright.config.ts

@@ -1,5 +1,4 @@
 import { defineConfig, devices } from "@playwright/test";
-
 /**
  * See https://playwright.dev/docs/test-configuration.
  */
@@ -29,7 +28,12 @@ export default defineConfig({
   projects: [
     {
       name: "chromium",
-      use: { ...devices["Desktop Chrome"] },
+      use: {
+        ...devices["Desktop Chrome"],
+        launchOptions: {
+          args: ["--disable-web-security", "--disable-features=IsolateOrigins,site-per-process"],
+        },
+      },
     },
     {
       name: "firefox",

src/Explorer/Controls/Dialog.tsx

@@ -214,8 +214,10 @@ export const Dialog: FC = () => {
       {contentHtml}
       {progressIndicatorProps && <ProgressIndicator {...progressIndicatorProps} />}
       <DialogFooter>
-        <PrimaryButton {...primaryButtonProps} />
+        <PrimaryButton {...primaryButtonProps} data-test={`DialogButton:${primaryButtonText}`} />
-        {secondaryButtonProps && <DefaultButton {...secondaryButtonProps} />}
+        {secondaryButtonProps && (
+          <DefaultButton {...secondaryButtonProps} data-test={`DialogButton:${secondaryButtonText}`} />
+        )}
       </DialogFooter>
     </FluentDialog>
   ) : (

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -773,16 +773,14 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     [_collection, _partitionKey],
   );
   const partitionKeyPropertyHeaders: string[] = useMemo(
-    () => _collection?.partitionKeyPropertyHeaders || partitionKey?.paths,
+    () => (partitionKey?.systemKey ? [] : _collection?.partitionKeyPropertyHeaders || partitionKey?.paths),
-    [_collection?.partitionKeyPropertyHeaders, partitionKey?.paths],
+    [_collection?.partitionKeyPropertyHeaders, partitionKey?.paths, partitionKey?.systemKey],
-  );
-  let partitionKeyProperties = useMemo(
-    () =>
-      partitionKeyPropertyHeaders?.map((partitionKeyPropertyHeader) =>
-        partitionKeyPropertyHeader.replace(/[/]+/g, ".").substring(1).replace(/[']+/g, ""),
-      ),
-    [partitionKeyPropertyHeaders],
   );
+  let partitionKeyProperties = useMemo(() => {
+    return partitionKeyPropertyHeaders?.map((partitionKeyPropertyHeader) =>
+      partitionKeyPropertyHeader.replace(/[/]+/g, ".").substring(1).replace(/[']+/g, ""),
+    );
+  }, [partitionKeyPropertyHeaders]);
 
   const getInitialColumnSelection = () => {
     const defaultColumnsIds = ["id"];
@@ -1046,6 +1044,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     );
 
     const selectedDocumentId = documentIds[clickedRowIndex as number];
+    const originalPartitionKeyValue = selectedDocumentId.partitionKeyValue;
     selectedDocumentId.partitionKeyValue = partitionKeyValueArray;
 
     onExecutionErrorChange(false);
@@ -1081,6 +1080,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);
@@ -1720,7 +1723,8 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     renderObjectForEditor = (value: unknown): string => MongoUtility.tojson(value, null, false);
 
     const _hasShardKeySpecified = (document: unknown): boolean => {
-      return Boolean(extractPartitionKeyValues(document, _getPartitionKeyDefinition() as PartitionKeyDefinition));
+      const partitionKeyDefinition: PartitionKeyDefinition = _getPartitionKeyDefinition() as PartitionKeyDefinition;
+      return partitionKeyDefinition.systemKey || Boolean(extractPartitionKeyValues(document, partitionKeyDefinition));
     };
 
     const _getPartitionKeyDefinition = (): DataModels.PartitionKey => {
@@ -1744,7 +1748,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       return partitionKey;
     };
 
-    partitionKeyProperties = partitionKeyProperties?.map((partitionKeyProperty, i) => {
+    partitionKeyProperties = partitionKeyProperties.map((partitionKeyProperty, i) => {
       if (partitionKeyProperty && ~partitionKeyProperty.indexOf(`"`)) {
         partitionKeyProperty = partitionKeyProperty.replace(/["]+/g, "");
       }
@@ -2094,8 +2098,8 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   return (
     <CosmosFluentProvider className={styles.container}>
-      <div className="tab-pane active" role="tabpanel" style={{ display: "flex" }}>
+      <div data-test={"DocumentsTab"} className="tab-pane active" role="tabpanel" style={{ display: "flex" }}>
-        <div className={styles.filterRow}>
+        <div data-test={"DocumentsTab/Filter"} className={styles.filterRow}>
           {!isPreferredApiMongoDB && <span> SELECT * FROM c </span>}
           <InputDataList
             dropdownOptions={getFilterChoices()}
@@ -2137,7 +2141,11 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
           }}
         >
           <Allotment.Pane preferredSize={`${tabStateData.leftPaneWidthPercent}%`} minSize={55}>
-            <div style={{ height: "100%", width: "100%", overflow: "hidden" }} ref={tableContainerRef}>
+            <div
+              data-test={"DocumentsTab/DocumentsPane"}
+              style={{ height: "100%", width: "100%", overflow: "hidden" }}
+              ref={tableContainerRef}
+            >
               <div className={styles.tableContainer}>
                 <div
                   style={
@@ -2191,7 +2199,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
             </div>
           </Allotment.Pane>
           <Allotment.Pane minSize={30}>
-            <div style={{ height: "100%", width: "100%" }}>
+            <div data-test={"DocumentsTab/ResultsPane"} style={{ height: "100%", width: "100%" }}>
               {isTabActive && selectedDocumentContent && selectedRows.size <= 1 && (
                 <EditorReact
                   language={"json"}

src/Explorer/Tabs/DocumentsTabV2/__snapshots__/DocumentsTabV2.test.tsx.snap

@@ -6,6 +6,7 @@ exports[`Documents tab (noSql API) when rendered should render the page 1`] = `
 >
   <div
     className="tab-pane active"
+    data-test="DocumentsTab"
     role="tabpanel"
     style={
       {
@@ -15,6 +16,7 @@ exports[`Documents tab (noSql API) when rendered should render the page 1`] = `
   >
     <div
       className="___11ktxfv_0000000 f1o614cb fy9rknc f22iagw fsnqrgy f1f5gg8d fjodcmx f122n59 f1f09k3d fg706s2 frpde29"
+      data-test="DocumentsTab/Filter"
     >
       <span>
          SELECT * FROM c 
@@ -65,6 +67,7 @@ exports[`Documents tab (noSql API) when rendered should render the page 1`] = `
         preferredSize="35%"
       >
         <div
+          data-test="DocumentsTab/DocumentsPane"
           style={
             {
               "height": "100%",
@@ -126,6 +129,7 @@ exports[`Documents tab (noSql API) when rendered should render the page 1`] = `
         minSize={30}
       >
         <div
+          data-test="DocumentsTab/ResultsPane"
           style={
             {
               "height": "100%",

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}]`;
       }
     });
-
+    const fullAccess = `${collectionAlias}${projectedProperty}`;
-    projections.push(`${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({});
     }
   });
 

test/fx.ts

@@ -1,5 +1,5 @@
 import { AzureCliCredential } from "@azure/identity";
-import { expect, Frame, Locator, Page } from "@playwright/test";
+import { Frame, Locator, Page, expect } from "@playwright/test";
 import crypto from "crypto";
 
 const RETRY_COUNT = 3;
@@ -26,7 +26,7 @@ export function getAzureCLICredentials(): AzureCliCredential {
 
 export async function getAzureCLICredentialsToken(): Promise<string> {
   const credentials = getAzureCLICredentials();
-  const token = (await credentials.getToken("https://management.core.windows.net//.default")).token;
+  const token = (await credentials.getToken("https://management.core.windows.net//.default"))?.token || "";
   return token;
 }
 
@@ -35,8 +35,10 @@ export enum TestAccount {
   Cassandra = "Cassandra",
   Gremlin = "Gremlin",
   Mongo = "Mongo",
+  MongoReadonly = "MongoReadOnly",
   Mongo32 = "Mongo32",
   SQL = "SQL",
+  SQLReadOnly = "SQLReadOnly",
 }
 
 export const defaultAccounts: Record<TestAccount, string> = {
@@ -44,8 +46,10 @@ export const defaultAccounts: Record<TestAccount, string> = {
   [TestAccount.Cassandra]: "github-e2etests-cassandra",
   [TestAccount.Gremlin]: "github-e2etests-gremlin",
   [TestAccount.Mongo]: "github-e2etests-mongo",
+  [TestAccount.MongoReadonly]: "github-e2etests-mongo-readonly",
   [TestAccount.Mongo32]: "github-e2etests-mongo32",
   [TestAccount.SQL]: "github-e2etests-sql",
+  [TestAccount.SQLReadOnly]: "github-e2etests-sql-readonly",
 };
 
 export const resourceGroupName = process.env.DE_TEST_RESOURCE_GROUP ?? "de-e2e-tests";
@@ -214,6 +218,25 @@ export class QueryTab {
   }
 }
 
+export class DocumentsTab {
+  documentsFilter: Locator;
+  documentsListPane: Locator;
+  documentResultsPane: Locator;
+  resultsEditor: Editor;
+
+  constructor(
+    public frame: Frame,
+    public tabId: string,
+    public tab: Locator,
+    public locator: Locator,
+  ) {
+    this.documentsFilter = this.locator.getByTestId("DocumentsTab/Filter");
+    this.documentsListPane = this.locator.getByTestId("DocumentsTab/DocumentsPane");
+    this.documentResultsPane = this.locator.getByTestId("DocumentsTab/ResultsPane");
+    this.resultsEditor = new Editor(this.frame, this.documentResultsPane.getByTestId("EditorReact/Host/Loaded"));
+  }
+}
+
 type PanelOpenOptions = {
   closeTimeout?: number;
 };
@@ -232,6 +255,12 @@ export class DataExplorer {
     return new QueryTab(this.frame, tabId, tab, queryTab);
   }
 
+  documentsTab(tabId: string): DocumentsTab {
+    const tab = this.tab(tabId);
+    const documentsTab = tab.getByTestId("DocumentsTab");
+    return new DocumentsTab(this.frame, tabId, tab, documentsTab);
+  }
+
   /** Select the primary global command button.
    *
    * There's only a single "primary" button, but we still require you to pass the label to confirm you're selecting the right button.
@@ -245,6 +274,10 @@ export class DataExplorer {
     return this.frame.getByTestId(`CommandBar/Button:${label}`).and(this.frame.locator("css=button"));
   }
 
+  dialogButton(label: string): Locator {
+    return this.frame.getByTestId(`DialogButton:${label}`).and(this.frame.locator("css=button"));
+  }
+
   /** Select the side panel with the specified title */
   panel(title: string): Locator {
     return this.frame.getByTestId(`Panel:${title}`);
@@ -294,6 +327,26 @@ export class DataExplorer {
     return await this.waitForNode(`${databaseId}/${containerId}`);
   }
 
+  async waitForContainerItemsNode(databaseId: string, containerId: string): Promise<TreeNode> {
+    return await this.waitForNode(`${databaseId}/${containerId}/Items`);
+  }
+
+  async waitForContainerDocumentsNode(databaseId: string, containerId: string): Promise<TreeNode> {
+    return await this.waitForNode(`${databaseId}/${containerId}/Documents`);
+  }
+
+  async waitForCommandBarButton(label: string, timeout?: number): Promise<Locator> {
+    const commandBar = this.commandBarButton(label);
+    await commandBar.waitFor({ state: "visible", timeout });
+    return commandBar;
+  }
+
+  async waitForDialogButton(label: string, timeout?: number): Promise<Locator> {
+    const dialogButton = this.dialogButton(label);
+    await dialogButton.waitFor({ timeout });
+    return dialogButton;
+  }
+
   /** Select the tree node with the specified id */
   treeNode(id: string): TreeNode {
     return new TreeNode(this.frame.getByTestId(`TreeNode:${id}`), this.frame, id);

test/mongo/document.spec.ts

@@ -0,0 +1,89 @@
+import { expect, test } from "@playwright/test";
+
+import { DataExplorer, DocumentsTab, TestAccount } from "../fx";
+import { retry, serializeMongoToJson, setPartitionKeys } from "../testData";
+import { documentTestCases } from "./testCases";
+
+let explorer: DataExplorer = null!;
+let documentsTab: DocumentsTab = null!;
+
+for (const { name, databaseId, containerId, documents } of documentTestCases) {
+  test.describe(`Test MongoRU Documents with ${name}`, () => {
+    test.beforeEach("Open documents tab", async ({ page }) => {
+      explorer = await DataExplorer.open(page, TestAccount.MongoReadonly);
+
+      const containerNode = await explorer.waitForContainerNode(databaseId, containerId);
+      await containerNode.expand();
+
+      const containerMenuNode = await explorer.waitForContainerDocumentsNode(databaseId, containerId);
+      await containerMenuNode.element.click();
+
+      documentsTab = explorer.documentsTab("tab0");
+
+      await documentsTab.documentsFilter.waitFor();
+      await documentsTab.documentsListPane.waitFor();
+      await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
+    });
+
+    for (const document of documents) {
+      const { documentId: docId, partitionKeys } = document;
+      test.describe(`Document ID: ${docId}`, () => {
+        test(`should load and view document ${docId}`, async () => {
+          const span = documentsTab.documentsListPane.getByText(docId, { exact: true }).nth(0);
+          await span.waitFor();
+          await expect(span).toBeVisible();
+
+          await span.click();
+          await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
+
+          const resultText = await documentsTab.resultsEditor.text();
+          const resultData = serializeMongoToJson(resultText!);
+          expect(resultText).not.toBeNull();
+          expect(resultData?._id).not.toBeNull();
+          expect(resultData?._id).toEqual(docId);
+        });
+        test(`should be able to create and delete new document from ${docId}`, async () => {
+          const span = documentsTab.documentsListPane.getByText(docId, { exact: true }).nth(0);
+          await span.waitFor();
+          await expect(span).toBeVisible();
+
+          await span.click();
+          let newDocumentId;
+          await retry(async () => {
+            const newDocumentButton = await explorer.waitForCommandBarButton("New Document", 5000);
+            await expect(newDocumentButton).toBeVisible();
+            await expect(newDocumentButton).toBeEnabled();
+            await newDocumentButton.click();
+
+            await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
+
+            newDocumentId = `${Date.now().toString()}-delete`;
+
+            const newDocument = {
+              _id: newDocumentId,
+              ...setPartitionKeys(partitionKeys || []),
+            };
+
+            await documentsTab.resultsEditor.setText(JSON.stringify(newDocument));
+            const saveButton = await explorer.waitForCommandBarButton("Save", 5000);
+            await saveButton.click({ timeout: 5000 });
+          }, 3);
+
+          const newSpan = documentsTab.documentsListPane.getByText(newDocumentId, { exact: true }).nth(0);
+          await newSpan.waitFor();
+          await newSpan.click();
+          await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
+
+          const deleteButton = await explorer.waitForCommandBarButton("Delete", 5000);
+          await deleteButton.click();
+
+          const deleteDialogButton = await explorer.waitForDialogButton("Delete", 5000);
+          await deleteDialogButton.click();
+
+          const deletedSpan = documentsTab.documentsListPane.getByText(newDocumentId, { exact: true }).nth(0);
+          await expect(deletedSpan).toHaveCount(0);
+        });
+      });
+    }
+  });
+}

test/mongo/testCases.ts

@@ -0,0 +1,31 @@
+import { DocumentTestCase } from "../testData";
+
+export const documentTestCases: DocumentTestCase[] = [
+  {
+    name: "Unsharded Collection",
+    databaseId: "e2etests-mongo-readonly",
+    containerId: "unsharded",
+    documents: [
+      {
+        documentId: "unsharded",
+        partitionKeys: [],
+      },
+    ],
+  },
+  {
+    name: "Sharded Collection",
+    databaseId: "e2etests-mongo-readonly",
+    containerId: "sharded",
+    documents: [
+      {
+        documentId: "sharded",
+        partitionKeys: [
+          {
+            key: "/shardKey",
+            value: "shardKey",
+          },
+        ],
+      },
+    ],
+  },
+];

test/sql/document.spec.ts

@@ -0,0 +1,93 @@
+import { expect, test } from "@playwright/test";
+
+import { DataExplorer, DocumentsTab, TestAccount } from "../fx";
+import { retry, setPartitionKeys } from "../testData";
+import { documentTestCases } from "./testCases";
+
+let explorer: DataExplorer = null!;
+let documentsTab: DocumentsTab = null!;
+
+for (const { name, databaseId, containerId, documents } of documentTestCases) {
+  test.describe(`Test SQL Documents with ${name}`, () => {
+    test.beforeEach("Open documents tab", async ({ page }) => {
+      explorer = await DataExplorer.open(page, TestAccount.SQLReadOnly);
+
+      const containerNode = await explorer.waitForContainerNode(databaseId, containerId);
+      await containerNode.expand();
+
+      const containerMenuNode = await explorer.waitForContainerItemsNode(databaseId, containerId);
+      await containerMenuNode.element.click();
+
+      documentsTab = explorer.documentsTab("tab0");
+
+      await documentsTab.documentsFilter.waitFor();
+      await documentsTab.documentsListPane.waitFor();
+      await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
+    });
+
+    for (const document of documents) {
+      const { documentId: docId, partitionKeys } = document;
+      test.describe(`Document ID: ${docId}`, () => {
+        test(`should load and view document ${docId}`, async () => {
+          const span = documentsTab.documentsListPane.getByText(docId, { exact: true }).nth(0);
+          await span.waitFor();
+          await expect(span).toBeVisible();
+
+          await span.click();
+          await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
+
+          const resultText = await documentsTab.resultsEditor.text();
+          const resultData = JSON.parse(resultText!);
+          expect(resultText).not.toBeNull();
+          expect(resultData?.id).toEqual(docId);
+        });
+        test(`should be able to create and delete new document from ${docId}`, async ({ page }) => {
+          const span = documentsTab.documentsListPane.getByText(docId, { exact: true }).nth(0);
+          await span.waitFor();
+          await expect(span).toBeVisible();
+
+          await span.click();
+          let newDocumentId;
+          await page.waitForTimeout(5000);
+          await retry(async () => {
+            // const discardButton = await explorer.waitForCommandBarButton("Discard", 5000);
+            // if (await discardButton.isEnabled()) {
+            //   await discardButton.click();
+            // }
+            const newDocumentButton = await explorer.waitForCommandBarButton("New Item", 5000);
+            await expect(newDocumentButton).toBeVisible();
+            await expect(newDocumentButton).toBeEnabled();
+            await newDocumentButton.click();
+
+            await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
+
+            newDocumentId = `${Date.now().toString()}-delete`;
+
+            const newDocument = {
+              id: newDocumentId,
+              ...setPartitionKeys(partitionKeys || []),
+            };
+
+            await documentsTab.resultsEditor.setText(JSON.stringify(newDocument));
+            const saveButton = await explorer.waitForCommandBarButton("Save", 5000);
+            await saveButton.click({ timeout: 5000 });
+          }, 3);
+
+          const newSpan = documentsTab.documentsListPane.getByText(newDocumentId, { exact: true }).nth(0);
+          await newSpan.waitFor();
+          await newSpan.click();
+          await expect(documentsTab.resultsEditor.locator).toBeAttached({ timeout: 60 * 1000 });
+
+          const deleteButton = await explorer.waitForCommandBarButton("Delete", 5000);
+          await deleteButton.click();
+
+          const deleteDialogButton = await explorer.waitForDialogButton("Delete", 5000);
+          await deleteDialogButton.click();
+
+          const deletedSpan = documentsTab.documentsListPane.getByText(newDocumentId, { exact: true }).nth(0);
+          await expect(deletedSpan).toHaveCount(0);
+        });
+      });
+    }
+  });
+}

test/sql/testCases.ts

@@ -0,0 +1,235 @@
+import { DocumentTestCase } from "../testData";
+
+export const documentTestCases: DocumentTestCase[] = [
+  {
+    name: "System Partition Key",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "systemPartitionKey",
+    documents: [{ documentId: "systempartition", partitionKeys: [] }],
+  },
+  {
+    name: "Single Partition Key",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "singlePartitionKey",
+    documents: [
+      {
+        documentId: "singlePartitionKey",
+        partitionKeys: [{ key: "/singlePartitionKey", value: "singlePartitionKey" }],
+      },
+      {
+        documentId: "singlePartitionKey_empty_string",
+        partitionKeys: [{ key: "/singlePartitionKey", value: "" }],
+      },
+      {
+        documentId: "singlePartitionKey_null",
+        partitionKeys: [{ key: "/singlePartitionKey", value: null }],
+      },
+      {
+        documentId: "singlePartitionKey_missing",
+        partitionKeys: [],
+      },
+    ],
+  },
+  {
+    name: "Single Nested Partition Key",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "singleNestedPartitionKey",
+    documents: [
+      {
+        documentId: "singlePartitionKey_nested",
+        partitionKeys: [{ key: "/singlePartitionKey/nested", value: "nestedValue" }],
+      },
+      {
+        documentId: "singlePartitionKey_nested_empty_string",
+        partitionKeys: [{ key: "/singlePartitionKey/nested", value: "" }],
+      },
+      {
+        documentId: "singlePartitionKey_nested_null",
+        partitionKeys: [{ key: "/singlePartitionKey/nested", value: null }],
+      },
+      {
+        documentId: "singlePartitionKey_nested_missing",
+        partitionKeys: [],
+      },
+    ],
+  },
+  {
+    name: "2-Level Hierarchical Partition Key",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "twoLevelPartitionKey",
+    documents: [
+      {
+        documentId: "twoLevelPartitionKey_value_empty",
+        partitionKeys: [
+          { key: "/twoLevelPartitionKey_1", value: "value" },
+          { key: "/twoLevelPartitionKey_2", value: "" },
+        ],
+      },
+      {
+        documentId: "twoLevelPartitionKey_value_null",
+        partitionKeys: [
+          { key: "/twoLevelPartitionKey_1", value: "value" },
+          { key: "/twoLevelPartitionKey_2", value: null },
+        ],
+      },
+      {
+        documentId: "twoLevelPartitionKey_value_missing",
+        partitionKeys: [{ key: "/twoLevelPartitionKey_1", value: "value" }],
+      },
+      {
+        documentId: "twoLevelPartitionKey_empty_null",
+        partitionKeys: [
+          { key: "/twoLevelPartitionKey_1", value: "" },
+          { key: "/twoLevelPartitionKey_2", value: null },
+        ],
+      },
+      {
+        documentId: "twoLevelPartitionKey_null_missing",
+        partitionKeys: [{ key: "/twoLevelPartitionKey_1", value: null }],
+      },
+      {
+        documentId: "twoLevelPartitionKey_missing_value",
+        partitionKeys: [{ key: "/twoLevelPartitionKey_2", value: "value" }],
+      },
+    ],
+  },
+  {
+    name: "2-Level Hierarchical Nested Partition Key",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "twoLevelNestedPartitionKey",
+    documents: [
+      {
+        documentId: "twoLevelNestedPartitionKey_nested_value_empty",
+        partitionKeys: [
+          { key: "/twoLevelNestedPartitionKey/nested", value: "value" },
+          { key: "/twoLevelNestedPartitionKey/nested_value/nested", value: "" },
+        ],
+      },
+      {
+        documentId: "twoLevelNestedPartitionKey_nested_null_missing",
+        partitionKeys: [{ key: "/twoLevelNestedPartitionKey/nested", value: null }],
+      },
+      {
+        documentId: "twoLevelNestedPartitionKey_nested_missing_value",
+        partitionKeys: [{ key: "/twoLevelNestedPartitionKey/nested_value/nested", value: "value" }],
+      },
+    ],
+  },
+  {
+    name: "3-Level Hierarchical Partition Key",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "threeLevelPartitionKey",
+    documents: [
+      {
+        documentId: "threeLevelPartitionKey_value_empty_null",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1", value: "value" },
+          { key: "/threeLevelPartitionKey_2", value: "" },
+          { key: "/threeLevelPartitionKey_3", value: null },
+        ],
+      },
+      {
+        documentId: "threeLevelPartitionKey_value_null_missing",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1", value: "value" },
+          { key: "/threeLevelPartitionKey_2", value: null },
+        ],
+      },
+      {
+        documentId: "threeLevelPartitionKey_value_missing_null",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1", value: "value" },
+          { key: "/threeLevelPartitionKey_3", value: null },
+        ],
+      },
+      {
+        documentId: "threeLevelPartitionKey_null_empty_value",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1", value: null },
+          { key: "/threeLevelPartitionKey_2", value: "" },
+          { key: "/threeLevelPartitionKey_3", value: "value" },
+        ],
+      },
+      {
+        documentId: "threeLevelPartitionKey_missing_value_value",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_2", value: "value" },
+          { key: "/threeLevelPartitionKey_3", value: "value" },
+        ],
+      },
+      {
+        documentId: "threeLevelPartitionKey_empty_value_missing",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1", value: "" },
+          { key: "/threeLevelPartitionKey_2", value: "value" },
+        ],
+      },
+    ],
+  },
+  {
+    name: "3-Level Nested Hierarchical Partition Key",
+    databaseId: "e2etests-sql-readonly",
+    containerId: "threeLevelNestedPartitionKey",
+    documents: [
+      {
+        documentId: "threeLevelNestedPartitionKey_nested_empty_value_null",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1/nested/nested_key", value: "" },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_key", value: "value" },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_3/nested_key", value: null },
+        ],
+      },
+      {
+        documentId: "threeLevelNestedPartitionKey_nested_null_value_missing",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1/nested/nested_key", value: null },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_key", value: "value" },
+        ],
+      },
+      {
+        documentId: "threeLevelNestedPartitionKey_nested_missing_value_null",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_key", value: "value" },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_3/nested_key", value: null },
+        ],
+      },
+      {
+        documentId: "threeLevelNestedPartitionKey_nested_null_empty_missing",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1/nested/nested_key", value: null },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_key", value: "" },
+        ],
+      },
+      {
+        documentId: "threeLevelNestedPartitionKey_nested_value_missing_empty",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1/nested/nested_key", value: "value" },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_3/nested_key", value: "" },
+        ],
+      },
+      {
+        documentId: "threeLevelNestedPartitionKey_nested_missing_null_empty",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_key", value: null },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_3/nested_key", value: "" },
+        ],
+      },
+      {
+        documentId: "threeLevelNestedPartitionKey_nested_empty_null_value",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1/nested/nested_key", value: "" },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_key", value: null },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_3/nested_key", value: "value" },
+        ],
+      },
+      {
+        documentId: "threeLevelNestedPartitionKey_nested_value_null_empty",
+        partitionKeys: [
+          { key: "/threeLevelPartitionKey_1/nested/nested_key", value: "value" },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_key", value: null },
+          { key: "/threeLevelPartitionKey_1/nested/nested_2/nested_3/nested_key", value: "" },
+        ],
+      },
+    ],
+  },
+];

test/testData.ts

@@ -16,6 +16,23 @@ export interface TestItem {
   randomData: string;
 }
 
+export interface DocumentTestCase {
+  name: string;
+  databaseId: string;
+  containerId: string;
+  documents: TestDocument[];
+}
+
+export interface TestDocument {
+  documentId: string;
+  partitionKeys?: PartitionKey[];
+}
+
+export interface PartitionKey {
+  key: string;
+  value: string | null;
+}
+
 const partitionCount = 4;
 
 // If we increase this number, we need to split bulk creates into multiple batches.
@@ -93,3 +110,46 @@ export async function createTestSQLContainer(includeTestData?: boolean) {
     throw e;
   }
 }
+
+export const setPartitionKeys = (partitionKeys: PartitionKey[]) => {
+  const result = {};
+
+  partitionKeys.forEach((partitionKey) => {
+    const { key: keyPath, value: keyValue } = partitionKey;
+    const cleanPath = keyPath.startsWith("/") ? keyPath.slice(1) : keyPath;
+    const keys = cleanPath.split("/");
+    let current = result;
+
+    keys.forEach((key, index) => {
+      if (index === keys.length - 1) {
+        current[key] = keyValue;
+      } else {
+        current[key] = current[key] || {};
+        current = current[key];
+      }
+    });
+  });
+
+  return result;
+};
+
+export const serializeMongoToJson = (text: string) => {
+  const normalized = text.replace(/ObjectId\("([0-9a-fA-F]{24})"\)/g, '"$1"');
+  return JSON.parse(normalized);
+};
+
+export async function retry<T>(fn: () => Promise<T>, retries = 3, delayMs = 1000): Promise<T> {
+  let lastError: unknown;
+  for (let i = 0; i < retries; i++) {
+    try {
+      return await fn();
+    } catch (error) {
+      lastError = error;
+      console.warn(`Retry ${i + 1}/${retries} failed: ${(error as Error).message}`);
+      if (i < retries - 1) {
+        await new Promise((res) => setTimeout(res, delayMs));
+      }
+    }
+  }
+  throw lastError;
+}