diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 4dcaaf4b0b..170d744358 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -308,6 +308,7 @@ export interface QueryPaginationExecutionPlan { pageLimit?: number; pageOffset?: number; countSql?: string; + exactQueryRowBound?: number; useAgentResultSession: boolean; } diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 14c9678c01..e8f5f0b008 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -3915,6 +3915,7 @@ export const useQueryStore = defineStore("query", () => { let pageLimit: number | undefined; let pageOffset: number | undefined; let countSql: string | undefined; + let exactQueryRowBound: number | undefined; let useAgentResultSession = false; let executionDispatched = false; let producedResult = false; @@ -4493,6 +4494,7 @@ export const useQueryStore = defineStore("query", () => { pageLimit = plan.pageLimit; pageOffset = plan.pageOffset; countSql = plan.countSql; + exactQueryRowBound = plan.exactQueryRowBound; useAgentResultSession = plan.useAgentResultSession; const hasBoundedPagination = typeof pageLimit === "number" && typeof pageOffset === "number"; if (options?.appendResult && !hasBoundedPagination && !useAgentResultSession) { @@ -4689,7 +4691,16 @@ export const useQueryStore = defineStore("query", () => { current.result.truncated = true; current.resultTotalRowCount = queryResultMaxRows; } - const totalKnownFromIncompletePage = !!current.result && typeof exactTotalFromIncompletePage(current.result, pageLimit, pageOffset, useAgentResultSession) === "number"; + const paginationPageResult = shouldAppendResult ? results[0] : current.result; + const exactIncompletePageTotal = paginationPageResult ? exactTotalFromIncompletePage(paginationPageResult, pageLimit, pageOffset, useAgentResultSession) : undefined; + const totalKnownFromIncompletePage = typeof exactIncompletePageTotal === "number"; + let totalRowCountResolved = false; + if (current.mode === "query" && current.result && !isQueryExecutionErrorResult(current.result) && typeof exactQueryRowBound === "number") { + const boundedTotal = capQueryResultTotal(exactQueryRowBound, queryResultMaxRows); + current.resultTotalRowCount = Math.min(boundedTotal, exactIncompletePageTotal ?? boundedTotal); + current.resultTotalRowCountLoading = false; + totalRowCountResolved = true; + } const dataCountTarget = current.mode === "data" ? (() => { @@ -4707,14 +4718,20 @@ export const useQueryStore = defineStore("query", () => { })() : undefined; const canAutoCalculateTotalRows = - !options?.appendResult && !!current.result && resultRowCount > 0 && !resultLimitReached && !totalKnownFromIncompletePage && settingsStore.editorSettings.autoCalculateTotalRows && ((current.mode === "query" && !!countSql) || (current.mode === "data" && !!dataCountTarget)); + !options?.appendResult && + !!current.result && + resultRowCount > 0 && + !resultLimitReached && + !totalKnownFromIncompletePage && + !totalRowCountResolved && + 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. - let totalRowCountResolved = false; - if (current.result && current.result.total_is_exact !== false && current.mode === "query" && typeof pageLimit === "number" && !countSql && typeof current.result.affected_rows === "number" && current.result.affected_rows > current.result.rows.length) { + if (!totalRowCountResolved && current.result && current.result.total_is_exact !== false && 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; diff --git a/crates/dbx-core/src/query_result_sql.rs b/crates/dbx-core/src/query_result_sql.rs index aea3a09c91..f0847b7e10 100644 --- a/crates/dbx-core/src/query_result_sql.rs +++ b/crates/dbx-core/src/query_result_sql.rs @@ -56,6 +56,8 @@ pub struct QueryPaginationExecutionPlan { pub page_offset: Option, #[serde(skip_serializing_if = "Option::is_none")] pub count_sql: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub exact_query_row_bound: Option, pub use_agent_result_session: bool, /// True when the statement cannot be paginated server-side and must be /// executed once with the whole result streamed back (single execution). @@ -115,12 +117,18 @@ pub struct SortedQuerySqlOptions { pub fn build_query_pagination_execution_plan( options: QueryPaginationExecutionPlanOptions, ) -> QueryPaginationExecutionPlan { + let exact_query_row_bound = if options.database_type == Some(DatabaseType::SqlServer) { + top_level_top_row_count(&options.query_base_sql) + } else { + None + }; let mut plan = QueryPaginationExecutionPlan { sql_to_execute: options.sql.clone(), page_sql: None, page_limit: None, page_offset: None, count_sql: None, + exact_query_row_bound, use_agent_result_session: false, single_execution: false, }; @@ -3315,6 +3323,39 @@ WHERE u.id = picked.id; assert_eq!(top_level_top_row_count("SELECT * FROM (SELECT TOP 5 * FROM events) t"), None); } + #[test] + fn sql_server_pagination_plan_exposes_only_safe_exact_top_bounds() { + for (database_type, sql, expected) in [ + (Some(DatabaseType::SqlServer), "SELECT TOP 100 * FROM events", Some(100)), + (Some(DatabaseType::SqlServer), "SELECT TOP(50) * FROM events", Some(50)), + (Some(DatabaseType::SqlServer), "SELECT TOP (250) * FROM events", Some(250)), + (Some(DatabaseType::SqlServer), "SELECT TOP 10 PERCENT * FROM events", None), + (Some(DatabaseType::SqlServer), "SELECT TOP (2) WITH TIES * FROM events", None), + (Some(DatabaseType::SqlServer), "SELECT TOP (100 + 1) * FROM events", None), + (Some(DatabaseType::SqlServer), "SELECT TOP (@row_count) * FROM events", None), + (Some(DatabaseType::SqlServer), "SELECT * FROM (SELECT TOP 5 * FROM events) t", None), + (Some(DatabaseType::SqlServer), "SELECT * FROM events", None), + (Some(DatabaseType::Postgres), "SELECT TOP 100 * FROM events", None), + (Some(DatabaseType::Kingbase), "SELECT TOP 100 * FROM events", None), + ] { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: sql.to_string(), + query_base_sql: sql.to_string(), + database_type, + pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, + use_agent_cursor: false, + first_page_uses_actual_sql: false, + }); + + assert_eq!(plan.exact_query_row_bound, expected, "unexpected bound for {database_type:?}: {sql}"); + let serialized = serde_json::to_value(&plan).unwrap(); + assert_eq!( + serialized.get("exactQueryRowBound").and_then(serde_json::Value::as_u64), + expected.map(|value| value as u64) + ); + } + } + #[test] fn non_kingbase_limit_offset_is_unaffected_by_top_keyword() { // A column literally named `top` must not be mistaken for a TOP clause diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index a97c327d45..0f1f9dbd16 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -6619,6 +6619,83 @@ test("query execution keeps automatically counting total rows in the background" } }); +for (const scenario of [ + { label: "TOP equal to one page", bound: 100, pageLimit: 100, pageOffset: 0, rowCount: 100, expectedTotal: 100 }, + { label: "TOP below one page", bound: 50, pageLimit: 100, pageOffset: 0, rowCount: 50, expectedTotal: 50 }, + { label: "TOP above the first page", bound: 250, pageLimit: 100, pageOffset: 0, rowCount: 100, expectedTotal: 250 }, + { label: "TOP above an intermediate page", bound: 250, pageLimit: 100, pageOffset: 100, rowCount: 100, expectedTotal: 250 }, + { label: "TOP reached on the final page", bound: 250, pageLimit: 100, pageOffset: 200, rowCount: 50, expectedTotal: 250 }, + { label: "TOP above the actual matching rows", bound: 500, pageLimit: 100, pageOffset: 200, rowCount: 37, expectedTotal: 237 }, + { label: "TOP above the enabled result cap", bound: 200_000, pageLimit: 100, pageOffset: 0, rowCount: 100, expectedTotal: 180, queryResultMaxRows: 180 }, + { label: "TOP overrides an unrelated exact backend total", bound: 100, pageLimit: 100, pageOffset: 0, rowCount: 100, expectedTotal: 100, affectedRows: 500 }, +]) { + test(`SQL Server ${scenario.label} resolves an exact bounded total without COUNT`, async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + let countRequests = 0; + + settingsStore.updateEditorSettings({ + autoCalculateTotalRows: true, + queryResultMaxRowsEnabled: true, + queryResultMaxRows: scenario.queryResultMaxRows ?? 100_000, + }); + connectionStore.addEphemeralConnection(sqlServerConn("sqlserver-1")); + const tabId = store.createTab("sqlserver-1", "db", "Query", "query", "dbo"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + + globalThis.fetch = withConnectionHealthMock(async (input) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + return Response.json({ + sqlToExecute: `SELECT TOP ${scenario.bound} 1 AS [id]`, + pageSql: `SELECT TOP ${scenario.bound} 1 AS [id]`, + pageLimit: scenario.pageLimit, + pageOffset: scenario.pageOffset, + countSql: "SELECT COUNT(*) FROM (SELECT 1 AS [id]) [dbx_count]", + exactQueryRowBound: scenario.bound, + useAgentResultSession: false, + }); + } + if (url === "/api/query/execute-multi") { + return Response.json([ + { + columns: ["id"], + rows: Array.from({ length: scenario.rowCount }, (_, index) => [scenario.pageOffset + index + 1]), + affected_rows: scenario.affectedRows ?? 0, + execution_time_ms: 1, + }, + ]); + } + if (url === "/api/query/execute") { + countRequests += 1; + return new Response("unexpected count", { status: 500 }); + } + if (url === "/api/query/analyze-editability") { + return Response.json({ editable: false, reason: "complex-source" }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + await store.executeTabSql(tabId, `SELECT TOP ${scenario.bound} 1 AS [id]`, { + pagination: { limit: scenario.pageLimit, offset: scenario.pageOffset }, + }); + + assert.equal(tab.resultTotalRowCount, scenario.expectedTotal); + assert.equal(tab.resultTotalRowCountLoading, false); + assert.equal(countRequests, 0); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } + }); +} + test("inexact backend totals do not become query pagination bounds", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia());