github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 06:00:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mongo): preserve BSON types in MongoDB-to-MongoDB transfer

Browse Source
This commit is contained in:
onenewcode
2026-07-01 16:31:39 +08:00
committed by GitHub
parent 4da237d955
commit e09ec3cd7a
9 changed files with 389 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agents/common/src/main/java/com/dbx/agent/AgentProtocol.java
+5 -1
View File
@@ -39,6 +39,10 @@ public final class AgentProtocol {
public static final String MONGO_METHOD_LIST_DATABASES = "list_databases";
public static final String MONGO_METHOD_LIST_COLLECTIONS = "list_collections";
public static final String MONGO_METHOD_FIND_DOCUMENTS = "find_documents";
/**
* MongoDB read path that returns documents as relaxed Extended JSON for transfer.
*/
public static final String MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON = "find_documents_extended_json";
public static final String MONGO_METHOD_SERVER_VERSION = "server_version";
public static final String MONGO_METHOD_INSERT_DOCUMENT = "insert_document";
public static final String MONGO_METHOD_UPDATE_DOCUMENT = "update_document";
@@ -115,10 +119,10 @@ public final class AgentProtocol {
MONGO_METHOD_LIST_DATABASES,
MONGO_METHOD_LIST_COLLECTIONS,
MONGO_METHOD_FIND_DOCUMENTS,
MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON,
MONGO_METHOD_SERVER_VERSION,
MONGO_METHOD_INSERT_DOCUMENT,
MONGO_METHOD_UPDATE_DOCUMENT,
MONGO_METHOD_UPDATE_DOCUMENTS,
MONGO_METHOD_DELETE_DOCUMENT
));
agents/common/src/main/resources/agent-protocol-v1.json
+1 -1
View File
@@ -68,10 +68,10 @@
"list_databases",
"list_collections",
"find_documents",
"find_documents_extended_json",
"server_version",
"insert_document",
"update_document",
"update_documents",
"delete_document"
],
"kvMethods": [
agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java
+48
View File
@@ -45,12 +45,17 @@ import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import org.bson.Document;
import org.bson.json.JsonMode;
import org.bson.json.JsonWriterSettings;
import org.bson.types.ObjectId;
public final class MongoAgent {
private static final Gson GSON = new Gson();
private static final DateTimeFormatter DATE_FORMAT = DateTimeFormatter.ISO_INSTANT.withZone(ZoneOffset.UTC);
private static final long JS_MAX_SAFE_INTEGER = 9_007_199_254_740_991L;
private static final JsonWriterSettings EXTENDED_JSON_SETTINGS = JsonWriterSettings.builder()
.outputMode(JsonMode.RELAXED)
.build();
private static MongoClient client;
private MongoAgent() {
@@ -383,6 +388,44 @@ public final class MongoAgent {
return result;
}
/**
* MongoDB Extended JSON read path for transfer; output follows the driver's
* relaxed Extended JSON representation rather than the UI display format.
*/
private static Object findDocumentsExtendedJson(JsonObject params) {
MongoClient c = requireClient();
String database = params.get("database").getAsString();
String collection = params.get("collection").getAsString();
long skip = params.has("skip") ? params.get("skip").getAsLong() : 0;
int limit = params.has("limit") ? params.get("limit").getAsInt() : 50;
Document filterDoc = documentOrNull(params, "filter");
Document projectionDoc = documentOrNull(params, "projection");
Document sortDoc = documentOrNull(params, "sort");
var col = c.getDatabase(database).getCollection(collection);
if (filterDoc == null) {
filterDoc = new Document();
}
long total = col.countDocuments(filterDoc);
var iterable = col.find(filterDoc).skip((int) skip).limit(limit);
if (projectionDoc != null) {
iterable = iterable.projection(projectionDoc);
}
if (sortDoc != null) {
iterable = iterable.sort(sortDoc);
}
List<JsonObject> documents = new ArrayList<>();
for (Document document : iterable) {
documents.add(bsonToExtendedJson(document));
}
Map<String, Object> result = new LinkedHashMap<>();
result.put("documents", documents);
result.put("total", total);
return result;
}
private static Object serverVersion(JsonObject params) {
MongoClient c = requireClient();
String database = defaultString(stringOrNull(params, "database"), "admin");
@@ -501,6 +544,10 @@ public final class MongoAgent {
return result;
}
static JsonObject bsonToExtendedJson(Document doc) {
return JsonParser.parseString(doc.toJson(EXTENDED_JSON_SETTINGS)).getAsJsonObject();
}
static Object convertValue(Object value) {
if (value == null) {
return null;
@@ -609,6 +656,7 @@ public final class MongoAgent {
case AgentProtocol.MONGO_METHOD_LIST_COLLECTIONS -> listCollections(params);
case AgentProtocol.METHOD_LIST_INDEXES -> listIndexes(params);
case AgentProtocol.MONGO_METHOD_FIND_DOCUMENTS -> findDocuments(params);
case AgentProtocol.MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON -> findDocumentsExtendedJson(params);
case AgentProtocol.MONGO_METHOD_SERVER_VERSION -> serverVersion(params);
case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENT -> insertDocument(params);
case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENT -> updateDocument(params);
agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java
+20 -1
View File
@@ -20,6 +20,7 @@ import java.security.PrivateKey;
import java.util.Base64;
import java.util.Date;
import org.bson.Document;
import org.bson.types.ObjectId;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
@@ -146,7 +147,6 @@ class MongoAgentTest {
assertEquals(10, json.get("id").getAsInt());
assertEquals("Not connected", json.getAsJsonObject("error").get("message").getAsString());
assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method"));
assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENTS));
}
@Test
@@ -355,6 +355,25 @@ class MongoAgentTest {
assertEquals("ISODate(\"2026-06-10T13:59:31.287Z\")", MongoAgent.convertValue(Date.from(java.time.Instant.parse("2026-06-10T13:59:31.287Z"))));
}
@Test
void convertValueKeepsObjectIdAsStringByDefault() {
assertEquals(
"507f1f77bcf86cd799439011",
MongoAgent.convertValue(new ObjectId("507f1f77bcf86cd799439011"))
);
}
@Test
void bsonToExtendedJsonUsesMongoExtendedJson() {
Document doc = new Document("_id", new ObjectId("507f1f77bcf86cd799439011"))
.append("created_at", Date.from(java.time.Instant.parse("2026-06-10T13:59:31.287Z")));
assertEquals(
"{\"_id\":{\"$oid\":\"507f1f77bcf86cd799439011\"},\"created_at\":{\"$date\":\"2026-06-10T13:59:31.287Z\"}}",
new com.google.gson.Gson().toJson(MongoAgent.bsonToExtendedJson(doc))
);
}
@Test
void documentForWriteParsesMongoShellIsoDateStrings() {
Document doc = MongoAgent.documentForWrite("{\"$set\":{\"CreateDate\":\"ISODate(\\\"2026-06-10T13:59:31.287Z\\\")\"}}");
crates/dbx-core/assets/agent-protocol-v1.json
+1 -1
View File
@@ -68,10 +68,10 @@
"list_databases",
"list_collections",
"find_documents",
"find_documents_extended_json",
"server_version",
"insert_document",
"update_document",
"update_documents",
"delete_document"
],
"kvMethods": [
crates/dbx-core/src/db/agent_driver.rs
+14 -1
View File
@@ -243,6 +243,7 @@ pub enum MongoAgentMethod {
ListDatabases,
ListCollections,
FindDocuments,
FindDocumentsExtendedJson,
ServerVersion,
InsertDocument,
UpdateDocument,
@@ -255,10 +256,10 @@ impl MongoAgentMethod {
Self::ListDatabases,
Self::ListCollections,
Self::FindDocuments,
Self::FindDocumentsExtendedJson,
Self::ServerVersion,
Self::InsertDocument,
Self::UpdateDocument,
Self::UpdateDocuments,
Self::DeleteDocument,
];
@@ -267,6 +268,7 @@ impl MongoAgentMethod {
Self::ListDatabases => "list_databases",
Self::ListCollections => "list_collections",
Self::FindDocuments => "find_documents",
Self::FindDocumentsExtendedJson => "find_documents_extended_json",
Self::ServerVersion => "server_version",
Self::InsertDocument => "insert_document",
Self::UpdateDocument => "update_document",
@@ -927,6 +929,14 @@ impl AgentDriverClient {
self.call_mongo_method(MongoAgentMethod::FindDocuments, 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,
params: Value,
) -> Result<T, String> {
self.call_mongo_method(MongoAgentMethod::FindDocumentsExtendedJson, params).await
}
pub async fn mongo_server_version<T: DeserializeOwned + Send + 'static>(
&mut self,
database: &str,
@@ -1481,6 +1491,7 @@ mod tests {
assert_eq!(MongoAgentMethod::ListDatabases.as_str(), "list_databases");
assert_eq!(MongoAgentMethod::ListCollections.as_str(), "list_collections");
assert_eq!(MongoAgentMethod::FindDocuments.as_str(), "find_documents");
assert_eq!(MongoAgentMethod::FindDocumentsExtendedJson.as_str(), "find_documents_extended_json");
assert_eq!(MongoAgentMethod::ServerVersion.as_str(), "server_version");
assert_eq!(MongoAgentMethod::InsertDocument.as_str(), "insert_document");
assert_eq!(MongoAgentMethod::UpdateDocument.as_str(), "update_document");
@@ -1521,6 +1532,8 @@ mod tests {
let _mongo_list_databases = AgentDriverClient::mongo_list_databases::<serde_json::Value>;
let _mongo_list_collections = AgentDriverClient::mongo_list_collections::<serde_json::Value>;
let _mongo_find_documents = AgentDriverClient::mongo_find_documents::<serde_json::Value>;
let _mongo_find_documents_extended_json =
AgentDriverClient::mongo_find_documents_extended_json::<serde_json::Value>;
let _mongo_server_version = AgentDriverClient::mongo_server_version::<serde_json::Value>;
let _mongo_insert_document = AgentDriverClient::mongo_insert_document::<serde_json::Value>;
let _mongo_update_document = AgentDriverClient::mongo_update_document::<serde_json::Value>;
crates/dbx-core/src/db/mongo_driver.rs
+139 -2
View File
@@ -245,8 +245,63 @@ pub async fn find_documents(
let mut documents = Vec::new();
while cursor.advance().await.map_err(|e| e.to_string())? {
let doc = cursor.deserialize_current().map_err(|e| e.to_string())?;
let json = bson_to_json(&Bson::Document(doc));
documents.push(json);
documents.push(bson_to_json(&Bson::Document(doc)));
}
Ok(MongoDocumentResult { documents, total })
}
/// Find MongoDB documents as relaxed Extended JSON for MongoDB transfer paths.
#[allow(clippy::too_many_arguments)]
pub async fn find_documents_extended_json(
client: &Client,
database: &str,
collection: &str,
skip: u64,
limit: i64,
filter: Option<&str>,
projection: Option<&str>,
sort: Option<&str>,
) -> Result<MongoDocumentResult, String> {
let col = client.database(database).collection::<Document>(collection);
let filter_doc: Document = match filter {
Some(f) if !f.trim().is_empty() => {
let json: serde_json::Value = serde_json::from_str(f).map_err(|e| format!("Invalid filter JSON: {e}"))?;
json_filter_to_document(&json)?
}
_ => doc! {},
};
let total = if filter_doc.is_empty() {
col.estimated_document_count().await.map_err(|e| e.to_string())?
} else {
col.count_documents(filter_doc.clone()).await.map_err(|e| e.to_string())?
};
let mut find = col.find(filter_doc).skip(skip).limit(limit);
if let Some(p) = projection {
if !p.trim().is_empty() {
let json: serde_json::Value =
serde_json::from_str(p).map_err(|e| format!("Invalid projection JSON: {e}"))?;
let projection_doc = json_object_to_document(&json).map_err(|e| format!("Invalid projection: {e}"))?;
find = find.projection(projection_doc);
}
}
if let Some(s) = sort {
if !s.trim().is_empty() {
let json: serde_json::Value = serde_json::from_str(s).map_err(|e| format!("Invalid sort JSON: {e}"))?;
let sort_doc = json_object_to_document(&json).map_err(|e| format!("Invalid sort: {e}"))?;
find = find.sort(sort_doc);
}
}
let mut cursor = find.await.map_err(|e| e.to_string())?;
let mut documents = Vec::new();
while cursor.advance().await.map_err(|e| e.to_string())? {
let doc = cursor.deserialize_current().map_err(|e| e.to_string())?;
documents.push(Bson::Document(doc).into_relaxed_extjson());
}
Ok(MongoDocumentResult { documents, total })
@@ -347,6 +402,28 @@ pub async fn insert_documents(
Ok(result.inserted_ids.len() as u64)
}
pub async fn insert_documents_extended_json(
client: &Client,
database: &str,
collection: &str,
docs_json: &str,
) -> Result<u64, String> {
let json: serde_json::Value = serde_json::from_str(docs_json).map_err(|e| format!("Invalid JSON: {e}"))?;
let docs = match json {
serde_json::Value::Array(values) => values
.into_iter()
.map(|value| json_object_to_document_extended_json(&value).map_err(|e| format!("Invalid document: {e}")))
.collect::<Result<Vec<Document>, String>>()?,
value => vec![json_object_to_document_extended_json(&value).map_err(|e| format!("Invalid document: {e}"))?],
};
if docs.is_empty() {
return Ok(0);
}
let col = client.database(database).collection::<Document>(collection);
let result = col.insert_many(docs).await.map_err(|e| e.to_string())?;
Ok(result.inserted_ids.len() as u64)
}
pub async fn update_document(
client: &Client,
database: &str,
@@ -479,6 +556,13 @@ pub fn json_object_to_document(value: &serde_json::Value) -> Result<Document, St
}
}
fn json_object_to_document_extended_json(value: &serde_json::Value) -> Result<Document, String> {
match Bson::try_from(value.clone()).map_err(|e| e.to_string())? {
Bson::Document(doc) => Ok(doc),
other => Err(format!("Expected a JSON object, got {other:?}")),
}
}
fn json_object_to_document_preserving_existing(
value: &serde_json::Value,
existing: Option<&Document>,
@@ -841,6 +925,35 @@ mod tests {
assert_eq!(value, serde_json::json!(42));
}
#[test]
fn bson_to_json_displays_object_id_as_string() {
let oid = ObjectId::parse_str("507f1f77bcf86cd799439011").unwrap();
let value = bson_to_json(&Bson::ObjectId(oid));
assert_eq!(value, serde_json::json!("507f1f77bcf86cd799439011"));
}
#[test]
fn bson_to_extended_json_preserves_nested_object_ids() {
let oid = ObjectId::parse_str("507f1f77bcf86cd799439011").unwrap();
let nested_oid = ObjectId::parse_str("507f191e810c19729de860ea").unwrap();
let value = Bson::Document(doc! {
"_id": Bson::ObjectId(oid),
"owner": { "id": Bson::ObjectId(nested_oid) },
"tags": [Bson::ObjectId(nested_oid)],
})
.into_relaxed_extjson();
assert_eq!(
value,
serde_json::json!({
"_id": { "$oid": "507f1f77bcf86cd799439011" },
"owner": { "id": { "$oid": "507f191e810c19729de860ea" } },
"tags": [{ "$oid": "507f191e810c19729de860ea" }],
})
);
}
#[test]
fn index_info_from_model_maps_mongodb_index_metadata() {
let model = IndexModel::builder()
@@ -893,6 +1006,30 @@ mod tests {
));
}
#[test]
fn json_object_to_document_parses_extended_json_object_id() {
let value = serde_json::json!({
"_id": { "$oid": "507f1f77bcf86cd799439011" },
});
let doc = json_object_to_document(&value).unwrap();
assert!(matches!(doc.get("_id"), Some(Bson::ObjectId(oid)) if oid.to_hex() == "507f1f77bcf86cd799439011"));
}
#[test]
fn json_object_to_document_extended_json_parses_official_wrappers() {
let value = serde_json::json!({
"_id": { "$oid": "507f1f77bcf86cd799439011" },
"created_at": { "$date": "2026-06-10T13:59:31.287Z" },
"count": { "$numberLong": "42" },
});
let doc = json_object_to_document_extended_json(&value).unwrap();
assert!(matches!(doc.get("_id"), Some(Bson::ObjectId(oid)) if oid.to_hex() == "507f1f77bcf86cd799439011"));
assert!(matches!(doc.get("created_at"), Some(Bson::DateTime(_))));
assert!(matches!(doc.get("count"), Some(Bson::Int64(42))));
}
#[test]
fn json_object_to_document_parses_find_projection() {
let value = serde_json::json!({
crates/dbx-core/src/mongo_ops.rs
+70
View File
@@ -96,6 +96,58 @@ pub async fn mongo_find_documents_core(
.await
}
/// Read MongoDB documents as relaxed Extended JSON for MongoDB transfer paths.
#[allow(clippy::too_many_arguments)]
pub async fn mongo_find_documents_extended_json_core(
state: &AppState,
connection_id: &str,
database: &str,
collection: &str,
skip: u64,
limit: i64,
filter: Option<&str>,
projection: Option<&str>,
sort: Option<&str>,
) -> Result<MongoDocumentResult, 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::find_documents_extended_json(
client, database, collection, skip, limit, filter, projection, sort,
)
.await
}
PoolKind::Agent(client) => {
let mut client = client.lock().await;
let mut params = serde_json::json!({
"database": database,
"collection": collection,
"skip": skip,
"limit": limit,
"filter": filter,
"sort": sort,
});
if let Some(projection) = projection {
params["projection"] = serde_json::json!(projection);
}
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") => {
client.mongo_find_documents(params).await
}
Err(error) => Err(error),
}
}
_ => Err("Not a MongoDB connection".to_string()),
}
}
fn is_unknown_agent_method_error(error: &str, method: &str) -> bool {
let lower = error.to_ascii_lowercase();
lower.contains(method) && (lower.contains("unknown method") || lower.contains("method not found"))
}
pub async fn mongo_aggregate_documents_core(
state: &AppState,
connection_id: &str,
@@ -160,6 +212,24 @@ pub async fn mongo_insert_documents_core(
}
}
pub async fn mongo_insert_documents_extended_json_core(
state: &AppState,
connection_id: &str,
database: &str,
collection: &str,
docs_json: &str,
) -> Result<u64, 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::insert_documents_extended_json(client, database, collection, docs_json).await
}
PoolKind::Agent(_) => Err("MongoDB legacy agent does not support bulk insertMany/insertOne writes".to_string()),
_ => Err("Not a MongoDB connection".to_string()),
}
}
pub async fn mongo_update_document_core(
state: &AppState,
connection_id: &str,
crates/dbx-core/src/transfer.rs
+91 -18
View File
@@ -2035,7 +2035,29 @@ fn sql_rows_to_mongo_documents(columns: &[String], rows: &[Vec<serde_json::Value
.collect()
}
async fn find_mongo_documents_for_transfer(
async fn find_mongo_documents_extended_json(
state: &AppState,
connection_id: &str,
database: &str,
collection: &str,
offset: u64,
batch_size: usize,
) -> Result<MongoDocumentResult, String> {
crate::mongo_ops::mongo_find_documents_extended_json_core(
state,
connection_id,
database,
collection,
offset,
batch_size as i64,
None,
None,
Some(r#"{"_id":1}"#),
)
.await
}
async fn find_mongo_documents_for_rows(
state: &AppState,
connection_id: &str,
database: &str,
@@ -2085,6 +2107,34 @@ async fn insert_mongo_documents_for_transfer(
}
}
async fn insert_mongo_documents_extended_json_for_transfer(
state: &AppState,
connection_id: &str,
database: &str,
collection: &str,
documents: &[serde_json::Value],
) -> Result<u64, String> {
if documents.is_empty() {
return Ok(0);
}
let docs_json = serde_json::to_string(documents).map_err(|e| format!("Failed to encode MongoDB documents: {e}"))?;
match crate::mongo_ops::mongo_insert_documents_extended_json_core(
state,
connection_id,
database,
collection,
&docs_json,
)
.await
{
Ok(count) => Ok(count),
Err(error) if error.to_ascii_lowercase().contains("legacy agent") => {
insert_mongo_documents_for_transfer(state, connection_id, database, collection, documents).await
}
Err(error) => Err(error),
}
}
async fn overwrite_mongo_collection_for_transfer(
state: &AppState,
connection_id: &str,
@@ -3178,15 +3228,27 @@ where
}
let documents = if is_mongodb_transfer_type(source_db_type) {
let result = find_mongo_documents_for_transfer(
state,
&request.source_connection_id,
&request.source_database,
table,
offset,
batch_size,
)
.await?;
let result = if is_mongodb_transfer_type(target_db_type) {
find_mongo_documents_extended_json(
state,
&request.source_connection_id,
&request.source_database,
table,
offset,
batch_size,
)
.await?
} else {
find_mongo_documents_for_rows(
state,
&request.source_connection_id,
&request.source_database,
table,
offset,
batch_size,
)
.await?
};
total_rows = Some(result.total);
result.documents
} else {
@@ -3224,14 +3286,25 @@ where
}
if is_mongodb_transfer_type(target_db_type) {
insert_mongo_documents_for_transfer(
state,
&request.target_connection_id,
&request.target_database,
&target_table,
&documents,
)
.await
if is_mongodb_transfer_type(source_db_type) {
insert_mongo_documents_extended_json_for_transfer(
state,
&request.target_connection_id,
&request.target_database,
&target_table,
&documents,
)
.await
} else {
insert_mongo_documents_for_transfer(
state,
&request.target_connection_id,
&request.target_database,
&target_table,
&documents,
)
.await
}
.map_err(|e| format!("Insert failed for MongoDB collection '{target_table}' at offset {offset}: {e}"))?;
} else {
if !sql_target_prepared {