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

src/Explorer/Controls/Dialog.tsx

@@ -214,8 +214,10 @@ export const Dialog: FC = () => {
       {contentHtml}
       {progressIndicatorProps && <ProgressIndicator {...progressIndicatorProps} />}
       <DialogFooter>
-        <PrimaryButton {...primaryButtonProps} />
-        {secondaryButtonProps && <DefaultButton {...secondaryButtonProps} />}
+        <PrimaryButton {...primaryButtonProps} data-test={`DialogButton:${primaryButtonText}`} />
+        {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,
-    [_collection?.partitionKeyPropertyHeaders, partitionKey?.paths],
-  );
-  let partitionKeyProperties = useMemo(
-    () =>
-      partitionKeyPropertyHeaders?.map((partitionKeyPropertyHeader) =>
-        partitionKeyPropertyHeader.replace(/[/]+/g, ".").substring(1).replace(/[']+/g, ""),
-      ),
-    [partitionKeyPropertyHeaders],
+    () => (partitionKey?.systemKey ? [] : _collection?.partitionKeyPropertyHeaders || partitionKey?.paths),
+    [_collection?.partitionKeyPropertyHeaders, partitionKey?.paths, partitionKey?.systemKey],
   );
+  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 className={styles.filterRow}>
+      <div data-test={"DocumentsTab"} className="tab-pane active" role="tabpanel" style={{ display: "flex" }}>
+        <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%",