mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-29 09:37:37 +08:00
feat: add usage daily rollups, incremental auto-vacuum, and sync-aware backup
- Add usage_daily_rollups table (schema v6) to aggregate proxy request logs into daily summaries, reducing query overhead for statistics - Add rollup_and_prune DAO that aggregates old detail logs (>N days) into rollup rows and deletes the originals - Update all usage stats queries to UNION detail logs with rollup data - Introduce incremental auto-vacuum for SQLite, with startup and periodic cleanup of old stream_check_logs and request log rollups - Split backup export/import into full vs sync variants: WebDAV sync now skips local-only table data (proxy_request_logs, stream_check_logs, provider_health, proxy_live_backup, usage_daily_rollups) while preserving them on import - Add enable_logging guard to skip request log writes when disabled - Apply cargo fmt formatting fixes across multiple modules
This commit is contained in:
@@ -15,6 +15,16 @@ use tempfile::NamedTempFile;
|
||||
|
||||
const CC_SWITCH_SQL_EXPORT_HEADER: &str = "-- CC Switch SQLite 导出";
|
||||
|
||||
/// Tables that are local-only (not synced via WebDAV snapshots).
|
||||
/// Schema is still exported, but data rows are skipped.
|
||||
const LOCAL_ONLY_TABLES: &[&str] = &[
|
||||
"proxy_request_logs",
|
||||
"stream_check_logs",
|
||||
"provider_health",
|
||||
"proxy_live_backup",
|
||||
"usage_daily_rollups",
|
||||
];
|
||||
|
||||
/// A database backup entry for the UI
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
@@ -25,10 +35,16 @@ pub struct BackupEntry {
|
||||
}
|
||||
|
||||
impl Database {
|
||||
/// 导出为 SQLite 兼容的 SQL 文本(内存字符串)
|
||||
/// 导出为 SQLite 兼容的 SQL 文本(内存字符串,完整导出)
|
||||
pub fn export_sql_string(&self) -> Result<String, AppError> {
|
||||
let snapshot = self.snapshot_to_memory()?;
|
||||
Self::dump_sql(&snapshot)
|
||||
Self::dump_sql(&snapshot, &[])
|
||||
}
|
||||
|
||||
/// Export SQL for sync (WebDAV), skipping local-only tables' data
|
||||
pub fn export_sql_string_for_sync(&self) -> Result<String, AppError> {
|
||||
let snapshot = self.snapshot_to_memory()?;
|
||||
Self::dump_sql(&snapshot, LOCAL_ONLY_TABLES)
|
||||
}
|
||||
|
||||
/// 导出为 SQLite 兼容的 SQL 文本
|
||||
@@ -58,12 +74,32 @@ impl Database {
|
||||
|
||||
/// 从 SQL 字符串导入,返回生成的备份 ID(若无备份则为空字符串)
|
||||
pub fn import_sql_string(&self, sql_raw: &str) -> Result<String, AppError> {
|
||||
self.import_sql_string_inner(sql_raw, &[])
|
||||
}
|
||||
|
||||
/// Import SQL generated for sync, then restore local-only tables from the
|
||||
/// current device snapshot before replacing the main database.
|
||||
pub(crate) fn import_sql_string_for_sync(&self, sql_raw: &str) -> Result<String, AppError> {
|
||||
self.import_sql_string_inner(sql_raw, LOCAL_ONLY_TABLES)
|
||||
}
|
||||
|
||||
fn import_sql_string_inner(
|
||||
&self,
|
||||
sql_raw: &str,
|
||||
preserve_tables: &[&str],
|
||||
) -> Result<String, AppError> {
|
||||
let sql_content = sql_raw.trim_start_matches('\u{feff}');
|
||||
Self::validate_cc_switch_sql_export(sql_content)?;
|
||||
|
||||
// 导入前备份现有数据库
|
||||
let backup_path = self.backup_database_file()?;
|
||||
|
||||
let local_snapshot = if preserve_tables.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(self.snapshot_to_memory()?)
|
||||
};
|
||||
|
||||
// 在临时数据库执行导入,确保失败不会污染主库
|
||||
let temp_file = NamedTempFile::new().map_err(|e| AppError::IoContext {
|
||||
context: "创建临时数据库文件失败".to_string(),
|
||||
@@ -81,6 +117,9 @@ impl Database {
|
||||
Self::create_tables_on_conn(&temp_conn)?;
|
||||
Self::apply_schema_migrations_on_conn(&temp_conn)?;
|
||||
Self::validate_basic_state(&temp_conn)?;
|
||||
if let Some(local_snapshot) = local_snapshot.as_ref() {
|
||||
Self::restore_tables(local_snapshot, &temp_conn, preserve_tables)?;
|
||||
}
|
||||
|
||||
// 使用 Backup 将临时库原子写回主库
|
||||
{
|
||||
@@ -129,6 +168,62 @@ impl Database {
|
||||
))
|
||||
}
|
||||
|
||||
fn restore_tables(
|
||||
source_conn: &Connection,
|
||||
target_conn: &Connection,
|
||||
tables: &[&str],
|
||||
) -> Result<(), AppError> {
|
||||
for table in tables {
|
||||
if !Self::table_exists(source_conn, table)? || !Self::table_exists(target_conn, table)?
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let columns = Self::get_table_columns(source_conn, table)?;
|
||||
if columns.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
target_conn
|
||||
.execute(&format!("DELETE FROM \"{table}\""), [])
|
||||
.map_err(|e| AppError::Database(format!("清空表 {table} 失败: {e}")))?;
|
||||
|
||||
let placeholders = (1..=columns.len())
|
||||
.map(|idx| format!("?{idx}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let cols = columns
|
||||
.iter()
|
||||
.map(|column| format!("\"{column}\""))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let insert_sql = format!("INSERT INTO \"{table}\" ({cols}) VALUES ({placeholders})");
|
||||
|
||||
let mut stmt = source_conn
|
||||
.prepare(&format!("SELECT * FROM \"{table}\""))
|
||||
.map_err(|e| AppError::Database(format!("读取表 {table} 失败: {e}")))?;
|
||||
let mut rows = stmt
|
||||
.query([])
|
||||
.map_err(|e| AppError::Database(format!("查询表 {table} 数据失败: {e}")))?;
|
||||
|
||||
while let Some(row) = rows.next().map_err(|e| AppError::Database(e.to_string()))? {
|
||||
let mut values = Vec::with_capacity(columns.len());
|
||||
for idx in 0..columns.len() {
|
||||
values.push(
|
||||
row.get::<_, rusqlite::types::Value>(idx)
|
||||
.map_err(|e| AppError::Database(e.to_string()))?,
|
||||
);
|
||||
}
|
||||
|
||||
target_conn
|
||||
.execute(&insert_sql, rusqlite::params_from_iter(values.iter()))
|
||||
.map_err(|e| AppError::Database(format!("恢复表 {table} 数据失败: {e}")))?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Periodic backup: create a new backup if the latest one is older than the configured interval
|
||||
pub(crate) fn periodic_backup_if_needed(&self) -> Result<(), AppError> {
|
||||
let interval_hours = crate::settings::effective_backup_interval_hours();
|
||||
@@ -165,6 +260,15 @@ impl Database {
|
||||
);
|
||||
self.backup_database_file()?;
|
||||
}
|
||||
|
||||
// Periodic cleanup: prune old logs alongside backup cycle
|
||||
if let Err(e) = self.cleanup_old_stream_check_logs(7) {
|
||||
log::warn!("Periodic stream_check_logs cleanup failed: {e}");
|
||||
}
|
||||
if let Err(e) = self.rollup_and_prune(30) {
|
||||
log::warn!("Periodic rollup_and_prune failed: {e}");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -258,7 +362,7 @@ impl Database {
|
||||
}
|
||||
|
||||
/// 导出数据库为 SQL 文本
|
||||
fn dump_sql(conn: &Connection) -> Result<String, AppError> {
|
||||
fn dump_sql(conn: &Connection, skip_tables: &[&str]) -> Result<String, AppError> {
|
||||
let mut output = String::new();
|
||||
let timestamp = Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
||||
let user_version: i64 = conn
|
||||
@@ -306,6 +410,9 @@ impl Database {
|
||||
|
||||
// 导出数据
|
||||
for table in tables {
|
||||
if skip_tables.iter().any(|t| *t == table) {
|
||||
continue;
|
||||
}
|
||||
let columns = Self::get_table_columns(conn, &table)?;
|
||||
if columns.is_empty() {
|
||||
continue;
|
||||
@@ -557,3 +664,96 @@ impl Database {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::Database;
|
||||
use crate::error::AppError;
|
||||
|
||||
#[test]
|
||||
fn sync_import_preserves_local_only_tables() -> Result<(), AppError> {
|
||||
let remote_db = Database::memory()?;
|
||||
{
|
||||
let conn = crate::database::lock_conn!(remote_db.conn);
|
||||
conn.execute(
|
||||
"INSERT INTO providers (id, app_type, name, settings_config, meta)
|
||||
VALUES ('remote-provider', 'claude', 'Remote Provider', '{}', '{}')",
|
||||
[],
|
||||
)?;
|
||||
}
|
||||
let remote_sql = remote_db.export_sql_string_for_sync()?;
|
||||
|
||||
let local_db = Database::memory()?;
|
||||
{
|
||||
let conn = crate::database::lock_conn!(local_db.conn);
|
||||
conn.execute(
|
||||
"INSERT INTO providers (id, app_type, name, settings_config, meta)
|
||||
VALUES ('local-provider', 'claude', 'Local Provider', '{}', '{}')",
|
||||
[],
|
||||
)?;
|
||||
conn.execute(
|
||||
"INSERT INTO proxy_request_logs (
|
||||
request_id, provider_id, app_type, model,
|
||||
input_tokens, output_tokens, total_cost_usd,
|
||||
latency_ms, status_code, created_at
|
||||
) VALUES ('req-1', 'local-provider', 'claude', 'claude-3', 100, 50, '0.01', 120, 200, 1000)",
|
||||
[],
|
||||
)?;
|
||||
conn.execute(
|
||||
"INSERT INTO usage_daily_rollups (
|
||||
date, app_type, provider_id, model, request_count, success_count,
|
||||
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
|
||||
total_cost_usd, avg_latency_ms
|
||||
) VALUES ('2026-03-01', 'claude', 'local-provider', 'claude-3', 7, 7, 700, 350, 0, 0, '0.07', 120)",
|
||||
[],
|
||||
)?;
|
||||
conn.execute(
|
||||
"INSERT INTO stream_check_logs (
|
||||
provider_id, provider_name, app_type, status, success, message,
|
||||
response_time_ms, http_status, model_used, retry_count, tested_at
|
||||
) VALUES ('local-provider', 'Local Provider', 'claude', 'operational', 1, 'ok', 42, 200, 'claude-3', 0, 1000)",
|
||||
[],
|
||||
)?;
|
||||
}
|
||||
|
||||
local_db.import_sql_string_for_sync(&remote_sql)?;
|
||||
|
||||
let remote_provider_exists: i64 = {
|
||||
let conn = crate::database::lock_conn!(local_db.conn);
|
||||
conn.query_row(
|
||||
"SELECT COUNT(*) FROM providers WHERE id = 'remote-provider' AND app_type = 'claude'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)?
|
||||
};
|
||||
assert_eq!(
|
||||
remote_provider_exists, 1,
|
||||
"remote config should be imported"
|
||||
);
|
||||
|
||||
let (request_logs, rollups, stream_logs): (i64, i64, i64) = {
|
||||
let conn = crate::database::lock_conn!(local_db.conn);
|
||||
let request_logs =
|
||||
conn.query_row("SELECT COUNT(*) FROM proxy_request_logs", [], |row| {
|
||||
row.get(0)
|
||||
})?;
|
||||
let rollups =
|
||||
conn.query_row("SELECT COUNT(*) FROM usage_daily_rollups", [], |row| {
|
||||
row.get(0)
|
||||
})?;
|
||||
let stream_logs =
|
||||
conn.query_row("SELECT COUNT(*) FROM stream_check_logs", [], |row| {
|
||||
row.get(0)
|
||||
})?;
|
||||
(request_logs, rollups, stream_logs)
|
||||
};
|
||||
assert_eq!(request_logs, 1, "local request logs should be preserved");
|
||||
assert_eq!(rollups, 1, "local rollups should be preserved");
|
||||
assert_eq!(
|
||||
stream_logs, 1,
|
||||
"local stream check logs should be preserved"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ pub mod settings;
|
||||
pub mod skills;
|
||||
pub mod stream_check;
|
||||
pub mod universal_providers;
|
||||
pub mod usage_rollup;
|
||||
|
||||
// 所有 DAO 方法都通过 Database impl 提供,无需单独导出
|
||||
// 导出 FailoverQueueItem 供外部使用
|
||||
|
||||
@@ -48,6 +48,23 @@ impl Database {
|
||||
}
|
||||
}
|
||||
|
||||
/// Delete stream check logs older than `retain_days` days.
|
||||
/// Returns the number of deleted rows.
|
||||
pub fn cleanup_old_stream_check_logs(&self, retain_days: i64) -> Result<u64, AppError> {
|
||||
let cutoff = chrono::Utc::now().timestamp() - retain_days * 86400;
|
||||
let conn = lock_conn!(self.conn);
|
||||
let deleted = conn
|
||||
.execute(
|
||||
"DELETE FROM stream_check_logs WHERE tested_at < ?1",
|
||||
[cutoff],
|
||||
)
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
if deleted > 0 {
|
||||
log::info!("Cleaned up {deleted} stream_check_logs older than {retain_days} days");
|
||||
}
|
||||
Ok(deleted as u64)
|
||||
}
|
||||
|
||||
/// 保存流式检查配置
|
||||
pub fn save_stream_check_config(&self, config: &StreamCheckConfig) -> Result<(), AppError> {
|
||||
let json = serde_json::to_string(config)
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
//! Usage rollup DAO
|
||||
//!
|
||||
//! Aggregates proxy_request_logs into daily rollups and prunes old detail rows.
|
||||
|
||||
use crate::database::{lock_conn, Database};
|
||||
use crate::error::AppError;
|
||||
|
||||
impl Database {
|
||||
/// Aggregate proxy_request_logs older than `retain_days` into usage_daily_rollups,
|
||||
/// then delete the aggregated detail rows.
|
||||
/// Returns the number of deleted detail rows.
|
||||
pub fn rollup_and_prune(&self, retain_days: i64) -> Result<u64, AppError> {
|
||||
let cutoff = chrono::Utc::now().timestamp() - retain_days * 86400;
|
||||
let conn = lock_conn!(self.conn);
|
||||
|
||||
// Check if there are any rows to process
|
||||
let count: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM proxy_request_logs WHERE created_at < ?1",
|
||||
[cutoff],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
|
||||
if count == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// Use a savepoint for atomicity
|
||||
conn.execute("SAVEPOINT rollup_prune;", [])
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
|
||||
let result = Self::do_rollup_and_prune(&conn, cutoff);
|
||||
|
||||
match result {
|
||||
Ok(deleted) => {
|
||||
conn.execute("RELEASE rollup_prune;", [])
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
if deleted > 0 {
|
||||
log::info!(
|
||||
"Rolled up and pruned {deleted} proxy_request_logs (retain={retain_days}d)"
|
||||
);
|
||||
}
|
||||
Ok(deleted)
|
||||
}
|
||||
Err(e) => {
|
||||
conn.execute("ROLLBACK TO rollup_prune;", []).ok();
|
||||
conn.execute("RELEASE rollup_prune;", []).ok();
|
||||
Err(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn do_rollup_and_prune(conn: &rusqlite::Connection, cutoff: i64) -> Result<u64, AppError> {
|
||||
// Aggregate old logs, merging with any pre-existing rollup rows via LEFT JOIN.
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO usage_daily_rollups
|
||||
(date, app_type, provider_id, model,
|
||||
request_count, success_count,
|
||||
input_tokens, output_tokens,
|
||||
cache_read_tokens, cache_creation_tokens,
|
||||
total_cost_usd, avg_latency_ms)
|
||||
SELECT
|
||||
d, a, p, m,
|
||||
COALESCE(old.request_count, 0) + new_req,
|
||||
COALESCE(old.success_count, 0) + new_succ,
|
||||
COALESCE(old.input_tokens, 0) + new_in,
|
||||
COALESCE(old.output_tokens, 0) + new_out,
|
||||
COALESCE(old.cache_read_tokens, 0) + new_cr,
|
||||
COALESCE(old.cache_creation_tokens, 0) + new_cc,
|
||||
CAST(COALESCE(CAST(old.total_cost_usd AS REAL), 0) + new_cost AS TEXT),
|
||||
CASE WHEN COALESCE(old.request_count, 0) + new_req > 0
|
||||
THEN (COALESCE(old.avg_latency_ms, 0) * COALESCE(old.request_count, 0)
|
||||
+ new_lat * new_req)
|
||||
/ (COALESCE(old.request_count, 0) + new_req)
|
||||
ELSE 0 END
|
||||
FROM (
|
||||
SELECT
|
||||
date(created_at, 'unixepoch', 'localtime') as d,
|
||||
app_type as a, provider_id as p, model as m,
|
||||
COUNT(*) as new_req,
|
||||
SUM(CASE WHEN status_code >= 200 AND status_code < 300 THEN 1 ELSE 0 END) as new_succ,
|
||||
COALESCE(SUM(input_tokens), 0) as new_in,
|
||||
COALESCE(SUM(output_tokens), 0) as new_out,
|
||||
COALESCE(SUM(cache_read_tokens), 0) as new_cr,
|
||||
COALESCE(SUM(cache_creation_tokens), 0) as new_cc,
|
||||
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) as new_cost,
|
||||
COALESCE(AVG(latency_ms), 0) as new_lat
|
||||
FROM proxy_request_logs WHERE created_at < ?1
|
||||
GROUP BY d, a, p, m
|
||||
) agg
|
||||
LEFT JOIN usage_daily_rollups old
|
||||
ON old.date = agg.d AND old.app_type = agg.a
|
||||
AND old.provider_id = agg.p AND old.model = agg.m",
|
||||
[cutoff],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("Rollup aggregation failed: {e}")))?;
|
||||
|
||||
// Delete the aggregated detail rows
|
||||
let deleted = conn
|
||||
.execute(
|
||||
"DELETE FROM proxy_request_logs WHERE created_at < ?1",
|
||||
[cutoff],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("Pruning old logs failed: {e}")))?;
|
||||
|
||||
Ok(deleted as u64)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::database::Database;
|
||||
use crate::error::AppError;
|
||||
|
||||
#[test]
|
||||
fn test_rollup_and_prune() -> Result<(), AppError> {
|
||||
let db = Database::memory()?;
|
||||
let now = chrono::Utc::now().timestamp();
|
||||
let old_ts = now - 40 * 86400; // 40 days ago
|
||||
let recent_ts = now - 5 * 86400; // 5 days ago
|
||||
|
||||
{
|
||||
let conn = crate::database::lock_conn!(db.conn);
|
||||
for i in 0..5 {
|
||||
conn.execute(
|
||||
"INSERT INTO proxy_request_logs (
|
||||
request_id, provider_id, app_type, model,
|
||||
input_tokens, output_tokens, total_cost_usd,
|
||||
latency_ms, status_code, created_at
|
||||
) VALUES (?1, 'p1', 'claude', 'claude-3', 100, 50, '0.01', 100, 200, ?2)",
|
||||
rusqlite::params![format!("old-{i}"), old_ts + i as i64],
|
||||
)?;
|
||||
}
|
||||
for i in 0..3 {
|
||||
conn.execute(
|
||||
"INSERT INTO proxy_request_logs (
|
||||
request_id, provider_id, app_type, model,
|
||||
input_tokens, output_tokens, total_cost_usd,
|
||||
latency_ms, status_code, created_at
|
||||
) VALUES (?1, 'p1', 'claude', 'claude-3', 200, 100, '0.02', 150, 200, ?2)",
|
||||
rusqlite::params![format!("recent-{i}"), recent_ts + i as i64],
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
let deleted = db.rollup_and_prune(30)?;
|
||||
assert_eq!(deleted, 5);
|
||||
|
||||
// Verify rollup data
|
||||
let conn = crate::database::lock_conn!(db.conn);
|
||||
let count: i64 = conn.query_row(
|
||||
"SELECT request_count FROM usage_daily_rollups WHERE app_type = 'claude'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
assert_eq!(count, 5);
|
||||
|
||||
// Verify recent logs untouched
|
||||
let remaining: i64 =
|
||||
conn.query_row("SELECT COUNT(*) FROM proxy_request_logs", [], |row| {
|
||||
row.get(0)
|
||||
})?;
|
||||
assert_eq!(remaining, 3);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rollup_noop_when_no_old_data() -> Result<(), AppError> {
|
||||
let db = Database::memory()?;
|
||||
assert_eq!(db.rollup_and_prune(30)?, 0);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rollup_merges_with_existing() -> Result<(), AppError> {
|
||||
let db = Database::memory()?;
|
||||
let now = chrono::Utc::now().timestamp();
|
||||
let old_ts = now - 40 * 86400;
|
||||
|
||||
{
|
||||
let conn = crate::database::lock_conn!(db.conn);
|
||||
let date_str = chrono::DateTime::from_timestamp(old_ts, 0)
|
||||
.unwrap()
|
||||
.format("%Y-%m-%d")
|
||||
.to_string();
|
||||
conn.execute(
|
||||
"INSERT INTO usage_daily_rollups
|
||||
(date, app_type, provider_id, model, request_count, success_count,
|
||||
input_tokens, output_tokens, total_cost_usd, avg_latency_ms)
|
||||
VALUES (?1, 'claude', 'p1', 'claude-3', 10, 10, 1000, 500, '0.10', 100)",
|
||||
[&date_str],
|
||||
)?;
|
||||
for i in 0..3 {
|
||||
conn.execute(
|
||||
"INSERT INTO proxy_request_logs (
|
||||
request_id, provider_id, app_type, model,
|
||||
input_tokens, output_tokens, total_cost_usd,
|
||||
latency_ms, status_code, created_at
|
||||
) VALUES (?1, 'p1', 'claude', 'claude-3', 100, 50, '0.01', 200, 200, ?2)",
|
||||
rusqlite::params![format!("merge-{i}"), old_ts + i as i64],
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
let deleted = db.rollup_and_prune(30)?;
|
||||
assert_eq!(deleted, 3);
|
||||
|
||||
let conn = crate::database::lock_conn!(db.conn);
|
||||
let (count, input): (i64, i64) = conn.query_row(
|
||||
"SELECT request_count, input_tokens FROM usage_daily_rollups
|
||||
WHERE app_type = 'claude' AND provider_id = 'p1'",
|
||||
[],
|
||||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||||
)?;
|
||||
assert_eq!(count, 13, "10 existing + 3 new");
|
||||
assert_eq!(input, 1300, "1000 existing + 300 new");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -44,7 +44,7 @@ use std::sync::Mutex;
|
||||
|
||||
/// 当前 Schema 版本号
|
||||
/// 每次修改表结构时递增,并在 schema.rs 中添加相应的迁移逻辑
|
||||
pub(crate) const SCHEMA_VERSION: i32 = 5;
|
||||
pub(crate) const SCHEMA_VERSION: i32 = 6;
|
||||
|
||||
/// 安全地序列化 JSON,避免 unwrap panic
|
||||
pub(crate) fn to_json_string<T: Serialize>(value: &T) -> Result<String, AppError> {
|
||||
@@ -89,6 +89,7 @@ impl Database {
|
||||
/// 数据库文件位于 `~/.cc-switch/cc-switch.db`
|
||||
pub fn init() -> Result<Self, AppError> {
|
||||
let db_path = get_app_config_dir().join("cc-switch.db");
|
||||
let db_exists = db_path.exists();
|
||||
|
||||
// 确保父目录存在
|
||||
if let Some(parent) = db_path.parent() {
|
||||
@@ -100,6 +101,12 @@ impl Database {
|
||||
// 启用外键约束
|
||||
conn.execute("PRAGMA foreign_keys = ON;", [])
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
if !db_exists {
|
||||
// For a brand-new database, configure incremental auto-vacuum
|
||||
// before creating any tables so no rebuild is needed later.
|
||||
conn.execute("PRAGMA auto_vacuum = INCREMENTAL;", [])
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
}
|
||||
register_db_change_hook(&conn);
|
||||
|
||||
let db = Self {
|
||||
@@ -123,8 +130,26 @@ impl Database {
|
||||
}
|
||||
|
||||
db.apply_schema_migrations()?;
|
||||
if let Err(e) = db.ensure_incremental_auto_vacuum() {
|
||||
log::warn!("Failed to ensure incremental auto-vacuum: {e}");
|
||||
}
|
||||
db.ensure_model_pricing_seeded()?;
|
||||
|
||||
// Startup cleanup: prune old logs and reclaim space
|
||||
if let Err(e) = db.cleanup_old_stream_check_logs(7) {
|
||||
log::warn!("Startup stream_check_logs cleanup failed: {e}");
|
||||
}
|
||||
if let Err(e) = db.rollup_and_prune(30) {
|
||||
log::warn!("Startup rollup_and_prune failed: {e}");
|
||||
}
|
||||
// Reclaim disk space after cleanup
|
||||
{
|
||||
let conn = lock_conn!(db.conn);
|
||||
if let Err(e) = conn.execute_batch("PRAGMA incremental_vacuum;") {
|
||||
log::warn!("Startup incremental vacuum failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
@@ -135,6 +160,8 @@ impl Database {
|
||||
// 启用外键约束
|
||||
conn.execute("PRAGMA foreign_keys = ON;", [])
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
conn.execute("PRAGMA auto_vacuum = INCREMENTAL;", [])
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
register_db_change_hook(&conn);
|
||||
|
||||
let db = Self {
|
||||
@@ -146,6 +173,79 @@ impl Database {
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
pub(crate) fn get_auto_vacuum_mode(conn: &Connection) -> Result<i32, AppError> {
|
||||
conn.query_row("PRAGMA auto_vacuum;", [], |row| row.get(0))
|
||||
.map_err(|e| AppError::Database(format!("读取 auto_vacuum 失败: {e}")))
|
||||
}
|
||||
|
||||
fn has_user_tables(conn: &Connection) -> Result<bool, AppError> {
|
||||
let count: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%'",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("读取表数量失败: {e}")))?;
|
||||
Ok(count > 0)
|
||||
}
|
||||
|
||||
pub(crate) fn ensure_incremental_auto_vacuum_on_conn(
|
||||
conn: &Connection,
|
||||
) -> Result<bool, AppError> {
|
||||
let mode = Self::get_auto_vacuum_mode(conn)?;
|
||||
if mode == 2 {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let has_tables = Self::has_user_tables(conn)?;
|
||||
conn.execute("PRAGMA auto_vacuum = INCREMENTAL;", [])
|
||||
.map_err(|e| AppError::Database(format!("设置 auto_vacuum 失败: {e}")))?;
|
||||
|
||||
if !has_tables {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
conn.execute("VACUUM;", [])
|
||||
.map_err(|e| AppError::Database(format!("执行 VACUUM 失败: {e}")))?;
|
||||
conn.execute("PRAGMA foreign_keys = ON;", [])
|
||||
.map_err(|e| AppError::Database(format!("恢复 foreign_keys 失败: {e}")))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub(crate) fn ensure_incremental_auto_vacuum(&self) -> Result<bool, AppError> {
|
||||
let mode = {
|
||||
let conn = lock_conn!(self.conn);
|
||||
Self::get_auto_vacuum_mode(&conn)?
|
||||
};
|
||||
if mode == 2 {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let has_tables = {
|
||||
let conn = lock_conn!(self.conn);
|
||||
Self::has_user_tables(&conn)?
|
||||
};
|
||||
if has_tables {
|
||||
log::info!(
|
||||
"Detected auto_vacuum={mode}, rebuilding database to enable incremental vacuum"
|
||||
);
|
||||
self.backup_database_file()?;
|
||||
}
|
||||
|
||||
let rebuilt = {
|
||||
let conn = lock_conn!(self.conn);
|
||||
Self::ensure_incremental_auto_vacuum_on_conn(&conn)?
|
||||
};
|
||||
|
||||
if rebuilt {
|
||||
log::info!("Incremental auto-vacuum enabled after database rebuild");
|
||||
} else {
|
||||
log::info!("Incremental auto-vacuum configured for new database");
|
||||
}
|
||||
|
||||
Ok(rebuilt)
|
||||
}
|
||||
|
||||
/// 检查 MCP 服务器表是否为空
|
||||
pub fn is_mcp_table_empty(&self) -> Result<bool, AppError> {
|
||||
let conn = lock_conn!(self.conn);
|
||||
|
||||
@@ -241,6 +241,27 @@ impl Database {
|
||||
)
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
|
||||
// 17. Usage Daily Rollups 表 (日聚合统计)
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS usage_daily_rollups (
|
||||
date TEXT NOT NULL,
|
||||
app_type TEXT NOT NULL,
|
||||
provider_id TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
request_count INTEGER NOT NULL DEFAULT 0,
|
||||
success_count INTEGER NOT NULL DEFAULT 0,
|
||||
input_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
output_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_cost_usd TEXT NOT NULL DEFAULT '0',
|
||||
avg_latency_ms INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (date, app_type, provider_id, model)
|
||||
)",
|
||||
[],
|
||||
)
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
|
||||
// 尝试添加 live_takeover_active 列到 proxy_config 表
|
||||
let _ = conn.execute(
|
||||
"ALTER TABLE proxy_config ADD COLUMN live_takeover_active INTEGER NOT NULL DEFAULT 0",
|
||||
@@ -360,6 +381,11 @@ impl Database {
|
||||
Self::migrate_v4_to_v5(conn)?;
|
||||
Self::set_user_version(conn, 5)?;
|
||||
}
|
||||
5 => {
|
||||
log::info!("迁移数据库从 v5 到 v6(使用量聚合表)");
|
||||
Self::migrate_v5_to_v6(conn)?;
|
||||
Self::set_user_version(conn, 6)?;
|
||||
}
|
||||
_ => {
|
||||
return Err(AppError::Database(format!(
|
||||
"未知的数据库版本 {version},无法迁移到 {SCHEMA_VERSION}"
|
||||
@@ -914,6 +940,32 @@ impl Database {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// v5 -> v6 迁移:添加使用量日聚合表
|
||||
fn migrate_v5_to_v6(conn: &Connection) -> Result<(), AppError> {
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS usage_daily_rollups (
|
||||
date TEXT NOT NULL,
|
||||
app_type TEXT NOT NULL,
|
||||
provider_id TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
request_count INTEGER NOT NULL DEFAULT 0,
|
||||
success_count INTEGER NOT NULL DEFAULT 0,
|
||||
input_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
output_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_cost_usd TEXT NOT NULL DEFAULT '0',
|
||||
avg_latency_ms INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (date, app_type, provider_id, model)
|
||||
)",
|
||||
[],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("创建 usage_daily_rollups 表失败: {e}")))?;
|
||||
|
||||
log::info!("v5 -> v6 迁移完成:已添加使用量日聚合表");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 插入默认模型定价数据
|
||||
/// 格式: (model_id, display_name, input, output, cache_read, cache_creation)
|
||||
/// 注意: model_id 使用短横线格式(如 claude-haiku-4-5),与 API 返回的模型名称标准化后一致
|
||||
|
||||
@@ -9,6 +9,7 @@ use indexmap::IndexMap;
|
||||
use rusqlite::{params, Connection};
|
||||
use serde_json::json;
|
||||
use std::collections::HashMap;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
const LEGACY_SCHEMA_SQL: &str = r#"
|
||||
CREATE TABLE providers (
|
||||
@@ -633,3 +634,32 @@ fn schema_model_pricing_is_seeded_on_init() {
|
||||
gemini_count
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ensure_incremental_auto_vacuum_rebuilds_existing_file_db() {
|
||||
let temp = NamedTempFile::new().expect("create temp db file");
|
||||
let path = temp.path().to_path_buf();
|
||||
|
||||
let conn = Connection::open(&path).expect("open temp db");
|
||||
conn.execute("PRAGMA auto_vacuum = NONE;", [])
|
||||
.expect("set none auto_vacuum");
|
||||
Database::create_tables_on_conn(&conn).expect("create tables");
|
||||
|
||||
assert_eq!(
|
||||
Database::get_auto_vacuum_mode(&conn).expect("auto_vacuum before rebuild"),
|
||||
0,
|
||||
"existing file db should start with NONE auto_vacuum"
|
||||
);
|
||||
|
||||
let rebuilt =
|
||||
Database::ensure_incremental_auto_vacuum_on_conn(&conn).expect("enable incremental mode");
|
||||
assert!(rebuilt, "existing db should require rebuild via VACUUM");
|
||||
drop(conn);
|
||||
|
||||
let reopened = Connection::open(&path).expect("reopen temp db");
|
||||
assert_eq!(
|
||||
Database::get_auto_vacuum_mode(&reopened).expect("auto_vacuum after rebuild"),
|
||||
2,
|
||||
"file db should persist INCREMENTAL auto_vacuum after VACUUM rebuild"
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user