perf(query): optimize result loading and connection reuse

apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts

@@ -2,28 +2,37 @@ import { createPinia, setActivePinia } from "pinia";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
 const executeMulti = vi.fn();
+const executeQuery = vi.fn();
 const analyzeEditableQueryEditability = vi.fn();
 const getColumns = vi.fn();
 const listIndexes = vi.fn();
 const getConnectionConfig = vi.fn();
 const buildSortedQuerySql = vi.fn();
+const buildDataGridCountSql = vi.fn();
+const prepareQueryPaginationExecutionPlan = vi.fn(async (options) => ({
+  sqlToExecute: options.sql,
+  pageSql: undefined,
+  pageLimit: undefined,
+  pageOffset: undefined,
+  countSql: undefined,
+  useAgentResultSession: false,
+}));
+const editorSettings = {
+  pageSize: 100,
+  autoCalculateTotalRows: false,
+};
 
 vi.mock("@/lib/backend/api", () => ({
   analyzeEditableQueryEditability,
+  buildDataGridCountSql,
   buildSortedQuerySql,
   closeClientConnectionSession: vi.fn().mockResolvedValue(undefined),
   closeQuerySession: vi.fn().mockResolvedValue(undefined),
   executeMulti,
+  executeQuery,
   getColumns,
   listIndexes,
-  prepareQueryPaginationExecutionPlan: vi.fn(async (options) => ({
-    sqlToExecute: options.sql,
-    pageSql: undefined,
-    pageLimit: undefined,
-    pageOffset: undefined,
-    countSql: undefined,
-    useAgentResultSession: false,
-  })),
+  prepareQueryPaginationExecutionPlan,
   saveOpenTabsState: vi.fn().mockResolvedValue(undefined),
 }));
 
