z.to
b71f76c8b3
- 实现SQLite缓存模块，支持高效查询和存储 - 添加缓存键生成策略（基于原文+语言对的SHA256哈希） - 集成缓存到Translator类，先查缓存再调用API - 添加缓存管理命令：cache clear, cache stats, cache cleanup - 实现组合缓存清理策略（数量限制+时间过期） - 添加完整的单元测试 - 更新配置文件模板，添加缓存配置 - 更新文档和版本记录 版本: v0.5.1
338 lines
8.6 KiB
Go
338 lines
8.6 KiB
Go
package cache
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"time"
|
||
|
||
_ "github.com/mattn/go-sqlite3"
|
||
)
|
||
|
||
// SQLiteCache SQLite缓存实现
|
||
type SQLiteCache struct {
|
||
db *sql.DB
|
||
config *CacheConfig
|
||
cleanupTTL time.Duration
|
||
}
|
||
|
||
// NewSQLiteCache 创建SQLite缓存实例
|
||
func NewSQLiteCache(config *CacheConfig) (*SQLiteCache, error) {
|
||
if config == nil {
|
||
config = NewCacheConfig()
|
||
}
|
||
|
||
// 展开路径中的~符号
|
||
dbPath, err := expandPath(config.DBPath)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("无效的数据库路径: %w", err)
|
||
}
|
||
|
||
// 确保目录存在
|
||
dir := filepath.Dir(dbPath)
|
||
if err := os.MkdirAll(dir, 0755); err != nil {
|
||
return nil, fmt.Errorf("创建缓存目录失败: %w", err)
|
||
}
|
||
|
||
// 打开数据库连接
|
||
db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_synchronous=NORMAL")
|
||
if err != nil {
|
||
return nil, fmt.Errorf("打开数据库失败: %w", err)
|
||
}
|
||
|
||
// 设置连接池参数
|
||
db.SetMaxOpenConns(1) // SQLite只支持单个写入连接
|
||
db.SetMaxIdleConns(1)
|
||
|
||
cache := &SQLiteCache{
|
||
db: db,
|
||
config: config,
|
||
cleanupTTL: time.Duration(config.ExpireDays) * 24 * time.Hour,
|
||
}
|
||
|
||
// 初始化数据库表
|
||
if err := cache.initTable(); err != nil {
|
||
db.Close()
|
||
return nil, fmt.Errorf("初始化缓存表失败: %w", err)
|
||
}
|
||
|
||
// 设置清理定时器
|
||
go cache.startCleanupTimer()
|
||
|
||
return cache, nil
|
||
}
|
||
|
||
// initTable 初始化缓存表
|
||
func (c *SQLiteCache) initTable() error {
|
||
query := `
|
||
CREATE TABLE IF NOT EXISTS translation_cache (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
cache_key TEXT NOT NULL UNIQUE,
|
||
original_text TEXT NOT NULL,
|
||
translated_text TEXT NOT NULL,
|
||
from_lang TEXT NOT NULL,
|
||
to_lang TEXT NOT NULL,
|
||
model TEXT NOT NULL,
|
||
prompt_name TEXT,
|
||
prompt_content TEXT,
|
||
prompt_tokens INTEGER DEFAULT 0,
|
||
completion_tokens INTEGER DEFAULT 0,
|
||
total_tokens INTEGER DEFAULT 0,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
last_used_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_cache_key ON translation_cache(cache_key);
|
||
CREATE INDEX IF NOT EXISTS idx_original_text ON translation_cache(original_text);
|
||
CREATE INDEX IF NOT EXISTS idx_created_at ON translation_cache(created_at);
|
||
CREATE INDEX IF NOT EXISTS idx_last_used_at ON translation_cache(last_used_at);
|
||
`
|
||
|
||
_, err := c.db.Exec(query)
|
||
return err
|
||
}
|
||
|
||
// Get 获取缓存
|
||
func (c *SQLiteCache) Get(ctx context.Context, key string) (*CacheEntry, error) {
|
||
query := `
|
||
SELECT
|
||
id, cache_key, original_text, translated_text, from_lang, to_lang,
|
||
model, prompt_name, prompt_content, prompt_tokens, completion_tokens,
|
||
total_tokens, created_at, last_used_at
|
||
FROM translation_cache
|
||
WHERE cache_key = ?
|
||
`
|
||
|
||
entry := &CacheEntry{}
|
||
var promptName, promptContent sql.NullString
|
||
var createdAt, lastUsedAt string
|
||
|
||
err := c.db.QueryRowContext(ctx, query, key).Scan(
|
||
&entry.ID, &entry.CacheKey, &entry.OriginalText, &entry.TranslatedText,
|
||
&entry.FromLang, &entry.ToLang, &entry.Model, &promptName, &promptContent,
|
||
&entry.PromptTokens, &entry.CompletionTokens, &entry.TotalTokens,
|
||
&createdAt, &lastUsedAt,
|
||
)
|
||
|
||
if err == sql.ErrNoRows {
|
||
return nil, nil // 缓存未命中
|
||
}
|
||
if err != nil {
|
||
return nil, fmt.Errorf("查询缓存失败: %w", err)
|
||
}
|
||
|
||
// 处理可空字段
|
||
if promptName.Valid {
|
||
entry.PromptName = promptName.String
|
||
}
|
||
if promptContent.Valid {
|
||
entry.PromptContent = promptContent.String
|
||
}
|
||
|
||
// 解析时间
|
||
entry.CreatedAt, _ = time.Parse("2006-01-02 15:04:05", createdAt)
|
||
entry.LastUsedAt, _ = time.Parse("2006-01-02 15:04:05", lastUsedAt)
|
||
|
||
// 更新最后使用时间
|
||
go c.updateLastUsed(context.Background(), key)
|
||
|
||
return entry, nil
|
||
}
|
||
|
||
// Set 设置缓存
|
||
func (c *SQLiteCache) Set(ctx context.Context, entry *CacheEntry) error {
|
||
// 开始事务
|
||
tx, err := c.db.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return fmt.Errorf("开始事务失败: %w", err)
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
// 插入或替换缓存
|
||
query := `
|
||
INSERT OR REPLACE INTO translation_cache
|
||
(cache_key, original_text, translated_text, from_lang, to_lang,
|
||
model, prompt_name, prompt_content, prompt_tokens, completion_tokens,
|
||
total_tokens, created_at, last_used_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
`
|
||
|
||
now := time.Now().Format("2006-01-02 15:04:05")
|
||
_, err = tx.ExecContext(ctx, query,
|
||
entry.CacheKey, entry.OriginalText, entry.TranslatedText,
|
||
entry.FromLang, entry.ToLang, entry.Model, entry.PromptName,
|
||
entry.PromptContent, entry.PromptTokens, entry.CompletionTokens,
|
||
entry.TotalTokens, now, now,
|
||
)
|
||
if err != nil {
|
||
return fmt.Errorf("插入缓存失败: %w", err)
|
||
}
|
||
|
||
// 提交事务
|
||
if err := tx.Commit(); err != nil {
|
||
return fmt.Errorf("提交事务失败: %w", err)
|
||
}
|
||
|
||
// 触发清理（异步）
|
||
go c.Cleanup(context.Background())
|
||
|
||
return nil
|
||
}
|
||
|
||
// Delete 删除缓存
|
||
func (c *SQLiteCache) Delete(ctx context.Context, key string) error {
|
||
query := `DELETE FROM translation_cache WHERE cache_key = ?`
|
||
_, err := c.db.ExecContext(ctx, query, key)
|
||
if err != nil {
|
||
return fmt.Errorf("删除缓存失败: %w", err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// Clear 清空缓存
|
||
func (c *SQLiteCache) Clear(ctx context.Context) error {
|
||
// 先删除所有记录
|
||
_, err := c.db.ExecContext(ctx, `DELETE FROM translation_cache`)
|
||
if err != nil {
|
||
return fmt.Errorf("清空缓存失败: %w", err)
|
||
}
|
||
|
||
// 然后执行VACUUM（不能在事务中执行）
|
||
_, err = c.db.ExecContext(ctx, `VACUUM`)
|
||
if err != nil {
|
||
return fmt.Errorf("清理数据库失败: %w", err)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// Stats 获取缓存统计信息
|
||
func (c *SQLiteCache) Stats(ctx context.Context) (*CacheStats, error) {
|
||
stats := &CacheStats{}
|
||
|
||
// 获取总记录数
|
||
err := c.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM translation_cache`).Scan(&stats.TotalRecords)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("查询缓存统计失败: %w", err)
|
||
}
|
||
|
||
// 如果没有记录，直接返回
|
||
if stats.TotalRecords == 0 {
|
||
return stats, nil
|
||
}
|
||
|
||
// 获取时间范围和平均tokens
|
||
var oldestStr, newestStr sql.NullString
|
||
var avgTokens sql.NullFloat64
|
||
err = c.db.QueryRowContext(ctx, `
|
||
SELECT
|
||
MIN(created_at),
|
||
MAX(created_at),
|
||
AVG(total_tokens)
|
||
FROM translation_cache
|
||
`).Scan(&oldestStr, &newestStr, &avgTokens)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("查询缓存时间范围失败: %w", err)
|
||
}
|
||
|
||
// 解析时间字符串
|
||
if oldestStr.Valid {
|
||
stats.OldestRecord, _ = time.Parse("2006-01-02 15:04:05", oldestStr.String)
|
||
}
|
||
if newestStr.Valid {
|
||
stats.NewestRecord, _ = time.Parse("2006-01-02 15:04:05", newestStr.String)
|
||
}
|
||
if avgTokens.Valid {
|
||
stats.AvgTokensPerRecord = avgTokens.Float64
|
||
}
|
||
|
||
// 计算数据库文件大小
|
||
dbPath, _ := expandPath(c.config.DBPath)
|
||
if info, err := os.Stat(dbPath); err == nil {
|
||
stats.TotalSizeBytes = info.Size()
|
||
}
|
||
|
||
return stats, nil
|
||
}
|
||
|
||
// Cleanup 清理过期缓存
|
||
func (c *SQLiteCache) Cleanup(ctx context.Context) error {
|
||
tx, err := c.db.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return fmt.Errorf("开始事务失败: %w", err)
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
// 清理过期缓存
|
||
if c.cleanupTTL > 0 {
|
||
expiredTime := time.Now().Add(-c.cleanupTTL).Format("2006-01-02 15:04:05")
|
||
_, err = tx.ExecContext(ctx, `DELETE FROM translation_cache WHERE last_used_at < ?`, expiredTime)
|
||
if err != nil {
|
||
return fmt.Errorf("清理过期缓存失败: %w", err)
|
||
}
|
||
} else if c.cleanupTTL == 0 {
|
||
// 如果过期时间为0，清理所有记录
|
||
_, err = tx.ExecContext(ctx, `DELETE FROM translation_cache`)
|
||
if err != nil {
|
||
return fmt.Errorf("清理所有缓存失败: %w", err)
|
||
}
|
||
}
|
||
|
||
// 清理超出数量限制的缓存
|
||
if c.config.MaxRecords > 0 {
|
||
_, err = tx.ExecContext(ctx, `
|
||
DELETE FROM translation_cache
|
||
WHERE id NOT IN (
|
||
SELECT id FROM translation_cache
|
||
ORDER BY last_used_at DESC
|
||
LIMIT ?
|
||
)
|
||
`, c.config.MaxRecords)
|
||
if err != nil {
|
||
return fmt.Errorf("清理超出数量限制的缓存失败: %w", err)
|
||
}
|
||
}
|
||
|
||
return tx.Commit()
|
||
}
|
||
|
||
// Close 关闭缓存
|
||
func (c *SQLiteCache) Close() error {
|
||
if c.db != nil {
|
||
return c.db.Close()
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// updateLastUsed 更新最后使用时间
|
||
func (c *SQLiteCache) updateLastUsed(ctx context.Context, key string) {
|
||
query := `UPDATE translation_cache SET last_used_at = ? WHERE cache_key = ?`
|
||
now := time.Now().Format("2006-01-02 15:04:05")
|
||
c.db.ExecContext(ctx, query, now, key)
|
||
}
|
||
|
||
// startCleanupTimer 启动清理定时器
|
||
func (c *SQLiteCache) startCleanupTimer() {
|
||
ticker := time.NewTicker(1 * time.Hour) // 每小时清理一次
|
||
defer ticker.Stop()
|
||
|
||
for range ticker.C {
|
||
c.Cleanup(context.Background())
|
||
}
|
||
}
|
||
|
||
// expandPath 展开路径中的~符号
|
||
func expandPath(path string) (string, error) {
|
||
if strings.HasPrefix(path, "~") {
|
||
home, err := os.UserHomeDir()
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
path = filepath.Join(home, path[1:])
|
||
}
|
||
return path, nil
|
||
}
|