fix(mongodb): support show databases shell commands

Closes #6226

crates/dbx-core/src/mongo_ops.rs

@@ -8,6 +8,9 @@ use crate::document_ops::CollectionInfo;
 use crate::mongo_shell::MongoCommand;
 use crate::types::{IndexInfo, QueryResult};
 
+pub const MONGO_SHOW_DATABASES_DATABASE: &str = "admin";
+pub const MONGO_SHOW_DATABASES_COMMAND_JSON: &str = r#"{"listDatabases":1}"#;
+
 async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<(), String> {
     state.get_or_create_pool(connection_id, None).await.map(|_| ())
 }
@@ -177,6 +180,10 @@ pub async fn mongo_run_command_core(
     }
 }
 
+pub async fn mongo_show_databases_core(state: &AppState, connection_id: &str) -> Result<MongoDocumentResult, String> {
+    mongo_run_command_core(state, connection_id, MONGO_SHOW_DATABASES_DATABASE, MONGO_SHOW_DATABASES_COMMAND_JSON).await
+}
+
 pub async fn mongo_collection_stats_core(
     state: &AppState,
     connection_id: &str,
@@ -918,6 +925,10 @@ pub async fn execute_mongo_command_core(
             .await
             .map(|version| scalar_query_result("version", Value::String(version))),
         MongoCommand::Use { database } => Ok(scalar_query_result("database", Value::String(database.clone()))),
+        MongoCommand::ShowDatabases => {
+            let result = mongo_show_databases_core(state, connection_id).await?;
+            mongo_show_databases_query_result(result.documents, max_rows)
+        }
         MongoCommand::RunCommand { command_json } => {
             let result = mongo_run_command_core(state, connection_id, database, command_json).await?;
             Ok(mongo_documents_query_result(result.documents))
@@ -1233,6 +1244,41 @@ fn mongo_documents_query_result(documents: Vec<serde_json::Value>) -> QueryResul
     query_result(columns, rows, 0)
 }
 
+pub fn mongo_show_databases_query_result(
+    documents: Vec<serde_json::Value>,
+    max_rows: usize,
+) -> Result<QueryResult, String> {
+    use serde_json::Value;
+
+    let databases = documents
+        .first()
+        .and_then(Value::as_object)
+        .and_then(|response| response.get("databases"))
+        .and_then(Value::as_array)
+        .ok_or("MongoDB listDatabases response is missing the databases array")?;
+    if databases.iter().any(|database| !database.is_object()) {
+        return Err("MongoDB listDatabases response contains an invalid database entry".to_string());
+    }
+
+    let total = databases.len();
+    let rows = databases
+        .iter()
+        .take(max_rows.max(1))
+        .map(|database| {
+            let database = database.as_object().expect("database entries were validated above");
+            ["name", "sizeOnDisk", "empty"].map(|field| database.get(field).cloned().unwrap_or(Value::Null)).to_vec()
+        })
+        .collect::<Vec<_>>();
+    let mut query_result = query_result(
+        vec!["name".to_string(), "sizeOnDisk".to_string(), "empty".to_string()],
+        rows,
+        u64::try_from(total).unwrap_or(u64::MAX),
+    );
+    query_result.truncated = query_result.rows.len() < total;
+    query_result.has_more = query_result.truncated;
+    Ok(query_result)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1261,6 +1307,63 @@ mod tests {
         assert_eq!(limited.extended_documents.unwrap(), vec![serde_json::json!(1), serde_json::json!(2)]);
     }
 
+    #[test]
+    fn show_databases_result_preserves_metadata_and_row_limit() {
+        let result = MongoDocumentResult {
+            documents: vec![serde_json::json!({
+                "databases": [
+                    {"name": "admin", "sizeOnDisk": 40960, "empty": false},
+                    {"name": "app", "sizeOnDisk": 8192, "empty": true},
+                    {"name": "logs", "sizeOnDisk": 1024, "empty": false}
+                ],
+                "totalSize": 50176,
+                "ok": 1
+            })],
+            raw_documents: None,
+            extended_documents: None,
+            total: 1,
+            total_is_exact: true,
+        };
+
+        let query_result = mongo_show_databases_query_result(result.documents, 2).unwrap();
+
+        assert_eq!(query_result.columns, ["name", "sizeOnDisk", "empty"]);
+        assert_eq!(
+            query_result.rows,
+            [
+                [serde_json::json!("admin"), serde_json::json!(40960), serde_json::json!(false)],
+                [serde_json::json!("app"), serde_json::json!(8192), serde_json::json!(true)],
+            ]
+        );
+        assert_eq!(query_result.affected_rows, 3);
+        assert!(query_result.truncated);
+        assert!(query_result.has_more);
+    }
+
+    #[test]
+    fn show_databases_result_handles_empty_and_rejects_malformed_responses() {
+        let empty = MongoDocumentResult {
+            documents: vec![serde_json::json!({"databases": [], "ok": 1})],
+            raw_documents: None,
+            extended_documents: None,
+            total: 1,
+            total_is_exact: true,
+        };
+        let query_result = mongo_show_databases_query_result(empty.documents, 100).unwrap();
+        assert_eq!(query_result.columns, ["name", "sizeOnDisk", "empty"]);
+        assert!(query_result.rows.is_empty());
+        assert!(!query_result.truncated);
+
+        let malformed = MongoDocumentResult {
+            documents: vec![serde_json::json!({"ok": 1})],
+            raw_documents: None,
+            extended_documents: None,
+            total: 1,
+            total_is_exact: true,
+        };
+        assert!(mongo_show_databases_query_result(malformed.documents, 100).unwrap_err().contains("databases array"));
+    }
+
     #[cfg(unix)]
     use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec};
     #[cfg(unix)]
@@ -1365,6 +1468,7 @@ EXPECTED_RESULT = json.loads({expected_result})
 CAPABILITIES = {capabilities}
 SERVER_VERSION = {server_version}
 EXPECTED_ERROR = {expected_error}
+expected_calls = 0
 
 print(json.dumps({{"ready": True}}), flush=True)
 for line in sys.stdin:
@@ -1382,6 +1486,10 @@ for line in sys.stdin:
     if request.get("method") != EXPECTED_METHOD or request.get("params") != EXPECTED_PARAMS:
         print(json.dumps({{"jsonrpc": "2.0", "id": request["id"], "error": {{"code": -1, "message": "unexpected MongoDB RPC"}}}}), flush=True)
         continue
+    expected_calls += 1
+    if expected_calls > 1:
+        print(json.dumps({{"jsonrpc": "2.0", "id": request["id"], "error": {{"code": -1, "message": "duplicate MongoDB RPC"}}}}), flush=True)
+        continue
     if EXPECTED_ERROR is not None:
         print(json.dumps({{"jsonrpc": "2.0", "id": request["id"], "error": {{"code": -1, "message": EXPECTED_ERROR}}}}), flush=True)
     else:
@@ -1477,6 +1585,37 @@ for line in sys.stdin:
         );
     }
 