@@ -37,7 +46,7 @@ vi.mock("@/stores/connectionStore", () => ({
 
 vi.mock("@/stores/settingsStore", () => ({
   useSettingsStore: () => ({
-    editorSettings: { pageSize: 100 },
+    editorSettings,
   }),
 }));
 
@@ -68,6 +77,23 @@ describe("queryStore hidden primary key editing", () => {
     listIndexes.mockResolvedValue([]);
     analyzeEditableQueryEditability.mockImplementation(async (sql: string) => queryAnalysis(sql));
     buildSortedQuerySql.mockImplementation(async (options) => ({ ok: true, sql: `${options.originalSql} ORDER BY ${options.column} ${options.direction.toUpperCase()}` }));
+    buildDataGridCountSql.mockResolvedValue("SELECT COUNT(*) FROM `users`");
+    prepareQueryPaginationExecutionPlan.mockImplementation(async (options) => ({
+      sqlToExecute: options.sql,
+      pageSql: undefined,
+      pageLimit: undefined,
+      pageOffset: undefined,
+      countSql: undefined,
+      useAgentResultSession: false,
+    }));
+    editorSettings.pageSize = 100;
+    editorSettings.autoCalculateTotalRows = false;
+    executeQuery.mockResolvedValue({
+      columns: ["row_count"],
+      rows: [[0]],
+      affected_rows: 0,
+      execution_time_ms: 1,
+    });
     executeMulti.mockResolvedValue([
       {
         columns: ["name", "__DBX_PK_0"],
@@ -339,4 +365,111 @@ describe("queryStore hidden primary key editing", () => {
     await vi.waitFor(() => expect(tab.queryEditabilityReason).toBe("primary-key-not-returned"));
     expect(tab.queryAnalysis).toBeUndefined();
   });
+
+  it("records the returned row count when a page is known to be incomplete without count sql", async () => {
+    prepareQueryPaginationExecutionPlan.mockResolvedValue({
+      sqlToExecute: "SELECT name FROM users LIMIT 100 OFFSET 0",
+      pageSql: "SELECT name FROM users LIMIT 100 OFFSET 0",
+      pageLimit: 100,
+      pageOffset: 0,
+      countSql: undefined,
+      useAgentResultSession: false,
+    });
+    executeMulti.mockResolvedValue([
+      {
+        columns: ["name"],
+        rows: Array.from({ length: 42 }, (_, index) => [`user-${index}`]),
+        affected_rows: 0,
+        execution_time_ms: 1,
+      },
+    ]);
+
+    const { useQueryStore } = await import("@/stores/queryStore");
+    const store = useQueryStore();
+    const tabId = store.createTab("mysql-1", "app", "Query");
+
+    await store.executeTabSql(tabId, "SELECT name FROM users");
+
+    const tab = store.tabs.find((item) => item.id === tabId)!;
+    expect(tab.resultTotalRowCount).toBe(42);
+    expect(tab.resultTotalRowCountLoading).toBe(false);
+    expect(executeQuery).not.toHaveBeenCalled();
+  });
+
+  it("does not treat an empty later page as the total row count", async () => {
+    prepareQueryPaginationExecutionPlan.mockResolvedValue({
+      sqlToExecute: "SELECT name FROM users LIMIT 100 OFFSET 200",
+      pageSql: "SELECT name FROM users LIMIT 100 OFFSET 200",
+      pageLimit: 100,
+      pageOffset: 200,
+      countSql: undefined,
+      useAgentResultSession: false,
+    });
+    executeMulti.mockResolvedValue([
+      {
+        columns: ["name"],
+        rows: [],
+        affected_rows: 0,
+        execution_time_ms: 1,
+      },
+    ]);
+
+    const { useQueryStore } = await import("@/stores/queryStore");
+    const store = useQueryStore();
+    const tabId = store.createTab("mysql-1", "app", "Query");
+
+    await store.executeTabSql(tabId, "SELECT name FROM users");
+
+    const tab = store.tabs.find((item) => item.id === tabId)!;
+    expect(tab.resultTotalRowCount).toBeUndefined();
+    expect(tab.resultTotalRowCountLoading).toBe(false);
+    expect(executeQuery).not.toHaveBeenCalled();
+  });
+
+  it("automatically counts table data totals when the setting is enabled", async () => {
+    editorSettings.autoCalculateTotalRows = true;
+    executeMulti.mockResolvedValue([
+      {
+        columns: ["id", "name"],
+        rows: Array.from({ length: 100 }, (_, index) => [index + 1, `user-${index + 1}`]),
+        affected_rows: 0,
+        execution_time_ms: 1,
+      },
+    ]);
+    executeQuery.mockResolvedValue({
+      columns: ["row_count"],
+      rows: [[123]],
+      affected_rows: 0,
+      execution_time_ms: 1,
+    });
+
+    const { useQueryStore } = await import("@/stores/queryStore");
+    const store = useQueryStore();
+    const tabId = store.createTab("mysql-1", "app", "users", "data", "public");
+    store.setTableMeta(tabId, {
+      schema: "public",
+      tableName: "users",
+      columns: [
+        { name: "id", data_type: "int", is_nullable: false, is_primary_key: true, column_default: null, extra: null },
+        { name: "name", data_type: "varchar", is_nullable: true, is_primary_key: false, column_default: null, extra: null },
+      ],
+      primaryKeys: ["id"],
+    });
+
+    await store.executeTabSql(tabId, "SELECT id, name FROM users LIMIT 100", {
+      pagination: { limit: 100, offset: 0 },
+    });
+
+    expect(buildDataGridCountSql).toHaveBeenCalledWith({
+      databaseType: "mysql",
+      catalog: undefined,
+      schema: "public",
+      tableName: "users",
+      whereInput: undefined,
+    });
+    await vi.waitFor(() => expect(executeQuery).toHaveBeenCalledWith("mysql-1", "app", "SELECT COUNT(*) FROM `users`", undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 })));
+    const tab = store.tabs.find((item) => item.id === tabId)!;
+    await vi.waitFor(() => expect(tab.resultTotalRowCount).toBe(123));
+    expect(tab.resultTotalRowCountLoading).toBe(false);
+  });
 });

apps/desktop/src/stores/queryStore.ts

@@ -2339,32 +2339,61 @@ export const useQueryStore = defineStore("query", () => {
 
   function setQueryTotalRowCountIfCurrent(tabId: string, executionId: string, result: QueryResult, totalRowCount: number | undefined) {
     const current = tabs.value.find((t) => t.id === tabId);
-    if (current?.mode !== "query") return;
+    if (!current || (current.mode !== "query" && current.mode !== "data")) return;
     if (current.executionId !== executionId && current.result !== result) return;
     current.resultTotalRowCount = totalRowCount;
     current.resultTotalRowCountLoading = false;
     syncActiveResultRunFromDisplayed(current);
   }
 
-  function countQueryTotalRowsInBackground(options: { tabId: string; connectionId: string; database: string; schema?: string; countSql?: string; result: QueryResult; pageLimit?: number; pageOffset?: number; executionId: string; traceId: string; elapsed: () => string; timeoutSecs: number }) {
+  type TotalRowCountSqlTarget = { sql: string; schema?: string };
+
+  function countQueryTotalRowsInBackground(options: {
+    tabId: string;
+    connectionId: string;
+    database: string;
+    schema?: string;
+    countSql?: string;
+    countSqlTarget?: () => Promise<TotalRowCountSqlTarget | undefined>;
+    result: QueryResult;
+    pageLimit?: number;
+    pageOffset?: number;
+    executionId: string;
+    traceId: string;
+    elapsed: () => string;
+    timeoutSecs: number;
+  }) {
     const resultRowCount = options.result.rows.length;
-    if (!options.countSql || resultRowCount <= 0) {
+    if (resultRowCount <= 0) {
       setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined);
       return;
     }
-    const countSql = options.countSql;
-    const clientSessionId = tabClientSessionId({ id: options.tabId }, "count");
-    const countExecutionId = `${options.executionId}:count`;
-
     if (typeof options.pageLimit === "number" && resultRowCount < options.pageLimit) {
       setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, (options.pageOffset ?? 0) + resultRowCount);
       return;
     }
 
+    // A full page was returned, so more rows may exist and determining the true
+    // total requires a potentially expensive COUNT(*) over the user's query.
+    // Only run it automatically when the user opted in; otherwise leave the
+    // total unknown and let them trigger it on demand from the result grid
+    // (matches DBeaver's default of not counting large result sets).
+    if (!useSettingsStore().editorSettings.autoCalculateTotalRows) {
+      setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined);
+      return;
+    }
+
+    const clientSessionId = tabClientSessionId({ id: options.tabId }, "count");
+    const countExecutionId = `${options.executionId}:count`;
     void (async () => {
       try {
+        const countTarget = options.countSql ? { sql: options.countSql, schema: options.schema } : await options.countSqlTarget?.();
+        if (!countTarget?.sql) {
+          setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined);
+          return;
+        }
         console.info("[DBX][executeTabSql:count:start]", { traceId: options.traceId, elapsed: options.elapsed() });
-        const countResult = await api.executeQuery(options.connectionId, options.database, countSql, options.schema, countExecutionId, {
+        const countResult = await api.executeQuery(options.connectionId, options.database, countTarget.sql, countTarget.schema, countExecutionId, {
           clientSessionId,
           timeoutSecs: options.timeoutSecs,
         });
@@ -2911,24 +2940,49 @@ export const useQueryStore = defineStore("query", () => {
         if (!options?.preserveTotalRowCountDuringExecution) {
           current.resultTotalRowCount = undefined;
         }
-        current.resultTotalRowCountLoading = current.mode === "query" && !!current.result && !!countSql;
+        const resultRowCount = current.result?.rows.length ?? 0;
+        const totalKnownFromIncompletePage = !!current.result && typeof pageLimit === "number" && resultRowCount < pageLimit;
+        const dataCountTarget =
+          current.mode === "data"
+            ? (() => {
+                const tableMeta = tableMetaForDataTab(current);
+                if (!tableMeta?.tableName) return undefined;
+                return {
+                  databaseType: effectiveDbType,
+                  catalog: tableMeta.catalog,
+                  schema: tableMeta.schema,
+                  tableName: tableMeta.tableName,
+                  whereInput: current.whereInput?.trim() || undefined,
+                };
+              })()
+            : undefined;
+        const canAutoCalculateTotalRows = !!current.result && resultRowCount > 0 && !totalKnownFromIncompletePage && settingsStore.editorSettings.autoCalculateTotalRows && ((current.mode === "query" && !!countSql) || (current.mode === "data" && !!dataCountTarget));
+        current.resultTotalRowCountLoading = canAutoCalculateTotalRows;
         // Server-side pagination without a countSql: the backend (currently
         // the Elasticsearch driver) already reports the true match total via
         // affected_rows. Use it directly so the result-grid can compute the
         // page count without issuing a separate COUNT query.
-        if (current.result && current.mode === "query" && typeof pageLimit === "number" && !countSql && typeof current.result.affected_rows === "number") {
+        let totalRowCountResolved = false;
+        if (current.result && current.mode === "query" && typeof pageLimit === "number" && !countSql && typeof current.result.affected_rows === "number" && current.result.affected_rows > current.result.rows.length) {
           current.resultTotalRowCount = current.result.affected_rows;
           current.resultTotalRowCountLoading = false;
+          totalRowCountResolved = true;
         }
         touchResult(current);
         syncDisplayedResultRun(current, queryBaseSql);
-        if (current.mode === "query" && current.result) {
+        if (!totalRowCountResolved && (current.mode === "query" || current.mode === "data") && current.result) {
           countQueryTotalRowsInBackground({
             tabId: id,
             connectionId: current.connectionId,
             database: current.database,
             schema: current.schema,
             countSql,
+            countSqlTarget: dataCountTarget
+              ? async () => ({
+                  sql: await api.buildDataGridCountSql(dataCountTarget),
+                  schema: undefined,
+                })
+              : undefined,
             result: current.result,
             pageLimit,
             pageOffset,

apps/desktop/src/stores/settingsStore.ts

@@ -388,6 +388,7 @@ export interface EditorSettings {
   pageSize: number;
   infiniteScroll: boolean;
   infiniteScrollMaxRows: number;
+  autoCalculateTotalRows: boolean;
   mongoViewMode: "document" | "table";
   showColumnCommentsInHeader: boolean;
   showColumnTypesInHeader: boolean;
@@ -521,6 +522,7 @@ export const DEFAULT_EDITOR_SETTINGS: EditorSettings = {
   pageSize: 100,
   infiniteScroll: false,
   infiniteScrollMaxRows: 5000,
+  autoCalculateTotalRows: false,
   mongoViewMode: "document",
   showColumnCommentsInHeader: true,
   showColumnTypesInHeader: true,
@@ -754,6 +756,7 @@ export function normalizeEditorSettings(settings: Partial<EditorSettings>, exist
     pageSize: normalizeResultPageSize(settings.pageSize),
     infiniteScroll: settings.infiniteScroll ?? DEFAULT_EDITOR_SETTINGS.infiniteScroll,
     infiniteScrollMaxRows: typeof settings.infiniteScrollMaxRows === "number" && settings.infiniteScrollMaxRows >= 1000 && settings.infiniteScrollMaxRows <= 50000 ? Math.round(settings.infiniteScrollMaxRows) : DEFAULT_EDITOR_SETTINGS.infiniteScrollMaxRows,
+    autoCalculateTotalRows: settings.autoCalculateTotalRows ?? DEFAULT_EDITOR_SETTINGS.autoCalculateTotalRows,
     mongoViewMode: settings.mongoViewMode === "table" ? "table" : DEFAULT_EDITOR_SETTINGS.mongoViewMode,
     showColumnCommentsInHeader: settings.showColumnCommentsInHeader ?? DEFAULT_EDITOR_SETTINGS.showColumnCommentsInHeader,
     showColumnTypesInHeader: settings.showColumnTypesInHeader ?? DEFAULT_EDITOR_SETTINGS.showColumnTypesInHeader,
@@ -999,6 +1002,7 @@ export const useSettingsStore = defineStore("settings", () => {
     if (partial.infiniteScroll !== undefined) editorSettings.value.infiniteScroll = partial.infiniteScroll;
     if (partial.infiniteScrollMaxRows !== undefined)
       editorSettings.value.infiniteScrollMaxRows = typeof partial.infiniteScrollMaxRows === "number" && partial.infiniteScrollMaxRows >= 1000 && partial.infiniteScrollMaxRows <= 50000 ? Math.round(partial.infiniteScrollMaxRows) : DEFAULT_EDITOR_SETTINGS.infiniteScrollMaxRows;
+    if (partial.autoCalculateTotalRows !== undefined) editorSettings.value.autoCalculateTotalRows = partial.autoCalculateTotalRows === true;
     if (partial.mongoViewMode !== undefined) editorSettings.value.mongoViewMode = partial.mongoViewMode;
     if (partial.showColumnCommentsInHeader !== undefined) editorSettings.value.showColumnCommentsInHeader = partial.showColumnCommentsInHeader;
     if (partial.showColumnTypesInHeader !== undefined) editorSettings.value.showColumnTypesInHeader = partial.showColumnTypesInHeader;