diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 003bade88a..7602468ef4 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -4898,8 +4898,8 @@ export async function vectorRenameCollection(connectionId: string, database: str await post("/api/vector/rename-collection", { connectionId, database, collection, newName }); } -export async function mongoFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, collation?: string, executionId?: string): Promise { - return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, collation, executionId); +export async function mongoFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, collation?: string, executionId?: string, includeTotal?: boolean): Promise { + return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, collation, executionId, undefined, undefined, includeTotal); } export async function mongoParseShellCommand(source: string): Promise { @@ -4932,6 +4932,7 @@ export async function documentFindDocuments( executionId?: string, cursor?: string, cursorPagination?: boolean, + includeTotal?: boolean, ): Promise { return post("/api/document-store/find-documents", { connectionId, @@ -4945,6 +4946,7 @@ export async function documentFindDocuments( collation, cursor, cursorPagination, + includeTotal, executionId, }); } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index bd75c4b52f..fe782022da 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -4828,8 +4828,8 @@ export async function vectorListCollections(connectionId: string, database?: str return documentListCollections(connectionId, database || "default"); } -export async function mongoFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, collation?: string, executionId?: string): Promise { - return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, collation, executionId); +export async function mongoFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, collation?: string, executionId?: string, includeTotal?: boolean): Promise { + return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, collation, executionId, undefined, undefined, includeTotal); } export async function mongoFindOne(connectionId: string, database: string, collection: string, filter?: string, projection?: string, options?: string, executionId?: string): Promise { @@ -4862,6 +4862,7 @@ export async function documentFindDocuments( executionId?: string, cursor?: string, cursorPagination?: boolean, + includeTotal?: boolean, ): Promise { return invoke("document_find_documents", { connectionId, @@ -4875,6 +4876,7 @@ export async function documentFindDocuments( collation, cursor, cursorPagination, + includeTotal, executionId, }); } diff --git a/apps/desktop/src/lib/mongo/mongoShellCommand.ts b/apps/desktop/src/lib/mongo/mongoShellCommand.ts index dba4d47bbf..31447b5c24 100644 --- a/apps/desktop/src/lib/mongo/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongo/mongoShellCommand.ts @@ -459,6 +459,11 @@ export function planMongoFindPagination(input: string, command: MongoFindCommand }; } +/** True when the first page already requests every row an explicit `.limit()` allows. */ +export function mongoFindPageCoversLogicalLimit(plan: Pick): boolean { + return plan.pageOffset === 0 && plan.logicalLimit !== undefined && plan.requestLimit >= plan.logicalLimit; +} + export function mongoFindLogicalTotal(total: number, plan: Pick): number { const afterSkip = Math.max(0, Math.trunc(total) - plan.logicalSkip); return plan.logicalLimit === undefined ? afterSkip : Math.min(afterSkip, plan.logicalLimit); diff --git a/apps/desktop/src/stores/__tests__/queryStore.mongoExecutionSummary.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.mongoExecutionSummary.spec.ts index 05b5e8dad3..80696c366f 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.mongoExecutionSummary.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.mongoExecutionSummary.spec.ts @@ -91,6 +91,37 @@ describe("queryStore MongoDB execution summary", () => { expect(tab.batchSqlExecution?.items[0]?.status).not.toBe("skipped"); }); + it("skips the total count when an explicit limit fits in the first page", async () => { + mocks.mongoFindDocuments.mockResolvedValue({ + documents: [{ _id: "1" }, { _id: "2" }], + extended_documents: [], + total: 2, + total_is_exact: false, + }); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mongo-1", "app", "Query"); + + await store.executeTabSql(tabId, "db.events.find({ type: 0 }).sort({ eventTime: -1 }).limit(50)"); + + expect(mocks.mongoFindDocuments).toHaveBeenCalledWith("mongo-1", "app", "events", 0, 50, expect.any(String), undefined, expect.any(String), undefined, expect.any(String), false); + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.result?.rows).toHaveLength(2); + expect(tab.result?.truncated).toBe(false); + expect(tab.resultTotalRowCount).toBe(2); + }); + + it("keeps counting when an explicit limit spans more than one page", async () => { + mocks.mongoFindDocuments.mockResolvedValue({ documents: [{ _id: "1" }], extended_documents: [], total: 500, total_is_exact: true }); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mongo-1", "app", "Query"); + + await store.executeTabSql(tabId, "db.events.find({ type: 0 }).limit(150)"); + + expect(mocks.mongoFindDocuments).toHaveBeenCalledWith("mongo-1", "app", "events", 0, 100, expect.any(String), undefined, undefined, undefined, expect.any(String), true); + }); + it("marks a failing mongo command as error instead of leaving it stuck at 'skipped'", async () => { mocks.mongoFindDocuments.mockRejectedValue(new Error("collection not found")); const { useQueryStore } = await import("@/stores/queryStore"); @@ -134,8 +165,8 @@ describe("queryStore MongoDB execution summary", () => { await store.executeTabSql(tabId, 'db.getSiblingDB("reports").events.find();\ndb.users.find()'); - expect(mocks.mongoFindDocuments).toHaveBeenNthCalledWith(1, "mongo-1", "reports", "events", expect.any(Number), expect.any(Number), "{}", undefined, undefined, undefined, expect.any(String)); - expect(mocks.mongoFindDocuments).toHaveBeenNthCalledWith(2, "mongo-1", "app", "users", expect.any(Number), expect.any(Number), "{}", undefined, undefined, undefined, expect.any(String)); + expect(mocks.mongoFindDocuments).toHaveBeenNthCalledWith(1, "mongo-1", "reports", "events", expect.any(Number), expect.any(Number), "{}", undefined, undefined, undefined, expect.any(String), true); + expect(mocks.mongoFindDocuments).toHaveBeenNthCalledWith(2, "mongo-1", "app", "users", expect.any(Number), expect.any(Number), "{}", undefined, undefined, undefined, expect.any(String), true); const tab = store.tabs.find((item) => item.id === tabId)!; expect(tab.batchSqlExecution).toMatchObject({ total: 2, diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 01651d2e61..3987d01340 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -63,6 +63,7 @@ import { describeMongoCommandParseFailure, mongoDroppedIndexesToQueryResult, mongoFindLogicalTotal, + mongoFindPageCoversLogicalLimit, mongoIndexesToQueryResult, planMongoFindPagination, mongoUseToQueryResult, @@ -7704,6 +7705,10 @@ export const useQueryStore = defineStore("query", () => { const paginationSource = siblingPrefix ? `db${siblingPrefix.rest}` : sourceStatement; const pagePlan = planMongoFindPagination(paginationSource, mongoCommand, options?.pagination?.offset ?? 0, normalizeResultPageSize(options?.pagination?.limit ?? settingsStore.editorSettings.pageSize)); if (!pagePlan) throw new Error(describeMongoCommandParseFailure(sourceStatement)); + // An explicit .limit() that fits in the first page returns every + // logical row at once, so the page itself is the exact total and + // counting every match would only add latency (#10121). + const totalFromPage = mongoFindPageCoversLogicalLimit(pagePlan); // A stale request can point past an explicit .limit() bound. Keep // the backend call bounded so limit(0) cannot become unbounded. const result = await commandTiming.run( @@ -7718,10 +7723,11 @@ export const useQueryStore = defineStore("query", () => { mongoCommand.sort, mongoCommand.collation, executionId, + !totalFromPage, ); const documents = pagePlan.requestLimit === 0 ? [] : result.documents; const extendedDocuments = pagePlan.requestLimit === 0 ? [] : result.extended_documents; - const totalIsExact = result.total_is_exact !== false; + const totalIsExact = totalFromPage || result.total_is_exact !== false; const reportedTotal = mongoFindLogicalTotal(result.total, pagePlan); const loadedLowerBound = pagePlan.pageOffset + documents.length; const total = totalIsExact ? reportedTotal : Math.max(reportedTotal, loadedLowerBound); diff --git a/crates/dbx-core/src/query/document_ops.rs b/crates/dbx-core/src/query/document_ops.rs index 9c16e40b73..940aff5298 100644 --- a/crates/dbx-core/src/query/document_ops.rs +++ b/crates/dbx-core/src/query/document_ops.rs @@ -454,6 +454,42 @@ pub async fn find_documents_core( collation: Option<&str>, cursor: Option<&str>, cursor_pagination: bool, +) -> Result { + find_documents_core_with_total( + state, + connection_id, + database, + collection, + skip, + limit, + filter, + projection, + sort, + collation, + cursor, + cursor_pagination, + true, + ) + .await +} + +/// `include_total: false` lets MongoDB callers skip the total-count query when they can derive +/// the total from the returned page. Other document stores ignore it. +#[allow(clippy::too_many_arguments)] +pub async fn find_documents_core_with_total( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + skip: u64, + limit: i64, + filter: Option<&str>, + projection: Option<&str>, + sort: Option<&str>, + collation: Option<&str>, + cursor: Option<&str>, + cursor_pagination: bool, + include_total: bool, ) -> Result { ensure_document_pool(state, connection_id).await?; let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; @@ -461,10 +497,17 @@ pub async fn find_documents_core( PoolKind::MongoDb(client) => { // Document browser responses must retain BSON type metadata so nested filters // can round-trip ObjectId, Date, and int64 values through Extended JSON. - mongo_driver::find_documents_extended_json( - client, database, collection, skip, limit, filter, projection, sort, collation, - ) - .await + if include_total { + mongo_driver::find_documents_extended_json( + client, database, collection, skip, limit, filter, projection, sort, collation, + ) + .await + } else { + mongo_driver::find_documents_extended_json_without_total( + client, database, collection, skip, limit, filter, projection, sort, collation, + ) + .await + } } PoolKind::DynamoDb(client) => { let client = client.clone(); diff --git a/crates/dbx-core/tests/live_mongodb_count.rs b/crates/dbx-core/tests/live_mongodb_count.rs index f82ccb4d1a..dd0284a1b9 100644 --- a/crates/dbx-core/tests/live_mongodb_count.rs +++ b/crates/dbx-core/tests/live_mongodb_count.rs @@ -33,5 +33,23 @@ async fn count_supports_fast_accurate_and_browser_total_paths() { assert!(!browser_result.total_is_exact); assert_eq!(browser_result.documents.len(), 2); + // Without the total query the result only reports the rows it saw. + let filtered_page = mongo_driver::find_documents_extended_json_without_total( + &client, + &database, + &collection, + 1, + 5, + Some(r#"{"n":{"$gte":2}}"#), + None, + None, + None, + ) + .await + .unwrap(); + assert_eq!(filtered_page.documents.len(), 1); + assert_eq!(filtered_page.total, 2); + assert!(!filtered_page.total_is_exact); + mongo_driver::drop_collection(&client, &database, &collection).await.unwrap(); } diff --git a/crates/dbx-driver-mongodb/src/mongo_driver.rs b/crates/dbx-driver-mongodb/src/mongo_driver.rs index 99be475a60..ce81b7384c 100644 --- a/crates/dbx-driver-mongodb/src/mongo_driver.rs +++ b/crates/dbx-driver-mongodb/src/mongo_driver.rs @@ -1510,6 +1510,46 @@ pub async fn find_documents_extended_json( projection: Option<&str>, sort: Option<&str>, collation: Option<&str>, +) -> Result { + find_documents_extended_json_with_total( + client, database, collection, skip, limit, filter, projection, sort, collation, true, + ) + .await +} + +/// Like [`find_documents_extended_json`], but without the total-count query. Callers that can +/// derive the total from the returned page (an explicit `.limit()` that fits in one page) use +/// this to avoid counting every match. +#[allow(clippy::too_many_arguments)] +pub async fn find_documents_extended_json_without_total( + client: &Client, + database: &str, + collection: &str, + skip: u64, + limit: i64, + filter: Option<&str>, + projection: Option<&str>, + sort: Option<&str>, + collation: Option<&str>, +) -> Result { + find_documents_extended_json_with_total( + client, database, collection, skip, limit, filter, projection, sort, collation, false, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn find_documents_extended_json_with_total( + client: &Client, + database: &str, + collection: &str, + skip: u64, + limit: i64, + filter: Option<&str>, + projection: Option<&str>, + sort: Option<&str>, + collation: Option<&str>, + include_total: bool, ) -> Result { let col = client.database(database).collection::(collection); @@ -1526,14 +1566,16 @@ pub async fn find_documents_extended_json( let count_is_exact = !filter_doc.is_empty(); // Count alongside the find, as in `find_documents_with_total`. let total = async { - if count_is_exact { + if !include_total { + None + } else if count_is_exact { let mut count = col.count_documents(filter_doc.clone()); if let Some(collation) = collation.clone() { count = count.collation(collation); } - count.await.map_err(|e| e.to_string()) + Some(count.await.map_err(|e| e.to_string())) } else { - estimated_document_count(client, database, collection).await + Some(estimated_document_count(client, database, collection).await) } }; let page = async { @@ -1561,7 +1603,10 @@ pub async fn find_documents_extended_json( }; let (total_result, page) = tokio::join!(total, page); let (documents, extended_documents) = page?; - let (total, total_is_exact) = resolve_mongo_find_total(total_result, count_is_exact, skip, documents.len()); + let (total, total_is_exact) = match total_result { + Some(total_result) => resolve_mongo_find_total(total_result, count_is_exact, skip, documents.len()), + None => (skip.saturating_add(u64::try_from(documents.len()).unwrap_or(u64::MAX)), false), + }; Ok(MongoDocumentResult { extended_documents: Some(extended_documents), diff --git a/crates/dbx-web/src/routes/document_store.rs b/crates/dbx-web/src/routes/document_store.rs index b4ae28abed..afc1835901 100644 --- a/crates/dbx-web/src/routes/document_store.rs +++ b/crates/dbx-web/src/routes/document_store.rs @@ -70,6 +70,7 @@ pub struct DocumentFindRequest { pub cursor: Option, #[serde(default)] pub cursor_pagination: bool, + pub include_total: Option, pub execution_id: Option, } @@ -342,7 +343,7 @@ pub async fn find_documents( let result = run_cancellable( &state, req.execution_id, - dbx_core::document_ops::find_documents_core( + dbx_core::document_ops::find_documents_core_with_total( &state.app, &req.connection_id, &req.database, @@ -355,6 +356,7 @@ pub async fn find_documents( req.collation.as_deref(), req.cursor.as_deref(), req.cursor_pagination, + req.include_total.unwrap_or(true), ), ) .await?; diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index 8754113643..937ce2ebe6 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -33,6 +33,7 @@ import { parseMongoGetIndexesCommand, parseMongoVersionCommand, planMongoFindPagination, + mongoFindPageCoversLogicalLimit, parseMongoWriteCommand, splitMongoCommands, splitMongoCommandRanges, @@ -255,6 +256,23 @@ test("planMongoFindPagination preserves explicit skip and limit bounds", () => { assert.equal(mongoFindLogicalTotal(824, plan!), 150); }); +test("mongoFindPageCoversLogicalLimit only skips counting when the first page holds every limited row", () => { + const plan = (source: string, pageOffset: number, pageLimit = 100) => { + const command = parseMongoFindCommand(source); + assert.ok(command); + const result = planMongoFindPagination(source, command, pageOffset, pageLimit); + assert.ok(result); + return result; + }; + + assert.equal(mongoFindPageCoversLogicalLimit(plan("db.events.find({ type: 0 }).sort({ eventTime: -1 }).limit(50)", 0)), true); + assert.equal(mongoFindPageCoversLogicalLimit(plan("db.events.find({}).skip(20).limit(100)", 0)), true); + assert.equal(mongoFindPageCoversLogicalLimit(plan("db.events.find({}).limit(150)", 0)), false); + assert.equal(mongoFindPageCoversLogicalLimit(plan("db.events.find({}).limit(150)", 100)), false); + assert.equal(mongoFindPageCoversLogicalLimit(plan("db.events.find({})", 0)), false); + assert.equal(mongoFindPageCoversLogicalLimit(plan("db.events.find({}).limit(0)", 0)), false); +}); + test("parseMongoFindCommand accepts line breaks before find and chained calls", () => { const command = parseMongoFindCommand(`db.getCollection("accounting_reconciliations") .find({ diff --git a/src-tauri/src/commands/document_cmd.rs b/src-tauri/src/commands/document_cmd.rs index bff5f6c437..0160a586c4 100644 --- a/src-tauri/src/commands/document_cmd.rs +++ b/src-tauri/src/commands/document_cmd.rs @@ -60,13 +60,14 @@ pub async fn document_find_documents( collation: Option, cursor: Option, cursor_pagination: Option, + include_total: Option, execution_id: Option, ) -> Result { let app = state.inner().clone(); run_cancellable( &app, execution_id, - dbx_core::document_ops::find_documents_core( + dbx_core::document_ops::find_documents_core_with_total( &app, &connection_id, &database, @@ -79,6 +80,7 @@ pub async fn document_find_documents( collation.as_deref(), cursor.as_deref(), cursor_pagination.unwrap_or(false), + include_total.unwrap_or(true), ), ) .await diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index e87e68a991..4eb926ebe1 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -140,6 +140,7 @@ pub async fn mongo_find_documents( collation, None, None, + None, execution_id, ) .await?;