use r2d2::Pool; use r2d2_sqlite::SqliteConnectionManager; use rand::Rng; use rusqlite::params; use std::collections::HashMap; use std::path::Path; use tracing; use crate::constants; use crate::error::AppErrorKind; pub type DbPool = Pool; pub fn db_path(data_dir: &str) -> String { std::fs::create_dir_all(data_dir).ok(); Path::new(data_dir) .join("file_metadata.db") .to_string_lossy() .to_string() } pub fn init_db(data_dir: &str) -> DbPool { let path = db_path(data_dir); let manager = SqliteConnectionManager::file(&path); let pool = Pool::builder() .max_size(10) .min_idle(Some(2)) .connection_customizer(Box::new(SqliteInitializer)) .build(manager) .expect("Failed to create database pool"); let conn = pool.get().expect("Failed to get connection for init"); conn.execute_batch( "CREATE TABLE IF NOT EXISTS files ( id INTEGER PRIMARY KEY AUTOINCREMENT, filename TEXT NOT NULL, file_id TEXT NOT NULL UNIQUE, filesize INTEGER NOT NULL, upload_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP, short_id TEXT UNIQUE );", ) .expect("Failed to create files table"); let has_short_id: bool = conn .prepare("PRAGMA table_info(files)") .unwrap() .query_map([], |row| row.get::<_, String>(1)) .unwrap() .any(|col| col.map_or(false, |c| c == "short_id")); if !has_short_id { tracing::info!("Migrating database: adding short_id column..."); if let Err(e) = conn.execute("ALTER TABLE files ADD COLUMN short_id TEXT", []) { tracing::error!("Migration warning: Failed to add short_id column: {}", e); } } conn.execute_batch( "CREATE UNIQUE INDEX IF NOT EXISTS idx_files_short_id ON files(short_id);", ) .ok(); conn.execute_batch( "CREATE INDEX IF NOT EXISTS idx_files_upload_date ON files(upload_date DESC);", ) .ok(); conn.execute_batch( "CREATE TABLE IF NOT EXISTS app_settings ( id INTEGER PRIMARY KEY CHECK (id = 1), bot_token TEXT, channel_name TEXT, pass_word TEXT, picgo_api_key TEXT, base_url TEXT );", ) .expect("Failed to create app_settings table"); conn.execute( "INSERT OR IGNORE INTO app_settings (id) VALUES (1)", [], ) .expect("Failed to init app_settings row"); // Migration: add session_token column if missing let has_session_token: bool = conn .prepare("PRAGMA table_info(app_settings)") .unwrap() .query_map([], |row| row.get::<_, String>(1)) .unwrap() .any(|col| col.map_or(false, |c| c == "session_token")); if !has_session_token { tracing::info!("Migrating database: adding session_token column..."); if let Err(e) = conn.execute("ALTER TABLE app_settings ADD COLUMN session_token TEXT", []) { tracing::error!("Migration warning: Failed to add session_token column: {}", e); } } tracing::info!("数据库已成功初始化"); pool } #[derive(Debug)] struct SqliteInitializer; impl r2d2::CustomizeConnection for SqliteInitializer { fn on_acquire(&self, conn: &mut rusqlite::Connection) -> Result<(), rusqlite::Error> { conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000;")?; Ok(()) } } fn generate_short_id(length: usize) -> String { const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; let mut rng = rand::thread_rng(); (0..length) .map(|_| CHARS[rng.gen_range(0..CHARS.len())] as char) .collect() } pub fn add_file_metadata( pool: &DbPool, filename: &str, file_id: &str, filesize: i64, ) -> Result { let conn = pool.get()?; for _ in 0..5 { let short_id = generate_short_id(constants::SHORT_ID_LENGTH); match conn.execute( "INSERT INTO files (filename, file_id, filesize, short_id) VALUES (?1, ?2, ?3, ?4)", params![filename, file_id, filesize, short_id], ) { Ok(_) => { // Only log a short_id prefix. `/d/` is a public, // unauthenticated download endpoint; the short_id is therefore // a bearer capability and must not be logged in full. let short_id_preview = short_id.chars().take(2).collect::(); tracing::info!( "已添加文件元数据: {}, short_id: {}***", filename, short_id_preview ); return Ok(short_id); } Err(rusqlite::Error::SqliteFailure(err, _)) if err.code == rusqlite::ErrorCode::ConstraintViolation => { let existing: Option = conn .query_row( "SELECT short_id FROM files WHERE file_id = ?1", params![file_id], |row| row.get(0), ) .ok(); if let Some(existing_sid) = existing { if !existing_sid.is_empty() { return Ok(existing_sid); } let new_sid = generate_short_id(constants::SHORT_ID_LENGTH); conn.execute( "UPDATE files SET short_id = ?1 WHERE file_id = ?2", params![new_sid, file_id], )?; return Ok(new_sid); } continue; } Err(e) => return Err(e.into()), } } Err(AppErrorKind::Other("Failed to generate unique short_id".into())) } #[derive(Debug, Clone, serde::Serialize)] pub struct FileMetadata { pub filename: String, pub file_id: String, pub filesize: i64, pub upload_date: String, pub short_id: Option, } pub fn get_all_files(pool: &DbPool) -> Result, AppErrorKind> { let conn = pool.get()?; let mut stmt = conn .prepare( "SELECT filename, file_id, filesize, upload_date, short_id FROM files ORDER BY upload_date DESC", )?; let files = stmt .query_map([], |row| { Ok(FileMetadata { filename: row.get(0)?, file_id: row.get(1)?, filesize: row.get(2)?, upload_date: row.get::<_, String>(3).unwrap_or_default(), short_id: row.get(4).ok(), }) })? .filter_map(|r| r.ok()) .collect(); Ok(files) } pub fn get_file_by_id(pool: &DbPool, identifier: &str) -> Result, AppErrorKind> { let conn = pool.get()?; let result = conn.query_row( "SELECT filename, filesize, upload_date, file_id, short_id FROM files WHERE short_id = ?1 OR file_id = ?1", params![identifier], |row| { Ok(FileMetadata { filename: row.get(0)?, filesize: row.get(1)?, upload_date: row.get::<_, String>(2).unwrap_or_default(), file_id: row.get(3)?, short_id: row.get(4).ok(), }) }, ); match result { Ok(meta) => Ok(Some(meta)), Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), Err(e) => Err(e.into()), } } pub fn delete_file_metadata(pool: &DbPool, file_id: &str) -> Result { let conn = pool.get()?; let rows = conn .execute("DELETE FROM files WHERE file_id = ?1", params![file_id])?; Ok(rows > 0) } pub fn delete_file_by_message_id(pool: &DbPool, message_id: i64) -> Result, AppErrorKind> { let conn = pool.get()?; let pattern = format!("{}:%", message_id); let file_id: Option = conn .query_row( "SELECT file_id FROM files WHERE file_id LIKE ?1", params![pattern], |row| row.get(0), ) .ok(); if let Some(ref fid) = file_id { conn.execute("DELETE FROM files WHERE file_id = ?1", params![fid])?; tracing::info!( "已从数据库中删除与消息ID {} 关联的文件: {}", message_id, fid ); } Ok(file_id) } pub fn get_app_settings_from_db( pool: &DbPool, ) -> Result>, AppErrorKind> { let conn = pool.get()?; let result = conn.query_row( "SELECT bot_token, channel_name, pass_word, picgo_api_key, base_url, session_token FROM app_settings WHERE id = 1", [], |row| { let mut map = HashMap::new(); map.insert("BOT_TOKEN".to_string(), row.get::<_, Option>(0)?); map.insert("CHANNEL_NAME".to_string(), row.get::<_, Option>(1)?); map.insert("PASS_WORD".to_string(), row.get::<_, Option>(2)?); map.insert("PICGO_API_KEY".to_string(), row.get::<_, Option>(3)?); map.insert("BASE_URL".to_string(), row.get::<_, Option>(4)?); map.insert("SESSION_TOKEN".to_string(), row.get::<_, Option>(5)?); Ok(map) }, ); match result { Ok(map) => Ok(map), Err(rusqlite::Error::QueryReturnedNoRows) => Ok(HashMap::new()), Err(e) => Err(e.into()), } } fn norm(v: Option<&str>) -> Option { v.map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()) } pub fn save_app_settings_to_db( pool: &DbPool, payload: &HashMap>, ) -> Result<(), AppErrorKind> { let conn = pool.get()?; conn.execute( "UPDATE app_settings SET bot_token = ?1, channel_name = ?2, pass_word = ?3, picgo_api_key = ?4, base_url = ?5, session_token = ?6 WHERE id = 1", params![ norm(payload.get("BOT_TOKEN").and_then(|v| v.as_deref())), norm(payload.get("CHANNEL_NAME").and_then(|v| v.as_deref())), norm(payload.get("PASS_WORD").and_then(|v| v.as_deref())), norm(payload.get("PICGO_API_KEY").and_then(|v| v.as_deref())), norm(payload.get("BASE_URL").and_then(|v| v.as_deref())), norm(payload.get("SESSION_TOKEN").and_then(|v| v.as_deref())), ], )?; Ok(()) } pub fn reset_app_settings_in_db(pool: &DbPool) -> Result<(), AppErrorKind> { let mut payload = HashMap::new(); payload.insert("BOT_TOKEN".to_string(), None); payload.insert("CHANNEL_NAME".to_string(), None); payload.insert("PASS_WORD".to_string(), None); payload.insert("PICGO_API_KEY".to_string(), None); payload.insert("BASE_URL".to_string(), None); payload.insert("SESSION_TOKEN".to_string(), None); save_app_settings_to_db(pool, &payload) }