From c0abc26fe822426004e2dc805a60c63f0687a440 Mon Sep 17 00:00:00 2001
From: Diego Fabricio <difagume@gmail.com>
Date: Thu, 27 Aug 2026 22:12:41 -0500
Subject: [PATCH] fix(grid): allow rerunning exact total row count from status

---
 apps/desktop/src/components/grid/DataGrid.vue | 25 +++++++-
 .../components/grid/DataGridPagination.vue    |  2 +-
 .../grid/__tests__/DataGridSurfaces.spec.ts   | 35 ++++++++++++
 .../__tests__/DataGridTotalRowCount.spec.ts   | 57 ++++++++++++++++++-
 .../__tests__/ContentArea.outputView.spec.ts  | 19 +++++++
 .../src/lib/dataGrid/dataGridPagination.ts    |  5 ++
 packages/app-tests/dataGridPagination.test.ts | 26 +++++++++
 7 files changed, 165 insertions(+), 4 deletions(-)

diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue
index 06ee9d2a10..4add0d6613 100644
--- a/apps/desktop/src/components/grid/DataGrid.vue
+++ b/apps/desktop/src/components/grid/DataGrid.vue
@@ -38,6 +38,7 @@ import {
   PanelBottom,
   PanelRight,
   RefreshCw,
+  RefreshCcw,
   TableProperties,
   UserRound,
   Database,
@@ -203,7 +204,7 @@ import {
   isToggleTransposeShortcut,
 } from "@/lib/editor/keyboardShortcuts";
 import { dataGridHeaderContentWidth, scrollbarGutterWidth } from "@/lib/dataGrid/dataGridScrollGutter";
-import { canFetchNextDataGridSegment, canGoNextDataGridPage, dataGridTotalRowCountLabelKey, dataGridTruncationHintKey, hasCompleteLocalDataGridResult, resolveDataGridPaginationTotal, type DataGridInexactTotalRowCountMode } from "@/lib/dataGrid/dataGridPagination";
+import { canFetchNextDataGridSegment, canGoNextDataGridPage, dataGridTotalRowCountLabelKey, dataGridTruncationHintKey, hasCompleteLocalDataGridResult, resolveDataGridPaginationTotal, showDataGridRerunTotalCountAction, type DataGridInexactTotalRowCountMode } from "@/lib/dataGrid/dataGridPagination";
 import { dataGridCountQueryOptions } from "@/lib/dataGrid/dataGridQueryOptions";
 import {
   createResultScopedPendingRequests,
@@ -3560,6 +3561,13 @@ watch(
 );
 const canCalculateTotalRowCount = computed(() => !!props.countTotalRows || (!!props.connectionId && (!!props.tableMeta || !!props.countSql)));
 const showExactTotalCountAction = computed(() => canCalculateTotalRowCount.value && (totalRowCountIsExact.value === false || typeof displayedTotalRowCount.value !== "number"));
+const showRerunTotalCountAction = computed(() =>
+  showDataGridRerunTotalCountAction({
+    canCalculateTotalRowCount: canCalculateTotalRowCount.value,
+    displayedTotalRowCount: displayedTotalRowCount.value,
+    totalRowCountIsExact: totalRowCountIsExact.value,
+  }),
+);
 watch(
   [
     () =>
@@ -13963,12 +13971,25 @@ function openGridSnapshot() {
             })
           }}
           <span v-if="typeof displayedTotalRowCount === 'number' && displayedTotalRowCount >= 0" class="text-muted-foreground/70">{{ t(totalRowCountLabelKey, { count: displayedTotalRowCount }) }}</span>
-          <span v-if="totalRowCountBusy" class="text-muted-foreground/70">
+          <span v-if="totalRowCountBusy && !(showRerunTotalCountAction && manualTotalRowCountLoading)" class="text-muted-foreground/70">
             {{ t("grid.totalRowCountLoading") }}
           </span>
           <button v-else-if="showExactTotalCountAction" type="button" class="text-muted-foreground/70 underline underline-offset-2 hover:text-foreground disabled:pointer-events-none" :disabled="manualTotalRowCountLoading" @click="calculateTotalRowCount">
             {{ t("grid.calculateTotalRowsInline") }}
           </button>
+          <button
+            v-else-if="showRerunTotalCountAction"
+            type="button"
+            class="ml-1 inline-flex h-3.5 w-3.5 shrink-0 items-center justify-center rounded-sm align-middle text-muted-foreground/70 hover:bg-gray-200 hover:text-foreground disabled:pointer-events-none disabled:opacity-50 dark:hover:bg-gray-800"
+            :disabled="manualTotalRowCountLoading"
+            :aria-busy="manualTotalRowCountLoading ? 'true' : undefined"
+            :title="manualTotalRowCountLoading ? t('grid.totalRowCountLoading') : t('grid.calculateTotalRows')"
+            :aria-label="manualTotalRowCountLoading ? t('grid.totalRowCountLoading') : t('grid.calculateTotalRows')"
+            @click="calculateTotalRowCount"
+          >
+            <Loader2 v-if="manualTotalRowCountLoading" aria-hidden="true" class="h-3 w-3 animate-spin" />
+            <RefreshCcw v-else aria-hidden="true" class="h-3 w-3" />
+          </button>
         </span>
         <span v-if="showTruncationWarning" class="shrink-0 text-amber-500 text-xs">(truncated)</span>
         <span v-if="!hasData" class="shrink-0">{{ t("grid.rowsAffected", { count: result.affected_rows }) }}</span>
diff --git a/apps/desktop/src/components/grid/DataGridPagination.vue b/apps/desktop/src/components/grid/DataGridPagination.vue
index 1c4651e5e0..f26beb770c 100644
--- a/apps/desktop/src/components/grid/DataGridPagination.vue
+++ b/apps/desktop/src/components/grid/DataGridPagination.vue
@@ -93,7 +93,7 @@ function handlePageInputKeydown(event: KeyboardEvent) {
         <span class="shrink-0">{{ t("grid.rows", { count: selectionSummary.rowCount }) }}</span>
       </div>
     </div>
-    <Loader2 v-if="loading" class="w-3 h-3 animate-spin text-muted-foreground" />
+    <Loader2 class="w-3 h-3 text-muted-foreground" :class="loading ? 'animate-spin' : 'invisible'" aria-hidden="true" />
     <template v-if="paginationEnabled && infiniteScrollEnabled">
       <span v-if="infiniteScrollAllLoaded" class="text-xs text-muted-foreground shrink-0">{{ t("grid.allLoaded") }}</span>
     </template>
diff --git a/apps/desktop/src/components/grid/__tests__/DataGridSurfaces.spec.ts b/apps/desktop/src/components/grid/__tests__/DataGridSurfaces.spec.ts
index e1cf49ce40..457a43bf2e 100644
--- a/apps/desktop/src/components/grid/__tests__/DataGridSurfaces.spec.ts
+++ b/apps/desktop/src/components/grid/__tests__/DataGridSurfaces.spec.ts
@@ -250,6 +250,41 @@ describe("DataGridPagination", () => {
     expect(summaryText).toContain("grid.rows");
   });
 
+  it("keeps the loading icon mounted with a constant footprint and only animates while loading", async () => {
+    const mounted = mountComponent(DataGridPagination, {
+      selectionSummary: null,
+      selectionSummarySumText: "",
+      selectionSummaryAverageText: "",
+      loading: false,
+      infiniteScrollEnabled: false,
+      infiniteScrollAllLoaded: false,
+      pageSize: 100,
+      customPageSizeInput: "",
+      pageSizeMenuItems: [],
+      exportMenuItems: [],
+      currentPage: 1,
+      canGoNextPage: false,
+      canJumpLastPage: false,
+    });
+    // Stable within this file's mountComponent+findAll format: Loader2 is the only Icon
+    // combining text-muted-foreground with w-3 h-3 in the pagination bar. Filtering by
+    // both avoids matching a future muted icon without introducing data-testid.
+    const findLoader = () => findAll(mounted.root, (node) => node.props["data-stub"] === "Icon" && String(node.props.class).includes("text-muted-foreground") && String(node.props.class).includes("w-3 h-3"));
+
+    const idleLoaders = findLoader();
+    expect(idleLoaders).toHaveLength(1);
+    expect(String(idleLoaders[0].props.class)).toContain("invisible");
+    expect(String(idleLoaders[0].props.class)).not.toContain("animate-spin");
+    expect(String(idleLoaders[0].props["aria-hidden"])).toContain("true");
+
+    await mounted.setProps({ loading: true });
+    const busyLoaders = findLoader();
+    expect(busyLoaders).toHaveLength(1);
+    expect(String(busyLoaders[0].props.class)).toContain("animate-spin");
+    expect(String(busyLoaders[0].props.class)).not.toContain("invisible");
+    expect(String(busyLoaders[0].props["aria-hidden"])).toContain("true");
+  });
+
   it("enforces first/previous/next/last disabled boundaries", async () => {
     const firstPage = vi.fn();
     const previousPage = vi.fn();
diff --git a/apps/desktop/src/components/grid/__tests__/DataGridTotalRowCount.spec.ts b/apps/desktop/src/components/grid/__tests__/DataGridTotalRowCount.spec.ts
index 8158106183..7eb4fdf627 100644
--- a/apps/desktop/src/components/grid/__tests__/DataGridTotalRowCount.spec.ts
+++ b/apps/desktop/src/components/grid/__tests__/DataGridTotalRowCount.spec.ts
@@ -1,9 +1,20 @@
 import { describe, expect, it } from "vitest";
-import { dataGridTotalRowCountLabelKey, dataGridTruncationHintKey } from "@/lib/dataGrid/dataGridPagination";
+import { dataGridTotalRowCountLabelKey, dataGridTruncationHintKey, showDataGridRerunTotalCountAction } from "@/lib/dataGrid/dataGridPagination";
 import DataGrid from "../DataGrid.vue";
 
 type VuePropDefinition = { default?: unknown };
 type VueComponentWithProps = { props?: Record<string, VuePropDefinition> };
+type VueComponentWithSsrRender = { ssrRender?: unknown };
+
+const statusRenderSource = (() => {
+  const component = DataGrid as unknown as VueComponentWithSsrRender;
+  return typeof component.ssrRender === "function" ? String(component.ssrRender) : "";
+})();
+
+const rerunBranch = (() => {
+  const start = statusRenderSource.indexOf("} else if ($setup.showRerunTotalCountAction) {");
+  return start === -1 ? "" : statusRenderSource.slice(start, statusRenderSource.indexOf("</button>", start));
+})();
 
 describe("DataGrid total row count exactness", () => {
   it("uses a VictoriaMetrics-specific truncation explanation", () => {
@@ -24,3 +35,47 @@ describe("DataGrid total row count exactness", () => {
     expect(dataGridTotalRowCountLabelKey(false, "estimated")).toBe("grid.totalRowCountEstimated");
   });
 });
+
+describe("DataGrid rerun total row count action", () => {
+  it("appears only for a displayed exact numeric total when counting is possible", () => {
+    expect(showDataGridRerunTotalCountAction({ canCalculateTotalRowCount: true, displayedTotalRowCount: 4200, totalRowCountIsExact: true })).toBe(true);
+    expect(showDataGridRerunTotalCountAction({ canCalculateTotalRowCount: true, displayedTotalRowCount: 0, totalRowCountIsExact: true })).toBe(true);
+    expect(showDataGridRerunTotalCountAction({ canCalculateTotalRowCount: true, totalRowCountIsExact: true })).toBe(false);
+    expect(showDataGridRerunTotalCountAction({ canCalculateTotalRowCount: true, displayedTotalRowCount: -1, totalRowCountIsExact: true })).toBe(false);
+    expect(showDataGridRerunTotalCountAction({ canCalculateTotalRowCount: false, displayedTotalRowCount: 4200, totalRowCountIsExact: true })).toBe(false);
+    expect(showDataGridRerunTotalCountAction({ canCalculateTotalRowCount: true, displayedTotalRowCount: 4200, totalRowCountIsExact: false })).toBe(false);
+    expect(showDataGridRerunTotalCountAction({ canCalculateTotalRowCount: true, displayedTotalRowCount: Infinity, totalRowCountIsExact: true })).toBe(false);
+    expect(showDataGridRerunTotalCountAction({ canCalculateTotalRowCount: true, displayedTotalRowCount: NaN, totalRowCountIsExact: true })).toBe(false);
+  });
+
+  it("keeps the numeric total outside the action chain and orders busy, inline link, then rerun icon", () => {
+    const numericTotal = statusRenderSource.indexOf('typeof $setup.displayedTotalRowCount === "number" && $setup.displayedTotalRowCount >= 0');
+    const busyBranch = statusRenderSource.indexOf("if ($setup.totalRowCountBusy && !($setup.showRerunTotalCountAction && $setup.manualTotalRowCountLoading)) {");
+    const inlineLink = statusRenderSource.indexOf("} else if ($setup.showExactTotalCountAction) {");
+    const rerunAction = statusRenderSource.indexOf("} else if ($setup.showRerunTotalCountAction) {");
+    expect(numericTotal).toBeGreaterThan(-1);
+    expect(busyBranch).toBeGreaterThan(numericTotal);
+    expect(inlineLink).toBeGreaterThan(busyBranch);
+    expect(rerunAction).toBeGreaterThan(inlineLink);
+  });
+
+  it("renders a spaced rerun icon labelled for the total-count action and disabled while counting", () => {
+    expect(rerunBranch).toContain('<button type="button" class="ml-1 inline-flex h-3.5 w-3.5');
+    expect(rerunBranch).toContain("disabled:pointer-events-none");
+    expect(rerunBranch).toContain("disabled:opacity-50");
+    // title/aria-label toggle between idle and busy copy
+    expect(rerunBranch).toContain('t("grid.calculateTotalRows")');
+    expect(rerunBranch).toContain('t("grid.totalRowCountLoading")');
+    expect(rerunBranch).toContain("disabled");
+    expect(rerunBranch).toContain("aria-busy");
+  });
+
+  it("swaps the rerun icon for a spinner only while the manual count is in flight and hides decorative icons", () => {
+    expect(rerunBranch).toContain("if ($setup.manualTotalRowCountLoading) {");
+    expect(rerunBranch).toContain("Loader2");
+    expect(rerunBranch).toContain("RefreshCcw");
+    expect(rerunBranch).toContain("h-3 w-3 animate-spin");
+    expect(rerunBranch).toContain("h-3 w-3");
+    expect(rerunBranch).toContain("aria-hidden");
+  });
+});
diff --git a/apps/desktop/src/components/layout/__tests__/ContentArea.outputView.spec.ts b/apps/desktop/src/components/layout/__tests__/ContentArea.outputView.spec.ts
index 06c95d56c4..788d916f66 100644
--- a/apps/desktop/src/components/layout/__tests__/ContentArea.outputView.spec.ts
+++ b/apps/desktop/src/components/layout/__tests__/ContentArea.outputView.spec.ts
@@ -15,3 +15,22 @@ describe("ContentArea Mongo tab reuse", () => {
     expect(contentAreaSource).toContain(':key="`${activeTab.id}:${activeTab.sql}`"');
   });
 });
+
+describe("ContentArea query result grid", () => {
+  it("passes connection and per-run count SQL to the query-tab grid so the total can be re-counted", () => {
+    const countSqlProp = ':count-sql="activeTab.resultCountSql"';
+    const countSqlIndex = contentAreaSource.indexOf(countSqlProp);
+    expect(countSqlIndex).toBeGreaterThan(-1);
+    // Robust string-matching: anchor on the per-run count SQL and walk back to
+    // the enclosing DataGrid so the test does not depend on DataGrid ordering
+    // while staying in readFileSync format.
+    const dataGridIndex = contentAreaSource.lastIndexOf("<DataGrid", countSqlIndex);
+    expect(dataGridIndex).toBeGreaterThan(-1);
+    const queryGridBlock = contentAreaSource.slice(dataGridIndex, countSqlIndex + 900);
+    expect(queryGridBlock).toContain("<DataGrid");
+    expect(queryGridBlock).toContain(':connection-id="activeResultConnectionId"');
+    expect(queryGridBlock).toContain(':count-sql="activeTab.resultCountSql"');
+    expect(queryGridBlock).toContain(':total-row-count="activeTab.resultTotalRowCount"');
+    expect(queryGridBlock).toContain(':total-row-count-loading="activeTab.resultTotalRowCountLoading"');
+  });
+});
diff --git a/apps/desktop/src/lib/dataGrid/dataGridPagination.ts b/apps/desktop/src/lib/dataGrid/dataGridPagination.ts
index 0abb9160fd..80348bbd67 100644
--- a/apps/desktop/src/lib/dataGrid/dataGridPagination.ts
+++ b/apps/desktop/src/lib/dataGrid/dataGridPagination.ts
@@ -42,6 +42,11 @@ export function dataGridTotalRowCountLabelKey(totalRowCountIsExact: boolean, ine
   return inexactMode === "estimated" ? "grid.totalRowCountEstimated" : "grid.totalRowCountAtLeast";
 }
 
+export function showDataGridRerunTotalCountAction(options: { canCalculateTotalRowCount: boolean; displayedTotalRowCount?: number; totalRowCountIsExact: boolean }): boolean {
+  if (!options.canCalculateTotalRowCount || options.totalRowCountIsExact !== true) return false;
+  return typeof options.displayedTotalRowCount === "number" && Number.isFinite(options.displayedTotalRowCount) && options.displayedTotalRowCount >= 0;
+}
+
 export function resolveDataGridPaginationTotal(options: { paginationTotalRowCount?: number; serverKnownTotalRowCount?: number; totalRowCountIsExact: boolean; maxRows?: number }): number | undefined {
   const total = options.paginationTotalRowCount ?? (options.totalRowCountIsExact ? options.serverKnownTotalRowCount : undefined);
   if (total === undefined || options.maxRows === undefined) return total;
diff --git a/packages/app-tests/dataGridPagination.test.ts b/packages/app-tests/dataGridPagination.test.ts
index b08209be6c..d84502be55 100644
--- a/packages/app-tests/dataGridPagination.test.ts
+++ b/packages/app-tests/dataGridPagination.test.ts
@@ -261,3 +261,29 @@ test("row number gutter width tracks the largest visible row index", () => {
   assert.match(source, /resolveDataGridMaxRowNumber/);
   assert.match(source, /rowNumberWidth,/);
 });
+
+test("rerun total-count visibility comes from the shared rule and triggers the manual count", () => {
+  const source = readFileSync("apps/desktop/src/components/grid/DataGrid.vue", "utf8");
+  const rerunComputed = source.match(/const showRerunTotalCountAction = computed\(\(\) =>[\s\S]*?\n\);/)?.[0] ?? "";
+  assert.match(rerunComputed, /showDataGridRerunTotalCountAction\(\{/);
+  assert.match(rerunComputed, /canCalculateTotalRowCount: canCalculateTotalRowCount\.value/);
+  assert.match(rerunComputed, /displayedTotalRowCount: displayedTotalRowCount\.value/);
+  assert.match(rerunComputed, /totalRowCountIsExact: totalRowCountIsExact\.value/);
+  const rerunButton = source.match(/<button\s+v-else-if="showRerunTotalCountAction"[\s\S]*?<\/button>/)?.[0] ?? "";
+  assert.match(rerunButton, /@click="calculateTotalRowCount"/);
+});
+
+test("manual rerun counts through props.countSql and resets when the query context changes", () => {
+  const source = readFileSync("apps/desktop/src/components/grid/DataGrid.vue", "utf8");
+  assert.match(source, /if \(props\.countSql\) return \{ sql: props\.countSql, schema: props\.schema \};/);
+  assert.match(source, /const serverKnownTotalRowCount = computed\(\(\) => \(typeof manualTotalRowCount\.value === "number" \? manualTotalRowCount\.value : props\.totalRowCount\)\)/);
+  assert.match(source, /watch\(\s*\(\) => \[props\.countSql \?\? ""/);
+});
+
+test("executed plans refresh the count SQL shared with the background count", () => {
+  const source = readFileSync("apps/desktop/src/stores/queryStore.ts", "utf8");
+  assert.ok((source.match(/current\.resultCountSql = countSql;/g)?.length ?? 0) >= 2, "every executed-plan path must refresh the tab count SQL");
+  assert.match(source, /countSql = sqlServerUseScript && plan\.countSql \? replaceSqlServerLeadingUseQuery\(queryBaseSql, sqlServerUseScript, plan\.countSql\) : plan\.countSql;/);
+  const backgroundCount = source.match(/countQueryTotalRowsInBackground\(\{[\s\S]*?\}\);/)?.[0] ?? "";
+  assert.match(backgroundCount, /\n\s+countSql,/);
+});