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 { 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 { 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( 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(()) }