diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.test.ts b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.test.ts
new file mode 100644
index 000000000..ecec95576
--- /dev/null
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.test.ts
@@ -0,0 +1,99 @@
+import { buildQuery, showPartitionKey } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabV2";
+import * as ko from "knockout";
+import { DatabaseAccount } from "../../../Contracts/DataModels";
+import * as ViewModels from "../../../Contracts/ViewModels";
+import { updateUserContext } from "../../../UserContext";
+import Explorer from "../../Explorer";
+
+describe("Documents tab", () => {
+  describe("buildQuery", () => {
+    it("should generate the right select query for SQL API", () => {
+      expect(buildQuery(false, "")).toContain("select");
+    });
+  });
+
+  describe("showPartitionKey", () => {
+    const explorer = new Explorer();
+    const mongoExplorer = new Explorer();
+    updateUserContext({
+      databaseAccount: {
+        properties: {
+          capabilities: [{ name: "EnableGremlin" }],
+        },
+      } as DatabaseAccount,
+    });
+
+    const collectionWithoutPartitionKey = <ViewModels.Collection>(<unknown>{
+      id: ko.observable<string>("foo"),
+      database: {
+        id: ko.observable<string>("foo"),
+      },
+      container: explorer,
+    });
+
+    const collectionWithSystemPartitionKey = <ViewModels.Collection>(<unknown>{
+      id: ko.observable<string>("foo"),
+      database: {
+        id: ko.observable<string>("foo"),
+      },
+      partitionKey: {
+        paths: ["/foo"],
+        kind: "Hash",
+        version: 2,
+        systemKey: true,
+      },
+      container: explorer,
+    });
+
+    const collectionWithNonSystemPartitionKey = <ViewModels.Collection>(<unknown>{
+      id: ko.observable<string>("foo"),
+      database: {
+        id: ko.observable<string>("foo"),
+      },
+      partitionKey: {
+        paths: ["/foo"],
+        kind: "Hash",
+        version: 2,
+        systemKey: false,
+      },
+      container: explorer,
+    });
+
+    const mongoCollectionWithSystemPartitionKey = <ViewModels.Collection>(<unknown>{
+      id: ko.observable<string>("foo"),
+      database: {
+        id: ko.observable<string>("foo"),
+      },
+      partitionKey: {
+        paths: ["/foo"],
+        kind: "Hash",
+        version: 2,
+        systemKey: true,
+      },
+      container: mongoExplorer,
+    });
+
+    it("should be false for null or undefined collection", () => {
+      expect(showPartitionKey(undefined, false)).toBe(false);
+      expect(showPartitionKey(null, false)).toBe(false);
+      expect(showPartitionKey(undefined, true)).toBe(false);
+      expect(showPartitionKey(null, true)).toBe(false);
+    });
+
+    it("should be false for null or undefined partitionKey", () => {
+      expect(showPartitionKey(collectionWithoutPartitionKey, false)).toBe(false);
+    });
+
+    it("should be true for non-Mongo accounts with system partitionKey", () => {
+      expect(showPartitionKey(collectionWithSystemPartitionKey, false)).toBe(true);
+    });
+
+    it("should be false for Mongo accounts with system partitionKey", () => {
+      expect(showPartitionKey(mongoCollectionWithSystemPartitionKey, true)).toBe(false);
+    });
+
+    it("should be true for non-system partitionKey", () => {
+      expect(showPartitionKey(collectionWithNonSystemPartitionKey, false)).toBe(true);
+    });
+  });
+});
diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
index a9a440c64..dd84cc976 100644
--- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
@@ -375,6 +375,37 @@ const _loadNextPageInternal = (
   return iterator.fetchNext().then((response) => response.resources);
 };
 
