Files
rust-work/mongo-proxy/src/cli.rs

279 lines
9.0 KiB
Rust

use anyhow::{Result, bail};
use serde::Deserialize;
use std::time::Instant;
const DEFAULT_SERVER: &str = "http://127.0.0.1:3309";
#[derive(Debug, Deserialize)]
struct ErrorResponse {
error: String,
}
#[derive(Debug, Deserialize)]
struct ConnectionsResponse {
connections: Vec<ConnectionInfo>,
}
#[derive(Debug, Deserialize)]
struct ConnectionInfo {
name: String,
database: String,
host: String,
status: String,
}
pub struct Cli {
server: String,
}
impl Cli {
pub fn new(server: Option<String>) -> Self {
Self { server: server.unwrap_or_else(|| DEFAULT_SERVER.to_string()) }
}
pub fn check_server(&self) -> Result<bool> {
match ureq::get(&format!("{}/health", self.server)).call() {
Ok(_) => Ok(true),
Err(_) => Ok(false),
}
}
pub fn find(&self, conn: &str, collection: &str, filter: &str, format: &str,
projection: Option<&str>, limit: Option<i64>, skip: Option<u64>) -> Result<()> {
let start = Instant::now();
let mut body = serde_json::json!({
"conn": conn,
"collection": collection,
"filter": serde_json::from_str::<serde_json::Value>(filter).unwrap_or(serde_json::json!({}))
});
if let Some(p) = projection {
body["projection"] = serde_json::from_str(p).unwrap_or(serde_json::Value::Null);
}
if let Some(l) = limit {
body["limit"] = serde_json::json!(l);
}
if let Some(s) = skip {
body["skip"] = serde_json::json!(s);
}
let resp = post_json(&format!("{}/find", self.server), &body)?;
let total_ms = start.elapsed().as_millis();
if let Ok(err) = serde_json::from_str::<ErrorResponse>(&resp) {
eprintln!("Error: {}", err.error);
return Ok(());
}
#[derive(Deserialize)]
struct FindResp {
documents: Vec<serde_json::Value>,
count: usize,
#[serde(rename = "durationMs")]
duration_ms: u64,
}
let result: FindResp = serde_json::from_str(&resp)?;
match format {
"json" => println!("{}", serde_json::to_string_pretty(&result.documents).unwrap_or_default()),
_ => {
for doc in &result.documents {
println!("{}", serde_json::to_string_pretty(doc).unwrap_or_default());
}
}
}
println!("\n{} documents ({}ms db, {}ms total)", result.count, result.duration_ms, total_ms);
Ok(())
}
pub fn insert(&self, conn: &str, collection: &str, documents: &str) -> Result<()> {
let start = Instant::now();
let docs: Vec<serde_json::Value> = serde_json::from_str(documents)
.map_err(|e| anyhow::anyhow!("Invalid JSON documents: {}", e))?;
let body = serde_json::json!({
"conn": conn,
"collection": collection,
"documents": docs
});
let resp = post_json(&format!("{}/insert", self.server), &body)?;
let total_ms = start.elapsed().as_millis();
if let Ok(err) = serde_json::from_str::<ErrorResponse>(&resp) {
eprintln!("Error: {}", err.error);
return Ok(());
}
#[derive(Deserialize)]
struct InsertResp {
#[serde(rename = "insertedCount")]
inserted_count: u64,
}
let result: InsertResp = serde_json::from_str(&resp)?;
println!("Inserted {} document(s) ({}ms total)", result.inserted_count, total_ms);
Ok(())
}
pub fn update(&self, conn: &str, collection: &str, filter: &str, update: &str, upsert: bool, multi: bool) -> Result<()> {
let start = Instant::now();
let mut body = serde_json::json!({
"conn": conn,
"collection": collection,
"filter": serde_json::from_str::<serde_json::Value>(filter)?,
"update": serde_json::from_str::<serde_json::Value>(update)?
});
if upsert { body["upsert"] = serde_json::json!(true); }
if multi { body["multi"] = serde_json::json!(true); }
let resp = post_json(&format!("{}/update", self.server), &body)?;
let total_ms = start.elapsed().as_millis();
if let Ok(err) = serde_json::from_str::<ErrorResponse>(&resp) {
eprintln!("Error: {}", err.error);
return Ok(());
}
#[derive(Deserialize)]
struct UpdateResp {
#[serde(rename = "matchedCount")]
matched_count: u64,
#[serde(rename = "modifiedCount")]
modified_count: u64,
}
let result: UpdateResp = serde_json::from_str(&resp)?;
println!("Matched: {}, Modified: {} ({}ms total)", result.matched_count, result.modified_count, total_ms);
Ok(())
}
pub fn delete(&self, conn: &str, collection: &str, filter: &str, multi: bool) -> Result<()> {
let start = Instant::now();
let mut body = serde_json::json!({
"conn": conn,
"collection": collection,
"filter": serde_json::from_str::<serde_json::Value>(filter)?
});
if multi { body["multi"] = serde_json::json!(true); }
let resp = post_json(&format!("{}/delete", self.server), &body)?;
let total_ms = start.elapsed().as_millis();
if let Ok(err) = serde_json::from_str::<ErrorResponse>(&resp) {
eprintln!("Error: {}", err.error);
return Ok(());
}
#[derive(Deserialize)]
struct DeleteResp {
#[serde(rename = "deletedCount")]
deleted_count: u64,
}
let result: DeleteResp = serde_json::from_str(&resp)?;
println!("Deleted {} document(s) ({}ms total)", result.deleted_count, total_ms);
Ok(())
}
pub fn aggregate(&self, conn: &str, collection: &str, pipeline: &str) -> Result<()> {
let start = Instant::now();
let body = serde_json::json!({
"conn": conn,
"collection": collection,
"pipeline": serde_json::from_str::<serde_json::Value>(pipeline)?
});
let resp = post_json(&format!("{}/aggregate", self.server), &body)?;
let total_ms = start.elapsed().as_millis();
if let Ok(err) = serde_json::from_str::<ErrorResponse>(&resp) {
eprintln!("Error: {}", err.error);
return Ok(());
}
#[derive(Deserialize)]
struct AggResp {
documents: Vec<serde_json::Value>,
#[serde(rename = "durationMs")]
duration_ms: u64,
}
let result: AggResp = serde_json::from_str(&resp)?;
for doc in &result.documents {
println!("{}", serde_json::to_string_pretty(doc).unwrap_or_default());
}
println!("\n{} result(s) ({}ms db, {}ms total)", result.documents.len(), result.duration_ms, total_ms);
Ok(())
}
pub fn count(&self, conn: &str, collection: &str, filter: Option<&str>) -> Result<()> {
let start = Instant::now();
let mut body = serde_json::json!({
"conn": conn,
"collection": collection
});
if let Some(f) = filter {
body["filter"] = serde_json::from_str::<serde_json::Value>(f).unwrap_or(serde_json::json!({}));
}
let resp = post_json(&format!("{}/count", self.server), &body)?;
let total_ms = start.elapsed().as_millis();
if let Ok(err) = serde_json::from_str::<ErrorResponse>(&resp) {
eprintln!("Error: {}", err.error);
return Ok(());
}
#[derive(Deserialize)]
struct CountResp {
count: u64,
#[serde(rename = "durationMs")]
duration_ms: u64,
}
let result: CountResp = serde_json::from_str(&resp)?;
println!("{} documents ({}ms db, {}ms total)", result.count, result.duration_ms, total_ms);
Ok(())
}
pub fn list_connections(&self) -> Result<()> {
let response = ureq::get(&format!("{}/connections", self.server)).call()?;
let body = response.into_string()?;
if let Ok(err) = serde_json::from_str::<ErrorResponse>(&body) {
bail!("{}", err.error);
}
let result: ConnectionsResponse = serde_json::from_str(&body)?;
println!("Connections:");
println!("{:<15} {:<20} {:<40} {:<10}", "Name", "Database", "Host", "Status");
println!("{}", "-".repeat(85));
for conn in result.connections {
println!("{:<15} {:<20} {:<40} {:<10}",
conn.name, conn.database, conn.host, conn.status);
}
Ok(())
}
}
fn post_json(url: &str, body: &serde_json::Value) -> Result<String> {
let data = serde_json::to_string(body)?;
let response = ureq::post(url)
.set("Content-Type", "application/json")
.send_string(&data);
match response {
Ok(r) => Ok(r.into_string()?),
Err(ureq::Error::Status(_, resp)) => Ok(resp.into_string()?),
Err(e) => bail!("HTTP error: {}", e),
}
}