feat(core): wire up attachment text extraction in IMAP sync pipeline

This commit is contained in:
rustmailer
2026-05-24 19:52:28 +08:00
parent ec3e842bbb
commit f9c2fc77ff
5 changed files with 148 additions and 45 deletions
+113 -41
View File
@@ -197,33 +197,37 @@ async fn extract_envelope_core(
let attachment_docs: Vec<TantivyDocument> = attachments
.iter()
.filter(|a| !a.inline || a.content_id.is_none())
.map(|a| AttachmentModel {
id: Uuid::new_v4().to_string(),
envelope_id: envelope_id.clone(),
account_id,
account_email: None,
mailbox_id,
mailbox_name: None,
subject: subject.clone(),
content_hash: a.content_hash.clone(),
from: from.clone(),
date,
ingest_at: now,
size: a.size as u64,
ext: a.get_extension(),
category: a.get_category().to_string(),
content_type: a.file_type.clone(),
shard_id: 0,
text: None,
has_text: false,
is_ocr: false,
page_count: None,
is_indexed: false,
is_message: a.is_message,
name: a.filename.clone(),
tags: None,
auto_tags: None,
}).map(|a|a.into_document())
.map(|a| {
let has_text = a.extracted_text.is_some();
AttachmentModel {
id: Uuid::new_v4().to_string(),
envelope_id: envelope_id.clone(),
account_id,
account_email: None,
mailbox_id,
mailbox_name: None,
subject: subject.clone(),
content_hash: a.content_hash.clone(),
from: from.clone(),
date,
ingest_at: now,
size: a.size as u64,
ext: a.get_extension(),
category: a.get_category().to_string(),
content_type: a.file_type.clone(),
shard_id: 0,
text: a.extracted_text.clone(),
has_text,
is_ocr: a.extracted_is_ocr,
page_count: a.extracted_page_count.map(|n| n as u64),
is_indexed: has_text,
is_message: a.is_message,
name: a.filename.clone(),
tags: None,
auto_tags: None,
}
})
.map(|a| a.into_document())
.collect();
let envelope = Envelope {
@@ -397,6 +401,16 @@ pub async fn detach_and_store_attachments(
ranges.sort_by(|a, b| b.0.cmp(&a.0));
let mut attachments = Vec::with_capacity(ranges.len());
// Collect candidates for text extraction (non-inline, known document types).
struct TextCandidate {
content_hash: String,
file_type: String,
ext: String,
bytes: Vec<u8>,
}
let mut text_candidates: Vec<TextCandidate> = Vec::new();
for (raw_start, raw_end, att) in ranges {
// Step 2: Extract raw bytes and store them as standalone documents
let raw_bytes = &original_body[raw_start..raw_end];
@@ -411,30 +425,88 @@ pub async fn detach_and_store_attachments(
let p_bytes = placeholder.as_bytes();
stripped_eml.splice(raw_start..raw_end, p_bytes.iter().cloned());
let inline = att
.content_disposition()
.map(|d| d.is_inline())
.unwrap_or(false);
let file_type = att
.content_type()
.map(|ct| {
format!(
"{}/{}",
ct.c_type.as_ref(),
ct.c_subtype.as_deref().unwrap_or("")
)
})
.unwrap_or_else(|| "application/octet-stream".to_string());
let has_cid = att.content_id().is_some();
let ext = att
.attachment_name()
.and_then(|n| {
std::path::Path::new(&n)
.extension()
.and_then(|e| e.to_str())
.map(|s| s.to_ascii_lowercase())
})
.unwrap_or_default();
if !inline || !has_cid {
let decoded_len = att.contents().len();
if decoded_len <= crate::ext::text_extractor::MAX_EXTRACT_BYTES
&& crate::ext::text_extractor::should_try_extract(&file_type, &ext)
{
text_candidates.push(TextCandidate {
content_hash: content_hash.clone(),
file_type: file_type.clone(),
ext: ext.clone(),
bytes: att.contents().to_vec(),
});
}
}
let info = AttachmentInfo {
filename: att.attachment_name().map(|n| n.to_string()),
size: att.contents().len(),
inline: att
.content_disposition()
.map(|d| d.is_inline())
.unwrap_or(false),
file_type: att
.content_type()
.map(|ct| {
format!(
"{}/{}",
ct.c_type.as_ref(),
ct.c_subtype.as_deref().unwrap_or("")
)
})
.unwrap_or_else(|| "application/octet-stream".to_string()),
inline,
file_type,
content_id: att.content_id().map(|id| id.to_string()),
content_hash: content_hash.clone(),
is_message: att.is_message(),
extracted_text: None,
extracted_page_count: None,
extracted_is_ocr: false,
};
attachment_infos.push(info);
}
// Run text extraction in a single spawn_blocking batch.
if !text_candidates.is_empty() {
if let Ok(mut extracted_map) = tokio::task::spawn_blocking(move || {
let mut map: std::collections::HashMap<
String,
(String, Option<u32>, bool),
> = std::collections::HashMap::new();
for c in text_candidates {
if let Some(r) =
crate::ext::text_extractor::extract_text(&c.file_type, &c.ext, &c.bytes)
{
map.insert(c.content_hash, (r.text, r.page_count, r.is_ocr));
}
}
map
})
.await
{
for info in &mut attachment_infos {
if let Some((text, pages, is_ocr)) = extracted_map.remove(&info.content_hash) {
info.extracted_text = Some(text);
info.extracted_page_count = pages;
info.extracted_is_ocr = is_ocr;
}
}
}
}
// Step 4: Store the final stripped EML content
BLOB_MANAGER
.queue(DetachedEmail {
+15
View File
@@ -53,6 +53,21 @@ pub fn set_extractor(extractor: Box<dyn AttachmentTextExtractor>) {
*EXTRACTOR.write().unwrap() = extractor;
}
/// Attachments larger than this are skipped (10 MiB). Avoids excessive memory
/// and CPU cost for huge files whose text is rarely useful for search.
pub const MAX_EXTRACT_BYTES: usize = 10 * 1024 * 1024;
/// Quick pre-filter: returns true for file types where text extraction may
/// produce useful results. Avoids cloning attachment bytes for images, videos,
/// archives, etc. when no registered extractor would handle them.
pub fn should_try_extract(content_type: &str, ext: &str) -> bool {
matches!(
ext,
"pdf" | "doc" | "docx" | "xls" | "xlsx" | "ppt" | "pptx"
| "txt" | "rtf" | "odt" | "ods" | "odp"
) || content_type.starts_with("text/")
}
/// Called by the attachment pipeline during IMAP sync.
/// The caller should wrap this in spawn_blocking for CPU-bound extraction.
pub fn extract_text(content_type: &str, ext: &str, bytes: &[u8]) -> Option<ExtractedText> {
+15 -2
View File
@@ -47,6 +47,13 @@ pub struct AttachmentInfo {
/// Hash of the content.
pub content_hash: String,
pub is_message: bool,
/// Text extracted from the attachment body (Pro/Enterprise feature).
/// Populated during IMAP sync; None for inline attachments and unsupported file types.
pub extracted_text: Option<String>,
/// Page count reported by the extractor, if any.
pub extracted_page_count: Option<u32>,
/// Whether the extracted text came from OCR.
pub extracted_is_ocr: bool,
}
impl AttachmentInfo {
@@ -222,13 +229,16 @@ pub fn retrieve_email_content(
let is_message = attachment.is_message();
let content_hash = compute_content_hash(attachment.contents());
attachments.push(AttachmentInfo {
filename: filename.or(Some(content_hash.clone())), // Fallback to content_hash as the default filename if it is not provided.
filename: filename.or(Some(content_hash.clone())),
size: attachment.contents().len(),
inline,
file_type,
is_message,
content_hash,
content_id: attachment.content_id().map(Into::into),
extracted_text: None,
extracted_page_count: None,
extracted_is_ocr: false,
});
}
let mut has_remote_content = false;
@@ -320,13 +330,16 @@ pub fn retrieve_nested_eml_content(
filename: attachment
.attachment_name()
.map(|n| n.to_string())
.or(Some(content_hash.clone())), // Fallback to content_hash as the default filename if it is not provided.
.or(Some(content_hash.clone())),
size: attachment.contents().len(),
inline: is_inline,
file_type,
content_hash,
is_message: attachment.is_message(),
content_id: cid.map(Into::into),
extracted_text: None,
extracted_page_count: None,
extracted_is_ocr: false,
});
}
+3
View File
@@ -114,6 +114,9 @@ pub fn detach_attachments_standalone(
content_id: att.content_id().map(|id| id.to_string()),
content_hash,
is_message: att.is_message(),
extracted_text: None,
extracted_page_count: None,
extracted_is_ocr: false,
});
}