+// Export for testing purposes
+export const showPartitionKey = (collection: ViewModels.CollectionBase, isPreferredApiMongoDB: boolean) => {
+  if (!collection) {
+    return false;
+  }
+
+  if (!collection.partitionKey) {
+    return false;
+  }
+
+  if (collection.partitionKey.systemKey && isPreferredApiMongoDB) {
+    return false;
+  }
+
+  return true;
+};
+
+// Export for testing purposes
+export const buildQuery = (
+  isMongo: boolean,
+  filter: string,
+  partitionKeyProperties?: string[],
+  partitionKey?: DataModels.PartitionKey,
+): string => {
+  if (isMongo) {
+    return filter || "{}";
+  }
+
+  return QueryUtils.buildDocumentsQuery(filter, partitionKeyProperties, partitionKey);
+};
+
 const DocumentsTabComponent: React.FunctionComponent<{
   isPreferredApiMongoDB: boolean;
   documentIds: DocumentId[]; // TODO: this contains ko observables. We need to convert them to React state.
@@ -882,18 +913,11 @@ const DocumentsTabComponent: React.FunctionComponent<{
     [isPreferredApiMongoDB],
   );
 
-  let buildQuery = useCallback(
-    (filter: string): string => {
-      return QueryUtils.buildDocumentsQuery(filter, partitionKeyProperties, partitionKey);
-    },
-    [partitionKeyProperties, partitionKey],
-  );
-
   const createIterator = useCallback((): QueryIterator<ItemDefinition & Resource> => {
     const _queryAbortController = new AbortController();
     setQueryAbortController(_queryAbortController);
     const filter: string = filterContent.trim();
-    const query: string = buildQuery(filter);
+    const query: string = buildQuery(isPreferredApiMongoDB, filter, partitionKeyProperties, partitionKey);
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const options: any = {};
     // TODO: Property 'enableCrossPartitionQuery' does not exist on type 'FeedOptions'.
@@ -908,7 +932,7 @@ const DocumentsTabComponent: React.FunctionComponent<{
     return isQueryCopilotSampleContainer
       ? querySampleDocuments(query, options)
       : queryDocuments(_collection.databaseId, _collection.id(), query, options);
-  }, [filterContent, buildQuery, resourceTokenPartitionKey, isQueryCopilotSampleContainer, _collection]);
+  }, [isPreferredApiMongoDB, filterContent, resourceTokenPartitionKey, isQueryCopilotSampleContainer, _collection]);
 
   /**
    * Query first page of documents
@@ -1115,23 +1139,6 @@ const DocumentsTabComponent: React.FunctionComponent<{
     }
   };
 
-  // TODO: use this when generating column headers
-  const showPartitionKey = (() => {
-    if (!_collection) {
-      return false;
-    }
-
-    if (!_collection.partitionKey) {
-      return false;
-    }
-
-    if (_collection.partitionKey.systemKey && isPreferredApiMongoDB) {
-      return false;
-    }
-
-    return true;
-  })();
-
   const _isQueryCopilotSampleContainer =
     _collection?.isSampleCollection &&
     _collection?.databaseId === QueryCopilotSampleDatabaseId &&
@@ -1259,7 +1266,7 @@ const DocumentsTabComponent: React.FunctionComponent<{
 
   const columnHeaders = {
     idHeader: isPreferredApiMongoDB ? "_id" : "id",
-    partitionKeyHeaders: (showPartitionKey && partitionKeyPropertyHeaders) || [],
+    partitionKeyHeaders: (showPartitionKey(_collection, isPreferredApiMongoDB) && partitionKeyPropertyHeaders) || [],
   };
 
   const onSelectedRowsChange = (selectedRows: Set<TableRowId>) => {
@@ -1549,15 +1556,11 @@ const DocumentsTabComponent: React.FunctionComponent<{
         .finally(() => setIsExecuting(false));
     };
 
-    buildQuery = (filter: string): string => {
-      return filter || "{}";
-    };
-
     loadNextPage = (): Promise<unknown> => {
       setIsExecuting(true);
       onExecutionErrorChange(false);
       const filter: string = filterContent.trim();
-      const query: string = buildQuery(filter);
+      const query: string = buildQuery(isPreferredApiMongoDB, filter);
 
       return MongoProxyClient.queryDocuments(
         _collection.databaseId,
@@ -1668,12 +1671,12 @@ const DocumentsTabComponent: React.FunctionComponent<{
       <div
         className="tab-pane active"
         /* data-bind="
-                                                                                  setTemplateReady: true,
-                                                                                  attr:{
-                                                                                      id: tabId
-                                                                                  },
-                                                                                  visible: isActive"
-                                                                                  */
+                                                                                            setTemplateReady: true,
+                                                                                            attr:{
+                                                                                                id: tabId
+                                                                                            },
+                                                                                            visible: isActive"
+                                                                                            */
         role="tabpanel"
         style={{ display: "flex" }}
       >
@@ -1766,9 +1769,9 @@ const DocumentsTabComponent: React.FunctionComponent<{
                         onClick={() => refreshDocumentsGrid(true)}
                         disabled={!applyFilterButton.enabled}
                         /* data-bind="
-                                                                                                              click: refreshDocumentsGrid.bind($data, true),
-                                                                                                              enable: applyFilterButton.enabled"
-                                                                                                    */
+                                                                                                                        click: refreshDocumentsGrid.bind($data, true),
+                                                                                                                        enable: applyFilterButton.enabled"
+                                                                                                              */
                         aria-label="Apply filter"
                         tabIndex={0}
                       >
@@ -1781,9 +1784,9 @@ const DocumentsTabComponent: React.FunctionComponent<{
                           style={filterButtonStyle}
                           appearance="primary"
                           /* data-bind="
-                                                                                                                visible: !isPreferredApiMongoDB && isExecuting,
-                                                                                                                click: onAbortQueryClick"
-                                                                                                      */
+                                                                                                                          visible: !isPreferredApiMongoDB && isExecuting,
+                                                                                                                          click: onAbortQueryClick"
+                                                                                                                */
                           aria-label="Cancel Query"
                           onClick={() => queryAbortController.abort()}
                           tabIndex={0}