use crate::colors::*;
use anyhow::{Context, Result};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tokio::io::AsyncWriteExt;
use url::Url;
#[derive(Debug, Deserialize)]
pub struct SnapshotInfo {
pub name: String,
pub creation_time: Option<String>,
pub size: u64,
}
#[derive(Debug, Deserialize)]
struct SnapshotListResponse {
result: Vec<SnapshotInfo>,
}
#[derive(Debug, Deserialize)]
struct SnapshotCreateResponse {
result: SnapshotInfo,
}
#[derive(Debug, Serialize)]
struct SnapshotRecoverRequest {
location: String,
#[serde(skip_serializing_if = "Option::is_none")]
priority: Option<String>,
}
pub async fn snapshot_create(collection: &str, url: &str) -> Result<SnapshotInfo> {
println!(
"{} Creating snapshot of '{}'...",
"→".cyan(),
collection.yellow()
);
let client = Client::new();
let endpoint = qdrant_endpoint(url, &["collections", collection, "snapshots"])?;
let response = client.post(&endpoint).send().await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
anyhow::bail!("Snapshot creation failed: {} - {}", status, body);
}
let result: SnapshotCreateResponse = response.json().await?;
println!("{} Snapshot created: {}", "✓".green(), result.result.name);
Ok(result.result)
}
pub async fn snapshot_list(collection: &str, url: &str) -> Result<Vec<SnapshotInfo>> {
let client = Client::new();
let endpoint = qdrant_endpoint(url, &["collections", collection, "snapshots"])?;
let response = client.get(&endpoint).send().await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
anyhow::bail!("Failed to list snapshots: {} - {}", status, body);
}
let result: SnapshotListResponse = response.json().await?;
Ok(result.result)
}
pub async fn snapshot_download(
collection: &str,
snapshot_name: &str,
output_path: &str,
url: &str,
) -> Result<()> {
println!(
"{} Downloading snapshot to '{}'...",
"→".cyan(),
output_path.yellow()
);
let client = Client::new();
let endpoint = qdrant_endpoint(
url,
&["collections", collection, "snapshots", snapshot_name],
)?;
let mut response = client.get(&endpoint).send().await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
anyhow::bail!("Snapshot download failed: {} - {}", status, body);
}
let mut file = tokio::fs::File::create(output_path).await?;
let mut bytes_written: u64 = 0;
while let Some(chunk) = response.chunk().await? {
file.write_all(&chunk).await?;
bytes_written += chunk.len() as u64;
}
file.flush().await?;
println!(
"{} Downloaded {} bytes to {}",
"✓".green(),
bytes_written,
output_path
);
Ok(())
}
pub async fn snapshot_restore(collection: &str, snapshot_location: &str, url: &str) -> Result<()> {
println!(
"{} Restoring '{}' from snapshot...",
"→".cyan(),
collection.yellow()
);
println!(" Location: {}", snapshot_location);
let client = Client::new();
let endpoint = qdrant_endpoint(url, &["collections", collection, "snapshots", "recover"])?;
let request = SnapshotRecoverRequest {
location: snapshot_location.to_string(),
priority: Some("snapshot".to_string()), };
let response = client.put(&endpoint).json(&request).send().await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
anyhow::bail!("Snapshot restore failed: {} - {}", status, body);
}
println!(
"{} Collection '{}' restored successfully!",
"✓".green(),
collection
);
Ok(())
}
pub async fn snapshot_delete(collection: &str, snapshot_name: &str, url: &str) -> Result<()> {
let client = Client::new();
let endpoint = qdrant_endpoint(
url,
&["collections", collection, "snapshots", snapshot_name],
)?;
let response = client.delete(&endpoint).send().await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
anyhow::bail!("Snapshot deletion failed: {} - {}", status, body);
}
println!("{} Snapshot '{}' deleted", "✓".green(), snapshot_name);
Ok(())
}
fn qdrant_endpoint(base_url: &str, segments: &[&str]) -> Result<String> {
let mut url = Url::parse(base_url).context("Invalid Qdrant URL")?;
{
let mut path = url
.path_segments_mut()
.map_err(|_| anyhow::anyhow!("Qdrant URL cannot be used as a path base"))?;
path.pop_if_empty();
for segment in segments {
path.push(segment);
}
}
Ok(url.to_string())
}
#[cfg(test)]
mod tests {
use super::qdrant_endpoint;
#[test]
fn qdrant_endpoint_encodes_collection_and_snapshot_segments() {
let endpoint = qdrant_endpoint(
"http://localhost:6333",
&["collections", "tenant/a?x", "snapshots", "snap 1.tar"],
)
.unwrap();
assert_eq!(
endpoint,
"http://localhost:6333/collections/tenant%2Fa%3Fx/snapshots/snap%201.tar"
);
}
#[test]
fn qdrant_endpoint_preserves_base_path_without_double_slash() {
let endpoint = qdrant_endpoint(
"http://localhost:6333/api/",
&["collections", "products", "snapshots"],
)
.unwrap();
assert_eq!(
endpoint,
"http://localhost:6333/api/collections/products/snapshots"
);
}
}