kulnoor/drive
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first commit

Browse Source
This commit is contained in:
gosangam
2026-10-18 23:48:22 +05:30
commit 732a06f65d
38 changed files with 8943 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
+119
View File
@@ -0,0 +1,119 @@
# Copy to .env for local runs. On Heroku set these as Config Vars instead.
# Placeholder values only - never commit real secrets.
# ---------------------------------------------------------------- Telegram
# user = your own account (full history, needs login); bot = bot admin of the channel
TELEGRAM_MODE=user
# From https://my.telegram.org -> API development tools (needed in both modes)
TELEGRAM_API_ID=1234567
TELEGRAM_API_HASH=your_api_hash_here
# From @BotFather: admin command interface (required in bot mode, recommended in user mode)
TELEGRAM_BOT_TOKEN=123456789:replace-with-your-bot-token
# User mode on Heroku: output of `python main.py login --session-string`.
# Locally you can leave it empty and use `python main.py login` (session file).
TELEGRAM_SESSION_STRING=
# Numeric Telegram user IDs allowed to use the commands (comma-separated)
ADMIN_USER_IDS=111111111
# Chats that receive progress/notifications (default: ADMIN_USER_IDS)
NOTIFY_CHAT_IDS=
# @username, https://t.me/... link, or -100... numeric ID
SOURCE_CHANNEL_ID=@your_channel
# existing | new | both
PROCESS_MODE=both
START_MESSAGE_ID=1
# Optional ISO date: only messages on/after this date (e.g. 2024-01-31)
START_DATE=
INCLUDE_DOCUMENT_VIDEOS=true
# Optional comma-separated filters (case-insensitive)
FILTER_TOPICS=
FILTER_BATCHES=
# Bot mode history probing (bots cannot list history)
BOT_SCAN_MAX_EMPTY=200
BOT_SCAN_END_ID=
# Telegram flood waits shorter than this (seconds) are slept through automatically
FLOOD_SLEEP_THRESHOLD=60
SESSION_DIR=data/sessions
# ---------------------------------------------------------------- Google Drive
# oauth (personal Drive) | service_account (Shared Drive)
GOOGLE_AUTH_MODE=oauth
# drive.file (least privilege; bot only sees folders it created) | drive (full access)
GOOGLE_DRIVE_SCOPE=drive.file
# Local OAuth files (used by `python main.py drive-auth`)
GOOGLE_CLIENT_SECRETS_FILE=credentials/client_secret.json
GOOGLE_TOKEN_FILE=credentials/token.json
# Heroku / file-less alternatives (raw JSON or base64 JSON):
GOOGLE_OAUTH_TOKEN_JSON=
GOOGLE_SERVICE_ACCOUNT_JSON=
GOOGLE_SERVICE_ACCOUNT_FILE=
# Folder that will contain the topic folders (empty = top level of My Drive)
GOOGLE_DRIVE_ROOT_FOLDER_ID=
DRIVE_SHARED_DRIVE_ID=
# none (private, default) | anyone_with_link
DRIVE_SHARE_MODE=none
# Multiple of 256KB; one chunk is held in memory at a time
UPLOAD_CHUNK_SIZE=16MB
# ---------------------------------------------------------------- Organisation & naming
BATCH_SUBFOLDERS_ENABLED=false
FALLBACK_TOPIC_FOLDER=Uncategorized
FALLBACK_BATCH_FOLDER=Unknown Batch
# suffix_number | suffix_message_id | skip | fail (never overwrites)
DUPLICATE_POLICY=suffix_number
# Reuse an existing same-name file in the folder when size and MD5 match
REUSE_IDENTICAL_FILES=true
# skip | upload : the same Telegram media posted again under a new message ID
DUPLICATE_MEDIA_POLICY=skip
# fallback (Telegram file name / telegram_<id>) | fail
MISSING_TITLE_POLICY=fallback
# keep | append : when the Title extension disagrees with the real container
EXTENSION_MISMATCH_POLICY=keep
# Replace characters Windows cannot store (<>:"|?*)
STRICT_FILENAMES=false
# Optional persistent sequence number prefix, e.g. "001 - Class 01.mp4"
SEQUENCE_PREFIX_ENABLED=false
SEQUENCE_PREFIX_FORMAT="{seq:03d} - "
# ---------------------------------------------------------------- Storage
# PostgreSQL (required on Heroku; Heroku Postgres sets it automatically)
DATABASE_URL=
# SQLite file used when DATABASE_URL is empty (local development)
DATABASE_PATH=data/bot.db
# Only set to true if you accept losing all progress on every Heroku restart
ALLOW_EPHEMERAL_DATABASE=false
DATA_DIR=data
DOWNLOAD_DIR=downloads
LOG_DIR=logs
# 0 = unlimited
MAX_TEMP_STORAGE=0
MIN_FREE_DISK_SPACE=1GB
KEEP_LOCAL_FILES=false
# Keep the local file of a video that FAILED (default: delete it; it can be downloaded again)
KEEP_FAILED_DOWNLOADS=false
# Seconds between sweeps that delete abandoned temp files (0 = only at startup / when disk is low)
TEMP_CLEANUP_INTERVAL=3600
COMPUTE_SHA256=true
# ---------------------------------------------------------------- Processing
MAX_RETRIES=4
RETRY_DELAY=30
RETRY_MAX_DELAY=1800
# Retries of a single HTTP request inside one attempt
HTTP_RETRIES=5
PROCESSING_DELAY=2
# true = halt the whole queue when any video fails permanently
STRICT_STOP=false
# true = start processing at every boot (otherwise /run; then remembered across restarts)
AUTO_START=false
# Seconds to wait for another worker to release the queue lock (default 0 locally, 90 on Heroku)
INSTANCE_LOCK_WAIT=
PROGRESS_INTERVAL=10
# ---------------------------------------------------------------- /upload tasks
# Minutes before an unfinished /upload setup (waiting for links or /confirm) expires
UPLOAD_SETUP_TIMEOUT=30
# Largest allowed range (end - start + 1) for one task
UPLOAD_MAX_RANGE=5000
# true = links must point to SOURCE_CHANNEL_ID; false = any channel the Telegram account/bot can access
UPLOAD_RESTRICT_TO_SOURCE_CHANNEL=true
LOG_LEVEL=info
.gitignore
+27
View File
@@ -0,0 +1,27 @@
# Secrets and credentials - never commit
.env
.env.*
!.env.example
credentials/
client_secret*.json
token.json
*service-account*.json
*.session
*.session-journal
session_string.txt
# Local state (SQLite, sessions, downloads, logs)
data/
downloads/
logs/
*.db
*.db-wal
*.db-shm
*.lock
# Python
.venv/
venv/
__pycache__/
*.pyc
.pytest_cache/
.python-version
+1
View File
@@ -0,0 +1 @@
3.12
ENV_SETUP.md
+388
View File
@@ -0,0 +1,388 @@
# How to Obtain Every Environment Variable
This guide explains where each configuration value comes from and how to get it. Work through the
sections in order: some steps (the Telegram session string and the Google token) need values from
earlier steps.
- **Local runs:** put the values in a `.env` file (`cp .env.example .env`).
- **Heroku:** set them as Config Vars, either in the dashboard (*App → Settings → Reveal Config Vars*)
or with `heroku config:set NAME=value --app YOUR_APP_NAME`.
> All values below are placeholders. Never commit real values or paste them into chats, issues or
> screenshots. `.env`, `credentials/` and `data/` are already in `.gitignore`.
---
## Quick checklist
| # | Variable | Required? | Where it comes from |
|---|---|---|---|
| 1 | `TELEGRAM_API_ID` | Always | my.telegram.org |
| 2 | `TELEGRAM_API_HASH` | Always | my.telegram.org |
| 3 | `TELEGRAM_BOT_TOKEN` | Always on Heroku (control commands); required in bot mode | @BotFather |
| 4 | `ADMIN_USER_IDS` | When a bot token is set | @userinfobot |
| 5 | `SOURCE_CHANNEL_ID` | Always | Channel username, link, or numeric ID |
| 6 | `TELEGRAM_MODE` | Optional (default `user`) | Your choice: `user` or `bot` |
| 7 | `TELEGRAM_SESSION_STRING` | `TELEGRAM_MODE=user` on Heroku | `python main.py login --session-string` |
| 8 | `GOOGLE_AUTH_MODE` | Optional (default `oauth`) | Your choice: `oauth` or `service_account` |
| 9 | `GOOGLE_OAUTH_TOKEN_JSON` | `oauth` on Heroku | `python main.py drive-auth` |
| 10 | `GOOGLE_SERVICE_ACCOUNT_JSON` | `service_account` on Heroku | Google Cloud Console key |
| 11 | `GOOGLE_DRIVE_ROOT_FOLDER_ID` | Recommended | Drive folder URL, or `python main.py drive-init` |
| 12 | `GOOGLE_DRIVE_SCOPE` | Optional (default `drive.file`) | Your choice |
| 13 | `DRIVE_SHARED_DRIVE_ID` | Shared Drive setups only | Shared Drive URL |
| 14 | `DATABASE_URL` | Required on Heroku | Heroku Postgres (set automatically) |
Everything else has a sensible default (see [Optional settings](#optional-settings)).
---
## Part 1 — Telegram
### 1–2. `TELEGRAM_API_ID` and `TELEGRAM_API_HASH`
These identify your Telegram API application. They are needed in **both** modes, because the bot talks
to Telegram through Telethon (MTProto).
1. Open <https://my.telegram.org> and log in with your phone number. Telegram sends a code to your
Telegram app, not by SMS.
2. Click **API development tools**.
3. If no app exists yet, fill in the form:
- *App title*: anything, for example `Drive Organizer`;
- *Short name*: 5–32 letters or digits, for example `driveorganizer`;
- *Platform*: `Desktop`;
- URL and description can stay empty.
4. Click **Create application**.
5. Copy the two values:
- **App api_id** (a number) → `TELEGRAM_API_ID`
- **App api_hash** (32 hex characters) → `TELEGRAM_API_HASH`
```env
TELEGRAM_API_ID=1234567
TELEGRAM_API_HASH=0123456789abcdef0123456789abcdef
```
> The API hash is a secret. If my.telegram.org shows an "ERROR" when you create the app, try another
> browser, disable ad blockers or VPNs, and use a short alphanumeric short name.
### 3. `TELEGRAM_BOT_TOKEN`
The bot provides the admin commands (`/run`, `/status`, …) and progress messages. In bot mode it also
reads the channel.
1. In Telegram, open [@BotFather](https://t.me/BotFather) and send `/newbot`.
2. Enter a display name, then a username ending in `bot` (for example `my_drive_organizer_bot`).
3. BotFather replies with a token like `123456789:AAH...`. Copy it → `TELEGRAM_BOT_TOKEN`.
4. **Open your new bot and press Start.** Bots cannot message you until you do, so without this you
get no progress messages.
```env
TELEGRAM_BOT_TOKEN=123456789:replace-with-your-bot-token
```
If the token leaks, send `/revoke` to @BotFather to get a new one.
### 4. `ADMIN_USER_IDS` (and optional `NOTIFY_CHAT_IDS`)
Only these numeric Telegram user IDs may control the bot.
1. Open [@userinfobot](https://t.me/userinfobot) and press Start.
2. It replies with your **Id** (a number such as `111111111`). Copy it.
3. For several admins, ask each person to do the same and separate the IDs with commas.
```env
ADMIN_USER_IDS=111111111,222222222
```
`NOTIFY_CHAT_IDS` is optional. It lists the chats that receive progress and error messages, and
defaults to `ADMIN_USER_IDS`. Each person listed must have pressed Start in the bot.
### 5. `SOURCE_CHANNEL_ID`
Any of these formats is accepted:
| Channel type | How to get it | Example value |
|---|---|---|
| Public channel | The `@username` shown in the channel info | `@my_channel` or `https://t.me/my_channel` |
| Private channel: post link | Right-click (or long-press) any post → **Copy Post Link** | `https://t.me/c/1234567890/55` |
| Private channel: numeric ID | Prefix the number from the post link with `-100` | `-1001234567890` |
| Private channel: Telegram Web | Open the channel at <https://web.telegram.org/k/>; the URL ends with `#-100…` | `-1001234567890` |
Another way: forward a channel post to [@JsonDumpBot](https://t.me/JsonDumpBot) and read
`forward_from_chat.id`.
```env
SOURCE_CHANNEL_ID=-1001234567890
```
> Bot mode with a private channel: the bot must be a channel **administrator**. If the numeric ID
> can't be resolved at the first start, post one message in the channel. The bot then learns the
> channel and remembers it in the database.
### 6. `TELEGRAM_MODE`
This is a choice, not a credential.
| Value | Use when | Needs |
|---|---|---|
| `user` (default) | You need the channel's **existing history**, or you are a member but not an admin | `TELEGRAM_SESSION_STRING` (Heroku) or a local login |
| `bot` | You only need **new posts**, and you can make the bot a channel admin | Bot added as channel administrator |
To add the bot as admin: open the channel → **Manage channel → Administrators → Add Admin** → search
for your bot's username → save. No extra rights are needed beyond the default.
### 7. `TELEGRAM_SESSION_STRING` (user mode only)
This lets the bot act as your Telegram account without a session file. Heroku's disk is wiped on
every restart, so a file wouldn't survive. You create it **once, on your own computer**.
1. Do the [local setup](README.md#19-local-development): create the virtualenv and install
requirements.
2. Create `.env` containing at least:
```env
TELEGRAM_API_ID=1234567
TELEGRAM_API_HASH=0123456789abcdef0123456789abcdef
```
3. Run:
```bash
python main.py login --session-string
```
4. Enter your phone number in international format (`+15551234567`), then the login code Telegram
sends to your app, then your 2FA password if you have one. **None of these are stored.**
5. The string is written to `data/sessions/session_string.txt` with owner-only permissions.
6. Put it on Heroku without printing it:
```bash
heroku config:set TELEGRAM_SESSION_STRING="$(cat data/sessions/session_string.txt)" --app YOUR_APP_NAME
```
For a local `.env`, paste the file's contents after `TELEGRAM_SESSION_STRING=`.
7. Delete the file: `rm data/sessions/session_string.txt`.
> The session string grants **full access to your Telegram account**. Treat it like a password. To
> revoke it: Telegram → *Settings → Devices* → terminate the session. Then generate a new string.
For local-only use you can skip the string and run `python main.py login`, which creates
`data/sessions/user.session` instead.
---
## Part 2 — Google Drive
Pick **one** authentication method:
| Your Drive | Use | Variable to obtain |
|---|---|---|
| Personal Google account (`@gmail.com`) or your own Workspace Drive | **OAuth** (`GOOGLE_AUTH_MODE=oauth`) | `GOOGLE_OAUTH_TOKEN_JSON` |
| Google Workspace **Shared Drive** | **Service account** (`GOOGLE_AUTH_MODE=service_account`) | `GOOGLE_SERVICE_ACCOUNT_JSON` |
> A service account **cannot** upload into a personal My Drive: it has no storage quota there. Use
> OAuth for personal Drives.
### Common first step: Google Cloud project and Drive API
1. Open <https://console.cloud.google.com/> and sign in with the Google account that owns the
destination Drive.
2. Use the project selector at the top → **New Project** → give it a name → **Create**, then select
it.
3. Go to **APIs & Services → Library**, search for **Google Drive API**, and click **Enable**.
### 9. `GOOGLE_OAUTH_TOKEN_JSON` (OAuth)
#### a) Configure the consent screen
1. Go to **APIs & Services → OAuth consent screen**. In the newer console this is **Google Auth
Platform → Branding / Audience**.
2. User type: **External** (or *Internal* if you use Workspace and only your organisation needs
access).
3. Fill in the app name, your support email and developer contact email, then save.
4. Under **Audience** (or *Test users*), add your own Google account.
5. **Publish the app** (*Publishing status → In production*).
> If the app stays in *Testing*, Google expires its refresh tokens after **7 days** and the bot
> will lose Drive access every week. With the default `drive.file` scope no Google verification
> is needed. With the full `drive` scope you'll see an "unverified app" warning during your own
> login: click *Advanced → Go to (app name)*.
#### b) Create the OAuth client
1. Go to **APIs & Services → Credentials → Create Credentials → OAuth client ID**.
2. Application type: **Desktop app**. Give it any name → **Create**.
3. Click **Download JSON** and save it in the project as `credentials/client_secret.json`. Create the
`credentials/` folder if needed; it is gitignored.
#### c) Choose the scope *before* authorizing
| `GOOGLE_DRIVE_SCOPE` | Meaning | Choose it if |
|---|---|---|
| `drive.file` (default, least privilege) | The bot only sees files and folders **it created** | You let the bot create its own root folder (step 11, option B) |
| `drive` | Full access to your Drive | You want uploads to go into a folder **you** created |
Set it in `.env` (for example `GOOGLE_DRIVE_SCOPE=drive.file`). Use the **same value on Heroku**.
#### d) Authorize and get the token
```bash
python main.py drive-auth # opens a browser
python main.py drive-auth --no-browser # prints a URL to open yourself
```
Sign in, approve access, and the token is saved to `credentials/token.json`. It contains the refresh
token, client ID and client secret. Then set it on Heroku:
```bash
heroku config:set GOOGLE_OAUTH_TOKEN_JSON="$(cat credentials/token.json)" --app YOUR_APP_NAME
```
If you prefer pasting into the Heroku dashboard, a single-line base64 version also works:
```bash
base64 < credentials/token.json | tr -d '\n'
```
```env
GOOGLE_AUTH_MODE=oauth
GOOGLE_OAUTH_TOKEN_JSON={"token": "...", "refresh_token": "...", "client_id": "...", "client_secret": "...", ...}
```
> If you change `GOOGLE_DRIVE_SCOPE` later, delete `credentials/token.json`, run `drive-auth` again,
> and update the Config Var. To revoke access: <https://myaccount.google.com/permissions>.
### 10. `GOOGLE_SERVICE_ACCOUNT_JSON` (Shared Drive)
1. In the Cloud Console: **IAM & Admin → Service Accounts → Create service account**. Give it a name;
no project roles are needed. Click **Done**.
2. Open the service account → **Keys → Add key → Create new key → JSON**. The file downloads once.
Store it safely.
3. Copy the `client_email` value from the JSON (for example
`uploader@your-project.iam.gserviceaccount.com`).
4. In Google Drive, open the **Shared Drive** → **Manage members** → add that email as
**Content manager**.
5. Set the variables:
```bash
heroku config:set GOOGLE_AUTH_MODE=service_account \
GOOGLE_DRIVE_SCOPE=drive \
GOOGLE_SERVICE_ACCOUNT_JSON="$(cat path/to/key.json)" --app YOUR_APP_NAME
```
Use `GOOGLE_DRIVE_SCOPE=drive`, because the service account must reach a Shared Drive folder it
did not create.
Locally you can instead set `GOOGLE_SERVICE_ACCOUNT_FILE=path/to/key.json`.
> Your organisation may block service-account key creation (policy
> `iam.disableServiceAccountKeyCreation`). In that case use OAuth.
### 11. `GOOGLE_DRIVE_ROOT_FOLDER_ID`
This is the folder that will contain all topic folders. If you leave it empty, topic folders are
created at the top level of My Drive.
**Option A: a folder you created** (needs `GOOGLE_DRIVE_SCOPE=drive`, or a Shared Drive folder)
1. Open the folder in Google Drive in a browser.
2. The URL looks like `https://drive.google.com/drive/folders/1AbCdEfGhIjKlMnOpQrStUvWxYz`.
3. The part after `/folders/` is the ID → `GOOGLE_DRIVE_ROOT_FOLDER_ID`.
**Option B: let the bot create it** (works with the default `drive.file` scope)
1. Run this after `drive-auth`:
```bash
python main.py drive-init --name "Telegram Videos"
```
2. It prints `Set GOOGLE_DRIVE_ROOT_FOLDER_ID=...`. Copy that ID.
```env
GOOGLE_DRIVE_ROOT_FOLDER_ID=1AbCdEfGhIjKlMnOpQrStUvWxYz
```
At startup the bot checks that it can see this folder. If it can't, it stops and explains how to fix
the problem.
### 13. `DRIVE_SHARED_DRIVE_ID` (Shared Drive only)
Open the Shared Drive itself (its top level). The URL is
`https://drive.google.com/drive/folders/<ID>`, and that `<ID>` is the Shared Drive ID.
```env
DRIVE_SHARED_DRIVE_ID=0AbCdEfGhIjKlUk9PVA
```
### `DRIVE_SHARE_MODE`
Leave it as `none` so uploads stay private. Set `anyone_with_link` only if every uploaded video
should be viewable by anyone who has its link.
---
## Part 3 — Database
### 14. `DATABASE_URL`
- **Heroku Deploy button:** the Heroku Postgres add-on is provisioned and `DATABASE_URL` is set
automatically. Nothing to do.
- **Manual Heroku setup:**
```bash
heroku addons:create heroku-postgresql:essential-0 --app YOUR_APP_NAME
heroku config:get DATABASE_URL --app YOUR_APP_NAME # confirm it is set
```
This is a paid plan; check Heroku's current pricing. Don't copy the URL anywhere else, because
Heroku may rotate it.
- **Other PostgreSQL providers:** use their connection string
(`postgresql://USER:PASSWORD@HOST:5432/DBNAME`). TLS is added automatically for remote hosts.
- **Local development:** leave it empty to use SQLite at `data/bot.db`.
On Heroku the bot refuses to start without `DATABASE_URL`, unless you set
`ALLOW_EPHEMERAL_DATABASE=true`. With that set, all progress is lost on every restart.
---
## Optional settings
These are preferences, not credentials. Defaults are shown; full descriptions are in
[`.env.example`](.env.example).
| Variable | Default | Purpose |
|---|---|---|
| `PROCESS_MODE` | `both` | `existing` (history), `new` (new posts), or `both` |
| `START_MESSAGE_ID` | `1` | First message ID to scan |
| `START_DATE` | empty | Only messages on or after this ISO date, for example `2024-01-31` |
| `INCLUDE_DOCUMENT_VIDEOS` | `true` | Also process videos sent as files |
| `FILTER_TOPICS` / `FILTER_BATCHES` | empty | Comma-separated allow-lists |
| `BATCH_SUBFOLDERS_ENABLED` | `false` | `Root/Batch/Topic` folder layout |
| `DUPLICATE_POLICY` | `suffix_number` | `suffix_number`, `suffix_message_id`, `skip`, `fail` |
| `MAX_RETRIES` / `RETRY_DELAY` / `RETRY_MAX_DELAY` | `4` / `30` / `1800` | Retry behaviour (seconds) |
| `PROGRESS_INTERVAL` | `10` | Seconds between progress-message edits |
| `DOWNLOAD_DIR` | `downloads` (`/tmp/telegram_drive_bot` on Heroku) | Temporary video storage |
| `UPLOAD_CHUNK_SIZE` | `16MB` (`8MB` on Heroku) | Upload chunk size (multiple of 256KB) |
| `MIN_FREE_DISK_SPACE` | `1GB` (`512MB` on Heroku) | Free-disk reserve |
| `AUTO_START` | `false` | Start processing at every boot |
| `LOG_LEVEL` | `info` | `debug`, `info`, `warning`, `error` |
| `BOT_SCAN_MAX_EMPTY` / `BOT_SCAN_END_ID` | `200` / empty | Bot-mode history probing limits |
---
## Verify your configuration
Locally, with a filled `.env`:
```bash
python main.py check-config
```
This prints `Configuration OK.` or lists every problem with a fix. On Heroku, run
`heroku logs --tail --app YOUR_APP_NAME` after deploying. A healthy start logs:
- `Database: PostgreSQL`
- `Google Drive credentials validated`
- `Source channel resolved`
The bot then sends you *Bot online* in Telegram.
## Minimal example (Heroku, user mode, OAuth)
```env
TELEGRAM_MODE=user
TELEGRAM_API_ID=1234567
TELEGRAM_API_HASH=0123456789abcdef0123456789abcdef
TELEGRAM_BOT_TOKEN=123456789:replace-with-your-bot-token
TELEGRAM_SESSION_STRING=1BVtsOK...replace-with-your-session-string
ADMIN_USER_IDS=111111111
SOURCE_CHANNEL_ID=-1001234567890
GOOGLE_AUTH_MODE=oauth
GOOGLE_DRIVE_SCOPE=drive.file
GOOGLE_OAUTH_TOKEN_JSON={"refresh_token": "...", "client_id": "...", "client_secret": "...", "token_uri": "https://oauth2.googleapis.com/token"}
GOOGLE_DRIVE_ROOT_FOLDER_ID=1AbCdEfGhIjKlMnOpQrStUvWxYz
# DATABASE_URL is set automatically by Heroku Postgres
```
Procfile
+1
View File
@@ -0,0 +1 @@
worker: python main.py run
README.md
+510
View File
@@ -0,0 +1,510 @@
# Telegram Channel → Google Drive Organizer
[![Deploy to Heroku](https://www.herokucdn.com/deploy/button.svg)](https://heroku.com/deploy?template=https://github.com/YOUR_GITHUB_USERNAME/YOUR_REPOSITORY_NAME)
> **Before the button works:** replace `https://github.com/YOUR_GITHUB_USERNAME/YOUR_REPOSITORY_NAME`
> in the line above with the URL of the **public** GitHub repository that contains this project
> (see [Publishing to GitHub](#publishing-to-github)). That URL is the only value to change.
> The button flow has not been tested against a live repository yet. Test it once after publishing.
A long-running worker that reads video posts from a Telegram channel, downloads each video, and
uploads it to Google Drive. The caption's **Title** becomes the file name and its **Topic** picks the
Drive folder. Videos are processed one at a time. All progress is stored in a database, so the bot
resumes after crashes, restarts and redeploys without uploading anything twice.
## Contents
1. [Overview](#1-overview) · 2. [Features](#2-features) · 3. [Caption format](#3-caption-format-and-metadata-extraction) ·
4. [Drive organization](#4-google-drive-topic-folder-organization) · 5. [Prerequisites](#5-prerequisites) ·
6. [One-click Heroku deployment](#6-one-click-heroku-deployment) · 7. [Telegram setup](#7-telegram-bot-setup) ·
8. [Historical access](#8-historical-channel-access-requirements) · 9. [Google Drive auth](#9-google-drive-authentication) ·
10. [Config Vars](#10-required-heroku-config-vars) · 11. [Database](#11-database-provisioning-and-configuration) ·
12. [First deployment](#12-first-deployment-and-startup) · 13. [Commands](#13-telegram-commands) · 14. [Logs](#14-checking-logs) ·
15. [Restart / redeploy](#15-restarting-and-redeploying) · 16. [Troubleshooting](#16-troubleshooting) ·
17. [Security](#17-security-and-credential-handling) · 18. [Limitations](#18-known-limitations-and-heroku-resource-constraints) ·
19. [Local development](#19-local-development)
---
## 1. Overview
```
Telegram channel ──(Telethon / MTProto)──► queue worker (one video at a time) ──► Google Drive
│ (resumable upload,
Admin chat ◄── bot commands & progress ─────────┤ verified by size + MD5)
▼
PostgreSQL (Heroku) / SQLite (local)
```
For each video: **discover → parse caption → resolve the topic folder → download to a temp disk →
upload (resumable) → verify size and MD5 in Drive → save the Drive file ID → mark COMPLETED → delete
the temp file → next video.**
## 2. Features
- Caption parsing for `Index` (one or many), `Title`, `Topic`, `Batch` and `Extracted By`.
- Topic folders are created once and reused. Optional `Batch/Topic` layout.
- Files are never overwritten. Name collisions get `(2)`, `(3)`, … or `[message id]`, or are skipped or failed, per your policy.
- No duplicate uploads, even if the process dies just after Drive accepted the file. Every upload carries a hidden `appProperties` tag (`channel:message`) that is checked before uploading.
- Resumable uploads that survive restarts. The upload session URI is stored in the database.
- Resumable downloads, continuing from the last 1 MiB boundary of a partial file.
- Explicit state machine: `PENDING → DOWNLOADING → DOWNLOADED → UPLOADING → COMPLETED`, plus `RETRY_WAIT`, `FAILED` and `SKIPPED`. Every transition is validated and logged in the `transitions` table.
- Bounded retries with exponential backoff, honouring Telegram flood waits and Drive rate limits. Permission errors are not retried.
- Disk-space checks before each download. A full disk, Drive quota or expired auth halts the queue cleanly instead of burning retries.
- Admin-only Telegram commands, plus one progress message that is edited in place rather than spammed.
- Graceful shutdown on `SIGTERM` (Heroku restarts) at a resumable checkpoint.
- A single-worker guard (PostgreSQL advisory lock) prevents two dynos from processing the same queue.
## 3. Caption format and metadata extraction
```
Index: 001
Index: 002
Title: Class 01 Aman sir.mp4
Topic: Misc• Grammar Practice
Batch: English ( SSC MAINS )
Extracted By: @PUTINxINDIA
```
| Field | Use |
|---|---|
| `Title` | Drive file name (`Class 01 Aman sir.mp4`) |
| `Topic` | Drive folder (`Misc• Grammar Practice`) |
| `Index` | All values kept in order (`["001", "002"]`); never used as the file name. The Telegram message ID is the unique identity. |
| `Batch` | Stored; optional `Batch/Topic` folders |
| `Extracted By` | Stored in the record |
Parsing rules:
- Fields can appear in any order.
- Labels are case-insensitive.
- Blank lines, emoji or bullet prefixes, extra spaces and full-width colons are allowed.
- Unicode text is preserved.
- Parsing is line-based, so one field never swallows the next.
- A flattened single-line caption is split at each known label.
Missing data never makes a video silently disappear:
| Case | Behaviour |
|---|---|
| No Topic | Folder `Uncategorized` (`FALLBACK_TOPIC_FOLDER`) |
| No Title | Telegram file name, else `telegram_<message id>.<ext>`. A warning is recorded. `MISSING_TITLE_POLICY=fail` marks it FAILED instead. |
| Title without extension | Extension taken from the real file's magic bytes, else from the MIME type |
| Title extension ≠ real container | Kept as written, with a warning. `EXTENSION_MISMATCH_POLICY=append` appends the real one. |
| `/` `\` or control characters | Replaced (`STRICT_FILENAMES=true` also replaces `<>:"\|?*`) |
## 4. Google Drive topic-folder organization
```
<GOOGLE_DRIVE_ROOT_FOLDER_ID>/
├── Misc• Grammar Practice/
│ ├── Class 01 Aman sir.mp4
│ └── Class 02 Aman sir.mp4
├── Error Detection/
│ └── Class 01 Aman sir.mp4
└── Uncategorized/
```
Folder lookup goes database cache → Drive search → create. If duplicate folders already exist, the
oldest is reused. If a folder-create response is lost, the bot searches again before retrying, so
it never creates a second copy. Set `BATCH_SUBFOLDERS_ENABLED=true` for `Root/Batch/Topic/`.
**Duplicate policy (`DUPLICATE_POLICY`)** applies when a *different* file already has the target
name in the folder:
| Value | Result |
|---|---|
| `suffix_number` (default) | `Class 01 Aman sir (2).mp4` |
| `suffix_message_id` | `Class 01 Aman sir [1234].mp4` |
| `skip` | Item marked SKIPPED with the reason |
| `fail` | Item marked FAILED |
An existing file is reused, not re-uploaded, only when its identity is verified. That means either
its hidden tag matches this Telegram message, or it has the same size and MD5
(`REUSE_IDENTICAL_FILES=true`). The same Telegram media re-posted under a new message ID is SKIPPED
by default (`DUPLICATE_MEDIA_POLICY`).
## 5. Prerequisites
- A Telegram account, plus API credentials from <https://my.telegram.org>.
- A Telegram bot from [@BotFather](https://t.me/BotFather).
- A Google account with a Google Cloud project (free) with the Drive API enabled.
- **A local computer with Python 3.12** for two one-time interactive logins. Telegram and Google logins need a human, so they can't run on Heroku.
- A Heroku account with billing enabled. The worker dyno and Heroku Postgres are paid resources; check current pricing on Heroku.
- Optional: the [Heroku CLI](https://devcenter.heroku.com/articles/heroku-cli) and the [GitHub CLI](https://cli.github.com/).
## 6. One-click Heroku deployment
The button (top of this page) opens Heroku's deploy form. It:
- creates the app from this repository's `app.json`;
- builds it with the official Python buildpack;
- provisions **Heroku Postgres (`essential-0`)**, which sets `DATABASE_URL`;
- starts **one `basic` worker dyno** (`worker: python main.py run`); there is no web process.
**It is not fully one-click.** Two secrets must be created on your computer first, because they need
interactive logins:
1. `TELEGRAM_SESSION_STRING`, only for `TELEGRAM_MODE=user` (see [§7](#7-telegram-bot-setup)).
2. `GOOGLE_OAUTH_TOKEN_JSON` or `GOOGLE_SERVICE_ACCOUNT_JSON` (see [§9](#9-google-drive-authentication)).
Prepare those, click the button, fill the form, then press **Deploy app**. If add-on provisioning
fails (for example because billing isn't verified or the plan isn't available), create the app
anyway and follow [§11](#11-database-provisioning-and-configuration) manually. The worker refuses
to start on Heroku without a database rather than silently losing progress.
Manual alternative with the CLI (from the repository root):
```bash
heroku create YOUR_APP_NAME --stack heroku-24
heroku addons:create heroku-postgresql:essential-0 --app YOUR_APP_NAME
heroku config:set TELEGRAM_API_ID=1234567 TELEGRAM_API_HASH=your_api_hash \
TELEGRAM_BOT_TOKEN=123456:your-bot-token ADMIN_USER_IDS=111111111 \
SOURCE_CHANNEL_ID=@your_channel --app YOUR_APP_NAME
heroku config:set TELEGRAM_SESSION_STRING="$(cat data/sessions/session_string.txt)" --app YOUR_APP_NAME
heroku config:set GOOGLE_OAUTH_TOKEN_JSON="$(cat credentials/token.json)" --app YOUR_APP_NAME
git push heroku main
heroku ps:scale worker=1 --app YOUR_APP_NAME
```
### Publishing to GitHub
This directory (the one containing `Procfile` and `app.json`) must be the **repository root**.
```bash
cd telegram_drive_bot
git init -b main
git add .
git status # confirm: no .env, credentials/, data/, *.session, downloads/
git commit -m "Telegram to Google Drive organizer"
gh repo create YOUR_REPOSITORY_NAME --public --source=. --push
# or create the repo on github.com, then:
# git remote add origin https://github.com/YOUR_GITHUB_USERNAME/YOUR_REPOSITORY_NAME.git && git push -u origin main
```
Then put that repository URL into the button link at the top of this README, commit, and push.
The Heroku button only works with **public** repositories.
## 7. Telegram bot setup
1. **API ID / hash:** log in at <https://my.telegram.org> → *API development tools* → create an app.
Copy `api_id` → `TELEGRAM_API_ID` and `api_hash` → `TELEGRAM_API_HASH`.
These are needed in both modes, because the bot talks MTProto through Telethon.
2. **Bot:** talk to @BotFather → `/newbot` → copy the token → `TELEGRAM_BOT_TOKEN`.
**Open your bot and press Start.** A bot can't message you until you do, so you won't get progress messages otherwise.
3. **Admin IDs:** message [@userinfobot](https://t.me/userinfobot) to get your numeric ID → `ADMIN_USER_IDS`
(comma-separate several). Only these users can run commands.
4. **Source channel (`SOURCE_CHANNEL_ID`):** any of these forms work:
- a public channel: `@channelname` or `https://t.me/channelname`;
- a private channel: the numeric ID `-100…`, or a post link `https://t.me/c/1234567890/55` (use *Copy Post Link* in Telegram);
- forwarding a post to @userinfobot also shows the channel ID.
5. **Choose a mode (`TELEGRAM_MODE`):**
- `user` (default): reads the channel as *your* account, which must be a member. Create the session string locally:
```bash
python main.py login --session-string
```
Telethon asks for your phone number, the login code and, only if enabled, your 2FA password.
None of these are stored. The resulting session string is written to
`data/sessions/session_string.txt` (mode 600). Copy it into `TELEGRAM_SESSION_STRING`, then delete the file.
- `bot`: add the bot to the channel as an **administrator**. No session string is needed.
## 8. Historical channel access requirements
| Capability | `user` mode | `bot` mode |
|---|---|---|
| New posts (live) | ✅ | ✅ (bot must be a channel admin) |
| Existing history | ✅ full history the account can see | ⚠️ best effort: bots can't list history, so message IDs are probed in batches until `BOT_SCAN_MAX_EMPTY` consecutive gaps (or up to `BOT_SCAN_END_ID`) |
| Large files | ✅ up to Telegram's file limit | ✅ uses MTProto, so the HTTP Bot API's 20 MB download limit does **not** apply |
| Private channel by numeric ID | ✅ | ⚠️ the bot learns the channel's access key from its first live post; the key is then stored in the database |
| Extra secret needed | `TELEGRAM_SESSION_STRING` | none |
`PROCESS_MODE` controls what is queued:
- `existing`: history only (from `START_MESSAGE_ID` / `START_DATE`).
- `new`: only posts after the first start.
- `both` (default).
The bot never deletes or edits Telegram messages and doesn't bypass channel restrictions. A video
in a channel with content protection (no forwarding/saving) may not be downloadable.
## 9. Google Drive authentication
Pick **one**.
### Option A — OAuth (personal Google Drive, recommended)
1. In the [Google Cloud Console](https://console.cloud.google.com/): create a project, then enable the **Google Drive API**.
2. Configure the **OAuth consent screen**: *External*, add yourself as a test user.
Then set the publishing status to **In production**. In *Testing* status, Google expires refresh
tokens after 7 days and the bot would lose Drive access weekly. With the default `drive.file`
scope no verification is required. With the full `drive` scope Google shows an "unverified app"
screen during your own login (*Advanced → continue*).
3. Under *Credentials*, create an **OAuth client ID** of type **Desktop app**. Download the JSON to
`credentials/client_secret.json` in this project.
4. Locally: `python main.py drive-auth`. This opens a browser; use `--no-browser` to get a URL to open yourself.
It writes `credentials/token.json` (mode 600).
5. Heroku: `heroku config:set GOOGLE_OAUTH_TOKEN_JSON="$(cat credentials/token.json)" --app YOUR_APP_NAME`.
Multi-line JSON works through the CLI. In the dashboard you can paste raw JSON, or a single-line
base64 version: `base64 < credentials/token.json | tr -d '\n'`.
**Scope and root folder.** `GOOGLE_DRIVE_SCOPE=drive.file` (least privilege) lets the bot see only
files and folders *it created*. So either:
- let the bot create the root folder: `python main.py drive-init --name "Telegram Videos"` prints the ID for `GOOGLE_DRIVE_ROOT_FOLDER_ID`; or
- use a folder you created yourself: set `GOOGLE_DRIVE_SCOPE=drive` **before** running `drive-auth`.
The folder ID is the part after `/folders/` in its URL. The bot checks at startup that it can see
the folder and stops with an explanation if it can't.
### Option B — Service account (Google Workspace Shared Drive)
A service account has **no storage quota in a personal My Drive**. Uploads into a folder that a
normal Gmail user merely shared with it fail with `storageQuotaExceeded`. Use it only with a
**Shared Drive**:
1. Create a service account, then create a JSON key for it.
2. In the Shared Drive, *Manage members* → add the service account's `client_email` as **Content manager**.
3. Set:
- `GOOGLE_AUTH_MODE=service_account`
- `GOOGLE_SERVICE_ACCOUNT_JSON="$(cat key.json)"`
- `DRIVE_SHARED_DRIVE_ID=<shared drive id>`
- `GOOGLE_DRIVE_ROOT_FOLDER_ID=<folder id inside it>`
- `GOOGLE_DRIVE_SCOPE=drive` (needed to use a folder the service account did not create).
Credentials from environment variables are parsed in memory. Nothing is written to the dyno's disk,
and they are validated by an authenticated Drive call at startup. Uploaded files are **private**
unless you set `DRIVE_SHARE_MODE=anyone_with_link`.
## 10. Required Heroku Config Vars
Set these in the dashboard (*Settings → Reveal Config Vars*) or with
`heroku config:set NAME=value --app YOUR_APP_NAME`.
| Variable | Required | Notes |
|---|---|---|
| `TELEGRAM_API_ID`, `TELEGRAM_API_HASH` | yes | my.telegram.org |
| `TELEGRAM_BOT_TOKEN` | yes | @BotFather; command interface (and channel access in bot mode) |
| `ADMIN_USER_IDS` | yes | numeric IDs, comma-separated |
| `SOURCE_CHANNEL_ID` | yes | `@name`, link, or `-100…` |
| `TELEGRAM_MODE` | yes (default `user`) | `user` or `bot` |
| `TELEGRAM_SESSION_STRING` | **when `TELEGRAM_MODE=user`** | from `main.py login --session-string` |
| `GOOGLE_AUTH_MODE` | yes (default `oauth`) | `oauth` or `service_account` |
| `GOOGLE_OAUTH_TOKEN_JSON` | **when `oauth`** | contents of `credentials/token.json` |
| `GOOGLE_SERVICE_ACCOUNT_JSON` | **when `service_account`** | key JSON |
| `GOOGLE_DRIVE_ROOT_FOLDER_ID` | recommended | empty = top level of My Drive |
| `GOOGLE_DRIVE_SCOPE` | no (`drive.file`) | must match the scope used in `drive-auth` |
| `DATABASE_URL` | **yes on Heroku** | set automatically by Heroku Postgres |
| `PROCESS_MODE` | no (`both`) | `existing` / `new` / `both` |
| `DUPLICATE_POLICY` | no (`suffix_number`) | see §4 |
| `BATCH_SUBFOLDERS_ENABLED` | no (`false`) | `Root/Batch/Topic` |
| `MAX_RETRIES` / `RETRY_DELAY` | no (`4` / `30`) | automatic retries / base delay in seconds |
| `PROGRESS_INTERVAL` | no (`10`) | seconds between progress-message edits |
| `DOWNLOAD_DIR` | no (`/tmp/telegram_drive_bot` in app.json) | temporary files |
| `UPLOAD_CHUNK_SIZE` | no (`8MB` in app.json) | multiple of 256KB; memory per chunk |
| `AUTO_START` | no (`false`) | start processing on every boot |
Step-by-step instructions for obtaining every value are in **[ENV_SETUP.md](ENV_SETUP.md)**. Every other supported variable is listed with an explanation in [`.env.example`](.env.example).
## 11. Database provisioning and configuration
Production state lives in **PostgreSQL**:
- video records with all caption metadata;
- states, attempts, errors, timestamps and Drive file IDs;
- upload session URIs;
- the folder map;
- discovery and completion cursors;
- duplicate-detection data;
- the transition history;
- the pause/run/stop choice.
- **Via the button:** `app.json` requests `heroku-postgresql:essential-0`, a paid plan.
- **Manually:** `heroku addons:create heroku-postgresql:essential-0 --app YOUR_APP_NAME`
(or any plan, or any PostgreSQL reachable through `DATABASE_URL`).
- The schema is created and migrated automatically at startup. The version is tracked in the `kv` table.
- TLS (`sslmode=require`) is added automatically for remote hosts.
- **Without `DATABASE_URL` on Heroku the worker refuses to start.** SQLite on a dyno would be wiped
at every restart, and the bot would re-upload everything. To accept that risk explicitly, set
`ALLOW_EPHEMERAL_DATABASE=true`; a warning is logged at every start.
- Locally, without `DATABASE_URL`, SQLite at `data/bot.db` is used. This is for development only.
## 12. First deployment and startup
1. Deploy (button or CLI). Watch the logs: `heroku logs --tail --app YOUR_APP_NAME`.
You should see the database backend, `Google Drive credentials validated`, and `Source channel resolved`.
2. The bot messages each admin: *Bot online (PostgreSQL); send /sync then /run to start.*
3. Send `/sync` to discover channel videos, then `/run`. Watch `/progress`.
4. After `/run`, the choice is stored. After any dyno restart the bot syncs and resumes on its own.
`/pause` and `/stop` are also remembered.
## 13. Telegram commands
Send these to the bot in a **private chat**. Users not in `ADMIN_USER_IDS` are refused.
| Command | Effect |
|---|---|
| `/start`, `/help` | Instructions |
| `/status` | Queue state, the current video, live progress |
| `/progress` | Processed / pending / retrying / failed / skipped / remaining counts |
| `/sync` | Discover unseen messages and queue them (idempotent). `/sync all` rescans from `START_MESSAGE_ID`. |
| `/run` | Start processing |
| `/pause` | Stop the current transfer at the next chunk boundary (resumable) and wait |
| `/resume` | Continue |
| `/stop` | Checkpoint and end the worker until `/run` |
| `/retry` | Re-queue FAILED items with a fresh retry budget |
| `/failed` | Failed items with short error summaries |
| `/folders` | Folder count and uploads per topic |
| `/check ID` | Full record for one Telegram message ID |
Progress is a single message, edited every `PROGRESS_INTERVAL` seconds:
```
Processing video 12 of 150
Topic: Misc• Grammar Practice
Title: Class 01 Aman sir.mp4
Status: Uploading
Download: 100% (1.2 GB / 1.2 GB)
Upload: 45% (540.0 MB / 1.2 GB) · 18.3 MB/s · ETA 36s
```
## 14. Checking logs
```bash
heroku logs --tail --app YOUR_APP_NAME # live
heroku logs -n 1500 --app YOUR_APP_NAME # recent history
heroku ps --app YOUR_APP_NAME # dyno status
```
On Heroku, logs go to stdout only. Locally they also go to `logs/bot.log`, rotated at 10 MB.
Tokens, hashes, session strings, private keys and the database URL are redacted from log output.
## 15. Restarting and redeploying
```bash
heroku ps:restart worker --app YOUR_APP_NAME
git push heroku main # or push to GitHub with automatic deploys enabled
heroku ps:scale worker=1 --app YOUR_APP_NAME
```
On `SIGTERM`, Heroku allows 30 seconds before it kills the process. The bot:
1. stops at the next download or upload chunk boundary;
2. saves state;
3. disconnects.
On the next start it:
1. returns interrupted downloads to PENDING;
2. checks Drive for any upload that may have finished, using the hidden tag;
3. re-downloads files lost with the old dyno's disk;
4. **continues the saved upload session**, when the re-downloaded file is byte-identical (MD5);
5. resumes processing if it was running.
**Scaling:** run **exactly one** worker. A second worker is blocked by a PostgreSQL advisory lock.
It waits `INSTANCE_LOCK_WAIT` seconds (90 on Heroku, which covers deploy overlap), then exits with
code 3. Scaling above 1 therefore only produces crash-looping dynos, not parallel processing. The
pipeline is intentionally sequential.
## 16. Troubleshooting
| Symptom | Fix |
|---|---|
| `DATABASE_URL is not set` | Attach Heroku Postgres (§11) |
| `TELEGRAM_SESSION_STRING is required…` | Run `python main.py login --session-string` locally (§7), or use `TELEGRAM_MODE=bot` |
| `Telegram user session is not authorized` | The session was revoked (e.g. *Terminate all sessions*); create a new session string |
| `Cannot resolve SOURCE_CHANNEL_ID` | Use `@username` for public channels; make sure the account/bot is a member/admin; in bot mode post once in the channel |
| `DRIVE_ROOT_FOLDER_ID is not visible…` | `drive.file` scope can't see folders you created: use `drive-init`, or switch to `GOOGLE_DRIVE_SCOPE=drive` and redo `drive-auth` |
| `Google authorization expired or was revoked` | Redo `drive-auth` and update `GOOGLE_OAUTH_TOKEN_JSON`; publish the consent screen to *In production* (§9) |
| `invalid_scope` on refresh | The token was created with another scope; delete `credentials/token.json` and run `drive-auth` again |
| `storageQuotaExceeded` | Drive is full, or you used a service account without a Shared Drive |
| `Not enough free disk space` | A video is larger than the dyno's free disk; lower `MIN_FREE_DISK_SPACE` with care, or process that video locally |
| No progress messages | Press *Start* in your bot chat; check `ADMIN_USER_IDS` / `NOTIFY_CHAT_IDS` |
| `Another worker is already processing` | `heroku ps:scale worker=1`; locally, stop the other process |
| Item stuck in FAILED | `/check ID` shows the error; fix the cause, then `/retry` |
| Long pause after `FloodWait` | Telegram rate limit; the bot waits the time Telegram asks for |
## 17. Security and credential handling
- All secrets come from environment variables / Config Vars; nothing is hard-coded.
- `.gitignore` excludes `.env`, `credentials/`, `token.json`, `*.session`, `session_string.txt`, `data/` and `downloads/`.
- The session string and OAuth token are full-access credentials. If one leaks:
- Telegram: *Settings → Devices → terminate the session*;
- Google: <https://myaccount.google.com/permissions>.
- Uploaded files stay private by default. Drive scope defaults to least privilege (`drive.file`).
- Commands are admin-only and private-chat only. Replies never include credentials, local paths or stack traces.
- Telegram login codes and 2FA passwords are entered interactively, locally, and never stored.
## 18. Known limitations and Heroku resource constraints
- **Ephemeral disk:** each video is staged on the dyno's temporary filesystem. Heroku doesn't
guarantee a specific capacity, and the disk is wiped on every restart (at least daily).
- The bot checks free space before each download and halts with a message if a video won't fit.
- Multi-GB videos may not fit on a dyno; run the bot locally for those.
- **Memory:** one upload chunk (`UPLOAD_CHUNK_SIZE`) is held in RAM at a time. `8MB` is comfortable on 512 MB dynos.
- **Daily restarts:** in-flight transfers resume from their checkpoint. A partially downloaded file
on the old dyno is lost and is downloaded again, but the Drive upload session is reused.
- **Bot mode:** history is best effort (ID probing), and private channels need one live post before the numeric ID resolves.
- **Telegram:** flood limits apply. Downloads go through Telegram's servers at whatever speed they allow.
- **Drive:** resumable sessions expire after about a week. Older sessions are discarded and the
upload restarts after a Drive check for an already-finished copy.
- **Captions edited after discovery** are not re-read. Use `/check ID` to see what was stored.
- **Cost:** the worker dyno and Postgres add-on are paid Heroku resources.
### What has and has not been verified
- **Automated tests (171):** run against SQLite **and** a real local PostgreSQL 16, with mocked Telegram and Google Drive.
- **Not tested:** real Telegram logins and downloads, real Google Drive uploads, and an actual Heroku deployment. Follow the manual end-to-end procedure below before relying on it.
## 19. Local development
```bash
python3.12 -m venv .venv && source .venv/bin/activate
pip install -r requirements-dev.txt
cp .env.example .env # fill in values
python main.py login # user mode: creates data/sessions/user.session
python main.py drive-auth # creates credentials/token.json
python main.py check-config
python main.py run # or ./start.sh
```
Other subcommands:
- `python main.py sync [--rescan]`
- `python main.py process` (headless: sync, process everything, exit)
- `python main.py status`
- `python main.py drive-init --name "Telegram Videos"`
### Tests
```bash
pytest
```
`pgserver` from `requirements-dev.txt` starts a throwaway PostgreSQL so every database test runs on
both backends. Without it, only SQLite is tested. The suite covers:
- caption parsing (the supplied example, multiple Index lines, missing fields, Unicode, flattened captions);
- naming and collisions;
- Drive folder reuse and dedupe;
- the resumable-upload protocol;
- state transitions and duplicate messages;
- download and upload failure recovery;
- restart during upload, and restart with a wiped disk;
- a successful upload followed by a database-write failure;
- sequential processing, pause/resume and graceful shutdown;
- admin-only commands;
- configuration and Heroku validation, env-var credentials, log redaction;
- `app.json` / Procfile consistency.
### Manual end-to-end test (real integrations)
1. Create a private test channel. Post four videos:
- two with `Topic: Test A`;
- one with `Topic: Test B`;
- one with no caption.
Give two of them the same `Title`.
2. Point the bot at a fresh Drive folder. Then `/sync`, `/run`, and watch `/progress`.
3. Check Drive:
- `Test A/` holds two files, one with a `(2)` suffix if the titles matched;
- `Test B/` holds one;
- `Uncategorized/` holds the caption-less video;
- file sizes match Telegram's.
4. Kill the process (`kill -9`, or `heroku ps:restart`) **during an upload** of a large video.
Restart it. Confirm the upload continues and Drive ends up with exactly one copy.
5. `/stop`, post a new video, `/run`: only the new video is processed.
6. Revoke Drive access temporarily and confirm the queue halts with a clear message. Restore access, then `/resume`.
app.json
+189
View File
@@ -0,0 +1,189 @@
{
"name": "Telegram Drive Organizer",
"description": "Telegram bot that downloads channel videos (whole channel or an /upload message range) into Google Drive topic folders.",
"keywords": [
"telegram",
"google-drive",
"telethon",
"video",
"worker"
],
"stack": "heroku-24",
"buildpacks": [
{
"url": "heroku/python"
}
],
"formation": {
"worker": {
"quantity": 1,
"size": "basic"
}
},
"addons": [
{
"plan": "heroku-postgresql:essential-0",
"as": "DATABASE"
}
],
"env": {
"TELEGRAM_MODE": {
"description": "How the bot reads the channel: 'user' (your Telegram account; full history; needs TELEGRAM_SESSION_STRING) or 'bot' (bot admin of the channel; new posts, best-effort history).",
"value": "user",
"required": true
},
"TELEGRAM_API_ID": {
"description": "Numeric API ID from https://my.telegram.org > API development tools. Needed in both modes.",
"required": true
},
"TELEGRAM_API_HASH": {
"description": "API hash from https://my.telegram.org > API development tools. Keep secret.",
"required": true
},
"TELEGRAM_BOT_TOKEN": {
"description": "Bot token from @BotFather. Provides the admin command interface (/run, /status, ...); in bot mode it also reads the channel.",
"required": true
},
"TELEGRAM_SESSION_STRING": {
"description": "Only for TELEGRAM_MODE=user: string session created locally with `python main.py login --session-string`. Grants access to your account - keep secret.",
"required": false
},
"ADMIN_USER_IDS": {
"description": "Comma-separated numeric Telegram user IDs allowed to run commands (get yours from @userinfobot).",
"required": true
},
"SOURCE_CHANNEL_ID": {
"description": "Source channel: @username, https://t.me/... link, or numeric -100... ID.",
"required": true
},
"GOOGLE_AUTH_MODE": {
"description": "'oauth' for a personal Google Drive (uses GOOGLE_OAUTH_TOKEN_JSON) or 'service_account' for a Shared Drive (uses GOOGLE_SERVICE_ACCOUNT_JSON).",
"value": "oauth",
"required": true
},
"GOOGLE_OAUTH_TOKEN_JSON": {
"description": "Required when GOOGLE_AUTH_MODE=oauth: contents of credentials/token.json produced locally by `python main.py drive-auth` (raw JSON or base64).",
"required": false
},
"GOOGLE_SERVICE_ACCOUNT_JSON": {
"description": "Required when GOOGLE_AUTH_MODE=service_account: the service-account key JSON (raw or base64). Service accounts only work with a Shared Drive they are a member of.",
"required": false
},
"GOOGLE_DRIVE_SCOPE": {
"description": "'drive.file' (least privilege; the bot only sees folders it created) or 'drive' (needed to use a folder you created yourself). Must match the scope used during drive-auth.",
"value": "drive.file",
"required": false
},
"GOOGLE_DRIVE_ROOT_FOLDER_ID": {
"description": "ID of the Drive folder that will contain the topic folders (the part after /folders/ in its URL). Empty = top level of My Drive.",
"required": false
},
"DRIVE_SHARED_DRIVE_ID": {
"description": "Shared Drive ID when the root folder lives in a Shared Drive (service-account setups).",
"required": false
},
"DRIVE_SHARE_MODE": {
"description": "'none' keeps uploads private (default). 'anyone_with_link' makes every uploaded file readable by anyone with the link.",
"value": "none",
"required": false
},
"PROCESS_MODE": {
"description": "'existing' (channel history only), 'new' (new posts only) or 'both'.",
"value": "both",
"required": false
},
"START_MESSAGE_ID": {
"description": "First channel message ID to consider when scanning history.",
"value": "1",
"required": false
},
"INCLUDE_DOCUMENT_VIDEOS": {
"description": "Also process videos that were sent as files/documents.",
"value": "true",
"required": false
},
"FILTER_TOPICS": {
"description": "Optional comma-separated list of Topic values to process; others are recorded as SKIPPED.",
"required": false
},
"FILTER_BATCHES": {
"description": "Optional comma-separated list of Batch values to process.",
"required": false
},
"BATCH_SUBFOLDERS_ENABLED": {
"description": "true = Root/Batch/Topic folders; false = Root/Topic (default).",
"value": "false",
"required": false
},
"DUPLICATE_POLICY": {
"description": "When a different file already has the same name in the topic folder: suffix_number ('Name (2).mp4'), suffix_message_id ('Name [123].mp4'), skip, or fail. Files are never overwritten.",
"value": "suffix_number",
"required": false
},
"MAX_RETRIES": {
"description": "Automatic retries for a failing video before it is marked FAILED.",
"value": "4",
"required": false
},
"RETRY_DELAY": {
"description": "Base retry delay in seconds (doubles on each retry, capped by RETRY_MAX_DELAY).",
"value": "30",
"required": false
},
"PROGRESS_INTERVAL": {
"description": "Seconds between edits of the Telegram progress message.",
"value": "10",
"required": false
},
"DOWNLOAD_DIR": {
"description": "Temporary download directory on the dyno's ephemeral disk.",
"value": "/tmp/telegram_drive_bot",
"required": false
},
"UPLOAD_CHUNK_SIZE": {
"description": "Resumable upload chunk size (multiple of 256KB). Held in memory per chunk; 8MB suits 512MB dynos.",
"value": "8MB",
"required": false
},
"MIN_FREE_DISK_SPACE": {
"description": "Free disk space to keep in reserve; a download is not started if it would go below this.",
"value": "512MB",
"required": false
},
"AUTO_START": {
"description": "true = start processing at every boot. Otherwise processing starts with /run and then resumes automatically after restarts.",
"value": "false",
"required": false
},
"LOG_LEVEL": {
"description": "debug, info, warning or error.",
"value": "info",
"required": false
},
"UPLOAD_RESTRICT_TO_SOURCE_CHANNEL": {
"description": "true = /upload links must point to SOURCE_CHANNEL_ID; false = any channel the Telegram account or bot can access.",
"value": "true",
"required": false
},
"UPLOAD_SETUP_TIMEOUT": {
"description": "Minutes before an unfinished /upload setup (waiting for links or /confirm) expires.",
"value": "30",
"required": false
},
"UPLOAD_MAX_RANGE": {
"description": "Largest number of message IDs one /upload task may cover.",
"value": "5000",
"required": false
},
"KEEP_FAILED_DOWNLOADS": {
"description": "Keep local files of failed videos on the dyno disk (default false: delete; they are downloaded again on retry).",
"value": "false",
"required": false
},
"TEMP_CLEANUP_INTERVAL": {
"description": "Seconds between sweeps that delete abandoned temporary files.",
"value": "3600",
"required": false
}
}
}
bot_handlers.py
+542
View File
@@ -0,0 +1,542 @@
"""Telegram control commands (admin only), the /upload conversation, and the notifier.
The /upload conversation is persisted in the ``upload_tasks`` table (statuses
AWAITING_START -> AWAITING_END -> AWAITING_CONFIRMATION), keyed by the admin's
Telegram user ID, so every administrator has an independent setup and a restart
does not lose it.
"""
from __future__ import annotations
import logging
import re
from datetime import datetime, timedelta
from database import iso, utcnow
from errors import BotError
from link_parser import EXAMPLE, LinkError, parse_message_link
from models import ACTIVE_TASK_STATUSES, SETUP_TASK_STATUSES, State, TaskStatus, UploadTask, VideoRecord
from progress import HtmlReply, esc
from queue_manager import QueueManager, public_error
log = logging.getLogger(__name__)
HELP = """Commands (administrators only):
/upload - process a range of channel messages (asks for start and end links)
/confirm - start the prepared upload task
/cancel - cancel your upload setup or your active task (/cancel ID for a specific task)
/taskstatus - status of upload tasks (/taskstatus ID for one task)
/uploadhistory - recent upload tasks and their results
/status - current processing state
/progress - processed / pending / failed counts
/run - process every discovered channel video
/pause - pause at the next safe checkpoint
/resume - continue after a pause or /stop
/stop - checkpoint the current video and stop
/sync - discover new channel videos and queue them
/retry - re-queue failed items
/failed - list failed items with error summaries
/folders - topic folders and uploaded files
/check ID - details for one Telegram message ID
/help - this message"""
START_PROMPT = (
"📥 <b>Start Upload Task</b>\n\n"
"Please send the Telegram link of the first message you want to process.\n\n"
"Example:\nhttps://t.me/c/3909337424/14"
)
END_PROMPT = (
"📌 <b>Enter the End Link</b>\n\n"
"Now send the Telegram link of the last message you want to process.\n\n"
"Example:\nhttps://t.me/c/3909337424/20"
)
_STEP = {
TaskStatus.AWAITING_START: "waiting for the start link",
TaskStatus.AWAITING_END: "waiting for the end link",
TaskStatus.AWAITING_CONFIRMATION: "waiting for /confirm",
}
_COMMAND_RE = re.compile(r"^/(?P<cmd>[A-Za-z_]+)(?:@\w+)?(?:\s+(?P<arg>.*))?$", re.DOTALL)
_LINK_HINT = re.compile(r"(?:t\.me|telegram\.me|telegram\.dog)/", re.IGNORECASE)
def _short(text: str | None, limit: int = 160) -> str:
if not text:
return "-"
text = " ".join(text.split())
# never echo filesystem paths back to chat
text = re.sub(r"(?:/[^\s/]+)+/([^\s/]+)", r"…/\1", text)
return text if len(text) <= limit else text[: limit - 1] + "…"
class CommandRouter:
def __init__(self, config, db, queue: QueueManager):
self.config = config
self.db = db
self.queue = queue
def is_admin(self, user_id: int | None) -> bool:
return user_id is not None and user_id in self.config.admin_ids
async def handle(self, user_id: int | None, text: str) -> str | None:
stripped = (text or "").strip()
match = _COMMAND_RE.match(stripped)
if not match:
if not self.is_admin(user_id):
return None # ignore chatter from non-admins
try:
return await self.on_text(user_id, stripped)
except Exception as exc:
log.exception("Handling a message from %s failed", user_id)
return f"Could not process that message: {public_error(exc)}"
if not self.is_admin(user_id):
log.warning("Rejected command from non-admin user %s", user_id)
return "You are not authorized to control this bot."
command = match.group("cmd").lower()
handler = getattr(self, f"cmd_{command}", None)
if handler is None:
return "Unknown command. Send /help for the list of commands."
try:
return await handler((match.group("arg") or "").strip(), user_id)
except Exception as exc:
log.exception("Command /%s failed", command)
return f"Command failed: {public_error(exc)}"
# -- /upload conversation --------------------------------------------------------------
def _active_setup(self, user_id: int) -> tuple[UploadTask | None, bool]:
"""This admin's unfinished setup, expiring it if it is older than the timeout."""
setups = self.db.tasks_with_status(*SETUP_TASK_STATUSES, admin_user_id=user_id)
if not setups:
return None, False
setup = setups[-1]
for stale in setups[:-1]: # defensive: only one setup per admin may exist
self.db.update_task(stale.id, status=TaskStatus.CANCELLED, last_error="superseded", completed_at=iso(utcnow()))
age = utcnow() - datetime.fromisoformat(setup.updated_at)
if age > timedelta(minutes=self.config.upload_setup_timeout):
self.db.update_task(
setup.id, status=TaskStatus.CANCELLED, last_error="setup expired", completed_at=iso(utcnow())
)
log.info("Upload setup #%s of admin %s expired", setup.id, user_id)
return None, True
return setup, False
async def on_text(self, user_id: int, text: str) -> str | None:
setup, expired = self._active_setup(user_id)
if setup is None:
if _LINK_HINT.search(text):
note = "Your previous upload setup expired. " if expired else ""
return f"{note}To process a message range, send /upload first."
return None
if setup.status == TaskStatus.AWAITING_START:
return await self._receive_start(setup, text)
if setup.status == TaskStatus.AWAITING_END:
return await self._receive_end(setup, text)
return f"Upload Task #{setup.id} is ready. Send /confirm to begin or /cancel to discard it."
async def _check_link(self, text: str, *, expected_channel_id: int | None = None):
"""Return (link, channel_id, title) or an error message."""
try:
link = parse_message_link(text)
except LinkError as exc:
return str(exc)
source = self.queue.source
if source is None:
return "Telegram access is not available right now."
try:
channel_id, title = await source.resolve_link(link)
except BotError as exc:
return exc.public
if expected_channel_id is not None and channel_id != expected_channel_id:
return "The end link belongs to a different channel than the start link. Both links must come from the same channel."
if (
expected_channel_id is None
and self.config.upload_restrict_to_source_channel
and channel_id != source.channel_id
):
return (
"This link is not from the configured source channel "
f"({source.channel_title or self.config.source_channel}). "
"Set UPLOAD_RESTRICT_TO_SOURCE_CHANNEL=false to allow other channels the bot can access."
)
try:
accessible = await source.message_accessible(channel_id, link.message_id)
except BotError as exc:
return exc.public
if not accessible:
return f"Message {link.message_id} was not found in that channel (deleted, or not visible to the bot's account)."
return link, channel_id, title
async def _receive_start(self, setup: UploadTask, text: str) -> str:
result = await self._check_link(text)
if isinstance(result, str):
return HtmlReply(f"❌ {esc(result)}\n\nPlease send a valid start link, for example:\n{EXAMPLE}")
link, channel_id, title = result
self.db.update_task(
setup.id,
status=TaskStatus.AWAITING_END,
channel_ref=link.channel_ref,
channel_id=channel_id,
channel_title=title,
start_message_id=link.message_id,
)
return HtmlReply(f"Start message: {link.message_id} ✔\n\n{END_PROMPT}")
async def _receive_end(self, setup: UploadTask, text: str) -> str:
result = await self._check_link(text, expected_channel_id=setup.channel_id)
if isinstance(result, str):
return HtmlReply(
f"❌ {esc(result)}\n\nThe start link (message {setup.start_message_id}) is kept. "
"Please send a valid end link."
)
link, _, _ = result
if link.message_id < setup.start_message_id:
return HtmlReply(
f"❌ The end message ID ({link.message_id}) must be greater than or equal to the start message ID "
f"({setup.start_message_id}). Please send a valid end link."
)
size = link.message_id - setup.start_message_id + 1
if size > self.config.upload_max_range:
return HtmlReply(
f"❌ That range covers {size} messages; the maximum is {self.config.upload_max_range} "
"(UPLOAD_MAX_RANGE). Please send a closer end link."
)
task = self.db.update_task(setup.id, status=TaskStatus.AWAITING_CONFIRMATION, end_message_id=link.message_id)
return self._confirmation(task)
def _channel_label(self, task: UploadTask) -> str:
source = self.queue.source
name = task.channel_title or str(task.channel_id)
if source is not None and task.channel_id == source.channel_id:
return f"{name} (configured source channel)"
return name
def _confirmation(self, task: UploadTask) -> HtmlReply:
return HtmlReply(
"✅ <b>Upload Task Ready</b>\n\n"
f"Start Message ID: {task.start_message_id}\n"
f"End Message ID: {task.end_message_id}\n"
f"Channel: {esc(self._channel_label(task))}\n\n"
"The bot will process eligible videos within this message range, organize them into Google Drive "
"topic folders, and clean up local temporary files after successful uploads.\n\n"
"Send /confirm to begin or /cancel to cancel this task."
)
# -- task commands ----------------------------------------------------------------------
async def cmd_upload(self, arg: str, user_id: int) -> str:
setup, expired = self._active_setup(user_id)
if setup is not None:
hint = {
TaskStatus.AWAITING_START: "Send the start link to continue.",
TaskStatus.AWAITING_END: f"Start message {setup.start_message_id} is set; send the end link to continue.",
TaskStatus.AWAITING_CONFIRMATION: "Send /confirm to begin.",
}[setup.status]
return (
f"⚠️ You already have an upload setup in progress (Task #{setup.id}, {_STEP[setup.status]}).\n"
f"{hint}\nOr send /cancel to discard it and start again."
)
self.db.create_task(user_id)
prefix = "Your previous unfinished setup expired.\n\n" if expired else ""
return HtmlReply(prefix + START_PROMPT)
async def cmd_confirm(self, arg: str, user_id: int) -> str:
setup, expired = self._active_setup(user_id)
if setup is None:
note = "Your previous setup expired. " if expired else ""
return f"{note}There is no upload task waiting for confirmation. Start one with /upload."
if setup.status != TaskStatus.AWAITING_CONFIRMATION:
return (
"Both the start and end links must be validated before /confirm. "
f"Task #{setup.id} is {_STEP[setup.status]}."
)
task = self.queue.submit_task(setup.id)
ahead = [t for t in self.db.tasks_with_status(*ACTIVE_TASK_STATUSES) if t.id < task.id]
lines = [f"🚀 <b>Upload Task #{task.id} confirmed</b>", f"Range: {task.range_text}"]
if self.queue.paused_flag:
lines.append("The queue is paused; send /resume to start processing.")
elif ahead:
lines.append("Queued after task(s) " + ", ".join(f"#{t.id}" for t in ahead) + ".")
else:
lines.append("Processing has started. Use /taskstatus to follow it.")
return HtmlReply("\n".join(lines))
async def cmd_cancel(self, arg: str, user_id: int) -> str:
if arg:
if not arg.isdigit():
return "Usage: /cancel or /cancel TASK_ID"
return self._cancel_reply(int(arg), await self.queue.cancel_task(int(arg)))
setup, _ = self._active_setup(user_id)
if setup is not None:
await self.queue.cancel_task(setup.id, reason="setup cancelled")
return f"Upload setup #{setup.id} cancelled. Nothing was downloaded or uploaded."
own = self.db.tasks_with_status(*ACTIVE_TASK_STATUSES, admin_user_id=user_id)
if own:
running = [t for t in own if t.status == TaskStatus.RUNNING]
target = running[0] if running else own[-1]
return self._cancel_reply(target.id, await self.queue.cancel_task(target.id))
others = self.db.tasks_with_status(*ACTIVE_TASK_STATUSES)
if others:
ids = ", ".join(f"#{t.id}" for t in others)
return f"You have nothing to cancel. Active task(s) {ids} belong to another administrator; use /cancel ID to cancel one explicitly."
return "Nothing to cancel."
@staticmethod
def _cancel_reply(task_id: int, outcome: str) -> str:
return {
"unknown": f"Task #{task_id} does not exist.",
"finished": f"Task #{task_id} has already finished.",
"setup_cancelled": f"Upload setup #{task_id} cancelled.",
"cancelled": (
f"🛑 Task #{task_id} cancelled at a safe checkpoint. Completed uploads are kept; "
"the remaining videos were not processed and their temporary files were removed."
),
}[outcome]
async def cmd_taskstatus(self, arg: str, user_id: int) -> str:
if arg:
if not arg.isdigit():
return "Usage: /taskstatus or /taskstatus TASK_ID"
task = self.db.get_task(int(arg))
return self._describe_task(task) if task else f"Task #{arg} does not exist."
parts = []
setup, expired = self._active_setup(user_id)
if setup is not None:
parts.append(
f"Your setup: Task #{setup.id}, {_STEP[setup.status]}"
+ (f" (start {setup.start_message_id})" if setup.start_message_id else "")
)
elif expired:
parts.append("Your previous setup expired.")
active = self.db.tasks_with_status(*ACTIVE_TASK_STATUSES)
parts += [self._describe_task(t) for t in active]
if not active:
recent = self.db.recent_tasks(1)
parts.append("No active upload tasks." + (f"\n\nLast task:\n{self._describe_task(recent[0])}" if recent else ""))
return HtmlReply("\n\n".join(parts))
def _describe_task(self, task: UploadTask) -> HtmlReply:
lines = [
f"<b>Task #{task.id}</b> — {esc(task.status.value)}",
f"Range: {task.range_text}",
f"Channel: {esc(self._channel_label(task)) if task.channel_id else '-'}",
f"Requested by: {task.admin_user_id}",
]
if task.status not in SETUP_TASK_STATUSES:
lines += [
f"Eligible videos: {task.eligible_count} · Uploaded: {task.uploaded_count} · "
f"Duplicates: {task.duplicate_count} · Failed: {task.failed_count}",
f"Non-video skipped: {task.non_video_count} · Other skipped: {task.skipped_count} · "
f"Missing: {task.missing_count}",
f"Temporary files deleted: {task.deleted_files_count}",
]
if task.current_message_id and task.status == TaskStatus.RUNNING:
lines.append(f"Current message: {task.current_message_id}")
board = self.queue.reporter.board if self.queue.reporter else None
if board and board.task and board.task.get("id") == task.id:
lines.append(f"Video: {esc(board.title)} · {esc(board.status)}")
if board.upload:
lines.append(f"Upload: {esc(board.upload.describe())}")
elif board.download:
lines.append(f"Download: {esc(board.download.describe())}")
for label, value in (("Created", task.created_at), ("Started", task.started_at), ("Finished", task.completed_at)):
if value:
lines.append(f"{label}: {value}")
if task.last_error:
lines.append(f"Note: {esc(_short(task.last_error))}")
return HtmlReply("\n".join(lines))
async def cmd_uploadhistory(self, arg: str, user_id: int) -> str:
tasks = self.db.recent_tasks(10)
if not tasks:
return "No upload tasks yet. Start one with /upload."
lines = ["<b>Recent upload tasks</b>"]
for t in tasks:
lines.append(
f"#{t.id} {t.range_text} {esc(t.status.value)} — uploaded {t.uploaded_count}, "
f"duplicates {t.duplicate_count}, failed {t.failed_count}, non-video {t.non_video_count} "
f"(by {t.admin_user_id})"
)
return HtmlReply("\n".join(lines))
# -- general commands ---------------------------------------------------------------------
async def cmd_start(self, arg: str, user_id: int) -> str:
return (
"Telegram -> Google Drive video organizer.\n"
"Videos are uploaded into Drive folders named after the caption Topic, "
"using the caption Title as the file name.\n\n"
"Process a message range: /upload, send the start and end links, then /confirm.\n"
"Process the whole channel: /sync then /run.\n\n" + HELP
)
async def cmd_help(self, arg: str, user_id: int) -> str:
return HELP
async def cmd_status(self, arg: str, user_id: int) -> str:
counts = self.db.counts()
lines = [f"Queue: {self.queue.status}"]
if self.queue.halt_reason:
lines.append(f"Halted because: {self.queue.halt_reason}")
active = self.db.tasks_with_status(*ACTIVE_TASK_STATUSES)
if active:
lines.append("Upload tasks: " + ", ".join(f"#{t.id} {t.status.value} ({t.range_text})" for t in active))
current = self.queue.current
if current and self.queue.reporter:
lines.append("")
lines.append(re.sub(r"</?b>", "", self.queue.reporter.board.render()))
elif current:
lines.append(f"Current: message {current.message_id} ({current.title or 'untitled'})")
lines.append("")
lines.append(
f"Completed {counts['COMPLETED']} · Pending {counts['PENDING'] + counts['RETRY_WAIT']} · "
f"Failed {counts['FAILED']} · Skipped {counts['SKIPPED']}"
)
lines.append(f"Whole-channel processing (/run): {'on' if self.queue.global_enabled else 'off'}")
lines.append(f"Live channel listener: {'on' if self.queue.listening else 'off'}")
return "\n".join(lines)
async def cmd_progress(self, arg: str, user_id: int) -> str:
c = self.db.counts()
total = sum(c.values())
processed = c["COMPLETED"]
remaining = self.queue.remaining()
considered = total - c["SKIPPED"]
pct = f"{100 * processed / considered:.1f}%" if considered else "n/a"
return (
f"Total discovered: {total}\n"
f"Processed (uploaded & verified): {processed} ({pct})\n"
f"Pending: {c['PENDING']}\n"
f"Waiting to retry: {c['RETRY_WAIT']}\n"
f"In progress: {c['DOWNLOADING'] + c['DOWNLOADED'] + c['UPLOADING']}\n"
f"Failed: {c['FAILED']}\n"
f"Skipped: {c['SKIPPED']}\n"
f"Remaining: {remaining}"
)
async def cmd_run(self, arg: str, user_id: int) -> str:
if self.queue.start():
pending = self.queue.remaining()
hint = "" if pending else "\nThe queue is empty - use /sync to discover channel videos."
return f"Processing started ({pending} item(s) waiting).{hint}"
return f"Processing is already {self.queue.status}."
async def cmd_pause(self, arg: str, user_id: int) -> str:
if self.queue.pause():
return "Pausing: the current transfer stops at the next resumable checkpoint. Use /resume to continue."
return f"Nothing to pause (queue is {self.queue.status})."
async def cmd_resume(self, arg: str, user_id: int) -> str:
if self.queue.resume():
return "Processing resumed."
return f"Nothing to resume (queue is {self.queue.status})."
async def cmd_stop(self, arg: str, user_id: int) -> str:
await self.queue.stop(timeout=60)
return (
"Stopped. Interrupted work was checkpointed; a running upload task was paused. "
"Use /resume to continue upload tasks or /run for whole-channel processing."
)
async def cmd_retry(self, arg: str, user_id: int) -> str:
count = self.queue.retry_failed()
if not count:
return "No failed items to retry."
started = self.queue.start()
return f"Re-queued {count} failed item(s)." + (" Processing started." if started else "")
async def cmd_failed(self, arg: str, user_id: int) -> str:
items = self.db.list_by_state(State.FAILED, limit=20)
if not items:
return "No failed items."
lines = [f"Failed items (showing {len(items)}):"]
for rec in items:
lines.append(f"• #{rec.message_id} {_short(rec.title, 60)} — {_short(rec.last_error)}")
lines.append("Use /retry to re-queue them, /check ID for details.")
return "\n".join(lines)
async def cmd_folders(self, arg: str, user_id: int) -> str:
topics = self.db.topic_counts()
counts = self.db.counts()
lines = [f"Drive folders created/reused: {self.db.folder_count()}", f"Files uploaded: {counts['COMPLETED']}"]
if topics:
lines.append("")
lines += [f"• {topic or self.config.fallback_topic}: {n}" for topic, n in topics]
return "\n".join(lines)
async def cmd_check(self, arg: str, user_id: int) -> str:
if not arg.lstrip("-").isdigit():
return "Usage: /check MESSAGE_ID"
records = self.db.find_by_message_id(int(arg))
if not records:
return f"Message {arg} is not in the database. Use /sync or /upload to discover messages."
return "\n\n".join(self._describe(r) for r in records)
async def cmd_sync(self, arg: str, user_id: int) -> str:
result = await self.queue.sync(rescan=arg.lower() == "all")
if self.queue.running:
self.queue.wake()
return "Sync complete.\n" + result.summary()
def _describe(self, rec: VideoRecord) -> str:
lines = [
f"Message {rec.message_id} — {rec.state}",
f"Title: {rec.title or '-'}",
f"Topic: {rec.topic or '- (uses ' + self.config.fallback_topic + ')'}",
f"Batch: {rec.batch or '-'}",
f"Index: {', '.join(rec.indexes) or '-'}",
f"Extracted by: {rec.extracted_by or '-'}",
f"Download: {rec.download_status or '-'} · Upload: {rec.upload_status or '-'}",
f"Size: {rec.expected_size or rec.local_size or '?'} bytes",
f"Attempts: {rec.attempts}",
]
if rec.task_id:
lines.append(f"Upload task: #{rec.task_id}")
if rec.final_name:
lines.append(f"Drive name: {rec.final_name}")
if rec.drive_file_id:
lines.append(f"Drive file ID: {rec.drive_file_id}")
if rec.next_retry_at and rec.state == State.RETRY_WAIT:
lines.append(f"Next retry: {rec.next_retry_at}")
if rec.skip_reason:
lines.append(f"Skip reason: {_short(rec.skip_reason)}")
if rec.last_error:
lines.append(f"Last error: {_short(rec.last_error)}")
if rec.warnings:
lines.append(f"Warnings: {_short('; '.join(rec.warnings), 300)}")
return "\n".join(lines)
def _parse_mode(text: str) -> str | None:
return "html" if isinstance(text, HtmlReply) else None
class TelethonNotifier:
def __init__(self, client):
self.client = client
async def send(self, chat_id: int, text: str) -> int | None:
message = await self.client.send_message(chat_id, str(text), parse_mode=_parse_mode(text), link_preview=False)
return message.id
async def edit(self, chat_id: int, message_id: int, text: str) -> None:
from telethon import errors
try:
await self.client.edit_message(chat_id, message_id, str(text), parse_mode=_parse_mode(text), link_preview=False)
except errors.MessageNotModifiedError:
pass
def register_command_handlers(client, router: CommandRouter) -> None:
from telethon import events
@client.on(events.NewMessage(incoming=True, func=lambda e: e.is_private))
async def _on_private_message(event) -> None: # pragma: no cover - needs a live Telegram connection
reply = await router.handle(event.sender_id, event.raw_text)
if not reply:
return
if isinstance(reply, HtmlReply):
await event.respond(str(reply), parse_mode="html", link_preview=False)
return
for start in range(0, len(reply), 4000):
await event.respond(reply[start : start + 4000], parse_mode=None, link_preview=False)
caption_parser.py
+108
View File
@@ -0,0 +1,108 @@
"""Extract structured metadata from a Telegram video caption.
Expected format (fields may appear in any order, with extra blank lines, emoji or
bullet decorations, and varying whitespace)::
Index: 001
Index: 002
Title: Class 01 Aman sir.mp4
Topic: Misc• Grammar Practice
Batch: English ( SSC MAINS )
Extracted By: @PUTINxINDIA
Parsing is line based, so the value of one field can never swallow the next field.
If a single line holds several known fields (captions flattened by a forwarder),
the line is split at each known label.
"""
from __future__ import annotations
import re
from models import ParsedCaption
REQUIRED_FIELDS = ("title", "topic")
_LABEL = r"(?P<label>index|title|topic|batch|extracted[ \t_]*by)"
# Optional decoration (emoji, bullets, markdown markers) before the label, then the
# label, optional spaces and an ASCII or full-width colon.
_LINE_FIELD = re.compile(rf"^[^\w@]*{_LABEL}[ \t]*[:：]", re.IGNORECASE)
_INLINE_FIELD = re.compile(rf"(?<=\s)[^\w\s@]*{_LABEL}[ \t]*[:：]", re.IGNORECASE)
_GENERIC_FIELD = re.compile(r"^(?P<key>[^:：]{1,40}?)[ \t]*[:：][ \t]*(?P<value>.*)$")
_LABEL_KEYS = {
"index": "index",
"title": "title",
"topic": "topic",
"batch": "batch",
"extractedby": "extracted_by",
}
_DISPLAY = {"title": "Title", "topic": "Topic", "batch": "Batch", "extracted_by": "Extracted By", "index": "Index"}
def _field_key(label: str) -> str:
return _LABEL_KEYS[re.sub(r"[\s_]+", "", label.lower())]
def _split_fields(line: str) -> list[tuple[str, str]] | None:
"""Return ``[(field, value), ...]`` if the line starts with a known label."""
first = _LINE_FIELD.match(line)
if not first:
return None
marks = [(first.end(), first.group("label"))]
starts = [first.start("label")]
for inline in _INLINE_FIELD.finditer(line, first.end()):
starts.append(inline.start())
marks.append((inline.end(), inline.group("label")))
result = []
for i, (value_start, label) in enumerate(marks):
is_last = i + 1 == len(marks)
value = line[value_start : len(line) if is_last else starts[i + 1]].strip()
if not is_last:
value = value.rstrip(" \t|,;").strip()
result.append((_field_key(label), value))
return result
def parse_caption(text: str | None) -> ParsedCaption:
raw = text or ""
parsed = ParsedCaption(raw=raw)
if not raw.strip():
parsed.warnings.append("caption is empty")
parsed.missing = list(REQUIRED_FIELDS)
return parsed
for line in raw.replace("\r\n", "\n").replace("\r", "\n").split("\n"):
line = line.strip()
if not line:
continue
fields = _split_fields(line)
if fields is None:
generic = _GENERIC_FIELD.match(line)
if generic and re.search(r"\w", generic.group("key")):
parsed.extra.setdefault(generic.group("key").strip(), generic.group("value").strip())
continue
for name, value in fields:
_assign(parsed, name, value)
parsed.missing = [name for name in REQUIRED_FIELDS if not getattr(parsed, name)]
for name in parsed.missing:
parsed.warnings.append(f"{_DISPLAY[name]} field is missing")
if not parsed.batch:
parsed.warnings.append("Batch field is missing")
return parsed
def _assign(parsed: ParsedCaption, name: str, value: str) -> None:
if name == "index":
if value:
parsed.indexes.append(value)
else:
parsed.warnings.append("empty Index field ignored")
return
if not value:
parsed.warnings.append(f"{_DISPLAY[name]} field is empty")
return
current = getattr(parsed, name)
if current is None:
setattr(parsed, name, value)
elif current != value:
parsed.warnings.append(f"duplicate {_DISPLAY[name]} field {value!r} ignored; kept {current!r}")
config.py
+438
View File
@@ -0,0 +1,438 @@
"""Configuration loading and validation.
Every setting comes from environment variables, optionally loaded from a ``.env``
file next to this module. No secret is hard-coded anywhere in the project.
"""
from __future__ import annotations
import os
import re
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Mapping
PROJECT_DIR = Path(__file__).resolve().parent
UPLOAD_CHUNK_ALIGN = 256 * 1024 # Google requires resumable chunks in multiples of 256 KiB
_TRUE = {"1", "true", "yes", "on"}
_FALSE = {"0", "false", "no", "off"}
_SIZE_RE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([KMGT]?I?B?)?\s*$", re.IGNORECASE)
_SIZE_UNITS = {"": 1, "B": 1, "K": 1024, "M": 1024**2, "G": 1024**3, "T": 1024**4}
class ConfigError(Exception):
def __init__(self, problems: list[str]):
self.problems = problems
super().__init__("Configuration problems:\n" + "\n".join(f" - {p}" for p in problems))
class _Reader:
"""Typed accessors that collect every problem instead of stopping at the first."""
def __init__(self, env: Mapping[str, str]):
self.env = env
self.problems: list[str] = []
def str(self, name: str, default: str | None = None) -> str | None:
raw = self.env.get(name)
if raw is None:
return default
raw = raw.strip()
return raw if raw else default
def int(self, name: str, default: int | None = None, *, minimum: int | None = None) -> int | None:
raw = self.str(name)
if raw is None:
return default
try:
value = int(raw)
except ValueError:
self.problems.append(f"{name} must be an integer (got {raw!r})")
return default
if minimum is not None and value < minimum:
self.problems.append(f"{name} must be >= {minimum} (got {value})")
return default
return value
def float(self, name: str, default: float, *, minimum: float = 0.0) -> float:
raw = self.str(name)
if raw is None:
return default
try:
value = float(raw)
except ValueError:
self.problems.append(f"{name} must be a number (got {raw!r})")
return default
if value < minimum:
self.problems.append(f"{name} must be >= {minimum} (got {value})")
return default
return value
def bool(self, name: str, default: bool) -> bool:
raw = self.str(name)
if raw is None:
return default
low = raw.lower()
if low in _TRUE:
return True
if low in _FALSE:
return False
self.problems.append(f"{name} must be true or false (got {raw!r})")
return default
def choice(self, name: str, default: str, choices: tuple[str, ...]) -> str:
raw = self.str(name, default)
value = raw.lower()
if value not in choices:
self.problems.append(f"{name} must be one of {', '.join(choices)} (got {raw!r})")
return default
return value
def int_list(self, name: str) -> tuple[int, ...]:
raw = self.str(name)
if raw is None:
return ()
values = []
for part in raw.split(","):
part = part.strip()
if not part:
continue
try:
values.append(int(part))
except ValueError:
self.problems.append(f"{name} must be a comma-separated list of numeric IDs (bad value {part!r})")
return tuple(values)
def text_set(self, name: str) -> frozenset[str]:
raw = self.str(name)
if raw is None:
return frozenset()
return frozenset(" ".join(p.split()).casefold() for p in raw.split(",") if p.strip())
def path(self, name: str, default: str) -> Path:
raw = self.str(name, default)
path = Path(raw).expanduser()
return path if path.is_absolute() else PROJECT_DIR / path
def size(self, name: str, default: int) -> int:
raw = self.str(name)
if raw is None:
return default
match = _SIZE_RE.match(raw)
if not match:
self.problems.append(f"{name} must be a size such as 512MB or 2GB (got {raw!r})")
return default
unit = (match.group(2) or "").upper().replace("I", "").replace("B", "")
return int(float(match.group(1)) * _SIZE_UNITS[unit or ""])
def datetime(self, name: str) -> datetime | None:
raw = self.str(name)
if raw is None:
return None
try:
value = datetime.fromisoformat(raw)
except ValueError:
self.problems.append(f"{name} must be an ISO date such as 2024-01-31 (got {raw!r})")
return None
return value if value.tzinfo else value.replace(tzinfo=timezone.utc)
@dataclass(frozen=True)
class Config:
# Telegram
telegram_mode: str
api_id: int | None
api_hash: str | None
bot_token: str | None
session_string: str | None
admin_ids: frozenset[int]
notify_chat_ids: tuple[int, ...]
source_channel: str | None
process_mode: str
start_message_id: int
start_date: datetime | None
include_document_videos: bool
filter_topics: frozenset[str]
filter_batches: frozenset[str]
bot_scan_max_empty: int
bot_scan_end_id: int | None
flood_sleep_threshold: int
session_dir: Path
# Google Drive
google_auth_mode: str
google_client_secrets_file: Path
google_token_file: Path
google_service_account_file: Path | None
google_service_account_json: str | None
google_oauth_token_json: str | None
drive_scope: str
drive_root_folder_id: str | None
shared_drive_id: str | None
share_mode: str
upload_chunk_size: int
# Organisation and naming
folder_structure: str
fallback_topic: str
fallback_batch: str
duplicate_name_policy: str
reuse_identical_files: bool
duplicate_media_policy: str
missing_title_policy: str
extension_mismatch_policy: str
strict_filenames: bool
sequence_prefix_enabled: bool
sequence_prefix_format: str
# Storage
data_dir: Path
database_url: str | None
allow_ephemeral_database: bool
is_heroku: bool
download_dir: Path
database_path: Path
log_dir: Path
max_temp_storage: int
min_free_disk_space: int
keep_local_files: bool
keep_failed_downloads: bool
temp_cleanup_interval: float
compute_sha256: bool
# Processing
max_attempts: int
retry_base_delay: float
retry_max_delay: float
http_retries: int
processing_delay: float
strict_stop: bool
auto_start: bool
instance_lock_wait: float
upload_setup_timeout: float
upload_max_range: int
upload_restrict_to_source_channel: bool
progress_interval: float
log_level: str
@classmethod
def from_env(cls, env: Mapping[str, str] | None = None, *, load_dotenv_file: bool = True) -> "Config":
if env is None:
if load_dotenv_file:
from dotenv import load_dotenv
load_dotenv(PROJECT_DIR / ".env")
env = os.environ
r = _Reader(env)
admin_ids = frozenset(r.int_list("ADMIN_USER_IDS"))
notify = r.int_list("NOTIFY_CHAT_IDS") or tuple(sorted(admin_ids))
sa_file = r.str("GOOGLE_SERVICE_ACCOUNT_FILE")
config = cls(
telegram_mode=r.choice("TELEGRAM_MODE", "user", ("user", "bot")),
api_id=r.int("TELEGRAM_API_ID"),
api_hash=r.str("TELEGRAM_API_HASH"),
bot_token=r.str("TELEGRAM_BOT_TOKEN"),
session_string=r.str("TELEGRAM_SESSION_STRING"),
admin_ids=admin_ids,
notify_chat_ids=notify,
source_channel=r.str("SOURCE_CHANNEL_ID"),
process_mode=r.choice("PROCESS_MODE", "both", ("existing", "new", "both")),
start_message_id=r.int("START_MESSAGE_ID", 1, minimum=1),
start_date=r.datetime("START_DATE"),
include_document_videos=r.bool("INCLUDE_DOCUMENT_VIDEOS", True),
filter_topics=r.text_set("FILTER_TOPICS"),
filter_batches=r.text_set("FILTER_BATCHES"),
bot_scan_max_empty=r.int("BOT_SCAN_MAX_EMPTY", 200, minimum=1),
bot_scan_end_id=r.int("BOT_SCAN_END_ID", None, minimum=1),
flood_sleep_threshold=r.int("FLOOD_SLEEP_THRESHOLD", 60, minimum=0),
session_dir=r.path("SESSION_DIR", "data/sessions"),
google_auth_mode=r.choice("GOOGLE_AUTH_MODE", "oauth", ("oauth", "service_account")),
google_client_secrets_file=r.path("GOOGLE_CLIENT_SECRETS_FILE", "credentials/client_secret.json"),
google_token_file=r.path("GOOGLE_TOKEN_FILE", "credentials/token.json"),
google_service_account_file=r.path("GOOGLE_SERVICE_ACCOUNT_FILE", sa_file) if sa_file else None,
google_service_account_json=r.str("GOOGLE_SERVICE_ACCOUNT_JSON"),
google_oauth_token_json=r.str("GOOGLE_OAUTH_TOKEN_JSON"),
drive_scope=r.choice("GOOGLE_DRIVE_SCOPE", "drive.file", ("drive.file", "drive")),
drive_root_folder_id=r.str("GOOGLE_DRIVE_ROOT_FOLDER_ID"),
shared_drive_id=r.str("DRIVE_SHARED_DRIVE_ID"),
share_mode=r.choice("DRIVE_SHARE_MODE", "none", ("none", "anyone_with_link")),
upload_chunk_size=r.size("UPLOAD_CHUNK_SIZE", 16 * 1024 * 1024),
folder_structure="batch_topic" if r.bool("BATCH_SUBFOLDERS_ENABLED", False) else "topic",
fallback_topic=r.str("FALLBACK_TOPIC_FOLDER", "Uncategorized"),
fallback_batch=r.str("FALLBACK_BATCH_FOLDER", "Unknown Batch"),
duplicate_name_policy=r.choice(
"DUPLICATE_POLICY", "suffix_number", ("suffix_number", "suffix_message_id", "skip", "fail")
),
reuse_identical_files=r.bool("REUSE_IDENTICAL_FILES", True),
duplicate_media_policy=r.choice("DUPLICATE_MEDIA_POLICY", "skip", ("skip", "upload")),
missing_title_policy=r.choice("MISSING_TITLE_POLICY", "fallback", ("fallback", "fail")),
extension_mismatch_policy=r.choice("EXTENSION_MISMATCH_POLICY", "keep", ("keep", "append")),
strict_filenames=r.bool("STRICT_FILENAMES", False),
sequence_prefix_enabled=r.bool("SEQUENCE_PREFIX_ENABLED", False),
# not stripped: trailing spaces are meaningful here
sequence_prefix_format=env.get("SEQUENCE_PREFIX_FORMAT") or "{seq:03d} - ",
data_dir=r.path("DATA_DIR", "data"),
database_url=_normalize_database_url(r.str("DATABASE_URL")),
allow_ephemeral_database=r.bool("ALLOW_EPHEMERAL_DATABASE", False),
is_heroku=bool(env.get("DYNO")),
download_dir=r.path("DOWNLOAD_DIR", "downloads"),
database_path=r.path("DATABASE_PATH", "data/bot.db"),
log_dir=r.path("LOG_DIR", "logs"),
max_temp_storage=r.size("MAX_TEMP_STORAGE", 0),
min_free_disk_space=r.size("MIN_FREE_DISK_SPACE", 1024**3),
keep_local_files=r.bool("KEEP_LOCAL_FILES", False),
keep_failed_downloads=r.bool("KEEP_FAILED_DOWNLOADS", False),
temp_cleanup_interval=r.float("TEMP_CLEANUP_INTERVAL", 3600.0, minimum=0.0),
compute_sha256=r.bool("COMPUTE_SHA256", True),
max_attempts=r.int("MAX_RETRIES", 4, minimum=0) + 1,
retry_base_delay=r.float("RETRY_DELAY", 30.0),
retry_max_delay=r.float("RETRY_MAX_DELAY", 1800.0),
http_retries=r.int("HTTP_RETRIES", 5, minimum=0),
processing_delay=r.float("PROCESSING_DELAY", 2.0),
strict_stop=r.bool("STRICT_STOP", False),
auto_start=r.bool("AUTO_START", False),
instance_lock_wait=r.float("INSTANCE_LOCK_WAIT", 90.0 if env.get("DYNO") else 0.0),
upload_setup_timeout=r.float("UPLOAD_SETUP_TIMEOUT", 30.0, minimum=1.0),
upload_max_range=r.int("UPLOAD_MAX_RANGE", 5000, minimum=1),
upload_restrict_to_source_channel=r.bool("UPLOAD_RESTRICT_TO_SOURCE_CHANNEL", True),
progress_interval=r.float("PROGRESS_INTERVAL", 10.0, minimum=1.0),
log_level=r.choice("LOG_LEVEL", "info", ("debug", "info", "warning", "error")).upper(),
)
if config.upload_chunk_size <= 0 or config.upload_chunk_size % UPLOAD_CHUNK_ALIGN:
r.problems.append("UPLOAD_CHUNK_SIZE must be a positive multiple of 256KB (e.g. 8MB, 16MB, 32MB)")
try:
config.sequence_prefix_format.format(seq=1)
except (KeyError, IndexError, ValueError):
r.problems.append("SEQUENCE_PREFIX_FORMAT must be a format string using {seq}, e.g. '{seq:03d} - '")
resolved = config.download_dir.resolve()
if resolved in {Path("/").resolve(), Path.home().resolve(), PROJECT_DIR.resolve()} or len(resolved.parts) < 3:
r.problems.append(
f"DOWNLOAD_DIR={config.download_dir} is not a dedicated directory; use e.g. 'downloads' or '/tmp/telegram_drive_bot'"
)
if r.problems:
raise ConfigError(r.problems)
return config
# -- section validation -------------------------------------------------------
def telegram_problems(self) -> list[str]:
problems = []
if not self.api_id or not self.api_hash:
problems.append(
"TELEGRAM_API_ID and TELEGRAM_API_HASH are required (create them at https://my.telegram.org -> API development tools)"
)
if self.telegram_mode == "bot" and not self.bot_token:
problems.append("TELEGRAM_BOT_TOKEN is required when TELEGRAM_MODE=bot (get one from @BotFather)")
if self.telegram_mode == "user" and self.is_heroku and not self.session_string:
problems.append(
"TELEGRAM_SESSION_STRING is required for TELEGRAM_MODE=user on Heroku (the filesystem is ephemeral, "
"so a session file would be lost). Generate it locally with `python main.py login --session-string`."
)
if self.bot_token and not self.admin_ids:
problems.append("ADMIN_USER_IDS must list at least one Telegram user ID allowed to control the bot")
if not self.source_channel:
problems.append("SOURCE_CHANNEL_ID is required (@username, https://t.me/... link, or -100... numeric ID)")
return problems
def drive_problems(self) -> list[str]:
problems = []
if self.google_auth_mode == "oauth":
if self.google_oauth_token_json:
problems += _json_problems(self.google_oauth_token_json, "GOOGLE_OAUTH_TOKEN_JSON", ("refresh_token", "client_id", "client_secret"))
elif self.is_heroku:
problems.append(
"GOOGLE_OAUTH_TOKEN_JSON is required on Heroku when GOOGLE_AUTH_MODE=oauth "
"(create it locally with `python main.py drive-auth`, see README)"
)
elif not self.google_token_file.exists() and not self.google_client_secrets_file.exists():
problems.append(
f"Google OAuth client file not found at {self.google_client_secrets_file.name!r} "
"(set GOOGLE_CLIENT_SECRETS_FILE or GOOGLE_OAUTH_TOKEN_JSON; see README)"
)
else:
if self.google_service_account_json:
problems += _json_problems(
self.google_service_account_json, "GOOGLE_SERVICE_ACCOUNT_JSON", ("client_email", "private_key")
)
elif not self.google_service_account_file or not self.google_service_account_file.exists():
problems.append(
"GOOGLE_AUTH_MODE=service_account needs GOOGLE_SERVICE_ACCOUNT_JSON (Heroku) "
"or GOOGLE_SERVICE_ACCOUNT_FILE pointing to the JSON key"
)
return problems
def database_problems(self) -> list[str]:
if self.is_heroku and not self.database_url and not self.allow_ephemeral_database:
return [
"DATABASE_URL is not set. On Heroku the local filesystem is wiped on every restart, so SQLite would "
"lose all progress and cause re-uploads. Attach Heroku Postgres (`heroku addons:create "
"heroku-postgresql:essential-0`) or set ALLOW_EPHEMERAL_DATABASE=true to accept volatile state."
]
return []
def validate(self, *sections: str) -> None:
problems: list[str] = []
if "telegram" in sections:
problems += self.telegram_problems()
if "drive" in sections:
problems += self.drive_problems()
if "database" in sections:
problems += self.database_problems()
if problems:
raise ConfigError(problems)
def secrets(self) -> list[str]:
"""Values that must never appear in logs."""
values = [self.api_hash, self.bot_token, self.session_string, self.database_url]
for blob in (self.google_oauth_token_json, self.google_service_account_json):
try:
data = load_json_secret(blob) if blob else {}
except ValueError:
data = {}
values += [data.get(k) for k in ("refresh_token", "client_secret", "private_key", "token")]
values.append(blob)
return [v for v in values if v and len(v) >= 8]
def load_json_secret(value: str) -> dict:
"""Parse a JSON credential given either as raw JSON or base64-encoded JSON.
Base64 is accepted because some dashboards mangle multi-line values; raw JSON
(including multi-line) works fine through `heroku config:set VAR="$(cat file)"`.
"""
import base64
import binascii
import json
text = value.strip()
if not text.startswith("{"):
try:
text = base64.b64decode(text, validate=False).decode("utf-8").strip()
except (binascii.Error, UnicodeDecodeError) as exc:
raise ValueError("value is neither JSON nor base64-encoded JSON") from exc
try:
data = json.loads(text)
except json.JSONDecodeError as exc:
raise ValueError(f"invalid JSON ({exc.msg} at line {exc.lineno})") from exc
if not isinstance(data, dict):
raise ValueError("expected a JSON object")
return data
def _json_problems(value: str, name: str, required_keys: tuple[str, ...]) -> list[str]:
try:
data = load_json_secret(value)
except ValueError as exc:
return [f"{name} could not be parsed: {exc}"]
missing = [k for k in required_keys if not data.get(k)]
return [f"{name} is missing keys: {', '.join(missing)}"] if missing else []
def _normalize_database_url(url: str | None) -> str | None:
"""Heroku provides postgres://...; enforce TLS for remote hosts."""
if not url:
return None
if url.startswith("postgres://"):
url = "postgresql://" + url[len("postgres://"):]
host_part = url.split("@")[-1]
remote = not (host_part.startswith(("localhost", "127.0.0.1", "/")) or "host=/" in url)
if remote and "sslmode=" not in url:
url += ("&" if "?" in url else "?") + "sslmode=require"
return url
database.py
+752
View File
@@ -0,0 +1,752 @@
"""Persistent state: video records, state transitions, folder cache, cursors.
Two backends share one SQL dialect subset:
* PostgreSQL (``DATABASE_URL``) — durable production storage, e.g. Heroku Postgres.
* SQLite (``DATABASE_PATH``) — local development. On Heroku the filesystem is
ephemeral, so SQLite there would lose all history on every dyno restart; the
configuration layer refuses that unless explicitly overridden.
SQL is written with ``?`` placeholders and translated to ``%s`` for psycopg.
"""
from __future__ import annotations
import json
import logging
import sqlite3
import threading
import time
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Iterator
from models import (
IN_PROGRESS_STATES,
MediaInfo,
ParsedCaption,
SourceMessage,
State,
TaskStatus,
UploadTask,
VideoRecord,
can_transition,
)
log = logging.getLogger(__name__)
# Each entry is one schema version; {id_pk} differs between backends.
MIGRATIONS: list[list[str]] = [
[
"""CREATE TABLE IF NOT EXISTS videos (
channel_id BIGINT NOT NULL,
message_id BIGINT NOT NULL,
state TEXT NOT NULL,
caption TEXT,
indexes TEXT NOT NULL DEFAULT '[]',
title TEXT,
topic TEXT,
batch TEXT,
extracted_by TEXT,
warnings TEXT NOT NULL DEFAULT '[]',
media_kind TEXT,
media_id BIGINT,
mime_type TEXT,
original_filename TEXT,
expected_size BIGINT,
message_date TEXT,
sequence_no BIGINT,
local_path TEXT,
local_size BIGINT,
sha256 TEXT,
md5 TEXT,
final_name TEXT,
drive_folder_id TEXT,
drive_file_id TEXT,
upload_session_uri TEXT,
upload_session_created_at TEXT,
download_status TEXT,
upload_status TEXT,
attempts INTEGER NOT NULL DEFAULT 0,
last_error TEXT,
error_kind TEXT,
next_retry_at TEXT,
skip_reason TEXT,
discovered_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
download_started_at TEXT,
downloaded_at TEXT,
upload_started_at TEXT,
completed_at TEXT,
PRIMARY KEY (channel_id, message_id)
)""",
"CREATE INDEX IF NOT EXISTS idx_videos_state ON videos(state)",
"CREATE INDEX IF NOT EXISTS idx_videos_media ON videos(media_id)",
"""CREATE TABLE IF NOT EXISTS folders (
parent_id TEXT NOT NULL,
name_key TEXT NOT NULL,
name TEXT NOT NULL,
folder_id TEXT NOT NULL,
created_at TEXT NOT NULL,
PRIMARY KEY (parent_id, name_key)
)""",
"""CREATE TABLE IF NOT EXISTS transitions (
id {id_pk},
channel_id BIGINT NOT NULL,
message_id BIGINT NOT NULL,
from_state TEXT,
to_state TEXT NOT NULL,
at TEXT NOT NULL,
note TEXT
)""",
"CREATE INDEX IF NOT EXISTS idx_transitions_msg ON transitions(channel_id, message_id)",
],
# v2: interactive /upload tasks
[
"ALTER TABLE videos ADD COLUMN task_id BIGINT",
"CREATE INDEX IF NOT EXISTS idx_videos_task ON videos(task_id)",
"""CREATE TABLE IF NOT EXISTS upload_tasks (
id {id_pk},
admin_user_id BIGINT NOT NULL,
status TEXT NOT NULL,
channel_ref TEXT,
channel_id BIGINT,
channel_title TEXT,
start_message_id BIGINT,
end_message_id BIGINT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
confirmed_at TEXT,
started_at TEXT,
scanned_at TEXT,
completed_at TEXT,
current_message_id BIGINT,
examined_count INTEGER NOT NULL DEFAULT 0,
eligible_count INTEGER NOT NULL DEFAULT 0,
uploaded_count INTEGER NOT NULL DEFAULT 0,
duplicate_count INTEGER NOT NULL DEFAULT 0,
non_video_count INTEGER NOT NULL DEFAULT 0,
skipped_count INTEGER NOT NULL DEFAULT 0,
missing_count INTEGER NOT NULL DEFAULT 0,
failed_count INTEGER NOT NULL DEFAULT 0,
deleted_files_count INTEGER NOT NULL DEFAULT 0,
last_error TEXT
)""",
"CREATE INDEX IF NOT EXISTS idx_tasks_status ON upload_tasks(status)",
"CREATE INDEX IF NOT EXISTS idx_tasks_admin ON upload_tasks(admin_user_id)",
"""CREATE TABLE IF NOT EXISTS task_items (
task_id BIGINT NOT NULL,
channel_id BIGINT NOT NULL,
message_id BIGINT NOT NULL,
result TEXT NOT NULL,
folder TEXT,
updated_at TEXT NOT NULL,
PRIMARY KEY (task_id, channel_id, message_id)
)""",
],
]
_COLUMNS = frozenset(VideoRecord.column_names())
_MUTABLE = _COLUMNS - {"channel_id", "message_id", "state", "discovered_at"}
# download_status / upload_status mirrored from the state for the metadata record.
_STAGE_STATUS = {
State.PENDING: ("pending", "waiting"),
State.DOWNLOADING: ("in_progress", "waiting"),
State.DOWNLOADED: ("done", "waiting"),
State.UPLOADING: ("done", "in_progress"),
State.COMPLETED: ("done", "done"),
State.SKIPPED: ("skipped", "skipped"),
}
_TASK_MUTABLE = frozenset(UploadTask.column_names()) - {"id", "admin_user_id", "created_at"}
_TASK_COUNTERS = frozenset({"deleted_files_count"})
TERMINAL_VIDEO_STATES = ("COMPLETED", "SKIPPED", "FAILED")
_UPSERT_KV = "INSERT INTO kv (key, value) VALUES (?, ?) ON CONFLICT (key) DO UPDATE SET value = excluded.value"
class InvalidTransition(RuntimeError):
pass
class DatabaseError(RuntimeError):
"""Backend-neutral wrapper so callers need not import sqlite3/psycopg."""
class InstanceLockError(RuntimeError):
pass
def utcnow() -> datetime:
return datetime.now(timezone.utc)
def iso(dt: datetime) -> str:
return dt.astimezone(timezone.utc).isoformat(timespec="seconds")
def retry_time(seconds: float) -> str:
return iso(utcnow() + timedelta(seconds=seconds))
# -- backends --------------------------------------------------------------------------
class _SQLiteBackend:
name = "sqlite"
id_pk = "INTEGER PRIMARY KEY AUTOINCREMENT"
for_update = ""
errors: tuple[type[BaseException], ...] = (sqlite3.Error,)
def __init__(self, path: Path):
path.parent.mkdir(parents=True, exist_ok=True)
self.path = path
self.conn = sqlite3.connect(str(path), isolation_level=None, check_same_thread=False, timeout=30)
self.conn.row_factory = sqlite3.Row
self.conn.execute("PRAGMA journal_mode=WAL")
self.conn.execute("PRAGMA synchronous=FULL")
def execute(self, sql: str, params: tuple = ()):
return self.conn.execute(sql, params)
@contextmanager
def transaction(self):
self.conn.execute("BEGIN IMMEDIATE")
try:
yield
except BaseException:
self.conn.execute("ROLLBACK")
raise
else:
self.conn.execute("COMMIT")
def close(self) -> None:
self.conn.close()
def describe(self) -> str:
return f"SQLite ({self.path.name})"
def insert_returning_id(self, sql: str, params: tuple) -> int:
return self.conn.execute(sql, params).lastrowid
class _PostgresBackend:
name = "postgres"
id_pk = "BIGSERIAL PRIMARY KEY"
for_update = " FOR UPDATE"
def __init__(self, url: str):
import psycopg
from psycopg.rows import dict_row
self._psycopg = psycopg
self._dict_row = dict_row
self.url = url
self.errors = (psycopg.Error,)
self.conn = self._connect()
def _connect(self):
return self._psycopg.connect(self.url, autocommit=True, row_factory=self._dict_row, connect_timeout=15)
def _ensure_connected(self) -> None:
if self.conn.closed or self.conn.broken:
log.warning("PostgreSQL connection lost; reconnecting")
self.conn = self._connect()
def execute(self, sql: str, params: tuple = ()):
self._ensure_connected()
return self.conn.execute(sql.replace("?", "%s"), params)
@contextmanager
def transaction(self):
self._ensure_connected()
with self.conn.transaction():
yield
def close(self) -> None:
self.conn.close()
def describe(self) -> str:
return "PostgreSQL"
def insert_returning_id(self, sql: str, params: tuple) -> int:
return self.execute(sql + " RETURNING id", params).fetchone()["id"]
# -- database --------------------------------------------------------------------------
class Database:
def __init__(self, *, path: Path | str | None = None, url: str | None = None):
if url:
self._backend: Any = _PostgresBackend(url)
elif path:
self._backend = _SQLiteBackend(Path(path))
else:
raise ValueError("either a SQLite path or a PostgreSQL URL is required")
self.backend_name = self._backend.name
self.errors = self._backend.errors
self._lock = threading.RLock()
self._migrate()
@classmethod
def from_config(cls, config) -> "Database":
return cls(url=config.database_url) if config.database_url else cls(path=config.database_path)
def describe(self) -> str:
return self._backend.describe()
def close(self) -> None:
with self._lock:
self._backend.close()
def _migrate(self) -> None:
with self._lock, self._backend.transaction():
self._backend.execute("CREATE TABLE IF NOT EXISTS kv (key TEXT PRIMARY KEY, value TEXT)")
row = self._backend.execute("SELECT value FROM kv WHERE key='schema_version'").fetchone()
version = int(row["value"]) if row else 0
for number, statements in enumerate(MIGRATIONS, start=1):
if number <= version:
continue
for statement in statements:
self._backend.execute(statement.format(id_pk=self._backend.id_pk))
self._backend.execute(_UPSERT_KV, ("schema_version", str(number)))
log.info("Database schema migrated to version %d", number)
@contextmanager
def _tx(self) -> Iterator[Any]:
with self._lock, self._backend.transaction():
yield self._backend
def _query(self, sql: str, params: tuple = ()) -> list[Any]:
with self._lock:
return self._backend.execute(sql, params).fetchall()
# -- videos ---------------------------------------------------------------------
def insert_video(
self,
msg: SourceMessage,
parsed: ParsedCaption,
*,
state: State = State.PENDING,
skip_reason: str | None = None,
sequence_no: int | None = None,
task_id: int | None = None,
) -> bool:
"""Insert a newly discovered video. Returns False if it was already known."""
media: MediaInfo = msg.media # type: ignore[assignment]
now = iso(utcnow())
dl, ul = _STAGE_STATUS[state]
with self._tx() as db:
cur = db.execute(
"""INSERT INTO videos (
channel_id, message_id, state, caption, indexes, title, topic, batch, extracted_by,
warnings, media_kind, media_id, mime_type, original_filename, expected_size,
message_date, sequence_no, skip_reason, download_status, upload_status,
discovered_at, updated_at, task_id)
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
ON CONFLICT (channel_id, message_id) DO NOTHING""",
(
msg.channel_id,
msg.message_id,
state.value,
msg.caption,
json.dumps(parsed.indexes, ensure_ascii=False),
parsed.title,
parsed.topic,
parsed.batch,
parsed.extracted_by,
json.dumps(parsed.warnings, ensure_ascii=False),
media.kind,
media.media_id,
media.mime_type,
media.file_name,
media.size,
iso(msg.date) if msg.date else None,
sequence_no,
skip_reason,
dl,
ul,
now,
now,
task_id,
),
)
inserted = cur.rowcount == 1
if inserted:
db.execute(
"INSERT INTO transitions (channel_id, message_id, from_state, to_state, at, note) VALUES (?,?,?,?,?,?)",
(msg.channel_id, msg.message_id, None, state.value, now, skip_reason or "discovered"),
)
return inserted
def exists(self, channel_id: int, message_id: int) -> bool:
return bool(self._query("SELECT 1 FROM videos WHERE channel_id=? AND message_id=?", (channel_id, message_id)))
def get(self, key: tuple[int, int]) -> VideoRecord | None:
rows = self._query("SELECT * FROM videos WHERE channel_id=? AND message_id=?", tuple(key))
return VideoRecord.from_row(rows[0]) if rows else None
def find_by_message_id(self, message_id: int) -> list[VideoRecord]:
rows = self._query("SELECT * FROM videos WHERE message_id=? ORDER BY channel_id", (message_id,))
return [VideoRecord.from_row(r) for r in rows]
def find_by_media_id(self, media_id: int) -> list[VideoRecord]:
rows = self._query("SELECT * FROM videos WHERE media_id=? ORDER BY message_id", (media_id,))
return [VideoRecord.from_row(r) for r in rows]
def transition(self, key: tuple[int, int], new_state: State, *, note: str | None = None, **fields) -> VideoRecord:
"""Atomically move an item to ``new_state`` (validated) and update fields."""
bad = set(fields) - _MUTABLE
if bad:
raise ValueError(f"unknown or immutable fields: {sorted(bad)}")
now = iso(utcnow())
with self._tx() as db:
row = db.execute(
f"SELECT state FROM videos WHERE channel_id=? AND message_id=?{self._backend.for_update}", tuple(key)
).fetchone()
if row is None:
raise KeyError(f"unknown video {key}")
current = State(row["state"])
if not can_transition(current, new_state):
raise InvalidTransition(f"{key}: {current} -> {new_state} is not allowed")
fields.setdefault("updated_at", now)
if new_state in _STAGE_STATUS:
fields.setdefault("download_status", _STAGE_STATUS[new_state][0])
fields.setdefault("upload_status", _STAGE_STATUS[new_state][1])
else:
failed_stage = "upload_status" if current in (State.DOWNLOADED, State.UPLOADING) else "download_status"
fields.setdefault(failed_stage, "failed" if new_state == State.FAILED else "retry_wait")
self._update(db, key, state=new_state.value, **fields)
db.execute(
"INSERT INTO transitions (channel_id, message_id, from_state, to_state, at, note) VALUES (?,?,?,?,?,?)",
(key[0], key[1], current.value, new_state.value, now, note),
)
return self.get(key) # type: ignore[return-value]
def update_fields(self, key: tuple[int, int], **fields) -> None:
bad = set(fields) - _MUTABLE
if bad:
raise ValueError(f"unknown or immutable fields: {sorted(bad)}")
fields.setdefault("updated_at", iso(utcnow()))
with self._tx() as db:
self._update(db, key, **fields)
@staticmethod
def _update(db, key: tuple[int, int], **fields) -> None:
for name in VideoRecord.JSON_FIELDS:
if name in fields and not isinstance(fields[name], str):
fields[name] = json.dumps(fields[name], ensure_ascii=False)
assignments = ", ".join(f"{name}=?" for name in fields)
db.execute(
f"UPDATE videos SET {assignments} WHERE channel_id=? AND message_id=?",
(*[v.value if isinstance(v, State) else v for v in fields.values()], *key),
)
def history(self, key: tuple[int, int]) -> list[tuple[str | None, str, str | None]]:
rows = self._query(
"SELECT from_state, to_state, note FROM transitions WHERE channel_id=? AND message_id=? ORDER BY id", tuple(key)
)
return [(r["from_state"], r["to_state"], r["note"]) for r in rows]
def next_item(self, now: datetime | None = None, *, task_id: int | None = None) -> VideoRecord | None:
"""Next item to work on: unfinished in-progress items first, then oldest message.
With ``task_id`` only that task's items are considered. Without it, items
belonging to a paused or queued /upload task are left alone.
"""
if task_id is not None:
scope, params = "task_id = ?", (task_id,)
else:
scope = (
"(task_id IS NULL OR task_id NOT IN "
"(SELECT id FROM upload_tasks WHERE status IN ('PAUSED','QUEUED')))"
)
params = ()
rows = self._query(
f"""SELECT * FROM videos
WHERE {scope}
AND (state IN ('UPLOADING','DOWNLOADED','DOWNLOADING','PENDING')
OR (state='RETRY_WAIT' AND next_retry_at <= ?))
ORDER BY CASE state WHEN 'UPLOADING' THEN 0 WHEN 'DOWNLOADED' THEN 1
WHEN 'DOWNLOADING' THEN 2 ELSE 3 END,
message_id
LIMIT 1""",
(*params, iso(now or utcnow())),
)
return VideoRecord.from_row(rows[0]) if rows else None
def earliest_retry_at(self) -> datetime | None:
rows = self._query("SELECT MIN(next_retry_at) AS t FROM videos WHERE state='RETRY_WAIT'")
value = rows[0]["t"] if rows else None
return datetime.fromisoformat(value) if value else None
def recover_stale(self) -> dict[str, int]:
"""Called once at startup, after the single-instance lock is held.
DOWNLOADING items go back to PENDING; a partial file, if it survived, is resumed.
UPLOADING items are left as they are: the queue reconciles them against Drive
(and resumes the saved upload session) before doing anything else.
"""
recovered = {}
for state in sorted(IN_PROGRESS_STATES):
rows = self._query("SELECT channel_id, message_id FROM videos WHERE state=?", (state.value,))
recovered[state.value] = len(rows)
if state == State.DOWNLOADING:
for r in rows:
self.transition((r["channel_id"], r["message_id"]), State.PENDING, note="recovered after restart")
if any(recovered.values()):
log.warning("Recovered interrupted items after restart: %s", recovered)
return recovered
def reset_failed(self, *, has_local) -> int:
"""Move FAILED items back into the queue with a fresh attempt budget."""
rows = self._query("SELECT * FROM videos WHERE state='FAILED' ORDER BY message_id")
for row in rows:
rec = VideoRecord.from_row(row)
target = State.DOWNLOADED if has_local(rec) else State.PENDING
self.transition(rec.key, target, note="manual retry", attempts=0, next_retry_at=None, error_kind=None)
return len(rows)
def counts(self) -> dict[str, int]:
counts = {s.value: 0 for s in State}
for r in self._query("SELECT state, COUNT(*) AS n FROM videos GROUP BY state"):
counts[r["state"]] = int(r["n"])
return counts
def list_by_state(self, state: State, limit: int = 20) -> list[VideoRecord]:
rows = self._query("SELECT * FROM videos WHERE state=? ORDER BY message_id LIMIT ?", (state.value, limit))
return [VideoRecord.from_row(r) for r in rows]
def topic_counts(self, limit: int = 30) -> list[tuple[str, int]]:
rows = self._query(
"""SELECT COALESCE(topic, '') AS topic, COUNT(*) AS n FROM videos WHERE state='COMPLETED'
GROUP BY COALESCE(topic, '') ORDER BY n DESC, topic LIMIT ?""",
(limit,),
)
return [(r["topic"], int(r["n"])) for r in rows]
def videos_for_task(self, task_id: int) -> list[VideoRecord]:
rows = self._query("SELECT * FROM videos WHERE task_id=? ORDER BY message_id", (task_id,))
return [VideoRecord.from_row(r) for r in rows]
def task_unfinished_count(self, task_id: int) -> int:
rows = self._query(
"SELECT COUNT(*) AS n FROM videos WHERE task_id=? AND state NOT IN ('COMPLETED','SKIPPED','FAILED')",
(task_id,),
)
return int(rows[0]["n"])
# -- upload tasks -----------------------------------------------------------------
def create_task(self, admin_user_id: int) -> UploadTask:
now = iso(utcnow())
with self._tx() as db:
task_id = db.insert_returning_id(
"INSERT INTO upload_tasks (admin_user_id, status, created_at, updated_at) VALUES (?,?,?,?)",
(admin_user_id, TaskStatus.AWAITING_START.value, now, now),
)
return self.get_task(task_id) # type: ignore[return-value]
def get_task(self, task_id: int) -> UploadTask | None:
rows = self._query("SELECT * FROM upload_tasks WHERE id=?", (task_id,))
return UploadTask.from_row(rows[0]) if rows else None
def update_task(self, task_id: int, **fields) -> UploadTask:
bad = set(fields) - _TASK_MUTABLE
if bad:
raise ValueError(f"unknown or immutable task fields: {sorted(bad)}")
fields.setdefault("updated_at", iso(utcnow()))
values = [v.value if isinstance(v, TaskStatus) else v for v in fields.values()]
assignments = ", ".join(f"{name}=?" for name in fields)
with self._tx() as db:
db.execute(f"UPDATE upload_tasks SET {assignments} WHERE id=?", (*values, task_id))
return self.get_task(task_id) # type: ignore[return-value]
def add_task_counter(self, task_id: int, column: str, amount: int) -> None:
if column not in _TASK_COUNTERS:
raise ValueError(column)
with self._tx() as db:
db.execute(f"UPDATE upload_tasks SET {column} = {column} + ? WHERE id=?", (amount, task_id))
def tasks_with_status(self, *statuses: TaskStatus, admin_user_id: int | None = None) -> list[UploadTask]:
marks = ",".join("?" for _ in statuses)
sql = f"SELECT * FROM upload_tasks WHERE status IN ({marks})"
params: tuple = tuple(s.value for s in statuses)
if admin_user_id is not None:
sql += " AND admin_user_id=?"
params += (admin_user_id,)
rows = self._query(sql + " ORDER BY id", params)
return [UploadTask.from_row(r) for r in rows]
def next_runnable_task(self) -> UploadTask | None:
"""The task the worker should run: an already RUNNING one, else the oldest QUEUED."""
rows = self._query(
"""SELECT * FROM upload_tasks WHERE status IN ('RUNNING','QUEUED')
ORDER BY CASE status WHEN 'RUNNING' THEN 0 ELSE 1 END, id LIMIT 1"""
)
return UploadTask.from_row(rows[0]) if rows else None
def recent_tasks(self, limit: int = 10) -> list[UploadTask]:
rows = self._query(
"SELECT * FROM upload_tasks WHERE status NOT IN ('AWAITING_START','AWAITING_END','AWAITING_CONFIRMATION') "
"ORDER BY id DESC LIMIT ?",
(limit,),
)
return [UploadTask.from_row(r) for r in rows]
def set_task_status(self, old: TaskStatus, new: TaskStatus) -> int:
with self._tx() as db:
cur = db.execute(
"UPDATE upload_tasks SET status=?, updated_at=? WHERE status=?", (new.value, iso(utcnow()), old.value)
)
return cur.rowcount
def add_task_item(self, task_id: int, channel_id: int, message_id: int, result: str) -> None:
"""Record a message examined by a task; an existing result is kept (rescans are idempotent)."""
with self._tx() as db:
db.execute(
"""INSERT INTO task_items (task_id, channel_id, message_id, result, updated_at) VALUES (?,?,?,?,?)
ON CONFLICT (task_id, channel_id, message_id) DO NOTHING""",
(task_id, channel_id, message_id, result, iso(utcnow())),
)
def set_task_item_result(
self, task_id: int, channel_id: int, message_id: int, result: str, *, folder: str | None = None
) -> None:
with self._tx() as db:
db.execute(
"""UPDATE task_items SET result=?, folder=COALESCE(?, folder), updated_at=?
WHERE task_id=? AND channel_id=? AND message_id=?""",
(result, folder, iso(utcnow()), task_id, channel_id, message_id),
)
def task_item_results(self, task_id: int) -> dict[str, int]:
rows = self._query("SELECT result, COUNT(*) AS n FROM task_items WHERE task_id=? GROUP BY result", (task_id,))
return {r["result"]: int(r["n"]) for r in rows}
def task_items(self, task_id: int) -> list[tuple[int, str]]:
rows = self._query("SELECT message_id, result FROM task_items WHERE task_id=? ORDER BY message_id", (task_id,))
return [(int(r["message_id"]), r["result"]) for r in rows]
def task_folders(self, task_id: int) -> list[str]:
rows = self._query(
"SELECT DISTINCT folder FROM task_items WHERE task_id=? AND folder IS NOT NULL ORDER BY folder", (task_id,)
)
return [r["folder"] for r in rows]
# -- folders --------------------------------------------------------------------
def get_folder(self, parent_id: str, name_key: str) -> str | None:
rows = self._query("SELECT folder_id FROM folders WHERE parent_id=? AND name_key=?", (parent_id, name_key))
return rows[0]["folder_id"] if rows else None
def put_folder(self, parent_id: str, name_key: str, name: str, folder_id: str) -> None:
with self._tx() as db:
db.execute(
"""INSERT INTO folders (parent_id, name_key, name, folder_id, created_at) VALUES (?,?,?,?,?)
ON CONFLICT (parent_id, name_key) DO UPDATE
SET name = excluded.name, folder_id = excluded.folder_id, created_at = excluded.created_at""",
(parent_id, name_key, name, folder_id, iso(utcnow())),
)
def delete_folder(self, parent_id: str, name_key: str) -> None:
with self._tx() as db:
db.execute("DELETE FROM folders WHERE parent_id=? AND name_key=?", (parent_id, name_key))
def folder_count(self) -> int:
return int(self._query("SELECT COUNT(*) AS n FROM folders")[0]["n"])
# -- key/value --------------------------------------------------------------------
def get_kv(self, key: str, default: str | None = None) -> str | None:
rows = self._query("SELECT value FROM kv WHERE key=?", (key,))
return rows[0]["value"] if rows else default
def set_kv(self, key: str, value: str) -> None:
with self._tx() as db:
db.execute(_UPSERT_KV, (key, value))
def get_int(self, key: str, default: int = 0) -> int:
value = self.get_kv(key)
return int(value) if value is not None else default
def raise_int(self, key: str, value: int) -> None:
"""Store ``value`` only if it is larger than the current one (cursor semantics)."""
with self._tx() as db:
row = db.execute(f"SELECT value FROM kv WHERE key=?{self._backend.for_update}", (key,)).fetchone()
if row is None or int(row["value"]) < value:
db.execute(_UPSERT_KV, (key, str(value)))
def next_sequence(self) -> int:
with self._tx() as db:
row = db.execute(f"SELECT value FROM kv WHERE key='sequence_counter'{self._backend.for_update}").fetchone()
value = (int(row["value"]) if row else 0) + 1
db.execute(_UPSERT_KV, ("sequence_counter", str(value)))
return value
# -- single-instance lock ----------------------------------------------------------------
_ADVISORY_LOCK_ID = 0x7464_6272 # arbitrary constant: "tdbr"
class InstanceLock:
"""Guarantee that only one worker processes the queue.
PostgreSQL: a session-level advisory lock held on a dedicated connection, so it
also protects against two Heroku dynos (e.g. worker scaled to 2, or an old dyno
still shutting down during a deploy). SQLite: an exclusive ``flock`` on a lock
file, which only protects processes on the same machine.
"""
def __init__(self, config, *, wait_seconds: float = 0, poll: float = 5.0):
self._config = config
self._wait = wait_seconds
self._poll = poll
self._conn = None
self._fh = None
def acquire(self) -> None:
deadline = time.monotonic() + self._wait
while True:
if self._try_acquire():
return
if time.monotonic() >= deadline:
raise InstanceLockError(
"Another worker is already processing this queue. Run exactly one worker "
"(Heroku: `heroku ps:scale worker=1`)."
)
log.warning("Another worker holds the queue lock; waiting...")
time.sleep(self._poll)
def _try_acquire(self) -> bool:
if self._config.database_url:
import psycopg
if self._conn is None or self._conn.closed:
self._conn = psycopg.connect(self._config.database_url, autocommit=True, connect_timeout=15)
return bool(self._conn.execute("SELECT pg_try_advisory_lock(%s)", (_ADVISORY_LOCK_ID,)).fetchone()[0])
import fcntl
path = Path(str(self._config.database_path) + ".lock")
path.parent.mkdir(parents=True, exist_ok=True)
if self._fh is None:
self._fh = open(path, "w")
try:
fcntl.flock(self._fh, fcntl.LOCK_EX | fcntl.LOCK_NB)
return True
except BlockingIOError:
return False
def release(self) -> None:
if self._conn is not None:
try:
self._conn.close() # closing the session releases the advisory lock
finally:
self._conn = None
if self._fh is not None:
self._fh.close()
self._fh = None
downloader.py
+246
View File
@@ -0,0 +1,246 @@
"""Stream Telegram media to disk with resume, disk-space checks and hashing."""
from __future__ import annotations
import asyncio
import errno
import hashlib
import logging
import os
import re
import shutil
from pathlib import Path
from typing import AsyncIterator, Callable, Protocol
from errors import DiskSpaceError, Interrupted, TransientError
from models import LocalFile, State, VideoRecord
log = logging.getLogger(__name__)
# Partial downloads are resumed from a 1 MiB boundary: Telegram's upload.getFile
# requires offsets aligned to the requested chunk size (512 KiB).
RESUME_ALIGNMENT = 1024 * 1024
HASH_BLOCK = 4 * 1024 * 1024
# Only files with exactly this shape are ever deleted by cleanup routines.
TEMP_NAME = re.compile(r"^(\d+)_(\d+)\.(part|video)$")
class MediaSource(Protocol):
def iter_media_chunks(
self, message_id: int, *, offset: int, expected_size: int | None, channel_id: int | None = None
) -> AsyncIterator[bytes]: ...
class _Hashes:
def __init__(self, sha256: bool):
self.md5 = hashlib.md5()
self.sha256 = hashlib.sha256() if sha256 else None
def update(self, data: bytes) -> None:
self.md5.update(data)
if self.sha256:
self.sha256.update(data)
def digest(self) -> tuple[str, str | None]:
return self.md5.hexdigest(), self.sha256.hexdigest() if self.sha256 else None
def hash_file(path: Path, *, sha256: bool, limit: int | None = None, hashes: _Hashes | None = None) -> _Hashes:
hashes = hashes or _Hashes(sha256)
remaining = limit
with open(path, "rb") as fh:
while remaining is None or remaining > 0:
block = fh.read(HASH_BLOCK if remaining is None else min(HASH_BLOCK, remaining))
if not block:
break
hashes.update(block)
if remaining is not None:
remaining -= len(block)
return hashes
class Downloader:
def __init__(
self,
source: MediaSource,
download_dir: Path,
*,
min_free_bytes: int = 0,
max_temp_bytes: int = 0,
compute_sha256: bool = True,
disk_usage: Callable[[str], tuple] = shutil.disk_usage,
):
self.source = source
self.download_dir = download_dir
self.min_free_bytes = min_free_bytes
self.max_temp_bytes = max_temp_bytes
self.compute_sha256 = compute_sha256
self._disk_usage = disk_usage
# Optional callback that frees space by removing abandoned temp files;
# returns the number of files deleted. Set by the queue manager.
self.reclaim: Callable[[], int] | None = None
download_dir.mkdir(parents=True, exist_ok=True)
self._root = download_dir.resolve()
# Local files are named by message identity, never by caption text, so
# captions cannot cause path problems.
def complete_path(self, rec: VideoRecord) -> Path:
return self.download_dir / f"{abs(rec.channel_id)}_{rec.message_id}.video"
def partial_path(self, rec: VideoRecord) -> Path:
return self.download_dir / f"{abs(rec.channel_id)}_{rec.message_id}.part"
def has_complete_local(self, rec: VideoRecord) -> bool:
path = self.complete_path(rec)
if not path.is_file():
return False
expected = rec.expected_size or rec.local_size
return expected is not None and path.stat().st_size == expected
async def local_file(self, rec: VideoRecord) -> LocalFile:
"""Describe an already complete download, reusing stored hashes if valid."""
path = self.complete_path(rec)
size = path.stat().st_size
if rec.md5 and rec.local_size == size:
return LocalFile(path, size, rec.md5, rec.sha256)
hashes = await asyncio.to_thread(hash_file, path, sha256=self.compute_sha256)
md5, sha = hashes.digest()
return LocalFile(path, size, md5, sha)
def check_space(self, needed: int, *, exclude: Path | None = None) -> None:
free = self._disk_usage(str(self.download_dir))[2]
if free - needed < self.min_free_bytes:
raise DiskSpaceError(
f"not enough disk space: need {needed} bytes plus {self.min_free_bytes} reserve, {free} free",
public="Not enough free disk space for the next video; free space and /resume.",
)
if self.max_temp_bytes:
used = sum(
p.stat().st_size for p in self.download_dir.iterdir() if p.is_file() and p != exclude
)
if used + needed > self.max_temp_bytes:
raise DiskSpaceError(
f"MAX_TEMP_STORAGE exceeded: {used} used + {needed} needed > {self.max_temp_bytes}",
public="Temporary storage limit (MAX_TEMP_STORAGE) reached; /retry or clean up downloads.",
)
async def download(
self,
rec: VideoRecord,
*,
on_progress: Callable[[int, int | None], None],
should_stop: Callable[[], bool],
) -> LocalFile:
if self.has_complete_local(rec):
return await self.local_file(rec)
part = self.partial_path(rec)
final = self.complete_path(rec)
expected = rec.expected_size
offset = part.stat().st_size if part.exists() else 0
offset -= offset % RESUME_ALIGNMENT
if expected is not None and offset > expected:
offset = 0
needed = (expected - offset) if expected else 0
try:
self.check_space(needed, exclude=part)
except DiskSpaceError:
if not (self.reclaim and self.reclaim()):
raise
self.check_space(needed, exclude=part)
hashes = _Hashes(self.compute_sha256)
if offset:
log.info("Resuming download of message %s at byte %d", rec.message_id, offset)
await asyncio.to_thread(hash_file, part, sha256=self.compute_sha256, limit=offset, hashes=hashes)
written = offset
chunks = self.source.iter_media_chunks(
rec.message_id, offset=offset, expected_size=expected, channel_id=rec.channel_id
)
try:
with open(part, "r+b" if part.exists() else "wb") as fh:
fh.truncate(offset)
fh.seek(offset)
on_progress(written, expected)
async for chunk in chunks:
fh.write(chunk)
hashes.update(chunk)
written += len(chunk)
on_progress(written, expected)
if expected is not None and written > expected:
break
if should_stop():
fh.flush()
raise Interrupted(f"download paused at byte {written}")
fh.flush()
os.fsync(fh.fileno())
except OSError as exc:
if exc.errno == errno.ENOSPC:
raise DiskSpaceError("disk full while downloading", public="Local disk is full; free space and /resume.") from exc
raise TransientError(f"download failed: {exc}", public="Download failed (I/O or network error)") from exc
finally:
await chunks.aclose()
size = part.stat().st_size
if expected is not None and size != expected:
if size > expected:
part.unlink(missing_ok=True)
raise TransientError(
f"incomplete download: {size} of {expected} bytes", public="Download incomplete or corrupted; will retry"
)
if size == 0:
part.unlink(missing_ok=True)
raise TransientError("downloaded file is empty", public="Downloaded file is empty; will retry")
os.replace(part, final)
md5, sha = hashes.digest()
return LocalFile(final, size, md5, sha)
def cleanup(self, rec: VideoRecord) -> int:
"""Delete this video's temporary files. Returns how many files were removed."""
return sum(self.safe_unlink(path) for path in (self.complete_path(rec), self.partial_path(rec)))
def safe_unlink(self, path: Path) -> bool:
"""Delete ``path`` only if it is a regular temp file directly inside the download dir."""
try:
if path.is_symlink() or path.parent.resolve() != self._root or not TEMP_NAME.match(path.name):
log.error("Refusing to delete %s: not an application temp file", path.name)
return False
if not path.exists():
return False
path.unlink()
log.info("Deleted temporary file %s", path.name)
return True
except OSError as exc:
log.warning("Could not delete temporary file %s: %s", path.name, exc)
return False
def cleanup_orphans(self, db, *, keep_failed: bool = False) -> int:
"""Remove temp files that no unfinished item needs any more.
A file is deleted only when it matches the application's temp-name pattern,
lives directly in the download directory (symlinks are never followed), and
its database record is missing or finished (COMPLETED / SKIPPED, or FAILED
unless ``keep_failed``). Leftover ``.part`` files of fully downloaded items
are removed too. Anything else in the directory is left untouched.
"""
deleted = 0
try:
entries = list(os.scandir(self._root))
except FileNotFoundError:
return 0
for entry in entries:
match = TEMP_NAME.match(entry.name)
if not match or entry.is_symlink() or not entry.is_file(follow_symlinks=False):
continue
channel_id, message_id, kind = -int(match.group(1)), int(match.group(2)), match.group(3)
rec = db.get((channel_id, message_id))
if rec is None or rec.state in (State.COMPLETED, State.SKIPPED):
reason = "record finished" if rec else "no database record"
elif rec.state == State.FAILED and not keep_failed:
reason = "item failed (can be downloaded again)"
elif kind == "part" and rec.state in (State.DOWNLOADED, State.UPLOADING):
reason = "stale partial of a completed download"
else:
continue
if self.safe_unlink(Path(entry.path)):
deleted += 1
log.info("Orphan cleanup: removed %s (%s)", entry.name, reason)
return deleted
drive_client.py
+620
View File
@@ -0,0 +1,620 @@
"""Google Drive access: authentication, folders, identity lookup, resumable uploads.
Metadata calls use the official ``google-api-python-client``. File uploads use
Google's documented resumable-upload protocol directly over
``google.auth.transport.requests.AuthorizedSession`` so that the session URI can be
persisted in SQLite and an interrupted upload can continue after a process restart.
https://developers.google.com/drive/api/guides/manage-uploads#resumable
"""
from __future__ import annotations
import json
import logging
import os
import random
import re
import threading
import time
from pathlib import Path
from typing import Any, Callable
from urllib.parse import quote
from errors import (
AuthError,
BotError,
DriveQuotaError,
Interrupted,
PermanentError,
TransientError,
UploadSessionExpired,
)
from models import DriveFile
from naming import normalize_key
log = logging.getLogger(__name__)
FOLDER_MIME = "application/vnd.google-apps.folder"
SCOPES = {
"drive.file": "https://www.googleapis.com/auth/drive.file",
"drive": "https://www.googleapis.com/auth/drive",
}
UPLOAD_URL = "https://www.googleapis.com/upload/drive/v3/files"
CHUNK_ALIGN = 256 * 1024
FILE_FIELDS = "id,name,size,md5Checksum,parents,mimeType,trashed,appProperties"
# appProperties key holding "<channel_id>:<message_id>"; used to recognise our own
# uploads after a crash, so a lost response never leads to a duplicate upload.
SOURCE_PROPERTY = "tg_source"
_RATE_LIMIT_REASONS = {"userRateLimitExceeded", "rateLimitExceeded", "sharingRateLimitExceeded"}
_QUOTA_REASONS = {"storageQuotaExceeded", "teamDriveFileLimitExceeded", "activeItemCreationLimitExceeded"}
def _network_errors() -> tuple[type[BaseException], ...]:
errors: list[type[BaseException]] = [OSError, TimeoutError] # requests' exceptions subclass OSError
try:
import httplib2
errors.append(httplib2.HttpLib2Error)
except ImportError: # pragma: no cover
pass
try:
from google.auth.exceptions import TransportError
errors.append(TransportError)
except ImportError: # pragma: no cover
pass
return tuple(errors)
NETWORK_ERRORS = _network_errors()
def escape_query(value: str) -> str:
"""Escape a literal for the Drive ``q`` search syntax."""
return value.replace("\\", "\\\\").replace("'", "\\'")
def backoff_delay(attempt: int, base: float, cap: float = 64.0) -> float:
return min(base * (2 ** (attempt - 1)), cap) * random.uniform(0.8, 1.2)
def parse_error_payload(payload: Any) -> tuple[str | None, str | None]:
"""Extract (reason, message) from a Google JSON error body."""
if isinstance(payload, (bytes, bytearray)):
payload = payload.decode("utf-8", "replace")
if isinstance(payload, str):
try:
payload = json.loads(payload)
except ValueError:
return None, payload[:200] or None
if not isinstance(payload, dict):
return None, None
error = payload.get("error")
if not isinstance(error, dict):
return None, str(error) if error else None
reason = None
details = error.get("errors")
if isinstance(details, list) and details and isinstance(details[0], dict):
reason = details[0].get("reason")
return reason or error.get("status"), error.get("message")
def classify_http_error(
status: int, reason: str | None = None, message: str | None = None, retry_after: float | None = None
) -> BotError:
text = f"Google Drive API error {status}" + (f" ({reason})" if reason else "") + (f": {message}" if message else "")
if status == 401:
return AuthError(text, public="Google authorization expired or was revoked; run `python main.py drive-auth`.")
if status == 403 and reason in _QUOTA_REASONS:
return DriveQuotaError(
text,
public="Google Drive storage quota exceeded (service accounts have no My Drive quota; use a Shared Drive).",
)
if status == 403 and reason == "dailyLimitExceeded":
return TransientError(text, retry_after=retry_after or 3600, rate_limited=True, public="Drive daily API limit reached")
if status == 429 or (status == 403 and reason in _RATE_LIMIT_REASONS):
return TransientError(text, retry_after=retry_after, rate_limited=True, public="Google Drive rate limit")
if status in (408, 500, 502, 503, 504) or status >= 500:
return TransientError(text, retry_after=retry_after, public=f"Google Drive temporary error {status}")
return PermanentError(text, public=f"Google Drive rejected the request ({status}{', ' + reason if reason else ''})")
# -- authentication --------------------------------------------------------------------
def load_credentials(config):
"""Build Google credentials without needing files on disk when env vars are set.
Service account: GOOGLE_SERVICE_ACCOUNT_JSON (parsed in memory) or a key file.
OAuth (personal Drive): GOOGLE_OAUTH_TOKEN_JSON (authorized-user JSON containing
the refresh token) or the local token file written by `main.py drive-auth`.
"""
from config import load_json_secret
scopes = [SCOPES[config.drive_scope]]
if config.google_auth_mode == "service_account":
from google.oauth2 import service_account
if config.google_service_account_json:
info = load_json_secret(config.google_service_account_json)
return service_account.Credentials.from_service_account_info(info, scopes=scopes)
return service_account.Credentials.from_service_account_file(str(config.google_service_account_file), scopes=scopes)
from google.auth.exceptions import RefreshError
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from_env = bool(config.google_oauth_token_json)
if from_env:
creds = Credentials.from_authorized_user_info(load_json_secret(config.google_oauth_token_json), scopes)
else:
token_file: Path = config.google_token_file
if not token_file.exists():
raise AuthError(
"Google token not found", public="Google Drive is not authorized yet; run `python main.py drive-auth`."
)
creds = Credentials.from_authorized_user_file(str(token_file), scopes)
if not creds.valid:
if not creds.refresh_token:
raise AuthError("Google token has no refresh token", public="Run `python main.py drive-auth` again.")
try:
creds.refresh(Request())
except RefreshError as exc:
raise AuthError(
f"Google token refresh failed: {exc}",
public="Google authorization expired or was revoked; run `python main.py drive-auth` and update the token.",
) from exc
if not from_env:
save_token(creds, config.google_token_file)
return creds
def save_token(creds, token_file: Path) -> None:
token_file.parent.mkdir(parents=True, exist_ok=True)
fd = os.open(token_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
with os.fdopen(fd, "w") as fh:
fh.write(creds.to_json())
def run_oauth_flow(config, *, open_browser: bool = True) -> Path:
from google_auth_oauthlib.flow import InstalledAppFlow
flow = InstalledAppFlow.from_client_secrets_file(str(config.google_client_secrets_file), [SCOPES[config.drive_scope]])
creds = flow.run_local_server(port=0, open_browser=open_browser)
save_token(creds, config.google_token_file)
return config.google_token_file
# -- client ------------------------------------------------------------------------------
class DriveClient:
def __init__(
self,
service,
session,
*,
chunk_size: int,
shared_drive_id: str | None = None,
http_retries: int = 5,
base_delay: float = 2.0,
sleep: Callable[[float], None] = time.sleep,
):
self._service = service
self._shared_drive_id = shared_drive_id
self._retries = http_retries
self._base_delay = base_delay
self._sleep = sleep
self._lock = threading.Lock() # httplib2 (used by the discovery client) is not thread-safe
self.uploader = ResumableUploader(
session, chunk_size=chunk_size, retries=http_retries, base_delay=base_delay, sleep=sleep
)
@classmethod
def from_config(cls, config) -> "DriveClient":
from google.auth.transport.requests import AuthorizedSession
from googleapiclient.discovery import build
creds = load_credentials(config)
service = build("drive", "v3", credentials=creds, cache_discovery=False)
return cls(
service,
AuthorizedSession(creds),
chunk_size=config.upload_chunk_size,
shared_drive_id=config.shared_drive_id,
http_retries=config.http_retries,
)
# -- low level ------------------------------------------------------------------------
def _execute(self, build_request: Callable[[], Any], *, retries: int | None = None, not_found_ok: bool = False):
from google.auth.exceptions import RefreshError
from googleapiclient.errors import HttpError
retries = self._retries if retries is None else retries
attempt = 0
while True:
attempt += 1
try:
with self._lock:
return build_request().execute()
except HttpError as exc:
status = int(exc.resp.status)
if status == 404 and not_found_ok:
return None
retry_after = exc.resp.get("retry-after")
error = classify_http_error(
status, *parse_error_payload(exc.content), retry_after=float(retry_after) if retry_after else None
)
except RefreshError as exc:
raise AuthError(
f"Google token refresh failed: {exc}",
public="Google authorization expired or was revoked; run `python main.py drive-auth`.",
) from exc
except NETWORK_ERRORS as exc:
error = TransientError(f"network error talking to Google Drive: {exc}", public="Network error (Drive)")
if not isinstance(error, TransientError) or attempt > retries:
raise error
delay = error.retry_after or backoff_delay(attempt, self._base_delay)
log.warning("Drive request failed (%s); retry %d/%d in %.1fs", error, attempt, retries, delay)
self._sleep(delay)
def _list(self, query: str, *, order_by: str | None = None) -> list[DriveFile]:
files: list[DriveFile] = []
page_token = None
while True:
kwargs: dict[str, Any] = dict(
q=query,
fields=f"nextPageToken, files({FILE_FIELDS})",
pageSize=100,
supportsAllDrives=True,
includeItemsFromAllDrives=True,
)
if order_by:
kwargs["orderBy"] = order_by
if self._shared_drive_id:
kwargs.update(corpora="drive", driveId=self._shared_drive_id)
if page_token:
kwargs["pageToken"] = page_token
response = self._execute(lambda: self._service.files().list(**kwargs))
files.extend(DriveFile.from_api(f) for f in response.get("files", []))
page_token = response.get("nextPageToken")
if not page_token:
return files
# -- queries ------------------------------------------------------------------------
def get_file(self, file_id: str) -> DriveFile | None:
data = self._execute(
lambda: self._service.files().get(fileId=file_id, fields=FILE_FIELDS, supportsAllDrives=True),
not_found_ok=True,
)
return DriveFile.from_api(data) if data else None
def find_folders(self, parent_id: str, name: str) -> list[DriveFile]:
query = (
f"'{escape_query(parent_id)}' in parents and name = '{escape_query(name)}' "
f"and mimeType = '{FOLDER_MIME}' and trashed = false"
)
return self._list(query, order_by="createdTime")
def find_by_name(self, parent_id: str, name: str) -> list[DriveFile]:
query = (
f"'{escape_query(parent_id)}' in parents and name = '{escape_query(name)}' "
f"and mimeType != '{FOLDER_MIME}' and trashed = false"
)
return self._list(query)
def find_by_source_key(self, key: str) -> list[DriveFile]:
query = (
f"appProperties has {{ key='{SOURCE_PROPERTY}' and value='{escape_query(key)}' }} and trashed = false"
)
return self._list(query)
# -- mutations ------------------------------------------------------------------------
def create_folder(self, parent_id: str, name: str) -> DriveFile:
# Single attempt: folder creation is not idempotent. FolderResolver searches
# again before retrying, so a lost response cannot create a duplicate.
body = {"name": name, "mimeType": FOLDER_MIME, "parents": [parent_id]}
data = self._execute(
lambda: self._service.files().create(body=body, fields=FILE_FIELDS, supportsAllDrives=True), retries=0
)
return DriveFile.from_api(data)
def share_with_link(self, file_id: str) -> None:
body = {"type": "anyone", "role": "reader"}
self._execute(
lambda: self._service.permissions().create(fileId=file_id, body=body, fields="id", supportsAllDrives=True)
)
def upload_file(
self,
path: Path,
*,
name: str,
parent_id: str,
mime_type: str,
size: int,
app_properties: dict[str, str],
description: str | None = None,
session_uri: str | None,
on_session: Callable[[str], None],
on_progress: Callable[[int, int], None],
should_stop: Callable[[], bool],
) -> DriveFile:
metadata: dict[str, Any] = {"name": name, "parents": [parent_id], "appProperties": app_properties}
if description:
metadata["description"] = description
data = self.uploader.upload(
path,
metadata=metadata,
mime_type=mime_type,
size=size,
session_uri=session_uri,
on_session=on_session,
on_progress=on_progress,
should_stop=should_stop,
)
return DriveFile.from_api(data)
def check_access(self) -> None:
"""Cheap authenticated call used at startup to validate credentials."""
self._execute(lambda: self._service.about().get(fields="user(displayName)"))
def check_root_folder(self, folder_id: str) -> DriveFile:
folder = self.get_file(folder_id)
if folder is None:
raise PermanentError(
f"Drive root folder {folder_id} not found",
public=(
"DRIVE_ROOT_FOLDER_ID is not visible to these credentials. With GOOGLE_DRIVE_SCOPE=drive.file the "
"bot only sees folders it created: run `python main.py drive-init`, or set GOOGLE_DRIVE_SCOPE=drive "
"and re-run `python main.py drive-auth`. For a service account, share the folder with it."
),
)
if folder.mime_type != FOLDER_MIME or folder.trashed:
raise PermanentError(f"Drive root {folder_id} is not an active folder", public="DRIVE_ROOT_FOLDER_ID is not a usable folder")
return folder
# -- resumable upload ---------------------------------------------------------------------
def parse_range(header: str | None) -> int:
"""'bytes=0-1048575' -> 1048576 bytes persisted; no header -> 0."""
if not header:
return 0
match = re.search(r"(\d+)-(\d+)", header)
return int(match.group(2)) + 1 if match else 0
def _response_error(resp) -> BotError:
try:
payload = resp.json()
except ValueError:
payload = getattr(resp, "text", None)
retry_after = resp.headers.get("Retry-After") if resp.headers else None
try:
retry = float(retry_after) if retry_after else None
except ValueError:
retry = None
return classify_http_error(resp.status_code, *parse_error_payload(payload), retry_after=retry)
class ResumableUploader:
def __init__(
self,
session,
*,
chunk_size: int,
retries: int = 5,
base_delay: float = 2.0,
sleep: Callable[[float], None] = time.sleep,
upload_url: str = UPLOAD_URL,
):
if chunk_size <= 0 or chunk_size % CHUNK_ALIGN:
raise ValueError("chunk_size must be a positive multiple of 256 KiB")
self._session = session
self._chunk_size = chunk_size
self._retries = retries
self._base_delay = base_delay
self._sleep = sleep
self._upload_url = upload_url
def upload(
self,
path: Path,
*,
metadata: dict[str, Any],
mime_type: str,
size: int,
session_uri: str | None,
on_session: Callable[[str], None],
on_progress: Callable[[int, int], None],
should_stop: Callable[[], bool],
) -> dict[str, Any]:
if size <= 0:
raise PermanentError("refusing to upload an empty file", public="Downloaded file is empty")
offset = 0
if session_uri:
status, value = self.query_status(session_uri, size)
if status == "done":
on_progress(size, size)
return value
if status == "expired":
log.info("Saved upload session expired; starting a new one")
session_uri = None
else:
offset = value
log.info("Resuming upload session at byte %d of %d", offset, size)
if not session_uri:
session_uri = self._initiate(metadata, mime_type, size)
on_session(session_uri)
offset = 0
on_progress(offset, size)
failures = 0
with open(path, "rb") as fh:
while True:
if should_stop():
raise Interrupted(f"upload paused at byte {offset} of {size}")
fh.seek(offset)
chunk = fh.read(min(self._chunk_size, size - offset))
if not chunk:
status, value = self.query_status(session_uri, size)
if status == "done":
return value
raise TransientError("upload finished locally but Drive did not finalize the file")
headers = {
"Content-Length": str(len(chunk)),
"Content-Range": f"bytes {offset}-{offset + len(chunk) - 1}/{size}",
}
try:
resp = self._session.request("PUT", session_uri, data=chunk, headers=headers, timeout=(30, 600))
except NETWORK_ERRORS as exc:
error: BotError = TransientError(f"network error during upload: {exc}", public="Network error during upload")
else:
if resp.status_code in (200, 201):
on_progress(size, size)
return resp.json()
if resp.status_code == 308:
offset = parse_range(resp.headers.get("Range"))
failures = 0
on_progress(offset, size)
continue
if resp.status_code in (404, 410):
raise UploadSessionExpired("upload session expired", public="Upload session expired; will restart")
error = _response_error(resp)
if not isinstance(error, TransientError):
raise error
failures += 1
if failures > self._retries:
raise error
delay = getattr(error, "retry_after", None) or backoff_delay(failures, self._base_delay)
log.warning("Upload chunk failed (%s); retry %d/%d in %.1fs", error, failures, self._retries, delay)
self._sleep(delay)
status, value = self.query_status(session_uri, size)
if status == "done":
on_progress(size, size)
return value
if status == "expired":
raise UploadSessionExpired("upload session expired", public="Upload session expired; will restart")
offset = value
def query_status(self, session_uri: str, size: int) -> tuple[str, Any]:
"""Ask Drive how much of a resumable session it has: done / partial / expired."""
attempt = 0
while True:
attempt += 1
try:
resp = self._session.request(
"PUT", session_uri, headers={"Content-Length": "0", "Content-Range": f"bytes */{size}"}, timeout=(30, 120)
)
except NETWORK_ERRORS as exc:
error: BotError = TransientError(f"network error querying upload status: {exc}", public="Network error during upload")
else:
if resp.status_code in (200, 201):
return "done", resp.json()
if resp.status_code == 308:
return "partial", parse_range(resp.headers.get("Range"))
if resp.status_code in (404, 410):
return "expired", None
error = _response_error(resp)
if not isinstance(error, TransientError):
raise error
if attempt > self._retries:
raise error
self._sleep(getattr(error, "retry_after", None) or backoff_delay(attempt, self._base_delay))
def _initiate(self, metadata: dict[str, Any], mime_type: str, size: int) -> str:
url = f"{self._upload_url}?uploadType=resumable&supportsAllDrives=true&fields={quote(FILE_FIELDS)}"
headers = {
"Content-Type": "application/json; charset=UTF-8",
"X-Upload-Content-Type": mime_type,
"X-Upload-Content-Length": str(size),
}
body = json.dumps(metadata).encode("utf-8")
attempt = 0
while True:
attempt += 1
try:
resp = self._session.request("POST", url, data=body, headers=headers, timeout=(30, 120))
except NETWORK_ERRORS as exc:
error: BotError = TransientError(f"network error starting upload: {exc}", public="Network error during upload")
else:
if resp.status_code == 200:
location = resp.headers.get("Location") or resp.headers.get("location")
if location:
return location
error = TransientError("Drive did not return an upload session URI")
else:
error = _response_error(resp)
if not isinstance(error, TransientError):
raise error
if attempt > self._retries:
raise error
self._sleep(getattr(error, "retry_after", None) or backoff_delay(attempt, self._base_delay))
# -- folder resolution --------------------------------------------------------------------
class FolderResolver:
"""Map folder names to Drive folder IDs without ever creating duplicates.
Lookup order: SQLite cache (verified once per process) -> Drive search (oldest
match wins if duplicates already exist) -> create. After a failed create the
search is repeated before trying again, so a create whose response was lost is
found instead of duplicated.
"""
def __init__(self, drive: DriveClient, db, root_id: str, *, attempts: int = 4, sleep=time.sleep):
self._drive = drive
self._db = db
self.root_id = root_id
self._attempts = attempts
self._sleep = sleep
self._verified: set[str] = set()
def resolve(self, names: list[str]) -> str:
parent = self.root_id
for name in names:
parent = self._ensure(parent, name)
return parent
def _ensure(self, parent_id: str, name: str) -> str:
key = normalize_key(name)
cached = self._db.get_folder(parent_id, key)
if cached:
if cached in self._verified:
return cached
folder = self._drive.get_file(cached)
if folder and not folder.trashed and folder.mime_type == FOLDER_MIME:
self._verified.add(cached)
return cached
log.warning("Cached folder %r (%s) is gone; looking it up again", name, cached)
self._db.delete_folder(parent_id, key)
attempt = 0
while True:
attempt += 1
matches = [f for f in self._drive.find_folders(parent_id, name) if normalize_key(f.name) == key]
if matches:
if len(matches) > 1:
log.warning("%d folders named %r exist; reusing the oldest (%s)", len(matches), name, matches[0].id)
return self._remember(parent_id, key, name, matches[0].id)
try:
created = self._drive.create_folder(parent_id, name)
except TransientError:
if attempt >= self._attempts:
raise
self._sleep(backoff_delay(attempt, 2.0))
continue
log.info("Created Drive folder %r (%s)", name, created.id)
return self._remember(parent_id, key, name, created.id)
def _remember(self, parent_id: str, key: str, name: str, folder_id: str) -> str:
self._db.put_folder(parent_id, key, name, folder_id)
self._verified.add(folder_id)
return folder_id
errors.py
+69
View File
@@ -0,0 +1,69 @@
"""Exception hierarchy.
The queue decides what to do with a failed item purely from the exception type:
* ``TransientError`` -> RETRY_WAIT with exponential backoff, bounded by MAX_ATTEMPTS.
* ``PermanentError`` -> FAILED immediately (can still be retried manually with /retry).
* ``QueueHaltError`` -> the item is parked for retry without consuming an attempt and
the whole queue halts (disk full, Drive quota, expired auth).
* ``SkipItem`` -> SKIPPED with a recorded reason (configured duplicate policy).
* ``Interrupted`` -> /pause or /stop reached a resumable checkpoint.
"""
from __future__ import annotations
class BotError(Exception):
"""Base error. ``public`` is a short summary that is safe to show in Telegram."""
def __init__(self, message: str, *, public: str | None = None):
super().__init__(message)
self.public = public or message
class TransientError(BotError):
def __init__(
self,
message: str,
*,
retry_after: float | None = None,
rate_limited: bool = False,
public: str | None = None,
):
super().__init__(message, public=public)
self.retry_after = retry_after
# Rate-limit waits do not consume one of the item's attempts.
self.rate_limited = rate_limited
class UploadSessionExpired(TransientError):
"""The saved resumable-upload session URI is no longer valid."""
class PermanentError(BotError):
"""A failure that will not fix itself by waiting."""
class QueueHaltError(BotError):
"""A failure that would affect every following item as well."""
class DiskSpaceError(QueueHaltError):
pass
class DriveQuotaError(QueueHaltError):
pass
class AuthError(QueueHaltError):
pass
class SkipItem(Exception):
def __init__(self, reason: str):
super().__init__(reason)
self.reason = reason
class Interrupted(Exception):
"""Raised at a resumable checkpoint after /pause or /stop."""
link_parser.py
+112
View File
@@ -0,0 +1,112 @@
"""Parse and validate Telegram message links.
Supported:
https://t.me/c/3909337424/14 private/internal channel link
https://t.me/c/3909337424/7/14 same, with a discussion-thread segment
https://t.me/channelname/14 public channel post
https://t.me/s/channelname/14 public web-preview link
Variants with ``www.``, ``telegram.me``, no scheme, a trailing slash, query
parameters (``?single``) or a fragment are accepted.
For ``/c/<id>/`` links the number is the channel's *internal* ID. The MTProto
"marked" peer ID used by Telethon is ``-(10**12 + id)`` (for 10-digit IDs this
looks like ``-100`` followed by the ID). It is only an identifier: access is
still verified by resolving it through the authenticated client.
"""
from __future__ import annotations
import re
from dataclasses import dataclass
from urllib.parse import urlsplit
_HOSTS = {"t.me", "telegram.me", "telegram.dog"}
_USERNAME = re.compile(r"^[A-Za-z][A-Za-z0-9_]{3,31}$")
# Paths on t.me that are not channel usernames.
_RESERVED = {
"addstickers", "addemoji", "addtheme", "addlist", "boost", "bg", "share", "proxy", "socks", "iv",
"login", "setlanguage", "invoice", "joinchat", "contact", "m", "confirmphone", "giftcode",
}
EXAMPLE = "https://t.me/c/3909337424/14"
class LinkError(ValueError):
"""Invalid link; the message is safe to show to the user."""
def marked_channel_id(internal_id: int) -> int:
return -(10**12 + internal_id)
def internal_channel_id(marked_id: int) -> int:
return -marked_id - 10**12
@dataclass(frozen=True)
class MessageLink:
kind: str # "private" or "public"
channel: str # internal numeric ID (private) or username (public)
message_id: int
url: str
@property
def channel_ref(self) -> str:
"""Reference understood by telegram_client.resolve_channel."""
if self.kind == "private":
return str(marked_channel_id(int(self.channel)))
return self.channel
@property
def same_channel_key(self) -> str:
return self.channel if self.kind == "private" else self.channel.lower()
def parse_message_link(text: str | None) -> MessageLink:
raw = (text or "").strip()
if not raw:
raise LinkError(f"Please send a Telegram message link, for example {EXAMPLE}")
if any(ch.isspace() for ch in raw):
raise LinkError("Send only the link, without any other text.")
url = raw if re.match(r"^[a-z]+://", raw, re.IGNORECASE) else "https://" + raw
try:
parts = urlsplit(url)
except ValueError as exc:
raise LinkError(f"That is not a valid URL. Example: {EXAMPLE}") from exc
host = (parts.hostname or "").lower()
if host.startswith("www."):
host = host[4:]
if parts.scheme.lower() not in ("http", "https") or host not in _HOSTS:
raise LinkError(f"That is not a Telegram (t.me) link. Example: {EXAMPLE}")
segments = [s for s in parts.path.split("/") if s]
if not segments:
raise LinkError(f"The link does not point to a message. Example: {EXAMPLE}")
head = segments[0]
if head.startswith("+") or head.lower() == "joinchat":
raise LinkError("Invite links cannot be used. Open the channel, then copy the link of a specific post.")
if head == "c":
if len(segments) not in (3, 4) or not all(s.isdigit() for s in segments[1:]):
raise LinkError(f"Private-channel links must look like {EXAMPLE} (channel ID, then message ID).")
channel, message = segments[1], segments[-1]
kind = "private"
if int(channel) == 0:
raise LinkError("The channel ID in the link is invalid.")
else:
if head == "s":
segments = segments[1:]
if len(segments) == 1:
raise LinkError("This link points to a channel, not to a message. Copy the link of a specific post.")
if len(segments) not in (2, 3) or not segments[-1].isdigit():
raise LinkError(f"Unsupported Telegram link format. Use a post link such as {EXAMPLE}")
channel, message = segments[0], segments[-1]
if channel.lower() in _RESERVED or not _USERNAME.match(channel):
raise LinkError("Unsupported Telegram link type; send the link of a channel post.")
if len(segments) == 3 and not segments[1].isdigit():
raise LinkError(f"Unsupported Telegram link format. Use a post link such as {EXAMPLE}")
kind = "public"
message_id = int(message)
if message_id <= 0:
raise LinkError("The message ID in the link must be a positive number.")
return MessageLink(kind=kind, channel=channel, message_id=message_id, url=raw)
logging_config.py
+48
View File
@@ -0,0 +1,48 @@
"""Logging setup: stdout (always) plus a rotating file locally, with secret redaction."""
from __future__ import annotations
import logging
import logging.handlers
from pathlib import Path
FORMAT = "%(asctime)s %(levelname)-7s %(name)s: %(message)s"
class RedactSecrets(logging.Filter):
"""Replace known secret values in every log record."""
def __init__(self, secrets: list[str]):
super().__init__()
self._secrets = sorted({s for s in secrets if s}, key=len, reverse=True)
def filter(self, record: logging.LogRecord) -> bool:
if not self._secrets:
return True
message = record.getMessage()
redacted = message
for secret in self._secrets:
redacted = redacted.replace(secret, "***")
if redacted != message:
record.msg, record.args = redacted, None
return True
def setup_logging(level: str = "INFO", *, log_dir: Path | None = None, secrets: list[str] | None = None) -> None:
"""``log_dir=None`` logs to stdout only (Heroku collects stdout via `heroku logs`)."""
root = logging.getLogger()
root.setLevel(level)
for handler in list(root.handlers):
root.removeHandler(handler)
redact = RedactSecrets(secrets or [])
handlers: list[logging.Handler] = [logging.StreamHandler()]
if log_dir is not None:
log_dir.mkdir(parents=True, exist_ok=True)
handlers.append(
logging.handlers.RotatingFileHandler(log_dir / "bot.log", maxBytes=10 * 1024 * 1024, backupCount=5, encoding="utf-8")
)
for handler in handlers:
handler.setFormatter(logging.Formatter(FORMAT))
handler.addFilter(redact)
root.addHandler(handler)
for noisy in ("telethon", "googleapiclient.discovery_cache", "urllib3", "google_auth_httplib2"):
logging.getLogger(noisy).setLevel(logging.WARNING)
main.py
+328
View File
@@ -0,0 +1,328 @@
"""Entry point.
python main.py run # long-running worker (default; used by the Heroku Procfile)
python main.py process # headless: sync, process the queue once, exit
python main.py sync [--rescan]
python main.py status
python main.py check-config
python main.py login [--session-string]
python main.py drive-auth [--no-browser]
python main.py drive-init --name "Telegram Videos"
"""
from __future__ import annotations
import argparse
import asyncio
import logging
import os
import signal
import sys
from dataclasses import dataclass, field
from config import Config, ConfigError
from database import Database, InstanceLock, InstanceLockError
from errors import BotError
from logging_config import setup_logging
log = logging.getLogger("main")
@dataclass
class App:
config: Config
db: Database
lock: InstanceLock
queue: object
source: object
clients: list = field(default_factory=list)
async def close(self) -> None:
for client in self.clients:
try:
await client.disconnect()
except Exception as exc: # pragma: no cover
log.debug("disconnect failed: %s", exc)
self.db.close()
self.lock.release()
async def build_app(config: Config, *, control: bool) -> App:
"""Wire every component. Fails fast with actionable errors."""
from telethon import events, utils
from bot_handlers import CommandRouter, TelethonNotifier, register_command_handlers
from downloader import Downloader
from drive_client import DriveClient, FolderResolver
from progress import LogNotifier, ProgressReporter
from queue_manager import CURSOR_KEY, QueueManager
from telegram_client import ACCESS_HASH_KEY, TelethonSource, build_client, resolve_channel
lock = InstanceLock(config, wait_seconds=config.instance_lock_wait)
await asyncio.to_thread(lock.acquire)
db = Database.from_config(config)
log.info("Database: %s", db.describe())
if config.is_heroku and not config.database_url:
log.warning("ALLOW_EPHEMERAL_DATABASE=true: progress is lost whenever the dyno restarts!")
db.recover_stale()
drive = await asyncio.to_thread(DriveClient.from_config, config)
await asyncio.to_thread(drive.check_access)
root_id = config.drive_root_folder_id or "root"
if config.drive_root_folder_id:
await asyncio.to_thread(drive.check_root_folder, root_id)
else:
log.warning("GOOGLE_DRIVE_ROOT_FOLDER_ID not set; topic folders are created in the top level of My Drive")
log.info("Google Drive credentials validated")
clients = []
if config.telegram_mode == "bot":
source_client = build_client(config, "bot")
await source_client.start(bot_token=config.bot_token)
control_client = source_client
clients.append(source_client)
else:
source_client = build_client(config, "user")
await source_client.connect()
clients.append(source_client)
if not await source_client.is_user_authorized():
raise ConfigError(["Telegram user session is not authorized. Run `python main.py login` (see README)."])
control_client = None
if config.bot_token and control:
control_client = build_client(config, "bot")
await control_client.start(bot_token=config.bot_token)
clients.append(control_client)
is_bot = config.telegram_mode == "bot"
try:
entity = await resolve_channel(source_client, config.source_channel, is_bot=is_bot, db=db)
channel_id = utils.get_peer_id(entity)
except BotError as exc:
if not is_bot:
raise
# A bot learns a private channel from its first live post (handled below).
log.warning("%s", exc.public)
entity = None
channel_id = int(config.source_channel) if config.source_channel.lstrip("-").isdigit() else None
if channel_id is None:
raise
log.info("Source channel resolved (id %s, mode %s)", channel_id, config.telegram_mode)
source = TelethonSource(source_client, entity, channel_id=channel_id, mode=config.telegram_mode, config=config, db=db)
downloader = Downloader(
source,
config.download_dir,
min_free_bytes=config.min_free_disk_space,
max_temp_bytes=config.max_temp_storage,
compute_sha256=config.compute_sha256,
)
notifier = TelethonNotifier(control_client) if control_client else LogNotifier()
reporter = ProgressReporter(notifier, config.notify_chat_ids, config.progress_interval)
queue = QueueManager(
config, db, drive, downloader, FolderResolver(drive, db, root_id), source=source, reporter=reporter
)
removed = downloader.cleanup_orphans(db, keep_failed=config.keep_failed_downloads)
log.info("Startup cleanup removed %d abandoned temporary file(s)", removed)
if config.process_mode in ("new", "both"):
queue.listening = True
@source_client.on(events.NewMessage(func=lambda e: e.chat_id == channel_id))
async def _on_channel_post(event) -> None: # pragma: no cover - needs live Telegram
if source.entity is None:
source.entity = await event.get_input_chat()
db.set_kv(ACCESS_HASH_KEY.format(source.entity.channel_id), str(source.entity.access_hash))
outcome = queue.register(source.convert(event.message))
log.info("Live post %s: %s", event.message.id, outcome)
if config.process_mode == "new" and db.get_int(CURSOR_KEY, 0) == 0 and entity is not None:
latest = await source.latest_message_id()
if latest:
db.raise_int(CURSOR_KEY, latest)
if control_client:
register_command_handlers(control_client, CommandRouter(config, db, queue))
return App(config=config, db=db, lock=lock, queue=queue, source=source, clients=clients)
async def run_service(config: Config) -> None:
app = await build_app(config, control=True)
queue = app.queue
stop_event = asyncio.Event()
loop = asyncio.get_running_loop()
for sig in (signal.SIGINT, signal.SIGTERM):
try:
loop.add_signal_handler(sig, stop_event.set)
except NotImplementedError: # pragma: no cover - Windows
pass
from models import ACTIVE_TASK_STATUSES
desired = queue.desired_state()
if desired != "paused" and (config.auto_start or desired == "running"):
if config.process_mode != "new" and app.source.entity is not None:
try:
result = await queue.sync()
log.info("Startup sync: %s", result.summary().replace("\n", "; "))
except BotError as exc:
log.error("Startup sync failed: %s", exc)
queue.start()
booted = queue.boot() # also resumes RUNNING/QUEUED /upload tasks from the database
active_tasks = app.db.tasks_with_status(*ACTIVE_TASK_STATUSES)
if booted == "paused":
state_text = "queue is paused; send /resume to continue"
elif booted == "running":
state_text = "processing resumed automatically"
else:
state_text = "send /upload to process a message range, or /sync then /run for the whole channel"
if active_tasks:
state_text += "; upload tasks: " + ", ".join(f"#{t.id} {t.status.value}" for t in active_tasks)
if queue.reporter:
await queue.reporter.notify(f"Bot online ({app.db.describe()}); {state_text}.")
waiters = [asyncio.ensure_future(stop_event.wait())]
waiters += [asyncio.ensure_future(client.disconnected) for client in app.clients]
await asyncio.wait(waiters, return_when=asyncio.FIRST_COMPLETED)
log.info("Shutting down: checkpointing current work")
# Heroku sends SIGTERM and allows 30 seconds before SIGKILL.
await queue.shutdown(timeout=20)
for waiter in waiters:
waiter.cancel()
await app.close()
log.info("Shutdown complete")
async def run_headless(config: Config, *, sync_only: bool = False, rescan: bool = False) -> None:
app = await build_app(config, control=False)
try:
if app.source.entity is not None:
result = await app.queue.sync(rescan=rescan)
print(result.summary())
if not sync_only:
app.queue.listening = False
app.queue.start()
await app.queue.wait()
finally:
await app.close()
async def login(config: Config, *, session_string: bool) -> None:
from telethon import TelegramClient
from telethon.sessions import StringSession
if session_string:
client = TelegramClient(StringSession(), config.api_id, config.api_hash)
else:
from telegram_client import build_client
client = build_client(config, "user")
# Telethon asks for the phone number, the login code and (only if enabled) the
# 2FA password interactively; none of them is stored by this project.
await client.start()
me = await client.get_me()
print(f"Logged in as user id {me.id}.")
if session_string:
path = config.session_dir / "session_string.txt"
config.session_dir.mkdir(parents=True, exist_ok=True)
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
with os.fdopen(fd, "w") as fh:
fh.write(client.session.save())
print(
f"Session string written to {path} (permissions 600). It grants full access to your Telegram "
"account: never share or commit it.\nSet it on Heroku with:\n"
f' heroku config:set TELEGRAM_SESSION_STRING="$(cat {path})" --app YOUR_APP_NAME\n'
"then delete the file."
)
await client.disconnect()
def drive_init(config: Config, name: str) -> None:
from drive_client import DriveClient
drive = DriveClient.from_config(config)
folder = drive.create_folder("root", name)
print(f"Created Drive folder {name!r}. Set GOOGLE_DRIVE_ROOT_FOLDER_ID={folder.id}")
def print_status(config: Config) -> None:
db = Database.from_config(config)
print(f"Database: {db.describe()}")
for state, count in db.counts().items():
print(f" {state:<12} {count}")
paused = db.get_kv("queue_paused", "0") == "1"
print(f"Queue: {'paused' if paused else ('running' if db.get_kv('queue_global_run', '0') == '1' else 'stopped')}")
for task in db.recent_tasks(5):
print(f" task #{task.id} {task.status.value} range {task.range_text} uploaded {task.uploaded_count}")
db.close()
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Telegram channel -> Google Drive organizer")
sub = parser.add_subparsers(dest="command")
sub.add_parser("run", help="run the bot worker (default)")
sub.add_parser("process", help="headless: sync and process the queue once, then exit")
p_sync = sub.add_parser("sync", help="discover channel videos and exit")
p_sync.add_argument("--rescan", action="store_true", help="scan again from START_MESSAGE_ID")
sub.add_parser("status", help="print queue counts")
sub.add_parser("check-config", help="validate configuration and exit")
p_login = sub.add_parser("login", help="log in the Telegram user account (user mode)")
p_login.add_argument("--session-string", action="store_true", help="create a TELEGRAM_SESSION_STRING for Heroku")
p_auth = sub.add_parser("drive-auth", help="authorize Google Drive (OAuth) and save the token")
p_auth.add_argument("--no-browser", action="store_true")
p_init = sub.add_parser("drive-init", help="create a Drive root folder owned by this app")
p_init.add_argument("--name", required=True)
args = parser.parse_args(argv)
command = args.command or "run"
try:
config = Config.from_env()
except ConfigError as exc:
print(exc, file=sys.stderr)
return 2
setup_logging(config.log_level, log_dir=None if config.is_heroku else config.log_dir, secrets=config.secrets())
try:
if command in ("run", "process", "sync", "check-config"):
config.validate("telegram", "drive", "database")
if command == "check-config":
print("Configuration OK.")
elif command == "run":
asyncio.run(run_service(config))
elif command == "process":
asyncio.run(run_headless(config))
elif command == "sync":
asyncio.run(run_headless(config, sync_only=True, rescan=args.rescan))
elif command == "status":
config.validate("database")
print_status(config)
elif command == "login":
if not config.api_id or not config.api_hash:
raise ConfigError(["TELEGRAM_API_ID and TELEGRAM_API_HASH are required for login"])
asyncio.run(login(config, session_string=args.session_string))
elif command == "drive-auth":
from drive_client import run_oauth_flow
path = run_oauth_flow(config, open_browser=not args.no_browser)
print(
f"Token saved to {path} (permissions 600).\nFor Heroku run:\n"
f' heroku config:set GOOGLE_OAUTH_TOKEN_JSON="$(cat {path})" --app YOUR_APP_NAME'
)
elif command == "drive-init":
config.validate("drive")
drive_init(config, args.name)
except ConfigError as exc:
log.error("%s", exc)
return 2
except InstanceLockError as exc:
log.error("%s", exc)
return 3
except BotError as exc:
log.error("%s", exc.public)
log.debug("details: %s", exc)
return 1
except KeyboardInterrupt:
return 130
return 0
if __name__ == "__main__":
sys.exit(main())
models.py
+258
View File
@@ -0,0 +1,258 @@
"""Shared data models and the processing state machine."""
from __future__ import annotations
import json
from dataclasses import dataclass, field, fields
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any
class State(str, Enum):
PENDING = "PENDING"
DOWNLOADING = "DOWNLOADING"
DOWNLOADED = "DOWNLOADED"
UPLOADING = "UPLOADING"
COMPLETED = "COMPLETED"
RETRY_WAIT = "RETRY_WAIT"
FAILED = "FAILED"
SKIPPED = "SKIPPED"
def __str__(self) -> str:
return self.value
# Allowed state changes. Anything else is a programming error and is rejected by the
# database layer, so a bug cannot silently mark an item COMPLETED.
TRANSITIONS: dict[State, frozenset[State]] = {
State.PENDING: frozenset(
{State.DOWNLOADING, State.DOWNLOADED, State.COMPLETED, State.SKIPPED, State.RETRY_WAIT, State.FAILED}
),
State.DOWNLOADING: frozenset({State.DOWNLOADED, State.PENDING, State.RETRY_WAIT, State.FAILED}),
State.DOWNLOADED: frozenset(
{State.UPLOADING, State.PENDING, State.COMPLETED, State.SKIPPED, State.RETRY_WAIT, State.FAILED}
),
State.UPLOADING: frozenset(
{State.COMPLETED, State.DOWNLOADED, State.PENDING, State.SKIPPED, State.RETRY_WAIT, State.FAILED}
),
State.RETRY_WAIT: frozenset({State.PENDING, State.DOWNLOADED, State.FAILED, State.SKIPPED}),
State.FAILED: frozenset({State.PENDING, State.DOWNLOADED}),
State.SKIPPED: frozenset({State.PENDING}),
State.COMPLETED: frozenset(),
}
IN_PROGRESS_STATES = frozenset({State.DOWNLOADING, State.UPLOADING})
def can_transition(current: State, new: State) -> bool:
return new in TRANSITIONS[current]
def source_key(channel_id: int, message_id: int) -> str:
"""Unique identity of a Telegram video, stored on the Drive file as appProperties."""
return f"{channel_id}:{message_id}"
@dataclass
class ParsedCaption:
raw: str
indexes: list[str] = field(default_factory=list)
title: str | None = None
topic: str | None = None
batch: str | None = None
extracted_by: str | None = None
extra: dict[str, str] = field(default_factory=dict)
warnings: list[str] = field(default_factory=list)
missing: list[str] = field(default_factory=list)
@dataclass(frozen=True)
class MediaInfo:
kind: str # "video" or "document"
size: int | None
mime_type: str | None
file_name: str | None
media_id: int | None
duration: float | None = None
@dataclass(frozen=True)
class SourceMessage:
channel_id: int
message_id: int
caption: str
media: MediaInfo | None
date: datetime | None = None
@dataclass(frozen=True)
class LocalFile:
path: Path
size: int
md5: str
sha256: str | None
@dataclass
class DriveFile:
id: str
name: str
size: int | None = None
md5: str | None = None
parents: tuple[str, ...] = ()
mime_type: str | None = None
trashed: bool = False
app_properties: dict[str, str] = field(default_factory=dict)
@classmethod
def from_api(cls, data: dict[str, Any]) -> "DriveFile":
size = data.get("size")
return cls(
id=data["id"],
name=data.get("name", ""),
size=int(size) if size is not None else None,
md5=data.get("md5Checksum"),
parents=tuple(data.get("parents") or ()),
mime_type=data.get("mimeType"),
trashed=bool(data.get("trashed", False)),
app_properties=dict(data.get("appProperties") or {}),
)
@dataclass
class VideoRecord:
channel_id: int
message_id: int
state: State
caption: str | None = None
indexes: list[str] = field(default_factory=list)
title: str | None = None
topic: str | None = None
batch: str | None = None
extracted_by: str | None = None
warnings: list[str] = field(default_factory=list)
media_kind: str | None = None
media_id: int | None = None
mime_type: str | None = None
original_filename: str | None = None
expected_size: int | None = None
message_date: str | None = None
sequence_no: int | None = None
local_path: str | None = None
local_size: int | None = None
sha256: str | None = None
md5: str | None = None
final_name: str | None = None
drive_folder_id: str | None = None
drive_file_id: str | None = None
upload_session_uri: str | None = None
upload_session_created_at: str | None = None
download_status: str | None = None
upload_status: str | None = None
attempts: int = 0
last_error: str | None = None
error_kind: str | None = None
next_retry_at: str | None = None
skip_reason: str | None = None
discovered_at: str | None = None
updated_at: str | None = None
download_started_at: str | None = None
downloaded_at: str | None = None
upload_started_at: str | None = None
completed_at: str | None = None
task_id: int | None = None
JSON_FIELDS = ("indexes", "warnings")
@property
def key(self) -> tuple[int, int]:
return (self.channel_id, self.message_id)
@property
def source_key(self) -> str:
return source_key(self.channel_id, self.message_id)
@classmethod
def column_names(cls) -> list[str]:
return [f.name for f in fields(cls)]
@classmethod
def from_row(cls, row: Any) -> "VideoRecord":
data = {name: row[name] for name in cls.column_names()}
for name in cls.JSON_FIELDS:
data[name] = json.loads(data[name]) if data[name] else []
data["state"] = State(data["state"])
return cls(**data)
# -- interactive /upload tasks ---------------------------------------------------------
class TaskStatus(str, Enum):
AWAITING_START = "AWAITING_START"
AWAITING_END = "AWAITING_END"
AWAITING_CONFIRMATION = "AWAITING_CONFIRMATION"
QUEUED = "QUEUED"
RUNNING = "RUNNING"
PAUSED = "PAUSED"
COMPLETED = "COMPLETED"
CANCELLED = "CANCELLED"
FAILED = "FAILED"
def __str__(self) -> str:
return self.value
SETUP_TASK_STATUSES = frozenset({TaskStatus.AWAITING_START, TaskStatus.AWAITING_END, TaskStatus.AWAITING_CONFIRMATION})
ACTIVE_TASK_STATUSES = frozenset({TaskStatus.QUEUED, TaskStatus.RUNNING, TaskStatus.PAUSED})
FINAL_TASK_STATUSES = frozenset({TaskStatus.COMPLETED, TaskStatus.CANCELLED, TaskStatus.FAILED})
# Per-message outcome inside a task (task_items.result).
TASK_ITEM_RESULTS = ("pending", "uploaded", "duplicate", "skipped", "non_video", "missing", "failed", "cancelled")
@dataclass
class UploadTask:
id: int
admin_user_id: int
status: TaskStatus
channel_ref: str | None = None
channel_id: int | None = None
channel_title: str | None = None
start_message_id: int | None = None
end_message_id: int | None = None
created_at: str | None = None
updated_at: str | None = None
confirmed_at: str | None = None
started_at: str | None = None
scanned_at: str | None = None
completed_at: str | None = None
current_message_id: int | None = None
examined_count: int = 0
eligible_count: int = 0
uploaded_count: int = 0
duplicate_count: int = 0
non_video_count: int = 0
skipped_count: int = 0
missing_count: int = 0
failed_count: int = 0
deleted_files_count: int = 0
last_error: str | None = None
@classmethod
def column_names(cls) -> list[str]:
return [f.name for f in fields(cls)]
@classmethod
def from_row(cls, row: Any) -> "UploadTask":
data = {name: row[name] for name in cls.column_names()}
data["status"] = TaskStatus(data["status"])
return cls(**data)
@property
def range_text(self) -> str:
if self.start_message_id is None:
return "-"
end = self.end_message_id if self.end_message_id is not None else "?"
return f"{self.start_message_id}–{end}"
naming.py
+182
View File
@@ -0,0 +1,182 @@
"""Filename and folder-name rules: sanitising, extension detection, collisions."""
from __future__ import annotations
import os
import re
import unicodedata
from pathlib import Path
from errors import PermanentError
VIDEO_EXTENSIONS = frozenset(
{".mp4", ".m4v", ".mov", ".mkv", ".webm", ".avi", ".flv", ".wmv", ".ts", ".m2ts", ".mts", ".3gp", ".mpg", ".mpeg", ".ogv"}
)
# Extensions that share a container format; a title saying ".mov" for an MP4 file is fine.
_FAMILIES = (
frozenset({".mp4", ".m4v", ".mov", ".3gp"}),
frozenset({".mkv", ".webm"}),
frozenset({".ts", ".m2ts", ".mts"}),
frozenset({".mpg", ".mpeg"}),
)
MIME_TO_EXTENSION = {
"video/mp4": ".mp4",
"video/x-m4v": ".m4v",
"video/quicktime": ".mov",
"video/x-matroska": ".mkv",
"video/webm": ".webm",
"video/x-msvideo": ".avi",
"video/x-flv": ".flv",
"video/x-ms-wmv": ".wmv",
"video/mp2t": ".ts",
"video/3gpp": ".3gp",
"video/mpeg": ".mpg",
"video/ogg": ".ogv",
}
EXTENSION_TO_MIME = {ext: mime for mime, ext in MIME_TO_EXTENSION.items()} | {
".mpeg": "video/mpeg",
".m2ts": "video/mp2t",
".mts": "video/mp2t",
}
MAX_NAME_LENGTH = 200
_ALWAYS_REPLACE = re.compile(r"[/\\]")
_STRICT_REPLACE = re.compile(r'[<>:"|?*]')
def normalize_key(text: str) -> str:
"""Comparison key for folder names: NFC, collapsed whitespace, case-insensitive."""
return " ".join(unicodedata.normalize("NFC", text).split()).casefold()
def clean_component(text: str | None, *, strict: bool = False) -> str:
"""Make a caption value safe as a Drive file/folder name while keeping Unicode.
Control characters and path separators are replaced; letters, bullets, dots,
parentheses and other punctuation are kept. ``strict`` additionally replaces
characters Windows cannot store (useful with Google Drive for desktop).
"""
if not text:
return ""
text = unicodedata.normalize("NFC", text)
text = "".join(" " if unicodedata.category(ch) == "Cc" else ch for ch in text)
text = _ALWAYS_REPLACE.sub("_", text)
if strict:
text = _STRICT_REPLACE.sub("_", text).rstrip(". ")
text = " ".join(text.split())
return truncate_name(text)
def split_video_extension(name: str) -> tuple[str, str]:
"""Split off a *video* extension only, so 'Lecture 1.5' keeps its '.5'."""
stem, ext = os.path.splitext(name)
if ext.lower() in VIDEO_EXTENSIONS and stem:
return stem, ext
return name, ""
def truncate_name(name: str, limit: int = MAX_NAME_LENGTH) -> str:
if len(name) <= limit:
return name
stem, ext = split_video_extension(name)
return stem[: limit - len(ext)].rstrip() + ext
def same_family(a: str, b: str) -> bool:
a, b = a.lower(), b.lower()
if a == b:
return True
return any(a in family and b in family for family in _FAMILIES)
def extension_from_mime(mime_type: str | None) -> str | None:
return MIME_TO_EXTENSION.get((mime_type or "").lower())
def mime_for_name(name: str, fallback: str | None = None) -> str:
_, ext = split_video_extension(name)
return EXTENSION_TO_MIME.get(ext.lower()) or fallback or "application/octet-stream"
def detect_video_extension(path: Path) -> str | None:
"""Identify the container from magic bytes; ``None`` if unknown."""
with open(path, "rb") as fh:
head = fh.read(4096)
if len(head) >= 12 and head[4:8] == b"ftyp":
brand = head[8:12]
if brand == b"qt ":
return ".mov"
if brand.startswith(b"3g"):
return ".3gp"
if brand.startswith(b"M4V"):
return ".m4v"
return ".mp4"
if head.startswith(b"\x1a\x45\xdf\xa3"):
return ".webm" if b"webm" in head[:64] else ".mkv"
if head[:4] == b"RIFF" and head[8:12] == b"AVI ":
return ".avi"
if head[:3] == b"FLV":
return ".flv"
if head[:8] == b"\x30\x26\xb2\x75\x8e\x66\xcf\x11":
return ".wmv"
if head[:4] in (b"\x00\x00\x01\xba", b"\x00\x00\x01\xb3"):
return ".mpg"
if head[:4] == b"OggS":
return ".ogv"
if len(head) > 376 and head[0] == 0x47 and head[188] == 0x47 and head[376] == 0x47:
return ".ts"
return None
def build_filename(
title: str | None,
*,
message_id: int,
original_filename: str | None,
mime_type: str | None,
detected_ext: str | None,
missing_title_policy: str = "fallback",
mismatch_policy: str = "keep",
strict: bool = False,
) -> tuple[str, list[str]]:
"""Return the Drive filename for a video plus any warnings worth recording.
The Title field is used verbatim (sanitised). The extension is only added when
the title has none; a title whose extension disagrees with the real container
is kept as-is unless EXTENSION_MISMATCH_POLICY=append.
"""
warnings: list[str] = []
name = clean_component(title, strict=strict)
if not name:
if missing_title_policy == "fail":
raise PermanentError(
f"message {message_id} has no usable Title", public="Caption has no usable Title (MISSING_TITLE_POLICY=fail)"
)
name = clean_component(original_filename, strict=strict)
if name:
warnings.append("Title missing; used the Telegram file name")
else:
name = f"telegram_{message_id}"
warnings.append("Title missing; used a generated name")
_, ext = split_video_extension(name)
_, original_ext = split_video_extension(original_filename or "")
actual = detected_ext or extension_from_mime(mime_type) or (original_ext.lower() or None)
if not ext:
if actual:
name = truncate_name(name, MAX_NAME_LENGTH - len(actual)) + actual
else:
warnings.append("could not determine the file extension")
elif detected_ext and not same_family(ext, detected_ext):
warnings.append(f"title extension {ext} does not match detected container {detected_ext}")
if mismatch_policy == "append":
name = truncate_name(name, MAX_NAME_LENGTH - len(detected_ext)) + detected_ext
return name, warnings
def add_suffix(name: str, suffix: str) -> str:
"""'Class 01.mp4' + '(2)' -> 'Class 01 (2).mp4'."""
stem, ext = split_video_extension(name)
if not ext:
stem, ext = os.path.splitext(name)
if len(ext) > 6 or not stem:
stem, ext = name, ""
return truncate_name(f"{stem} {suffix}{ext}")
progress.py
+259
View File
@@ -0,0 +1,259 @@
"""Progress tracking and throttled Telegram progress messages.
Transfer callbacks only update in-memory counters (cheap, thread-safe enough for
plain attribute writes). A background task renders the board every
PROGRESS_INTERVAL seconds and edits one Telegram message per chat, so progress never
floods the chat with new messages.
"""
from __future__ import annotations
import asyncio
import html
import logging
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Callable, Protocol
log = logging.getLogger(__name__)
class HtmlReply(str):
"""A message whose text is Telegram HTML (dynamic parts already escaped)."""
def esc(value) -> str:
return html.escape(str(value), quote=False)
class Notifier(Protocol):
async def send(self, chat_id: int, text: str) -> int | None: ...
async def edit(self, chat_id: int, message_id: int, text: str) -> None: ...
class LogNotifier:
"""Notifier used when no control bot is configured (headless mode)."""
async def send(self, chat_id: int, text: str) -> int | None:
log.info("%s", text.replace("\n", " | "))
return None
async def edit(self, chat_id: int, message_id: int, text: str) -> None:
log.info("%s", text.replace("\n", " | "))
def format_bytes(value: float | None) -> str:
if value is None:
return "?"
for unit in ("B", "KB", "MB", "GB"):
if abs(value) < 1024:
return f"{value:.0f} {unit}" if unit == "B" else f"{value:.1f} {unit}"
value /= 1024
return f"{value:.2f} TB"
def format_duration(seconds: float | None) -> str:
if seconds is None:
return "?"
seconds = int(round(seconds))
hours, rem = divmod(seconds, 3600)
minutes, secs = divmod(rem, 60)
if hours:
return f"{hours}h {minutes}m"
if minutes:
return f"{minutes}m {secs}s"
return f"{secs}s"
@dataclass
class TransferProgress:
total: int | None
done: int = 0
clock: Callable[[], float] = time.monotonic
window: float = 30.0
_samples: deque = field(default_factory=deque)
def update(self, done: int, total: int | None = None) -> None:
if total is not None:
self.total = total
self.done = done
now = self.clock()
self._samples.append((now, done))
while len(self._samples) > 2 and now - self._samples[0][0] > self.window:
self._samples.popleft()
@property
def percent(self) -> float | None:
if not self.total:
return None
return min(100.0, 100.0 * self.done / self.total)
@property
def speed(self) -> float | None:
if len(self._samples) < 2:
return None
(t0, b0), (t1, b1) = self._samples[0], self._samples[-1]
if t1 - t0 < 1.0 or b1 < b0:
return None
return (b1 - b0) / (t1 - t0)
@property
def eta(self) -> float | None:
speed = self.speed
if not speed or not self.total:
return None
return max(0.0, (self.total - self.done) / speed)
def describe(self) -> str:
pct = self.percent
parts = [f"{pct:.0f}%" if pct is not None else format_bytes(self.done)]
parts.append(f"({format_bytes(self.done)} / {format_bytes(self.total)})")
text = " ".join(parts)
if self.speed:
text += f" · {format_bytes(self.speed)}/s"
if self.eta is not None and self.done < (self.total or 0):
text += f" · ETA {format_duration(self.eta)}"
return text
@dataclass
class ProgressBoard:
position: int = 0
total: int = 0
topic: str = ""
title: str = ""
status: str = "Idle"
download: TransferProgress | None = None
upload: TransferProgress | None = None
clock: Callable[[], float] = time.monotonic
# Set while an /upload task runs: id, start, end, current, completed, skipped, failed
task: dict | None = None
def begin_item(self, position: int, total: int, topic: str, title: str) -> None:
self.position, self.total, self.topic, self.title = position, total, topic, title
self.status = "Preparing"
self.download = None
self.upload = None
def start_download(self, total: int | None) -> None:
self.status = "Downloading"
self.download = TransferProgress(total, clock=self.clock)
def start_upload(self, total: int) -> None:
self.status = "Uploading"
self.upload = TransferProgress(total, clock=self.clock)
_TASK_STATUS = {
"Downloading": "Downloading from Telegram",
"Uploading": "Uploading to Google Drive",
"Verifying": "Verifying upload in Google Drive",
}
def render(self) -> str:
if self.task is not None:
return self._render_task()
if self.status == "Idle":
return "Queue idle"
lines = [
f"Processing video {self.position} of {self.total}",
f"Topic: {self.topic}",
f"Title: {self.title}",
f"Status: {self.status}",
f"Download: {self.download.describe() if self.download else 'Waiting'}",
f"Upload: {self.upload.describe() if self.upload else 'Waiting'}",
]
return "\n".join(lines)
def _render_task(self) -> str:
t = self.task or {}
status = self._TASK_STATUS.get(self.status, self.status)
lines = [
"🚀 <b>Upload Task Running</b>",
"",
f"Task: #{esc(t.get('id'))}",
f"Range: {esc(t.get('start'))}–{esc(t.get('end'))}",
f"Current Message: {esc(t.get('current') or '-')}",
f"Current Video: {esc(self.title or '-')}",
f"Topic: {esc(self.topic or '-')}",
f"Status: {esc(status)}",
]
if self.download:
lines.append(f"Download: {esc(self.download.describe())}")
if self.upload:
lines.append(f"Upload: {esc(self.upload.describe())}")
lines += [
"",
f"Completed: {t.get('completed', 0)}",
f"Skipped: {t.get('skipped', 0)}",
f"Failed: {t.get('failed', 0)}",
]
return HtmlReply("\n".join(lines))
class ProgressReporter:
def __init__(
self,
notifier: Notifier,
chat_ids: tuple[int, ...] | list[int],
interval: float,
*,
clock: Callable[[], float] = time.monotonic,
):
self.notifier = notifier
self.chat_ids = tuple(chat_ids) or (0,)
self.interval = interval
self.board = ProgressBoard(clock=clock)
self._message_ids: dict[int, int] = {}
self._last_text: str | None = None
self._task: asyncio.Task | None = None
async def start(self) -> None:
self._message_ids.clear()
self._last_text = None
if self._task is None or self._task.done():
self._task = asyncio.create_task(self._loop(), name="progress-reporter")
async def stop(self) -> None:
if self._task:
self._task.cancel()
try:
await self._task
except asyncio.CancelledError:
pass
self._task = None
await self.refresh()
async def _loop(self) -> None:
while True:
await asyncio.sleep(self.interval)
await self.refresh()
async def refresh(self) -> None:
text = self.board.render()
if text == self._last_text:
return
self._last_text = text
for chat_id in self.chat_ids:
try:
message_id = self._message_ids.get(chat_id)
if message_id is None:
sent = await self.notifier.send(chat_id, text)
if sent is not None:
self._message_ids[chat_id] = sent
else:
await self.notifier.edit(chat_id, message_id, text)
except Exception as exc: # progress must never break processing
log.debug("Progress update to %s failed: %s", chat_id, exc)
def new_message(self) -> None:
"""Start a fresh progress message (e.g. for each /upload task)."""
self._message_ids.clear()
self._last_text = None
async def notify(self, text: str, *, extra_chat_ids: tuple[int, ...] = ()) -> None:
for chat_id in dict.fromkeys((*self.chat_ids, *extra_chat_ids)):
try:
await self.notifier.send(chat_id, text)
except Exception as exc:
log.warning("Notification to %s failed: %s", chat_id, exc)
pytest.ini
+4
View File
@@ -0,0 +1,4 @@
[pytest]
pythonpath = .
testpaths = tests
addopts = -q
queue_manager.py
+1003
View File
File diff suppressed because it is too large Load Diff
requirements-dev.txt
+5
View File
@@ -0,0 +1,5 @@
-r requirements.txt
pytest==9.1.1
cryptography>=42
# Optional: runs the database tests against a throwaway local PostgreSQL as well.
pgserver==0.1.4
requirements.txt
+10
View File
@@ -0,0 +1,10 @@
# Runtime dependencies (pinned to the versions the test suite was run with).
Telethon==1.45.0
cryptg==0.5.0.post0
google-api-python-client==2.201.0
google-auth==2.62.0
google-auth-httplib2==0.4.4
google-auth-oauthlib==1.5.0
requests==2.34.2
python-dotenv==1.2.4
psycopg[binary]==3.3.6
start.sh
Executable
+16
View File
@@ -0,0 +1,16 @@
#!/usr/bin/env bash
# Local convenience launcher: creates a virtualenv, installs dependencies, runs main.py.
# Usage: ./start.sh (runs the bot worker)
# ./start.sh login (any main.py subcommand)
set -euo pipefail
cd "$(dirname "$0")"
PYTHON="${PYTHON:-python3.12}"
if [ ! -d .venv ]; then
"$PYTHON" -m venv .venv
fi
.venv/bin/pip install -q -r requirements.txt
if [ ! -f .env ]; then
echo "No .env found. Copy .env.example to .env and fill it in first." >&2
exit 2
fi
exec .venv/bin/python main.py "${@:-run}"
telegram_client.py
+345
View File
@@ -0,0 +1,345 @@
"""Telegram access through Telethon (MTProto).
Both operating modes use MTProto, not the HTTP Bot API:
* TELEGRAM_MODE=user — your own account (TELEGRAM_SESSION_STRING or a local session
file). Can read full channel history the account can see; downloads up to the
largest file size Telegram allows.
* TELEGRAM_MODE=bot — a bot (TELEGRAM_BOT_TOKEN) that is an admin of the channel.
Receives new posts live. Bots cannot list history; past posts are fetched by
probing message IDs (best effort). Because this is MTProto rather than the HTTP
Bot API, the HTTP API's 20 MB getFile limit does not apply.
"""
from __future__ import annotations
import asyncio
import logging
import os
import re
from datetime import datetime
from typing import AsyncIterator
from errors import BotError, PermanentError, TransientError
from link_parser import marked_channel_id
from models import MediaInfo, SourceMessage
from naming import VIDEO_EXTENSIONS, split_video_extension
log = logging.getLogger(__name__)
REQUEST_SIZE = 512 * 1024 # Telegram's maximum upload.getFile chunk
ACCESS_HASH_KEY = "channel_access_hash:{}"
def classify_media(message, include_documents: bool) -> MediaInfo | None:
"""Return MediaInfo for a supported video message, else None.
Excludes photos, stickers, GIF animations, round video notes, voice/audio and
non-video documents (unless they are video files and documents are enabled).
"""
from telethon.tl import types
media = getattr(message, "media", None)
if not isinstance(media, types.MessageMediaDocument):
return None
doc = media.document
if not isinstance(doc, types.Document):
return None
video_attr = None
file_name = None
for attr in doc.attributes:
if isinstance(attr, (types.DocumentAttributeAnimated, types.DocumentAttributeSticker)):
return None
if isinstance(attr, types.DocumentAttributeVideo):
video_attr = attr
elif isinstance(attr, types.DocumentAttributeFilename):
file_name = attr.file_name
mime = (doc.mime_type or "").lower()
if video_attr is not None:
if getattr(video_attr, "round_message", False):
return None
kind = "video"
elif include_documents and (
mime.startswith("video/") or (file_name and split_video_extension(file_name)[1].lower() in VIDEO_EXTENSIONS)
):
kind = "document"
else:
return None
return MediaInfo(
kind=kind,
size=doc.size,
mime_type=doc.mime_type,
file_name=file_name,
media_id=doc.id,
duration=getattr(video_attr, "duration", None),
)
def to_source_message(message, channel_id: int, include_documents: bool) -> SourceMessage:
return SourceMessage(
channel_id=channel_id,
message_id=message.id,
caption=getattr(message, "message", None) or "",
media=classify_media(message, include_documents),
date=getattr(message, "date", None),
)
def map_telegram_error(exc: BaseException) -> BotError:
from telethon import errors
if isinstance(exc, errors.FloodWaitError):
return TransientError(
f"Telegram flood wait {exc.seconds}s", retry_after=exc.seconds, rate_limited=True,
public=f"Telegram rate limit; waiting {exc.seconds}s",
)
if isinstance(exc, (errors.ChannelPrivateError, errors.ChatAdminRequiredError, errors.ChannelInvalidError)):
return PermanentError(str(exc), public="No access to the source channel (check membership/admin rights)")
if isinstance(exc, (errors.AuthKeyError, errors.UnauthorizedError)):
return PermanentError(str(exc), public="Telegram session is no longer authorized; log in again")
if isinstance(exc, errors.RPCError):
return TransientError(f"Telegram error: {exc}", public=f"Telegram error ({type(exc).__name__})")
return TransientError(f"Telegram connection error: {exc}", public="Telegram connection problem")
def build_client(config, role: str):
"""Create a Telethon client. role: 'user' or 'bot'."""
from telethon import TelegramClient
from telethon.sessions import StringSession
if role == "user" and config.session_string:
session = StringSession(config.session_string)
else:
config.session_dir.mkdir(parents=True, exist_ok=True)
os.chmod(config.session_dir, 0o700)
session = str(config.session_dir / role)
return TelegramClient(
session,
config.api_id,
config.api_hash,
flood_sleep_threshold=config.flood_sleep_threshold,
connection_retries=10,
retry_delay=5,
auto_reconnect=True,
)
def parse_channel_reference(ref: str) -> str | int:
ref = ref.strip()
if re.fullmatch(r"-?\d+", ref):
return int(ref)
private = re.match(r"(?:https?://)?t\.me/c/(\d+)", ref)
if private:
return marked_channel_id(int(private.group(1)))
public = re.match(r"(?:https?://)?t\.me/([A-Za-z0-9_]+)", ref)
if public:
return public.group(1)
return ref.lstrip("@")
async def resolve_channel(client, ref: str, *, is_bot: bool, db=None):
"""Resolve the configured channel to a Telethon entity.
The channel's access hash is cached in the database, so a bot (which cannot
list dialogs) still resolves a private channel after a restart on an ephemeral
filesystem.
"""
from telethon import utils
from telethon.tl import types
target = parse_channel_reference(ref)
try:
entity = await client.get_entity(target)
except (ValueError, TypeError) as first_error:
entity = None
if isinstance(target, int):
real_id, _ = utils.resolve_id(target)
cached = db.get_kv(ACCESS_HASH_KEY.format(real_id)) if db else None
if cached:
entity = await client.get_entity(types.InputPeerChannel(real_id, int(cached)))
elif not is_bot:
await client.get_dialogs() # fills the entity cache for chats the account is in
entity = await client.get_entity(target)
if entity is None:
raise PermanentError(
f"cannot resolve channel {ref!r}: {first_error}",
public=(
"Cannot resolve SOURCE_CHANNEL_ID. Use the @username for public channels. For a private channel "
"in bot mode, make the bot an admin and post one message in the channel so the bot learns it."
),
) from first_error
if db is not None and getattr(entity, "access_hash", None):
db.set_kv(ACCESS_HASH_KEY.format(entity.id), str(entity.access_hash))
return entity
class TelethonSource:
def __init__(self, client, entity, *, channel_id: int, mode: str, config, db=None):
self.client = client
self.entity = entity
self.channel_id = channel_id
self.mode = mode
self.config = config
self.db = db
self._entities: dict[int, object] = {}
@property
def channel_title(self) -> str | None:
return getattr(self.entity, "title", None)
def convert(self, message, channel_id: int | None = None) -> SourceMessage:
return to_source_message(message, channel_id or self.channel_id, self.config.include_document_videos)
async def _entity_for(self, channel_id: int | None):
if channel_id is None or channel_id == self.channel_id:
if self.entity is not None:
return self.entity
if channel_id in self._entities:
return self._entities[channel_id]
entity = await resolve_channel(self.client, str(channel_id), is_bot=self.mode == "bot", db=self.db)
self._entities[channel_id] = entity
return entity
# -- /upload task support ----------------------------------------------------------
async def resolve_link(self, link) -> tuple[int, str | None]:
"""Resolve a MessageLink's channel through the authenticated client.
Raises PermanentError if this account/bot cannot access the channel.
"""
from telethon import utils
try:
entity = await resolve_channel(self.client, link.channel_ref, is_bot=self.mode == "bot", db=self.db)
except BotError:
raise
except Exception as exc:
raise PermanentError(
f"cannot resolve {link.url}: {exc}",
public="This channel is not accessible to the bot's Telegram account (join it, or make the bot an admin).",
) from exc
channel_id = utils.get_peer_id(entity)
self._entities[channel_id] = entity
return channel_id, getattr(entity, "title", None)
async def message_accessible(self, channel_id: int, message_id: int) -> bool:
entity = await self._entity_for(channel_id)
try:
message = await self.client.get_messages(entity, ids=message_id)
except Exception as exc:
raise map_telegram_error(exc) from exc
return message is not None and message.__class__.__name__ != "MessageEmpty"
async def iter_range(self, channel_id: int, start_id: int, end_id: int) -> AsyncIterator[tuple[int, SourceMessage | None]]:
"""Yield (message_id, SourceMessage or None) for every ID in [start_id, end_id], ascending.
Uses messages.getMessages / channels.getMessages by ID, which both user
accounts and bots (channel admins) may call; no history listing is needed.
"""
entity = await self._entity_for(channel_id)
current = start_id
try:
while current <= end_id:
ids = list(range(current, min(current + 100, end_id + 1)))
messages = await self.client.get_messages(entity, ids=ids)
for message_id, message in zip(ids, messages):
if message is None or message.__class__.__name__ in ("MessageEmpty", "MessageService"):
yield message_id, None
else:
yield message_id, self.convert(message, channel_id)
current = ids[-1] + 1
except (asyncio.CancelledError, GeneratorExit):
raise
except BotError:
raise
except Exception as exc:
raise map_telegram_error(exc) from exc
# -- discovery for the configured channel ------------------------------------------
async def iter_history(self, *, min_id: int, start_date: datetime | None) -> AsyncIterator[SourceMessage]:
try:
if self.mode == "user":
kwargs = {"reverse": True, "min_id": max(min_id - 1, 0)}
if start_date:
kwargs["offset_date"] = start_date
async for message in self.client.iter_messages(self.entity, **kwargs):
yield self.convert(message)
else:
async for message in self._probe_history(min_id, start_date):
yield self.convert(message)
except (asyncio.CancelledError, GeneratorExit):
raise
except BotError:
raise
except Exception as exc:
raise map_telegram_error(exc) from exc
async def _probe_history(self, start_id: int, start_date: datetime | None):
"""Bots cannot call messages.getHistory, but may fetch messages by ID."""
current, empty = max(start_id, 1), 0
end = self.config.bot_scan_end_id
while True:
ids = list(range(current, current + 100))
if end:
ids = [i for i in ids if i <= end]
if not ids:
return
messages = await self.client.get_messages(self.entity, ids=ids)
for message in messages:
if message is None or message.__class__.__name__ == "MessageEmpty":
empty += 1
continue
empty = 0
if start_date and message.date and message.date < start_date:
continue
yield message
if not end and empty >= self.config.bot_scan_max_empty:
return
current += 100
async def _fetch(self, message_id: int, channel_id: int | None = None):
entity = await self._entity_for(channel_id)
try:
message = await self.client.get_messages(entity, ids=message_id)
except Exception as exc:
raise map_telegram_error(exc) from exc
if message is None or getattr(message, "media", None) is None or not getattr(message, "document", None):
raise PermanentError(
f"message {message_id} no longer exists or has no media", public="Source message was deleted or has no video"
)
return message
async def iter_media_chunks(
self, message_id: int, *, offset: int, expected_size: int | None, channel_id: int | None = None
) -> AsyncIterator[bytes]:
from telethon import errors
message = await self._fetch(message_id, channel_id)
position, refreshes = offset, 0
while True:
try:
async for chunk in self.client.iter_download(
message.document, offset=position, request_size=REQUEST_SIZE, file_size=message.document.size
):
position += len(chunk)
yield chunk
return
except errors.FileReferenceExpiredError:
refreshes += 1
if refreshes > 3:
raise TransientError("file reference kept expiring", public="Telegram file reference expired")
message = await self._fetch(message_id, channel_id)
except (asyncio.CancelledError, GeneratorExit):
raise
except BotError:
raise
except OSError:
raise
except Exception as exc:
raise map_telegram_error(exc) from exc
async def latest_message_id(self) -> int | None:
if self.mode != "user":
return None
messages = await self.client.get_messages(self.entity, limit=1)
return messages[0].id if messages else None
tests/conftest.py
+373
View File
@@ -0,0 +1,373 @@
"""Shared fixtures and in-memory fakes for Telegram and Google Drive.
No test needs real credentials or network access. Database-backed tests run against
SQLite and, when the optional ``pgserver`` dev package is installed, a throwaway
local PostgreSQL server as well.
"""
from __future__ import annotations
import asyncio
import hashlib
import itertools
import os
import shutil
import tempfile
import threading
import time
from pathlib import Path
import pytest
from config import Config
from database import Database
from downloader import Downloader
from drive_client import FOLDER_MIME, SOURCE_PROPERTY
from errors import Interrupted
from models import DriveFile, MediaInfo, SourceMessage
CHANNEL_ID = -1001234567890
EXAMPLE_CAPTION = """Index: 001
Index: 002
Title: Class 01 Aman sir.mp4
Topic: Misc• Grammar Practice
Batch: English ( SSC MAINS )
Extracted By: @PUTINxINDIA"""
def make_config(tmp_path: Path, **env) -> Config:
base = {
"TELEGRAM_API_ID": "12345",
"TELEGRAM_API_HASH": "0123456789abcdef0123456789abcdef",
"TELEGRAM_BOT_TOKEN": "123456:TEST-TOKEN-VALUE",
"ADMIN_USER_IDS": "111",
"SOURCE_CHANNEL_ID": str(CHANNEL_ID),
"DATABASE_PATH": str(tmp_path / "bot.db"),
"DOWNLOAD_DIR": str(tmp_path / "downloads"),
"SESSION_DIR": str(tmp_path / "sessions"),
"LOG_DIR": str(tmp_path / "logs"),
"MIN_FREE_DISK_SPACE": "0",
"RETRY_DELAY": "0",
"PROCESSING_DELAY": "0",
"MAX_RETRIES": "3",
"UPLOAD_CHUNK_SIZE": "256KB",
}
base.update({k: str(v) for k, v in env.items()})
return Config.from_env(base)
def video_bytes(size: int, seed: int = 0) -> bytes:
"""Bytes that look like an MP4 container (ftyp box) followed by filler."""
header = b"\x00\x00\x00\x18ftypisom\x00\x00\x02\x00isomiso2"
filler = hashlib.sha256(str(seed).encode()).digest()
body = (filler * (size // len(filler) + 1))[: size - len(header)]
return header + body
# -- PostgreSQL (optional) ---------------------------------------------------------------
_PG_SERVER = None
def _pg_uri() -> str | None:
global _PG_SERVER
try:
import pgserver
except ImportError:
return None
if _PG_SERVER is None:
_PG_SERVER = pgserver.get_server(tempfile.mkdtemp(prefix="tdb-pg-"), cleanup_mode="stop")
return _PG_SERVER.get_uri()
_pg_counter = itertools.count()
def _fresh_pg_url() -> str:
import psycopg
uri = _pg_uri()
name = f"test_{os.getpid()}_{next(_pg_counter)}"
with psycopg.connect(uri, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {name}")
return uri.replace("/postgres?", f"/{name}?")
BACKENDS = ["sqlite"] + (["postgres"] if _pg_uri() else [])
@pytest.fixture(params=BACKENDS)
def backend(request):
return request.param
@pytest.fixture
def db_url(backend):
return _fresh_pg_url() if backend == "postgres" else None
@pytest.fixture
def config(tmp_path, db_url):
env = {"DATABASE_URL": db_url} if db_url else {}
return make_config(tmp_path, **env)
@pytest.fixture
def db(config):
database = Database.from_config(config)
yield database
database.close()
# -- fakes -------------------------------------------------------------------------------
class FakeSource:
CHUNK = 512 * 1024
def __init__(self, channel_id: int = CHANNEL_ID):
self.channel_id = channel_id
self.channel_title = "Test Channel"
# Other channels the fake account can access: marked id -> title, plus usernames.
self.accessible = {channel_id: self.channel_title}
self.usernames = {"test_channel": channel_id}
self.messages: dict[int, tuple[SourceMessage, bytes]] = {}
self.failures: dict[int, list] = {}
self.download_offsets: list[tuple[int, int]] = []
self.events: list[tuple[str, int]] = []
def add(self, message_id, caption, content: bytes | None = None, *, media_id=None, kind="video", mime="video/mp4", file_name=None):
content = content if content is not None else video_bytes(700_000, seed=message_id)
media = MediaInfo(kind, len(content), mime, file_name, media_id if media_id is not None else 10_000 + message_id)
msg = SourceMessage(self.channel_id, message_id, caption, media)
self.messages[message_id] = (msg, content)
return msg
def add_text(self, message_id, text):
msg = SourceMessage(self.channel_id, message_id, text, None)
self.messages[message_id] = (msg, b"")
return msg
async def iter_history(self, *, min_id, start_date):
for message_id in sorted(self.messages):
if message_id >= min_id:
yield self.messages[message_id][0]
async def resolve_link(self, link):
from errors import PermanentError
if link.kind == "private":
channel_id = int(link.channel_ref)
else:
channel_id = self.usernames.get(link.channel.lower())
if channel_id not in self.accessible:
raise PermanentError("no access", public="This channel is not accessible to the bot's Telegram account.")
return channel_id, self.accessible[channel_id]
async def message_accessible(self, channel_id, message_id):
return channel_id == self.channel_id and message_id in self.messages
async def iter_range(self, channel_id, start_id, end_id):
self.events.append(("scan", (start_id, end_id)))
for message_id in range(start_id, end_id + 1):
entry = self.messages.get(message_id) if channel_id == self.channel_id else None
yield message_id, (entry[0] if entry else None)
await asyncio.sleep(0)
async def iter_media_chunks(self, message_id, *, offset, expected_size, channel_id=None):
assert channel_id in (None, self.channel_id)
self.download_offsets.append((message_id, offset))
self.events.append(("download", message_id))
content = self.messages[message_id][1]
plan = self.failures.get(message_id) or []
failure = plan.pop(0) if plan else None
sent = 0
for start in range(offset, len(content), self.CHUNK):
if isinstance(failure, tuple) and failure[0] == "after" and sent >= failure[1]:
raise failure[2]
if isinstance(failure, BaseException):
raise failure
chunk = content[start : start + self.CHUNK]
sent += len(chunk)
yield chunk
await asyncio.sleep(0)
class SimulatedCrash(BaseException):
"""Stands in for the process being killed (not caught by `except Exception`)."""
class FakeDrive:
CHUNK = 256 * 1024
def __init__(self):
self.files: dict[str, DriveFile] = {}
self.sessions: dict[str, dict] = {}
self._ids = itertools.count(1)
self._lock = threading.Lock()
self.active_uploads = 0
self.max_active_uploads = 0
self.upload_failures: list[BaseException] = []
self.crash_at_bytes: int | None = None
self.lose_response_once = False
self.resumed: list[tuple[str, int]] = []
self.upload_calls = 0
self.events_sink: list | None = None
self.created_folders = 0
def _id(self, prefix):
return f"{prefix}{next(self._ids)}"
def add_folder(self, parent, name):
folder = DriveFile(self._id("fold"), name, parents=(parent,), mime_type=FOLDER_MIME)
self.files[folder.id] = folder
return folder
def add_file(self, parent, name, content: bytes, app_properties=None):
f = DriveFile(
self._id("file"), name, size=len(content), md5=hashlib.md5(content).hexdigest(), parents=(parent,),
mime_type="video/mp4", app_properties=dict(app_properties or {}),
)
self.files[f.id] = f
return f
def uploaded(self):
return [f for f in self.files.values() if f.mime_type != FOLDER_MIME and not f.trashed]
# DriveClient interface
def get_file(self, file_id):
return self.files.get(file_id)
def find_folders(self, parent_id, name):
return [f for f in self.files.values() if f.mime_type == FOLDER_MIME and parent_id in f.parents and f.name == name and not f.trashed]
def create_folder(self, parent_id, name):
self.created_folders += 1
return self.add_folder(parent_id, name)
def find_by_name(self, parent_id, name):
return [f for f in self.uploaded() if parent_id in f.parents and f.name == name]
def find_by_source_key(self, key):
return [f for f in self.uploaded() if f.app_properties.get(SOURCE_PROPERTY) == key]
def share_with_link(self, file_id):
pass
def upload_file(self, path, *, name, parent_id, mime_type, size, app_properties, description=None,
session_uri, on_session, on_progress, should_stop):
with self._lock:
self.active_uploads += 1
self.max_active_uploads = max(self.max_active_uploads, self.active_uploads)
self.upload_calls += 1
if self.events_sink is not None:
self.events_sink.append(("upload", app_properties[SOURCE_PROPERTY]))
try:
data = Path(path).read_bytes()
assert len(data) == size
session = self.sessions.get(session_uri) if session_uri else None
if session is not None:
self.resumed.append((session_uri, len(session["data"])))
if session.get("file"):
return session["file"]
else:
session_uri = f"https://upload.example/session/{self._id('s')}"
session = {"data": b"", "name": name, "parent": parent_id, "props": app_properties}
self.sessions[session_uri] = session
on_session(session_uri)
if self.upload_failures:
raise self.upload_failures.pop(0)
while len(session["data"]) < size:
if should_stop():
raise Interrupted("upload paused")
start = len(session["data"])
session["data"] += data[start : start + self.CHUNK]
on_progress(len(session["data"]), size)
time.sleep(0.001)
if self.crash_at_bytes is not None and len(session["data"]) >= self.crash_at_bytes:
self.crash_at_bytes = None
raise SimulatedCrash()
f = self.add_file(session["parent"], session["name"], session["data"], session["props"])
session["file"] = f
if self.lose_response_once:
self.lose_response_once = False
from errors import TransientError
raise TransientError("connection reset after upload completed")
return f
finally:
with self._lock:
self.active_uploads -= 1
@pytest.fixture
def source():
return FakeSource()
@pytest.fixture
def drive():
return FakeDrive()
class RecordingNotifier:
def __init__(self):
self.sent: list[tuple[int, str]] = []
self.edits: list[tuple[int, int, str]] = []
self._ids = itertools.count(1)
async def send(self, chat_id, text):
self.sent.append((chat_id, text))
return next(self._ids)
async def edit(self, chat_id, message_id, text):
self.edits.append((chat_id, message_id, text))
def build_queue(config, db, source, drive, *, reporter=None, disk_usage=None, root="root-folder"):
from drive_client import FolderResolver
from queue_manager import QueueManager
kwargs = {"disk_usage": disk_usage} if disk_usage else {}
downloader = Downloader(
source, config.download_dir, min_free_bytes=config.min_free_disk_space,
max_temp_bytes=config.max_temp_storage, compute_sha256=config.compute_sha256, **kwargs,
)
if root not in drive.files:
drive.files[root] = DriveFile(root, "Root", mime_type=FOLDER_MIME)
resolver = FolderResolver(drive, db, root, sleep=lambda s: None)
return QueueManager(config, db, drive, downloader, resolver, source=source, reporter=reporter, db_error_delay=0)
def run(coro):
return asyncio.run(coro)
@pytest.fixture(autouse=True)
def _no_dotenv(monkeypatch):
# Never pick up a developer's real .env during tests.
monkeypatch.setattr("dotenv.load_dotenv", lambda *a, **k: False)
@pytest.fixture(autouse=True)
def _restore_logging():
import logging
root = logging.getLogger()
handlers, level = list(root.handlers), root.level
yield
root.handlers[:] = handlers
root.setLevel(level)
def pytest_sessionfinish(session, exitstatus):
if _PG_SERVER is not None:
try:
_PG_SERVER.cleanup()
except Exception:
pass
shutil.rmtree(_PG_SERVER.pgdata, ignore_errors=True)
tests/test_bot_handlers.py
+69
View File
@@ -0,0 +1,69 @@
import pytest
from bot_handlers import CommandRouter
from conftest import CHANNEL_ID, EXAMPLE_CAPTION, build_queue, run
from models import State
ADMIN, STRANGER = 111, 999
COMMANDS = ["/start", "/help", "/status", "/run", "/pause", "/resume", "/stop", "/retry", "/failed",
"/progress", "/folders", "/check 1", "/sync"]
@pytest.fixture
def router(config, db, source, drive):
source.add(1, EXAMPLE_CAPTION)
queue = build_queue(config, db, source, drive)
return CommandRouter(config, db, queue)
@pytest.mark.parametrize("command", COMMANDS)
def test_non_admin_is_rejected(router, command):
reply = run(router.handle(STRANGER, command))
assert reply == "You are not authorized to control this bot."
assert router.queue.running is False
assert router.db.counts()["PENDING"] == 0 # /sync did not run
def test_none_sender_is_rejected(router):
assert "not authorized" in run(router.handle(None, "/status"))
def test_non_commands_are_ignored(router):
assert run(router.handle(ADMIN, "hello")) is None
def test_admin_flow(router, drive):
async def flow():
replies = {}
replies["help"] = await router.handle(ADMIN, "/help")
replies["sync"] = await router.handle(ADMIN, "/sync@MyBot")
replies["check"] = await router.handle(ADMIN, "/check 1")
replies["run"] = await router.handle(ADMIN, "/run")
await router.queue.wait()
replies["progress"] = await router.handle(ADMIN, "/progress")
replies["folders"] = await router.handle(ADMIN, "/folders")
replies["status"] = await router.handle(ADMIN, "/status")
replies["failed"] = await router.handle(ADMIN, "/failed")
replies["unknown"] = await router.handle(ADMIN, "/nope")
replies["badcheck"] = await router.handle(ADMIN, "/check abc")
return replies
r = run(flow())
assert "/status" in r["help"]
assert "New videos queued: 1" in r["sync"]
assert "Class 01 Aman sir.mp4" in r["check"] and "001, 002" in r["check"]
assert r["run"].startswith("Processing started")
assert "Processed (uploaded & verified): 1" in r["progress"]
assert "Misc• Grammar Practice: 1" in r["folders"]
assert "Queue: idle" in r["status"]
assert r["failed"] == "No failed items."
assert "Unknown command" in r["unknown"]
assert r["badcheck"].startswith("Usage")
def test_failed_listing_hides_paths(router, db):
run(router.handle(ADMIN, "/sync"))
db.transition((CHANNEL_ID, 1), State.FAILED, last_error="cannot open /Users/someone/secret/dir/file.part")
reply = run(router.handle(ADMIN, "/failed"))
assert "/Users/someone" not in reply and "file.part" in reply
tests/test_caption_parser.py
+102
View File
@@ -0,0 +1,102 @@
from caption_parser import parse_caption
from conftest import EXAMPLE_CAPTION
def test_supplied_example():
p = parse_caption(EXAMPLE_CAPTION)
assert p.indexes == ["001", "002"]
assert p.title == "Class 01 Aman sir.mp4"
assert p.topic == "Misc• Grammar Practice"
assert p.batch == "English ( SSC MAINS )"
assert p.extracted_by == "@PUTINxINDIA"
assert p.missing == []
assert p.warnings == []
assert p.raw == EXAMPLE_CAPTION
def test_index_is_never_used_as_title():
p = parse_caption("Index: 007\nIndex: 008\nTopic: Nouns")
assert p.indexes == ["007", "008"]
assert p.title is None
assert "title" in p.missing
def test_multiple_index_order_preserved():
p = parse_caption("Index: 3\nTitle: A\nIndex: 1\nTopic: T\nIndex: 2")
assert p.indexes == ["3", "1", "2"]
def test_fields_in_any_order_and_whitespace_variants():
caption = " topic : Error Detection \n\n\nTITLE:Class 02.mp4\r\n batch: X \nExtracted by :@someone\n"
p = parse_caption(caption)
assert p.topic == "Error Detection"
assert p.title == "Class 02.mp4"
assert p.batch == "X"
assert p.extracted_by == "@someone"
def test_decorations_and_fullwidth_colon():
p = parse_caption("🎬 Title： Class 05.mp4\n📚 Topic: Reading Comprehension\n• Batch - ignored\n**Batch:** B1")
assert p.title == "Class 05.mp4"
assert p.topic == "Reading Comprehension"
assert p.batch == "** B1" # value text is preserved verbatim apart from trimming
def test_missing_title_and_topic():
p = parse_caption("Index: 001\nBatch: English")
assert p.missing == ["title", "topic"]
assert "Title field is missing" in p.warnings
assert "Topic field is missing" in p.warnings
def test_empty_and_none_caption():
for value in (None, "", " \n "):
p = parse_caption(value)
assert p.missing == ["title", "topic"]
assert "caption is empty" in p.warnings
def test_empty_field_values_are_reported():
p = parse_caption("Title:\nTopic: \nIndex:")
assert p.title is None and p.topic is None
assert "Title field is empty" in p.warnings
assert "empty Index field ignored" in p.warnings
def test_unicode_preserved():
p = parse_caption("Title: कक्षा 01 — व्याकरण (भाग १).mp4\nTopic: हिंदी • व्याकरण\nBatch: Ünïcödé ( 2024 )")
assert p.title == "कक्षा 01 — व्याकरण (भाग १).mp4"
assert p.topic == "हिंदी • व्याकरण"
assert p.batch == "Ünïcödé ( 2024 )"
def test_next_field_is_not_swallowed():
p = parse_caption("Title: Class 01\nTopic: Grammar\nBatch: SSC")
assert p.title == "Class 01"
assert p.topic == "Grammar"
def test_flattened_single_line_caption():
p = parse_caption("Index: 001 Title: Class 01 Aman sir.mp4 Topic: Misc• Grammar Practice | Batch: English ( SSC MAINS )")
assert p.indexes == ["001"]
assert p.title == "Class 01 Aman sir.mp4"
assert p.topic == "Misc• Grammar Practice"
assert p.batch == "English ( SSC MAINS )"
def test_duplicate_field_keeps_first_and_warns():
p = parse_caption("Title: First.mp4\nTitle: Second.mp4\nTopic: T")
assert p.title == "First.mp4"
assert any("duplicate Title" in w for w in p.warnings)
def test_unknown_fields_and_free_text_kept_in_extra_or_ignored():
p = parse_caption("Join: @channel\nSome free text line\nTitle: A.mp4\nTopic: T\nSubtitle: not a title")
assert p.title == "A.mp4"
assert p.extra["Join"] == "@channel"
assert p.extra["Subtitle"] == "not a title"
def test_label_must_be_whole_word():
p = parse_caption("Topics: wrong\nTopic: right")
assert p.topic == "right"
tests/test_config.py
+140
View File
@@ -0,0 +1,140 @@
import base64
import json
import logging
import pytest
from config import Config, ConfigError, load_json_secret
from conftest import make_config
from logging_config import RedactSecrets
def _rsa_pem() -> str:
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
return key.private_bytes(
serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()
).decode()
SA_INFO = {
"type": "service_account",
"project_id": "example-project",
"private_key_id": "abc123",
"client_email": "uploader@example-project.iam.gserviceaccount.com",
"client_id": "1",
"token_uri": "https://oauth2.googleapis.com/token",
}
TOKEN_INFO = {
"token": "ya29.placeholder-access-token",
"refresh_token": "1//placeholder-refresh-token",
"client_id": "placeholder.apps.googleusercontent.com",
"client_secret": "placeholder-client-secret",
"token_uri": "https://oauth2.googleapis.com/token",
"expiry": "2999-01-01T00:00:00Z",
}
def test_defaults_and_mapped_names(tmp_path):
c = make_config(tmp_path, MAX_RETRIES="4", RETRY_DELAY="15", DUPLICATE_POLICY="suffix_message_id",
BATCH_SUBFOLDERS_ENABLED="yes", GOOGLE_DRIVE_ROOT_FOLDER_ID="FOLDER123")
assert c.max_attempts == 5 and c.retry_base_delay == 15
assert c.duplicate_name_policy == "suffix_message_id"
assert c.folder_structure == "batch_topic"
assert c.drive_root_folder_id == "FOLDER123"
assert c.source_channel == "-1001234567890"
assert c.fallback_topic == "Uncategorized"
assert c.share_mode == "none" # never public by default
assert c.drive_scope == "drive.file"
def test_invalid_values_are_all_reported(tmp_path):
with pytest.raises(ConfigError) as err:
make_config(tmp_path, MAX_RETRIES="many", PROCESS_MODE="sometimes", UPLOAD_CHUNK_SIZE="1000", KEEP_LOCAL_FILES="maybe")
text = str(err.value)
for name in ("MAX_RETRIES", "PROCESS_MODE", "UPLOAD_CHUNK_SIZE", "KEEP_LOCAL_FILES"):
assert name in text
def test_missing_required_telegram_settings():
config = Config.from_env({})
with pytest.raises(ConfigError) as err:
config.validate("telegram")
text = str(err.value)
assert "TELEGRAM_API_ID" in text and "SOURCE_CHANNEL_ID" in text
def test_heroku_requires_durable_database_and_session_string(tmp_path):
c = make_config(tmp_path, DYNO="worker.1", TELEGRAM_MODE="user")
problems = c.database_problems() + c.telegram_problems() + c.drive_problems()
joined = " ".join(problems)
assert "DATABASE_URL" in joined
assert "TELEGRAM_SESSION_STRING" in joined
assert "GOOGLE_OAUTH_TOKEN_JSON" in joined
ok = make_config(
tmp_path, DYNO="worker.1", TELEGRAM_MODE="user", TELEGRAM_SESSION_STRING="1AaBbCc-placeholder",
DATABASE_URL="postgres://u:p@db.example.com:5432/d", GOOGLE_OAUTH_TOKEN_JSON=json.dumps(TOKEN_INFO),
)
ok.validate("telegram", "drive", "database")
assert ok.database_url == "postgresql://u:p@db.example.com:5432/d?sslmode=require"
assert ok.instance_lock_wait > 0
def test_ephemeral_database_must_be_explicit(tmp_path):
c = make_config(tmp_path, DYNO="worker.1", ALLOW_EPHEMERAL_DATABASE="true")
assert c.database_problems() == []
def test_json_secrets_raw_multiline_and_base64():
raw = json.dumps(TOKEN_INFO, indent=2)
assert load_json_secret(raw)["refresh_token"] == TOKEN_INFO["refresh_token"]
encoded = base64.b64encode(raw.encode()).decode()
assert load_json_secret(encoded)["client_id"] == TOKEN_INFO["client_id"]
with pytest.raises(ValueError):
load_json_secret("{not json")
def test_bad_credential_json_reported(tmp_path):
c = make_config(tmp_path, GOOGLE_AUTH_MODE="service_account", GOOGLE_SERVICE_ACCOUNT_JSON='{"type": "service_account"}')
assert any("missing keys" in p for p in c.drive_problems())
def test_credentials_load_from_env_without_files(tmp_path):
from drive_client import load_credentials
sa = dict(SA_INFO, private_key=_rsa_pem())
c = make_config(tmp_path, GOOGLE_AUTH_MODE="service_account", GOOGLE_SERVICE_ACCOUNT_JSON=json.dumps(sa))
assert not c.google_token_file.exists() and not c.google_client_secrets_file.exists()
c.validate("drive")
creds = load_credentials(c)
assert creds.service_account_email == SA_INFO["client_email"]
assert creds.scopes == ["https://www.googleapis.com/auth/drive.file"]
c = make_config(tmp_path, GOOGLE_OAUTH_TOKEN_JSON=base64.b64encode(json.dumps(TOKEN_INFO).encode()).decode())
c.validate("drive")
creds = load_credentials(c)
assert creds.refresh_token == TOKEN_INFO["refresh_token"] and creds.valid
assert not c.google_token_file.exists() # nothing written to the ephemeral disk
def test_secrets_are_redacted_from_logs(tmp_path, caplog):
sa = dict(SA_INFO, private_key="-----BEGIN PRIVATE KEY-----placeholder")
c = make_config(tmp_path, GOOGLE_SERVICE_ACCOUNT_JSON=json.dumps(sa), GOOGLE_OAUTH_TOKEN_JSON=json.dumps(TOKEN_INFO))
secrets = c.secrets()
assert c.bot_token in secrets and c.api_hash in secrets
assert TOKEN_INFO["refresh_token"] in secrets and TOKEN_INFO["client_secret"] in secrets
record = logging.LogRecord("x", logging.INFO, __file__, 1, "token=%s refresh=%s", (c.bot_token, TOKEN_INFO["refresh_token"]), None)
RedactSecrets(secrets).filter(record)
assert c.bot_token not in record.getMessage()
assert TOKEN_INFO["refresh_token"] not in record.getMessage()
def test_main_exits_with_clear_error_when_unconfigured(monkeypatch, tmp_path):
import main
monkeypatch.setenv("LOG_DIR", str(tmp_path / "logs"))
for name in ("TELEGRAM_API_ID", "TELEGRAM_API_HASH", "TELEGRAM_BOT_TOKEN", "SOURCE_CHANNEL_ID", "DYNO", "DATABASE_URL"):
monkeypatch.delenv(name, raising=False)
assert main.main(["check-config"]) == 2
tests/test_database.py
+155
View File
@@ -0,0 +1,155 @@
from datetime import timedelta
import pytest
from caption_parser import parse_caption
from conftest import CHANNEL_ID, EXAMPLE_CAPTION, FakeSource
from database import Database, InstanceLock, InstanceLockError, InvalidTransition, utcnow
from models import State
def _insert(db, message_id, caption=EXAMPLE_CAPTION):
msg = FakeSource().add(message_id, caption)
return msg, db.insert_video(msg, parse_caption(caption))
def test_insert_and_metadata_roundtrip(db):
msg, inserted = _insert(db, 5)
assert inserted
rec = db.get((CHANNEL_ID, 5))
assert rec.state == State.PENDING
assert rec.indexes == ["001", "002"]
assert rec.title == "Class 01 Aman sir.mp4"
assert rec.topic == "Misc• Grammar Practice"
assert rec.batch == "English ( SSC MAINS )"
assert rec.extracted_by == "@PUTINxINDIA"
assert rec.caption == EXAMPLE_CAPTION
assert rec.expected_size == msg.media.size
assert rec.download_status == "pending"
def test_duplicate_message_is_not_inserted_twice(db):
_, first = _insert(db, 7)
_, second = _insert(db, 7)
assert first and not second
assert db.counts()["PENDING"] == 1
def test_valid_state_transitions_and_history(db):
_insert(db, 1)
key = (CHANNEL_ID, 1)
for state in (State.DOWNLOADING, State.DOWNLOADED, State.UPLOADING, State.COMPLETED):
rec = db.transition(key, state)
assert rec.state == state
assert rec.download_status == "done" and rec.upload_status == "done"
assert [h[1] for h in db.history(key)] == ["PENDING", "DOWNLOADING", "DOWNLOADED", "UPLOADING", "COMPLETED"]
@pytest.mark.parametrize(
"path",
[
[State.COMPLETED, State.PENDING], # completed is terminal
[State.UPLOADING], # cannot upload before downloading
[State.DOWNLOADING, State.COMPLETED], # cannot complete without upload/reconcile
],
)
def test_invalid_transitions_rejected(db, path):
_insert(db, 1)
key = (CHANNEL_ID, 1)
with pytest.raises(InvalidTransition):
for state in path:
db.transition(key, state)
def test_failure_status_fields(db):
_insert(db, 1)
key = (CHANNEL_ID, 1)
db.transition(key, State.DOWNLOADING)
rec = db.transition(key, State.RETRY_WAIT, attempts=1, last_error="boom")
assert rec.download_status == "retry_wait" and rec.attempts == 1 and rec.last_error == "boom"
db.transition(key, State.DOWNLOADED)
db.transition(key, State.UPLOADING)
rec = db.transition(key, State.FAILED)
assert rec.upload_status == "failed"
def test_unknown_field_rejected(db):
_insert(db, 1)
with pytest.raises(ValueError):
db.update_fields((CHANNEL_ID, 1), message_id=3)
def test_recover_stale(db):
for i in (1, 2, 3):
_insert(db, i)
db.transition((CHANNEL_ID, 1), State.DOWNLOADING)
db.transition((CHANNEL_ID, 2), State.DOWNLOADING)
db.transition((CHANNEL_ID, 2), State.DOWNLOADED)
db.transition((CHANNEL_ID, 2), State.UPLOADING)
recovered = db.recover_stale()
assert recovered == {"DOWNLOADING": 1, "UPLOADING": 1}
assert db.get((CHANNEL_ID, 1)).state == State.PENDING
assert db.get((CHANNEL_ID, 2)).state == State.UPLOADING # reconciled by the queue
def test_next_item_prefers_unfinished_then_oldest(db):
for i in (10, 20, 30):
_insert(db, i)
assert db.next_item().message_id == 10
db.transition((CHANNEL_ID, 30), State.DOWNLOADING)
db.transition((CHANNEL_ID, 30), State.DOWNLOADED)
assert db.next_item().message_id == 30
db.transition((CHANNEL_ID, 30), State.RETRY_WAIT, next_retry_at=(utcnow() + timedelta(hours=1)).isoformat(timespec="seconds"))
assert db.next_item().message_id == 10
assert db.next_item(utcnow() + timedelta(hours=2)).message_id == 10
assert db.earliest_retry_at() is not None
def test_folders_kv_sequence(db):
assert db.get_folder("root", "misc") is None
db.put_folder("root", "misc", "Misc", "F1")
db.put_folder("root", "misc", "Misc", "F2")
assert db.get_folder("root", "misc") == "F2"
assert db.folder_count() == 1
db.raise_int("cursor", 5)
db.raise_int("cursor", 3)
assert db.get_int("cursor") == 5
assert [db.next_sequence() for _ in range(3)] == [1, 2, 3]
def test_state_survives_reopen(config, db):
_insert(db, 1)
db.set_kv("queue_desired_state", "running")
db.close()
again = Database.from_config(config)
assert again.get((CHANNEL_ID, 1)).title == "Class 01 Aman sir.mp4"
assert again.get_kv("queue_desired_state") == "running"
again.close()
def test_instance_lock_blocks_second_worker(config, db):
first = InstanceLock(config)
first.acquire()
second = InstanceLock(config, wait_seconds=0)
try:
if config.database_url:
with pytest.raises(InstanceLockError):
second.acquire()
else:
# flock is per open file description, so test from a second process
import subprocess
import sys
code = (
"import fcntl,sys; f=open(sys.argv[1],'w');\n"
"try:\n fcntl.flock(f, fcntl.LOCK_EX|fcntl.LOCK_NB); print('acquired')\n"
"except BlockingIOError: print('blocked')"
)
out = subprocess.run([sys.executable, "-c", code, str(config.database_path) + ".lock"], capture_output=True, text=True)
assert out.stdout.strip() == "blocked"
finally:
first.release()
second.release()
third = InstanceLock(config)
third.acquire()
third.release()
tests/test_deployment.py
+77
View File
@@ -0,0 +1,77 @@
"""Static deployment-readiness checks (these do NOT deploy anything to Heroku)."""
import json
import re
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
# Top-level keys documented for app.json (https://devcenter.heroku.com/articles/app-json-schema)
APP_JSON_KEYS = {
"name", "description", "keywords", "website", "repository", "logo", "success_url", "scripts", "env",
"formation", "image", "addons", "buildpacks", "environments", "stack",
}
ENV_KEYS = {"description", "value", "required", "generator"}
def _config_vars() -> set[str]:
source = (ROOT / "config.py").read_text()
return set(re.findall(r'r\.\w+\(\s*"([A-Z0-9_]+)"', source)) | set(re.findall(r'env\.get\("([A-Z0-9_]+)"\)', source))
def _app_json() -> dict:
return json.loads((ROOT / "app.json").read_text())
def test_app_json_structure():
manifest = _app_json()
assert set(manifest) <= APP_JSON_KEYS
assert len(manifest["name"]) <= 30 and manifest["description"]
assert manifest["buildpacks"] == [{"url": "heroku/python"}]
assert manifest["formation"] == {"worker": {"quantity": 1, "size": "basic"}}
assert manifest["addons"][0]["plan"].startswith("heroku-postgresql:")
for name, spec in manifest["env"].items():
assert set(spec) <= ENV_KEYS, name
assert spec["description"], name
assert isinstance(spec.get("required", False), bool), name
if "value" in spec:
assert isinstance(spec["value"], str), name
def test_app_json_only_lists_supported_variables():
supported = _config_vars()
unsupported = set(_app_json()["env"]) - supported
assert not unsupported, f"app.json lists variables the app does not read: {unsupported}"
def test_env_example_documents_every_variable():
documented = set(re.findall(r"^([A-Z0-9_]+)=", (ROOT / ".env.example").read_text(), re.MULTILINE))
supported = _config_vars() - {"DYNO"} # DYNO is set by Heroku itself
assert supported - documented == set()
assert documented - supported == set()
def test_procfile_matches_entry_point():
procfile = (ROOT / "Procfile").read_text().strip().splitlines()
assert procfile == ["worker: python main.py run"]
assert set(_app_json()["formation"]) == {line.split(":")[0] for line in procfile}
main_source = (ROOT / "main.py").read_text()
assert 'sub.add_parser("run"' in main_source and "signal.SIGTERM" in main_source
def test_runtime_and_requirements():
assert (ROOT / ".python-version").read_text().strip() == "3.12"
assert not (ROOT / "runtime.txt").exists() # deprecated in favour of .python-version
reqs = (ROOT / "requirements.txt").read_text()
for package in ("Telethon", "google-api-python-client", "google-auth", "google-auth-oauthlib", "psycopg[binary]", "python-dotenv"):
assert re.search(rf"^{re.escape(package)}==", reqs, re.MULTILINE), package
def test_gitignore_protects_secrets_and_state():
ignored = (ROOT / ".gitignore").read_text().splitlines()
for pattern in (".env", "*.session", "credentials/", "token.json", "session_string.txt", "data/", "downloads/", "*.db"):
assert pattern in ignored, pattern
def test_readme_deploy_button():
readme = (ROOT / "README.md").read_text()
assert "https://www.herokucdn.com/deploy/button.svg" in readme
assert re.search(r"\]\(https://heroku\.com/deploy\?template=https://github\.com/[^)\s]+\)", readme)
tests/test_drive_client.py
+305
View File
@@ -0,0 +1,305 @@
import hashlib
import json
import re
import httplib2
import pytest
from googleapiclient.errors import HttpError
from conftest import FakeDrive
from drive_client import (
DriveClient,
FolderResolver,
ResumableUploader,
classify_http_error,
escape_query,
parse_range,
)
from errors import AuthError, DriveQuotaError, Interrupted, PermanentError, TransientError, UploadSessionExpired
CHUNK = 256 * 1024
class Resp:
def __init__(self, status, headers=None, body=None):
self.status_code = status
self.headers = headers or {}
self._body = body
def json(self):
if self._body is None:
raise ValueError("no body")
return self._body
class FakeUploadServer:
"""Implements Google's resumable upload protocol in memory."""
def __init__(self):
self.sessions = {}
self.fail_next_put = [] # exceptions or Resp objects
self.initiations = 0
self.put_count = 0
def request(self, method, url, data=None, headers=None, timeout=None):
headers = headers or {}
if method == "POST":
self.initiations += 1
uri = f"https://upload/session/{self.initiations}"
meta = json.loads(data)
self.sessions[uri] = {"data": b"", "size": int(headers["X-Upload-Content-Length"]), "meta": meta, "done": None}
return Resp(200, {"Location": uri})
session = self.sessions.get(url)
if session is None:
return Resp(404)
content_range = headers["Content-Range"]
if content_range.startswith("bytes */"):
return self._status(session)
self.put_count += 1
if self.fail_next_put:
failure = self.fail_next_put.pop(0)
if isinstance(failure, BaseException):
# simulate a connection drop after the server stored half the chunk
session["data"] += data[: len(data) // 2]
raise failure
return failure
start = int(re.match(r"bytes (\d+)-", content_range).group(1))
assert start == len(session["data"]), "client must continue exactly where the server is"
session["data"] += data
return self._status(session)
def _status(self, session):
if len(session["data"]) >= session["size"]:
if session["done"] is None:
data = session["data"]
session["done"] = {
"id": "file-1",
"name": session["meta"]["name"],
"size": str(len(data)),
"md5Checksum": hashlib.md5(data).hexdigest(),
"parents": session["meta"]["parents"],
}
return Resp(200, body=session["done"])
headers = {"Range": f"bytes=0-{len(session['data']) - 1}"} if session["data"] else {}
return Resp(308, headers)
def _upload(server, path, *, session_uri=None, stop_after=None, retries=3):
sessions, progress = [], []
calls = {"n": 0}
def should_stop():
calls["n"] += 1
return stop_after is not None and calls["n"] > stop_after
uploader = ResumableUploader(server, chunk_size=CHUNK, retries=retries, sleep=lambda s: None)
result = uploader.upload(
path,
metadata={"name": "Class 01.mp4", "parents": ["F"], "appProperties": {"tg_source": "1:2"}},
mime_type="video/mp4",
size=path.stat().st_size,
session_uri=session_uri,
on_session=sessions.append,
on_progress=lambda d, t: progress.append(d),
should_stop=should_stop,
)
return result, sessions, progress
@pytest.fixture
def big_file(tmp_path):
path = tmp_path / "video.bin"
path.write_bytes(bytes(range(256)) * (CHUNK * 3 // 256 + 100)) # 3+ chunks
return path
def test_full_upload_in_chunks(big_file):
server = FakeUploadServer()
result, sessions, progress = _upload(server, big_file)
assert result["size"] == str(big_file.stat().st_size)
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
assert len(sessions) == 1 and server.put_count == 4
assert progress[-1] == big_file.stat().st_size
def test_pause_then_resume_same_session(big_file):
server = FakeUploadServer()
with pytest.raises(Interrupted):
_upload(server, big_file, stop_after=2)
uri = next(iter(server.sessions))
stored = len(server.sessions[uri]["data"])
assert 0 < stored < big_file.stat().st_size
result, sessions, _ = _upload(server, big_file, session_uri=uri)
assert sessions == [] # no new session was started
assert server.initiations == 1
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
def test_network_error_mid_chunk_resumes_from_server_offset(big_file):
server = FakeUploadServer()
server.fail_next_put = [ConnectionError("reset")]
result, _, _ = _upload(server, big_file)
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
def test_server_errors_are_retried_then_raised(big_file):
server = FakeUploadServer()
server.fail_next_put = [Resp(503), Resp(500)]
result, _, _ = _upload(server, big_file)
assert result["id"] == "file-1"
server = FakeUploadServer()
server.fail_next_put = [Resp(503)] * 10
with pytest.raises(TransientError):
_upload(server, big_file, retries=2)
def test_expired_session(big_file):
server = FakeUploadServer()
# unknown saved session -> a new session is created transparently
result, sessions, _ = _upload(server, big_file, session_uri="https://upload/session/expired")
assert len(sessions) == 1 and result["id"] == "file-1"
server = FakeUploadServer()
server.fail_next_put = [Resp(404)]
with pytest.raises(UploadSessionExpired):
_upload(server, big_file)
def test_quota_and_permission_errors_are_not_retried(big_file):
server = FakeUploadServer()
server.fail_next_put = [Resp(403, body={"error": {"errors": [{"reason": "storageQuotaExceeded"}], "message": "full"}})]
with pytest.raises(DriveQuotaError):
_upload(server, big_file)
server = FakeUploadServer()
server.fail_next_put = [Resp(403, body={"error": {"errors": [{"reason": "insufficientFilePermissions"}]}})]
with pytest.raises(PermanentError):
_upload(server, big_file)
def test_chunk_size_must_be_aligned():
with pytest.raises(ValueError):
ResumableUploader(None, chunk_size=1000)
def test_helpers():
assert escape_query("It's a \\ test") == "It\\'s a \\\\ test"
assert parse_range("bytes=0-1048575") == 1048576
assert parse_range(None) == 0
assert isinstance(classify_http_error(401), AuthError)
assert isinstance(classify_http_error(429), TransientError)
rate = classify_http_error(403, "userRateLimitExceeded")
assert isinstance(rate, TransientError) and rate.rate_limited
assert isinstance(classify_http_error(404), PermanentError)
assert isinstance(classify_http_error(502), TransientError)
# -- DriveClient retry wrapper ---------------------------------------------------------
class FakeRequest:
def __init__(self, outcome):
self.outcome = outcome
def execute(self):
if isinstance(self.outcome, BaseException):
raise self.outcome
return self.outcome
def _http_error(status, reason=None):
body = json.dumps({"error": {"errors": [{"reason": reason}] if reason else [], "message": "x"}}).encode()
return HttpError(httplib2.Response({"status": status}), body)
def test_execute_retries_transient_then_succeeds():
outcomes = [_http_error(503), OSError("net down"), {"ok": True}]
client = DriveClient(None, None, chunk_size=CHUNK, sleep=lambda s: None)
assert client._execute(lambda: FakeRequest(outcomes.pop(0))) == {"ok": True}
def test_execute_does_not_retry_permission_errors():
calls = []
client = DriveClient(None, None, chunk_size=CHUNK, sleep=lambda s: None)
def build():
calls.append(1)
return FakeRequest(_http_error(403, "insufficientPermissions"))
with pytest.raises(PermanentError):
client._execute(build)
assert len(calls) == 1
assert client._execute(lambda: FakeRequest(_http_error(404)), not_found_ok=True) is None
# -- FolderResolver ----------------------------------------------------------------------
class MemoryFolderDB:
def __init__(self):
self.folders = {}
def get_folder(self, parent, key):
return self.folders.get((parent, key))
def put_folder(self, parent, key, name, folder_id):
self.folders[(parent, key)] = folder_id
def delete_folder(self, parent, key):
self.folders.pop((parent, key), None)
def _resolver(drive, db=None):
return FolderResolver(drive, db or MemoryFolderDB(), "root", sleep=lambda s: None)
def test_folder_created_once_and_reused():
drive = FakeDrive()
resolver = _resolver(drive)
first = resolver.resolve(["Misc• Grammar Practice"])
second = resolver.resolve(["Misc• Grammar Practice"])
assert first == second and drive.created_folders == 1
# a fresh process (empty in-memory verification) still reuses the DB mapping
db = resolver._db
again = FolderResolver(drive, db, "root").resolve(["Misc• Grammar Practice"])
assert again == first and drive.created_folders == 1
def test_existing_drive_folder_is_reused_and_duplicates_pick_oldest():
drive = FakeDrive()
oldest = drive.add_folder("root", "Error Detection")
drive.add_folder("root", "Error Detection")
resolver = _resolver(drive)
assert resolver.resolve(["Error Detection"]) == oldest.id
assert drive.created_folders == 0
def test_deleted_cached_folder_is_recreated():
drive = FakeDrive()
resolver = _resolver(drive)
folder_id = resolver.resolve(["Topic"])
drive.files[folder_id].trashed = True
fresh = FolderResolver(drive, resolver._db, "root")
new_id = fresh.resolve(["Topic"])
assert new_id != folder_id and drive.created_folders == 2
def test_lost_create_response_does_not_duplicate_folder():
drive = FakeDrive()
real_create = drive.create_folder
def flaky_create(parent, name):
real_create(parent, name) # Drive created it...
raise TransientError("...but the response was lost")
drive.create_folder = flaky_create
resolver = _resolver(drive)
folder_id = resolver.resolve(["Reading Comprehension"])
assert len(drive.find_folders("root", "Reading Comprehension")) == 1
assert drive.files[folder_id].name == "Reading Comprehension"
def test_nested_batch_topic_path():
drive = FakeDrive()
resolver = _resolver(drive)
leaf = resolver.resolve(["English ( SSC MAINS )", "Misc• Grammar Practice"])
batch = drive.find_folders("root", "English ( SSC MAINS )")[0]
assert drive.files[leaf].parents == (batch.id,)
tests/test_link_parser.py
+66
View File
@@ -0,0 +1,66 @@
import pytest
from link_parser import LinkError, internal_channel_id, marked_channel_id, parse_message_link
@pytest.mark.parametrize(
"url",
[
"https://t.me/c/3909337424/14",
"https://t.me/c/3909337424/14/",
"https://t.me/c/3909337424/14?single",
"https://www.t.me/c/3909337424/14",
"http://t.me/c/3909337424/14#frag",
"t.me/c/3909337424/14",
"https://telegram.me/c/3909337424/14",
"https://t.me/c/3909337424/7/14", # thread segment
" https://t.me/c/3909337424/14 ",
],
)
def test_private_links(url):
link = parse_message_link(url)
assert link.kind == "private"
assert link.channel == "3909337424"
assert link.message_id == 14
assert link.channel_ref == "-1003909337424"
@pytest.mark.parametrize("url", ["https://t.me/my_channel/14", "https://t.me/s/my_channel/14", "t.me/My_Channel/14?single"])
def test_public_links(url):
link = parse_message_link(url)
assert link.kind == "public" and link.message_id == 14
assert link.channel.lower() == "my_channel"
assert link.channel_ref.lower() == "my_channel"
@pytest.mark.parametrize(
"text,fragment",
[
("", "send a Telegram message link"),
("hello", "not a Telegram"),
("https://example.com/c/1/2", "not a Telegram"),
("https://t.me/c/3909337424", "must look like"),
("https://t.me/c/abc/14", "must look like"),
("https://t.me/c/3909337424/0", "positive"),
("https://t.me/c/3909337424/x14", "must look like"),
("https://t.me/+AbCdEf123", "Invite links"),
("https://t.me/joinchat/AbCdEf", "Invite links"),
("https://t.me/my_channel", "not to a message"),
("https://t.me/addstickers/pack/1", "Unsupported"),
("https://t.me/ab/14", "Unsupported"),
("ftp://t.me/c/1/2", "not a Telegram"),
("https://t.me/c/3909337424/14 https://t.me/c/3909337424/15", "only the link"),
],
)
def test_rejected_links(text, fragment):
with pytest.raises(LinkError) as err:
parse_message_link(text)
assert fragment.lower() in str(err.value).lower()
def test_marked_id_follows_telegram_convention():
# -100 prefix is only a coincidence of 10-digit IDs; the real rule is -(10**12 + id)
assert marked_channel_id(3909337424) == -1003909337424
assert marked_channel_id(12345) == -1000000012345
assert internal_channel_id(-1000000012345) == 12345
assert parse_message_link("https://t.me/c/12345/9").channel_ref == "-1000000012345"
tests/test_naming.py
+84
View File
@@ -0,0 +1,84 @@
import pytest
from conftest import video_bytes
from errors import PermanentError
from naming import add_suffix, build_filename, clean_component, detect_video_extension, normalize_key, same_family
def _name(title, **kw):
defaults = dict(message_id=42, original_filename=None, mime_type="video/mp4", detected_ext=".mp4")
defaults.update(kw)
return build_filename(title, **defaults)
def test_title_used_verbatim():
assert _name("Class 01 Aman sir.mp4") == ("Class 01 Aman sir.mp4", [])
def test_extension_added_when_missing():
assert _name("Class 01 Aman sir")[0] == "Class 01 Aman sir.mp4"
# '.5' is not a video extension, so it stays part of the name
assert _name("Lecture 1.5", detected_ext=".mkv")[0] == "Lecture 1.5.mkv"
def test_extension_from_mime_when_undetected():
assert _name("Clip", detected_ext=None, mime_type="video/x-matroska")[0] == "Clip.mkv"
def test_compatible_extension_kept():
assert _name("Talk.mov", detected_ext=".mp4") == ("Talk.mov", [])
def test_mismatch_kept_by_default_but_reported():
name, warnings = _name("Talk.mp4", detected_ext=".mkv")
assert name == "Talk.mp4"
assert any("does not match" in w for w in warnings)
def test_mismatch_append_policy():
assert _name("Talk.mp4", detected_ext=".mkv", mismatch_policy="append")[0] == "Talk.mp4.mkv"
def test_missing_title_fallbacks():
name, warnings = _name(None, original_filename="orig_video.mp4")
assert name == "orig_video.mp4" and warnings
name, warnings = _name(" ", original_filename=None)
assert name == "telegram_42.mp4" and warnings
def test_missing_title_fail_policy():
with pytest.raises(PermanentError):
_name(None, missing_title_policy="fail")
def test_sanitizing():
assert clean_component("a/b\\c\x00d\n e") == "a_b_c d e"
assert clean_component("Misc• Grammar Practice") == "Misc• Grammar Practice"
assert clean_component('Q: what? "x" <y>|z*', strict=True) == "Q_ what_ _x_ _y__z_"
assert len(clean_component("x" * 500 + ".mp4")) == 200
assert clean_component("x" * 500 + ".mp4").endswith(".mp4")
def test_add_suffix():
assert add_suffix("Class 01 Aman sir.mp4", "(2)") == "Class 01 Aman sir (2).mp4"
assert add_suffix("Class 01", "[77]") == "Class 01 [77]"
def test_detect_container(tmp_path):
cases = {
".mp4": video_bytes(1000),
".mkv": b"\x1a\x45\xdf\xa3" + b"\x00" * 30 + b"matroska" + b"\x00" * 100,
".webm": b"\x1a\x45\xdf\xa3" + b"\x00" * 30 + b"webm" + b"\x00" * 100,
".avi": b"RIFF\x00\x00\x00\x00AVI LIST",
None: b"hello world, not a video",
}
for expected, data in cases.items():
path = tmp_path / "f"
path.write_bytes(data)
assert detect_video_extension(path) == expected
def test_families_and_keys():
assert same_family(".MP4", ".mov")
assert not same_family(".mp4", ".mkv")
assert normalize_key(" Misc• Grammar Practice ") == normalize_key("misc• grammar practice")
tests/test_progress_and_telegram.py
+107
View File
@@ -0,0 +1,107 @@
import asyncio
from telethon import errors
from telethon.tl import types
from conftest import RecordingNotifier
from errors import PermanentError, TransientError
from progress import ProgressReporter, TransferProgress, format_bytes, format_duration
from telegram_client import classify_media, map_telegram_error, parse_channel_reference, to_source_message
# -- progress ------------------------------------------------------------------------------
def test_formatting():
assert format_bytes(512) == "512 B"
assert format_bytes(1536) == "1.5 KB"
assert format_bytes(5 * 1024**3) == "5.0 GB"
assert format_duration(42) == "42s"
assert format_duration(125) == "2m 5s"
assert format_duration(3 * 3600 + 120) == "3h 2m"
def test_speed_and_eta():
now = [0.0]
p = TransferProgress(total=100 * 1024 * 1024, clock=lambda: now[0])
p.update(0)
now[0] = 10.0
p.update(50 * 1024 * 1024)
assert round(p.speed) == 5 * 1024 * 1024
assert round(p.eta) == 10
assert p.percent == 50
text = p.describe()
assert "50%" in text and "MB/s" in text and "ETA 10s" in text
def test_reporter_edits_single_message():
notifier = RecordingNotifier()
reporter = ProgressReporter(notifier, [1, 2], interval=60)
async def go():
reporter.board.begin_item(12, 150, "Misc• Grammar Practice", "Class 01 Aman sir.mp4")
reporter.board.start_download(1000)
reporter.board.download.update(640)
await reporter.refresh()
await reporter.refresh() # unchanged -> nothing sent
reporter.board.download.update(900)
await reporter.refresh()
asyncio.run(go())
assert len(notifier.sent) == 2 # one per chat
assert len(notifier.edits) == 2 # then edited
first = notifier.sent[0][1]
assert first.splitlines()[0] == "Processing video 12 of 150"
assert "Topic: Misc• Grammar Practice" in first and "Download: 64%" in first and "Upload: Waiting" in first
# -- telegram media classification ----------------------------------------------------------
def _doc(attributes, mime="video/mp4"):
return types.Document(id=77, access_hash=1, file_reference=b"", date=None, mime_type=mime, size=1234, dc_id=1,
attributes=attributes)
def _msg(media, text="Title: A.mp4"):
return types.Message(id=5, peer_id=types.PeerChannel(1), date=None, message=text, media=media)
def test_classify_video_message():
video = types.DocumentAttributeVideo(duration=60, w=1280, h=720)
info = classify_media(_msg(types.MessageMediaDocument(document=_doc([video]))), include_documents=False)
assert info.kind == "video" and info.size == 1234 and info.media_id == 77
sm = to_source_message(_msg(types.MessageMediaDocument(document=_doc([video]))), -100, False)
assert sm.caption == "Title: A.mp4" and sm.message_id == 5
def test_classify_document_video_only_when_enabled():
doc = _doc([types.DocumentAttributeFilename("lecture.mkv")], mime="application/octet-stream")
message = _msg(types.MessageMediaDocument(document=doc))
assert classify_media(message, include_documents=False) is None
assert classify_media(message, include_documents=True).file_name == "lecture.mkv"
def test_unsupported_media_is_ignored():
gif = _doc([types.DocumentAttributeVideo(duration=1, w=1, h=1), types.DocumentAttributeAnimated()])
note = _doc([types.DocumentAttributeVideo(duration=1, w=1, h=1, round_message=True)])
pdf = _doc([types.DocumentAttributeFilename("notes.pdf")], mime="application/pdf")
for doc in (gif, note, pdf):
assert classify_media(_msg(types.MessageMediaDocument(document=doc)), include_documents=True) is None
photo = types.MessageMediaPhoto(photo=types.PhotoEmpty(id=1))
assert classify_media(_msg(photo), include_documents=True) is None
assert classify_media(_msg(None, text="text only"), include_documents=True) is None
def test_channel_reference_parsing():
assert parse_channel_reference("-1001234567890") == -1001234567890
assert parse_channel_reference("@my_channel") == "my_channel"
assert parse_channel_reference("https://t.me/my_channel") == "my_channel"
assert parse_channel_reference("https://t.me/c/1234567890/55") == -1001234567890
def test_telegram_error_mapping():
flood = map_telegram_error(errors.FloodWaitError(request=None, capture=42))
assert isinstance(flood, TransientError) and flood.retry_after == 42 and flood.rate_limited
assert isinstance(map_telegram_error(errors.ChannelPrivateError(request=None)), PermanentError)
assert isinstance(map_telegram_error(ConnectionError("x")), TransientError)
tests/test_queue_manager.py
+438
View File
@@ -0,0 +1,438 @@
import asyncio
import hashlib
import shutil
import pytest
from conftest import (
CHANNEL_ID,
EXAMPLE_CAPTION,
FakeDrive,
FakeSource,
RecordingNotifier,
SimulatedCrash,
build_queue,
make_config,
run,
video_bytes,
)
from database import Database
from drive_client import SOURCE_PROPERTY
from errors import PermanentError, TransientError
from models import State
from progress import ProgressReporter
def caption(title, topic, batch="English ( SSC MAINS )", index="001"):
return f"Index: {index}\nTitle: {title}\nTopic: {topic}\nBatch: {batch}\nExtracted By: @someone"
async def sync_and_run(queue):
await queue.sync()
queue.start()
await queue.wait()
def key(message_id):
return (CHANNEL_ID, message_id)
def drive_names(drive, folder_name=None):
names = []
for f in drive.uploaded():
parent = drive.files[f.parents[0]].name
if folder_name is None or parent == folder_name:
names.append(f.name)
return sorted(names)
# -- end to end ---------------------------------------------------------------------------
def test_full_pipeline_topics_and_metadata(config, db, source, drive):
source.add(1, EXAMPLE_CAPTION)
source.add(2, caption("Class 02 Aman sir.mp4", "Misc• Grammar Practice"))
source.add(3, caption("Class 01 Aman sir.mp4", "Error Detection"))
source.add_text(4, "just an announcement")
queue = build_queue(config, db, source, drive)
run(sync_and_run(queue))
assert db.counts()["COMPLETED"] == 3
assert drive_names(drive, "Misc• Grammar Practice") == ["Class 01 Aman sir.mp4", "Class 02 Aman sir.mp4"]
assert drive_names(drive, "Error Detection") == ["Class 01 Aman sir.mp4"]
assert drive.created_folders == 2 # topic folder reused for the second video
rec = db.get(key(1))
assert rec.indexes == ["001", "002"]
assert rec.batch == "English ( SSC MAINS )" and rec.extracted_by == "@PUTINxINDIA"
assert rec.final_name == "Class 01 Aman sir.mp4"
uploaded = drive.files[rec.drive_file_id]
assert uploaded.app_properties[SOURCE_PROPERTY] == f"{CHANNEL_ID}:1"
assert uploaded.size == rec.local_size == rec.expected_size
assert uploaded.md5 == rec.md5
assert rec.sha256 and rec.completed_at
assert rec.download_status == "done" and rec.upload_status == "done"
# temporary files are removed only after completion
assert list(config.download_dir.iterdir()) == []
assert queue.stats.uploaded == 3
assert db.get_int("last_completed_message_id") == 3
assert db.get_int("last_discovered_message_id") == 4
def test_missing_topic_and_title_use_fallbacks(config, db, source, drive):
source.add(1, "Title: Lonely.mp4")
source.add(2, "Topic: Grammar", file_name="original_upload.mp4")
source.add(3, "")
queue = build_queue(config, db, source, drive)
run(sync_and_run(queue))
assert drive_names(drive, "Uncategorized") == ["Lonely.mp4", "telegram_3.mp4"]
assert drive_names(drive, "Grammar") == ["original_upload.mp4"]
assert any("Title missing" in w for w in db.get(key(2)).warnings)
def test_batch_subfolders(tmp_path, source, drive):
config = make_config(tmp_path, BATCH_SUBFOLDERS_ENABLED="true")
db = Database.from_config(config)
source.add(1, EXAMPLE_CAPTION)
run(sync_and_run(build_queue(config, db, source, drive)))
f = drive.uploaded()[0]
topic_folder = drive.files[f.parents[0]]
assert topic_folder.name == "Misc• Grammar Practice"
assert drive.files[topic_folder.parents[0]].name == "English ( SSC MAINS )"
# -- duplicates ----------------------------------------------------------------------------
def test_duplicate_message_detection(config, db, source, drive):
msg = source.add(1, EXAMPLE_CAPTION, media_id=999)
forwarded = source.add(2, EXAMPLE_CAPTION, video_bytes(700_000, seed=1), media_id=999)
queue = build_queue(config, db, source, drive)
assert queue.register(msg) == "new"
assert queue.register(msg) == "known"
assert queue.register(forwarded) == "skipped"
assert "duplicate of message 1" in db.get(key(2)).skip_reason
run(sync_and_run(queue))
assert len(drive.uploaded()) == 1
def test_filename_collision_gets_suffix(config, db, source, drive):
for i in (1, 2, 3):
source.add(i, caption("Class 01 Aman sir.mp4", "Misc• Grammar Practice"))
run(sync_and_run(build_queue(config, db, source, drive)))
assert drive_names(drive) == ["Class 01 Aman sir (2).mp4", "Class 01 Aman sir (3).mp4", "Class 01 Aman sir.mp4"]
def test_existing_different_file_is_never_overwritten(config, db, source, drive):
queue = build_queue(config, db, source, drive)
folder_id = queue.folders.resolve(["Topic"])
existing = drive.add_file(folder_id, "Class 01.mp4", b"someone else's video")
source.add(1, caption("Class 01.mp4", "Topic"))
run(sync_and_run(queue))
assert drive.files[existing.id].md5 == hashlib.md5(b"someone else's video").hexdigest()
assert drive_names(drive) == ["Class 01 (2).mp4", "Class 01.mp4"]
@pytest.mark.parametrize(
"policy,expected_state,expected_names",
[
("suffix_message_id", State.COMPLETED, ["X [2].mp4", "X.mp4"]),
("skip", State.SKIPPED, ["X.mp4"]),
("fail", State.FAILED, ["X.mp4"]),
],
)
def test_duplicate_policies(tmp_path, source, drive, policy, expected_state, expected_names):
config = make_config(tmp_path, DUPLICATE_POLICY=policy)
db = Database.from_config(config)
source.add(1, caption("X.mp4", "T"))
source.add(2, caption("X.mp4", "T"))
run(sync_and_run(build_queue(config, db, source, drive)))
assert db.get(key(2)).state == expected_state
assert drive_names(drive) == expected_names
def test_identical_existing_file_is_reused_not_duplicated(config, db, source, drive):
content = video_bytes(700_000, seed=5)
queue = build_queue(config, db, source, drive)
folder_id = queue.folders.resolve(["T"])
existing = drive.add_file(folder_id, "X.mp4", content)
source.add(1, caption("X.mp4", "T"), content)
run(sync_and_run(queue))
assert db.get(key(1)).drive_file_id == existing.id
assert len(drive.uploaded()) == 1 and drive.upload_calls == 0
# -- failures and recovery -----------------------------------------------------------------
def test_failed_download_is_retried_and_resumed(config, db, source, drive):
content = video_bytes(3 * 1024 * 1024 + 123, seed=9)
source.add(1, caption("Big.mp4", "T"), content)
source.failures[1] = [("after", 2 * 1024 * 1024, ConnectionError("network dropped"))]
queue = build_queue(config, db, source, drive)
run(sync_and_run(queue))
rec = db.get(key(1))
assert rec.state == State.COMPLETED and rec.attempts == 1
offsets = [o for m, o in source.download_offsets if m == 1]
assert offsets[0] == 0 and offsets[1] == 2 * 1024 * 1024 # resumed, not restarted
assert drive.uploaded()[0].md5 == hashlib.md5(content).hexdigest()
assert State.RETRY_WAIT.value in [h[1] for h in db.history(key(1))]
def test_failed_upload_is_retried(config, db, source, drive):
source.add(1, caption("A.mp4", "T"))
drive.upload_failures = [TransientError("Drive 503")]
run(sync_and_run(build_queue(config, db, source, drive)))
assert db.get(key(1)).state == State.COMPLETED
assert len(drive.uploaded()) == 1
def test_retries_are_bounded(config, db, source, drive):
source.add(1, caption("A.mp4", "T"))
source.add(2, caption("B.mp4", "T"))
drive.upload_failures = [TransientError("Drive 503")] * config.max_attempts
run(sync_and_run(build_queue(config, db, source, drive)))
rec = db.get(key(1))
assert rec.state == State.FAILED and rec.attempts == config.max_attempts
assert "Drive 503" in rec.last_error
assert db.get(key(2)).state == State.COMPLETED # the queue was not blocked
def test_permanent_error_fails_without_blocking_queue(config, db, source, drive):
source.add(1, caption("Gone.mp4", "T"))
source.add(2, caption("B.mp4", "T"))
source.failures[1] = [PermanentError("message deleted")]
run(sync_and_run(build_queue(config, db, source, drive)))
assert db.get(key(1)).state == State.FAILED and db.get(key(1)).attempts == 0
assert db.get(key(2)).state == State.COMPLETED
def test_strict_stop_halts_after_failure(tmp_path, source, drive):
config = make_config(tmp_path, STRICT_STOP="true")
db = Database.from_config(config)
source.add(1, caption("Gone.mp4", "T"))
source.add(2, caption("B.mp4", "T"))
source.failures[1] = [PermanentError("message deleted")]
queue = build_queue(config, db, source, drive)
run(sync_and_run(queue))
assert queue.status == "halted"
assert db.get(key(2)).state == State.PENDING
def test_retry_command_requeues_failed(config, db, source, drive):
source.add(1, caption("A.mp4", "T"))
source.failures[1] = [PermanentError("temporarily missing")]
queue = build_queue(config, db, source, drive)
run(sync_and_run(queue))
assert db.get(key(1)).state == State.FAILED
async def retry():
assert queue.retry_failed() == 1
queue.start()
await queue.wait()
run(retry())
assert db.get(key(1)).state == State.COMPLETED
def test_disk_full_halts_queue_without_using_attempts(config, db, source, drive):
source.add(1, caption("A.mp4", "T"))
queue = build_queue(config, db, source, drive, disk_usage=lambda p: (100, 100, 10))
run(sync_and_run(queue))
rec = db.get(key(1))
assert queue.status == "halted" and "disk space" in queue.halt_reason
assert rec.state == State.RETRY_WAIT and rec.attempts == 0
def test_restart_during_upload_resumes_session(config, source, drive):
source.add(1, caption("Big.mp4", "T"), video_bytes(2_000_000, seed=3))
db = Database.from_config(config)
queue = build_queue(config, db, source, drive)
drive.crash_at_bytes = 700_000
async def crash_run():
await queue.sync()
with pytest.raises(SimulatedCrash):
await queue.process_item(db.next_item())
run(crash_run())
rec = db.get(key(1))
assert rec.state == State.UPLOADING and rec.upload_session_uri
db.close()
# "new process"
db2 = Database.from_config(config)
db2.recover_stale()
queue2 = build_queue(config, db2, source, drive)
async def resume():
queue2.start()
await queue2.wait()
run(resume())
assert db2.get(key(1)).state == State.COMPLETED
assert drive.resumed and drive.resumed[0][1] >= 700_000
assert len(drive.uploaded()) == 1
assert len([1 for m, _ in source.download_offsets if m == 1]) == 1 # local file reused, no re-download
def test_restart_with_wiped_disk_redownloads_and_resumes_upload(config, source, drive):
"""Heroku dyno restart: ephemeral disk is gone but the DB and upload session survive."""
source.add(1, caption("Big.mp4", "T"), video_bytes(2_000_000, seed=4))
db = Database.from_config(config)
queue = build_queue(config, db, source, drive)
drive.crash_at_bytes = 1_000_000
async def crash_run():
await queue.sync()
with pytest.raises(SimulatedCrash):
await queue.process_item(db.next_item())
run(crash_run())
shutil.rmtree(config.download_dir)
db2 = Database.from_config(config)
db2.recover_stale()
queue2 = build_queue(config, db2, source, drive)
run(_start_and_wait(queue2))
assert db2.get(key(1)).state == State.COMPLETED
assert len([1 for m, _ in source.download_offsets if m == 1]) == 2
assert drive.resumed and drive.resumed[0][1] >= 1_000_000
assert len(drive.uploaded()) == 1
async def _start_and_wait(queue):
queue.start()
await queue.wait()
def test_upload_success_then_database_write_failure(config, source, drive):
source.add(1, caption("A.mp4", "T"))
db = Database.from_config(config)
real_transition = db.transition
failed = {"done": False}
def flaky_transition(k, state, **kw):
if state == State.COMPLETED and not failed["done"]:
failed["done"] = True
raise db.errors[0]("database is locked")
return real_transition(k, state, **kw)
db.transition = flaky_transition
queue = build_queue(config, db, source, drive)
async def first():
await queue.sync()
await queue.process_item(db.next_item())
run(first())
rec = db.get(key(1))
assert rec.state == State.UPLOADING # completion was not recorded
assert len(drive.uploaded()) == 1
assert (config.download_dir / f"{abs(CHANNEL_ID)}_1.video").exists() # not cleaned up
# Same process continues (or a restart): it reconciles instead of re-uploading.
run(_start_and_wait(queue))
rec = db.get(key(1))
assert rec.state == State.COMPLETED
assert len(drive.uploaded()) == 1 and drive.upload_calls == 1
assert rec.drive_file_id == drive.uploaded()[0].id
def test_lost_upload_response_does_not_duplicate(config, db, source, drive):
source.add(1, caption("A.mp4", "T"))
drive.lose_response_once = True
run(sync_and_run(build_queue(config, db, source, drive)))
assert db.get(key(1)).state == State.COMPLETED
assert len(drive.uploaded()) == 1
# -- control / concurrency ------------------------------------------------------------------
def test_sequential_processing_never_overlaps(config, db, source, drive):
for i in range(1, 6):
source.add(i, caption(f"Class {i}.mp4", "T"))
events = []
source.events = events
drive.events_sink = events
queue = build_queue(config, db, source, drive)
async def go():
await queue.sync()
assert queue.start() is True
assert queue.start() is False # a second worker is never created
await queue.wait()
run(go())
assert drive.max_active_uploads == 1 and queue.max_concurrent_uploads == 1
order = [(kind, str(ident).split(":")[-1]) for kind, ident in events]
expected = []
for i in range(1, 6):
expected += [("download", str(i)), ("upload", str(i))]
assert order == expected
def test_pause_checkpoint_and_resume(config, db, source, drive):
source.add(1, caption("Big.mp4", "T"), video_bytes(3_000_000, seed=7))
queue = build_queue(config, db, source, drive)
async def go():
await queue.sync()
queue.start()
while not (config.download_dir / f"{abs(CHANNEL_ID)}_1.part").exists():
await asyncio.sleep(0)
assert queue.pause()
for _ in range(200):
await asyncio.sleep(0.005)
if db.get(key(1)).state == State.PENDING:
break
assert db.get(key(1)).state == State.PENDING
assert queue.status == "paused"
assert queue.desired_state() == "paused"
assert queue.resume()
while db.get(key(1)).state != State.COMPLETED:
await asyncio.sleep(0.01)
await queue.stop()
assert queue.desired_state() == "stopped"
run(go())
assert len(drive.uploaded()) == 1
def test_graceful_shutdown_keeps_desired_state_and_recovers(config, source, drive):
source.add(1, caption("Big.mp4", "T"), video_bytes(3_000_000, seed=8))
db = Database.from_config(config)
queue = build_queue(config, db, source, drive)
async def go():
await queue.sync()
queue.start()
while db.get(key(1)).state != State.UPLOADING:
await asyncio.sleep(0.001)
await queue.shutdown(timeout=10)
run(go())
rec = db.get(key(1))
assert rec.state == State.UPLOADING and rec.upload_session_uri
assert queue.desired_state() == "running" # auto-resumes after restart
db.close()
db2 = Database.from_config(config)
db2.recover_stale()
queue2 = build_queue(config, db2, source, drive)
run(_start_and_wait(queue2))
assert db2.get(key(1)).state == State.COMPLETED
assert len(drive.uploaded()) == 1 and drive.resumed
def test_progress_messages_are_edited_not_spammed(config, db, source, drive):
for i in (1, 2):
source.add(i, caption(f"Class {i}.mp4", "Misc• Grammar Practice"))
notifier = RecordingNotifier()
reporter = ProgressReporter(notifier, [111], interval=1)
queue = build_queue(config, db, source, drive, reporter=reporter)
run(sync_and_run(queue))
texts = [t for _, t in notifier.sent]
progress_messages = [t for t in texts if t.startswith("Processing video") or t == "Queue idle"]
assert len(progress_messages) <= 1 # one progress message, then edits
assert texts[-1].startswith("Run finished") and "Uploaded: 2" in texts[-1]
tests/test_upload_tasks.py
+447
View File
@@ -0,0 +1,447 @@
"""Acceptance tests for the interactive /upload workflow (mocked Telegram and Drive)."""
import asyncio
from datetime import timedelta
import pytest
from bot_handlers import CommandRouter
from conftest import (
CHANNEL_ID,
FakeDrive,
FakeSource,
RecordingNotifier,
SimulatedCrash,
build_queue,
make_config,
run,
video_bytes,
)
from database import Database, iso, utcnow
from errors import PermanentError
from link_parser import internal_channel_id
from models import State, TaskStatus
from progress import ProgressReporter
A, B, STRANGER = 111, 222, 999
INTERNAL = internal_channel_id(CHANNEL_ID) # 1234567890
OTHER_CHANNEL = -1000999999999 # accessible, but not the configured channel
def link(message_id, internal=INTERNAL):
return f"https://t.me/c/{internal}/{message_id}"
def caption(title, topic):
return f"Index: 001\nTitle: {title}\nTopic: {topic}\nBatch: SSC\nExtracted By: @x"
def populate(source, size=600_000):
"""Messages 13..21; 15 is text, 18 is missing; 13 and 21 lie outside the 14-20 range."""
source.add(13, caption("Outside Before.mp4", "Topic X"), video_bytes(size, 13))
source.add(14, caption("Class 01 Aman sir.mp4", "Misc• Grammar Practice"), video_bytes(size, 14))
source.add_text(15, "Announcement: new classes tomorrow")
source.add(16, caption("Class 01 Aman sir.mp4", "Error Detection"), video_bytes(size, 16))
source.add(17, caption("Class 02 Aman sir.mp4", "Misc• Grammar Practice"), video_bytes(size, 17))
source.add(19, caption("Class 03 Aman sir", "Misc• Grammar Practice"), video_bytes(size, 19))
source.add(20, "", video_bytes(size, 20))
source.add(21, caption("Outside After.mp4", "Topic X"), video_bytes(size, 21))
source.accessible[OTHER_CHANNEL] = "Other Channel"
class Env:
def __init__(self, config, db, source, drive, notifier=None):
self.config, self.db, self.source, self.drive = config, db, source, drive
self.notifier = notifier or RecordingNotifier()
reporter = ProgressReporter(self.notifier, [A], interval=60)
self.queue = build_queue(config, db, source, drive, reporter=reporter)
self.router = CommandRouter(config, db, self.queue)
async def say(self, user, *texts):
replies = []
for text in texts:
replies.append(await self.router.handle(user, text))
return replies
def setup(self, user):
return self.db.tasks_with_status(
TaskStatus.AWAITING_START, TaskStatus.AWAITING_END, TaskStatus.AWAITING_CONFIRMATION, admin_user_id=user
)
def drive_names(self):
return sorted((self.drive.files[f.parents[0]].name, f.name) for f in self.drive.uploaded())
@pytest.fixture
def env(tmp_path, db_url):
extra = {"DATABASE_URL": db_url} if db_url else {}
config = make_config(tmp_path, ADMIN_USER_IDS=f"{A},{B}", **extra)
db = Database.from_config(config)
source = FakeSource()
populate(source)
e = Env(config, db, source, FakeDrive())
yield e
db.close()
async def full_task(env, start=14, end=20, user=A):
replies = await env.say(user, "/upload", link(start), link(end), "/confirm")
await env.queue.wait()
return replies
# 1-4, 7-10: conversation, inclusive ascending range, metadata, skips ------------------------
def test_upload_conversation_and_range_processing(env):
async def go():
r = await env.say(A, "/upload")
assert "📥 <b>Start Upload Task</b>" in r[0] and "https://t.me/c/3909337424/14" in r[0]
assert env.setup(A)[0].status == TaskStatus.AWAITING_START
r = await env.say(A, link(14))
assert "📌 <b>Enter the End Link</b>" in r[0]
assert env.setup(A)[0].start_message_id == 14
r = await env.say(A, link(20))
assert "✅ <b>Upload Task Ready</b>" in r[0]
assert "Start Message ID: 14" in r[0] and "End Message ID: 20" in r[0]
assert "Test Channel (configured source channel)" in r[0]
assert env.source.download_offsets == [] # nothing downloaded before /confirm
r = await env.say(A, "/confirm")
assert "Upload Task #1 confirmed" in r[0]
await env.queue.wait()
run(go())
task = env.db.get_task(1)
assert task.status == TaskStatus.COMPLETED
# inclusive range 14..20, nothing outside it
assert [m for m, _ in env.db.task_items(1)] == [14, 15, 16, 17, 18, 19, 20]
downloaded = [m for kind, m in env.source.events if kind == "download"]
assert downloaded == [14, 16, 17, 19, 20] # ascending, videos only
assert not env.db.exists(CHANNEL_ID, 13) and not env.db.exists(CHANNEL_ID, 21)
# caption metadata -> folders and filenames
assert env.drive_names() == [
("Error Detection", "Class 01 Aman sir.mp4"),
("Misc• Grammar Practice", "Class 01 Aman sir.mp4"),
("Misc• Grammar Practice", "Class 02 Aman sir.mp4"),
("Misc• Grammar Practice", "Class 03 Aman sir.mp4"),
("Uncategorized", "telegram_20.mp4"),
]
# non-video and missing messages are skipped, not failed
assert dict(env.db.task_items(1))[15] == "non_video"
assert dict(env.db.task_items(1))[18] == "missing"
assert (task.uploaded_count, task.failed_count, task.non_video_count, task.missing_count) == (5, 0, 1, 1)
assert task.examined_count == 6 and task.eligible_count == 5
assert task.deleted_files_count == 5
assert list(env.config.download_dir.iterdir()) == []
summary = [t for _, t in env.notifier.sent if "Upload Task Completed" in t][0]
for fragment in ("Requested range: 14–20", "Eligible videos found: 5", "Uploaded: 5", "Duplicates skipped: 0",
"Non-video messages skipped: 1", "Failed: 0", "Local temporary files deleted: 5",
"Misc• Grammar Practice", "Error Detection"):
assert fragment in summary
def test_confirm_requires_both_validated_links(env):
async def go():
assert "no upload task waiting" in (await env.say(A, "/confirm"))[0]
await env.say(A, "/upload")
assert "Both the start and end links" in (await env.say(A, "/confirm"))[0]
await env.say(A, link(14))
assert "Both the start and end links" in (await env.say(A, "/confirm"))[0]
assert env.db.next_runnable_task() is None
run(go())
assert env.source.download_offsets == []
@pytest.mark.parametrize(
"bad,fragment",
[
("example.com/post/14", "not a Telegram"),
("https://t.me/+invite", "Invite links"),
("https://t.me/c/1234567890", "must look like"),
("https://t.me/c/555/14", "not accessible"), # channel the account cannot access
(link(14, internal_channel_id(OTHER_CHANNEL)), "not from the configured source channel"),
(link(18), "not found"), # deleted message
],
)
def test_invalid_start_links_rejected(env, bad, fragment):
replies = run(env.say(A, "/upload", bad))
assert replies[1].startswith("❌") and fragment in replies[1]
assert env.setup(A)[0].status == TaskStatus.AWAITING_START
def test_end_link_errors_keep_start_link(env):
replies = run(env.say(
A, "/upload", link(14),
link(20, internal_channel_id(OTHER_CHANNEL)), # different channel
link(13), # before start
"https://t.me/c/1234567890/abc", # malformed
))
assert "different channel" in replies[2]
assert "greater than or equal" in replies[3]
assert replies[4].startswith("❌")
setup = env.setup(A)[0]
assert setup.status == TaskStatus.AWAITING_END and setup.start_message_id == 14
assert "Upload Task Ready" in run(env.say(A, link(14)))[0] # single-message range is allowed
def test_other_channels_allowed_when_unrestricted(tmp_path):
config = make_config(tmp_path, ADMIN_USER_IDS=str(A), UPLOAD_RESTRICT_TO_SOURCE_CHANNEL="false")
db = Database.from_config(config)
source = FakeSource()
populate(source)
e = Env(config, db, source, FakeDrive())
async def other_exists(channel_id, message_id):
return True
source.message_accessible = other_exists
replies = run(e.say(A, "/upload", link(1, internal_channel_id(OTHER_CHANNEL)), link(5, internal_channel_id(OTHER_CHANNEL))))
assert "Channel: Other Channel" in replies[2]
def test_range_size_limit(tmp_path):
config = make_config(tmp_path, ADMIN_USER_IDS=str(A), UPLOAD_MAX_RANGE="3")
e = Env(config, Database.from_config(config), FakeSource(), FakeDrive())
populate(e.source)
replies = run(e.say(A, "/upload", link(14), link(20)))
assert "maximum is 3" in replies[2]
def test_non_admins_cannot_use_upload_commands(env):
for command in ("/upload", "/confirm", "/cancel", "/taskstatus", "/uploadhistory"):
assert "not authorized" in run(env.say(STRANGER, command))[0]
assert run(env.say(STRANGER, link(14)))[0] is None # plain text from non-admins is ignored
assert env.db.recent_tasks() == [] and env.setup(STRANGER) == []
# 15: administrators are isolated ------------------------------------------------------
def test_admin_setups_are_isolated(env):
async def go():
await env.say(A, "/upload", link(14))
r = await env.say(B, link(16))
assert "send /upload first" in r[0]
await env.say(B, "/upload", link(16))
assert env.setup(A)[0].start_message_id == 14
assert env.setup(B)[0].start_message_id == 16
r = await env.say(B, "/cancel")
assert "cancelled" in r[0]
assert env.setup(B) == []
assert env.setup(A)[0].status == TaskStatus.AWAITING_END # untouched
await env.say(A, link(17), "/confirm")
r = await env.say(B, "/cancel")
assert "belong to another administrator" in r[0]
await env.queue.wait()
run(go())
assert env.db.get_task(1).status == TaskStatus.COMPLETED
def test_second_upload_does_not_overwrite_setup(env):
replies = run(env.say(A, "/upload", link(14), "/upload"))
assert "already have an upload setup" in replies[2]
setup = env.setup(A)
assert len(setup) == 1 and setup[0].start_message_id == 14
# setup persistence / expiry ------------------------------------------------------------
def test_setup_survives_restart(env):
run(env.say(A, "/upload", link(14)))
# new process: fresh queue/router on the same durable database
again = Env(env.config, env.db, env.source, env.drive)
replies = run(again.say(A, link(20)))
assert "Upload Task Ready" in replies[0]
def test_stale_setup_expires(env):
run(env.say(A, "/upload", link(14)))
setup = env.setup(A)[0]
env.db.update_task(setup.id, updated_at=iso(utcnow() - timedelta(minutes=env.config.upload_setup_timeout + 1)))
reply = run(env.say(A, "/confirm"))[0]
assert "expired" in reply
assert env.db.get_task(setup.id).status == TaskStatus.CANCELLED
assert "Start Upload Task" in run(env.say(A, "/upload"))[0]
# 11, 12: cleanup ordering and failures -------------------------------------------------
def test_local_file_deleted_only_after_verified_upload_is_recorded(env):
seen = []
real_cleanup = env.queue.downloader.cleanup
def spy(rec):
current = env.db.get(rec.key)
if current.state == State.COMPLETED:
assert current.drive_file_id in env.drive.files
assert env.drive.files[current.drive_file_id].size == current.local_size
assert env.config.download_dir.joinpath(f"{abs(CHANNEL_ID)}_{rec.message_id}.video").exists()
seen.append((rec.message_id, current.state))
return real_cleanup(rec)
env.queue.downloader.cleanup = spy
run(full_task(env))
assert seen and all(state == State.COMPLETED for _, state in seen)
def test_failed_upload_is_not_completed(env):
env.drive.upload_failures = [PermanentError("insufficientFilePermissions")]
run(full_task(env))
rec = env.db.get((CHANNEL_ID, 14))
assert rec.state == State.FAILED and rec.drive_file_id is None
task = env.db.get_task(1)
assert task.status == TaskStatus.COMPLETED and task.failed_count == 1 and task.uploaded_count == 4
assert dict(env.db.task_items(1))[14] == "failed"
assert not any(p.name.startswith(f"{abs(CHANNEL_ID)}_14.") for p in env.config.download_dir.iterdir())
def test_failed_download_file_can_be_kept(tmp_path):
config = make_config(tmp_path, ADMIN_USER_IDS=str(A), KEEP_FAILED_DOWNLOADS="true")
e = Env(config, Database.from_config(config), FakeSource(), FakeDrive())
populate(e.source)
e.drive.upload_failures = [PermanentError("denied")]
run(full_task(e))
assert (config.download_dir / f"{abs(CHANNEL_ID)}_14.video").exists()
# 14: repeating a task ------------------------------------------------------------------
def test_repeating_completed_task_does_not_duplicate(env):
run(full_task(env))
calls = env.drive.upload_calls
run(full_task(env))
second = env.db.get_task(2)
assert second.status == TaskStatus.COMPLETED
assert second.uploaded_count == 0 and second.duplicate_count == 5
assert env.drive.upload_calls == calls and len(env.drive.uploaded()) == 5
# 16: cancellation ----------------------------------------------------------------------
def test_cancel_running_task_keeps_completed_work(tmp_path):
config = make_config(tmp_path, ADMIN_USER_IDS=str(A))
db = Database.from_config(config)
source = FakeSource()
populate(source, size=3_000_000)
e = Env(config, db, source, FakeDrive())
part16 = config.download_dir / f"{abs(CHANNEL_ID)}_16.part"
async def go():
await e.say(A, "/upload", link(14), link(20), "/confirm")
while not part16.exists():
await asyncio.sleep(0.001)
reply = (await e.say(A, "/cancel"))[0]
assert "cancelled at a safe checkpoint" in reply
await e.queue.wait()
run(go())
task = db.get_task(1)
assert task.status == TaskStatus.CANCELLED
assert db.get((CHANNEL_ID, 14)).state == State.COMPLETED
assert len(e.drive.uploaded()) == 1 # completed upload preserved in Drive
for message_id in (16, 17, 19, 20):
rec = db.get((CHANNEL_ID, message_id))
assert rec.state == State.SKIPPED and rec.skip_reason.startswith("cancelled:")
assert list(config.download_dir.iterdir()) == []
assert "Upload Task Cancelled" in "".join(t for _, t in e.notifier.sent)
# the same range can be requested again: cancelled items are re-queued, done ones are not re-uploaded
run(full_task(e))
second = db.get_task(2)
assert second.uploaded_count == 4 and second.duplicate_count == 1
assert len(e.drive.uploaded()) == 5
def test_cancel_queued_task_and_explicit_id(env):
async def go():
await env.say(A, "/pause") # nothing running -> nothing to pause
env.queue._persist_desired("paused")
await env.say(A, "/upload", link(14), link(20), "/confirm")
assert env.db.get_task(1).status == TaskStatus.QUEUED
reply = (await env.say(B, "/cancel 1"))[0]
assert "Task #1 cancelled" in reply
assert "already finished" in (await env.say(A, "/cancel 1"))[0]
run(go())
assert env.db.get_task(1).status == TaskStatus.CANCELLED
assert env.source.download_offsets == []
# 13, 18: restart recovery ----------------------------------------------------------------
def test_task_recovers_after_crash_and_restart(tmp_path, db_url):
extra = {"DATABASE_URL": db_url} if db_url else {}
config = make_config(tmp_path, ADMIN_USER_IDS=str(A), **extra)
source = FakeSource()
populate(source, size=1_500_000)
drive = FakeDrive()
db = Database.from_config(config)
e = Env(config, db, source, drive)
drive.crash_at_bytes = 700_000 # the process is killed half-way through the first upload
async def first_run():
await e.say(A, "/upload", link(14), link(20), "/confirm")
with pytest.raises(SimulatedCrash):
await e.queue.wait()
run(first_run())
assert db.get_task(1).status == TaskStatus.RUNNING
assert db.get((CHANNEL_ID, 14)).state == State.UPLOADING
db.close()
# Heroku restarts the worker: new process, same durable database
db2 = Database.from_config(config)
db2.recover_stale()
e2 = Env(config, db2, source, drive)
async def second_run():
assert e2.queue.boot() == "running"
await e2.queue.wait()
run(second_run())
task = db2.get_task(1)
assert task.status == TaskStatus.COMPLETED and task.uploaded_count == 5
assert len(drive.uploaded()) == 5 # no duplicates
assert drive.resumed # the interrupted upload session was resumed
def test_pause_resume_task_and_paused_boot(env):
async def go():
await env.say(A, "/upload", link(14), link(20), "/confirm")
while env.db.get((CHANNEL_ID, 14)) is None or env.db.get((CHANNEL_ID, 14)).state != State.COMPLETED:
await asyncio.sleep(0.001)
assert "Pausing" in (await env.say(A, "/pause"))[0]
await asyncio.sleep(0.05)
assert env.db.get_task(1).status == TaskStatus.PAUSED
# a restart while paused stays paused
assert Env(env.config, env.db, env.source, env.drive).queue.boot() == "paused"
assert "resumed" in (await env.say(A, "/resume"))[0]
await env.queue.wait()
run(go())
assert env.db.get_task(1).status == TaskStatus.COMPLETED
def test_status_and_history_commands(env):
async def go():
await env.say(A, "/upload", link(14))
status = (await env.say(A, "/taskstatus"))[0]
assert "waiting for the end link" in status
await env.say(A, link(20), "/confirm")
await env.queue.wait()
status = (await env.say(A, "/taskstatus 1"))[0]
assert "Task #1</b> — COMPLETED" in status and "Range: 14–20" in status
history = (await env.say(A, "/uploadhistory"))[0]
assert "#1 14–20 COMPLETED — uploaded 5" in history
assert "Upload tasks" not in (await env.say(A, "/status"))[0]
run(go())