+    #[cfg(unix)]
+    #[tokio::test]
+    async fn mongo_show_databases_uses_one_admin_agent_command() {
+        let expected_result = serde_json::json!({
+            "documents": [{
+                "databases": [{"name": "admin", "sizeOnDisk": 40960, "empty": false}],
+                "ok": 1
+            }],
+            "total": 1,
+        });
+        let (state, _directory) = legacy_mongo_state_with_capabilities(
+            "run_command",
+            serde_json::json!({
+                "database": "admin",
+                "command_json": "{\"listDatabases\":1}",
+            }),
+            expected_result,
+            &[AgentCapability::MongoRunCommand.as_str()],
+        )
+        .await;
+
+        let query_result =
+            execute_mongo_command_core(&state, "legacy", "ignored-current-database", &MongoCommand::ShowDatabases, 100)
+                .await
+                .unwrap();
+
+        assert_eq!(query_result.columns, ["name", "sizeOnDisk", "empty"]);
+        assert_eq!(query_result.rows.len(), 1);
+        assert_eq!(query_result.rows[0][0], "admin");
+    }
+
     #[cfg(unix)]
     #[tokio::test]
     async fn mongo_run_command_requires_an_explicit_legacy_agent_capability() {

crates/dbx-core/src/mongo_shell.rs

@@ -8,6 +8,8 @@ pub enum MongoCommand {
     Version,
     #[serde(rename = "use")]
     Use { database: String },
+    #[serde(rename = "showDatabases")]
+    ShowDatabases,
     #[serde(rename = "runCommand")]
     RunCommand {
         #[serde(rename = "commandJson")]
@@ -355,6 +357,10 @@ fn mongo_field_predicate_is_exists_true(value: &serde_json::Value) -> bool {
 }
 
 pub fn parse(input: &str) -> Result<MongoCommand, String> {
+    let show_source = trim_mongo_outer_comments(input).trim_end_matches(';').trim();
+    if parse_show_databases(show_source) {
+        return Ok(MongoCommand::ShowDatabases);
+    }
     let source = input.trim().trim_end_matches(';').trim();
     if source.eq_ignore_ascii_case("db.version()") {
         return Ok(MongoCommand::Version);
@@ -645,6 +651,86 @@ pub fn parse(input: &str) -> Result<MongoCommand, String> {
     Err("Unsupported MongoDB shell command.".to_string())
 }
 
+fn parse_show_databases(source: &str) -> bool {
+    let mut words = source.split_whitespace();
+    words.next().is_some_and(|word| word.eq_ignore_ascii_case("show"))
+        && words.next().is_some_and(|word| word.eq_ignore_ascii_case("dbs") || word.eq_ignore_ascii_case("databases"))
+        && words.next().is_none()
+}
+
+fn trim_mongo_outer_comments(mut source: &str) -> &str {
+    loop {
+        source = source.trim_start();
+        if source.starts_with("//") || source.starts_with("--") {
+            source = source
+                .char_indices()
+                .find_map(|(index, character)| (character == '\n' || character == '\r').then_some(&source[index + 1..]))
+                .unwrap_or("");
+            continue;
+        }
+        if source.starts_with("/*") {
+            let Some(end) = source.find("*/") else {
+                return source;
+            };
+            source = &source[end + 2..];
+            continue;
+        }
+        break;
+    }
+
+    let source = source.trim_end();
+    let mut body_end = 0;
+    let mut quote = None;
+    let mut escaped = false;
+    let mut line_comment = false;
+    let mut block_comment = false;
+    let mut characters = source.char_indices().peekable();
+    while let Some((index, character)) = characters.next() {
+        let next = characters.peek().map(|(_, character)| *character);
+        if line_comment {
+            if character == '\n' || character == '\r' {
+                line_comment = false;
+            }
+            continue;
+        }
+        if block_comment {
+            if character == '*' && next == Some('/') {
+                block_comment = false;
+                characters.next();
+            }
+            continue;
+        }
+        if let Some(delimiter) = quote {
+            body_end = index + character.len_utf8();
+            if escaped {
+                escaped = false;
+            } else if character == '\\' {
+                escaped = true;
+            } else if character == delimiter {
+                quote = None;
+            }
+            continue;
+        }
+        if (character == '/' && next == Some('/')) || (character == '-' && next == Some('-')) {
+            line_comment = true;
+            characters.next();
+            continue;
+        }
+        if character == '/' && next == Some('*') {
+            block_comment = true;
+            characters.next();
+            continue;
+        }
+        if matches!(character, '"' | '\'' | '`') {
+            quote = Some(character);
+        }
+        if !character.is_whitespace() {
+            body_end = index + character.len_utf8();
+        }
+    }
+    source[..body_end].trim_end()
+}
+
 fn parse_collection_prefix(source: &str) -> Result<(String, usize), String> {
     if !source.get(..3).is_some_and(|prefix| prefix.eq_ignore_ascii_case("db.")) {
         return Err("MongoDB command must start with db.<collection>.".to_string());
@@ -1064,6 +1150,27 @@ mod tests {
         assert_eq!(validate_safety(&command, true, true, false), Ok(()));
     }
 
+    #[test]
+    fn parses_show_databases_aliases_as_read_only() {
+        for source in [
+            "show dbs",
+            "SHOW DATABASES;",
+            "ShOw DbS ;",
+            "/* databases */ show dbs -- list",
+            "// databases\nshow databases; /* list */",
+        ] {
+            let command = parse(source).unwrap();
+            assert_eq!(command, MongoCommand::ShowDatabases, "{source}");
+            assert!(!command.is_mutating(), "{source}");
+            assert!(!command.is_dangerous(), "{source}");
+            assert_eq!(validate_safety(&command, false, false, true), Ok(()), "{source}");
+        }
+
+        for source in ["show dbs extra", "show database", "show collections", "show"] {
+            assert!(parse(source).is_err(), "{source}");
+        }
+    }
+
     #[test]
     fn rejects_unsupported_run_command_shapes() {
         for source in [
@@ -1230,6 +1337,8 @@ mod tests {
         let run_command = serde_json::to_value(parse("db.runCommand({ping: 1})").unwrap()).unwrap();
         assert_eq!(run_command["kind"], "runCommand");
         assert_eq!(run_command["commandJson"], r#"{"ping":1}"#);
+        let show_databases = serde_json::to_value(parse("show dbs").unwrap()).unwrap();
+        assert_eq!(show_databases["kind"], "showDatabases");
     }
 
     #[test]