feat(mongo): add collection clone
This commit is contained in:
1 parent
e80942de17
commit
5fe3b5b435
34 files changed
+1217
-21
No files matched your search
@@ -2,7 +2,7 @@
|
||||
"protocolVersion": 1,
|
||||
"handshakeMethod": "handshake",
|
||||
"handshakeResponseFields": ["protocolVersion", "agentProtocolVersion", "capabilities"],
|
||||
"allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database"],
|
||||
"allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection"],
|
||||
"capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
|
||||
"defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
|
||||
"commonMethods": [
|
||||
@@ -39,6 +39,6 @@
|
||||
"disconnect",
|
||||
"shutdown"
|
||||
],
|
||||
"mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "drop_database", "insert_document", "update_document", "update_documents", "delete_document", "delete_documents"],
|
||||
"mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "update_document", "update_documents", "delete_document", "delete_documents"],
|
||||
"kvMethods": ["kv_list_prefix", "kv_get", "kv_put", "kv_delete", "kv_rename", "kv_history", "kv_status", "etcd_compact", "etcd_defrag", "etcd_watch_start", "etcd_watch_poll", "etcd_watch_stop", "etcd_lease_list", "etcd_lease_get", "etcd_lease_grant", "etcd_lease_keepalive_once", "etcd_lease_revoke", "etcd_auth_user_list", "etcd_auth_user_get", "etcd_auth_user_add", "etcd_auth_user_delete", "etcd_auth_user_change_password", "etcd_auth_user_grant_role", "etcd_auth_user_revoke_role", "etcd_auth_role_list", "etcd_auth_role_get", "etcd_auth_role_add", "etcd_auth_role_delete", "etcd_auth_role_grant_permission", "etcd_auth_role_revoke_permission"]
|
||||
}
|
||||
@@ -2,7 +2,7 @@
|
||||
"protocolVersion": 2,
|
||||
"handshakeMethod": "handshake",
|
||||
"handshakeResponseFields": ["protocolVersion", "agentProtocolVersion", "capabilities"],
|
||||
"allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "multi_session", "structured_error_v1"],
|
||||
"allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "multi_session", "structured_error_v1"],
|
||||
"capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session"],
|
||||
"defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session", "structured_error_v1"],
|
||||
"commonMethods": [
|
||||
@@ -43,7 +43,7 @@
|
||||
"disconnect",
|
||||
"shutdown"
|
||||
],
|
||||
"mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "drop_database", "insert_document", "update_document", "update_documents", "delete_document", "delete_documents"],
|
||||
"mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "update_document", "update_documents", "delete_document", "delete_documents"],
|
||||
"kvMethods": ["kv_list_prefix", "kv_get", "kv_put", "kv_delete", "kv_rename", "kv_history", "kv_status", "etcd_compact", "etcd_defrag", "etcd_watch_start", "etcd_watch_poll", "etcd_watch_stop", "etcd_lease_list", "etcd_lease_get", "etcd_lease_grant", "etcd_lease_keepalive_once", "etcd_lease_revoke", "etcd_auth_user_list", "etcd_auth_user_get", "etcd_auth_user_add", "etcd_auth_user_delete", "etcd_auth_user_change_password", "etcd_auth_user_grant_role", "etcd_auth_user_revoke_role", "etcd_auth_role_list", "etcd_auth_role_get", "etcd_auth_role_add", "etcd_auth_role_delete", "etcd_auth_role_grant_permission", "etcd_auth_role_revoke_permission"],
|
||||
"sessionField": "agentSessionId",
|
||||
"cursorSessionField": "sessionId"
|
||||
|
||||
@@ -972,6 +972,7 @@ pub enum AgentCapability {
|
||||
EtcdLease,
|
||||
EtcdAuth,
|
||||
MongoDropDatabase,
|
||||
MongoCloneCollection,
|
||||
MultiSession,
|
||||
StructuredErrorV1,
|
||||
}
|
||||
@@ -1054,7 +1055,7 @@ fn parse_agent_rpc_error_header(header: &str) -> (Option<i64>, String) {
|
||||
}
|
||||
|
||||
impl AgentCapability {
|
||||
pub const ALL: [Self; 21] = [
|
||||
pub const ALL: [Self; 22] = [
|
||||
Self::Connect,
|
||||
Self::TestConnection,
|
||||
Self::Metadata,
|
||||
@@ -1074,6 +1075,7 @@ impl AgentCapability {
|
||||
Self::EtcdLease,
|
||||
Self::EtcdAuth,
|
||||
Self::MongoDropDatabase,
|
||||
Self::MongoCloneCollection,
|
||||
Self::MultiSession,
|
||||
Self::StructuredErrorV1,
|
||||
];
|
||||
@@ -1099,6 +1101,7 @@ impl AgentCapability {
|
||||
Self::EtcdLease => "etcd_lease",
|
||||
Self::EtcdAuth => "etcd_auth",
|
||||
Self::MongoDropDatabase => "mongo_drop_database",
|
||||
Self::MongoCloneCollection => "mongo_clone_collection",
|
||||
Self::MultiSession => "multi_session",
|
||||
Self::StructuredErrorV1 => "structured_error_v1",
|
||||
}
|
||||
@@ -1282,6 +1285,7 @@ pub enum MongoAgentMethod {
|
||||
CreateUser,
|
||||
DropIndexes,
|
||||
DropCollection,
|
||||
CloneCollection,
|
||||
DropDatabase,
|
||||
InsertDocument,
|
||||
UpdateDocument,
|
||||
@@ -1291,7 +1295,7 @@ pub enum MongoAgentMethod {
|
||||
}
|
||||
|
||||
impl MongoAgentMethod {
|
||||
pub const ALL: [Self; 19] = [
|
||||
pub const ALL: [Self; 20] = [
|
||||
Self::ListDatabases,
|
||||
Self::ListCollections,
|
||||
Self::FindDocuments,
|
||||
@@ -1305,6 +1309,7 @@ impl MongoAgentMethod {
|
||||
Self::CreateUser,
|
||||
Self::DropIndexes,
|
||||
Self::DropCollection,
|
||||
Self::CloneCollection,
|
||||
Self::DropDatabase,
|
||||
Self::InsertDocument,
|
||||
Self::UpdateDocument,
|
||||
@@ -1328,6 +1333,7 @@ impl MongoAgentMethod {
|
||||
Self::CreateUser => "create_user",
|
||||
Self::DropIndexes => "drop_indexes",
|
||||
Self::DropCollection => "drop_collection",
|
||||
Self::CloneCollection => "clone_collection",
|
||||
Self::DropDatabase => "drop_database",
|
||||
Self::InsertDocument => "insert_document",
|
||||
Self::UpdateDocument => "update_document",
|
||||
@@ -2544,6 +2550,13 @@ impl AgentDriverClient {
|
||||
self.call_mongo_method(MongoAgentMethod::DropCollection, params).await
|
||||
}
|
||||
|
||||
pub async fn mongo_clone_collection<T: DeserializeOwned + Send + 'static>(
|
||||
&mut self,
|
||||
params: Value,
|
||||
) -> Result<T, String> {
|
||||
self.call_mongo_method(MongoAgentMethod::CloneCollection, params).await
|
||||
}
|
||||
|
||||
pub async fn mongo_drop_database<T: DeserializeOwned + Send + 'static>(
|
||||
&mut self,
|
||||
params: Value,
|
||||
@@ -2724,6 +2737,7 @@ pub fn agent_supports_capability(handshake: Option<&AgentHandshake>, capability:
|
||||
| AgentCapability::EtcdLease
|
||||
| AgentCapability::EtcdAuth
|
||||
| AgentCapability::MongoDropDatabase
|
||||
| AgentCapability::MongoCloneCollection
|
||||
) {
|
||||
return handshake.map(|value| value.supports(capability)).unwrap_or(false);
|
||||
}
|
||||
@@ -4307,9 +4321,10 @@ for line in sys.stdin:
|
||||
assert_eq!(AgentCapability::EtcdLease.as_str(), "etcd_lease");
|
||||
assert_eq!(AgentCapability::EtcdAuth.as_str(), "etcd_auth");
|
||||
assert_eq!(AgentCapability::MongoDropDatabase.as_str(), "mongo_drop_database");
|
||||
assert_eq!(AgentCapability::MongoCloneCollection.as_str(), "mongo_clone_collection");
|
||||
assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session");
|
||||
assert_eq!(AgentCapability::StructuredErrorV1.as_str(), "structured_error_v1");
|
||||
assert_eq!(AgentCapability::ALL.len(), 21);
|
||||
assert_eq!(AgentCapability::ALL.len(), 22);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -4361,6 +4376,7 @@ for line in sys.stdin:
|
||||
assert_eq!(MongoAgentMethod::CreateUser.as_str(), "create_user");
|
||||
assert_eq!(MongoAgentMethod::DropIndexes.as_str(), "drop_indexes");
|
||||
assert_eq!(MongoAgentMethod::DropCollection.as_str(), "drop_collection");
|
||||
assert_eq!(MongoAgentMethod::CloneCollection.as_str(), "clone_collection");
|
||||
assert_eq!(MongoAgentMethod::DropDatabase.as_str(), "drop_database");
|
||||
assert_eq!(MongoAgentMethod::InsertDocument.as_str(), "insert_document");
|
||||
assert_eq!(MongoAgentMethod::UpdateDocument.as_str(), "update_document");
|
||||
@@ -4429,6 +4445,7 @@ for line in sys.stdin:
|
||||
let _mongo_create_index = AgentDriverClient::mongo_create_index::<serde_json::Value>;
|
||||
let _mongo_drop_indexes = AgentDriverClient::mongo_drop_indexes::<serde_json::Value>;
|
||||
let _mongo_drop_collection = AgentDriverClient::mongo_drop_collection::<serde_json::Value>;
|
||||
let _mongo_clone_collection = AgentDriverClient::mongo_clone_collection::<serde_json::Value>;
|
||||
let _mongo_drop_database = AgentDriverClient::mongo_drop_database::<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>;
|
||||
@@ -4662,10 +4679,18 @@ for line in sys.stdin:
|
||||
assert!(!agent_supports_capability(Some(&handshake), AgentCapability::KvHistory));
|
||||
assert!(!agent_supports_capability(None, AgentCapability::MongoDropDatabase));
|
||||
assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoDropDatabase));
|
||||
assert!(!agent_supports_capability(None, AgentCapability::MongoCloneCollection));
|
||||
assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoCloneCollection));
|
||||
|
||||
let mongo_handshake =
|
||||
AgentHandshake { capabilities: vec![AgentCapability::MongoDropDatabase.as_str().to_string()], ..handshake };
|
||||
assert!(agent_supports_capability(Some(&mongo_handshake), AgentCapability::MongoDropDatabase));
|
||||
|
||||
let mongo_clone_handshake = AgentHandshake {
|
||||
capabilities: vec![AgentCapability::MongoCloneCollection.as_str().to_string()],
|
||||
..mongo_handshake
|
||||
};
|
||||
assert!(agent_supports_capability(Some(&mongo_clone_handshake), AgentCapability::MongoCloneCollection));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -58,6 +58,13 @@ pub struct MongoCollectionStatsResult {
|
||||
pub nindexes: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Result counts returned after cloning a regular MongoDB collection.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct MongoCloneCollectionResult {
|
||||
pub documents_copied: u64,
|
||||
pub indexes_copied: u64,
|
||||
}
|
||||
|
||||
pub async fn connect(url: &str, timeout: Duration, idle_timeout: Duration) -> Result<Client, String> {
|
||||
let url = normalize_mongo_uri_direct_connection(url);
|
||||
let is_multi_host = is_multi_host_mongo_uri(&url);
|
||||
@@ -645,6 +652,122 @@ pub async fn rename_collection(client: &Client, database: &str, old_name: &str,
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clone a regular collection within one database without relying on MongoDB's
|
||||
/// deprecated clone commands. The individual commands used here have been
|
||||
/// available across the MongoDB versions supported by the native driver.
|
||||
pub async fn clone_collection(
|
||||
client: &Client,
|
||||
database: &str,
|
||||
source_name: &str,
|
||||
target_name: &str,
|
||||
) -> Result<MongoCloneCollectionResult, String> {
|
||||
validate_clone_collection_names(database, source_name, target_name)?;
|
||||
|
||||
let database = client.database(database);
|
||||
let source_spec = find_collection_specification(&database, source_name).await?;
|
||||
if !matches!(source_spec.collection_type, mongodb::results::CollectionType::Collection) {
|
||||
return Err(
|
||||
"Only regular MongoDB collections can be cloned; views and time-series collections are not supported"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
// Create explicitly before copying data so an existing target fails rather
|
||||
// than being silently merged with or overwritten by the source documents.
|
||||
database
|
||||
.create_collection(target_name)
|
||||
.with_options(source_spec.options.clone())
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
|
||||
let source = database.collection::<Document>(source_name);
|
||||
let target = database.collection::<Document>(target_name);
|
||||
let needs_validation_bypass = source_spec.options.validator.is_some()
|
||||
|| source_spec.options.validation_level.is_some()
|
||||
|| source_spec.options.validation_action.is_some();
|
||||
let mut cursor = source.find(doc! {}).await.map_err(|error| error.to_string())?;
|
||||
let mut batch = Vec::with_capacity(1_000);
|
||||
let mut documents_copied = 0_u64;
|
||||
|
||||
while let Some(document) = cursor.try_next().await.map_err(|error| error.to_string())? {
|
||||
batch.push(document);
|
||||
if batch.len() == 1_000 {
|
||||
documents_copied += insert_clone_batch(&target, &mut batch, needs_validation_bypass).await?;
|
||||
}
|
||||
}
|
||||
if !batch.is_empty() {
|
||||
documents_copied += insert_clone_batch(&target, &mut batch, needs_validation_bypass).await?;
|
||||
}
|
||||
|
||||
// The target gets its _id index during createCollection. Recreating every
|
||||
// other source index after the data copy avoids needless index maintenance.
|
||||
let mut index_cursor = source.list_indexes().await.map_err(|error| error.to_string())?;
|
||||
let mut indexes_copied = 0_u64;
|
||||
while let Some(index) = index_cursor.try_next().await.map_err(|error| error.to_string())? {
|
||||
if is_automatic_id_index(&index) {
|
||||
continue;
|
||||
}
|
||||
target.create_index(index).await.map_err(|error| error.to_string())?;
|
||||
indexes_copied += 1;
|
||||
}
|
||||
|
||||
Ok(MongoCloneCollectionResult { documents_copied, indexes_copied })
|
||||
}
|
||||
|
||||
async fn find_collection_specification(
|
||||
database: &Database,
|
||||
source_name: &str,
|
||||
) -> Result<mongodb::results::CollectionSpecification, String> {
|
||||
let mut cursor = database.list_collections().await.map_err(|error| error.to_string())?;
|
||||
while let Some(specification) = cursor.try_next().await.map_err(|error| error.to_string())? {
|
||||
if specification.name == source_name {
|
||||
return Ok(specification);
|
||||
}
|
||||
}
|
||||
Err(format!("MongoDB collection '{source_name}' was not found"))
|
||||
}
|
||||
|
||||
async fn insert_clone_batch(
|
||||
target: &mongodb::Collection<Document>,
|
||||
batch: &mut Vec<Document>,
|
||||
bypass_document_validation: bool,
|
||||
) -> Result<u64, String> {
|
||||
let documents = std::mem::take(batch);
|
||||
let result = if bypass_document_validation {
|
||||
target.insert_many(documents).bypass_document_validation(true).await
|
||||
} else {
|
||||
target.insert_many(documents).await
|
||||
}
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(result.inserted_ids.len() as u64)
|
||||
}
|
||||
|
||||
/// `createCollection` always creates the `_id` index, but it need not be named
|
||||
/// `_id_` (for example, on a clustered collection). Compare the key instead
|
||||
/// of the name so cloning never attempts to create it a second time.
|
||||
fn is_automatic_id_index(index: &IndexModel) -> bool {
|
||||
index.keys.len() == 1
|
||||
&& (matches!(index.keys.get("_id"), Some(Bson::Int32(1) | Bson::Int64(1)))
|
||||
|| matches!(index.keys.get("_id"), Some(Bson::Double(value)) if *value == 1.0))
|
||||
}
|
||||
|
||||
pub(crate) fn validate_clone_collection_names(
|
||||
database: &str,
|
||||
source_name: &str,
|
||||
target_name: &str,
|
||||
) -> Result<(), String> {
|
||||
validate_mongo_namespace_name(database, "Database")?;
|
||||
validate_mongo_namespace_name(source_name, "Source collection")?;
|
||||
validate_mongo_namespace_name(target_name, "Target collection")?;
|
||||
if source_name == target_name {
|
||||
return Err("Target collection name must differ from the source collection name".to_string());
|
||||
}
|
||||
if source_name.starts_with("system.") || target_name.starts_with("system.") {
|
||||
return Err("System collections cannot be cloned".to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_indexes(client: &Client, database: &str, collection: &str) -> Result<Vec<IndexInfo>, String> {
|
||||
let col = client.database(database).collection::<Document>(collection);
|
||||
let mut cursor = col.list_indexes().await.map_err(|e| e.to_string())?;
|
||||
@@ -2568,6 +2691,33 @@ mod tests {
|
||||
assert!(!command.contains_key("dropTarget"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clone_collection_validation_preserves_identifiers_and_rejects_unsafe_targets() {
|
||||
validate_clone_collection_names("app", " users ", " users_backup ").unwrap();
|
||||
assert!(validate_clone_collection_names("app", "users", "users").unwrap_err().contains("differ"));
|
||||
assert!(validate_clone_collection_names("app", "system.users", "users_backup")
|
||||
.unwrap_err()
|
||||
.contains("System collections"));
|
||||
assert!(validate_clone_collection_names("app", "users", "system.users_backup")
|
||||
.unwrap_err()
|
||||
.contains("System collections"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clone_collection_skips_the_automatic_id_index_by_key_not_name() {
|
||||
let automatic_id = IndexModel::builder()
|
||||
.keys(doc! { "_id": 1 })
|
||||
.options(IndexOptions::builder().name("custom_id_name".to_string()).build())
|
||||
.build();
|
||||
let ordinary_index = IndexModel::builder()
|
||||
.keys(doc! { "external_id": 1 })
|
||||
.options(IndexOptions::builder().name("external_id_1".to_string()).build())
|
||||
.build();
|
||||
|
||||
assert!(is_automatic_id_index(&automatic_id));
|
||||
assert!(!is_automatic_id_index(&ordinary_index));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mongo_collection_kind_as_str_is_stable_wire_value() {
|
||||
assert_eq!(MongoCollectionKind::Collection.as_str(), "collection");
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
use crate::connection::{AppState, PoolKind};
|
||||
use crate::db::agent_driver::AgentCapability;
|
||||
use crate::db::mongo_driver::{
|
||||
self, MongoCollectionStatsResult, MongoDocumentResult, MongoDropIndexFailure, MongoDropIndexesResult,
|
||||
self, MongoCloneCollectionResult, MongoCollectionStatsResult, MongoDocumentResult, MongoDropIndexFailure,
|
||||
MongoDropIndexesResult,
|
||||
};
|
||||
use crate::document_ops::CollectionInfo;
|
||||
use crate::mongo_shell::MongoCommand;
|
||||
@@ -96,6 +97,40 @@ pub async fn mongo_rename_collection_core(
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn mongo_clone_collection_core(
|
||||
state: &AppState,
|
||||
connection_id: &str,
|
||||
database: &str,
|
||||
source_collection: &str,
|
||||
target_collection: &str,
|
||||
) -> Result<MongoCloneCollectionResult, String> {
|
||||
mongo_driver::validate_clone_collection_names(database, source_collection, target_collection)?;
|
||||
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::clone_collection(client, database, source_collection, target_collection).await
|
||||
}
|
||||
PoolKind::Agent(client) => {
|
||||
let mut client = client.lock().await;
|
||||
if !client.supports_capability(AgentCapability::MongoCloneCollection) {
|
||||
return Err(
|
||||
"MongoDB Legacy Agent does not support clone collection; upgrade or reinstall the MongoDB Legacy driver"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
client
|
||||
.mongo_clone_collection(serde_json::json!({
|
||||
"database": database,
|
||||
"source_collection": source_collection,
|
||||
"target_collection": target_collection,
|
||||
}))
|
||||
.await
|
||||
}
|
||||
_ => Err("Not a MongoDB connection".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn mongo_server_version_core(
|
||||
state: &AppState,
|
||||
connection_id: &str,
|
||||
@@ -1328,6 +1363,47 @@ for line in sys.stdin:
|
||||
mongo_drop_collection_core(&state, "legacy", "app", "users").await.unwrap();
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn mongo_clone_collection_routes_legacy_connections_to_the_agent() {
|
||||
let (state, _directory) = legacy_mongo_state_with_capabilities(
|
||||
"clone_collection",
|
||||
serde_json::json!({
|
||||
"database": "app",
|
||||
"source_collection": "users",
|
||||
"target_collection": "users_copy",
|
||||
}),
|
||||
serde_json::json!({ "documents_copied": 2, "indexes_copied": 1 }),
|
||||
&[AgentCapability::MongoCloneCollection.as_str()],
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = mongo_clone_collection_core(&state, "legacy", "app", "users", "users_copy").await.unwrap();
|
||||
|
||||
assert_eq!(result, MongoCloneCollectionResult { documents_copied: 2, indexes_copied: 1 });
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn mongo_clone_collection_requires_an_explicit_legacy_agent_capability() {
|
||||
let (state, _directory) = legacy_mongo_state_with_capabilities(
|
||||
"clone_collection",
|
||||
serde_json::json!({
|
||||
"database": "app",
|
||||
"source_collection": "users",
|
||||
"target_collection": "users_copy",
|
||||
}),
|
||||
serde_json::json!({ "documents_copied": 2, "indexes_copied": 1 }),
|
||||
&[],
|
||||
)
|
||||
.await;
|
||||
|
||||
let error = mongo_clone_collection_core(&state, "legacy", "app", "users", "users_copy").await.unwrap_err();
|
||||
|
||||
assert!(error.contains("upgrade or reinstall"), "{error}");
|
||||
assert!(!error.contains("Unknown method"), "{error}");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn mongo_list_collections_requests_legacy_collection_types() {
|
||||
|
||||
@@ -1,7 +1,12 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use dbx_core::db::mongo_driver;
|
||||
use mongodb::bson::{doc, Bson, DateTime};
|
||||
use futures::TryStreamExt;
|
||||
use mongodb::{
|
||||
bson::{doc, Bson, DateTime, Document},
|
||||
options::IndexOptions,
|
||||
IndexModel,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires DBX_LIVE_MONGODB_URL pointing at a writable MongoDB database"]
|
||||
@@ -76,6 +81,54 @@ 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 clone_collection_copies_options_documents_and_non_id_indexes() {
|
||||
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_clone".to_string());
|
||||
let source_name = format!("clone_source_{}", std::process::id());
|
||||
let target_name = format!("clone_target_{}", std::process::id());
|
||||
let database_ref = client.database(&database);
|
||||
|
||||
database_ref.create_collection(&source_name).validator(doc! { "kind": { "$type": "string" } }).await.unwrap();
|
||||
let source = database_ref.collection::<Document>(&source_name);
|
||||
source
|
||||
.insert_many(vec![
|
||||
doc! { "ordinal": 1, "kind": "record", "value": "first" },
|
||||
doc! { "ordinal": 2, "kind": "record", "value": "second" },
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
source
|
||||
.create_index(
|
||||
IndexModel::builder()
|
||||
.keys(doc! { "ordinal": 1 })
|
||||
.options(IndexOptions::builder().name("ordinal_unique".to_string()).unique(true).build())
|
||||
.build(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = mongo_driver::clone_collection(&client, &database, &source_name, &target_name).await.unwrap();
|
||||
assert_eq!(result.documents_copied, 2);
|
||||
assert_eq!(result.indexes_copied, 1);
|
||||
|
||||
let target = database_ref.collection::<Document>(&target_name);
|
||||
let source_documents =
|
||||
source.find(doc! {}).sort(doc! { "ordinal": 1 }).await.unwrap().try_collect::<Vec<Document>>().await.unwrap();
|
||||
let target_documents =
|
||||
target.find(doc! {}).sort(doc! { "ordinal": 1 }).await.unwrap().try_collect::<Vec<Document>>().await.unwrap();
|
||||
assert_eq!(target_documents, source_documents);
|
||||
assert!(target.insert_one(doc! { "ordinal": 3 }).await.is_err(), "target must retain the validator");
|
||||
|
||||
let target_indexes = mongo_driver::list_indexes(&client, &database, &target_name).await.unwrap();
|
||||
assert!(target_indexes.iter().any(|index| index.name == "ordinal_unique" && index.is_unique));
|
||||
|
||||
mongo_driver::drop_collection(&client, &database, &target_name).await.unwrap();
|
||||
mongo_driver::drop_collection(&client, &database, &source_name).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() {
|
||||
|
||||
@@ -605,6 +605,7 @@ async fn main() {
|
||||
.route("/mongo/drop-database", post(routes::mongo::drop_database))
|
||||
.route("/mongo/drop-collection", post(routes::mongo::drop_collection))
|
||||
.route("/mongo/rename-collection", post(routes::mongo::rename_collection))
|
||||
.route("/mongo/clone-collection", post(routes::mongo::clone_collection))
|
||||
.route("/document-store/list-databases", post(routes::document_store::list_databases))
|
||||
.route("/document-store/list-collections", post(routes::document_store::list_collections))
|
||||
.route("/document-store/find-documents", post(routes::document_store::find_documents))
|
||||
|
||||
@@ -86,6 +86,15 @@ pub struct MongoRenameCollectionRequest {
|
||||
pub new_name: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct MongoCloneCollectionRequest {
|
||||
pub connection_id: String,
|
||||
pub database: String,
|
||||
pub source_collection: String,
|
||||
pub target_collection: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct MongoFindRequest {
|
||||
@@ -408,6 +417,26 @@ pub async fn rename_collection(
|
||||
Ok(Json(serde_json::json!({ "ok": true })))
|
||||
}
|
||||
|
||||
pub async fn clone_collection(
|
||||
State(state): State<Arc<WebState>>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<MongoCloneCollectionRequest>,
|
||||
) -> Result<Json<dbx_core::db::mongo_driver::MongoCloneCollectionResult>, AppError> {
|
||||
super::mcp_policy::ensure_dangerous_write(&state, &headers, &req.connection_id, &req.database, "Clone collection")
|
||||
.await?;
|
||||
ensure_writable(&state.app, &req.connection_id, "Clone collection").await?;
|
||||
let result = dbx_core::mongo_ops::mongo_clone_collection_core(
|
||||
&state.app,
|
||||
&req.connection_id,
|
||||
&req.database,
|
||||
&req.source_collection,
|
||||
&req.target_collection,
|
||||
)
|
||||
.await
|
||||
.map_err(AppError::from)?;
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
pub async fn find_documents(
|
||||
State(state): State<Arc<WebState>>,
|
||||
headers: HeaderMap,
|
||||
@@ -874,7 +903,10 @@ async fn ensure_find_one_write_policy(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{drop_database, ensure_find_one_write_policy, MongoCollectionRequest};
|
||||
use super::{
|
||||
clone_collection, drop_database, ensure_find_one_write_policy, MongoCloneCollectionRequest,
|
||||
MongoCollectionRequest,
|
||||
};
|
||||
use crate::state::WebState;
|
||||
use axum::{
|
||||
extract::State,
|
||||
@@ -1012,4 +1044,38 @@ mod tests {
|
||||
drop(state);
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn clone_collection_requires_mcp_dangerous_write_approval() {
|
||||
let (state, dir) = test_web_state().await;
|
||||
let connection = mongo_config(false);
|
||||
state.app.storage.save_connections(std::slice::from_ref(&connection)).await.unwrap();
|
||||
state
|
||||
.app
|
||||
.storage
|
||||
.save_mcp_global_policy(&McpGlobalPolicy {
|
||||
read_only: false,
|
||||
allow_dangerous_sql: false,
|
||||
allowed_connection_ids: Some(vec![connection.id.clone()]),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = clone_collection(
|
||||
State(state.clone()),
|
||||
mcp_headers(),
|
||||
Json(MongoCloneCollectionRequest {
|
||||
connection_id: connection.id,
|
||||
database: "app".to_string(),
|
||||
source_collection: "users".to_string(),
|
||||
target_collection: "users_copy".to_string(),
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.message.starts_with("SQL_BLOCKED:"), "{}", error.message);
|
||||
|
||||
drop(state);
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user