YoursFunny/TelegramTwitterMediaBot
1
0
Fork 0
mirror of https://github.com/TheFunny/TelegramTwitterMediaBot.git synced 2026-09-23 23:32:05 +00:00
Code Releases Activity

feat: add admin-only /clear_cache command

Browse Source 
/clear_cache with no argument wipes the whole link_cache table;
with a post URL it removes that single entry (normalized via
site::cache_key so fxtwitter/mobile/photo variants collide with
the write-side key). Non-admins get 'Admin only.'. LinkCache gains
clear(Option<&str>) -> usize reporting removed rows.
This commit is contained in:
YoursFunny
2026-08-07 16:10:25 +08:00
parent b0ced34b4c
commit 063e910473
2 changed files with 201 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/xmedia-bot/src/handlers.rs
+116 -21
View File
@@ -5,20 +5,19 @@ use crate::send::{self, MediaItemPayload, Task};
use crate::state::{ChatData, ChatStore, unix_now};
use std::collections::HashSet;
use std::sync::LazyLock;
use teloxide::RequestError;
use teloxide::prelude::*;
use tokio::sync::Semaphore;
use teloxide::types::{
CallbackQuery, ChatAction, ChatId, ChatKind, InlineQuery, InlineQueryResult,
InlineQueryResultMpeg4Gif, InlineQueryResultPhoto, InlineQueryResultVideo, Message,
MessageEntityKind, MessageId, ParseMode, Recipient, ReplyParameters,
};
use teloxide::utils::command::BotCommands;
use teloxide::RequestError;
use tokio::sync::Semaphore;
use x_media::media::Media;
pub static CHAT_STORE: LazyLock<ChatStore> = LazyLock::new(|| {
ChatStore::open("data/task_queue.db").expect("failed to open chat store")
});
pub static CHAT_STORE: LazyLock<ChatStore> =
LazyLock::new(|| ChatStore::open("data/task_queue.db").expect("failed to open chat store"));
pub static TASK_QUEUE: LazyLock<PersistentTaskQueue> =
LazyLock::new(|| PersistentTaskQueue::new("data/task_queue.db"));
pub static LINK_CACHE: LazyLock<LinkCache> =
@@ -34,24 +33,38 @@ pub static CONFIG: LazyLock<Config> = LazyLock::new(Config::load);
static URL_TASKS: LazyLock<Semaphore> = LazyLock::new(|| Semaphore::new(8));
#[derive(BotCommands, Clone)]
#[command(rename_rule = "snake_case", description = "Turn X/Pixiv/Bluesky links into media messages")]
#[command(
rename_rule = "snake_case",
description = "Turn X/Pixiv/Bluesky links into media messages"
)]
enum Command {
#[command(description = "Get started")]
Start,
#[command(description = "Show command help")]
Help,
#[command(description = "Set forward channel (@channel or ID)", parse_with = "split")]
#[command(
description = "Set forward channel (@channel or ID)",
parse_with = "split"
)]
SetForwardChannel(String),
#[command(description = "Remove forward channel")]
RemoveForwardChannel,
#[command(description = "Toggle edit-before-forward")]
EditBeforeForward,
#[command(description = "Reply with [] to save as template", parse_with = "split")]
#[command(
description = "Reply with [] to save as template",
parse_with = "split"
)]
SetTemplate(String),
#[command(description = "Show chat state (debug)")]
BotDict,
#[command(description = "Set site caption format", parse_with = "split")]
SetFormat(String),
#[command(
description = "Clear link cache (admin; optional URL, else all)",
parse_with = "split"
)]
ClearCache(String),
}
async fn reply<T>(bot: Bot, message: Message, text: T) -> Result<Message, RequestError>
@@ -197,7 +210,11 @@ async fn set_forward_channel_handler(
Ok(channel_id)
}
async fn execute_command(bot: &Bot, message: &Message, command: Command) -> Result<(), RequestError> {
async fn execute_command(
bot: &Bot,
message: &Message,
command: Command,
) -> Result<(), RequestError> {
match command {
Command::Start => {
bot.send_message(message.chat.id, "Hello!").await?;
@@ -215,7 +232,8 @@ async fn execute_command(bot: &Bot, message: &Message, command: Command) -> Resu
"Add successfully.".to_string()
}
Err(SetForwardChannelError::EmptyParameter) => {
"Receive empty parameter.\nYou should enter a channel id or username".to_string()
"Receive empty parameter.\nYou should enter a channel id or username"
.to_string()
}
Err(SetForwardChannelError::NotChannel) => {
"Given id / username is not a channel".to_string()
@@ -292,7 +310,9 @@ async fn execute_command(bot: &Bot, message: &Message, command: Command) -> Resu
Command::SetFormat(arg) => {
let chat_id = message.chat.id.0;
let (site, format) = match arg.split_once(char::is_whitespace) {
Some((site, format)) if !format.trim().is_empty() => (site.trim(), format.trim().to_string()),
Some((site, format)) if !format.trim().is_empty() => {
(site.trim(), format.trim().to_string())
}
_ => {
reply(
bot.clone(),
@@ -317,10 +337,62 @@ async fn execute_command(bot: &Bot, message: &Message, command: Command) -> Resu
CHAT_STORE.set(chat_id, &chat_data).await;
reply(bot.clone(), message.clone(), "Format set.").await?;
}
Command::ClearCache(arg) => {
let sender_id = message
.from
.as_ref()
.map(|user| user.id.0 as i64)
.unwrap_or(-1);
if !CONFIG.admin_ids.contains(&sender_id) {
reply(bot.clone(), message.clone(), "Admin only.").await?;
return Ok(());
}
let arg = arg.trim();
if arg.is_empty() {
let removed = LINK_CACHE.clear(None).await;
log::info!("cache cleared by {sender_id}: {removed} entries");
reply(
bot.clone(),
message.clone(),
format!("Cleared {removed} cached entr{}.", plural(removed)),
)
.await?;
} else {
let key = match x_media::site::cache_key(arg) {
Some(key) => key,
None => {
reply(
bot.clone(),
message.clone(),
"Unrecognized link. Use a twitter/x, pixiv or bsky post URL.",
)
.await?;
return Ok(());
}
};
let removed = LINK_CACHE.clear(Some(&key)).await;
log::info!("cache entry cleared by {sender_id}: {key} ({removed} rows)");
reply(
bot.clone(),
message.clone(),
format!(
"Cleared cache for {arg} ({} entr{}).",
removed,
plural(removed)
),
)
.await?;
}
}
}
Ok(())
}
/// `""` for one, `"ies"` for anything else — "1 entry" / "2 entries".
fn plural(n: usize) -> &'static str {
if n == 1 { "" } else { "ies" }
}
/// For locally produced media (encoded ugoira MP4) the thumbnail URL is a
/// hotlink-protected remote URL Telegram may not fetch; let Telegram generate
/// its own thumbnail instead.
@@ -383,7 +455,10 @@ async fn dispatch_send(bot: Bot, message: &Message, task: &Task, url: &str) {
log::info!("sent {} message(s) for {url}", message_ids.len());
send::post_send_actions(&bot, task, message_ids).await;
}
Err(send::SendError::Retryable { delay_seconds, task }) => {
Err(send::SendError::Retryable {
delay_seconds,
task,
}) => {
log::info!("send for {url} failed, queued for retry in {delay_seconds:.1}s");
enqueue_retry(task, delay_seconds).await;
let _ = reply(bot, message.clone(), "Send failed. Task queued for retry.").await;
@@ -444,7 +519,10 @@ fn build_send_task(
async fn url_media(bot: Bot, message: &Message, url: &str) {
let chat_id = message.chat.id.0;
if let Err(e) = bot.send_chat_action(ChatId(chat_id), ChatAction::Typing).await {
if let Err(e) = bot
.send_chat_action(ChatId(chat_id), ChatAction::Typing)
.await
{
log::error!("send_chat_action failed: {e}");
}
@@ -520,7 +598,12 @@ async fn url_media(bot: Bot, message: &Message, url: &str) {
// Retries exhausted: notify the user (Rust-only requirement 3).
Err(e) => {
log::error!("fetch {url}: {e}");
let _ = reply(bot, message.clone(), "Failed to fetch media from this link.").await;
let _ = reply(
bot,
message.clone(),
"Failed to fetch media from this link.",
)
.await;
}
Ok(Some(fetched)) => {
if fetched.media.is_empty() {
@@ -542,8 +625,9 @@ async fn url_media(bot: Bot, message: &Message, url: &str) {
let caption = fetched.caption_with(&format);
// Raw render data for the link cache; the send fills in the
// Telegram file ids and persists the entry.
let cache_data = fetched.render_fields().map(|(author, author_url, title, tags)| {
CachedPost {
let cache_data = fetched
.render_fields()
.map(|(author, author_url, title, tags)| CachedPost {
url: fetched.source_url.clone(),
caption: fetched.caption.clone(),
title: title.to_string(),
@@ -552,7 +636,6 @@ async fn url_media(bot: Bot, message: &Message, url: &str) {
tags: tags.to_string(),
sensitive: fetched.sensitive,
media: vec![],
}
});
let items: Vec<MediaItemPayload> = fetched
.media
@@ -583,7 +666,10 @@ pub async fn message_handler(bot: Bot, message: Message) -> Result<(), RequestEr
.text()
.map(|t| if t.len() > 120 { &t[..120] } else { t })
.unwrap_or("<no text>");
log::info!("message from {sender} in {} (private={is_private}): {text_preview}", message.chat.id);
log::info!(
"message from {sender} in {} (private={is_private}): {text_preview}",
message.chat.id
);
// URL/edit flows only run in private chats; commands run in any chat.
if is_private && edit_message_handler(&bot, &message).await {
return respond(());
@@ -686,7 +772,10 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<()
let mut chat_data = CHAT_STORE.get(chat_id).await;
let edit = chat_data.edit_message.get(&prompt_message_id).cloned();
let Some(edit) = edit else {
log::info!("callback from {}: no edit record for prompt {prompt_message_id}", chat_id);
log::info!(
"callback from {}: no edit record for prompt {prompt_message_id}",
chat_id
);
bot.answer_callback_query(callback_query_id)
.text("Expired")
.await?;
@@ -705,7 +794,10 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<()
let Some(data) = data else {
return respond(());
};
log::info!("callback from {} on prompt {prompt_message_id}: {data}", chat_id);
log::info!(
"callback from {} on prompt {prompt_message_id}: {data}",
chat_id
);
if data == "forward" {
match chat_data.forward_channel_id {
Some(channel_id) => {
@@ -731,7 +823,10 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<()
chat_data.edit_message.remove(&prompt_message_id);
CHAT_STORE.set(chat_id, &chat_data).await;
}
Err(send::SendError::Retryable { delay_seconds, task }) => {
Err(send::SendError::Retryable {
delay_seconds,
task,
}) => {
log::info!("forward queued for retry in {delay_seconds:.1}s");
enqueue_retry(task, delay_seconds).await;
bot.answer_callback_query(callback_query_id)
crates/xmedia-bot/src/link_cache.rs
+81 -10
View File
@@ -8,7 +8,7 @@
//! [`Config::link_cache_ttl`]; a stale entry is dropped lazily on read and
//! by the periodic prune in `main`.
use rusqlite::{params, Connection};
use rusqlite::{Connection, params};
use serde::{Deserialize, Serialize};
use std::time::Duration;
@@ -143,6 +143,24 @@ impl LinkCache {
}
}
}
/// Deletes one entry (by normalized cache key) or the whole cache when
/// `key` is `None`. Returns how many rows were removed.
pub async fn clear(&self, key: Option<&str>) -> usize {
let key = key.map(str::to_string);
let result = crate::db::with_conn(&self.db_path, move |conn| match &key {
Some(key) => conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key]),
None => conn.execute("DELETE FROM link_cache", []),
})
.await;
match result {
Ok(n) => n,
Err(e) => {
log::error!("link cache clear failed: {e}");
0
}
}
}
}
fn now_f64() -> f64 {
@@ -192,14 +210,21 @@ mod tests {
// Force the row into the past so a 1s TTL expires it.
{
let conn = Connection::open(dir.path().join("c.db")).unwrap();
conn.execute(
"UPDATE link_cache SET created_at = created_at - 100",
[],
)
conn.execute("UPDATE link_cache SET created_at = created_at - 100", [])
.unwrap();
}
assert!(cache.get("twitter:1", Duration::from_secs(1)).await.is_none());
assert!(cache.get("twitter:1", Duration::from_secs(3600)).await.is_none());
assert!(
cache
.get("twitter:1", Duration::from_secs(1))
.await
.is_none()
);
assert!(
cache
.get("twitter:1", Duration::from_secs(3600))
.await
.is_none()
);
}
#[tokio::test]
@@ -209,14 +234,60 @@ mod tests {
cache.put("twitter:1", &entry()).await;
cache.put("pixiv:2", &entry()).await;
cache.remove("twitter:1").await;
assert!(cache.get("twitter:1", Duration::from_secs(3600)).await.is_none());
assert!(cache.get("pixiv:2", Duration::from_secs(3600)).await.is_some());
assert!(
cache
.get("twitter:1", Duration::from_secs(3600))
.await
.is_none()
);
assert!(
cache
.get("pixiv:2", Duration::from_secs(3600))
.await
.is_some()
);
{
let conn = Connection::open(dir.path().join("c.db")).unwrap();
conn.execute("UPDATE link_cache SET created_at = created_at - 100", [])
.unwrap();
}
assert_eq!(cache.prune(Duration::from_secs(1)).await, 1);
assert!(cache.get("pixiv:2", Duration::from_secs(3600)).await.is_none());
assert!(
cache
.get("pixiv:2", Duration::from_secs(3600))
.await
.is_none()
);
}
#[tokio::test]
async fn clear_one_entry_or_all() {
let dir = tempfile::tempdir().unwrap();
let cache = LinkCache::open(dir.path().join("c.db").to_str().unwrap());
cache.put("twitter:1", &entry()).await;
cache.put("pixiv:2", &entry()).await;
// By key: only the matching row is removed.
assert_eq!(cache.clear(Some("twitter:1")).await, 1);
assert!(
cache
.get("twitter:1", Duration::from_secs(3600))
.await
.is_none()
);
assert!(
cache
.get("pixiv:2", Duration::from_secs(3600))
.await
.is_some()
);
// Whole cache: nothing left; removing an absent key deletes 0 rows.
assert_eq!(cache.clear(None).await, 1);
assert!(
cache
.get("pixiv:2", Duration::from_secs(3600))
.await
.is_none()
);
assert_eq!(cache.clear(None).await, 0);
}
}