tauri_plugin_mongoose/db/

files.rs

use std::path::Path;

use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use futures::io::{AsyncReadExt, Cursor};
use mongodb::{
    bson::{doc, oid::ObjectId, Bson, DateTime, Document},
    gridfs::FilesCollectionDocument,
    options::{GridFsBucketOptions, GridFsUploadOptions},
};
use serde::Deserialize;
use serde_json::Value;

use crate::db::state::{get_client, get_db_name};

/// Arguments accepted by the `save_file` command.
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SaveFileRequest {
    /// GridFS bucket / collection name. Maps to `<bucket>.files` and `<bucket>.chunks` collections.
    pub collection: String,
    /// Absolute path of a file to read from disk.
    pub path: Option<Value>,
    /// Raw bytes to be written. Useful when the caller already has the buffer in memory.
    pub data: Option<Vec<u8>>,
    /// Optional filename stored in GridFS. Defaults to the basename of `path` or `"untitled"`.
    pub filename: Option<String>,
    /// Arbitrary JSON metadata stored alongside the file document.
    pub metadata: Option<Value>,
    /// Optional chunk size in bytes. Defaults to MongoDB GridFS default (255 KiB).
    pub chunk_size_bytes: Option<u32>,
}

/// Store a file in MongoDB GridFS using either a filesystem path or an in-memory buffer.
/// Returns the created files document as JSON (including the `_id`).
pub async fn save_file(args: SaveFileRequest) -> Result<Value, String> {
    let SaveFileRequest {
        collection,
        path,
        data,
        filename,
        metadata,
        chunk_size_bytes,
    } = args;

    if path.is_none() && data.is_none() {
        return Err("Either 'path' or 'data' must be provided".to_string());
    }

    let client = get_client().await?;
    let db_name = get_db_name().await;
    let db = client.database(&db_name);

    let bucket_options = GridFsBucketOptions::builder()
        .bucket_name(collection.clone())
        .chunk_size_bytes(chunk_size_bytes)
        .build();
    let bucket = db.gridfs_bucket(bucket_options);

    // Load bytes and determine filename.
    let (bytes, final_filename) = if let Some(path_value) = path {
        let path_str_opt = match path_value {
            Value::String(s) => Some(s),
            Value::Object(map) => map
                .get("path")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string())
                .or_else(|| {
                    map.get("filePath")
                        .and_then(|v| v.as_str())
                        .map(|s| s.to_string())
                }),
            _ => None,
        };

        let file_path = path_str_opt.ok_or_else(|| {
            "Invalid 'path': expected a string or object with a 'path' field".to_string()
        })?;

        let bytes = tokio::fs::read(&file_path)
            .await
            .map_err(|e| format!("Failed to read file at {}: {}", file_path, e))?;
        let derived_name = Path::new(&file_path)
            .file_name()
            .and_then(|n| n.to_str())
            .map(|s| s.to_string());
        (
            bytes,
            filename
                .or(derived_name)
                .unwrap_or_else(|| "untitled".to_string()),
        )
    } else if let Some(bytes) = data {
        (bytes, filename.unwrap_or_else(|| "untitled".to_string()))
    } else {
        // Unreachable because of the earlier check.
        return Err("No data supplied".to_string());
    };

    let length = bytes.len() as i64;
    let chunk_size = chunk_size_bytes.unwrap_or(255 * 1024);

    let metadata_doc: Option<Document> = match metadata {
        Some(value) => Some(
            mongodb::bson::to_document(&value).map_err(|e| format!("Invalid metadata: {}", e))?,
        ),
        None => None,
    };

    let upload_options = GridFsUploadOptions::builder()
        .chunk_size_bytes(chunk_size_bytes)
        .metadata(metadata_doc.clone())
        .build();

    let cursor = Cursor::new(bytes);
    let file_id = bucket
        .upload_from_futures_0_3_reader(final_filename.clone(), cursor, upload_options)
        .await
        .map_err(|e| e.to_string())?;

    // Try to fetch the stored files document for richer output.
    let files_collection =
        db.collection::<FilesCollectionDocument>(&format!("{}.files", collection));
    let stored = files_collection
        .find_one(doc! {"_id": file_id.clone()}, None)
        .await
        .map_err(|e| e.to_string())?;

    let mut response_doc = if let Some(file) = stored {
        mongodb::bson::to_document(&file).map_err(|e| e.to_string())?
    } else {
        // Fallback when the files document could not be read immediately.
        let mut fallback = doc! {
            "_id": file_id.clone(),
            "length": length,
            "chunkSizeBytes": chunk_size as i32,
            "uploadDate": DateTime::now(),
            "filename": final_filename.clone(),
        };
        if let Some(meta) = metadata_doc.clone() {
            fallback.insert("metadata", meta);
        }
        fallback
    };

    response_doc.insert("_id", file_id.to_hex());
    response_doc.insert("bucket", collection.clone());

    mongodb::bson::from_document::<Value>(response_doc).map_err(|e| e.to_string())
}

