//
// Copyright (c) 2025-2026 rustmailer.com (https://rustmailer.com)
//
// This file is part of the Bichon Email Archiving Project
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see .
use chrono::{NaiveDateTime, Utc};
use duckdb::{params, types::Value, DuckdbConnectionManager};
use refinery::Runner;
use std::{
collections::{HashMap, HashSet},
path::PathBuf,
sync::OnceLock,
};
use crate::{
modules::{
account::migration::AccountModel,
context::Initialize,
dashboard::{DashboardStats, Group, LargestEmail, TimeBucket},
duckdb::{build::build_record_batch, refinery::DuckDBConnection},
error::{code::ErrorCode, BichonResult},
indexer::{
attachment::ATTACHMENT_INDEX_MANAGER, eml::EML_INDEX_MANAGER, envelope::Envelope,
manager::ENVELOPE_INDEX_MANAGER,
},
message::{
attachment::AttachmentMetadata,
content::{AttachmentDetail, AttachmentInfo},
search::{SearchFilter, SortBy},
tags::TagCount,
},
rest::response::DataPage,
settings::{cli::SETTINGS, dir::DATA_DIR_MANAGER},
},
raise_error,
};
pub type DuckDBConn = r2d2::PooledConnection;
pub static DUCKDBMANAGER: OnceLock = OnceLock::new();
pub mod duckdb_tables {
refinery::embed_migrations!("src/modules/duckdb/migrations");
}
pub fn duckdb() -> BichonResult<&'static DuckDBManager> {
DUCKDBMANAGER.get().ok_or_else(|| {
raise_error!(
"DuckDB manager is not initialized".into(),
ErrorCode::InternalError
)
})
}
fn debug_sql(sql: &str, args: &[Value]) -> String {
let mut result = String::new();
let mut parts = sql.split('?');
for (i, part) in parts.by_ref().enumerate() {
result.push_str(part);
if i < args.len() {
result.push_str(&format!("{:?}", args[i]));
}
}
result
}
pub struct DuckDBManager {
pool: r2d2::Pool,
}
impl Initialize for DuckDBManager {
async fn initialize() -> BichonResult<()> {
tracing::debug!("Initializing databases");
if !&DATA_DIR_MANAGER.envelope_dir.exists() {
std::fs::create_dir_all(&DATA_DIR_MANAGER.envelope_dir)
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
}
let pool = init(
&DATA_DIR_MANAGER.envelope_dir.join("envelopes.db"),
duckdb_tables::migrations::runner(),
)?;
let _ = DUCKDBMANAGER.set(DuckDBManager { pool });
Ok(())
}
}
impl DuckDBManager {
pub fn conn(&self) -> BichonResult {
Ok(self
.pool
.get()
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?)
}
pub fn shutdown(&self) -> BichonResult<()> {
self.conn()?
.execute("FORCE CHECKPOINT", [])
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
tracing::info!("Shutting down");
Ok(())
}
pub fn validate_regex(&self, pattern: &str) -> BichonResult<()> {
let conn = self.conn()?;
let check_sql = "SELECT regexp_matches('', ?)";
if let Err(e) = conn
.prepare(check_sql)
.and_then(|mut stmt| stmt.execute([pattern]))
{
return Err(raise_error!(
format!("Invalid Regular Expression for DuckDB: {}", e).into(),
ErrorCode::InvalidParameter
));
}
Ok(())
}
pub fn delete_account_envelopes_with_orphans(
&self,
account_id: u64,
) -> BichonResult> {
let mut conn = self.conn()?;
let tx = conn
.transaction()
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?;
let orphan_sql = r#"
WITH target_hashes AS (
SELECT content_hash FROM envelopes WHERE account_id = ?
UNION
SELECT content_hash FROM envelope_attachments WHERE account_id = ?
),
active_hashes AS (
SELECT content_hash FROM envelopes WHERE account_id != ?
UNION
SELECT content_hash FROM envelope_attachments WHERE account_id != ?
)
SELECT content_hash FROM target_hashes
EXCEPT
SELECT content_hash FROM active_hashes
"#;
let mut stmt = tx
.prepare(orphan_sql)
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?;
let orphan_hashes: Vec = stmt
.query_map([account_id, account_id, account_id, account_id], |row| {
row.get::<_, String>(0)
})
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?
.collect::, _>>()
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?;
tx.execute(
"DELETE FROM envelope_attachments WHERE account_id = ?",
[account_id],
)
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?;
let count = tx
.execute("DELETE FROM envelopes WHERE account_id = ?", [account_id])
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?;
tx.commit()
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?;
if count > 0 {
tracing::info!(
"Account {} data cleared. Deleted {} envelopes, {} orphan hashes identified.",
account_id,
count,
orphan_hashes.len()
);
}
Ok(orphan_hashes)
}
pub fn delete_mailbox_envelopes_with_orphans(
&self,
account_id: u64,
mailbox_ids: Vec,
) -> BichonResult> {
if mailbox_ids.is_empty() {
return Ok(vec![]);
}
let mut conn = self.conn()?;
let tx = conn
.transaction()
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
let placeholders = mailbox_ids
.iter()
.map(|_| "?")
.collect::>()
.join(", ");
let mut sql_params: Vec = vec![account_id.into()];
sql_params.extend(mailbox_ids.iter().map(|&id| duckdb::types::Value::from(id)));
let param_iter = duckdb::params_from_iter(sql_params);
let orphan_sql = format!(
r#"
WITH target_hashes AS (
SELECT content_hash FROM envelopes WHERE account_id = ? AND mailbox_id IN ({0})
UNION
SELECT content_hash FROM envelope_attachments WHERE account_id = ? AND mailbox_id IN ({0})
),
active_hashes AS (
SELECT content_hash FROM envelopes WHERE NOT (account_id = ? AND mailbox_id IN ({0}))
UNION
SELECT content_hash FROM envelope_attachments WHERE NOT (account_id = ? AND mailbox_id IN ({0}))
)
SELECT content_hash FROM target_hashes
EXCEPT
SELECT content_hash FROM active_hashes
"#,
placeholders
);
let mut query_params: Vec = Vec::new();
for _ in 0..4 {
query_params.push(account_id.into());
query_params.extend(mailbox_ids.iter().map(|&id| duckdb::types::Value::from(id)));
}
let mut stmt = tx
.prepare(&orphan_sql)
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?;
let orphan_hashes: Vec = stmt
.query_map(duckdb::params_from_iter(query_params), |row| {
row.get::<_, String>(0)
})
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?
.collect::, _>>()
.map_err(|e| raise_error!(e.to_string(), ErrorCode::InternalError))?;
let delete_att_sql = format!(
"DELETE FROM envelope_attachments WHERE account_id = ? AND mailbox_id IN ({})",
placeholders
);
tx.execute(&delete_att_sql, param_iter.clone())
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
let delete_env_sql = format!(
"DELETE FROM envelopes WHERE account_id = ? AND mailbox_id IN ({})",
placeholders
);
let count = tx
.execute(&delete_env_sql, param_iter)
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
tx.commit()
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
if count > 0 {
tracing::info!(
"Deleted {} emails. Found {} orphan hashes to clean up.",
count,
orphan_hashes.len()
);
}
Ok(orphan_hashes)
}
pub fn append_envelopes_with_attachments(
&self,
items: &[(Envelope, Vec)],
) -> BichonResult<()> {
let mut conn = self.conn()?;
let tx = conn
.transaction()
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
{
let mut env_appender = tx
.appender("envelopes")
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
let mut att_appender = tx
.appender("envelope_attachments")
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
for (env, atts) in items {
env_appender
.append_record_batch(build_record_batch(std::slice::from_ref(env)))
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
for att in atts {
att_appender
.append_row(params![
env.id,
env.account_id,
env.mailbox_id,
att.filename,
att.is_message,
att.inline,
att.content_id,
att.get_extension(),
att.get_category(),
att.file_type.to_ascii_lowercase(),
att.size as u64,
att.content_hash,
0
])
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
}
}
env_appender
.flush()
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
att_appender
.flush()
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
}
tx.commit()
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
Ok(())
}
pub fn total_emails(
&self,
accounts: Option>,
) -> BichonResult {
let conn = self.conn()?;
let total: u64 = match accounts {
None => conn
.query_row("SELECT COUNT(*) FROM envelopes", [], |row| row.get(0))
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?,
Some(accts) if !accts.is_empty() => {
let placeholders = accts.iter().map(|_| "?").collect::>().join(", ");
let sql = format!(
"SELECT COUNT(*) FROM envelopes WHERE account_id IN ({})",
placeholders
);
conn.query_row(&sql, duckdb::params_from_iter(accts), |row| row.get(0))
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?
}
_ => 0,
};
Ok(total)
}
pub fn num_messages_in_thread(&self, account_id: u64, thread_id: String) -> BichonResult {
let conn = self.conn()?;
let count: u64 = conn
.query_row(
"SELECT COUNT(*) FROM envelopes WHERE account_id = ? AND thread_id = ?;",
params![account_id, thread_id],
|row| row.get(0),
)
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
Ok(count)
}
pub fn num_messages_in_mailbox(&self, account_id: u64, mailbox_id: u64) -> BichonResult {
let conn = self.conn()?;
let count: u64 = conn
.query_row(
"SELECT COUNT(*) FROM envelopes WHERE account_id = ? AND mailbox_id = ?;",
params![account_id, mailbox_id],
|row| row.get(0),
)
.map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
Ok(count)
}
pub fn get_max_uid(&self, account_id: u64, mailbox_id: u64) -> BichonResult