diff --git a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java index 08f61c694c..e537c1d063 100644 --- a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java +++ b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java @@ -15,6 +15,12 @@ import com.mongodb.ServerAddress; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoClients; import com.mongodb.client.MongoCollection; +import com.mongodb.client.model.Collation; +import com.mongodb.client.model.CollationAlternate; +import com.mongodb.client.model.CollationCaseFirst; +import com.mongodb.client.model.CollationMaxVariable; +import com.mongodb.client.model.CollationStrength; +import com.mongodb.client.model.CountOptions; import com.mongodb.client.model.UpdateOptions; import com.mongodb.client.result.UpdateResult; import java.io.BufferedReader; @@ -423,12 +429,13 @@ public final class MongoAgent { Document filterDoc = documentOrNull(params, "filter"); Document projectionDoc = documentOrNull(params, "projection"); Document sortDoc = documentOrNull(params, "sort"); + Collation collation = collationOrNull(documentOrNull(params, "collation")); var col = c.getDatabase(database).getCollection(collection); if (filterDoc == null) { filterDoc = new Document(); } - CollectionTotal total = collectionTotal(col, filterDoc); + CollectionTotal total = collectionTotal(col, filterDoc, collation); var iterable = col.find(filterDoc).skip((int) skip).limit(limit); if (projectionDoc != null) { @@ -437,6 +444,9 @@ public final class MongoAgent { if (sortDoc != null) { iterable = iterable.sort(sortDoc); } + if (collation != null) { + iterable = iterable.collation(collation); + } List> documents = new ArrayList<>(); for (Document document : iterable) { @@ -458,12 +468,13 @@ public final class MongoAgent { Document filterDoc = documentOrNull(params, "filter"); Document projectionDoc = documentOrNull(params, "projection"); Document sortDoc = documentOrNull(params, "sort"); + Collation collation = collationOrNull(documentOrNull(params, "collation")); var col = c.getDatabase(database).getCollection(collection); if (filterDoc == null) { filterDoc = new Document(); } - CollectionTotal total = collectionTotal(col, filterDoc); + CollectionTotal total = collectionTotal(col, filterDoc, collation); var iterable = col.find(filterDoc).skip((int) skip).limit(limit); if (projectionDoc != null) { @@ -472,6 +483,9 @@ public final class MongoAgent { if (sortDoc != null) { iterable = iterable.sort(sortDoc); } + if (collation != null) { + iterable = iterable.collation(collation); + } List documents = new ArrayList<>(); for (Document document : iterable) { @@ -480,11 +494,89 @@ public final class MongoAgent { return documentQueryResult(documents, total); } + static Collation collationOrNull(Document document) { + if (document == null) { + return null; + } + Set supported = Set.of( + "locale", "strength", "caseLevel", "caseFirst", "numericOrdering", + "alternate", "maxVariable", "normalization", "backwards" + ); + for (String key : document.keySet()) { + if (!supported.contains(key)) { + throw new IllegalArgumentException("Unsupported collation option: " + key); + } + } + String locale = document.getString("locale"); + if (locale == null || locale.isBlank()) { + throw new IllegalArgumentException("Invalid collation: locale must not be empty"); + } + + Collation.Builder builder = Collation.builder().locale(locale); + if (document.containsKey("strength")) { + Object strength = document.get("strength"); + if (!(strength instanceof Number number) || number.doubleValue() != Math.rint(number.doubleValue())) { + throw new IllegalArgumentException("Invalid collation option strength: expected an integer from 1 to 5"); + } + int strengthValue = number.intValue(); + if (strengthValue < 1 || strengthValue > 5) { + throw new IllegalArgumentException("Invalid collation option strength: expected an integer from 1 to 5"); + } + builder.collationStrength(CollationStrength.fromInt(strengthValue)); + } + if (document.containsKey("caseLevel")) { + builder.caseLevel(collationBoolean(document, "caseLevel")); + } + if (document.containsKey("caseFirst")) { + builder.collationCaseFirst(CollationCaseFirst.fromString(collationString(document, "caseFirst"))); + } + if (document.containsKey("numericOrdering")) { + builder.numericOrdering(collationBoolean(document, "numericOrdering")); + } + if (document.containsKey("alternate")) { + builder.collationAlternate(CollationAlternate.fromString(collationString(document, "alternate"))); + } + if (document.containsKey("maxVariable")) { + builder.collationMaxVariable(CollationMaxVariable.fromString(collationString(document, "maxVariable"))); + } + if (document.containsKey("normalization")) { + builder.normalization(collationBoolean(document, "normalization")); + } + if (document.containsKey("backwards")) { + builder.backwards(collationBoolean(document, "backwards")); + } + return builder.build(); + } + + private static boolean collationBoolean(Document document, String key) { + Object value = document.get(key); + if (value instanceof Boolean booleanValue) { + return booleanValue; + } + throw new IllegalArgumentException("Invalid collation option " + key + ": expected a boolean"); + } + + private static String collationString(Document document, String key) { + Object value = document.get(key); + if (value instanceof String stringValue) { + return stringValue; + } + throw new IllegalArgumentException("Invalid collation option " + key + ": expected a string"); + } + static CollectionTotal collectionTotal(MongoCollection collection, Document filter) { + return collectionTotal(collection, filter, null); + } + + static CollectionTotal collectionTotal(MongoCollection collection, Document filter, Collation collation) { if (filter.isEmpty()) { return new CollectionTotal(collection.estimatedDocumentCount(), false); } - return new CollectionTotal(collection.countDocuments(filter), true); + CountOptions options = new CountOptions(); + if (collation != null) { + options.collation(collation); + } + return new CollectionTotal(collection.countDocuments(filter, options), true); } static Map documentQueryResult(List documents, CollectionTotal total) { diff --git a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java index 0f78d5a885..9602bd12a5 100644 --- a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java +++ b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java @@ -16,6 +16,9 @@ import com.mongodb.MongoClientSettings; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.Collation; +import com.mongodb.client.model.CollationStrength; +import com.mongodb.client.model.CountOptions; import com.mongodb.client.model.UpdateOptions; import com.mongodb.client.result.UpdateResult; import java.io.FileInputStream; @@ -193,6 +196,45 @@ class MongoAgentTest { assertEquals(List.of("countDocuments:{\"status\": \"active\"}"), calls); } + @Test + void parsesFindCollationAndUsesItForExactCounts() { + Collation collation = MongoAgent.collationOrNull(Document.parse( + "{\"locale\":\"en\",\"strength\":1,\"caseLevel\":false,\"numericOrdering\":true}" + )); + assertNotNull(collation); + assertEquals("en", collation.getLocale()); + assertEquals(CollationStrength.PRIMARY, collation.getStrength()); + assertEquals(false, collation.getCaseLevel()); + assertEquals(true, collation.getNumericOrdering()); + + List calls = new ArrayList<>(); + MongoCollection collection = recordingCountCollection(calls); + MongoAgent.CollectionTotal total = MongoAgent.collectionTotal( + collection, + new Document("name", "xxx"), + collation + ); + + assertEquals(42L, total.value()); + assertEquals(List.of("countDocuments:{\"name\": \"xxx\"}:collation=en/1"), calls); + } + + @Test + void rejectsInvalidFindCollationOptions() { + assertThrows( + IllegalArgumentException.class, + () -> MongoAgent.collationOrNull(Document.parse("{\"strength\":1}")) + ); + assertThrows( + IllegalArgumentException.class, + () -> MongoAgent.collationOrNull(Document.parse("{\"locale\":\"en\",\"unknown\":true}")) + ); + assertThrows( + IllegalArgumentException.class, + () -> MongoAgent.collationOrNull(Document.parse("{\"locale\":\"en\",\"strength\":1.5}")) + ); + } + @Test void estimatedDocumentQueryResultMarksTotalAsInexact() { Map result = MongoAgent.documentQueryResult( @@ -822,7 +864,12 @@ class MongoAgentTest { } if ("countDocuments".equals(method.getName())) { Document filter = (Document) args[0]; - calls.add("countDocuments:" + filter.toJson()); + String call = "countDocuments:" + filter.toJson(); + if (args.length > 1 && args[1] instanceof CountOptions options && options.getCollation() != null) { + Collation collation = options.getCollation(); + call += ":collation=" + collation.getLocale() + "/" + collation.getStrength().getIntRepresentation(); + } + calls.add(call); return 42L; } throw new UnsupportedOperationException(method.getName()); diff --git a/apps/desktop/src/components/document/DocumentBrowser.vue b/apps/desktop/src/components/document/DocumentBrowser.vue index 6ba9e80eb2..e1e716aceb 100644 --- a/apps/desktop/src/components/document/DocumentBrowser.vue +++ b/apps/desktop/src/components/document/DocumentBrowser.vue @@ -904,7 +904,7 @@ async function load(options: { page?: number } = {}) { } const sort = currentDocumentSortJson(sortInput.value); const skip = requestPage * pageSize.value; - const result = await api.documentFindDocuments(connectionId, database, collection, skip, documentRequestLimit.value, filter, undefined, sort, executionId); + const result = await api.documentFindDocuments(connectionId, database, collection, skip, documentRequestLimit.value, filter, undefined, sort, undefined, executionId); if (documentLoadExecutionId.value !== executionId) return; if (connectionId !== props.connectionId || database !== props.database || collection !== props.collection || storeKind !== documentStoreProvider.value.kind) return; const nextDocuments = diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index b0e454ed8a..ad0837511c 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -2937,8 +2937,8 @@ export async function vectorGetCollectionDetail(connectionId: string, database: }); } -export async function mongoFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, executionId?: string): Promise { - return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, 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): Promise { + return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, collation, executionId); } export async function mongoParseShellCommand(source: string): Promise { @@ -2958,7 +2958,7 @@ export async function mongoFindOne(connectionId: string, database: string, colle }); } -export async function documentFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, executionId?: string): Promise { +export async function documentFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, collation?: string, executionId?: string): Promise { return post("/api/document-store/find-documents", { connectionId, database, @@ -2968,6 +2968,7 @@ export async function documentFindDocuments(connectionId: string, database: stri filter, projection, sort, + collation, executionId, }); } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 86546fc0ba..c7e400c488 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -2895,8 +2895,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, executionId?: string): Promise { - return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, 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): Promise { + return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, collation, executionId); } export async function mongoFindOne(connectionId: string, database: string, collection: string, filter?: string, projection?: string, options?: string, executionId?: string): Promise { @@ -2916,7 +2916,7 @@ export async function mongoParseShellCommand(source: string): Promise { +export async function documentFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, collation?: string, executionId?: string): Promise { return invoke("document_find_documents", { connectionId, database, @@ -2926,6 +2926,7 @@ export async function documentFindDocuments(connectionId: string, database: stri filter, projection, sort, + collation, executionId, }); } diff --git a/apps/desktop/src/lib/mongo/mongoShellCommand.ts b/apps/desktop/src/lib/mongo/mongoShellCommand.ts index 4598b9eaef..08a5757201 100644 --- a/apps/desktop/src/lib/mongo/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongo/mongoShellCommand.ts @@ -24,6 +24,7 @@ export interface MongoFindCommand { skip: number; limit: number; sort?: string; + collation?: string; } export interface MongoFindPaginationPlan { @@ -162,6 +163,14 @@ export function parseMongoFindCommand(input: string): MongoFindCommand | null { sort = parsedSort; } + const collationArg = readChainedCallArgument(chain, "collation"); + let collation: string | undefined; + if (collationArg !== undefined) { + const parsedCollation = normalizeJsonArgument(collationArg); + if (!parsedCollation) return null; + collation = parsedCollation; + } + const skip = readChainedIntegerArgument(chain, "skip", 0); const limit = readChainedIntegerArgument(chain, "limit", DEFAULT_LIMIT); if (skip === null || limit === null) return null; @@ -173,6 +182,7 @@ export function parseMongoFindCommand(input: string): MongoFindCommand | null { skip, limit, sort, + ...(collation ? { collation } : {}), }; } diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 9ffbed95bd..df875e34c6 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -3536,7 +3536,7 @@ export const useQueryStore = defineStore("query", () => { if (!pagePlan) throw new Error(describeMongoCommandParseFailure(sourceStatement)); // A stale request can point past an explicit .limit() bound. Keep // the backend call bounded so limit(0) cannot become unbounded. - const result = await api.mongoFindDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, pagePlan.requestSkip, Math.max(1, pagePlan.requestLimit), mongoCommand.filter, mongoCommand.projection, mongoCommand.sort, executionId); + const result = await api.mongoFindDocuments(tab.connectionId, currentDatabase, mongoCommand.collection, pagePlan.requestSkip, Math.max(1, pagePlan.requestLimit), mongoCommand.filter, mongoCommand.projection, mongoCommand.sort, mongoCommand.collation, executionId); const documents = pagePlan.requestLimit === 0 ? [] : result.documents; const extendedDocuments = pagePlan.requestLimit === 0 ? [] : result.extended_documents; const totalIsExact = result.total_is_exact !== false; @@ -5046,7 +5046,7 @@ export const useQueryStore = defineStore("query", () => { if (!plan) throw new Error(QUERY_RESULT_EXPORT_UNSUPPORTED_ERROR); if (plan.requestLimit === 0) break; - const result = await api.mongoFindDocuments(tab.connectionId, tab.database, mongoCommand.collection, plan.requestSkip, plan.requestLimit, mongoCommand.filter, mongoCommand.projection, mongoCommand.sort, exportExecutionId); + const result = await api.mongoFindDocuments(tab.connectionId, tab.database, mongoCommand.collection, plan.requestSkip, plan.requestLimit, mongoCommand.filter, mongoCommand.projection, mongoCommand.sort, mongoCommand.collation, exportExecutionId); const pageDocuments = result.documents.slice(0, plan.requestLimit); documents.push(...pageDocuments); diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index 0566534423..780668a46c 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -1,6 +1,6 @@ use mongodb::{ bson::{doc, oid::ObjectId, Bson, DateTime, Document}, - options::{ClientOptions, GridFsBucketOptions, IndexOptions, UpdateModifications}, + options::{ClientOptions, Collation, GridFsBucketOptions, IndexOptions, UpdateModifications}, Client, Cursor, Database, IndexModel, }; use serde::{Deserialize, Serialize}; @@ -16,6 +16,18 @@ pub use super::document_result::DocumentQueryResult; /// Backward-compatible name for callers of Mongo-specific APIs. pub type MongoDocumentResult = DocumentQueryResult; +const MONGO_COLLATION_FIELDS: &[&str] = &[ + "locale", + "strength", + "caseLevel", + "caseFirst", + "numericOrdering", + "alternate", + "maxVariable", + "normalization", + "backwards", +]; + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MongoDropIndexesResult { pub dropped_names: Vec, @@ -636,6 +648,23 @@ fn index_info_from_model(model: IndexModel) -> IndexInfo { } } +fn parse_find_collation(value: Option<&str>) -> Result, String> { + let Some(raw) = value.map(str::trim).filter(|value| !value.is_empty()) else { + return Ok(None); + }; + let json: serde_json::Value = + serde_json::from_str(raw).map_err(|error| format!("Invalid collation JSON: {error}"))?; + let object = json.as_object().ok_or_else(|| "Invalid collation: expected an options object".to_string())?; + if let Some(field) = object.keys().find(|field| !MONGO_COLLATION_FIELDS.contains(&field.as_str())) { + return Err(format!("Unsupported collation option: {field}")); + } + let collation: Collation = serde_json::from_value(json).map_err(|error| format!("Invalid collation: {error}"))?; + if collation.locale.trim().is_empty() { + return Err("Invalid collation: locale must not be empty".to_string()); + } + Ok(Some(collation)) +} + pub async fn find_documents( client: &Client, database: &str, @@ -645,6 +674,7 @@ pub async fn find_documents( filter: Option<&str>, projection: Option<&str>, sort: Option<&str>, + collation: Option<&str>, ) -> Result { let col = client.database(database).collection::(collection); @@ -656,9 +686,14 @@ pub async fn find_documents( _ => doc! {}, }; + let collation = parse_find_collation(collation)?; let count_is_exact = !filter_doc.is_empty(); let total_result = if count_is_exact { - col.count_documents(filter_doc.clone()).await.map_err(|e| e.to_string()) + 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()) } else { col.estimated_document_count().await.map_err(|e| e.to_string()) }; @@ -679,6 +714,9 @@ pub async fn find_documents( find = find.sort(sort_doc); } } + if let Some(collation) = collation { + find = find.collation(collation); + } let mut cursor = find.await.map_err(|e| e.to_string())?; @@ -791,6 +829,7 @@ pub async fn find_documents_extended_json( filter: Option<&str>, projection: Option<&str>, sort: Option<&str>, + collation: Option<&str>, ) -> Result { let col = client.database(database).collection::(collection); @@ -802,9 +841,14 @@ pub async fn find_documents_extended_json( _ => doc! {}, }; + let collation = parse_find_collation(collation)?; let count_is_exact = !filter_doc.is_empty(); let total_result = if count_is_exact { - col.count_documents(filter_doc.clone()).await.map_err(|e| e.to_string()) + 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()) } else { col.estimated_document_count().await.map_err(|e| e.to_string()) }; @@ -825,6 +869,9 @@ pub async fn find_documents_extended_json( find = find.sort(sort_doc); } } + if let Some(collation) = collation { + find = find.collation(collation); + } let mut cursor = find.await.map_err(|e| e.to_string())?; @@ -2088,6 +2135,37 @@ fn expand_object_id_string_array(items: &[serde_json::Value]) -> Bson { mod tests { use super::*; + #[test] + fn parses_find_collation_options() { + let collation = parse_find_collation(Some( + r#"{"locale":"en","strength":1,"caseLevel":true,"caseFirst":"upper","numericOrdering":true,"alternate":"shifted","maxVariable":"space","normalization":true,"backwards":false}"#, + )) + .unwrap() + .unwrap(); + + assert_eq!(collation.locale, "en"); + assert!(matches!(collation.strength, Some(mongodb::options::CollationStrength::Primary))); + assert_eq!(collation.case_level, Some(true)); + assert!(matches!(collation.case_first, Some(mongodb::options::CollationCaseFirst::Upper))); + assert_eq!(collation.numeric_ordering, Some(true)); + assert!(matches!(collation.alternate, Some(mongodb::options::CollationAlternate::Shifted))); + assert!(matches!(collation.max_variable, Some(mongodb::options::CollationMaxVariable::Space))); + assert_eq!(collation.normalization, Some(true)); + assert_eq!(collation.backwards, Some(false)); + } + + #[test] + fn rejects_invalid_find_collation_options() { + assert!(parse_find_collation(Some(r#"{"strength":1}"#)).unwrap_err().contains("locale")); + assert!(parse_find_collation(Some(r#"{"locale":""}"#)).unwrap_err().contains("locale")); + assert!(parse_find_collation(Some(r#"{"locale":"en","unknown":true}"#)) + .unwrap_err() + .contains("Unsupported collation option")); + assert!(parse_find_collation(Some(r#"{"locale":"en","strength":"primary"}"#)) + .unwrap_err() + .contains("Invalid collation")); + } + #[test] fn mongo_find_count_failure_returns_loaded_lower_bound() { let error = "invalid type: floating point `2053278871.0`, expected u64".to_string(); diff --git a/crates/dbx-core/src/document_ops.rs b/crates/dbx-core/src/document_ops.rs index c573b2625e..d6aa446697 100644 --- a/crates/dbx-core/src/document_ops.rs +++ b/crates/dbx-core/src/document_ops.rs @@ -402,6 +402,7 @@ pub async fn find_documents_core( filter: Option<&str>, projection: Option<&str>, sort: Option<&str>, + collation: Option<&str>, ) -> Result { ensure_document_pool(state, connection_id).await?; let connections = state.connections.read().await; @@ -410,7 +411,7 @@ pub async fn find_documents_core( // 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, + client, database, collection, skip, limit, filter, projection, sort, collation, ) .await } @@ -443,6 +444,9 @@ pub async fn find_documents_core( if let Some(projection) = projection { params["projection"] = serde_json::json!(projection); } + if let Some(collation) = collation { + params["collation"] = serde_json::json!(collation); + } match client.mongo_find_documents_extended_json(params.clone()).await { Ok(result) => Ok(result), Err(error) if is_unknown_agent_method_error(&error, "find_documents_extended_json") => { diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 313927ef34..9f34994ef0 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -138,6 +138,7 @@ pub async fn mongo_find_documents_core( filter: Option<&str>, projection: Option<&str>, sort: Option<&str>, + collation: Option<&str>, ) -> Result { crate::document_ops::find_documents_core( state, @@ -149,6 +150,7 @@ pub async fn mongo_find_documents_core( filter, projection, sort, + collation, ) .await } @@ -236,7 +238,7 @@ pub async fn mongo_find_documents_extended_json_core( match connections.get(connection_id).ok_or("Not found")? { PoolKind::MongoDb(client) => { mongo_driver::find_documents_extended_json( - client, database, collection, skip, limit, filter, projection, sort, + client, database, collection, skip, limit, filter, projection, sort, None, ) .await } diff --git a/crates/dbx-core/src/mongo_shell.rs b/crates/dbx-core/src/mongo_shell.rs index 48c51803b2..5cd5ff1cc9 100644 --- a/crates/dbx-core/src/mongo_shell.rs +++ b/crates/dbx-core/src/mongo_shell.rs @@ -9,7 +9,15 @@ pub enum MongoCommand { #[serde(rename = "use")] Use { database: String }, #[serde(rename = "find")] - Find { collection: String, filter: String, projection: Option, sort: Option, skip: u64, limit: i64 }, + Find { + collection: String, + filter: String, + projection: Option, + sort: Option, + collation: Option, + skip: u64, + limit: i64, + }, #[serde(rename = "findOne")] FindOne { collection: String, filter: String, projection: Option, options: Option }, #[serde(rename = "countDocuments")] @@ -339,11 +347,18 @@ pub fn parse(input: &str) -> Result { return Err("MongoDB find() accepts at most filter and projection arguments.".to_string()); } let mut sort = None; + let mut collation = None; let mut skip = 0; let mut limit = 100; for (name, call_args) in chained_calls(&tail)? { match name.as_str() { "sort" => sort = Some(normalized_json(call_args.first().map(String::as_str).unwrap_or("{}"))?), + "collation" => { + if call_args.len() != 1 { + return Err("MongoDB collation() requires one options object.".to_string()); + } + collation = Some(normalized_json(&call_args[0])?); + } "skip" => skip = parse_integer(&call_args, "skip")? as u64, "limit" => limit = parse_integer(&call_args, "limit")?, "count" if call_args.is_empty() => { @@ -352,7 +367,7 @@ pub fn parse(input: &str) -> Result { _ => return Err(format!("Unsupported MongoDB find() chain: {name}()")), } } - return Ok(MongoCommand::Find { collection, filter, projection, sort, skip, limit }); + return Ok(MongoCommand::Find { collection, filter, projection, sort, collation, skip, limit }); } if let Some((args, tail)) = method_call(source, prefix_end, "findOne") { @@ -822,12 +837,29 @@ mod tests { filter: r#"{"_id":{"$oid":"507f1f77bcf86cd799439011"}}"#.to_string(), projection: Some(r#"{"title":1,"_id":0}"#.to_string()), sort: Some(r#"{"title":1}"#.to_string()), + collation: None, skip: 0, limit: 1, } ); } + #[test] + fn parses_find_with_collation_chain() { + assert_eq!( + parse(r#"db.t_user.find({name: 'xxx'}).collation({ locale: "en", strength: 1 }).limit(20)"#).unwrap(), + MongoCommand::Find { + collection: "t_user".to_string(), + filter: r#"{"name":"xxx"}"#.to_string(), + projection: None, + sort: None, + collation: Some(r#"{"locale":"en","strength":1}"#.to_string()), + skip: 0, + limit: 20, + } + ); + } + #[test] fn parses_get_collection_and_count() { assert_eq!( diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index 56e8057da5..d516c676db 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -3716,6 +3716,7 @@ async fn find_mongo_documents_for_rows( None, None, Some(r#"{"_id":1}"#), + None, ) .await } diff --git a/crates/dbx-core/tests/live_mongodb_find_one.rs b/crates/dbx-core/tests/live_mongodb_find_one.rs index dd9bb704a4..9ff3a945e5 100644 --- a/crates/dbx-core/tests/live_mongodb_find_one.rs +++ b/crates/dbx-core/tests/live_mongodb_find_one.rs @@ -54,10 +54,19 @@ async fn find_documents_returns_type_preserving_copy_documents() { .await .unwrap(); - let result = - mongo_driver::find_documents(&client, database, &collection, 0, 10, Some("{}"), Some(r#"{"_id":0}"#), None) - .await - .unwrap(); + let result = mongo_driver::find_documents( + &client, + database, + &collection, + 0, + 10, + Some("{}"), + Some(r#"{"_id":0}"#), + None, + None, + ) + .await + .unwrap(); assert_eq!(result.documents[0]["lastUpdatedDate"], serde_json::json!("ISODate(\"2025-05-06T08:35:32Z\")")); let extended = result.extended_documents.expect("extended documents"); @@ -66,3 +75,41 @@ async fn find_documents_returns_type_preserving_copy_documents() { mongo_driver::drop_collection(&client, database, &collection).await.unwrap(); } + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MONGODB_URL pointing at a writable MongoDB database"] +async fn find_documents_applies_collation_to_results_total_and_pagination() { + let url = std::env::var("DBX_LIVE_MONGODB_URL").expect("DBX_LIVE_MONGODB_URL"); + let client = mongo_driver::connect(&url, Duration::from_secs(10), Duration::from_secs(60)).await.unwrap(); + let database = std::env::var("DBX_LIVE_MONGODB_DATABASE").unwrap_or_else(|_| "dbx_live_find_one".to_string()); + let collection = format!("collation_{}", std::process::id()); + + mongo_driver::insert_documents( + &client, + &database, + &collection, + r#"[{"name":"xxx","rank":1},{"name":"XXX","rank":2},{"name":"yyy","rank":3}]"#, + ) + .await + .unwrap(); + + let result = mongo_driver::find_documents_extended_json( + &client, + &database, + &collection, + 1, + 1, + Some(r#"{"name":"xxx"}"#), + Some(r#"{"_id":0}"#), + Some(r#"{"rank":1}"#), + Some(r#"{"locale":"en","strength":1}"#), + ) + .await + .unwrap(); + + assert_eq!(result.total, 2); + assert!(result.total_is_exact); + assert_eq!(result.documents, vec![serde_json::json!({ "name": "XXX", "rank": 2 })]); + + mongo_driver::drop_collection(&client, &database, &collection).await.unwrap(); +} diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index b982f6a603..9e5fe8b2f0 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -440,7 +440,7 @@ impl DbxBackend for LocalBackend { .await .map(|version| scalar_query_result("version", Value::String(version))), MongoCommand::Use { database } => Ok(scalar_query_result("database", Value::String(database.clone()))), - MongoCommand::Find { collection, filter, projection, sort, skip, limit } => { + MongoCommand::Find { collection, filter, projection, sort, collation, skip, limit } => { let result = mongo_ops::mongo_find_documents_core( &self.state, connection_id, @@ -451,6 +451,7 @@ impl DbxBackend for LocalBackend { Some(filter), projection.as_deref(), sort.as_deref(), + collation.as_deref(), ) .await?; Ok(mongo_documents_query_result(result.documents)) @@ -973,7 +974,7 @@ impl DbxBackend for WebBackend { Ok(scalar_query_result("version", Value::String(version))) } MongoCommand::Use { database } => Ok(scalar_query_result("database", Value::String(database.clone()))), - MongoCommand::Find { collection, filter, projection, sort, skip, limit } => { + MongoCommand::Find { collection, filter, projection, sort, collation, skip, limit } => { let result = self .request( reqwest::Method::POST, @@ -987,6 +988,7 @@ impl DbxBackend for WebBackend { "filter": filter, "projection": projection, "sort": sort, + "collation": collation, })), ) .await? diff --git a/crates/dbx-web/src/routes/document_store.rs b/crates/dbx-web/src/routes/document_store.rs index 71bba40535..96309958ee 100644 --- a/crates/dbx-web/src/routes/document_store.rs +++ b/crates/dbx-web/src/routes/document_store.rs @@ -66,6 +66,7 @@ pub struct DocumentFindRequest { pub filter: Option, pub projection: Option, pub sort: Option, + pub collation: Option, pub execution_id: Option, } @@ -191,6 +192,7 @@ pub async fn find_documents( req.filter.as_deref(), req.projection.as_deref(), req.sort.as_deref(), + req.collation.as_deref(), ), ) .await?; diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index 8ed509b212..87431dddca 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -97,6 +97,7 @@ pub struct MongoFindRequest { pub filter: Option, pub projection: Option, pub sort: Option, + pub collation: Option, pub execution_id: Option, } @@ -401,6 +402,7 @@ pub async fn find_documents( req.filter.as_deref(), req.projection.as_deref(), req.sort.as_deref(), + req.collation.as_deref(), ), ) .await?; diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index 82a081520b..752e15a072 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -59,6 +59,19 @@ test("parseMongoFindCommand parses getCollection find with chained sort skip and }); }); +test("parseMongoFindCommand preserves chained collation for execution and pagination", () => { + const command = parseMongoFindCommand(`db.t_user.find({name: 'xxx'}) + .collation({ locale: "en", strength: 1 }) + .limit(20)`); + + assert.ok(command); + assert.equal(command.collection, "t_user"); + assert.deepEqual(JSON.parse(command.filter), { name: "xxx" }); + assert.deepEqual(JSON.parse(command.collation ?? "null"), { locale: "en", strength: 1 }); + assert.equal(command.skip, 0); + assert.equal(command.limit, 20); +}); + test("planMongoFindPagination pages unbounded find queries", () => { const command = parseMongoFindCommand("db.users.find({})"); assert.ok(command); diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index 2358a19cb5..0a34cec697 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -3390,7 +3390,7 @@ test("MongoDB query result export pages find commands through the document API", const originalFetch = globalThis.fetch; const findBodies: any[] = []; const progress: Array<{ rowsExported: number; totalRows: number | null }> = []; - const command = 'db.permissions.find({"role":"admin"},{"name":1,"active":1}).sort({"createdTime":-1}).skip(3).limit(205)'; + const command = 'db.permissions.find({"role":"admin"},{"name":1,"active":1}).collation({locale:"en",strength:1}).sort({"createdTime":-1}).skip(3).limit(205)'; const documents = Array.from({ length: 205 }, (_, index) => (index === 100 ? { _id: index + 4, active: true } : { _id: index + 4, name: `user-${index + 4}` })); const copyDocuments = documents.map((document) => ({ ...document, _id: { $numberInt: String(document._id) } })); @@ -3439,7 +3439,7 @@ test("MongoDB query result export pages find commands through the document API", { skip: 203, limit: 5 }, ], ); - assert.ok(findBodies.every((body) => body.collection === "permissions" && body.filter === '{"role":"admin"}' && body.projection === '{"name":1,"active":1}' && body.sort === '{"createdTime":-1}')); + assert.ok(findBodies.every((body) => body.collection === "permissions" && body.filter === '{"role":"admin"}' && body.projection === '{"name":1,"active":1}' && body.sort === '{"createdTime":-1}' && JSON.stringify(JSON.parse(body.collation)) === JSON.stringify({ locale: "en", strength: 1 }))); assert.equal(new Set(findBodies.map((body) => body.executionId)).size, 1); assert.ok(findBodies[0]?.executionId); assert.deepEqual(exported?.columns, ["_id", "name", "active"]); @@ -3911,18 +3911,20 @@ test("mongo find execution uses editor page size and supports server pagination" try { const tabId = store.createTab("mongo-page-1", "dbx_test", "Query", "query", ""); - await store.executeTabSql(tabId, "db.issue_4566.find({})"); + const sql = 'db.issue_4566.find({name:"xxx"}).collation({locale:"en",strength:1})'; + await store.executeTabSql(tabId, sql); const tab = store.tabs.find((item) => item.id === tabId); assert.equal(findBodies[0]?.collection, "issue_4566"); assert.equal(findBodies[0]?.skip, 0); assert.equal(findBodies[0]?.limit, 100); + assert.deepEqual(JSON.parse(findBodies[0]?.collation), { locale: "en", strength: 1 }); assert.equal(tab?.result?.rows.length, 100); assert.equal(tab?.resultPageLimit, 100); assert.equal(tab?.resultPageOffset, 0); assert.equal(tab?.resultTotalRowCount, 824); - await store.executeTabSql(tabId, "db.issue_4566.find({})", { + await store.executeTabSql(tabId, sql, { pagination: { offset: 100, limit: 100 }, preserveResultDuringExecution: true, preserveTotalRowCountDuringExecution: true, @@ -3931,6 +3933,7 @@ test("mongo find execution uses editor page size and supports server pagination" assert.equal(findBodies[1]?.collection, "issue_4566"); assert.equal(findBodies[1]?.skip, 100); assert.equal(findBodies[1]?.limit, 100); + assert.deepEqual(JSON.parse(findBodies[1]?.collation), { locale: "en", strength: 1 }); assert.equal(tab?.result?.rows[0]?.[0], 101); assert.equal(tab?.resultPageOffset, 100); assert.equal(tab?.resultTotalRowCount, 824); diff --git a/src-tauri/src/commands/document_cmd.rs b/src-tauri/src/commands/document_cmd.rs index f0aa71a6eb..569a9ce27d 100644 --- a/src-tauri/src/commands/document_cmd.rs +++ b/src-tauri/src/commands/document_cmd.rs @@ -57,6 +57,7 @@ pub async fn document_find_documents( filter: Option, projection: Option, sort: Option, + collation: Option, execution_id: Option, ) -> Result { let app = state.inner().clone(); @@ -73,6 +74,7 @@ pub async fn document_find_documents( filter.as_deref(), projection.as_deref(), sort.as_deref(), + collation.as_deref(), ), ) .await diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index f42596ab75..f8556376c4 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -60,6 +60,7 @@ struct MongoFindDocumentsRequest { filter: Option, projection: Option, sort: Option, + collation: Option, } #[derive(Deserialize)] @@ -1181,6 +1182,7 @@ async fn handle_mongo_find_documents_data(state: &Arc, body: &str, str req.filter.as_deref(), req.projection.as_deref(), req.sort.as_deref(), + req.collation.as_deref(), ) .await { diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index 61e28a8f43..43e9d37491 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -110,6 +110,7 @@ pub async fn mongo_find_documents( filter: Option, projection: Option, sort: Option, + collation: Option, execution_id: Option, mcp_request: Option, ) -> Result { @@ -126,6 +127,7 @@ pub async fn mongo_find_documents( filter, projection, sort, + collation, execution_id, ) .await