Files
email-unlimit/server/src/db.rs
T

527 lines
21 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
use crate::config::database::{DatabaseConfig, DatabaseType};
use sea_orm::{
ConnectOptions, ConnectionTrait, Database, DatabaseConnection, DbBackend, EntityName,
EntityTrait, Schema, Statement,
};
use std::time::Duration;
/// 数据库连接池(SeaORM 统一接口)
pub type DbPool = DatabaseConnection;
/// 创建数据库连接池
pub async fn create_pool(config: &DatabaseConfig) -> anyhow::Result<DbPool> {
let url = config
.build_url()
.map_err(|e| anyhow::anyhow!("数据库配置错误: {}", e))?;
tracing::debug!("数据库连接 URL: {}", url);
let mut opt = ConnectOptions::new(&url);
opt.max_connections(config.max_connections)
.min_connections(1)
.connect_timeout(Duration::from_secs(8))
.idle_timeout(Duration::from_secs(8))
.max_lifetime(Duration::from_secs(7200))
.sqlx_logging(true);
let pool = Database::connect(opt)
.await
.map_err(|e| anyhow::anyhow!("数据库连接失败: {}", e))?;
tracing::info!("已连接到数据库: {}", sanitize_url(&url));
Ok(pool)
}
/// 隐藏 URL 中的敏感信息(用于日志输出)
fn sanitize_url(url: &str) -> String {
// 隐藏密码:mysql://user:password@host -> mysql://user:***@host
if let Some(at_pos) = url.find('@') {
if let Some(scheme_end) = url.find("://") {
if scheme_end < at_pos {
return format!("{}***@{}", &url[..scheme_end + 3], &url[at_pos + 1..]);
}
}
}
url.to_string()
}
/// 健康检查(保持向后兼容)
pub async fn health_check(pool: &DbPool) -> anyhow::Result<()> {
// 使用官方推荐的 ping 方法
pool.ping()
.await
.map_err(|e| anyhow::anyhow!("数据库健康检查失败: {}", e))
}
/// 初始化数据库和表结构
/// 每次启动时检查数据库和表是否存在,不存在则创建
pub async fn init_database(config: &DatabaseConfig) -> anyhow::Result<DatabaseConnection> {
match config.database_type {
DatabaseType::MySQL => {
init_mysql_database(config).await?;
}
DatabaseType::PostgreSQL => {
init_postgresql_database(config).await?;
}
DatabaseType::SQLite => {
// 确保 SQLite 数据库文件的目录存在
init_sqlite_database(config).await?;
}
}
// 连接到数据库
let pool = create_pool(config).await?;
// 创建表
create_tables(&pool).await?;
migrate_existing_tables(&pool).await?;
create_indexes(&pool).await?;
Ok(pool)
}
async fn create_indexes(db: &DatabaseConnection) -> anyhow::Result<()> {
let backend = db.get_database_backend();
let statements = match backend {
DbBackend::MySql => vec![
"CREATE INDEX idx_users_email ON users(email)",
"CREATE INDEX idx_users_status_created ON users(status, created_at)",
"CREATE INDEX idx_email_logs_user_created ON email_logs(user_id, created_at)",
"CREATE INDEX idx_email_logs_status_created ON email_logs(status, created_at)",
"CREATE INDEX idx_mailboxes_user ON mailboxes(user_id)",
"CREATE INDEX idx_mailboxes_status ON mailboxes(status)",
"CREATE INDEX idx_mailboxes_expires_at ON mailboxes(expires_at)",
"CREATE INDEX idx_mailboxes_status_expires ON mailboxes(status, expires_at)",
"CREATE INDEX idx_emails_mailbox ON emails(mailbox_id)",
"CREATE INDEX idx_emails_user ON emails(user_id)",
"CREATE INDEX idx_emails_received_at ON emails(received_at)",
"CREATE INDEX idx_emails_expires_at ON emails(expires_at)",
"CREATE INDEX idx_emails_status_received ON emails(status, received_at)",
"CREATE INDEX idx_emails_source_received ON emails(source_ip, received_at)",
"CREATE INDEX idx_email_attachments_email ON email_attachments(email_id)",
"CREATE INDEX idx_credit_transactions_user_created ON credit_transactions(user_id, created_at)",
"CREATE INDEX idx_credit_transactions_reason_created ON credit_transactions(reason, created_at)",
"CREATE INDEX idx_credit_rule_changes_created ON credit_rule_changes(created_at)",
"CREATE INDEX idx_blocked_senders_value ON blocked_senders(value)",
"CREATE INDEX idx_blocked_ips_ip ON blocked_ips(ip)",
"CREATE INDEX idx_greylist_triplet ON greylist_entries(sender_ip, mail_from, rcpt_to)",
"CREATE INDEX idx_audit_logs_created ON audit_logs(created_at)",
"CREATE INDEX idx_audit_logs_source_ip ON audit_logs(source_ip)",
"CREATE INDEX idx_audit_logs_event_type ON audit_logs(event_type)",
"CREATE INDEX idx_audit_logs_action_created ON audit_logs(action, created_at)",
"CREATE INDEX idx_audit_logs_operator_created ON audit_logs(operator_id, created_at)",
"CREATE INDEX idx_abuse_events_rule ON abuse_events(rule_id)",
"CREATE INDEX idx_abuse_events_created ON abuse_events(created_at)",
],
_ => vec![
"CREATE INDEX IF NOT EXISTS idx_users_email ON users(email)",
"CREATE INDEX IF NOT EXISTS idx_users_status_created ON users(status, created_at)",
"CREATE INDEX IF NOT EXISTS idx_email_logs_user_created ON email_logs(user_id, created_at)",
"CREATE INDEX IF NOT EXISTS idx_email_logs_status_created ON email_logs(status, created_at)",
"CREATE INDEX IF NOT EXISTS idx_mailboxes_user ON mailboxes(user_id)",
"CREATE INDEX IF NOT EXISTS idx_mailboxes_status ON mailboxes(status)",
"CREATE INDEX IF NOT EXISTS idx_mailboxes_expires_at ON mailboxes(expires_at)",
"CREATE INDEX IF NOT EXISTS idx_mailboxes_status_expires ON mailboxes(status, expires_at)",
"CREATE INDEX IF NOT EXISTS idx_emails_mailbox ON emails(mailbox_id)",
"CREATE INDEX IF NOT EXISTS idx_emails_user ON emails(user_id)",
"CREATE INDEX IF NOT EXISTS idx_emails_received_at ON emails(received_at)",
"CREATE INDEX IF NOT EXISTS idx_emails_expires_at ON emails(expires_at)",
"CREATE INDEX IF NOT EXISTS idx_emails_status_received ON emails(status, received_at)",
"CREATE INDEX IF NOT EXISTS idx_emails_source_received ON emails(source_ip, received_at)",
"CREATE INDEX IF NOT EXISTS idx_email_attachments_email ON email_attachments(email_id)",
"CREATE INDEX IF NOT EXISTS idx_credit_transactions_user_created ON credit_transactions(user_id, created_at)",
"CREATE INDEX IF NOT EXISTS idx_credit_transactions_reason_created ON credit_transactions(reason, created_at)",
"CREATE INDEX IF NOT EXISTS idx_credit_rule_changes_created ON credit_rule_changes(created_at)",
"CREATE INDEX IF NOT EXISTS idx_blocked_senders_value ON blocked_senders(value)",
"CREATE INDEX IF NOT EXISTS idx_blocked_ips_ip ON blocked_ips(ip)",
"CREATE INDEX IF NOT EXISTS idx_greylist_triplet ON greylist_entries(sender_ip, mail_from, rcpt_to)",
"CREATE INDEX IF NOT EXISTS idx_audit_logs_created ON audit_logs(created_at)",
"CREATE INDEX IF NOT EXISTS idx_audit_logs_source_ip ON audit_logs(source_ip)",
"CREATE INDEX IF NOT EXISTS idx_audit_logs_event_type ON audit_logs(event_type)",
"CREATE INDEX IF NOT EXISTS idx_audit_logs_action_created ON audit_logs(action, created_at)",
"CREATE INDEX IF NOT EXISTS idx_audit_logs_operator_created ON audit_logs(operator_id, created_at)",
"CREATE INDEX IF NOT EXISTS idx_abuse_events_rule ON abuse_events(rule_id)",
"CREATE INDEX IF NOT EXISTS idx_abuse_events_created ON abuse_events(created_at)",
],
};
for statement in statements {
if let Err(error) = db.execute(Statement::from_string(backend, statement)).await {
let message = error.to_string().to_lowercase();
if !message.contains("duplicate") && !message.contains("already exists") {
return Err(anyhow::anyhow!("创建索引失败: {error}"));
}
}
}
Ok(())
}
async fn migrate_existing_tables(db: &DatabaseConnection) -> anyhow::Result<()> {
let backend = db.get_database_backend();
let datetime_type = match backend {
DbBackend::Postgres => "TIMESTAMP NULL",
_ => "DATETIME NULL",
};
let varchar_default = |default: &str| -> String {
match backend {
DbBackend::Sqlite => format!("TEXT NOT NULL DEFAULT '{}'", default),
_ => format!("VARCHAR(32) NOT NULL DEFAULT '{}'", default),
}
};
// 兼容旧表:users 增加 deleted_at(原有逻辑)
add_column_if_missing(db, backend, "users", "deleted_at", datetime_type).await?;
// 兼容旧表:users 增加角色与状态字段
add_column_if_missing(db, backend, "users", "role", &varchar_default("user")).await?;
add_column_if_missing(db, backend, "users", "status", &varchar_default("active")).await?;
let text_type = "TEXT NULL";
let nullable_identifier = match backend {
DbBackend::MySql | DbBackend::Postgres => "VARCHAR(255) NULL",
DbBackend::Sqlite => "TEXT NULL",
};
add_column_if_missing(db, backend, "credit_transactions", "description", text_type).await?;
add_column_if_missing(
db,
backend,
"credit_transactions",
"operator_id",
nullable_identifier,
)
.await?;
add_column_if_missing(
db,
backend,
"audit_logs",
"target_type",
nullable_identifier,
)
.await?;
add_column_if_missing(db, backend, "audit_logs", "target_id", nullable_identifier).await?;
add_column_if_missing(db, backend, "audit_logs", "metadata_json", text_type).await?;
Ok(())
}
/// 若列不存在则添加(忽略 duplicate / already exists 错误)
async fn add_column_if_missing(
db: &DatabaseConnection,
backend: DbBackend,
table: &str,
column: &str,
type_decl: &str,
) -> anyhow::Result<()> {
let statement = format!("ALTER TABLE {} ADD COLUMN {} {}", table, column, type_decl);
if let Err(error) = db.execute(Statement::from_string(backend, statement)).await {
let message = error.to_string().to_lowercase();
if !message.contains("duplicate") && !message.contains("already exists") {
return Err(anyhow::anyhow!("迁移 {}.{} 失败: {error}", table, column));
}
}
Ok(())
}
/// 获取端口号(根据数据库类型返回默认值)
fn get_database_port(config: &DatabaseConfig) -> u16 {
config.port.unwrap_or(match config.database_type {
DatabaseType::MySQL => 3306,
DatabaseType::PostgreSQL => 5432,
DatabaseType::SQLite => 0,
})
}
/// 为 MySQL 创建数据库(如果不存在)
async fn init_mysql_database(config: &DatabaseConfig) -> anyhow::Result<()> {
let database_name = config
.database
.as_ref()
.ok_or_else(|| anyhow::anyhow!("MySQL 需要配置 database.database"))?;
validate_database_name(database_name)?;
let host = config
.host
.as_ref()
.ok_or_else(|| anyhow::anyhow!("MySQL 需要配置 database.host"))?;
let user = config
.user
.as_ref()
.ok_or_else(|| anyhow::anyhow!("MySQL 需要配置 database.user"))?;
let password = config
.password
.as_ref()
.ok_or_else(|| anyhow::anyhow!("MySQL 需要配置 database.password"))?;
// 连接到 MySQL 服务器(不指定数据库)
let url = format!(
"mysql://{}:{}@{}:{}",
user,
password,
host,
get_database_port(config)
);
let mut opt = ConnectOptions::new(&url);
opt.max_connections(1)
.connect_timeout(Duration::from_secs(8))
.sqlx_logging(true);
let conn = Database::connect(opt)
.await
.map_err(|e| anyhow::anyhow!("连接 MySQL 服务器失败: {}", e))?;
// 检查数据库是否存在,不存在则创建
let query = format!(
"CREATE DATABASE IF NOT EXISTS `{}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci",
database_name
);
conn.execute(Statement::from_string(
sea_orm::DatabaseBackend::MySql,
query,
))
.await
.map_err(|e| anyhow::anyhow!("创建 MySQL 数据库失败: {}", e))?;
tracing::info!("✅ MySQL 数据库 '{}' 检查完成", database_name);
Ok(())
}
/// 为 PostgreSQL 创建数据库(如果不存在)
async fn init_postgresql_database(config: &DatabaseConfig) -> anyhow::Result<()> {
let database_name = config
.database
.as_ref()
.ok_or_else(|| anyhow::anyhow!("PostgreSQL 需要配置 database.database"))?;
validate_database_name(database_name)?;
let host = config
.host
.as_ref()
.ok_or_else(|| anyhow::anyhow!("PostgreSQL 需要配置 database.host"))?;
let user = config
.user
.as_ref()
.ok_or_else(|| anyhow::anyhow!("PostgreSQL 需要配置 database.user"))?;
let password = config
.password
.as_ref()
.ok_or_else(|| anyhow::anyhow!("PostgreSQL 需要配置 database.password"))?;
// 连接到 PostgreSQL 默认数据库(postgres
let url = format!(
"postgresql://{}:{}@{}:{}/postgres",
user,
password,
host,
get_database_port(config)
);
let mut opt = ConnectOptions::new(&url);
opt.max_connections(1)
.connect_timeout(Duration::from_secs(8))
.sqlx_logging(true);
let conn = Database::connect(opt)
.await
.map_err(|e| anyhow::anyhow!("连接 PostgreSQL 服务器失败: {}", e))?;
// 检查数据库是否存在,不存在则创建
// PostgreSQL 不支持 CREATE DATABASE IF NOT EXISTS,需要先查询
let check_query = format!(
"SELECT 1 FROM pg_database WHERE datname='{}'",
database_name
);
let result = conn
.query_one(Statement::from_string(
sea_orm::DatabaseBackend::Postgres,
check_query,
))
.await
.map_err(|e| anyhow::anyhow!("检查 PostgreSQL 数据库失败: {e}"))?;
match result {
Some(_) => {
tracing::info!("PostgreSQL 数据库 '{}' 已存在", database_name);
}
None => {
// 数据库不存在,创建它
let create_query = format!(
"CREATE DATABASE {} WITH ENCODING 'UTF8' LC_COLLATE='en_US.UTF-8' LC_CTYPE='en_US.UTF-8'",
database_name
);
conn.execute(Statement::from_string(
sea_orm::DatabaseBackend::Postgres,
create_query,
))
.await
.map_err(|e| anyhow::anyhow!("创建 PostgreSQL 数据库失败: {}", e))?;
tracing::info!("✅ PostgreSQL 数据库 '{}' 创建成功", database_name);
}
}
Ok(())
}
fn validate_database_name(name: &str) -> anyhow::Result<()> {
if name.is_empty()
|| !name
.bytes()
.all(|value| value.is_ascii_alphanumeric() || value == b'_')
{
anyhow::bail!("数据库名称只能包含字母、数字和下划线");
}
Ok(())
}
/// 为 SQLite 确保数据库文件目录存在
async fn init_sqlite_database(config: &DatabaseConfig) -> anyhow::Result<()> {
let path = config
.path
.as_ref()
.ok_or_else(|| anyhow::anyhow!("SQLite 需要配置 database.path"))?;
// 如果是相对路径,转换为绝对路径
let absolute_path = if path.is_absolute() {
path.clone()
} else {
std::env::current_dir()
.map_err(|e| anyhow::anyhow!("获取当前目录失败: {}", e))?
.join(path)
};
tracing::info!("SQLite 数据库路径: {}", absolute_path.display());
// 获取数据库文件的父目录
if let Some(parent) = absolute_path.parent() {
// 如果父目录不存在,则创建
if !parent.exists() {
std::fs::create_dir_all(parent)
.map_err(|e| anyhow::anyhow!("创建 SQLite 数据库目录失败: {}", e))?;
tracing::info!("✅ SQLite 数据库目录创建成功: {}", parent.display());
} else {
tracing::info!("SQLite 数据库目录已存在: {}", parent.display());
}
}
// 如果数据库文件不存在,创建空文件
if !absolute_path.exists() {
std::fs::File::create(&absolute_path)
.map_err(|e| anyhow::anyhow!("创建 SQLite 数据库文件失败: {}", e))?;
tracing::info!("✅ SQLite 数据库文件创建成功: {}", absolute_path.display());
} else {
tracing::info!("SQLite 数据库文件已存在: {}", absolute_path.display());
}
Ok(())
}
/// 辅助函数:创建单个表(如果不存在)
async fn create_single_table<E>(
db: &DatabaseConnection,
schema: &Schema,
builder: &DbBackend,
entity: E,
table_name: &str,
) -> anyhow::Result<()>
where
E: EntityName + EntityTrait,
{
let create_table = schema.create_table_from_entity(entity);
let sql = match builder {
DbBackend::MySql => {
use sea_orm::sea_query::MysqlQueryBuilder;
create_table.to_string(MysqlQueryBuilder {})
}
DbBackend::Postgres => {
use sea_orm::sea_query::PostgresQueryBuilder;
create_table.to_string(PostgresQueryBuilder {})
}
DbBackend::Sqlite => {
use sea_orm::sea_query::SqliteQueryBuilder;
create_table.to_string(SqliteQueryBuilder {})
}
};
let sql = sql.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS");
match db.execute(Statement::from_string(*builder, sql)).await {
Ok(_) => {
tracing::info!("✅ {}检查完成", table_name);
}
Err(e) => {
let err_msg = e.to_string();
if err_msg.contains("already exists")
|| (err_msg.contains("table") && err_msg.contains("exists"))
{
tracing::info!("✅ {}已存在", table_name);
} else {
return Err(anyhow::anyhow!("创建{}失败: {}", table_name, e));
}
}
}
Ok(())
}
/// 创建数据库表结构
async fn create_tables(db: &DatabaseConnection) -> anyhow::Result<()> {
tracing::info!("检查数据库表结构...");
let builder = db.get_database_backend();
let schema = Schema::new(builder);
// 导入所有 entities
use crate::domain::entities::{
abuse_event, abuse_rule, audit_log, blocked_ip, blocked_sender, credit_account,
credit_check_in, credit_rule, credit_rule_change, credit_transaction, daily_quota, email,
email_attachment, email_logs, greylist_entry, mailbox, user_profiles, users,
};
// 创建所有表(添加新表只需一行!)
create_single_table(db, &schema, &builder, users::Entity, "用户表").await?;
create_single_table(db, &schema, &builder, user_profiles::Entity, "用户资料表").await?;
create_single_table(db, &schema, &builder, mailbox::Entity, "邮箱表").await?;
create_single_table(db, &schema, &builder, email::Entity, "邮件表").await?;
create_single_table(db, &schema, &builder, email_attachment::Entity, "附件表").await?;
create_single_table(db, &schema, &builder, email_logs::Entity, "邮件日志表").await?;
create_single_table(db, &schema, &builder, credit_account::Entity, "积分账户表").await?;
create_single_table(
db,
&schema,
&builder,
credit_transaction::Entity,
"积分流水表",
)
.await?;
create_single_table(db, &schema, &builder, credit_rule::Entity, "积分规则表").await?;
create_single_table(
db,
&schema,
&builder,
credit_rule_change::Entity,
"积分规则变更表",
)
.await?;
create_single_table(db, &schema, &builder, credit_check_in::Entity, "每日签到表").await?;
create_single_table(
db,
&schema,
&builder,
blocked_sender::Entity,
"发件黑名单表",
)
.await?;
create_single_table(db, &schema, &builder, blocked_ip::Entity, "IP黑名单表").await?;
create_single_table(db, &schema, &builder, greylist_entry::Entity, "灰名单表").await?;
create_single_table(db, &schema, &builder, audit_log::Entity, "审计日志表").await?;
create_single_table(db, &schema, &builder, daily_quota::Entity, "每日额度表").await?;
create_single_table(db, &schema, &builder, abuse_rule::Entity, "滥用规则表").await?;
create_single_table(db, &schema, &builder, abuse_event::Entity, "滥用事件表").await?;
tracing::info!("✅ 数据库表结构检查完成");
Ok(())
}