Add a vault for saving whole forwarded posts
Completes TODO #5. Forwarding a post now asks whether to save it whole or extract its media; sending media directly is unchanged. Built on forwarding rather than copying file_ids. The bot already has the post in its private chat with the owner, so it stores message ids and forwards them back. forwardMessages takes a list and keeps an album intact, which is what made the caption-plus-album case tractable; ids must be strictly ascending, so they are sorted on the way out. - Vault rows live in their own table and every query is scoped to an owner_id, so a guessed id cannot reach someone else's saved post. Nothing in the vault is searchable inline, moderated or federated, and the separate table makes that structural rather than a flag. - Albums arrive as separate messages with no completion signal, so they are buffered by media_group_id behind a 1s debounce and prompt once. - The add path now takes a list of items, so extracting an album stores every piece under one set of keywords instead of asking repeatedly. - /vault searches keywords and the original caption, since the point of a vault is finding something you never got round to tagging. - Retrieval failures are reported: a message id is a pointer into the chat, not a self-contained file_id, so clearing history breaks it. Three near-identical keyword collectors are now one helper. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
goyban
@@ -7,10 +7,10 @@ Progress on branch `develop` (not pushed):
|
||||
- [x] **#6** full media coverage — `efc2be5`
|
||||
- [x] **#3** moderation queue — `0a8c81d`
|
||||
- [x] **#2** public mode + membership gate — `82b5da4`
|
||||
- [ ] **#5** save vs. extract (vault)
|
||||
- [x] **#5** save vs. extract (vault) — forward-based
|
||||
- [ ] **#4** federation
|
||||
|
||||
Test suites: `test_db.py`, `test_bot.py`, `test_access.py` — 20 cases, all passing.
|
||||
Test suites: `test_db.py`, `test_bot.py`, `test_access.py`, `test_vault.py` — 29 cases, all passing.
|
||||
|
||||
---
|
||||
|
||||
@@ -98,29 +98,45 @@ With `PUBLIC_MODE=true`, non-admin submissions land as `status='pending'` from #
|
||||
|
||||
---
|
||||
|
||||
## #5 — Save vs. extract on forward
|
||||
## #5 — Save vs. extract on forward — **implemented**
|
||||
|
||||
On a forwarded message, ask with inline buttons: **[📥 Save whole post]** / **[✂️ Extract media]**. Extract = today's behavior. Save = the private vault.
|
||||
Built on forwarding rather than the file_id copy originally planned here, which
|
||||
removed most of the work.
|
||||
|
||||
**Albums are the hard part.** A post with caption + multiple media arrives as *several separate messages* sharing a `media_group_id`, with no "album complete" signal. Standard approach: buffer by `media_group_id` and debounce ~1s after the last message before prompting. Ask once per album, not once per file.
|
||||
**The bot already holds the post.** When you forward something, Telegram puts a
|
||||
copy in your private chat with the bot, so the bot can forward *that* back later
|
||||
by `message_id`. It never needs to reach the original channel, store media, or
|
||||
re-upload anything.
|
||||
|
||||
**Schema (migration v4)** — a separate table, not a `media` column, since one vault item holds N files:
|
||||
**`forwardMessages` (plural) preserves albums.** It takes up to 100 message ids
|
||||
and returns them as one intact album, so the caption + N media case that made
|
||||
this item large is handled by Telegram itself. Ids must be strictly ascending.
|
||||
|
||||
**One table, not two:**
|
||||
|
||||
```sql
|
||||
vault_items(id, owner_id, caption TEXT, keywords TEXT, source_chat TEXT, created_at)
|
||||
vault_files(item_id, file_id, file_unique_id, media_type, position)
|
||||
vault_items(id, owner_id, chat_id, message_ids TEXT, caption, keywords, created_at)
|
||||
```
|
||||
|
||||
Vault rows are **always private** and never touched by the federation sync in #4 — enforce that structurally by keeping them in a different table, so no `visibility` flag can ever be mis-set.
|
||||
Every vault query takes an `owner_id` and filters on it, so a guessed id cannot
|
||||
reach another user's post. Vault rows live in their own table and are never
|
||||
searched by inline mode, never moderated, and never federated.
|
||||
|
||||
**Search should cover the caption, not just keywords** — the point of a vault is finding a post you saved without having tagged it. This is where the current `LIKE '%,kw%'` approach (`db.py:43`) stops being adequate; use an FTS5 virtual table over `caption + keywords`.
|
||||
**Flow.** A forwarded message prompts *Save whole post* / *Extract media*
|
||||
(extract only appears when there is media). Save collects keywords and stores
|
||||
the pointer; extract runs the existing add path, now generalised to a list so a
|
||||
whole album shares one set of keywords. Albums are buffered by `media_group_id`
|
||||
with a 1s debounce, since Telegram gives no "album complete" signal.
|
||||
|
||||
**Open question — recall UX.** Telegram **inline mode cannot send an album**; one inline result sends exactly one message. So a multi-file vault item can't come back through `@bot query` intact. Two options:
|
||||
**Recall is inside the bot only**, as decided: `/vault [search]` lists matching
|
||||
posts with a retrieve button and a delete button. Matching covers keywords and
|
||||
the original caption. The save confirmation tells the user plainly that vault
|
||||
posts do not appear in inline search.
|
||||
|
||||
1. Inline returns the item's *first* media with the caption attached — works everywhere, but loses the other files.
|
||||
2. A `/vault <query>` command in the bot's private chat that replies with the full album, which you then forward manually — preserves everything, costs one extra step.
|
||||
|
||||
These aren't exclusive; doing both (1 for speed, 2 for fidelity) is probably right, but it's your call. Single-file vault items are unaffected either way.
|
||||
**The tradeoff accepted:** a `message_id` is a pointer into the bot's chat with
|
||||
the owner, not self-contained like a `file_id`. If that history is cleared, the
|
||||
pointer breaks; retrieval catches the failure and tells the user to forward the
|
||||
post again.
|
||||
|
||||
---
|
||||
|
||||
@@ -164,14 +180,14 @@ The hub is a separate deployable (small FastAPI + SQLite would do). Do you want
|
||||
| 6 | Full media coverage | 0 (only for the refactor) | S |
|
||||
| 3 | Moderation queue | 0 | M |
|
||||
| 2 | Public mode + membership gate | 0, 3 | M |
|
||||
| 5 | Save vs. extract + vault | 0 | L |
|
||||
| 5 | Save vs. extract + vault | 0 | M (was L) |
|
||||
| 4 | Federation | 0, 3, hub built | XL |
|
||||
|
||||
#6 is shippable on its own and touches nothing else — good first commit on `develop`. #4 should stay last: it's the only item that's hard to reverse once other people are running instances against your hub.
|
||||
|
||||
## Decisions — settled 2026-09-01
|
||||
|
||||
1. **#5 recall UX** — inline cannot send an album. For multi-file vault items, the inline result sends the **link to the original post** where one exists; otherwise it tells the user to open the bot and pull the full album there. Single-file items send directly as today.
|
||||
1. **#5 recall UX** — settled further during implementation: the vault is **forward-based and retrievable only inside the bot**. Inline search does not surface vault posts at all; the save confirmation says so.
|
||||
2. **#4 hub location** — **separate repository.** Channel name and hub hosting are the owner's call, not decided here.
|
||||
3. **#6 video_note** — **store as `video`**, accepting the loss of round presentation, so round videos stay inline-searchable. See the risk note below.
|
||||
4. **#4 architecture** — **shared federation channel accepted.** Media rides the channel, metadata rides the hub API.
|
||||
|
||||
@@ -59,10 +59,28 @@ docker compose up -d --build
|
||||
| دستور | توضیح |
|
||||
|---|---|
|
||||
| `/list` | نمایش همه رسانههای ذخیرهشده |
|
||||
| `/vault [search]` | بازیابی پستی که کامل ذخیره کردهاید |
|
||||
| `/review` | بررسی یکبهیک ارسالهای در انتظار تأیید (فقط ادمین) |
|
||||
| `/delete <id>` | حذف یک مورد (ID را از `/list` بگیرید) |
|
||||
| `/cancel` | لغو عملیات جاری |
|
||||
|
||||
### گاوصندوق — ذخیره کامل یک پست
|
||||
|
||||
اگر رسانهای را مستقیم برای ربات بفرستید، مثل همیشه رفتار میکند: رسانه را جدا میکند و کلیدواژه میخواهد.
|
||||
|
||||
اما اگر پستی را **فوروارد** کنید، ربات میپرسد چه میخواهید:
|
||||
|
||||
- **ذخیره کل پست** — پست را دستنخورده با کپشن در گاوصندوق نگه میدارد. آلبوم هم آلبوم میماند.
|
||||
- **جدا کردن رسانه** — همان رفتار قبلی؛ رسانه به کتابخانه قابل جستجو اضافه میشود.
|
||||
|
||||
با `/vault [search]` پستهای ذخیرهشده را پیدا کنید؛ جستجو هم روی کلیدواژهها و هم روی متن کپشن اصلی انجام میشود. کنار هر نتیجه یک دکمه برای فرستادن دوباره پست و یک دکمه برای حذف آن هست.
|
||||
|
||||
> **پستهای گاوصندوق فقط داخل خود ربات برمیگردند.** حالت inline تلگرام نمیتواند پیام موجود را فوروارد کند، پس `@yourbotusername keyword` به گاوصندوق دسترسی ندارد. پست را داخل ربات بگیرید و بعد هرجا خواستید فوروارد کنید.
|
||||
|
||||
گاوصندوق برای هر کاربر جداگانه و خصوصی است: در جدولی جدا از کتابخانه میمها ذخیره میشود، هر جستجو محدود به صاحب آن است، و هیچچیزِ آن بررسی یا به اشتراک گذاشته نمیشود.
|
||||
|
||||
ذخیرهسازی بهجای کپی رسانه، اشارهای به پیام موجود در چت شما با ربات را نگه میدارد. اگر تاریخچه چتتان با ربات را پاک کنید، آن پستها دیگر قابل بازیابی نیستند — ربات در این حالت به شما اطلاع میدهد تا دوباره فوروارد کنید.
|
||||
|
||||
### عمومیکردن ربات
|
||||
|
||||
با تنظیم `PUBLIC_MODE=true` ربات بهجای فهرست `ALLOWED_USERS`، به هر کسی که شرط عضویت را داشته باشد اجازه دسترسی میدهد:
|
||||
|
||||
@@ -65,11 +65,29 @@ Telegram suggests these commands as you type `/`, with a short description besid
|
||||
| `/start` | Start the bot and show usage instructions |
|
||||
| `/menu` | Show the available commands |
|
||||
| `/list` | Show all saved media |
|
||||
| `/vault [search]` | Retrieve a whole post you saved |
|
||||
| `/review` | Review pending submissions one by one (admin only) |
|
||||
| `/edit [search]` | Select a meme and edit its keywords (admin only) |
|
||||
| `/delete <id>` | Remove an entry (get ID from `/list`; admin only) |
|
||||
| `/cancel` | Cancel current add operation |
|
||||
|
||||
### The vault — saving whole posts
|
||||
|
||||
Send the bot media directly and it behaves as it always has: it extracts the media and asks for keywords.
|
||||
|
||||
**Forward** a post to it instead and it asks what you want:
|
||||
|
||||
- **Save whole post** — keeps the post intact in your vault, caption and all. An album stays an album.
|
||||
- **Extract media** — the existing behaviour, adding the media to your searchable library.
|
||||
|
||||
Retrieve saved posts with `/vault [search]`, which matches your keywords and the original caption. Each result has a button to forward it back to you, and one to remove it.
|
||||
|
||||
> **Vault posts only come back inside this bot.** Telegram's inline mode cannot forward an existing message, so `@yourbotusername keyword` does not reach the vault. Retrieve the post in the bot chat, then forward it wherever you want.
|
||||
|
||||
The vault is per-user and private: it is stored separately from the meme library, every lookup is scoped to its owner, and nothing in it is ever moderated or shared.
|
||||
|
||||
Saving stores a pointer to the message in your chat with the bot rather than a copy of the media. If you clear your chat history with the bot, saved posts can no longer be retrieved — the bot tells you when that happens so you can forward them again.
|
||||
|
||||
### Making the bot public
|
||||
|
||||
Set `PUBLIC_MODE=true` to open the bot to anyone who meets a membership requirement, instead of the `ALLOWED_USERS` list:
|
||||
|
||||
@@ -79,6 +79,7 @@ USER_COMMANDS = [
|
||||
BotCommand(command="start", description="Start the bot"),
|
||||
BotCommand(command="menu", description="Show available commands"),
|
||||
BotCommand(command="list", description="Browse saved memes"),
|
||||
BotCommand(command="vault", description="Retrieve a saved post"),
|
||||
BotCommand(command="cancel", description="Cancel current operation"),
|
||||
]
|
||||
ADMIN_COMMANDS = [
|
||||
@@ -238,12 +239,17 @@ class EditMedia(StatesGroup):
|
||||
waiting_for_new_keywords = State()
|
||||
|
||||
|
||||
class VaultSave(StatesGroup):
|
||||
waiting_for_keywords = State()
|
||||
|
||||
|
||||
def menu_text(user_id: int) -> str:
|
||||
if is_admin(user_id):
|
||||
return (
|
||||
"Send me any <b>photo, GIF, video, voice, music, sticker, or file</b>.\n"
|
||||
"I'll ask you for keywords next.\n\n"
|
||||
"/list — show all saved memes\n"
|
||||
"/vault [search] — retrieve a whole post you saved\n"
|
||||
"/review — review pending submissions\n"
|
||||
"/edit [search] — edit keywords of a meme\n"
|
||||
"/delete <id> — remove a meme by ID\n"
|
||||
@@ -257,7 +263,8 @@ def menu_text(user_id: int) -> str:
|
||||
)
|
||||
return (
|
||||
"Use me inline: type <code>@botname keyword</code> in any chat.\n\n"
|
||||
"/list — browse all saved memes"
|
||||
"/list — browse all saved memes\n"
|
||||
"/vault [search] — retrieve a whole post you saved"
|
||||
)
|
||||
|
||||
|
||||
@@ -303,23 +310,28 @@ async def cmd_cancel(msg: Message, state: FSMContext):
|
||||
|
||||
# ── step 1: receive media (admin only) ───────────────────────────────────────
|
||||
|
||||
async def _receive_media(msg: Message, state: FSMContext, file_id: str, file_unique_id: str,
|
||||
media_type: str, status: str = "approved"):
|
||||
async def start_keyword_prompt(chat_id: int, state: FSMContext, items: list[dict],
|
||||
status: str = "approved") -> None:
|
||||
"""Begin collecting keywords for one or more media items. Everything in
|
||||
`items` ends up sharing the keywords, which is what makes extracting a
|
||||
whole album a single interaction."""
|
||||
await state.set_state(AddMedia.waiting_for_keywords)
|
||||
await state.update_data(
|
||||
file_id=file_id,
|
||||
file_unique_id=file_unique_id,
|
||||
media_type=media_type,
|
||||
keywords=[],
|
||||
status=status,
|
||||
)
|
||||
await msg.reply(
|
||||
"Got it! Now send keywords — one per message or comma-separated.\n"
|
||||
await state.update_data(items=items, keywords=[], status=status)
|
||||
count = f"{len(items)} items" if len(items) > 1 else "it"
|
||||
await bot.send_message(
|
||||
chat_id,
|
||||
f"Got {count}! Now send keywords — one per message or comma-separated.\n"
|
||||
"Send /done when finished.",
|
||||
parse_mode="HTML",
|
||||
)
|
||||
|
||||
|
||||
async def _receive_media(msg: Message, state: FSMContext, file_id: str, file_unique_id: str,
|
||||
media_type: str, status: str = "approved"):
|
||||
item = {"file_id": file_id, "file_unique_id": file_unique_id, "media_type": media_type}
|
||||
await start_keyword_prompt(msg.chat.id, state, [item], status)
|
||||
|
||||
|
||||
# Ordered: Telegram sets `document` alongside `animation` for backward
|
||||
# compatibility, so animations must be matched before documents or every GIF
|
||||
# would be filed as a file. Round videos are stored as plain "video" so they
|
||||
@@ -350,7 +362,7 @@ def extract_media(msg: Message) -> tuple[str, str, str] | None:
|
||||
return None
|
||||
|
||||
|
||||
@dp.message(MEDIA_FILTER, StateFilter(None))
|
||||
@dp.message(MEDIA_FILTER, StateFilter(None), F.forward_origin.is_(None))
|
||||
async def on_media(msg: Message, state: FSMContext):
|
||||
user_id = msg.from_user.id
|
||||
if is_admin(user_id):
|
||||
@@ -376,27 +388,32 @@ async def cmd_done(msg: Message, state: FSMContext):
|
||||
await msg.answer("You haven't added any keywords yet. Send some, then /done.")
|
||||
return
|
||||
status = data.get("status", "approved")
|
||||
items = data.get("items", [])
|
||||
for item in items:
|
||||
await db.add_media(
|
||||
msg.from_user.id,
|
||||
data["file_id"],
|
||||
data["file_unique_id"],
|
||||
data["media_type"],
|
||||
item["file_id"],
|
||||
item["file_unique_id"],
|
||||
item["media_type"],
|
||||
keywords,
|
||||
status=status,
|
||||
)
|
||||
await state.clear()
|
||||
saved = f"{len(items)} items" if len(items) > 1 else "it"
|
||||
if status == "pending":
|
||||
await msg.answer(
|
||||
f"Submitted with keywords: <b>{', '.join(keywords)}</b>\n"
|
||||
f"Submitted {saved} with keywords: <b>{', '.join(keywords)}</b>\n"
|
||||
"A moderator will review it before it becomes searchable.",
|
||||
parse_mode="HTML",
|
||||
)
|
||||
else:
|
||||
await msg.answer(f"Saved with keywords: <b>{', '.join(keywords)}</b>", parse_mode="HTML")
|
||||
await msg.answer(f"Saved {saved} with keywords: <b>{', '.join(keywords)}</b>",
|
||||
parse_mode="HTML")
|
||||
|
||||
|
||||
@dp.message(AddMedia.waiting_for_keywords, F.text)
|
||||
async def on_keyword(msg: Message, state: FSMContext):
|
||||
async def accumulate_keywords(msg: Message, state: FSMContext) -> None:
|
||||
"""Shared by adding, editing and vault saves: merge in whatever was typed
|
||||
and echo the running list back."""
|
||||
new_kws = [k.strip().lower() for k in msg.text.split(",") if k.strip()]
|
||||
if not new_kws:
|
||||
return
|
||||
@@ -407,6 +424,11 @@ async def on_keyword(msg: Message, state: FSMContext):
|
||||
await msg.reply(f"Keywords so far: <b>{', '.join(merged)}</b>\nSend more or /done.", parse_mode="HTML")
|
||||
|
||||
|
||||
@dp.message(AddMedia.waiting_for_keywords, F.text)
|
||||
async def on_keyword(msg: Message, state: FSMContext):
|
||||
await accumulate_keywords(msg, state)
|
||||
|
||||
|
||||
@dp.message(AddMedia.waiting_for_keywords, MEDIA_FILTER)
|
||||
async def on_media_during_keywords(msg: Message):
|
||||
await msg.reply("Still waiting for keywords. Send /done to save the current item, or /cancel to discard.")
|
||||
@@ -560,14 +582,254 @@ async def cmd_edit_cancel(msg: Message, state: FSMContext):
|
||||
|
||||
@dp.message(EditMedia.waiting_for_new_keywords, F.text)
|
||||
async def on_edit_keyword(msg: Message, state: FSMContext):
|
||||
new_kws = [k.strip().lower() for k in msg.text.split(",") if k.strip()]
|
||||
if not new_kws:
|
||||
await accumulate_keywords(msg, state)
|
||||
|
||||
|
||||
# ── forwards: save whole, or extract the media ───────────────────────────────
|
||||
#
|
||||
# A forwarded album arrives as several separate messages sharing a
|
||||
# media_group_id, with no signal that the last one has landed, so they are
|
||||
# buffered briefly and treated as one post.
|
||||
|
||||
ALBUM_DEBOUNCE = 1.0
|
||||
PENDING_FORWARD_LIMIT = 200
|
||||
|
||||
_album_buffers: dict[str, dict] = {}
|
||||
_pending_forwards: dict[str, dict] = {}
|
||||
|
||||
|
||||
def _new_forward_record(msg: Message) -> dict:
|
||||
return {
|
||||
"user_id": msg.from_user.id,
|
||||
"chat_id": msg.chat.id,
|
||||
"message_ids": [],
|
||||
"caption": "",
|
||||
"media": [],
|
||||
}
|
||||
|
||||
|
||||
def _absorb(record: dict, msg: Message) -> None:
|
||||
record["message_ids"].append(msg.message_id)
|
||||
# In an album only one message carries the caption.
|
||||
text = msg.caption or msg.text or ""
|
||||
if text and not record["caption"]:
|
||||
record["caption"] = text
|
||||
found = extract_media(msg)
|
||||
if found is not None:
|
||||
media_type, file_id, file_unique_id = found
|
||||
record["media"].append(
|
||||
{"media_type": media_type, "file_id": file_id, "file_unique_id": file_unique_id}
|
||||
)
|
||||
|
||||
|
||||
async def prompt_forward_choice(record: dict) -> None:
|
||||
token = uuid.uuid4().hex[:12]
|
||||
# Bounded so a long-running bot cannot accumulate abandoned prompts.
|
||||
while len(_pending_forwards) >= PENDING_FORWARD_LIMIT:
|
||||
_pending_forwards.pop(next(iter(_pending_forwards)))
|
||||
_pending_forwards[token] = record
|
||||
|
||||
record["message_ids"].sort()
|
||||
parts = [f"{len(record['message_ids'])} message(s)"]
|
||||
if record["media"]:
|
||||
parts.append(f"{len(record['media'])} media")
|
||||
summary = ", ".join(parts)
|
||||
|
||||
row = [InlineKeyboardButton(text="📥 Save whole post", callback_data=f"fwd:save:{token}")]
|
||||
if record["media"]:
|
||||
row.append(InlineKeyboardButton(text="✂️ Extract media", callback_data=f"fwd:extract:{token}"))
|
||||
|
||||
await bot.send_message(
|
||||
record["chat_id"],
|
||||
f"Forwarded post received ({summary}).\n\n"
|
||||
"<b>Save whole post</b> keeps it intact in your vault — retrievable only "
|
||||
"inside this bot, with /vault.\n"
|
||||
"<b>Extract media</b> adds the media to your searchable library as usual.",
|
||||
parse_mode="HTML",
|
||||
reply_markup=InlineKeyboardMarkup(inline_keyboard=[row]),
|
||||
)
|
||||
|
||||
|
||||
async def _flush_album(group_id: str) -> None:
|
||||
"""Wait until an album has stopped arriving, then prompt once for all of it."""
|
||||
while True:
|
||||
record = _album_buffers.get(group_id)
|
||||
if record is None:
|
||||
return
|
||||
remaining = record["deadline"] - time.monotonic()
|
||||
if remaining <= 0:
|
||||
break
|
||||
await asyncio.sleep(remaining)
|
||||
|
||||
record = _album_buffers.pop(group_id, None)
|
||||
if record is not None:
|
||||
await prompt_forward_choice(record)
|
||||
|
||||
|
||||
@dp.message(F.forward_origin, StateFilter(None))
|
||||
async def on_forwarded(msg: Message, state: FSMContext):
|
||||
if not await has_access(msg.from_user.id):
|
||||
await deny(msg)
|
||||
return
|
||||
|
||||
group_id = msg.media_group_id
|
||||
if group_id is None:
|
||||
record = _new_forward_record(msg)
|
||||
_absorb(record, msg)
|
||||
await prompt_forward_choice(record)
|
||||
return
|
||||
|
||||
record = _album_buffers.get(group_id)
|
||||
if record is None:
|
||||
record = _new_forward_record(msg)
|
||||
_album_buffers[group_id] = record
|
||||
record["deadline"] = time.monotonic() + ALBUM_DEBOUNCE
|
||||
asyncio.create_task(_flush_album(group_id))
|
||||
_absorb(record, msg)
|
||||
record["deadline"] = time.monotonic() + ALBUM_DEBOUNCE
|
||||
|
||||
|
||||
@dp.callback_query(F.data.startswith("fwd:"))
|
||||
async def on_forward_choice(callback: CallbackQuery, state: FSMContext):
|
||||
_, action, token = callback.data.split(":", 2)
|
||||
record = _pending_forwards.pop(token, None)
|
||||
|
||||
if record is None:
|
||||
await callback.answer("That prompt expired — forward the post again.", show_alert=True)
|
||||
return
|
||||
if record["user_id"] != callback.from_user.id:
|
||||
await callback.answer("Not yours.", show_alert=True)
|
||||
return
|
||||
|
||||
await callback.message.edit_reply_markup(reply_markup=None)
|
||||
|
||||
if action == "extract":
|
||||
if not await has_access(callback.from_user.id):
|
||||
await callback.answer("Not allowed.", show_alert=True)
|
||||
return
|
||||
status = "approved" if is_admin(callback.from_user.id) else "pending"
|
||||
await start_keyword_prompt(record["chat_id"], state, record["media"], status)
|
||||
await callback.answer()
|
||||
return
|
||||
|
||||
await state.set_state(VaultSave.waiting_for_keywords)
|
||||
await state.update_data(
|
||||
chat_id=record["chat_id"],
|
||||
message_ids=record["message_ids"],
|
||||
caption=record["caption"],
|
||||
keywords=[],
|
||||
)
|
||||
hint = "Send keywords, then /done." if record["caption"] else \
|
||||
"Send keywords, then /done. This post has no caption, so keywords are how you will find it."
|
||||
await callback.message.answer(hint)
|
||||
await callback.answer()
|
||||
|
||||
|
||||
@dp.message(VaultSave.waiting_for_keywords, Command("done"))
|
||||
async def cmd_vault_done(msg: Message, state: FSMContext):
|
||||
data = await state.get_data()
|
||||
existing: list[str] = data.get("keywords", [])
|
||||
merged = existing + [k for k in new_kws if k not in existing]
|
||||
await state.update_data(keywords=merged)
|
||||
await msg.reply(f"Keywords so far: <b>{', '.join(merged)}</b>\nSend more or /done.", parse_mode="HTML")
|
||||
keywords: list[str] = data.get("keywords", [])
|
||||
caption: str = data.get("caption", "")
|
||||
|
||||
if not keywords and not caption:
|
||||
await msg.answer("Add at least one keyword — there is no caption to find this by.")
|
||||
return
|
||||
|
||||
item_id = await db.add_vault_item(
|
||||
msg.from_user.id, data["chat_id"], data["message_ids"], caption, keywords
|
||||
)
|
||||
await state.clear()
|
||||
await msg.answer(
|
||||
f"Saved to your vault as <b>#{item_id}</b>.\n\n"
|
||||
"Vault posts come back whole, exactly as forwarded — but only here in "
|
||||
"this chat, with /vault. Inline search does not reach them.",
|
||||
parse_mode="HTML",
|
||||
)
|
||||
|
||||
|
||||
@dp.message(VaultSave.waiting_for_keywords, F.text)
|
||||
async def on_vault_keyword(msg: Message, state: FSMContext):
|
||||
await accumulate_keywords(msg, state)
|
||||
|
||||
|
||||
# ── /vault ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def vault_label(item: dict) -> str:
|
||||
text = item["keywords"].replace(",", ", ") or item["caption"]
|
||||
text = " ".join(text.split())
|
||||
if len(text) > 40:
|
||||
text = text[:39] + "…"
|
||||
count = len(db.unpack_message_ids(item["message_ids"]))
|
||||
suffix = f" ({count})" if count > 1 else ""
|
||||
return f"#{item['id']} {text or 'untitled'}{suffix}"
|
||||
|
||||
|
||||
@dp.message(Command("vault"))
|
||||
async def cmd_vault(msg: Message):
|
||||
if not await has_access(msg.from_user.id):
|
||||
await deny(msg)
|
||||
return
|
||||
|
||||
parts = msg.text.split(maxsplit=1)
|
||||
query = parts[1].strip() if len(parts) > 1 else ""
|
||||
items = await db.search_vault(msg.from_user.id, query)
|
||||
|
||||
if not items:
|
||||
if query:
|
||||
await msg.answer(f"Nothing in your vault matching '{query}'.")
|
||||
else:
|
||||
await msg.answer(
|
||||
"Your vault is empty. Forward a post to me and choose "
|
||||
"<b>Save whole post</b> to put it here.",
|
||||
parse_mode="HTML",
|
||||
)
|
||||
return
|
||||
|
||||
rows = [
|
||||
[
|
||||
InlineKeyboardButton(text=vault_label(item), callback_data=f"vault:get:{item['id']}"),
|
||||
InlineKeyboardButton(text="🗑", callback_data=f"vault:del:{item['id']}"),
|
||||
]
|
||||
for item in items
|
||||
]
|
||||
total = await db.count_vault(msg.from_user.id)
|
||||
header = f"Showing {len(items)} of {total} saved posts." if not query else \
|
||||
f"{len(items)} match '{query}'."
|
||||
await msg.answer(header, reply_markup=InlineKeyboardMarkup(inline_keyboard=rows))
|
||||
|
||||
|
||||
@dp.callback_query(F.data.startswith("vault:"))
|
||||
async def on_vault_action(callback: CallbackQuery):
|
||||
_, action, raw_id = callback.data.split(":", 2)
|
||||
# Scoped by owner, so a guessed id cannot reach someone else's vault.
|
||||
item = await db.get_vault_item(callback.from_user.id, int(raw_id))
|
||||
if item is None:
|
||||
await callback.answer("Not found.", show_alert=True)
|
||||
return
|
||||
|
||||
if action == "del":
|
||||
await db.delete_vault_item(callback.from_user.id, item["id"])
|
||||
await callback.answer("Removed from vault.")
|
||||
await callback.message.answer(f"Removed #{item['id']} from your vault.")
|
||||
return
|
||||
|
||||
message_ids = db.unpack_message_ids(item["message_ids"])
|
||||
try:
|
||||
# forwardMessages keeps an album together and preserves the original
|
||||
# "forwarded from" header. Telegram requires ascending ids.
|
||||
await bot.forward_messages(
|
||||
chat_id=callback.message.chat.id,
|
||||
from_chat_id=item["chat_id"],
|
||||
message_ids=sorted(message_ids),
|
||||
)
|
||||
await callback.answer()
|
||||
except TelegramAPIError:
|
||||
log.exception("could not forward vault item %s", item["id"])
|
||||
await callback.answer("Could not retrieve that post.", show_alert=True)
|
||||
await callback.message.answer(
|
||||
f"#{item['id']} could not be retrieved — the original messages may have "
|
||||
"been deleted from this chat. Forward the post again to re-save it."
|
||||
)
|
||||
|
||||
|
||||
# ── /review ──────────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -4,7 +4,7 @@ import os
|
||||
DB_PATH = os.getenv("DB_PATH", "/data/media.db")
|
||||
|
||||
# Bump this and add a matching _migrate_to_vN() when the schema changes.
|
||||
SCHEMA_VERSION = 3
|
||||
SCHEMA_VERSION = 4
|
||||
|
||||
|
||||
# ── schema / migrations ──────────────────────────────────────────────────────
|
||||
@@ -69,9 +69,28 @@ async def _migrate_to_v3(db) -> None:
|
||||
await db.execute("CREATE INDEX IF NOT EXISTS idx_status ON media(status)")
|
||||
|
||||
|
||||
async def _migrate_to_v4(db) -> None:
|
||||
"""The vault: posts kept whole, as pointers to the messages the bot already
|
||||
holds in its private chat with the owner. Nothing here is ever shared, so it
|
||||
lives in its own table rather than as a flag on media."""
|
||||
await db.execute("""
|
||||
CREATE TABLE IF NOT EXISTS vault_items (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
owner_id INTEGER NOT NULL,
|
||||
chat_id INTEGER NOT NULL,
|
||||
message_ids TEXT NOT NULL,
|
||||
caption TEXT NOT NULL DEFAULT '',
|
||||
keywords TEXT NOT NULL DEFAULT '',
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""")
|
||||
await db.execute("CREATE INDEX IF NOT EXISTS idx_vault_owner ON vault_items(owner_id)")
|
||||
|
||||
|
||||
MIGRATIONS = {
|
||||
2: _migrate_to_v2,
|
||||
3: _migrate_to_v3,
|
||||
4: _migrate_to_v4,
|
||||
}
|
||||
|
||||
|
||||
@@ -246,3 +265,89 @@ async def set_status(media_id: int, status: str, reviewer_id: int) -> bool:
|
||||
)
|
||||
await db.commit()
|
||||
return cur.rowcount > 0
|
||||
|
||||
|
||||
# ── vault ────────────────────────────────────────────────────────────────────
|
||||
#
|
||||
# Every query here takes an owner_id and filters on it. A vault item belongs to
|
||||
# one person and is never searchable, never moderated and never federated; the
|
||||
# separate table is what makes that structural rather than a flag to remember.
|
||||
|
||||
VAULT_COLUMNS = "id, owner_id, chat_id, message_ids, caption, keywords, created_at"
|
||||
|
||||
|
||||
def pack_message_ids(message_ids: list[int]) -> str:
|
||||
return ",".join(str(m) for m in message_ids)
|
||||
|
||||
|
||||
def unpack_message_ids(packed: str) -> list[int]:
|
||||
return [int(m) for m in packed.split(",") if m.strip()]
|
||||
|
||||
|
||||
async def add_vault_item(owner_id: int, chat_id: int, message_ids: list[int],
|
||||
caption: str, keywords: list[str]) -> int:
|
||||
kw = ",".join(k.strip().lower() for k in keywords if k.strip())
|
||||
async with aiosqlite.connect(DB_PATH) as db:
|
||||
cur = await db.execute(
|
||||
"""
|
||||
INSERT INTO vault_items (owner_id, chat_id, message_ids, caption, keywords)
|
||||
VALUES (?,?,?,?,?)
|
||||
""",
|
||||
(owner_id, chat_id, pack_message_ids(message_ids), caption or "", kw),
|
||||
)
|
||||
await db.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
|
||||
async def search_vault(owner_id: int, query: str, limit: int = 20) -> list[dict]:
|
||||
"""Match keywords by prefix or the original caption anywhere in its text —
|
||||
the point of a vault is finding a post you never got round to tagging."""
|
||||
query = query.strip().lower()
|
||||
params: list = [owner_id]
|
||||
clause = ""
|
||||
if query:
|
||||
clause = "AND ((',' || keywords || ',') LIKE ? OR LOWER(caption) LIKE ?)"
|
||||
params += [f"%,{query}%", f"%{query}%"]
|
||||
params.append(limit)
|
||||
|
||||
async with aiosqlite.connect(DB_PATH) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
cursor = await db.execute(
|
||||
f"""
|
||||
SELECT {VAULT_COLUMNS} FROM vault_items
|
||||
WHERE owner_id = ? {clause}
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
params,
|
||||
)
|
||||
return [dict(r) for r in await cursor.fetchall()]
|
||||
|
||||
|
||||
async def get_vault_item(owner_id: int, item_id: int) -> dict | None:
|
||||
"""Scoped by owner so one user cannot fetch another user's saved post."""
|
||||
async with aiosqlite.connect(DB_PATH) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
cursor = await db.execute(
|
||||
f"SELECT {VAULT_COLUMNS} FROM vault_items WHERE id = ? AND owner_id = ?",
|
||||
(item_id, owner_id),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
|
||||
async def delete_vault_item(owner_id: int, item_id: int) -> bool:
|
||||
async with aiosqlite.connect(DB_PATH) as db:
|
||||
cur = await db.execute(
|
||||
"DELETE FROM vault_items WHERE id = ? AND owner_id = ?", (item_id, owner_id)
|
||||
)
|
||||
await db.commit()
|
||||
return cur.rowcount > 0
|
||||
|
||||
|
||||
async def count_vault(owner_id: int) -> int:
|
||||
async with aiosqlite.connect(DB_PATH) as db:
|
||||
cursor = await db.execute(
|
||||
"SELECT COUNT(*) FROM vault_items WHERE owner_id = ?", (owner_id,)
|
||||
)
|
||||
return (await cursor.fetchone())[0]
|
||||
|
||||
+5
-2
@@ -47,12 +47,15 @@ async def test_legacy_upgrade(tmp: str):
|
||||
db = load_db(path)
|
||||
await db.init_db()
|
||||
|
||||
assert user_version(path) == 3
|
||||
# Pinned to the module so adding a migration does not require editing this.
|
||||
assert user_version(path) == db.SCHEMA_VERSION
|
||||
rows = sorted((r["id"], r["keywords"]) for r in await db.search_media(""))
|
||||
assert rows == [(7, "cat,funny"), (9, "dog")], f"rows not preserved: {rows}"
|
||||
with sqlite3.connect(path) as c:
|
||||
statuses = {r[0] for r in c.execute("SELECT status FROM media")}
|
||||
tables = {r[0] for r in c.execute("SELECT name FROM sqlite_master WHERE type='table'")}
|
||||
assert statuses == {"approved"}, f"upgraded rows must stay live: {statuses}"
|
||||
assert "vault_items" in tables, "upgrading should create the vault table too"
|
||||
print(" legacy v1 database upgrades with ids, keywords and visibility intact")
|
||||
|
||||
|
||||
@@ -91,7 +94,7 @@ async def test_fresh_and_idempotent(tmp: str):
|
||||
with sqlite3.connect(path) as c:
|
||||
sql = c.execute("SELECT sql FROM sqlite_master WHERE name='media'").fetchone()[0]
|
||||
assert "UNIQUE(owner_id, file_unique_id)" in sql
|
||||
assert user_version(path) == 3
|
||||
assert user_version(path) == db.SCHEMA_VERSION
|
||||
|
||||
await db.init_db()
|
||||
await db.init_db()
|
||||
|
||||
+176
@@ -0,0 +1,176 @@
|
||||
"""Vault tests. Run with: python test_vault.py
|
||||
|
||||
The vault stores pointers to messages the bot already holds, so the things
|
||||
worth pinning down are that albums are grouped into one post, that retrieval
|
||||
asks Telegram for the right message ids, and that one user can never reach
|
||||
another user's saved posts.
|
||||
"""
|
||||
import asyncio
|
||||
import datetime
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
os.environ.setdefault("BOT_TOKEN", "123456789:AAHtestTokenPlaceholder_ABCDEFGHIJKLMNOP")
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from aiogram.types import Chat, Message, PhotoSize, User
|
||||
|
||||
logging.disable(logging.CRITICAL)
|
||||
|
||||
CHAT = Chat(id=500, type="private")
|
||||
USER = User(id=111, is_bot=False, first_name="Owner")
|
||||
NOW = datetime.datetime.now()
|
||||
|
||||
|
||||
def load(path: str, **env):
|
||||
for key in ("PUBLIC_MODE", "REQUIRED_CHATS", "ALLOWED_USERS", "ADMIN_USERS"):
|
||||
os.environ.pop(key, None)
|
||||
os.environ["DB_PATH"] = path
|
||||
os.environ.update(env)
|
||||
import db, bot
|
||||
importlib.reload(db)
|
||||
B = importlib.reload(bot)
|
||||
B.db = db
|
||||
return db, B
|
||||
|
||||
|
||||
def photo_msg(message_id: int, caption: str | None = None, group: str | None = None) -> Message:
|
||||
photo = PhotoSize(file_id=f"F{message_id}", file_unique_id=f"U{message_id}", width=1, height=1)
|
||||
return Message(message_id=message_id, date=NOW, chat=CHAT, from_user=USER,
|
||||
photo=[photo], caption=caption, media_group_id=group)
|
||||
|
||||
|
||||
# ── database ─────────────────────────────────────────────────────────────────
|
||||
|
||||
async def test_roundtrip(db):
|
||||
item_id = await db.add_vault_item(111, 500, [12, 10, 11], "A trip to the sea", ["sea"])
|
||||
item = await db.get_vault_item(111, item_id)
|
||||
assert db.unpack_message_ids(item["message_ids"]) == [12, 10, 11]
|
||||
assert item["caption"] == "A trip to the sea"
|
||||
print(" a saved post round-trips with its message ids and caption")
|
||||
|
||||
|
||||
async def test_search_by_keyword_and_caption(db):
|
||||
await db.add_vault_item(111, 500, [20], "Recipe for lasagne", ["food"])
|
||||
assert len(await db.search_vault(111, "food")) == 1, "keyword search"
|
||||
assert len(await db.search_vault(111, "lasagne")) == 1, "caption search"
|
||||
assert len(await db.search_vault(111, "lasagne")) == 1
|
||||
assert await db.search_vault(111, "nonsense") == []
|
||||
print(" posts are findable by keyword and by words in the original caption")
|
||||
|
||||
|
||||
async def test_vault_is_private_to_its_owner(db):
|
||||
mine = await db.add_vault_item(111, 500, [30], "private note", ["secret"])
|
||||
|
||||
assert await db.get_vault_item(222, mine) is None, "another user fetched my post"
|
||||
assert await db.search_vault(222, "secret") == [], "another user searched my vault"
|
||||
assert await db.delete_vault_item(222, mine) is False, "another user deleted my post"
|
||||
assert await db.get_vault_item(111, mine) is not None, "owner lost their own post"
|
||||
print(" one user cannot read, search or delete another user's vault")
|
||||
|
||||
|
||||
async def test_delete(db):
|
||||
item_id = await db.add_vault_item(111, 500, [40], "", ["bye"])
|
||||
assert await db.delete_vault_item(111, item_id) is True
|
||||
assert await db.get_vault_item(111, item_id) is None
|
||||
assert await db.delete_vault_item(111, item_id) is False
|
||||
print(" deleting removes the post and is safe to repeat")
|
||||
|
||||
|
||||
# ── album grouping ───────────────────────────────────────────────────────────
|
||||
|
||||
async def test_album_becomes_one_post(B):
|
||||
B.ALBUM_DEBOUNCE = 0.05
|
||||
B._album_buffers.clear()
|
||||
B._pending_forwards.clear()
|
||||
sent: list = []
|
||||
|
||||
async def send_message(chat_id, text, **kw):
|
||||
sent.append(text)
|
||||
B.bot.send_message = send_message
|
||||
|
||||
# Telegram delivers an album as separate messages; only one has the caption.
|
||||
for msg in (photo_msg(10, caption="beach day", group="g1"),
|
||||
photo_msg(11, group="g1"),
|
||||
photo_msg(12, group="g1")):
|
||||
await B.on_forwarded(msg, None)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
assert len(sent) == 1, f"an album should prompt once, got {len(sent)} prompts"
|
||||
assert len(B._pending_forwards) == 1
|
||||
record = next(iter(B._pending_forwards.values()))
|
||||
assert record["message_ids"] == [10, 11, 12]
|
||||
assert record["caption"] == "beach day", "the one caption in the album was lost"
|
||||
assert len(record["media"]) == 3, "every item in the album should be extractable"
|
||||
print(" an album is buffered into one post, keeping its caption and all media")
|
||||
|
||||
|
||||
async def test_single_forward_prompts_immediately(B):
|
||||
B._album_buffers.clear()
|
||||
B._pending_forwards.clear()
|
||||
sent: list = []
|
||||
|
||||
async def send_message(chat_id, text, **kw):
|
||||
sent.append(text)
|
||||
B.bot.send_message = send_message
|
||||
|
||||
await B.on_forwarded(photo_msg(20, caption="just one"), None)
|
||||
assert len(sent) == 1, "a lone forward should not wait for an album"
|
||||
record = next(iter(B._pending_forwards.values()))
|
||||
assert record["message_ids"] == [20]
|
||||
print(" a single forwarded post prompts straight away")
|
||||
|
||||
|
||||
async def test_pending_prompts_are_bounded(B):
|
||||
B._pending_forwards.clear()
|
||||
B.bot.send_message = lambda *a, **kw: asyncio.sleep(0)
|
||||
|
||||
for i in range(B.PENDING_FORWARD_LIMIT + 25):
|
||||
await B.prompt_forward_choice(
|
||||
{"user_id": 1, "chat_id": 500, "message_ids": [i], "caption": "", "media": []}
|
||||
)
|
||||
assert len(B._pending_forwards) <= B.PENDING_FORWARD_LIMIT
|
||||
print(" abandoned prompts cannot grow without bound")
|
||||
|
||||
|
||||
async def test_retrieval_forwards_ascending_ids(db, B):
|
||||
item_id = await db.add_vault_item(111, 500, [12, 10, 11], "out of order", ["x"])
|
||||
forwarded: dict = {}
|
||||
|
||||
async def forward_messages(chat_id, from_chat_id, message_ids):
|
||||
forwarded.update(chat_id=chat_id, from_chat_id=from_chat_id, message_ids=message_ids)
|
||||
B.bot.forward_messages = forward_messages
|
||||
|
||||
class Cb:
|
||||
data = f"vault:get:{item_id}"
|
||||
from_user = type("U", (), {"id": 111})()
|
||||
message = type("M", (), {"chat": CHAT})()
|
||||
async def answer(self, *a, **kw): pass
|
||||
|
||||
await B.on_vault_action(Cb())
|
||||
# Telegram requires strictly increasing ids, and the album must stay whole.
|
||||
assert forwarded["message_ids"] == [10, 11, 12], forwarded
|
||||
assert forwarded["from_chat_id"] == 500
|
||||
print(" retrieval forwards the whole post with ids in ascending order")
|
||||
|
||||
|
||||
async def main():
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db, B = load(os.path.join(tmp, "vault.db"), ADMIN_USERS="111")
|
||||
await db.init_db()
|
||||
|
||||
for test in (test_roundtrip, test_search_by_keyword_and_caption,
|
||||
test_vault_is_private_to_its_owner, test_delete):
|
||||
await test(db)
|
||||
for test in (test_album_becomes_one_post, test_single_forward_prompts_immediately,
|
||||
test_pending_prompts_are_bounded):
|
||||
await test(B)
|
||||
await test_retrieval_forwards_ascending_ids(db, B)
|
||||
print("\nAll vault tests passed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user