use super::system::AppState;
use crate::{
error::DbError,
storage::query_cache,
sync::{LogEntry, Operation},
};
use axum::{
body::Bytes,
extract::{Path, State},
response::Json,
};
use serde::Deserialize;
use serde_json::Value;
#[derive(Deserialize)]
pub struct CreateUploadRequest {
pub file_name: Option<String>,
pub mime_type: Option<String>,
pub total_size: u64,
pub chunk_size: Option<u32>,
}
pub async fn create_upload_session(
State(state): State<AppState>,
Path((db_name, coll_name)): Path<(String, String)>,
Json(body): Json<CreateUploadRequest>,
) -> Result<Json<Value>, DbError> {
let database = state.storage.get_database(&db_name)?;
match database.get_collection(&coll_name) {
Ok(coll) => {
if coll.get_type() != "blob" {
return Err(DbError::BadRequest(format!(
"Collection '{}' is not a blob collection.",
coll_name
)));
}
}
Err(DbError::CollectionNotFound(_)) => {
tracing::info!("Auto-creating blob collection {}/{}", db_name, coll_name);
database.create_collection(coll_name.clone(), Some("blob".to_string()))?;
}
Err(e) => return Err(e),
}
if body.total_size == 0 {
return Err(DbError::BadRequest(
"total_size must be greater than 0".to_string(),
));
}
let info = state.upload_session_store.create(
db_name,
coll_name,
body.file_name,
body.mime_type,
body.total_size,
body.chunk_size,
)?;
Ok(Json(serde_json::json!({
"upload_id": info.upload_id,
"blob_key": info.blob_key,
"chunk_size": info.chunk_size,
"total_chunks": info.total_chunks,
})))
}
pub async fn upload_chunk(
State(state): State<AppState>,
Path((_db_name, _coll_name, upload_id, chunk_index)): Path<(String, String, String, u32)>,
body: Bytes,
) -> Result<Json<Value>, DbError> {
let (db_name, coll_name, total_chunks, received_count, missing) = {
let mut session = state
.upload_session_store
.get_mut(&upload_id)
.ok_or_else(|| {
DbError::DocumentNotFound("Upload session not found or expired".to_string())
})?;
if chunk_index >= session.total_chunks {
return Err(DbError::BadRequest(format!(
"chunk_index {} out of range (total_chunks: {})",
chunk_index, session.total_chunks
)));
}
let database = state.storage.get_database(&session.db_name)?;
let collection = database.get_collection(&session.collection_name)?;
collection.put_blob_chunk_tmp(&upload_id, chunk_index, &body)?;
if !session.received_chunks[chunk_index as usize] {
session.bytes_received += body.len() as u64;
}
session.received_chunks[chunk_index as usize] = true;
session.last_activity = std::time::Instant::now();
let received_count = session.received_chunks.iter().filter(|&&r| r).count() as u32;
let missing: Vec<u32> = session
.received_chunks
.iter()
.enumerate()
.filter(|(_, &r)| !r)
.map(|(i, _)| i as u32)
.collect();
(
session.db_name.clone(),
session.collection_name.clone(),
session.total_chunks,
received_count,
missing,
)
};
Ok(Json(serde_json::json!({
"upload_id": upload_id,
"chunk_index": chunk_index,
"received_chunks": received_count,
"total_chunks": total_chunks,
"complete": missing.is_empty(),
"missing_chunks": missing,
"db": db_name,
"collection": coll_name,
})))
}
pub async fn get_upload_status(
State(state): State<AppState>,
Path((_db_name, _coll_name, upload_id)): Path<(String, String, String)>,
) -> Result<Json<Value>, DbError> {
let session = state.upload_session_store.get(&upload_id).ok_or_else(|| {
DbError::DocumentNotFound("Upload session not found or expired".to_string())
})?;
let received_count = session.received_chunks.iter().filter(|&&r| r).count() as u32;
let missing: Vec<u32> = session
.received_chunks
.iter()
.enumerate()
.filter(|(_, &r)| !r)
.map(|(i, _)| i as u32)
.collect();
Ok(Json(serde_json::json!({
"upload_id": upload_id,
"blob_key": session.blob_key,
"db": session.db_name,
"collection": session.collection_name,
"file_name": session.file_name,
"mime_type": session.mime_type,
"total_size": session.total_size,
"chunk_size": session.chunk_size,
"total_chunks": session.total_chunks,
"received_chunks": received_count,
"bytes_received": session.bytes_received,
"complete": missing.is_empty(),
"missing_chunks": missing,
})))
}
pub async fn abort_upload(
State(state): State<AppState>,
Path((_db_name, _coll_name, upload_id)): Path<(String, String, String)>,
) -> Result<Json<Value>, DbError> {
let session = state
.upload_session_store
.remove(&upload_id)
.ok_or_else(|| {
DbError::DocumentNotFound("Upload session not found or expired".to_string())
})?;
if let Ok(db) = state.storage.get_database(&session.db_name) {
if let Ok(coll) = db.get_collection(&session.collection_name) {
if let Err(e) = coll.delete_upload_chunks(&upload_id) {
tracing::warn!(
"Failed to clean up temp chunks for aborted upload {}: {}",
upload_id,
e
);
}
}
}
Ok(Json(serde_json::json!({
"upload_id": upload_id,
"status": "aborted",
})))
}
pub async fn complete_upload(
State(state): State<AppState>,
Path((_db_name, _coll_name, upload_id)): Path<(String, String, String)>,
) -> Result<Json<Value>, DbError> {
let (db_name, coll_name, blob_key, file_name, mime_type, total_size, total_chunks) = {
let session = state.upload_session_store.get(&upload_id).ok_or_else(|| {
DbError::DocumentNotFound("Upload session not found or expired".to_string())
})?;
let missing: Vec<u32> = session
.received_chunks
.iter()
.enumerate()
.filter(|(_, &r)| !r)
.map(|(i, _)| i as u32)
.collect();
if !missing.is_empty() {
return Err(DbError::ConflictError(
serde_json::json!({
"error": "Upload incomplete",
"missing_chunks": missing,
"received": session.total_chunks - missing.len() as u32,
"total": session.total_chunks,
})
.to_string(),
));
}
(
session.db_name.clone(),
session.collection_name.clone(),
session.blob_key.clone(),
session.file_name.clone(),
session.mime_type.clone(),
session.total_size,
session.total_chunks,
)
};
let database = state.storage.get_database(&db_name)?;
let collection = database.get_collection(&coll_name)?;
collection.finalize_blob_upload(&upload_id, &blob_key, total_chunks)?;
let mut metadata = serde_json::Map::new();
metadata.insert("_key".to_string(), Value::String(blob_key.clone()));
if let Some(name) = file_name {
metadata.insert("name".to_string(), Value::String(name));
}
if let Some(mt) = mime_type {
metadata.insert("type".to_string(), Value::String(mt));
}
metadata.insert("size".to_string(), Value::Number(total_size.into()));
metadata.insert("chunks".to_string(), Value::Number(total_chunks.into()));
metadata.insert(
"created".to_string(),
Value::String(chrono::Utc::now().to_rfc3339()),
);
let doc_value = Value::Object(metadata);
if collection.get_type() == "blob" {
if let Some(ref coordinator) = state.shard_coordinator {
let mut chunks_buffer: Vec<(u32, Vec<u8>)> = Vec::new();
for i in 0..total_chunks {
if let Some(data) = collection.get_blob_chunk(&blob_key, i)? {
chunks_buffer.push((i, data));
}
}
super::blobs::distribute_blob_chunks_across_cluster(
coordinator,
&db_name,
&coll_name,
&blob_key,
&chunks_buffer,
&doc_value,
&state.storage,
)
.await?;
} else {
collection.insert(doc_value.clone())?;
}
} else {
collection.insert(doc_value.clone())?;
}
if let Some(ref log) = state.replication_log {
let entry = LogEntry {
sequence: 0,
node_id: "".to_string(),
database: db_name.clone(),
collection: coll_name.clone(),
operation: Operation::Insert,
key: blob_key.clone(),
data: serde_json::to_vec(&doc_value).ok(),
timestamp: chrono::Utc::now().timestamp_millis() as u64,
origin_sequence: None,
};
let _ = log.append(entry);
}
state.upload_session_store.remove(&upload_id);
query_cache::get_query_cache().invalidate_collection(&coll_name);
Ok(Json(doc_value))
}