/// Resolve a file by id and return it as a data URI string.
pub async fn get_file_url(
    collection: String,
    id: String,
    mime_type: Option<String>,
) -> Result<String, String> {
    let client = get_client().await?;
    let db_name = get_db_name().await;
    let db = client.database(&db_name);

    let files_collection = db.collection::<Document>(&format!("{}.files", collection));
    let file_filter = match ObjectId::parse_str(&id) {
        Ok(oid) => doc! {
            "$or": [
                { "_id": oid },
                { "_id": id.clone() }
            ]
        },
        Err(_) => doc! { "_id": id.clone() },
    };

    let file_doc = files_collection
        .find_one(file_filter, None)
        .await
        .map_err(|e| e.to_string())?
        .ok_or_else(|| format!("File not found for id '{}'", id))?;

    let file_id = file_doc
        .get("_id")
        .cloned()
        .ok_or_else(|| "Stored file is missing _id".to_string())?;
    let filename = file_doc
        .get("filename")
        .and_then(Bson::as_str)
        .map(|name| name.to_string());
    let metadata = match file_doc.get("metadata") {
        Some(Bson::Document(doc)) => Some(doc),
        _ => None,
    };

    let bucket_options = GridFsBucketOptions::builder()
        .bucket_name(collection)
        .build();
    let bucket = db.gridfs_bucket(bucket_options);
    let mut stream = bucket
        .open_download_stream(file_id)
        .await
        .map_err(|e| e.to_string())?;

    let mut bytes = Vec::new();
    stream
        .read_to_end(&mut bytes)
        .await
        .map_err(|e| e.to_string())?;

    let resolved_mime = resolve_mime_type(filename.as_deref(), metadata, mime_type);
    let encoded = BASE64_STANDARD.encode(bytes);

    Ok(format!("data:{};base64,{}", resolved_mime, encoded))
}

fn resolve_mime_type(
    filename: Option<&str>,
    metadata: Option<&Document>,
    explicit_mime_type: Option<String>,
) -> String {
    if let Some(explicit) = explicit_mime_type {
        if !explicit.trim().is_empty() {
            return explicit;
        }
    }

    if let Some(meta) = metadata {
        if let Some(Bson::String(content_type)) = meta.get("contentType") {
            return content_type.clone();
        }
        if let Some(Bson::String(content_type)) = meta.get("mimeType") {
            return content_type.clone();
        }
    }

    if let Some(name) = filename {
        if let Some(ext) = Path::new(name).extension().and_then(|value| value.to_str()) {
            match ext.to_ascii_lowercase().as_str() {
                "png" => return "image/png".to_string(),
                "jpg" | "jpeg" => return "image/jpeg".to_string(),
                "gif" => return "image/gif".to_string(),
                "webp" => return "image/webp".to_string(),
                "svg" => return "image/svg+xml".to_string(),
                "bmp" => return "image/bmp".to_string(),
                "tif" | "tiff" => return "image/tiff".to_string(),
                "pdf" => return "application/pdf".to_string(),
                "txt" => return "text/plain".to_string(),
                "json" => return "application/json".to_string(),
                _ => {}
            }
        }
    }

    "application/octet-stream".to_string()
}