fix(mongodb): support find explain commands

Closes #5619
This commit is contained in:
t8y2 committed 2026-08-07 18:51:04 +08:00
1 parent bac8f4166f
commit 996ce42e80
11 files changed
+544 -2

No files matched your search

+12 -1
View File
@@ -1273,6 +1273,7 @@ pub enum MongoAgentMethod {
ListCollections,
FindDocuments,
FindOne,
ExplainFind,
FindDocumentsExtendedJson,
CountDocuments,
ServerVersion,
@@ -1288,11 +1289,12 @@ pub enum MongoAgentMethod {
}
impl MongoAgentMethod {
pub const ALL: [Self; 16] = [
pub const ALL: [Self; 17] = [
Self::ListDatabases,
Self::ListCollections,
Self::FindDocuments,
Self::FindOne,
Self::ExplainFind,
Self::FindDocumentsExtendedJson,
Self::CountDocuments,
Self::ServerVersion,
@@ -1313,6 +1315,7 @@ impl MongoAgentMethod {
Self::ListCollections => "list_collections",
Self::FindDocuments => "find_documents",
Self::FindOne => "find_one",
Self::ExplainFind => "explain_find",
Self::FindDocumentsExtendedJson => "find_documents_extended_json",
Self::CountDocuments => "count_documents",
Self::ServerVersion => "server_version",
@@ -2471,6 +2474,13 @@ impl AgentDriverClient {
self.call_mongo_method(MongoAgentMethod::FindOne, params).await
}
pub async fn mongo_explain_find<T: DeserializeOwned + Send + 'static>(
&mut self,
params: Value,
) -> Result<T, String> {
self.call_mongo_method(MongoAgentMethod::ExplainFind, params).await
}
/// Calls the Mongo agent read method that returns MongoDB relaxed Extended JSON.
pub async fn mongo_find_documents_extended_json<T: DeserializeOwned + Send + 'static>(
&mut self,
@@ -4322,6 +4332,7 @@ for line in sys.stdin:
assert_eq!(MongoAgentMethod::ListCollections.as_str(), "list_collections");
assert_eq!(MongoAgentMethod::FindDocuments.as_str(), "find_documents");
assert_eq!(MongoAgentMethod::FindOne.as_str(), "find_one");
assert_eq!(MongoAgentMethod::ExplainFind.as_str(), "explain_find");
assert_eq!(MongoAgentMethod::FindDocumentsExtendedJson.as_str(), "find_documents_extended_json");
assert_eq!(MongoAgentMethod::CountDocuments.as_str(), "count_documents");
assert_eq!(MongoAgentMethod::ServerVersion.as_str(), "server_version");
+95
View File
@@ -732,6 +732,70 @@ pub async fn find_documents_without_total(
.await
}
#[allow(clippy::too_many_arguments)]
pub async fn explain_find(
client: &Client,
database: &str,
collection: &str,
skip: u64,
limit: i64,
filter: Option<&str>,
projection: Option<&str>,
sort: Option<&str>,
collation: Option<&str>,
verbosity: &str,
) -> Result<serde_json::Value, String> {
let command = build_find_explain_command(collection, skip, limit, filter, projection, sort, collation, verbosity)?;
let result = client.database(database).run_command(command).await.map_err(|error| error.to_string())?;
Ok(bson_to_json(&Bson::Document(result)))
}
#[allow(clippy::too_many_arguments)]
fn build_find_explain_command(
collection: &str,
skip: u64,
limit: i64,
filter: Option<&str>,
projection: Option<&str>,
sort: Option<&str>,
collation: Option<&str>,
verbosity: &str,
) -> Result<Document, String> {
let mut find = doc! {
"find": collection,
"filter": parse_optional_filter_document(filter)?.unwrap_or_default(),
};
if let Some(projection) = parse_optional_json_document(projection, "projection")? {
find.insert("projection", projection);
}
if let Some(sort) = parse_optional_json_document(sort, "sort")? {
find.insert("sort", sort);
}
if let Some(collation) = parse_find_collation(collation)? {
find.insert(
"collation",
mongodb::bson::to_document(&collation).map_err(|error| format!("Invalid collation: {error}"))?,
);
}
if skip > 0 {
find.insert("skip", i64::try_from(skip).map_err(|_| "MongoDB skip exceeds the supported range")?);
}
if limit > 0 {
find.insert("limit", limit);
}
Ok(doc! {
"explain": find,
"verbosity": validate_find_explain_verbosity(verbosity)?,
})
}
fn validate_find_explain_verbosity(verbosity: &str) -> Result<&str, String> {
match verbosity {
"queryPlanner" | "executionStats" | "allPlansExecution" => Ok(verbosity),
_ => Err("MongoDB explain verbosity must be queryPlanner, executionStats, or allPlansExecution.".to_string()),
}
}
#[allow(clippy::too_many_arguments)]
async fn find_documents_with_total(
client: &Client,
@@ -2240,6 +2304,37 @@ mod tests {
.contains("Invalid collation"));
}
#[test]
fn builds_find_explain_command_with_all_find_options() {
let command = build_find_explain_command(
"im_msg",
2,
5,
Some(r#"{"active":true}"#),
Some(r#"{"email":1}"#),
Some(r#"{"email":1}"#),
Some(r#"{"locale":"en","strength":1}"#),
"executionStats",
)
.unwrap();
let find = command.get_document("explain").unwrap();
assert_eq!(find.get_str("find").unwrap(), "im_msg");
assert_eq!(find.get_document("filter").unwrap().get_bool("active").unwrap(), true);
assert_eq!(find.get_document("projection").unwrap().get_i64("email").unwrap(), 1);
assert_eq!(find.get_document("sort").unwrap().get_i64("email").unwrap(), 1);
assert_eq!(find.get_i64("skip").unwrap(), 2);
assert_eq!(find.get_i64("limit").unwrap(), 5);
assert_eq!(find.get_document("collation").unwrap().get_str("locale").unwrap(), "en");
assert_eq!(command.get_str("verbosity").unwrap(), "executionStats");
}
#[test]
fn rejects_invalid_find_explain_verbosity() {
let error = build_find_explain_command("items", 0, 0, None, None, None, None, "invalid").unwrap_err();
assert!(error.contains("queryPlanner, executionStats, or allPlansExecution"));
}
#[test]
fn mongo_find_count_failure_returns_loaded_lower_bound() {
let error = "invalid type: floating point `2053278871.0`, expected u64".to_string();
+61
View File
@@ -232,6 +232,49 @@ pub async fn mongo_find_one_core(
}
}
#[allow(clippy::too_many_arguments)]
pub async fn mongo_explain_find_core(
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>,
verbosity: &str,
) -> Result<serde_json::Value, String> {
ensure_document_pool(state, connection_id).await?;
let connections = state.connections.read().await;
match connections.get(connection_id).ok_or("Not found")? {
PoolKind::MongoDb(client) => {
mongo_driver::explain_find(
client, database, collection, skip, limit, filter, projection, sort, collation, verbosity,
)
.await
}
PoolKind::Agent(client) => {
let mut client = client.lock().await;
client
.mongo_explain_find(serde_json::json!({
"database": database,
"collection": collection,
"skip": skip,
"limit": limit,
"filter": filter,
"projection": projection,
"sort": sort,
"collation": collation,
"verbosity": verbosity,
}))
.await
}
_ => Err("Not a MongoDB connection".to_string()),
}
}
pub async fn mongo_count_documents_core(
state: &AppState,
connection_id: &str,
@@ -712,6 +755,24 @@ pub async fn execute_mongo_command_core(
.await?;
Ok(mongo_documents_query_result(result.documents))
}
MongoCommand::FindExplain { collection, filter, projection, sort, collation, skip, limit, verbosity } => {
let limit = bounded_mongo_find_limit(*limit, max_rows);
let plan = mongo_explain_find_core(
state,
connection_id,
database,
collection,
*skip,
limit,
Some(filter),
projection.as_deref(),
sort.as_deref(),
collation.as_deref(),
verbosity,
)
.await?;
Ok(mongo_documents_query_result(vec![plan]))
}
MongoCommand::FindOne { collection, filter, projection, options } => {
let result = mongo_find_one_core(
state,
+78 -1
View File
@@ -18,6 +18,17 @@ pub enum MongoCommand {
skip: u64,
limit: i64,
},
#[serde(rename = "findExplain")]
FindExplain {
collection: String,
filter: String,
projection: Option<String>,
sort: Option<String>,
collation: Option<String>,
skip: u64,
limit: i64,
verbosity: String,
},
#[serde(rename = "findOne")]
FindOne { collection: String, filter: String, projection: Option<String>, options: Option<String> },
#[serde(rename = "countDocuments")]
@@ -350,7 +361,9 @@ pub fn parse(input: &str) -> Result<MongoCommand, String> {
let mut collation = None;
let mut skip = 0;
let mut limit = 100;
for (name, call_args) in chained_calls(&tail)? {
let calls = chained_calls(&tail)?;
let call_count = calls.len();
for (index, (name, call_args)) in calls.into_iter().enumerate() {
match name.as_str() {
"sort" => sort = Some(normalized_json(call_args.first().map(String::as_str).unwrap_or("{}"))?),
"collation" => {
@@ -364,6 +377,21 @@ pub fn parse(input: &str) -> Result<MongoCommand, String> {
"count" if call_args.is_empty() => {
return Ok(MongoCommand::Count { collection, filter, accurate: false });
}
"explain" => {
if index + 1 != call_count {
return Err("MongoDB explain() must be the final find() chain operation.".to_string());
}
return Ok(MongoCommand::FindExplain {
collection,
filter,
projection,
sort,
collation,
skip,
limit,
verbosity: parse_explain_verbosity(&call_args)?,
});
}
_ => return Err(format!("Unsupported MongoDB find() chain: {name}()")),
}
}
@@ -656,6 +684,20 @@ fn parse_string_arg(arg: &str) -> Result<String, String> {
value.as_str().map(ToOwned::to_owned).ok_or_else(|| "MongoDB argument must be a string.".to_string())
}
fn parse_explain_verbosity(args: &[String]) -> Result<String, String> {
if args.len() > 1 {
return Err("MongoDB explain() accepts at most one verbosity string.".to_string());
}
let verbosity = match args.first() {
Some(value) => parse_string_arg(value)?,
None => "queryPlanner".to_string(),
};
match verbosity.as_str() {
"queryPlanner" | "executionStats" | "allPlansExecution" => Ok(verbosity),
_ => Err("MongoDB explain() verbosity must be queryPlanner, executionStats, or allPlansExecution.".to_string()),
}
}
fn normalized_json(input: &str) -> Result<String, String> {
let transformed = transform_shell_constructors(input.trim())?;
let value: Value =
@@ -860,6 +902,41 @@ mod tests {
);
}
#[test]
fn parses_find_explain_with_query_options_and_verbosity() {
let command = parse(
r#"db.im_msg.find({active: true}, {email: 1}).sort({email: 1}).collation({locale: "en", strength: 1}).skip(2).limit(5).explain("executionStats")"#,
)
.unwrap();
assert_eq!(
serde_json::to_value(command).unwrap(),
serde_json::json!({
"kind": "findExplain",
"collection": "im_msg",
"filter": "{\"active\":true}",
"projection": "{\"email\":1}",
"sort": "{\"email\":1}",
"collation": "{\"locale\":\"en\",\"strength\":1}",
"skip": 2,
"limit": 5,
"verbosity": "executionStats"
})
);
}
#[test]
fn find_explain_defaults_and_validates_verbosity() {
let default = serde_json::to_value(parse("db.items.find({}).explain()").unwrap()).unwrap();
assert_eq!(default["verbosity"], "queryPlanner");
let all_plans = serde_json::to_value(parse("db.items.find({}).explain('allPlansExecution')").unwrap()).unwrap();
assert_eq!(all_plans["verbosity"], "allPlansExecution");
assert!(parse("db.items.find({}).explain('invalid')").unwrap_err().contains("verbosity"));
assert!(parse("db.items.find({}).explain('executionStats').limit(1)").unwrap_err().contains("final"));
}
#[test]
fn parses_get_collection_and_count() {
assert_eq!(