Co-authored-by: aider (openai/andrew/openrouter/deepseek/deepseek-chat-v3.1) <aider@aider.chat>
254 lines
8.5 KiB
Rust
254 lines
8.5 KiB
Rust
use crate::services::error::CoreError;
|
|
use crate::services::item_service::ItemService;
|
|
use crate::services::types::{ItemWithContent, ItemWithMeta};
|
|
use rusqlite::Connection;
|
|
use std::collections::HashMap;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
use tokio::sync::Mutex;
|
|
|
|
/// An asynchronous wrapper around the `ItemService` for use in async contexts like the web server.
|
|
/// It uses `tokio::task::spawn_blocking` to run synchronous database and filesystem operations
|
|
/// on a dedicated thread pool, preventing them from blocking the async runtime.
|
|
#[allow(dead_code)]
|
|
pub struct AsyncItemService {
|
|
pub data_dir: PathBuf,
|
|
db: Arc<Mutex<Connection>>,
|
|
item_service: Arc<ItemService>,
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
impl AsyncItemService {
|
|
pub fn new(data_dir: PathBuf, db: Arc<Mutex<Connection>>, item_service: Arc<ItemService>) -> Self {
|
|
Self { data_dir, db, item_service }
|
|
}
|
|
|
|
pub async fn get_item(&self, id: i64) -> Result<ItemWithMeta, CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let conn = db.blocking_lock();
|
|
item_service.get_item(&conn, id)
|
|
})
|
|
.await
|
|
.unwrap() // Propagate panics from spawn_blocking
|
|
}
|
|
|
|
pub async fn get_item_content(&self, id: i64) -> Result<ItemWithContent, CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let conn = db.blocking_lock();
|
|
item_service.get_item_content(&conn, id)
|
|
})
|
|
.await
|
|
.unwrap()
|
|
}
|
|
|
|
pub async fn get_item_content_info(&self, id: i64) -> Result<(Vec<u8>, String, bool), CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let conn = db.blocking_lock();
|
|
item_service.get_item_content_info(&conn, id)
|
|
})
|
|
.await
|
|
.unwrap()
|
|
}
|
|
|
|
|
|
pub async fn stream_item_content_by_id(
|
|
&self,
|
|
item_id: i64,
|
|
allow_binary: bool,
|
|
offset: u64,
|
|
length: u64,
|
|
) -> Result<(std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<tokio_util::bytes::Bytes, std::io::Error>> + Send>>, String), CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
// Get item content
|
|
let content = tokio::task::spawn_blocking(move || {
|
|
let conn = db.blocking_lock();
|
|
let item_with_content = item_service.get_item_content(&conn, item_id)?;
|
|
Ok::<_, CoreError>(item_with_content.content)
|
|
})
|
|
.await
|
|
.unwrap()?;
|
|
|
|
// Get metadata to determine MIME type and binary status
|
|
let (mime_type, is_binary) = {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
tokio::task::spawn_blocking(move || {
|
|
let conn = db.blocking_lock();
|
|
let item_with_meta = item_service.get_item(&conn, item_id)?;
|
|
let metadata = item_with_meta.meta_as_map();
|
|
|
|
let mime_type = metadata
|
|
.get("mime_type")
|
|
.map(|s| s.to_string())
|
|
.unwrap_or_else(|| "application/octet-stream".to_string());
|
|
|
|
let is_binary = if let Some(binary_val) = metadata.get("binary") {
|
|
binary_val == "true"
|
|
} else {
|
|
crate::common::is_binary::is_binary(&content)
|
|
};
|
|
|
|
Ok::<_, CoreError>((mime_type, is_binary))
|
|
})
|
|
.await
|
|
.unwrap()?
|
|
};
|
|
|
|
// Check if content is binary when allow_binary is false
|
|
if !allow_binary && is_binary {
|
|
return Err(CoreError::InvalidInput("Binary content not allowed".to_string()));
|
|
}
|
|
|
|
// Create a stream that reads only the requested portion
|
|
let content_len = content.len() as u64;
|
|
|
|
// Apply offset and length constraints
|
|
let start = std::cmp::min(offset, content_len);
|
|
let end = if length > 0 {
|
|
std::cmp::min(start + length, content_len)
|
|
} else {
|
|
content_len
|
|
};
|
|
|
|
let stream = if start < content_len {
|
|
let chunk = tokio_util::bytes::Bytes::from(content[start as usize..end as usize].to_vec());
|
|
Box::pin(tokio_stream::iter(vec![Ok(chunk)]))
|
|
} else {
|
|
Box::pin(tokio_stream::iter(vec![]))
|
|
};
|
|
|
|
Ok((stream, mime_type))
|
|
}
|
|
|
|
pub async fn stream_item_content_by_id_with_metadata(
|
|
&self,
|
|
item_id: i64,
|
|
metadata: &HashMap<String, String>,
|
|
allow_binary: bool,
|
|
offset: u64,
|
|
length: u64,
|
|
) -> Result<(std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<tokio_util::bytes::Bytes, std::io::Error>> + Send>>, String), CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
// Get item content
|
|
let content = tokio::task::spawn_blocking(move || {
|
|
let conn = db.blocking_lock();
|
|
let item_with_content = item_service.get_item_content(&conn, item_id)?;
|
|
Ok::<_, CoreError>(item_with_content.content)
|
|
})
|
|
.await
|
|
.unwrap()?;
|
|
|
|
// Use provided metadata to determine MIME type and binary status
|
|
let mime_type = metadata
|
|
.get("mime_type")
|
|
.map(|s| s.to_string())
|
|
.unwrap_or_else(|| "application/octet-stream".to_string());
|
|
|
|
let is_binary = if let Some(binary_val) = metadata.get("binary") {
|
|
binary_val == "true"
|
|
} else {
|
|
crate::common::is_binary::is_binary(&content)
|
|
};
|
|
|
|
// Check if content is binary when allow_binary is false
|
|
if !allow_binary && is_binary {
|
|
return Err(CoreError::InvalidInput("Binary content not allowed".to_string()));
|
|
}
|
|
|
|
// Create a stream that reads only the requested portion
|
|
let content_len = content.len() as u64;
|
|
|
|
// Apply offset and length constraints
|
|
let start = std::cmp::min(offset, content_len);
|
|
let end = if length > 0 {
|
|
std::cmp::min(start + length, content_len)
|
|
} else {
|
|
content_len
|
|
};
|
|
|
|
let stream = if start < content_len {
|
|
let chunk = tokio_util::bytes::Bytes::from(content[start as usize..end as usize].to_vec());
|
|
Box::pin(tokio_stream::iter(vec![Ok(chunk)]))
|
|
} else {
|
|
Box::pin(tokio_stream::iter(vec![]))
|
|
};
|
|
|
|
Ok((stream, mime_type))
|
|
}
|
|
|
|
pub async fn find_item(
|
|
&self,
|
|
ids: Vec<i64>,
|
|
tags: Vec<String>,
|
|
meta: HashMap<String, String>,
|
|
) -> Result<ItemWithMeta, CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let conn = db.blocking_lock();
|
|
item_service.find_item(&conn, &ids, &tags, &meta)
|
|
})
|
|
.await
|
|
.unwrap()
|
|
}
|
|
|
|
pub async fn list_items(
|
|
&self,
|
|
tags: Vec<String>,
|
|
meta: HashMap<String, String>,
|
|
) -> Result<Vec<ItemWithMeta>, CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let conn = db.blocking_lock();
|
|
item_service.list_items(&conn, &tags, &meta)
|
|
})
|
|
.await
|
|
.unwrap()
|
|
}
|
|
|
|
pub async fn delete_item(&self, id: i64) -> Result<(), CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let mut conn = db.blocking_lock();
|
|
item_service.delete_item(&mut conn, id)
|
|
})
|
|
.await
|
|
.unwrap()
|
|
}
|
|
|
|
pub async fn save_item_from_mcp(
|
|
&self,
|
|
content: Vec<u8>,
|
|
tags: Vec<String>,
|
|
metadata: HashMap<String, String>,
|
|
) -> Result<ItemWithMeta, CoreError> {
|
|
let db = self.db.clone();
|
|
let item_service = self.item_service.clone();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let mut conn = db.blocking_lock();
|
|
item_service.save_item_from_mcp(&content, &tags, &metadata, &mut conn)
|
|
})
|
|
.await
|
|
.unwrap()
|
|
}
|
|
}
|