diff --git a/core/database.go b/core/database.go index f9134cf..cd2f4bd 100644 --- a/core/database.go +++ b/core/database.go @@ -44,8 +44,11 @@ func NewDatabase() (*Database, error) { func (d *Database) createTables() error { queries := []string{ `CREATE TABLE IF NOT EXISTS keywords - (id INTEGER PRIMARY KEY, keyword TEXT UNIQUE)`, + (id INTEGER PRIMARY KEY, keyword TEXT UNIQUE, is_link BOOLEAN, is_auto_added BOOLEAN, added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)`, `CREATE INDEX IF NOT EXISTS idx_keyword ON keywords(keyword)`, + `CREATE INDEX IF NOT EXISTS idx_added_at ON keywords(added_at)`, + `CREATE INDEX IF NOT EXISTS idx_is_link ON keywords(is_link)`, + `CREATE INDEX IF NOT EXISTS idx_is_auto_added ON keywords(is_auto_added)`, `CREATE TABLE IF NOT EXISTS whitelist (id INTEGER PRIMARY KEY, domain TEXT UNIQUE)`, `CREATE INDEX IF NOT EXISTS idx_domain ON whitelist(domain)`, @@ -82,8 +85,9 @@ func (d *Database) executeQuery(query string, args ...interface{}) ([]string, er return results, nil } -func (d *Database) AddKeyword(keyword string) error { - _, err := d.db.Exec("INSERT OR IGNORE INTO keywords (keyword) VALUES (?)", keyword) +func (d *Database) AddKeyword(keyword string, isLink bool, isAutoAdded bool) error { + _, err := d.db.Exec("INSERT OR IGNORE INTO keywords (keyword, is_link, is_auto_added, added_at) VALUES (?, ?, ?, ?)", + keyword, isLink, isAutoAdded, time.Now()) if err != nil { return err } @@ -104,6 +108,16 @@ func (d *Database) RemoveKeyword(keyword string) (bool, error) { return rowsAffected > 0, nil } +func (d *Database) CleanupExpiredLinks() error { + twoMonthsAgo := time.Now().AddDate(0, -2, 0) + _, err := d.db.Exec("DELETE FROM keywords WHERE is_link = TRUE AND is_auto_added = TRUE AND added_at < ?", twoMonthsAgo) + if err != nil { + return err + } + d.invalidateCache("keywords") + return nil +} + func (d *Database) GetAllKeywords() ([]string, error) { d.mu.Lock() defer d.mu.Unlock() @@ -119,6 +133,13 @@ func (d *Database) GetAllKeywords() ([]string, error) { return d.keywordsCache, nil } +func (d *Database) GetAllManualKeywords() ([]string, error) { + return d.executeQuery("SELECT keyword FROM keywords WHERE is_auto_added = ?", false) +} + +func (d *Database) GetAllAutoAddedLinks() ([]string, error) { + return d.executeQuery("SELECT keyword FROM keywords WHERE is_link = ? AND is_auto_added = ?", true, true) +} func (d *Database) RemoveKeywordsContaining(substring string) ([]string, error) { // 首先获取要删除的关键词列表 @@ -308,3 +329,116 @@ func (d *Database) invalidateCache(cacheType string) { func (d *Database) Close() error { return d.db.Close() } + +// 迁移现有关键词 +// +// +// +// +// +// +// +// +// +// +// +// +// +// + +func (d *Database) MigrateExistingKeywords() error { + // 检查是否已经执行过迁移 + var migrationDone bool + err := d.db.QueryRow("SELECT value FROM config WHERE key = 'keywords_migrated'").Scan(&migrationDone) + if err != nil && err != sql.ErrNoRows { + return err + } + + if migrationDone { + return nil // 迁移已经完成,无需再次执行 + } + + // 开始事务 + tx, err := d.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + // 获取所有现有的关键词 + rows, err := tx.Query("SELECT keyword FROM keywords") + if err != nil { + return err + } + defer rows.Close() + + // 准备插入语句 + stmt, err := tx.Prepare("INSERT INTO keywords (keyword, is_link, is_auto_added) VALUES (?, ?, ?)") + if err != nil { + return err + } + defer stmt.Close() + + // 迁移现有关键词 + for rows.Next() { + var keyword string + if err := rows.Scan(&keyword); err != nil { + return err + } + + // 这里我们假设所有现有的关键词都是手动添加的非链接关键词 + _, err = stmt.Exec(keyword, false, false) + if err != nil { + return err + } + } + + // 删除旧的关键词表 + _, err = tx.Exec("DROP TABLE IF EXISTS old_keywords") + if err != nil { + return err + } + + // 重命名现有的关键词表 + _, err = tx.Exec("ALTER TABLE keywords RENAME TO old_keywords") + if err != nil { + return err + } + + // 创建新的关键词表 + _, err = tx.Exec(`CREATE TABLE keywords ( + id INTEGER PRIMARY KEY, + keyword TEXT UNIQUE, + is_link BOOLEAN, + is_auto_added BOOLEAN, + added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + )`) + if err != nil { + return err + } + + // 将数据从旧表复制到新表 + _, err = tx.Exec("INSERT INTO keywords SELECT id, keyword, is_link, is_auto_added, added_at FROM old_keywords") + if err != nil { + return err + } + + // 删除旧表 + _, err = tx.Exec("DROP TABLE old_keywords") + if err != nil { + return err + } + + // 更新配置,标记迁移已完成 + _, err = tx.Exec("INSERT OR REPLACE INTO config (key, value) VALUES ('keywords_migrated', 'true')") + if err != nil { + return err + } + + // 提交事务 + if err := tx.Commit(); err != nil { + return err + } + + return nil +} diff --git a/core/init.go b/core/init.go index 6a37fc1..a4daa6d 100644 --- a/core/init.go +++ b/core/init.go @@ -45,6 +45,12 @@ func Init() error { return fmt.Errorf("初始化数据库失败: %v", err) } + // 执行数据迁移 + err = DB.MigrateExistingKeywords() + if err != nil { + return fmt.Errorf("迁移现有关键词失败: %v", err) + } + // 从环境变量中读取调试模式设置 DEBUG_MODE = os.Getenv("DEBUG_MODE") == "true" diff --git a/main.go b/main.go index ca44895..474a2a2 100644 --- a/main.go +++ b/main.go @@ -23,4 +23,7 @@ func main() { if err != nil { log.Fatalf("Error in RunMessageHandler: %v", err) } + + // 启动定期任务 + service.StartScheduledTasks() } diff --git a/service/link_filter/link_filter.go b/service/link_filter/link_filter.go index cd64b53..456b810 100644 --- a/service/link_filter/link_filter.go +++ b/service/link_filter/link_filter.go @@ -124,7 +124,7 @@ func addNewKeyword(keyword string) error { return fmt.Errorf("检查关键词时发生错误: %v", err) } if !exists { - err = core.DB.AddKeyword(keyword) + err = core.DB.AddKeyword(keyword, true, true) // isLink = true, isAutoAdded = true if err != nil { return fmt.Errorf("添加关键词时发生错误: %v", err) } @@ -132,6 +132,7 @@ func addNewKeyword(keyword string) error { } return nil } + func containsKeyword(text string, linkFilter *LinkFilter) bool { linkFilter.Mu.RLock() defer linkFilter.Mu.RUnlock() diff --git a/service/message_handler_utils.go b/service/message_handler_utils.go index 2245b7c..605f238 100644 --- a/service/message_handler_utils.go +++ b/service/message_handler_utils.go @@ -27,23 +27,41 @@ func HandleKeywordCommand(bot *tgbotapi.BotAPI, message *tgbotapi.Message, comma } func handleListKeywords(bot *tgbotapi.BotAPI, message *tgbotapi.Message) { - keywords, err := core.DB.GetAllKeywords() + manualKeywords, err := core.DB.GetAllManualKeywords() if err != nil { - core.SendErrorMessage(bot, message.Chat.ID, "获取关键词列表时发生错误。") - logger.Printf("Failed to get keywords: %v", err) + core.SendErrorMessage(bot, message.Chat.ID, "获取手动添加的关键词列表时发生错误。") + logger.Printf("Failed to get manual keywords: %v", err) return } - if len(keywords) == 0 { - core.SendMessage(bot, message.Chat.ID, "关键词列表为空。") - } else { - // 对关键词进行排序 - sort.Strings(keywords) + autoAddedLinks, err := core.DB.GetAllAutoAddedLinks() + if err != nil { + core.SendErrorMessage(bot, message.Chat.ID, "获取自动添加的链接列表时发生错误。") + logger.Printf("Failed to get auto-added links: %v", err) + return + } - // 直接发送排序后的关键词列表 - err := core.SendLongMessage(bot, message.Chat.ID, "当前关键词列表(按字母顺序排序):", keywords) - if err != nil { - core.SendErrorMessage(bot, message.Chat.ID, "发送关键词列表时发生错误。") + if len(manualKeywords) == 0 && len(autoAddedLinks) == 0 { + core.SendMessage(bot, message.Chat.ID, "关键词和链接列表为空。") + } else { + // 对关键词和链接进行排序 + sort.Strings(manualKeywords) + sort.Strings(autoAddedLinks) + + // 发送手动添加的关键词列表 + if len(manualKeywords) > 0 { + err := core.SendLongMessage(bot, message.Chat.ID, "手动添加的关键词列表(按字母顺序排序):", manualKeywords) + if err != nil { + core.SendErrorMessage(bot, message.Chat.ID, "发送手动添加的关键词列表时发生错误。") + } + } + + // 发送自动添加的链接列表 + if len(autoAddedLinks) > 0 { + err := core.SendLongMessage(bot, message.Chat.ID, "自动添加的链接列表(按字母顺序排序):", autoAddedLinks) + if err != nil { + core.SendErrorMessage(bot, message.Chat.ID, "发送自动添加的链接列表时发生错误。") + } } } } @@ -61,7 +79,7 @@ func handleAddKeyword(bot *tgbotapi.BotAPI, message *tgbotapi.Message, keyword s return } if !exists { - err = core.DB.AddKeyword(keyword) + err = core.DB.AddKeyword(keyword, false, false) // isLink = false, isAutoAdded = false if err != nil { core.SendErrorMessage(bot, message.Chat.ID, "添加关键词时发生错误。") logger.Printf("Failed to add keyword: %v", err) diff --git a/service/scheduled_tasks.go b/service/scheduled_tasks.go new file mode 100644 index 0000000..8e9f113 --- /dev/null +++ b/service/scheduled_tasks.go @@ -0,0 +1,29 @@ +package service + +import ( + "log" + "time" + + "github.com/woodchen-ink/Q58Bot/core" +) + +func StartScheduledTasks() { + go periodicCleanup() +} + +func periodicCleanup() { + ticker := time.NewTicker(24 * time.Hour) // 每天执行一次清理 + defer ticker.Stop() + + for { + select { + case <-ticker.C: + err := core.DB.CleanupExpiredLinks() + if err != nil { + log.Printf("清理过期链接时发生错误: %v", err) + } else { + log.Println("已成功清理过期链接") + } + } + } +}