Compare commits
85
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ec5979437d | ||
|
|
59f171a62d | ||
|
|
e3559da2a7 | ||
|
|
dcd6786ca5 | ||
|
|
32d58193f3 | ||
|
|
59d0e2c8ad | ||
|
|
850d39ad6f | ||
|
|
b8a38092eb | ||
|
|
2836dcb3a9 | ||
|
|
01cf01b629 | ||
|
|
7f124059be | ||
|
|
5eacabd101 | ||
|
|
0362777a51 | ||
|
|
dfc726976b | ||
|
|
cb01d7dff8 | ||
|
|
f646f29ff9 | ||
|
|
945adcf879 | ||
|
|
d213dfdf64 | ||
|
|
566ee09431 | ||
|
|
09c5f78061 | ||
|
|
897421e7c4 | ||
|
|
3223f62dea | ||
|
|
c9c088ed73 | ||
|
|
580927da0a | ||
|
|
06ae3e928f | ||
|
|
55634984fd | ||
|
|
f2c245707a | ||
|
|
c20e132872 | ||
|
|
4d16105e92 | ||
|
|
b1ffebe858 | ||
|
|
2512a7988d | ||
|
|
56146e4e8d | ||
|
|
5fb8f8a534 | ||
|
|
98ab8a579b | ||
|
|
a5d10146fd | ||
|
|
5f1f18b343 | ||
|
|
8eb31ed4d3 | ||
|
|
e67dbdb325 | ||
|
|
869f8291fb | ||
|
|
283eab1eb0 | ||
|
|
7f4d500302 | ||
|
|
cf48d2db3e | ||
|
|
114058e6ae | ||
|
|
eb8187ff6e | ||
|
|
53dc805fc3 | ||
|
|
ed8489616e | ||
|
|
b60815d042 | ||
|
|
468f72ed11 | ||
|
|
59d3f214e8 | ||
|
|
34dcbaac2b | ||
|
|
c8e5cff871 | ||
|
|
fc4758aa0b | ||
|
|
b0ed041c2e | ||
|
|
4bd031f21e | ||
|
|
089a31db8d | ||
|
|
81ebbdc93e | ||
|
|
a25f79543a | ||
|
|
c61ecfe8f2 | ||
|
|
3d0b15bd2c | ||
|
|
0e54786672 | ||
|
|
5778611a38 | ||
|
|
b138533d88 | ||
|
|
0b6cbac945 | ||
|
|
ff80141785 | ||
|
|
d206f7e5ca | ||
|
|
4ebe6de60b | ||
|
|
02fb567c1f | ||
|
|
ccbc167005 | ||
|
|
5ec0212af5 | ||
|
|
ab59ef4525 | ||
|
|
676680319b | ||
|
|
6a6e63ff75 | ||
|
|
1fa4f18a9b | ||
|
|
9c12628304 | ||
|
|
cf277c56cf | ||
|
|
8f7d47476c | ||
|
|
c832513e5c | ||
|
|
0399145791 | ||
|
|
0b8e90159f | ||
|
|
995322b17a | ||
|
|
6ae28cb321 | ||
|
|
e833de245a | ||
|
|
ef739838a5 | ||
|
|
3e08698b62 | ||
|
|
683b9a31a3 |
No files matched your search
@@ -4,7 +4,7 @@
|
||||
{
|
||||
"hooks": [
|
||||
{
|
||||
"command": "bun fix --skip=correctness/noUnusedImports",
|
||||
"command": "cd frontend && bun run fix --skip=correctness/noUnusedImports",
|
||||
"type": "command"
|
||||
}
|
||||
],
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
_commit: v0.7.1
|
||||
_src_path: https://git.kotikot.com/templates/backend-python.git
|
||||
author_email: h@kotikot.com
|
||||
author_name: h
|
||||
backend_services:
|
||||
- api
|
||||
- kurigram_userbot
|
||||
database: postgres
|
||||
dynamic_config: false
|
||||
install_deps: false
|
||||
layout: part
|
||||
project_description: A read-only Telegram client that keeps everything
|
||||
project_name: beavergram
|
||||
project_slug: beavergram
|
||||
serve_spa: true
|
||||
use_llm: false
|
||||
use_redis: false
|
||||
@@ -0,0 +1,5 @@
|
||||
_commit: v0.1.4
|
||||
_src_path: https://git.kotikot.com/templates/caddy.git
|
||||
project_slug: beavergram
|
||||
routes:
|
||||
- api
|
||||
@@ -0,0 +1,9 @@
|
||||
_commit: v0.8.0
|
||||
_src_path: https://git.kotikot.com/templates/frontend-svelte.git
|
||||
install_deps: false
|
||||
layout: part
|
||||
linter: biome
|
||||
project_name: beavergram
|
||||
project_slug: beavergram
|
||||
serve: api
|
||||
ui_library: none
|
||||
@@ -0,0 +1,12 @@
|
||||
_commit: v0.6.0
|
||||
_src_path: https://git.kotikot.com/templates/infra.git
|
||||
init_git: false
|
||||
layout: parts
|
||||
mcp_servers:
|
||||
- context7
|
||||
- svelte
|
||||
parts:
|
||||
- backend
|
||||
- frontend
|
||||
project_description: A read-only Telegram client that keeps everything
|
||||
project_name: beavergram
|
||||
+2
-16
@@ -1,16 +1,2 @@
|
||||
COMPOSE_PROFILES=db,userbot,api
|
||||
RUN_ENVIRONMENT=prod
|
||||
|
||||
# Leave empty for *.localhost.
|
||||
FRONTEND_DEV_HOST=
|
||||
|
||||
DB__HOST=postgres
|
||||
DB__PORT=5432
|
||||
DB__USER=beavergram
|
||||
DB__PASSWORD=beavergram
|
||||
DB__DB_NAME=beavergram
|
||||
DB__MIN_POOL_SIZE=5
|
||||
DB__MAX_POOL_SIZE=20
|
||||
DB__SCRIPTS_CONNECTION_URL=postgresql://beavergram:beavergram@localhost:5433/beavergram
|
||||
|
||||
STORAGE__ROOT=./storage
|
||||
COMPOSE_FILE=docker-compose.yml:backend/docker-compose.override.yml:frontend/docker-compose.override.yml
|
||||
COMPOSE_PROFILES=external,services
|
||||
+8
-3
@@ -1,6 +1,11 @@
|
||||
.idea
|
||||
docker-compose.override.yml
|
||||
.env
|
||||
.env.*
|
||||
!.env.example
|
||||
docker-compose.override.yml
|
||||
.DS_Store
|
||||
.idea
|
||||
docs
|
||||
storage
|
||||
backend/sessions
|
||||
|
||||
# Local playground
|
||||
/t/
|
||||
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"mcpServers": {
|
||||
"svelte": {
|
||||
"type": "http",
|
||||
"url": "https://mcp.svelte.dev/mcp"
|
||||
},
|
||||
"context7": {
|
||||
"type": "http",
|
||||
"url": "https://mcp.context7.com/mcp"
|
||||
}
|
||||
}
|
||||
"mcpServers": {
|
||||
"context7": {
|
||||
"type": "http",
|
||||
"url": "https://mcp.context7.com/mcp"
|
||||
},
|
||||
"svelte": {
|
||||
"type": "http",
|
||||
"url": "https://mcp.svelte.dev/mcp"
|
||||
}
|
||||
}
|
||||
}
|
||||
+18
-22
@@ -1,31 +1,27 @@
|
||||
repos:
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.15.15
|
||||
hooks:
|
||||
- id: ruff-check
|
||||
types_or: [ python, pyi ]
|
||||
args: [ --fix ]
|
||||
- id: ruff-format
|
||||
types_or: [ python, pyi ]
|
||||
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: ultracite
|
||||
name: ultracite
|
||||
entry: bash -c 'cd frontend && bun x ultracite fix'
|
||||
- id: backend-fmt
|
||||
name: backend fmt
|
||||
entry: make -C backend fmt
|
||||
language: system
|
||||
files: ^backend/
|
||||
pass_filenames: false
|
||||
- id: backend-check
|
||||
name: backend check
|
||||
entry: make -C backend check
|
||||
language: system
|
||||
files: ^backend/
|
||||
pass_filenames: false
|
||||
- id: frontend-fmt
|
||||
name: frontend fmt
|
||||
entry: make -C frontend fmt
|
||||
language: system
|
||||
files: ^frontend/
|
||||
pass_filenames: false
|
||||
|
||||
- id: svelte-check
|
||||
name: svelte check
|
||||
entry: bash -c 'cd frontend && bun check'
|
||||
- id: frontend-check
|
||||
name: frontend check
|
||||
entry: make -C frontend check
|
||||
language: system
|
||||
files: ^frontend/
|
||||
pass_filenames: false
|
||||
|
||||
- id: ty
|
||||
name: ty check
|
||||
entry: uvx ty check
|
||||
language: python
|
||||
types_or: [ python, pyi ]
|
||||
@@ -6,10 +6,11 @@
|
||||
- **Reusability**: Create reusable components and functions
|
||||
- **Less Code = Less Debt**: Minimize code footprint
|
||||
- **NEVER write comments** - code should be self-documenting
|
||||
- **Docstrings ONLY if explicitly asked**
|
||||
|
||||
## Checking commands
|
||||
After writing code, always run:
|
||||
```shell
|
||||
bun check # svelte-check
|
||||
bun fix # biome linter
|
||||
make fmt
|
||||
make check
|
||||
```
|
||||
@@ -1,30 +1,34 @@
|
||||
.PHONY: recreate down restart rebuild deploy migrate session-create frontend
|
||||
PARTS := $(patsubst %/Makefile,%,$(wildcard */Makefile))
|
||||
|
||||
.PHONY: env recreate rebuild restart logs down fmt check build
|
||||
|
||||
env:
|
||||
@fresh=; test -f .env || { awk 'FNR == 1 && NR > 1 { print "" } 1' .env.example $(addsuffix /.env.example,$(PARTS)) > .env; fresh=1; }; \
|
||||
for part in $(PARTS); do \
|
||||
override=$$part/docker-compose.override.yml; \
|
||||
test -f $$override && continue; \
|
||||
cp $$override.example $$override; \
|
||||
for port in $$(sed -n 's/.*127\.0\.0\.1:\([0-9]*\):.*/\1/p' $$override | sort -u); do \
|
||||
random=$$((20000 + $$(od -An -N2 -tu2 /dev/urandom) % 40000)); \
|
||||
sed "s/127\.0\.0\.1:$$port:/127.0.0.1:$$random:/" $$override > $$override.tmp && mv $$override.tmp $$override; \
|
||||
test -z "$$fresh" || { sed "s|localhost:$$port/|localhost:$$random/|" .env > .env.tmp && mv .env.tmp .env; }; \
|
||||
done; \
|
||||
done
|
||||
|
||||
recreate:
|
||||
docker compose up -d --force-recreate
|
||||
|
||||
rebuild: build
|
||||
docker compose up -d
|
||||
|
||||
restart:
|
||||
docker compose restart
|
||||
|
||||
logs:
|
||||
docker compose logs -f --tail=100
|
||||
|
||||
down:
|
||||
docker compose down
|
||||
|
||||
restart:
|
||||
$(MAKE) recreate
|
||||
|
||||
rebuild:
|
||||
docker compose build
|
||||
docker compose up -d
|
||||
|
||||
frontend:
|
||||
docker compose run --rm --no-deps frontend-dev sh -c "bun install && bun run build"
|
||||
|
||||
deploy:
|
||||
$(MAKE) frontend
|
||||
$(MAKE) rebuild
|
||||
|
||||
migrate:
|
||||
docker compose --profile db --profile migrate run --rm migrator $(filter-out $@,$(MAKECMDGOALS))
|
||||
|
||||
session-create:
|
||||
cd backend && uv run python scripts/session/create.py
|
||||
|
||||
%:
|
||||
@:
|
||||
fmt check build:
|
||||
@for part in $(PARTS); do $(MAKE) -C $$part $@ || exit 1; done
|
||||
@@ -1,2 +1,46 @@
|
||||
# beavergram
|
||||
[](https://sladge.net)
|
||||
|
||||
A read-only Telegram client that keeps everything. A userbot on your own account logs every incoming message into Postgres, including the ones later deleted or edited, along with media, reactions, presence and stories. A web client reads that archive back, and an MCP server hands it to agents.
|
||||
|
||||
Useful on its own: a Telegram where nothing disappears. In the [beaver](https://git.kotikot.com/beaver) setup it is one of the agent's hands - the dispatcher reads chat history by the `chat_id` stored in a person's card.
|
||||
|
||||
## What it keeps
|
||||
|
||||
- **Messages** with edits as versions and deletions as flags, so a message removed in Telegram is still in the archive with its history.
|
||||
- **Media** on disk, content-addressed, with the versions of a replaced photo.
|
||||
- **Voice** transcribed through Telegram's own transcription (Premium), queued when the quota runs out.
|
||||
- **Presence** as a time series (TimescaleDB), read receipts, stories.
|
||||
- **Capture policy** from your Telegram folders with per-chat overrides: channels nothing, groups messages only, private chats everything, by default.
|
||||
- **Annotations and watches**: your notes on chats and rules that fire alerts, both yours to write in the web client.
|
||||
- **Watched bots**: bot tokens you add in the web client. Each one logs in over MTProto (`auth.importBotAuthorization`, the same door a bot-login client uses), backfills the history of the bot's dialogs once and then keeps the incoming stream: who writes to the bot, what, and when. A separate authorization, so it does not take `getUpdates` away from whoever runs that bot over the Bot API.
|
||||
|
||||
Nothing is ever deleted from the archive.
|
||||
|
||||
## Parts
|
||||
|
||||
- `backend/` - `userbot` (kurigram, the capture loop and jobs), `watcher` (kurigram again, one bot-token session per watched bot) and `api` (FastAPI: the web API, the MCP server, and the built frontend as static files).
|
||||
- `frontend/` - SvelteKit, `adapter-static`; a chat client over the archive: folders, chat list, thread with media, versions and the deleted messages inline, search, stories, presence graphs.
|
||||
- `docker-compose.yml` - includes `backend/` (Postgres with TimescaleDB, `userbot`, `api`) and `frontend/` (the dev server); profiles `external` and `services`.
|
||||
|
||||
## MCP
|
||||
|
||||
Read-only: `list_watched_bots`, `get_bot_stats`, `get_bot_updates`, `search_messages_tool`, `list_chats`, `get_chat_history` (with deleted messages included), `get_deleted_messages`, `get_message_versions`, `get_media`, `get_callbacks`, `presence_history`, `get_peer_history`, `get_stories`, `get_annotations`, `list_watches`, `list_alerts`. The one write is `set_watch`. There are no send, edit or delete tools and there will not be.
|
||||
|
||||
## Run
|
||||
|
||||
```sh
|
||||
make env # .env and docker-compose.override.yml from the examples
|
||||
make -C backend session # log the account in once; the session file stays in backend/sessions
|
||||
make -C backend migrate # the schema
|
||||
make recreate # build and start the services in COMPOSE_PROFILES
|
||||
```
|
||||
|
||||
`COMPOSE_PROFILES` in `.env` selects what runs on the machine, `docker-compose.override.yml` in each part holds the published ports and the `caddy` network aliases. The web client and the MCP endpoint are served by `api`; the bearer for both is `AUTH__TOKEN`.
|
||||
|
||||
## Develop
|
||||
|
||||
```sh
|
||||
make fmt
|
||||
make check
|
||||
make -C frontend dev
|
||||
```
|
||||
@@ -0,0 +1,7 @@
|
||||
.venv
|
||||
.env
|
||||
__pycache__
|
||||
*.pyc
|
||||
.ruff_cache
|
||||
.pytest_cache
|
||||
sessions
|
||||
+27
-9
@@ -1,14 +1,32 @@
|
||||
TG__SESSION_NAME=beavergram
|
||||
TG__SESSIONS_DIR=sessions
|
||||
|
||||
API__HOST=0.0.0.0
|
||||
API__PORT=8080
|
||||
|
||||
AUTH__TOKEN=<AUTH__TOKEN>
|
||||
|
||||
STORAGE__SHARD_DEPTH=2
|
||||
RUN_ENVIRONMENT=prod
|
||||
TZ=UTC
|
||||
|
||||
LOG__LEVEL=INFO
|
||||
LOG__LEVEL_EXTERNAL=WARNING
|
||||
LOG__SHOW_TIME=false
|
||||
LOG__SHOW_LOCALS=false
|
||||
LOG__CONSOLE_WIDTH=150
|
||||
|
||||
DB__HOST=postgres
|
||||
DB__PORT=5432
|
||||
DB__USER=beavergram
|
||||
DB__PASSWORD=beavergram
|
||||
DB__DB_NAME=beavergram
|
||||
DB__MIN_POOL_SIZE=5
|
||||
DB__MAX_POOL_SIZE=20
|
||||
DB__SCRIPTS_CONNECTION_URL=postgresql://beavergram:beavergram@localhost:5432/beavergram
|
||||
|
||||
TG__SESSION_NAME=beavergram
|
||||
TG__SESSIONS_DIR=sessions
|
||||
|
||||
BOTS__SESSIONS_DIR=sessions/bots
|
||||
BOTS__BACKFILL_CHATS=0
|
||||
BOTS__BACKFILL_PER_CHAT=0
|
||||
|
||||
API__HOST=0.0.0.0
|
||||
API__PORT=8080
|
||||
|
||||
AUTH__TOKEN=
|
||||
|
||||
STORAGE__ROOT=./storage
|
||||
STORAGE__SHARD_DEPTH=2
|
||||
@@ -0,0 +1,8 @@
|
||||
## Checking commands
|
||||
After writing code, always run from `backend/`:
|
||||
```shell
|
||||
make fmt
|
||||
make check
|
||||
```
|
||||
|
||||
Mind startup imports: `make check` runs `scripts/check_imports.py`.
|
||||
+6
-5
@@ -1,4 +1,6 @@
|
||||
FROM ghcr.io/astral-sh/uv:python3.13-bookworm-slim
|
||||
FROM ghcr.io/astral-sh/uv:python3.13-trixie-slim
|
||||
|
||||
RUN apt-get update && apt-get install -y build-essential && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
@@ -7,10 +9,6 @@ ENV PATH="/app/.venv/bin:$PATH" \
|
||||
UV_COMPILE_BYTECODE=1 \
|
||||
UV_LINK_MODE=copy
|
||||
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends build-essential \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY pyproject.toml uv.lock ./
|
||||
RUN uv sync --frozen --no-install-project --no-dev
|
||||
|
||||
@@ -20,4 +18,7 @@ RUN uv sync --frozen --no-dev
|
||||
COPY alembic.ini ./
|
||||
COPY migrations ./migrations
|
||||
|
||||
COPY scripts ./scripts
|
||||
|
||||
ENTRYPOINT ["python", "-m"]
|
||||
CMD ["api"]
|
||||
@@ -0,0 +1,28 @@
|
||||
COMPOSE := cd .. && docker compose
|
||||
|
||||
.PHONY: fmt check build migrate revision session
|
||||
|
||||
fmt:
|
||||
uv run ruff format
|
||||
uv run ruff check --fix
|
||||
|
||||
check:
|
||||
uv run ruff format --check
|
||||
uv run ruff check
|
||||
uv run ty check
|
||||
uv run python scripts/check_imports.py
|
||||
|
||||
build:
|
||||
$(COMPOSE) build api
|
||||
|
||||
migrate:
|
||||
$(COMPOSE) run --rm migrator $(or $(filter-out $@,$(MAKECMDGOALS)),upgrade head)
|
||||
|
||||
revision:
|
||||
$(COMPOSE) run --rm migrator revision --autogenerate -m "$(m)"
|
||||
|
||||
%:
|
||||
@:
|
||||
|
||||
session:
|
||||
uv run python scripts/session/create.py
|
||||
@@ -0,0 +1,17 @@
|
||||
services:
|
||||
api:
|
||||
networks:
|
||||
default: {}
|
||||
caddy:
|
||||
aliases:
|
||||
- beavergram-api
|
||||
# ports:
|
||||
# - "127.0.0.1:8080:${API__PORT}"
|
||||
|
||||
postgres:
|
||||
# ports:
|
||||
# - "127.0.0.1:5432:${DB__PORT}"
|
||||
|
||||
networks:
|
||||
caddy:
|
||||
external: true
|
||||
@@ -0,0 +1,133 @@
|
||||
services:
|
||||
api:
|
||||
build: backend
|
||||
image: beavergram/backend
|
||||
profiles: [api, services]
|
||||
restart: unless-stopped
|
||||
env_file:
|
||||
- path: .env
|
||||
required: false
|
||||
- path: backend/.env
|
||||
required: false
|
||||
environment:
|
||||
RUN_ENVIRONMENT: prod
|
||||
TZ: ${TZ:-UTC}
|
||||
STORAGE__ROOT: /app/storage
|
||||
volumes:
|
||||
- ./backend/src:/app/src
|
||||
- ./backend/scripts:/app/scripts
|
||||
- ./backend/sessions:/app/sessions
|
||||
- ./frontend/build:/app/static:ro
|
||||
- ${STORAGE__ROOT:-./storage}:/app/storage
|
||||
command: [api]
|
||||
depends_on:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
required: false
|
||||
networks:
|
||||
default:
|
||||
database:
|
||||
|
||||
userbot:
|
||||
build: backend
|
||||
image: beavergram/backend
|
||||
profiles: [userbot, services]
|
||||
restart: unless-stopped
|
||||
env_file:
|
||||
- path: .env
|
||||
required: false
|
||||
- path: backend/.env
|
||||
required: false
|
||||
environment:
|
||||
RUN_ENVIRONMENT: prod
|
||||
TZ: ${TZ:-UTC}
|
||||
STORAGE__ROOT: /app/storage
|
||||
volumes:
|
||||
- ./backend/src:/app/src
|
||||
- ./backend/scripts:/app/scripts
|
||||
- ./backend/sessions:/app/sessions
|
||||
- ${STORAGE__ROOT:-./storage}:/app/storage
|
||||
command: [userbot]
|
||||
depends_on:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
required: false
|
||||
networks:
|
||||
database:
|
||||
|
||||
watcher:
|
||||
build: backend
|
||||
image: beavergram/backend
|
||||
profiles: [watcher, services]
|
||||
restart: unless-stopped
|
||||
env_file:
|
||||
- path: .env
|
||||
required: false
|
||||
- path: backend/.env
|
||||
required: false
|
||||
environment:
|
||||
RUN_ENVIRONMENT: prod
|
||||
TZ: ${TZ:-UTC}
|
||||
volumes:
|
||||
- ./backend/src:/app/src
|
||||
- ./backend/sessions:/app/sessions
|
||||
command: [watcher]
|
||||
depends_on:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
required: false
|
||||
networks:
|
||||
database:
|
||||
|
||||
postgres:
|
||||
image: timescale/timescaledb:2.30.2-pg17@sha256:b346edcdb51a1fd6020e3965e0bd1c9f3406fa6d5fbce1e28f4852587ef934e2
|
||||
profiles: [postgres, external]
|
||||
restart: unless-stopped
|
||||
environment:
|
||||
POSTGRES_USER: ${DB__USER:-beavergram}
|
||||
POSTGRES_PASSWORD: ${DB__PASSWORD:-beavergram}
|
||||
POSTGRES_DB: ${DB__DB_NAME:-beavergram}
|
||||
PGPORT: ${DB__PORT:-5432}
|
||||
volumes:
|
||||
- pgdata:/var/lib/postgresql/data
|
||||
- ./backend/migrations/init:/docker-entrypoint-initdb.d:ro
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U ${DB__USER:-beavergram}"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 10
|
||||
networks:
|
||||
database:
|
||||
aliases:
|
||||
- ${DB__HOST:-postgres}
|
||||
|
||||
migrator:
|
||||
build: backend
|
||||
image: beavergram/backend
|
||||
profiles: [migrate]
|
||||
env_file:
|
||||
- path: .env
|
||||
required: false
|
||||
- path: backend/.env
|
||||
required: false
|
||||
environment:
|
||||
RUN_ENVIRONMENT: prod
|
||||
TZ: ${TZ:-UTC}
|
||||
volumes:
|
||||
- ./backend/src:/app/src
|
||||
- ./backend/migrations:/app/migrations
|
||||
- ./backend/alembic.ini:/app/alembic.ini
|
||||
entrypoint: [alembic]
|
||||
command: [upgrade, head]
|
||||
depends_on:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
required: false
|
||||
networks:
|
||||
database:
|
||||
|
||||
volumes:
|
||||
pgdata:
|
||||
|
||||
networks:
|
||||
database:
|
||||
@@ -20,10 +20,15 @@ def _async_url() -> str:
|
||||
return env.db.connection_url.replace("postgresql://", "postgresql+asyncpg://", 1)
|
||||
|
||||
|
||||
def include_name(name: str | None, type_: str, _parents: object) -> bool:
|
||||
return not (type_ == "table" and name is not None and name.startswith("taskiq_"))
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
context.configure(
|
||||
url=_async_url(),
|
||||
target_metadata=target_metadata,
|
||||
include_name=include_name,
|
||||
literal_binds=True,
|
||||
dialect_opts={"paramstyle": "named"},
|
||||
)
|
||||
@@ -32,7 +37,11 @@ def run_migrations_offline() -> None:
|
||||
|
||||
|
||||
def _do_run_migrations(connection) -> None: # noqa: ANN001
|
||||
context.configure(connection=connection, target_metadata=target_metadata)
|
||||
context.configure(
|
||||
connection=connection,
|
||||
target_metadata=target_metadata,
|
||||
include_name=include_name,
|
||||
)
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
"""read path indexes
|
||||
|
||||
Revision ID: a6d2f8b41c9e
|
||||
Revises: f2b8d3c9a51e
|
||||
Create Date: 2026-09-01 12:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "a6d2f8b41c9e"
|
||||
down_revision: str | None = "f2b8d3c9a51e"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("CREATE INDEX ix_jobs_account_id ON jobs (account_id, id DESC)")
|
||||
op.execute(
|
||||
"CREATE INDEX ix_messages_deleted ON messages "
|
||||
"(account_id, deleted_at DESC) WHERE deleted_at IS NOT NULL"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_messages_pinned ON messages "
|
||||
"(account_id, chat_id, date DESC, message_id DESC) "
|
||||
"WHERE raw->>'service' LIKE '%PINNED_MESSAGE%'"
|
||||
)
|
||||
op.execute("CREATE INDEX ix_watches_account ON watches (account_id, id DESC)")
|
||||
op.execute("CREATE INDEX ix_alerts_account_ts ON alerts (account_id, ts DESC)")
|
||||
op.execute(
|
||||
"CREATE INDEX ix_stories_account_date ON stories "
|
||||
"(account_id, date DESC NULLS LAST)"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_annotations_account_msg ON annotations "
|
||||
"(account_id, chat_id, message_id)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP INDEX IF EXISTS ix_annotations_account_msg")
|
||||
op.execute("DROP INDEX IF EXISTS ix_stories_account_date")
|
||||
op.execute("DROP INDEX IF EXISTS ix_alerts_account_ts")
|
||||
op.execute("DROP INDEX IF EXISTS ix_watches_account")
|
||||
op.execute("DROP INDEX IF EXISTS ix_messages_pinned")
|
||||
op.execute("DROP INDEX IF EXISTS ix_messages_deleted")
|
||||
op.execute("DROP INDEX IF EXISTS ix_jobs_account_id")
|
||||
@@ -0,0 +1,25 @@
|
||||
"""retry bot backlog
|
||||
|
||||
Revision ID: a7f3d2b81c04
|
||||
Revises: e2c9a41f7b68
|
||||
Create Date: 2026-09-18 17:30:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "a7f3d2b81c04"
|
||||
down_revision: str | None = "e2c9a41f7b68"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""No drain has succeeded yet, so let every bot try the queue once more."""
|
||||
op.execute("UPDATE watched_bots SET backfilled_at = NULL")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
@@ -0,0 +1,34 @@
|
||||
"""bot backlog cursor
|
||||
|
||||
Revision ID: b6d1f4a82e57
|
||||
Revises: a7f3d2b81c04
|
||||
Create Date: 2026-09-19 10:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b6d1f4a82e57"
|
||||
down_revision: str | None = "a7f3d2b81c04"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(
|
||||
"ALTER TABLE watched_bots "
|
||||
"ADD COLUMN backfill_pts integer, "
|
||||
"ADD COLUMN backfill_qts integer, "
|
||||
"ADD COLUMN backfill_date integer"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute(
|
||||
"ALTER TABLE watched_bots "
|
||||
"DROP COLUMN IF EXISTS backfill_pts, "
|
||||
"DROP COLUMN IF EXISTS backfill_qts, "
|
||||
"DROP COLUMN IF EXISTS backfill_date"
|
||||
)
|
||||
@@ -0,0 +1,85 @@
|
||||
"""bot watcher
|
||||
|
||||
Revision ID: c8f3b6a2d571
|
||||
Revises: a6d2f8b41c9e
|
||||
Create Date: 2026-09-18 12:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "c8f3b6a2d571"
|
||||
down_revision: str | None = "a6d2f8b41c9e"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_BOTS = """
|
||||
CREATE TABLE watched_bots (
|
||||
bot_id bigint PRIMARY KEY,
|
||||
username text,
|
||||
title text,
|
||||
enabled boolean NOT NULL DEFAULT true,
|
||||
status text NOT NULL DEFAULT 'pending',
|
||||
status_detail text,
|
||||
backfilled_at timestamptz,
|
||||
checked_at timestamptz,
|
||||
added_at timestamptz NOT NULL DEFAULT now(),
|
||||
updated_at timestamptz NOT NULL DEFAULT now()
|
||||
)
|
||||
"""
|
||||
|
||||
_UPDATES = """
|
||||
CREATE TABLE bot_updates (
|
||||
bot_id bigint NOT NULL,
|
||||
kind text NOT NULL,
|
||||
chat_id bigint NOT NULL,
|
||||
item_id bigint NOT NULL,
|
||||
ts timestamptz NOT NULL,
|
||||
source text NOT NULL,
|
||||
outgoing boolean NOT NULL DEFAULT false,
|
||||
sender_id bigint,
|
||||
sender_username text,
|
||||
sender_name text,
|
||||
sender_is_bot boolean,
|
||||
chat_type text,
|
||||
chat_title text,
|
||||
text text,
|
||||
media_kind text,
|
||||
edited_at timestamptz,
|
||||
received_at timestamptz NOT NULL DEFAULT now(),
|
||||
raw jsonb NOT NULL,
|
||||
PRIMARY KEY (bot_id, kind, chat_id, item_id, ts)
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(_BOTS)
|
||||
op.execute(_UPDATES)
|
||||
op.execute(
|
||||
"SELECT create_hypertable('bot_updates', by_range('ts', INTERVAL '1 month'))"
|
||||
)
|
||||
op.execute(
|
||||
"ALTER TABLE bot_updates SET ("
|
||||
"timescaledb.enable_columnstore = true, "
|
||||
"timescaledb.segmentby = 'bot_id', "
|
||||
"timescaledb.orderby = 'ts DESC')"
|
||||
)
|
||||
op.execute(
|
||||
"CALL add_columnstore_policy('bot_updates', after => INTERVAL '30 days')"
|
||||
)
|
||||
op.execute("CREATE INDEX ix_bot_updates_recent ON bot_updates (bot_id, ts DESC)")
|
||||
op.execute(
|
||||
"CREATE INDEX ix_bot_updates_sender ON bot_updates "
|
||||
"(bot_id, sender_id, ts DESC) WHERE NOT outgoing"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_bot_updates_chat ON bot_updates (bot_id, chat_id, ts DESC)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP TABLE IF EXISTS bot_updates")
|
||||
op.execute("DROP TABLE IF EXISTS watched_bots")
|
||||
@@ -0,0 +1,57 @@
|
||||
"""media trash
|
||||
|
||||
Revision ID: c9e2a7f4b1d8
|
||||
Revises: b6d1f4a82e57
|
||||
Create Date: 2026-10-14 03:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "c9e2a7f4b1d8"
|
||||
down_revision: str | None = "b6d1f4a82e57"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_TRASH = """
|
||||
CREATE TABLE media_trash (
|
||||
storage_key text PRIMARY KEY,
|
||||
account_id integer NOT NULL,
|
||||
chat_id bigint NOT NULL,
|
||||
message_id bigint NOT NULL,
|
||||
kind text NOT NULL,
|
||||
mime text,
|
||||
file_name text NOT NULL,
|
||||
file_size bigint,
|
||||
trashed_at timestamptz NOT NULL DEFAULT now()
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(
|
||||
"ALTER TABLE media ADD COLUMN evicted_at timestamptz, "
|
||||
"ADD COLUMN evicted_key text"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_media_evicted_unique ON media (unique_id) "
|
||||
"WHERE evicted_at IS NOT NULL"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_media_evicted_key ON media (evicted_key) "
|
||||
"WHERE evicted_key IS NOT NULL"
|
||||
)
|
||||
op.execute(_TRASH)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_media_trash_account ON media_trash "
|
||||
"(account_id, trashed_at DESC)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP TABLE IF EXISTS media_trash")
|
||||
op.execute("DROP INDEX IF EXISTS ix_media_evicted_key")
|
||||
op.execute("DROP INDEX IF EXISTS ix_media_evicted_unique")
|
||||
op.execute("ALTER TABLE media DROP COLUMN evicted_key, DROP COLUMN evicted_at")
|
||||
@@ -0,0 +1,75 @@
|
||||
"""drop scheduled message duplicates
|
||||
|
||||
Revision ID: e1c7a4b62d90
|
||||
Revises: d4a7e2b91f38
|
||||
Create Date: 2026-08-13 10:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "e1c7a4b62d90"
|
||||
down_revision: str | None = "d4a7e2b91f38"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_STAGE = """
|
||||
CREATE TEMP TABLE scheduled_keys ON COMMIT DROP AS
|
||||
SELECT DISTINCT s.account_id, s.chat_id, s.message_id
|
||||
FROM messages s
|
||||
WHERE s.raw->>'scheduled' = 'true'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM messages r
|
||||
WHERE r.account_id = s.account_id AND r.chat_id = s.chat_id
|
||||
AND r.message_id = s.message_id
|
||||
AND r.raw->>'scheduled' IS DISTINCT FROM 'true'
|
||||
)
|
||||
"""
|
||||
|
||||
_DELETE_CHILD = """
|
||||
DELETE FROM {table} t
|
||||
USING scheduled_keys k
|
||||
WHERE t.account_id = k.account_id AND t.chat_id = k.chat_id
|
||||
AND t.message_id = k.message_id
|
||||
"""
|
||||
|
||||
_DELETE_MESSAGES = """
|
||||
DELETE FROM messages m
|
||||
USING scheduled_keys k
|
||||
WHERE m.account_id = k.account_id AND m.chat_id = k.chat_id
|
||||
AND m.message_id = k.message_id AND m.raw->>'scheduled' = 'true'
|
||||
"""
|
||||
|
||||
_RECOUNT = """
|
||||
UPDATE chat_stats cs
|
||||
SET message_count = fresh.message_count
|
||||
FROM (
|
||||
SELECT k.account_id, k.chat_id,
|
||||
(SELECT count(*) FROM messages m
|
||||
WHERE m.account_id = k.account_id AND m.chat_id = k.chat_id)
|
||||
AS message_count
|
||||
FROM (SELECT DISTINCT account_id, chat_id FROM scheduled_keys) k
|
||||
) fresh
|
||||
WHERE cs.account_id = fresh.account_id AND cs.chat_id = fresh.chat_id
|
||||
"""
|
||||
|
||||
_CHILD_TABLES = ("media", "media_versions", "message_versions", "links", "callbacks")
|
||||
|
||||
_ALLOW_DECOMPRESSION = (
|
||||
"SET LOCAL timescaledb.max_tuples_decompressed_per_dml_transaction = 0"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(_ALLOW_DECOMPRESSION)
|
||||
op.execute(_STAGE)
|
||||
for table in _CHILD_TABLES:
|
||||
op.execute(_DELETE_CHILD.format(table=table))
|
||||
op.execute(_DELETE_MESSAGES)
|
||||
op.execute(_RECOUNT)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
@@ -0,0 +1,61 @@
|
||||
"""bot accounts
|
||||
|
||||
Revision ID: e2c9a41f7b68
|
||||
Revises: c8f3b6a2d571
|
||||
Create Date: 2026-09-18 16:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "e2c9a41f7b68"
|
||||
down_revision: str | None = "c8f3b6a2d571"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_ACCOUNTS = """
|
||||
INSERT INTO accounts
|
||||
(tg_user_id, label, phone, session_name, is_active, is_bot, raw)
|
||||
SELECT b.bot_id,
|
||||
coalesce(b.title, '@' || b.username, b.bot_id::text),
|
||||
NULL,
|
||||
b.bot_id::text,
|
||||
b.enabled,
|
||||
true,
|
||||
jsonb_build_object('id', b.bot_id, 'username', b.username,
|
||||
'first_name', b.title)
|
||||
FROM watched_bots b
|
||||
ON CONFLICT (tg_user_id) DO NOTHING
|
||||
"""
|
||||
|
||||
_LINK = """
|
||||
UPDATE watched_bots b SET account_id = a.account_id
|
||||
FROM accounts a WHERE a.tg_user_id = b.bot_id
|
||||
"""
|
||||
|
||||
_MERGE_MESSAGES = """
|
||||
INSERT INTO messages
|
||||
(account_id, chat_id, message_id, date, sender_id, text, raw, has_media)
|
||||
SELECT b.account_id, u.chat_id, u.item_id, u.ts, u.sender_id, u.text, u.raw,
|
||||
u.media_kind IS NOT NULL
|
||||
FROM bot_updates u
|
||||
JOIN watched_bots b ON b.bot_id = u.bot_id
|
||||
WHERE u.kind = 'message' AND b.account_id IS NOT NULL
|
||||
ON CONFLICT (account_id, chat_id, message_id, date) DO NOTHING
|
||||
"""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("ALTER TABLE accounts ADD COLUMN is_bot boolean NOT NULL DEFAULT false")
|
||||
op.execute("ALTER TABLE watched_bots ADD COLUMN account_id integer")
|
||||
op.execute(_ACCOUNTS)
|
||||
op.execute(_LINK)
|
||||
op.execute(_MERGE_MESSAGES)
|
||||
op.execute("DROP TABLE IF EXISTS bot_updates")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("ALTER TABLE watched_bots DROP COLUMN IF EXISTS account_id")
|
||||
op.execute("ALTER TABLE accounts DROP COLUMN IF EXISTS is_bot")
|
||||
@@ -0,0 +1,73 @@
|
||||
"""file shares
|
||||
|
||||
Revision ID: f2b8d3c9a51e
|
||||
Revises: e1c7a4b62d90
|
||||
Create Date: 2026-08-13 10:30:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "f2b8d3c9a51e"
|
||||
down_revision: str | None = "e1c7a4b62d90"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_SHARES = """
|
||||
CREATE TABLE file_shares (
|
||||
id serial PRIMARY KEY,
|
||||
account_id integer NOT NULL,
|
||||
token text NOT NULL UNIQUE,
|
||||
kind text NOT NULL,
|
||||
storage_key text NOT NULL,
|
||||
file_name text NOT NULL,
|
||||
mime text,
|
||||
file_size bigint,
|
||||
title text,
|
||||
chat_id bigint,
|
||||
message_id bigint,
|
||||
peer_id bigint,
|
||||
story_id bigint,
|
||||
expires_at timestamptz,
|
||||
max_downloads integer,
|
||||
download_count integer NOT NULL DEFAULT 0,
|
||||
last_download_at timestamptz,
|
||||
revoked_at timestamptz,
|
||||
created_at timestamptz NOT NULL DEFAULT now()
|
||||
)
|
||||
"""
|
||||
|
||||
_HITS = """
|
||||
CREATE TABLE file_share_hits (
|
||||
id bigserial PRIMARY KEY,
|
||||
share_id integer NOT NULL REFERENCES file_shares (id) ON DELETE CASCADE,
|
||||
ts timestamptz NOT NULL DEFAULT now(),
|
||||
method text NOT NULL,
|
||||
ip text,
|
||||
user_agent text,
|
||||
counted boolean NOT NULL DEFAULT false
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(_SHARES)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_file_shares_account ON file_shares "
|
||||
"(account_id, created_at DESC)"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_file_shares_subject ON file_shares "
|
||||
"(account_id, storage_key) WHERE revoked_at IS NULL"
|
||||
)
|
||||
op.execute(_HITS)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_file_share_hits_share ON file_share_hits (share_id, ts DESC)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP TABLE IF EXISTS file_share_hits")
|
||||
op.execute("DROP TABLE IF EXISTS file_shares")
|
||||
@@ -23,22 +23,23 @@ dependencies = [
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["uv_build>=0.11.14,<0.12.0"]
|
||||
requires = ["uv_build>=0.13.0,<0.14.0"]
|
||||
build-backend = "uv_build"
|
||||
|
||||
[tool.uv.build-backend]
|
||||
module-name = ["userbot", "api", "utils", "dependencies"]
|
||||
module-name = ["userbot", "api", "watcher", "utils", "dependencies"]
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py313"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["ALL"]
|
||||
ignore = ["CPY", "D1", "D203", "D212", "COM812", "BLE", "S101", "PT", "RUF001", "RUF002", "RUF003", "TRY301"]
|
||||
ignore = ["CPY", "D1", "D203", "D212", "COM812", "BLE", "S101", "PT", "PLR0917", "RUF001", "RUF002", "RUF003", "TRY301"]
|
||||
unfixable = ["F401"]
|
||||
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
"scripts/*" = ["INP"]
|
||||
"scripts/check_imports.py" = ["T201"]
|
||||
"t/*" = ["ALL"]
|
||||
"*.lock" = ["ALL"]
|
||||
"migrations/*" = ["INP", "D", "E501"]
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
HEAVY = ["aiogram", "pyrogram", "fastapi", "google.genai"]
|
||||
SERVICES = {
|
||||
"api": (["api.app"], {"fastapi", "pyrogram"}),
|
||||
"userbot": (["userbot.runner", "userbot.handlers"], {"pyrogram"}),
|
||||
"watcher": (["watcher.runner"], {"pyrogram"}),
|
||||
}
|
||||
PROBE = """
|
||||
import importlib, sys
|
||||
for module in sys.argv[1].split(","):
|
||||
importlib.import_module(module)
|
||||
print(" ".join(name for name in sys.argv[2:] if name in sys.modules))
|
||||
"""
|
||||
ENV = {**os.environ, "AUTH__TOKEN": "check"}
|
||||
|
||||
|
||||
def loaded(modules: list[str]) -> tuple[set[str], str]:
|
||||
result = subprocess.run( # noqa: S603
|
||||
[sys.executable, "-c", PROBE, ",".join(modules), *HEAVY],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd="src",
|
||||
env=ENV,
|
||||
check=False,
|
||||
)
|
||||
error = (result.stderr.strip().splitlines() or ["exit code"])[-1]
|
||||
return set(result.stdout.split()) & set(HEAVY), error if result.returncode else ""
|
||||
|
||||
|
||||
def main() -> int:
|
||||
failed = False
|
||||
for service, (modules, allowed) in SERVICES.items():
|
||||
found, error = loaded(modules)
|
||||
extra = sorted(found - allowed)
|
||||
status = error or ("imports " + ", ".join(extra) if extra else "")
|
||||
print(f"{service:10s} {'FAIL: ' + status if status else 'ok'}")
|
||||
failed |= bool(status)
|
||||
return int(failed)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,23 @@
|
||||
import logging
|
||||
import re
|
||||
|
||||
QUERYLESS_PREFIXES = ("/mcp",)
|
||||
_PATH_ARG = 2
|
||||
_TOKEN = re.compile(r"(?<=[?&]token=)[^&]+")
|
||||
|
||||
|
||||
class ScrubQueryFilter(logging.Filter):
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
args = record.args
|
||||
if isinstance(args, tuple) and len(args) > _PATH_ARG:
|
||||
path = str(args[_PATH_ARG])
|
||||
if path.startswith(QUERYLESS_PREFIXES):
|
||||
path = path.partition("?")[0]
|
||||
else:
|
||||
path = _TOKEN.sub("***", path)
|
||||
record.args = (*args[:_PATH_ARG], path, *args[_PATH_ARG + 1 :])
|
||||
return True
|
||||
|
||||
|
||||
def scrub_access_log() -> None:
|
||||
logging.getLogger("uvicorn.access").addFilter(ScrubQueryFilter())
|
||||
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
@@ -9,6 +10,7 @@ from fastapi.responses import FileResponse, RedirectResponse
|
||||
from fastmcp.utilities.lifespan import combine_lifespans
|
||||
from starlette.applications import Starlette
|
||||
|
||||
from api.access_log import scrub_access_log
|
||||
from api.auth import BearerAuthMiddleware
|
||||
from api.mcp.server import mcp
|
||||
from api.realtime import hub
|
||||
@@ -18,10 +20,12 @@ from api.routers import (
|
||||
annotations,
|
||||
avatars,
|
||||
backfill,
|
||||
bots,
|
||||
chats,
|
||||
custom_emoji,
|
||||
discover,
|
||||
events,
|
||||
files,
|
||||
folders,
|
||||
media,
|
||||
peers,
|
||||
@@ -29,19 +33,25 @@ from api.routers import (
|
||||
presence,
|
||||
profile,
|
||||
search,
|
||||
shares,
|
||||
social,
|
||||
storage,
|
||||
stories,
|
||||
watches,
|
||||
)
|
||||
from dependencies.container import container
|
||||
from utils.cache import DAY_HEADERS, IMMUTABLE_HEADERS, NO_STORE_HEADERS
|
||||
from utils.env import env
|
||||
from utils.storage import ContentAddressedStorage
|
||||
from utils.trash import purge_loop
|
||||
|
||||
if env.auth.token is None:
|
||||
msg = "AUTH__TOKEN is required for the api process"
|
||||
raise RuntimeError(msg)
|
||||
_token = env.auth.token.get_secret_value()
|
||||
|
||||
scrub_access_log()
|
||||
|
||||
mcp_app = mcp.http_app(path="/")
|
||||
|
||||
|
||||
@@ -49,7 +59,11 @@ mcp_app = mcp.http_app(path="/")
|
||||
async def lifespan(app_: Starlette) -> AsyncGenerator[None]:
|
||||
pool = await container.get(asyncpg.Pool)
|
||||
await hub.start(pool)
|
||||
purger = asyncio.create_task(
|
||||
purge_loop(pool, await container.get(ContentAddressedStorage))
|
||||
)
|
||||
yield
|
||||
purger.cancel()
|
||||
await hub.stop()
|
||||
await app_.state.dishka_container.close()
|
||||
|
||||
@@ -71,6 +85,7 @@ async def health(pool: FromDishka[asyncpg.Pool]) -> dict[str, bool]:
|
||||
|
||||
app.include_router(accounts.router)
|
||||
app.include_router(analytics.router)
|
||||
app.include_router(bots.router)
|
||||
app.include_router(policy.router)
|
||||
app.include_router(folders.router)
|
||||
app.include_router(backfill.router)
|
||||
@@ -88,6 +103,9 @@ app.include_router(peers.router)
|
||||
app.include_router(discover.router)
|
||||
app.include_router(annotations.router)
|
||||
app.include_router(watches.router)
|
||||
app.include_router(shares.router)
|
||||
app.include_router(storage.router)
|
||||
app.include_router(files.router)
|
||||
|
||||
app.mount("/mcp", mcp_app)
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ from pyrogram.errors import AuthTokenExpired, SessionPasswordNeeded
|
||||
from pyrogram.qrlogin import QRLogin
|
||||
from pyrogram.types import User
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.client import PyroClient
|
||||
from utils.env import env
|
||||
|
||||
LOGIN_TTL_SECONDS = 900
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.server.dependencies import get_http_request
|
||||
from fastmcp.utilities.types import Image
|
||||
from mcp.types import TextContent
|
||||
|
||||
from utils.files import expiry_from_seconds
|
||||
from utils.jobs import enqueue
|
||||
from utils.read import chats, media, shares
|
||||
from utils.read.models import FileShareView, MediaView
|
||||
from utils.storage import ContentAddressedStorage
|
||||
|
||||
LINK_TTL_SECONDS = 24 * 3600
|
||||
MAX_WAIT_SECONDS = 120
|
||||
_POLL_SECONDS = 1.0
|
||||
_INLINE_IMAGE_BYTES = 5 * 1024 * 1024
|
||||
_INLINE_TEXT_BYTES = 256 * 1024
|
||||
_INLINE_IMAGE_MIMES = {"image/jpeg", "image/png", "image/gif", "image/webp"}
|
||||
_TEXT_MIMES = {"application/json", "application/xml", "application/x-yaml"}
|
||||
_FAILED_JOB = ("failed", "canceled")
|
||||
|
||||
|
||||
async def _job_error(pool: asyncpg.Pool, job_id: int) -> str | None:
|
||||
row = await pool.fetchrow("SELECT status, error FROM jobs WHERE id = $1", job_id)
|
||||
if row is None or row["status"] in _FAILED_JOB:
|
||||
return (row and row["error"]) or "download job is gone"
|
||||
return None
|
||||
|
||||
|
||||
async def fetch_media(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
message_id: int,
|
||||
wait_seconds: float,
|
||||
) -> MediaView | None:
|
||||
item = await media.get_message_media(pool, account_id, chat_id, message_id)
|
||||
if item is None:
|
||||
msg = f"message {chat_id}/{message_id} has no media"
|
||||
raise ToolError(msg)
|
||||
if item.downloaded and item.storage_key:
|
||||
return item
|
||||
job_id = await enqueue(
|
||||
pool, account_id, "fetch_media", {"chat_id": chat_id, "message_id": message_id}
|
||||
)
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + wait_seconds
|
||||
while loop.time() < deadline:
|
||||
await asyncio.sleep(_POLL_SECONDS)
|
||||
item = await media.get_message_media(pool, account_id, chat_id, message_id)
|
||||
if item is not None and item.downloaded and item.storage_key:
|
||||
return item
|
||||
error = await _job_error(pool, job_id)
|
||||
if error is not None:
|
||||
msg = f"download failed: {error}"
|
||||
raise ToolError(msg)
|
||||
return None
|
||||
|
||||
|
||||
def _origin() -> str:
|
||||
request = get_http_request()
|
||||
proto = request.headers.get("x-forwarded-proto", request.url.scheme)
|
||||
return f"{proto}://{request.url.netloc}"
|
||||
|
||||
|
||||
async def share_media(pool: asyncpg.Pool, item: MediaView, mime: str) -> FileShareView:
|
||||
storage_key = item.storage_key or ""
|
||||
share = await shares.find_active_share(pool, item.account_id, storage_key)
|
||||
if share is None:
|
||||
chat = await chats.get_chat(pool, item.account_id, item.chat_id)
|
||||
share = await shares.create_share(
|
||||
pool,
|
||||
item.account_id,
|
||||
"media",
|
||||
storage_key,
|
||||
item.file_name or f"media_{item.id}",
|
||||
mime=mime,
|
||||
file_size=item.file_size,
|
||||
title=chat.title if chat else None,
|
||||
chat_id=item.chat_id,
|
||||
message_id=item.message_id,
|
||||
expires_at=expiry_from_seconds(LINK_TTL_SECONDS),
|
||||
)
|
||||
return share
|
||||
|
||||
|
||||
def _is_text(mime: str) -> bool:
|
||||
return mime.startswith("text/") or mime in _TEXT_MIMES
|
||||
|
||||
|
||||
def inline_content(storage: ContentAddressedStorage, key: str, mime: str) -> list[Any]:
|
||||
size = Path(storage.url(key)).stat().st_size
|
||||
if mime in _INLINE_IMAGE_MIMES and size <= _INLINE_IMAGE_BYTES:
|
||||
return [Image(data=storage.get(key), format=mime.rsplit("/", maxsplit=1)[-1])]
|
||||
if _is_text(mime) and size <= _INLINE_TEXT_BYTES:
|
||||
text = storage.get(key).decode("utf-8", errors="replace")
|
||||
return [TextContent(type="text", text=text)]
|
||||
return []
|
||||
|
||||
|
||||
def describe(item: MediaView, mime: str, share: FileShareView) -> str:
|
||||
size = f"{item.file_size} bytes" if item.file_size else "size unknown"
|
||||
expires = share.expires_at.isoformat() if share.expires_at else "never"
|
||||
return (
|
||||
f"{item.file_name} ({item.kind}, {mime}, {size})\n"
|
||||
f"Download: {_origin()}/f/{share.token}/{share.url_name}\n"
|
||||
f"Public link, expires {expires}; fetch it with curl/wget or any HTTP "
|
||||
"client, no auth needed."
|
||||
)
|
||||
@@ -4,13 +4,24 @@ from typing import Any
|
||||
|
||||
import asyncpg
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
|
||||
from fastmcp.tools.base import ToolResult
|
||||
from mcp.types import CallToolRequestParams, TextContent
|
||||
from pydantic import BaseModel
|
||||
|
||||
from api.mcp.download import (
|
||||
MAX_WAIT_SECONDS,
|
||||
describe,
|
||||
fetch_media,
|
||||
inline_content,
|
||||
share_media,
|
||||
)
|
||||
from api.mcp.format import build_transcript, load_notes, load_photos, resolve_names
|
||||
from dependencies.container import container
|
||||
from utils.jobs import enqueue
|
||||
from utils.files import resolve_mime
|
||||
from utils.read import annotations, chats, media, peers, presence, social, watches
|
||||
from utils.read.accounts import self_user_id
|
||||
from utils.read.accounts import self_user_id, unknown_account_hint
|
||||
from utils.read.models import DEFAULT_LIMIT, Page
|
||||
from utils.search.models import SearchFilters
|
||||
from utils.search.repository import search_messages
|
||||
@@ -23,6 +34,9 @@ peer history) into Postgres and exposes it read-only over these tools.
|
||||
Account scoping:
|
||||
- Every tool takes `account_id`. It selects which archived Telegram account to
|
||||
read. Unless the user names a different one, always pass `account_id=1`.
|
||||
- `account_id` is beavergram's own small serial (1, 2, ...), never the archived
|
||||
account's Telegram user id. Passing a Telegram id returns an error naming the
|
||||
right `account_id`.
|
||||
|
||||
Identifiers:
|
||||
- `chat_id` and `peer_id` are Telegram IDs (negative for groups/channels,
|
||||
@@ -37,7 +51,7 @@ Typical flows:
|
||||
`before_id` (first id you saw), to newer with `after_id` (last id you saw);
|
||||
begin a full forward walk with `after_id=0`. Each transcript line starts
|
||||
with its `#message_id`, and photos show as `[photo #message_id]` — pass that
|
||||
id to `get_media` (with `fetch=True` to download) or to annotation tools.
|
||||
id to `download_media` or to annotation tools.
|
||||
Lines marked
|
||||
"📝 [your private note, NOT in Telegram]" are the user's own annotations
|
||||
attached locally via the web UI; they never existed in Telegram. Treat them
|
||||
@@ -46,12 +60,21 @@ Typical flows:
|
||||
transcripts; supports chat/sender/date filters and regex).
|
||||
- Forensics: `get_deleted_messages` and `get_message_versions` recover content
|
||||
removed or edited in Telegram but kept in the archive.
|
||||
- Media: `get_media` returns metadata; pass `fetch=True` to enqueue a lazy
|
||||
download if the file isn't stored yet.
|
||||
- Files: `download_media` gets any message attachment (photo, video, voice,
|
||||
document, ...). It pulls the file from Telegram if the archive lacks it and
|
||||
returns a public link valid 24h that needs no auth (`curl -O <link>`), plus
|
||||
the content itself for small images and text files. `get_media` returns
|
||||
metadata only.
|
||||
- Monitoring: `set_watch` creates a local rule, `list_watches` lists rules, and
|
||||
`list_alerts` reads what those rules fired.
|
||||
|
||||
Writes: everything is read-only except `set_watch`, the only allowed write.
|
||||
Watched bots:
|
||||
- a bot whose token the user added is just another archived account: it has its
|
||||
own `account_id` and its chats, messages and media read through the same
|
||||
tools as a user account. Nothing bot-specific to call.
|
||||
|
||||
Writes: everything is read-only except `set_watch` and the share links
|
||||
`download_media` creates (visible and revocable in the web UI).
|
||||
|
||||
Paging: list/search tools accept `limit` and `offset`.
|
||||
"""
|
||||
@@ -63,6 +86,25 @@ async def _pool() -> asyncpg.Pool:
|
||||
return await container.get(asyncpg.Pool)
|
||||
|
||||
|
||||
class AccountGuard(Middleware):
|
||||
"""Reject tool calls scoped to an account_id that does not exist."""
|
||||
|
||||
async def on_call_tool(
|
||||
self,
|
||||
context: MiddlewareContext[CallToolRequestParams],
|
||||
call_next: CallNext[CallToolRequestParams, ToolResult],
|
||||
) -> ToolResult:
|
||||
account_id = (context.message.arguments or {}).get("account_id")
|
||||
if isinstance(account_id, int) and not isinstance(account_id, bool):
|
||||
hint = await unknown_account_hint(await _pool(), account_id)
|
||||
if hint is not None:
|
||||
raise ToolError(hint)
|
||||
return await call_next(context)
|
||||
|
||||
|
||||
mcp.add_middleware(AccountGuard())
|
||||
|
||||
|
||||
def _dump(items: Sequence[BaseModel]) -> list[dict[str, Any]]:
|
||||
return [item.model_dump(mode="json") for item in items]
|
||||
|
||||
@@ -208,21 +250,40 @@ async def get_message_versions(
|
||||
|
||||
@mcp.tool
|
||||
async def get_media(
|
||||
account_id: int, chat_id: int, message_id: int, fetch: bool = False
|
||||
account_id: int, chat_id: int, message_id: int
|
||||
) -> dict[str, Any] | None:
|
||||
"""Get media metadata for a message; set fetch=True to enqueue lazy download."""
|
||||
"""Get media metadata for a message; use download_media to get the file."""
|
||||
item = await media.get_message_media(await _pool(), account_id, chat_id, message_id)
|
||||
return item.model_dump(mode="json") if item else None
|
||||
|
||||
|
||||
@mcp.tool(output_schema=None)
|
||||
async def download_media(
|
||||
account_id: int, chat_id: int, message_id: int, wait_seconds: int = 45
|
||||
) -> list[Any]:
|
||||
"""Download the file attached to a message: any photo, video, voice, document.
|
||||
|
||||
Pulls it from Telegram first if the archive doesn't have it yet, waiting up
|
||||
to `wait_seconds`. Returns a public download link (valid 24h, no auth) and
|
||||
inlines small images and text files.
|
||||
"""
|
||||
pool = await _pool()
|
||||
item = await media.get_message_media(pool, account_id, chat_id, message_id)
|
||||
if item is None:
|
||||
return None
|
||||
if fetch and not item.downloaded:
|
||||
await enqueue(
|
||||
pool,
|
||||
account_id,
|
||||
"fetch_media",
|
||||
{"chat_id": chat_id, "message_id": message_id},
|
||||
)
|
||||
return item.model_dump(mode="json")
|
||||
wait = min(max(wait_seconds, 0), MAX_WAIT_SECONDS)
|
||||
item = await fetch_media(pool, account_id, chat_id, message_id, wait)
|
||||
if item is None or item.storage_key is None:
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text="Still downloading from Telegram; call download_media again.",
|
||||
)
|
||||
]
|
||||
mime = resolve_mime(item.kind, item.mime, item.file_name)
|
||||
share = await share_media(pool, item, mime)
|
||||
storage = await container.get(ContentAddressedStorage)
|
||||
return [
|
||||
TextContent(type="text", text=describe(item, mime, share)),
|
||||
*inline_content(storage, item.storage_key, mime),
|
||||
]
|
||||
|
||||
|
||||
@mcp.tool
|
||||
|
||||
@@ -9,7 +9,7 @@ from pyrogram.errors import FloodWait, RPCError
|
||||
from pyrogram.types import User
|
||||
|
||||
from api.login import LoginError, login_manager
|
||||
from userbot import DEVICE_MODEL_LIMIT
|
||||
from userbot.modules.client import DEVICE_MODEL_LIMIT
|
||||
from utils.read import accounts
|
||||
from utils.read.models import AccountView
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ from fastapi import APIRouter, HTTPException, Query
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils.jobs import enqueue
|
||||
from utils.read.models import DEFAULT_LIMIT, Page
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["backfill"], route_class=DishkaRoute)
|
||||
|
||||
@@ -130,16 +131,26 @@ async def list_jobs(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
status: Annotated[str | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[JobView]:
|
||||
page = Page(limit=limit, offset=offset)
|
||||
if status is None:
|
||||
rows = await pool.fetch(
|
||||
"SELECT * FROM jobs WHERE account_id = $1 ORDER BY id DESC", account_id
|
||||
"SELECT * FROM jobs WHERE account_id = $1 "
|
||||
"ORDER BY id DESC LIMIT $2 OFFSET $3",
|
||||
account_id,
|
||||
page.capped_limit,
|
||||
page.offset,
|
||||
)
|
||||
else:
|
||||
rows = await pool.fetch(
|
||||
"SELECT * FROM jobs WHERE account_id = $1 AND status = $2 ORDER BY id DESC",
|
||||
"SELECT * FROM jobs WHERE account_id = $1 AND status = $2 "
|
||||
"ORDER BY id DESC LIMIT $3 OFFSET $4",
|
||||
account_id,
|
||||
status,
|
||||
page.capped_limit,
|
||||
page.offset,
|
||||
)
|
||||
return [_to_view(row) for row in rows]
|
||||
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils.read import bots
|
||||
from utils.read.models import BotView
|
||||
from watcher import tokens
|
||||
|
||||
router = APIRouter(prefix="/api/bots", tags=["bots"], route_class=DishkaRoute)
|
||||
|
||||
INVALID_TOKEN = "token must look like 123456:ABC-DEF..." # noqa: S105
|
||||
|
||||
|
||||
class BotAdd(BaseModel):
|
||||
token: str
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def list_bots(pool: FromDishka[asyncpg.Pool]) -> list[BotView]:
|
||||
return await bots.list_bots(pool)
|
||||
|
||||
|
||||
@router.post("", status_code=201)
|
||||
async def add_bot(pool: FromDishka[asyncpg.Pool], body: BotAdd) -> BotView:
|
||||
bot_id = tokens.bot_id_of(body.token)
|
||||
if bot_id is None:
|
||||
raise HTTPException(status_code=422, detail=INVALID_TOKEN)
|
||||
tokens.save(bot_id, body.token)
|
||||
await bots.register_bot(pool, bot_id)
|
||||
await bots.notify_bots_changed(pool)
|
||||
bot = await bots.get_bot(pool, bot_id)
|
||||
if bot is None:
|
||||
raise HTTPException(status_code=500, detail="bot was not registered")
|
||||
return bot
|
||||
|
||||
|
||||
@router.delete("/{bot_id}", status_code=204)
|
||||
async def remove_bot(pool: FromDishka[asyncpg.Pool], bot_id: int) -> None:
|
||||
if not await bots.remove_bot(pool, bot_id):
|
||||
raise HTTPException(status_code=404, detail="unknown bot")
|
||||
await bots.notify_bots_changed(pool)
|
||||
@@ -10,6 +10,7 @@ from utils.policy import repository
|
||||
from utils.read import chats
|
||||
from utils.read.models import (
|
||||
DEFAULT_LIMIT,
|
||||
MAX_LIMIT,
|
||||
ChatListItem,
|
||||
MessageVersionView,
|
||||
MessageView,
|
||||
@@ -49,6 +50,14 @@ async def list_chats(
|
||||
)
|
||||
|
||||
|
||||
@router.get("/chats/batch")
|
||||
async def get_chats_batch(
|
||||
pool: FromDishka[asyncpg.Pool], account_id: AccountId, ids: Annotated[str, Query()]
|
||||
) -> list[ChatListItem]:
|
||||
parsed = [int(part) for part in ids.split(",") if part.strip()]
|
||||
return await chats.get_chats(pool, account_id, parsed[:MAX_LIMIT])
|
||||
|
||||
|
||||
@router.get("/chats/{chat_id}")
|
||||
async def get_chat(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
|
||||
|
||||
@@ -13,7 +13,7 @@ from utils.policy import repository as policy_repository
|
||||
from utils.policy.defaults import TRACKING
|
||||
from utils.policy.models import ScopeType
|
||||
from utils.read import discover
|
||||
from utils.read.models import DiscoverItem
|
||||
from utils.read.models import MAX_LIMIT, DiscoverItem
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["discover"], route_class=DishkaRoute)
|
||||
|
||||
@@ -68,6 +68,7 @@ async def discover_peers(
|
||||
) -> list[DiscoverItem]:
|
||||
if not query.strip():
|
||||
return []
|
||||
limit = min(limit, MAX_LIMIT)
|
||||
items = await discover.search(pool, account_id, query, limit)
|
||||
if not remote:
|
||||
return items
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
from datetime import UTC, datetime
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, Query, Request
|
||||
from fastapi.responses import FileResponse, PlainTextResponse
|
||||
|
||||
from utils.files import (
|
||||
content_disposition,
|
||||
counts_as_download,
|
||||
is_active_mime,
|
||||
is_inline_mime,
|
||||
resolve_mime,
|
||||
)
|
||||
from utils.read import shares
|
||||
from utils.storage import ContentAddressedStorage
|
||||
|
||||
router = APIRouter(tags=["files"], route_class=DishkaRoute)
|
||||
|
||||
_GONE = "This link is no longer available.\n"
|
||||
_MISSING = "Not found.\n"
|
||||
|
||||
_NO_STORE = {
|
||||
"Cache-Control": "private, no-store",
|
||||
"X-Robots-Tag": "noindex, nofollow",
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
}
|
||||
_SANDBOX = {"Content-Security-Policy": "sandbox"}
|
||||
|
||||
|
||||
def _client_ip(request: Request) -> str | None:
|
||||
forwarded = request.headers.get("x-forwarded-for")
|
||||
if forwarded:
|
||||
return forwarded.split(",")[0].strip()
|
||||
return request.client.host if request.client else None
|
||||
|
||||
|
||||
def _expired(row: asyncpg.Record) -> bool:
|
||||
expires_at = row["expires_at"]
|
||||
return expires_at is not None and expires_at <= datetime.now(UTC)
|
||||
|
||||
|
||||
def _exhausted(row: asyncpg.Record) -> bool:
|
||||
limit = row["max_downloads"]
|
||||
return limit is not None and row["download_count"] >= limit
|
||||
|
||||
|
||||
async def _serve(
|
||||
request: Request,
|
||||
pool: asyncpg.Pool,
|
||||
storage: ContentAddressedStorage,
|
||||
token: str,
|
||||
*,
|
||||
dl: bool,
|
||||
) -> FileResponse | PlainTextResponse:
|
||||
row = await shares.share_by_token(pool, token)
|
||||
if row is None:
|
||||
return PlainTextResponse(_MISSING, status_code=404, headers=_NO_STORE)
|
||||
|
||||
method = request.method
|
||||
agent = request.headers.get("user-agent")
|
||||
ip = _client_ip(request)
|
||||
counts = counts_as_download(method, agent, request.headers.get("range"))
|
||||
|
||||
if row["revoked_at"] is not None or _expired(row) or _exhausted(row):
|
||||
await shares.record_hit(pool, row["id"], method, ip, agent, counted=False)
|
||||
return PlainTextResponse(_GONE, status_code=410, headers=_NO_STORE)
|
||||
|
||||
if not storage.exists(row["storage_key"]):
|
||||
return PlainTextResponse(_MISSING, status_code=404, headers=_NO_STORE)
|
||||
|
||||
if counts and not await shares.consume_download(pool, row["id"]):
|
||||
await shares.record_hit(pool, row["id"], method, ip, agent, counted=False)
|
||||
return PlainTextResponse(_GONE, status_code=410, headers=_NO_STORE)
|
||||
|
||||
await shares.record_hit(pool, row["id"], method, ip, agent, counted=counts)
|
||||
mime = resolve_mime(row["kind"], row["mime"], row["file_name"])
|
||||
attachment = dl or not is_inline_mime(mime)
|
||||
return FileResponse(
|
||||
storage.url(row["storage_key"]),
|
||||
media_type=mime,
|
||||
headers={
|
||||
**_NO_STORE,
|
||||
**(_SANDBOX if is_active_mime(mime) else {}),
|
||||
"Content-Disposition": content_disposition(
|
||||
row["file_name"], attachment=attachment
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.api_route("/f/{token}", methods=["GET", "HEAD"], response_model=None)
|
||||
async def serve_shared_file(
|
||||
request: Request,
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
token: str,
|
||||
dl: Annotated[bool, Query()] = False,
|
||||
) -> FileResponse | PlainTextResponse:
|
||||
return await _serve(request, pool, storage, token, dl=dl)
|
||||
|
||||
|
||||
@router.api_route("/f/{token}/{name}", methods=["GET", "HEAD"], response_model=None)
|
||||
async def serve_shared_file_named(
|
||||
request: Request,
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
token: str,
|
||||
name: str, # noqa: ARG001
|
||||
dl: Annotated[bool, Query()] = False,
|
||||
) -> FileResponse | PlainTextResponse:
|
||||
return await _serve(request, pool, storage, token, dl=dl)
|
||||
@@ -6,6 +6,7 @@ from fastapi import APIRouter, HTTPException, Query
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from utils.cache import DAY_HEADERS, IMMUTABLE_HEADERS
|
||||
from utils.files import content_disposition, media_file_name, resolve_mime
|
||||
from utils.read.media import (
|
||||
get_media,
|
||||
get_media_version,
|
||||
@@ -17,6 +18,15 @@ from utils.storage import ContentAddressedStorage
|
||||
|
||||
router = APIRouter(prefix="/api/media", tags=["media"], route_class=DishkaRoute)
|
||||
|
||||
Download = Annotated[bool, Query()]
|
||||
|
||||
|
||||
def _attachment(headers: dict[str, str], file_name: str) -> dict[str, str]:
|
||||
return {
|
||||
**headers,
|
||||
"Content-Disposition": content_disposition(file_name, attachment=True),
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{media_id}/meta")
|
||||
async def media_meta(pool: FromDishka[asyncpg.Pool], media_id: int) -> MediaView:
|
||||
@@ -41,14 +51,20 @@ async def serve_media_version(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
version_id: int,
|
||||
download: Download = False,
|
||||
) -> FileResponse:
|
||||
version = await get_media_version(pool, version_id)
|
||||
if version is None:
|
||||
if version is None or not storage.exists(version.storage_key):
|
||||
raise HTTPException(status_code=404, detail="media version not found")
|
||||
headers = IMMUTABLE_HEADERS
|
||||
if download:
|
||||
headers = _attachment(
|
||||
headers, media_file_name(version.kind, version.mime, version_id)
|
||||
)
|
||||
return FileResponse(
|
||||
storage.url(version.storage_key),
|
||||
media_type=version.mime or "application/octet-stream",
|
||||
headers=IMMUTABLE_HEADERS,
|
||||
media_type=resolve_mime(version.kind, version.mime),
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
||||
@@ -70,6 +86,7 @@ async def serve_media(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
media_id: int,
|
||||
download: Download = False,
|
||||
) -> FileResponse:
|
||||
media = await get_media(pool, media_id)
|
||||
if media is None:
|
||||
@@ -79,8 +96,13 @@ async def serve_media(
|
||||
status_code=409,
|
||||
detail="media not downloaded; enqueue fetch via POST /api/media/fetch",
|
||||
)
|
||||
if not storage.exists(media.storage_key):
|
||||
raise HTTPException(status_code=404, detail="media file is missing")
|
||||
headers = DAY_HEADERS
|
||||
if download:
|
||||
headers = _attachment(headers, media.file_name or f"media_{media_id}")
|
||||
return FileResponse(
|
||||
storage.url(media.storage_key),
|
||||
media_type=media.mime or "application/octet-stream",
|
||||
headers=DAY_HEADERS,
|
||||
media_type=resolve_mime(media.kind, media.mime, media.file_name),
|
||||
headers=headers,
|
||||
)
|
||||
@@ -5,13 +5,15 @@ from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils.policy import repository
|
||||
from utils.policy import overview, repository
|
||||
from utils.policy.models import (
|
||||
CaptureToggles,
|
||||
ChatKind,
|
||||
ChatMeta,
|
||||
PolicyCreate,
|
||||
PolicyRecord,
|
||||
PolicyRuleView,
|
||||
PolicyScopeView,
|
||||
)
|
||||
from utils.policy.resolver import resolve
|
||||
|
||||
@@ -50,6 +52,32 @@ async def effective_policy(
|
||||
return resolve(chat, folders, policies)
|
||||
|
||||
|
||||
@router.get("/overview")
|
||||
async def policy_overview(
|
||||
pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()]
|
||||
) -> list[PolicyRuleView]:
|
||||
return await overview.list_rules(pool, account_id)
|
||||
|
||||
|
||||
@router.get("/scope")
|
||||
async def policy_scope(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
chat_id: Annotated[int | None, Query()] = None,
|
||||
folder_id: Annotated[int | None, Query()] = None,
|
||||
) -> PolicyScopeView:
|
||||
if chat_id is not None:
|
||||
return await overview.chat_scope(pool, account_id, chat_id)
|
||||
scope = (
|
||||
await overview.folder_scope(pool, account_id, folder_id)
|
||||
if folder_id is not None
|
||||
else None
|
||||
)
|
||||
if scope is None:
|
||||
raise HTTPException(status_code=404, detail="scope not found")
|
||||
return scope
|
||||
|
||||
|
||||
@router.post("", status_code=201)
|
||||
async def create_policy(
|
||||
pool: FromDishka[asyncpg.Pool], body: PolicyCreate
|
||||
|
||||
@@ -52,9 +52,15 @@ async def chat_links(
|
||||
|
||||
@router.get("/calendar")
|
||||
async def chat_calendar(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
chat_id: int,
|
||||
account_id: AccountId,
|
||||
date_from: Annotated[datetime | None, Query()] = None,
|
||||
date_to: Annotated[datetime | None, Query()] = None,
|
||||
) -> list[DayCount]:
|
||||
return await profile.daily_counts(pool, account_id, chat_id)
|
||||
return await profile.daily_counts(
|
||||
pool, account_id, chat_id, date_from=date_from, date_to=date_to
|
||||
)
|
||||
|
||||
|
||||
@router.get("/message-at")
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils.files import (
|
||||
expiry_from_seconds,
|
||||
media_file_name,
|
||||
resolve_mime,
|
||||
story_file_name,
|
||||
)
|
||||
from utils.read import chats, peers, shares
|
||||
from utils.read.media import get_media, get_media_version
|
||||
from utils.read.models import DEFAULT_LIMIT, FileShareHitView, FileShareView, Page
|
||||
|
||||
router = APIRouter(prefix="/api/shares", tags=["shares"], route_class=DishkaRoute)
|
||||
|
||||
|
||||
class ShareCreate(BaseModel):
|
||||
account_id: int
|
||||
kind: str
|
||||
media_id: int | None = None
|
||||
version_id: int | None = None
|
||||
peer_id: int | None = None
|
||||
story_id: int | None = None
|
||||
expires_in_seconds: int | None = None
|
||||
max_downloads: int | None = None
|
||||
|
||||
|
||||
class ShareUpdate(BaseModel):
|
||||
expires_in_seconds: int | None = None
|
||||
max_downloads: int | None = None
|
||||
keep_expiry: bool = False
|
||||
|
||||
|
||||
class Subject(BaseModel):
|
||||
storage_key: str
|
||||
file_name: str
|
||||
mime: str | None
|
||||
file_size: int | None
|
||||
title: str | None
|
||||
chat_id: int | None = None
|
||||
message_id: int | None = None
|
||||
peer_id: int | None = None
|
||||
story_id: int | None = None
|
||||
|
||||
|
||||
_NOT_STORED = "file is not downloaded yet"
|
||||
|
||||
|
||||
async def _chat_title(pool: asyncpg.Pool, account_id: int, chat_id: int) -> str | None:
|
||||
chat = await chats.get_chat(pool, account_id, chat_id)
|
||||
return chat.title if chat else None
|
||||
|
||||
|
||||
async def _media_subject(pool: asyncpg.Pool, media_id: int) -> Subject:
|
||||
media = await get_media(pool, media_id)
|
||||
if media is None:
|
||||
raise HTTPException(status_code=404, detail="media not found")
|
||||
if not media.downloaded or media.storage_key is None:
|
||||
raise HTTPException(status_code=409, detail=_NOT_STORED)
|
||||
file_name = media.file_name or f"media_{media.id}"
|
||||
return Subject(
|
||||
storage_key=media.storage_key,
|
||||
file_name=file_name,
|
||||
mime=resolve_mime(media.kind, media.mime, file_name),
|
||||
file_size=media.file_size,
|
||||
title=await _chat_title(pool, media.account_id, media.chat_id),
|
||||
chat_id=media.chat_id,
|
||||
message_id=media.message_id,
|
||||
)
|
||||
|
||||
|
||||
async def _version_subject(pool: asyncpg.Pool, version_id: int) -> Subject:
|
||||
version = await get_media_version(pool, version_id)
|
||||
if version is None:
|
||||
raise HTTPException(status_code=404, detail="media version not found")
|
||||
version_name = media_file_name(version.kind, version.mime, version_id)
|
||||
return Subject(
|
||||
storage_key=version.storage_key,
|
||||
file_name=version_name,
|
||||
mime=resolve_mime(version.kind, version.mime, version_name),
|
||||
file_size=version.file_size,
|
||||
title=None,
|
||||
chat_id=None,
|
||||
message_id=None,
|
||||
)
|
||||
|
||||
|
||||
async def _story_subject(
|
||||
pool: asyncpg.Pool, account_id: int, peer_id: int, story_id: int
|
||||
) -> Subject:
|
||||
story = await peers.get_story(pool, account_id, peer_id, story_id)
|
||||
if story is None:
|
||||
raise HTTPException(status_code=404, detail="story not found")
|
||||
if not story.downloaded or story.storage_key is None:
|
||||
raise HTTPException(status_code=409, detail=_NOT_STORED)
|
||||
return Subject(
|
||||
storage_key=story.storage_key,
|
||||
file_name=story_file_name(peer_id, story_id, story.media_kind),
|
||||
mime="video/mp4" if story.media_kind == "video" else "image/jpeg",
|
||||
file_size=None,
|
||||
title=await _chat_title(pool, account_id, peer_id),
|
||||
peer_id=peer_id,
|
||||
story_id=story_id,
|
||||
)
|
||||
|
||||
|
||||
async def _resolve(pool: asyncpg.Pool, body: ShareCreate) -> Subject:
|
||||
if body.kind == "media" and body.media_id is not None:
|
||||
return await _media_subject(pool, body.media_id)
|
||||
if body.kind == "media_version" and body.version_id is not None:
|
||||
return await _version_subject(pool, body.version_id)
|
||||
if body.kind == "story" and body.peer_id is not None and body.story_id is not None:
|
||||
return await _story_subject(pool, body.account_id, body.peer_id, body.story_id)
|
||||
raise HTTPException(status_code=422, detail="unsupported share subject")
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def list_shares(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
active_only: Annotated[bool, Query()] = False,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[FileShareView]:
|
||||
return await shares.list_shares(
|
||||
pool, account_id, Page(limit=limit, offset=offset), active_only=active_only
|
||||
)
|
||||
|
||||
|
||||
@router.post("", status_code=201)
|
||||
async def create_share(
|
||||
pool: FromDishka[asyncpg.Pool], body: ShareCreate
|
||||
) -> FileShareView:
|
||||
subject = await _resolve(pool, body)
|
||||
existing = await shares.find_active_share(
|
||||
pool, body.account_id, subject.storage_key
|
||||
)
|
||||
if existing is not None:
|
||||
return existing
|
||||
return await shares.create_share(
|
||||
pool,
|
||||
body.account_id,
|
||||
body.kind,
|
||||
subject.storage_key,
|
||||
subject.file_name,
|
||||
mime=subject.mime,
|
||||
file_size=subject.file_size,
|
||||
title=subject.title,
|
||||
chat_id=subject.chat_id,
|
||||
message_id=subject.message_id,
|
||||
peer_id=subject.peer_id,
|
||||
story_id=subject.story_id,
|
||||
expires_at=expiry_from_seconds(body.expires_in_seconds),
|
||||
max_downloads=body.max_downloads,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/lookup")
|
||||
async def lookup_share(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
kind: Annotated[str, Query()],
|
||||
media_id: Annotated[int | None, Query()] = None,
|
||||
version_id: Annotated[int | None, Query()] = None,
|
||||
peer_id: Annotated[int | None, Query()] = None,
|
||||
story_id: Annotated[int | None, Query()] = None,
|
||||
) -> FileShareView | None:
|
||||
body = ShareCreate(
|
||||
account_id=account_id,
|
||||
kind=kind,
|
||||
media_id=media_id,
|
||||
version_id=version_id,
|
||||
peer_id=peer_id,
|
||||
story_id=story_id,
|
||||
)
|
||||
subject = await _resolve(pool, body)
|
||||
return await shares.find_active_share(pool, account_id, subject.storage_key)
|
||||
|
||||
|
||||
@router.get("/{share_id}/hits")
|
||||
async def share_hits(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
share_id: int,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[FileShareHitView]:
|
||||
return await shares.list_hits(pool, share_id, Page(limit=limit, offset=offset))
|
||||
|
||||
|
||||
@router.patch("/{share_id}")
|
||||
async def update_share(
|
||||
pool: FromDishka[asyncpg.Pool], share_id: int, body: ShareUpdate
|
||||
) -> FileShareView:
|
||||
current = await shares.get_share(pool, share_id)
|
||||
if current is None:
|
||||
raise HTTPException(status_code=404, detail="share not found")
|
||||
expires_at = (
|
||||
current.expires_at
|
||||
if body.keep_expiry
|
||||
else expiry_from_seconds(body.expires_in_seconds)
|
||||
)
|
||||
updated = await shares.update_share(
|
||||
pool, share_id, expires_at=expires_at, max_downloads=body.max_downloads
|
||||
)
|
||||
if updated is None:
|
||||
raise HTTPException(status_code=404, detail="share not found")
|
||||
return updated
|
||||
|
||||
|
||||
@router.post("/{share_id}/revoke")
|
||||
async def revoke_share(pool: FromDishka[asyncpg.Pool], share_id: int) -> FileShareView:
|
||||
share = await shares.revoke_share(pool, share_id)
|
||||
if share is None:
|
||||
raise HTTPException(status_code=404, detail="share not found")
|
||||
return share
|
||||
|
||||
|
||||
@router.post("/{share_id}/reissue")
|
||||
async def reissue_share(pool: FromDishka[asyncpg.Pool], share_id: int) -> FileShareView:
|
||||
share = await shares.rotate_token(pool, share_id)
|
||||
if share is None:
|
||||
raise HTTPException(status_code=404, detail="share not found")
|
||||
return share
|
||||
|
||||
|
||||
@router.delete("/{share_id}", status_code=204)
|
||||
async def delete_share(pool: FromDishka[asyncpg.Pool], share_id: int) -> None:
|
||||
if not await shares.delete_share(pool, share_id):
|
||||
raise HTTPException(status_code=404, detail="share not found")
|
||||
@@ -0,0 +1,166 @@
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from fastapi.responses import FileResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils import trash
|
||||
from utils.cache import NO_STORE_HEADERS
|
||||
from utils.files import resolve_mime
|
||||
from utils.read import storage as storage_read
|
||||
from utils.read.models import (
|
||||
DEFAULT_LIMIT,
|
||||
Page,
|
||||
StorageChatView,
|
||||
StorageFileView,
|
||||
StorageScopeView,
|
||||
StorageSummaryView,
|
||||
TrashItemView,
|
||||
TrashOutcome,
|
||||
)
|
||||
from utils.storage import ContentAddressedStorage
|
||||
|
||||
router = APIRouter(prefix="/api/storage", tags=["storage"], route_class=DishkaRoute)
|
||||
|
||||
|
||||
class TrashRequest(BaseModel):
|
||||
account_id: int
|
||||
media_ids: list[int]
|
||||
|
||||
|
||||
class RestoreRequest(BaseModel):
|
||||
account_id: int
|
||||
keys: list[str] | None = None
|
||||
|
||||
|
||||
class ScopeRequest(BaseModel):
|
||||
account_id: int
|
||||
chat_id: int | None = None
|
||||
folder_id: int | None = None
|
||||
|
||||
|
||||
class PurgeRequest(BaseModel):
|
||||
account_id: int
|
||||
keys: list[str] | None = None
|
||||
|
||||
|
||||
@router.get("/summary")
|
||||
async def summary(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
chat_id: Annotated[int | None, Query()] = None,
|
||||
) -> StorageSummaryView:
|
||||
return await storage_read.get_summary(pool, account_id, chat_id)
|
||||
|
||||
|
||||
@router.get("/files")
|
||||
async def files(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
group: Annotated[str | None, Query()] = None,
|
||||
chat_id: Annotated[int | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[StorageFileView]:
|
||||
page = Page(limit=limit, offset=offset)
|
||||
return await storage_read.list_files(pool, account_id, group, page, chat_id)
|
||||
|
||||
|
||||
@router.get("/chats")
|
||||
async def chats(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
group: Annotated[str | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[StorageChatView]:
|
||||
page = Page(limit=limit, offset=offset)
|
||||
return await storage_read.list_chats(pool, account_id, group, page)
|
||||
|
||||
|
||||
async def _chat_ids(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int | None, folder_id: int | None
|
||||
) -> list[int]:
|
||||
chat_ids = await storage_read.scope_chat_ids(pool, account_id, chat_id, folder_id)
|
||||
if chat_ids is None:
|
||||
raise HTTPException(status_code=404, detail="scope not found")
|
||||
return chat_ids
|
||||
|
||||
|
||||
@router.get("/scope")
|
||||
async def scope(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
chat_id: Annotated[int | None, Query()] = None,
|
||||
folder_id: Annotated[int | None, Query()] = None,
|
||||
) -> StorageScopeView:
|
||||
chat_ids = await _chat_ids(pool, account_id, chat_id, folder_id)
|
||||
return await storage_read.get_scope(pool, account_id, chat_ids)
|
||||
|
||||
|
||||
@router.get("/trash")
|
||||
async def trash_items(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[TrashItemView]:
|
||||
page = Page(limit=limit, offset=offset)
|
||||
return await storage_read.list_trash(pool, account_id, page)
|
||||
|
||||
|
||||
@router.post("/trash/scope")
|
||||
async def trash_scope(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
body: ScopeRequest,
|
||||
) -> TrashOutcome:
|
||||
chat_ids = await _chat_ids(pool, body.account_id, body.chat_id, body.folder_id)
|
||||
blobs = await pool.fetch(storage_read.SCOPE_BLOBS, body.account_id, chat_ids)
|
||||
free = [blob for blob in blobs if blob["free"]]
|
||||
return await trash.trash_blobs(pool, storage, body.account_id, free)
|
||||
|
||||
|
||||
@router.post("/trash")
|
||||
async def move_to_trash(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
body: TrashRequest,
|
||||
) -> TrashOutcome:
|
||||
return await trash.move_to_trash(pool, storage, body.account_id, body.media_ids)
|
||||
|
||||
|
||||
@router.post("/trash/restore")
|
||||
async def restore(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
body: RestoreRequest,
|
||||
) -> TrashOutcome:
|
||||
return await trash.restore(pool, storage, body.account_id, body.keys)
|
||||
|
||||
|
||||
@router.post("/trash/purge")
|
||||
async def purge(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
body: PurgeRequest,
|
||||
) -> TrashOutcome:
|
||||
return await trash.purge(pool, storage, body.account_id, body.keys)
|
||||
|
||||
|
||||
@router.get("/trash/{storage_key}/file")
|
||||
async def serve_trashed(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
storage_key: str,
|
||||
) -> FileResponse:
|
||||
item = await storage_read.get_trash_item(pool, storage_key)
|
||||
if item is None or not storage.trashed(storage_key):
|
||||
raise HTTPException(status_code=404, detail="file is not in trash")
|
||||
return FileResponse(
|
||||
storage.trash_url(storage_key),
|
||||
media_type=resolve_mime(item["kind"], item["mime"], item["file_name"]),
|
||||
headers=NO_STORE_HEADERS,
|
||||
)
|
||||
@@ -5,6 +5,7 @@ from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from utils.files import content_disposition, story_file_name
|
||||
from utils.read import peers
|
||||
from utils.read.models import DEFAULT_LIMIT, Page, StoryView
|
||||
from utils.storage import ContentAddressedStorage
|
||||
@@ -36,13 +37,20 @@ async def serve_story_media(
|
||||
peer_id: int,
|
||||
story_id: int,
|
||||
account_id: AccountId,
|
||||
download: Annotated[bool, Query()] = False,
|
||||
) -> FileResponse:
|
||||
story = await peers.get_story(pool, account_id, peer_id, story_id)
|
||||
if story is None:
|
||||
raise HTTPException(status_code=404, detail="story not found")
|
||||
if not story.downloaded or story.storage_key is None:
|
||||
raise HTTPException(status_code=409, detail="story media not downloaded")
|
||||
headers = {}
|
||||
if download:
|
||||
headers["Content-Disposition"] = content_disposition(
|
||||
story_file_name(peer_id, story_id, story.media_kind), attachment=True
|
||||
)
|
||||
return FileResponse(
|
||||
storage.url(story.storage_key),
|
||||
media_type=_STORY_MIME.get(story.media_kind or "", "application/octet-stream"),
|
||||
headers=headers,
|
||||
)
|
||||
@@ -1,3 +0,0 @@
|
||||
from userbot.modules.client import DEVICE_MODEL, DEVICE_MODEL_LIMIT, PyroClient
|
||||
|
||||
__all__ = ["DEVICE_MODEL", "DEVICE_MODEL_LIMIT", "PyroClient"]
|
||||
@@ -1,8 +1,8 @@
|
||||
from pyrogram.types import Message
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture import repository
|
||||
from userbot.modules.capture.repository import CHANNEL_ID_THRESHOLD
|
||||
from userbot.modules.client import PyroClient
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
|
||||
@@ -30,4 +30,4 @@ async def on_deleted_messages(client: PyroClient, messages: list[Message]) -> No
|
||||
)
|
||||
|
||||
|
||||
handlers = on_deleted_messages.handlers
|
||||
handlers = on_deleted_messages.handlers # ty: ignore[unresolved-attribute]
|
||||
@@ -1,17 +1,24 @@
|
||||
from pyrogram.types import Message
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture import repository
|
||||
from userbot.modules.capture.chat_meta import meta_from_chat
|
||||
from userbot.modules.capture.message import sender_id
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
|
||||
@PyroClient.on_edited_message()
|
||||
async def on_edited_message(client: PyroClient, message: Message) -> None:
|
||||
ctx = client.capture
|
||||
if ctx is None or message.empty or message.chat is None or message.date is None:
|
||||
if (
|
||||
ctx is None
|
||||
or message.empty
|
||||
or message.scheduled
|
||||
or message.chat is None
|
||||
or message.date is None
|
||||
):
|
||||
return
|
||||
chat = message.chat
|
||||
chat_id = chat.id or 0
|
||||
@@ -27,7 +34,7 @@ async def on_edited_message(client: PyroClient, message: Message) -> None:
|
||||
message.date,
|
||||
sender_id(message),
|
||||
message.text or message.caption,
|
||||
str(message),
|
||||
raw_json(message),
|
||||
message.edit_date,
|
||||
media_unique_id(message),
|
||||
has_media=message.media is not None,
|
||||
@@ -42,4 +49,4 @@ async def on_edited_message(client: PyroClient, message: Message) -> None:
|
||||
)
|
||||
|
||||
|
||||
handlers = on_edited_message.handlers
|
||||
handlers = on_edited_message.handlers # ty: ignore[unresolved-attribute]
|
||||
@@ -1,8 +1,8 @@
|
||||
from pyrogram.types import Message
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture import capture_message
|
||||
from userbot.modules.capture.chat_meta import meta_from_chat
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.stt import is_transcribable
|
||||
from userbot.modules.stt.gate import safe_transcribe
|
||||
from utils.events import notify_bg_event
|
||||
@@ -11,7 +11,13 @@ from utils.events import notify_bg_event
|
||||
@PyroClient.on_message()
|
||||
async def on_message(client: PyroClient, message: Message) -> None:
|
||||
ctx = client.capture
|
||||
if ctx is None or message.empty or message.chat is None or message.date is None:
|
||||
if (
|
||||
ctx is None
|
||||
or message.empty
|
||||
or message.scheduled
|
||||
or message.chat is None
|
||||
or message.date is None
|
||||
):
|
||||
return
|
||||
meta = meta_from_chat(message.chat, ctx.contacts.ids)
|
||||
await ctx.watches.on_text(meta.chat_id, message.id, message.text or message.caption)
|
||||
@@ -30,4 +36,4 @@ async def on_message(client: PyroClient, message: Message) -> None:
|
||||
await safe_transcribe(client, ctx, meta.chat_id, message.id)
|
||||
|
||||
|
||||
handlers = on_message.handlers
|
||||
handlers = on_message.handlers # ty: ignore[unresolved-attribute]
|
||||
@@ -2,8 +2,9 @@ from datetime import UTC, datetime
|
||||
|
||||
from pyrogram.types import User
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.presence import repository
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
|
||||
@@ -20,10 +21,10 @@ async def on_user_status(client: PyroClient, user: User) -> None:
|
||||
user.status.name.lower(),
|
||||
user.last_online_date,
|
||||
user.next_offline_date,
|
||||
str(user.raw),
|
||||
raw_json(user.raw),
|
||||
)
|
||||
await ctx.watches.on_status(user.id, is_online=user.status.name.lower() == "online")
|
||||
await notify_bg_event(ctx.pool, "presence", ctx.account_id, chat_id=user.id)
|
||||
|
||||
|
||||
handlers = on_user_status.handlers
|
||||
handlers = on_user_status.handlers # ty: ignore[unresolved-attribute]
|
||||
@@ -1,6 +1,6 @@
|
||||
from pyrogram import raw
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.client import PyroClient
|
||||
|
||||
HANDLES = (raw.types.UpdateContactsReset, raw.types.UpdateUser)
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from pyrogram import raw
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.client import PyroClient
|
||||
|
||||
HANDLES = (
|
||||
raw.types.UpdateDialogFilter,
|
||||
|
||||
@@ -2,7 +2,6 @@ from collections.abc import Awaitable, Callable
|
||||
|
||||
from pyrogram import raw
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.handlers.raw import (
|
||||
contacts,
|
||||
dialog_filters,
|
||||
@@ -12,6 +11,7 @@ from userbot.handlers.raw import (
|
||||
read_receipts,
|
||||
transcribed,
|
||||
)
|
||||
from userbot.modules.client import PyroClient
|
||||
|
||||
RawHandler = Callable[..., Awaitable[None]]
|
||||
|
||||
@@ -24,8 +24,9 @@ _MODULES = (
|
||||
read_receipts,
|
||||
transcribed,
|
||||
)
|
||||
_REGISTRY: dict[type, RawHandler] = {
|
||||
update_type: module.handle for module in _MODULES for update_type in module.HANDLES
|
||||
_REGISTRY: dict[type, list[RawHandler]] = {
|
||||
update_type: [module.handle for module in _MODULES if update_type in module.HANDLES]
|
||||
for update_type in {kind for module in _MODULES for kind in module.HANDLES}
|
||||
}
|
||||
|
||||
|
||||
@@ -33,9 +34,8 @@ _REGISTRY: dict[type, RawHandler] = {
|
||||
async def dispatch(
|
||||
client: PyroClient, update: raw.base.Update, users: dict, chats: dict
|
||||
) -> None:
|
||||
handler = _REGISTRY.get(type(update))
|
||||
if handler is not None:
|
||||
for handler in _REGISTRY.get(type(update), ()):
|
||||
await handler(client, update, users, chats)
|
||||
|
||||
|
||||
handlers = dispatch.handlers
|
||||
handlers = dispatch.handlers # ty: ignore[unresolved-attribute]
|
||||
@@ -2,11 +2,12 @@ from dataclasses import replace
|
||||
|
||||
from pyrogram import raw
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.avatars import capture_avatar
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.profiles import active_username, snapshot_from_user
|
||||
from userbot.modules.profiles.parse import ProfileFields
|
||||
from userbot.modules.profiles.repository import get_peer, write_profile
|
||||
from userbot.modules.raw_json import raw_json
|
||||
|
||||
HANDLES = (raw.types.UpdateUserName, raw.types.UpdateUser, raw.types.UpdateUserPhone)
|
||||
|
||||
@@ -20,11 +21,13 @@ async def _handle_user(
|
||||
raw_user = users.get(update.user_id)
|
||||
if raw_user is None:
|
||||
return
|
||||
fields, photo_file_id, photo_unique_id = snapshot_from_user(client, raw_user)
|
||||
fields, photo_file_id, photo_unique_id = await snapshot_from_user(client, raw_user)
|
||||
current = await get_peer(ctx.pool, ctx.account_id, update.user_id)
|
||||
if current == fields:
|
||||
return
|
||||
await write_profile(ctx.pool, ctx.account_id, update.user_id, fields, str(raw_user))
|
||||
await write_profile(
|
||||
ctx.pool, ctx.account_id, update.user_id, fields, raw_json(raw_user)
|
||||
)
|
||||
changed_photo = current is None or current.photo_unique_id != photo_unique_id
|
||||
if photo_file_id and photo_unique_id and changed_photo:
|
||||
await capture_avatar(
|
||||
@@ -34,7 +37,7 @@ async def _handle_user(
|
||||
"peer",
|
||||
photo_file_id,
|
||||
photo_unique_id,
|
||||
str(raw_user),
|
||||
raw_json(raw_user),
|
||||
)
|
||||
|
||||
|
||||
@@ -65,9 +68,9 @@ async def handle(
|
||||
"last_name": update.last_name or None,
|
||||
"username": active_username(update.usernames),
|
||||
},
|
||||
str(update),
|
||||
raw_json(update),
|
||||
)
|
||||
elif isinstance(update, raw.types.UpdateUserPhone):
|
||||
await _merge_partial(
|
||||
client, update.user_id, {"phone": update.phone or None}, str(update)
|
||||
client, update.user_id, {"phone": update.phone or None}, raw_json(update)
|
||||
)
|
||||
@@ -1,8 +1,8 @@
|
||||
from pyrogram import raw, utils
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture import repository
|
||||
from userbot.modules.capture.chat_meta import meta_from_peer
|
||||
from userbot.modules.client import PyroClient
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
HANDLES = (raw.types.UpdateMessageReactions,)
|
||||
|
||||
@@ -2,8 +2,9 @@ from datetime import UTC, datetime
|
||||
|
||||
from pyrogram import raw, utils
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture.chat_meta import meta_from_chat_id
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from userbot.modules.read_receipts import repository as receipts
|
||||
from userbot.modules.stt import repository
|
||||
from userbot.modules.stt.gate import safe_transcribe
|
||||
@@ -41,7 +42,7 @@ async def handle(
|
||||
cand_chat_id,
|
||||
datetime.now(UTC),
|
||||
message_id,
|
||||
str(update),
|
||||
raw_json(update),
|
||||
)
|
||||
continue
|
||||
if not untranscribed:
|
||||
|
||||
@@ -2,7 +2,8 @@ from datetime import UTC, datetime
|
||||
|
||||
from pyrogram import raw, utils
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from userbot.modules.read_receipts import repository
|
||||
from utils.events import notify_bg_event
|
||||
|
||||
@@ -23,7 +24,7 @@ async def handle(
|
||||
chat_id,
|
||||
datetime.now(UTC),
|
||||
update.max_id,
|
||||
str(update),
|
||||
raw_json(update),
|
||||
)
|
||||
await notify_bg_event(
|
||||
ctx.pool, "receipt", ctx.account_id, chat_id=chat_id, message_id=update.max_id
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from pyrogram import raw
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture.chat_meta import meta_from_peer
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.stt import repository
|
||||
|
||||
HANDLES = (raw.types.UpdateTranscribedAudio,)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from pyrogram.types import Story
|
||||
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.stories.service import save_story
|
||||
|
||||
|
||||
@@ -11,4 +11,4 @@ async def on_story(client: PyroClient, story: Story) -> None:
|
||||
await save_story(client, client.capture, story)
|
||||
|
||||
|
||||
handlers = on_story.handlers
|
||||
handlers = on_story.handlers # ty: ignore[unresolved-attribute]
|
||||
@@ -1,9 +1,9 @@
|
||||
from io import BytesIO
|
||||
|
||||
from pyrogram import Client
|
||||
|
||||
from userbot.modules.avatars import repository
|
||||
from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.download import download_bytes
|
||||
from utils.logging import logger
|
||||
|
||||
|
||||
async def capture_avatar( # noqa: PLR0913
|
||||
@@ -20,9 +20,12 @@ async def capture_avatar( # noqa: PLR0913
|
||||
storage_key: str | None = None
|
||||
file_size: int | None = None
|
||||
downloaded = False
|
||||
buffer = await client.download_media(file_id, in_memory=True)
|
||||
if isinstance(buffer, BytesIO):
|
||||
data = buffer.getvalue()
|
||||
try:
|
||||
data = await download_bytes(client, file_id)
|
||||
except TimeoutError:
|
||||
logger.warning(f"[yellow]Avatar download timed out for {owner_id}.[/]")
|
||||
data = None
|
||||
if data is not None:
|
||||
storage_key = ctx.storage.put(data)
|
||||
file_size = len(data)
|
||||
downloaded = True
|
||||
|
||||
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
|
||||
from userbot.modules.capture.chat_meta import chat_kind
|
||||
from userbot.modules.profiles.parse import ProfileFields, snapshot_from_high_level
|
||||
from userbot.modules.profiles.repository import get_peer, write_profile
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from utils.policy.models import ChatKind
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -40,9 +41,7 @@ class ChatMetaCache:
|
||||
chat_id: int,
|
||||
meta: tuple[str | None, str | None],
|
||||
) -> bool:
|
||||
from userbot.modules.groups.repository import ( # noqa: PLC0415
|
||||
get_latest_chat_meta,
|
||||
)
|
||||
from userbot.modules.groups.repository import get_latest_chat_meta # noqa: PLC0415
|
||||
|
||||
if chat_id not in self._cache:
|
||||
self._cache[chat_id] = await get_latest_chat_meta(pool, account_id, chat_id)
|
||||
@@ -59,7 +58,7 @@ async def _capture_peer(message: Message, ctx: CaptureContext) -> None:
|
||||
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
|
||||
if not await ctx.peer_identity.changed(ctx.pool, ctx.account_id, user.id, fields):
|
||||
return
|
||||
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user))
|
||||
await write_profile(ctx.pool, ctx.account_id, user.id, fields, raw_json(user))
|
||||
if photo_file_id and photo_unique_id:
|
||||
from userbot.modules.avatars import note_avatar # noqa: PLC0415
|
||||
|
||||
@@ -95,7 +94,7 @@ async def _capture_chat(message: Message, ctx: CaptureContext) -> None:
|
||||
photo_unique_id,
|
||||
None,
|
||||
message.date,
|
||||
str(message),
|
||||
raw_json(message),
|
||||
)
|
||||
if photo_file_id and photo_unique_id:
|
||||
from userbot.modules.avatars import note_avatar # noqa: PLC0415
|
||||
|
||||
@@ -6,6 +6,7 @@ from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.links import extract_links
|
||||
from userbot.modules.links import repository as links_repository
|
||||
from userbot.modules.media import capture_media, self_destruct_ttl
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from utils.policy.models import CaptureToggles
|
||||
|
||||
|
||||
@@ -47,7 +48,7 @@ async def capture_message(
|
||||
message.date,
|
||||
sender_id(message),
|
||||
message.text or message.caption,
|
||||
str(message),
|
||||
raw_json(message),
|
||||
has_media=message.media is not None,
|
||||
is_self_destruct=self_destruct_ttl(message) is not None,
|
||||
)
|
||||
@@ -68,5 +69,5 @@ async def capture_message(
|
||||
links = extract_links(message)
|
||||
if links:
|
||||
await links_repository.insert_links(
|
||||
ctx.pool, ctx.account_id, chat_id, message.id, links, str(message)
|
||||
ctx.pool, ctx.account_id, chat_id, message.id, links, raw_json(message)
|
||||
)
|
||||
@@ -30,6 +30,16 @@ ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
|
||||
edited_at = now()
|
||||
"""
|
||||
|
||||
_FIND_DELETED = """
|
||||
SELECT chat_id, message_id, date FROM messages
|
||||
WHERE account_id = $1 AND message_id = ANY($2::bigint[]) AND deleted_at IS NULL
|
||||
"""
|
||||
|
||||
_MARK_DELETED = """
|
||||
UPDATE messages SET deleted_at = now()
|
||||
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND date = $4
|
||||
"""
|
||||
|
||||
_INSERT_VERSION = """
|
||||
INSERT INTO message_versions
|
||||
(account_id, chat_id, message_id, observed_at, edit_date, text, raw)
|
||||
@@ -78,8 +88,8 @@ ON CONFLICT (account_id, chat_id, message_id, storage_key) DO NOTHING
|
||||
_INSERT_MEDIA = """
|
||||
INSERT INTO media
|
||||
(account_id, chat_id, message_id, kind, storage_key, file_size, mime,
|
||||
ttl_seconds, downloaded, unique_id)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
ttl_seconds, downloaded, unique_id, evicted_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, CASE WHEN $11 THEN now() END)
|
||||
ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET
|
||||
kind = EXCLUDED.kind,
|
||||
storage_key = EXCLUDED.storage_key,
|
||||
@@ -87,7 +97,27 @@ ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET
|
||||
mime = EXCLUDED.mime,
|
||||
ttl_seconds = EXCLUDED.ttl_seconds,
|
||||
downloaded = EXCLUDED.downloaded,
|
||||
unique_id = EXCLUDED.unique_id
|
||||
unique_id = EXCLUDED.unique_id,
|
||||
evicted_at = CASE
|
||||
WHEN EXCLUDED.downloaded THEN NULL
|
||||
WHEN media.unique_id IS DISTINCT FROM EXCLUDED.unique_id
|
||||
THEN EXCLUDED.evicted_at
|
||||
ELSE coalesce(media.evicted_at, EXCLUDED.evicted_at) END,
|
||||
evicted_key = CASE WHEN EXCLUDED.downloaded
|
||||
OR media.unique_id IS DISTINCT FROM EXCLUDED.unique_id
|
||||
THEN NULL ELSE media.evicted_key END
|
||||
"""
|
||||
|
||||
_IS_EVICTED = """
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM media WHERE unique_id = $1 AND evicted_at IS NOT NULL
|
||||
)
|
||||
"""
|
||||
|
||||
_UNEVICT = """
|
||||
UPDATE media SET storage_key = $2, file_size = $3, downloaded = true,
|
||||
evicted_at = NULL, evicted_key = NULL
|
||||
WHERE evicted_key = $2 OR (unique_id = $1 AND evicted_at IS NOT NULL)
|
||||
"""
|
||||
|
||||
|
||||
@@ -122,36 +152,40 @@ async def max_message_id(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
) -> int | None:
|
||||
return await pool.fetchval(
|
||||
"SELECT max(message_id) FROM messages WHERE account_id = $1 AND chat_id = $2",
|
||||
"SELECT last_message_id FROM chat_stats WHERE account_id = $1 AND chat_id = $2",
|
||||
account_id,
|
||||
chat_id,
|
||||
)
|
||||
|
||||
|
||||
async def _mark_deleted(
|
||||
pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
|
||||
) -> None:
|
||||
for row in rows:
|
||||
await pool.execute(
|
||||
_MARK_DELETED, account_id, row["chat_id"], row["message_id"], row["date"]
|
||||
)
|
||||
|
||||
|
||||
async def mark_deleted_box(
|
||||
pool: asyncpg.Pool, account_id: int, message_ids: list[int]
|
||||
) -> None:
|
||||
await pool.execute(
|
||||
"UPDATE messages SET deleted_at = now() "
|
||||
"WHERE account_id = $1 AND message_id = ANY($2::bigint[]) "
|
||||
"AND chat_id > $3 AND deleted_at IS NULL",
|
||||
rows = await pool.fetch(
|
||||
f"{_FIND_DELETED} AND chat_id > $3",
|
||||
account_id,
|
||||
message_ids,
|
||||
CHANNEL_ID_THRESHOLD,
|
||||
)
|
||||
await _mark_deleted(pool, account_id, rows)
|
||||
|
||||
|
||||
async def mark_deleted_channel(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_ids: list[int]
|
||||
) -> None:
|
||||
await pool.execute(
|
||||
"UPDATE messages SET deleted_at = now() "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = ANY($3::bigint[]) "
|
||||
"AND deleted_at IS NULL",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_ids,
|
||||
rows = await pool.fetch(
|
||||
f"{_FIND_DELETED} AND chat_id = $3", account_id, message_ids, chat_id
|
||||
)
|
||||
await _mark_deleted(pool, account_id, rows)
|
||||
|
||||
|
||||
async def add_version( # noqa: PLR0913
|
||||
@@ -215,6 +249,7 @@ async def insert_media( # noqa: PLR0913
|
||||
unique_id: str | None,
|
||||
*,
|
||||
downloaded: bool,
|
||||
evicted: bool = False,
|
||||
) -> None:
|
||||
async with pool.acquire() as conn, conn.transaction():
|
||||
await conn.execute(
|
||||
@@ -229,6 +264,7 @@ async def insert_media( # noqa: PLR0913
|
||||
ttl_seconds,
|
||||
downloaded,
|
||||
unique_id,
|
||||
evicted,
|
||||
)
|
||||
if storage_key is not None:
|
||||
await conn.execute(
|
||||
@@ -245,6 +281,22 @@ async def insert_media( # noqa: PLR0913
|
||||
)
|
||||
|
||||
|
||||
async def is_evicted(pool: asyncpg.Pool, unique_id: str | None) -> bool:
|
||||
return unique_id is not None and await pool.fetchval(_IS_EVICTED, unique_id)
|
||||
|
||||
|
||||
async def unevict(
|
||||
pool: asyncpg.Pool, unique_id: str | None, storage_key: str, file_size: int | None
|
||||
) -> bool:
|
||||
async with pool.acquire() as conn, conn.transaction():
|
||||
await conn.execute(_UNEVICT, unique_id, storage_key, file_size)
|
||||
dropped = await conn.fetchval(
|
||||
"DELETE FROM media_trash WHERE storage_key = $1 RETURNING storage_key",
|
||||
storage_key,
|
||||
)
|
||||
return dropped is not None
|
||||
|
||||
|
||||
async def insert_callbacks(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
|
||||
@@ -2,6 +2,8 @@ from typing import TYPE_CHECKING
|
||||
|
||||
from pyrogram import Client, enums
|
||||
|
||||
from utils.telegram import API_HASH, API_ID
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from userbot.modules.capture import CaptureContext
|
||||
|
||||
@@ -11,19 +13,23 @@ DEVICE_MODEL_LIMIT = 32
|
||||
|
||||
|
||||
class PyroClient(Client):
|
||||
def __init__(
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
workdir: str = "sessions",
|
||||
device_model: str | None = None,
|
||||
load_handlers: bool = True,
|
||||
bot_token: str | None = None,
|
||||
skip_updates: bool = True,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
name,
|
||||
workdir=workdir,
|
||||
api_id=2040,
|
||||
api_hash="b18441a1ff607e10a989891a5462e627",
|
||||
api_id=API_ID,
|
||||
api_hash=API_HASH,
|
||||
bot_token=bot_token,
|
||||
skip_updates=skip_updates,
|
||||
device_model=device_model or DEVICE_MODEL,
|
||||
system_version="Windows 11 x64",
|
||||
app_version="7.0.8 x64",
|
||||
@@ -31,6 +37,7 @@ class PyroClient(Client):
|
||||
lang_code="en",
|
||||
system_lang_code="en-US",
|
||||
client_platform=enums.ClientPlatform.DESKTOP,
|
||||
max_concurrent_transmissions=4,
|
||||
)
|
||||
self.capture: CaptureContext | None = None
|
||||
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
from io import BytesIO
|
||||
|
||||
from pyrogram import Client
|
||||
from pyrogram.types import (
|
||||
Animation,
|
||||
Audio,
|
||||
Document,
|
||||
Message,
|
||||
Photo,
|
||||
Sticker,
|
||||
Story,
|
||||
Video,
|
||||
VideoNote,
|
||||
Voice,
|
||||
)
|
||||
|
||||
from utils.logging import logger
|
||||
|
||||
type Downloadable = (
|
||||
str
|
||||
| Message
|
||||
| Story
|
||||
| Audio
|
||||
| Document
|
||||
| Photo
|
||||
| Sticker
|
||||
| Animation
|
||||
| Video
|
||||
| Voice
|
||||
| VideoNote
|
||||
)
|
||||
|
||||
STOP_TIMEOUT = 5.0
|
||||
|
||||
_reset_lock = asyncio.Lock()
|
||||
|
||||
|
||||
async def _reset_media_sessions(client: Client) -> None:
|
||||
async with _reset_lock:
|
||||
sessions = list(client.media_sessions.values())
|
||||
client.media_sessions.clear()
|
||||
for session in sessions:
|
||||
with contextlib.suppress(Exception):
|
||||
await asyncio.wait_for(session.stop(), STOP_TIMEOUT)
|
||||
logger.warning(f"[yellow]Dropped {len(sessions)} stale media session(s).[/]")
|
||||
|
||||
|
||||
async def download_bytes(client: Client, target: Downloadable) -> bytes | None:
|
||||
try:
|
||||
buffer = await client.download_media(target, in_memory=True)
|
||||
except (TimeoutError, OSError):
|
||||
await _reset_media_sessions(client)
|
||||
buffer = await client.download_media(target, in_memory=True)
|
||||
return buffer.getvalue() if isinstance(buffer, BytesIO) else None
|
||||
@@ -5,6 +5,7 @@ from userbot.modules.avatars import capture_avatar
|
||||
from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.capture.message import sender_id
|
||||
from userbot.modules.groups import repository
|
||||
from userbot.modules.raw_json import raw_json
|
||||
|
||||
|
||||
async def capture_service(
|
||||
@@ -28,7 +29,7 @@ async def capture_service(
|
||||
None,
|
||||
actor,
|
||||
ts,
|
||||
str(message),
|
||||
raw_json(message),
|
||||
)
|
||||
elif service is enums.MessageServiceType.NEW_CHAT_PHOTO:
|
||||
photo = message.new_chat_photo
|
||||
@@ -43,11 +44,17 @@ async def capture_service(
|
||||
unique_id,
|
||||
actor,
|
||||
ts,
|
||||
str(message),
|
||||
raw_json(message),
|
||||
)
|
||||
if photo is not None and unique_id is not None:
|
||||
await capture_avatar(
|
||||
client, ctx, chat_id, "chat", photo.file_id, unique_id, str(message)
|
||||
client,
|
||||
ctx,
|
||||
chat_id,
|
||||
"chat",
|
||||
photo.file_id,
|
||||
unique_id,
|
||||
raw_json(message),
|
||||
)
|
||||
elif service is enums.MessageServiceType.DELETE_CHAT_PHOTO:
|
||||
await repository.insert_chat_history(
|
||||
@@ -60,7 +67,7 @@ async def capture_service(
|
||||
None,
|
||||
actor,
|
||||
ts,
|
||||
str(message),
|
||||
raw_json(message),
|
||||
)
|
||||
elif service is enums.MessageServiceType.NEW_CHAT_MEMBERS:
|
||||
for member in message.new_chat_members or []:
|
||||
@@ -73,7 +80,7 @@ async def capture_service(
|
||||
"join",
|
||||
actor,
|
||||
ts,
|
||||
str(message),
|
||||
raw_json(message),
|
||||
)
|
||||
elif service is enums.MessageServiceType.LEFT_CHAT_MEMBER:
|
||||
member = message.left_chat_member
|
||||
@@ -87,5 +94,5 @@ async def capture_service(
|
||||
"leave",
|
||||
actor,
|
||||
ts,
|
||||
str(message),
|
||||
raw_json(message),
|
||||
)
|
||||
@@ -69,6 +69,9 @@ class JobConsumer:
|
||||
requeued = await repository.requeue_running(self.pool, self.account_id)
|
||||
if requeued:
|
||||
logger.info(f"[yellow]Requeued {requeued} stale running job(s).[/]")
|
||||
pruned = await repository.prune_finished(self.pool, self.account_id)
|
||||
if pruned:
|
||||
logger.info(f"[yellow]Pruned {pruned} finished job(s).[/]")
|
||||
conn = await asyncpg.connect(dsn=env.db.connection_url)
|
||||
await conn.add_listener(JOBS_CHANGED_CHANNEL, lambda *_: self._wake.set())
|
||||
logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from pyrogram import Client
|
||||
from pyrogram.errors import PeerIdInvalid
|
||||
from pyrogram.errors import FloodPremiumWait, FloodWait, PeerIdInvalid
|
||||
from pyrogram.types import Message
|
||||
|
||||
from userbot.modules.capture import capture_message
|
||||
@@ -11,6 +11,7 @@ from userbot.modules.jobs.registry import register
|
||||
from userbot.modules.stt import repository as stt_repo
|
||||
from userbot.modules.stt import should_transcribe_on_backfill
|
||||
from userbot.modules.stt.gate import safe_transcribe
|
||||
from utils.logging import logger
|
||||
from utils.policy.models import CaptureToggles
|
||||
|
||||
SAVE_EVERY = 100
|
||||
@@ -61,21 +62,30 @@ async def backfill(ctx: JobContext) -> None:
|
||||
min_id = await resolve_min_id(ctx, chat_id)
|
||||
await ctx.save_cursor({"max_id": max_id, "min_id": min_id})
|
||||
processed = ctx.job.progress.get("processed", 0)
|
||||
errors = ctx.job.progress.get("errors", 0)
|
||||
self_id = client.me.id if client.me else None
|
||||
try:
|
||||
async for message in client.get_chat_history(
|
||||
chat_id, max_id=max_id, min_id=min_id
|
||||
):
|
||||
await capture_message(client, message, capture, toggles)
|
||||
await maybe_transcribe(client, capture, chat_id, message, self_id)
|
||||
try:
|
||||
await capture_message(client, message, capture, toggles)
|
||||
await maybe_transcribe(client, capture, chat_id, message, self_id)
|
||||
except (FloodWait, FloodPremiumWait, PeerIdInvalid):
|
||||
raise
|
||||
except Exception:
|
||||
errors += 1
|
||||
logger.exception(f"Backfill capture failed for {chat_id}/{message.id}")
|
||||
processed += 1
|
||||
if processed % SAVE_EVERY == 0:
|
||||
next_max = message.id - 1
|
||||
await ctx.save_cursor({"max_id": next_max, "min_id": min_id})
|
||||
await ctx.report_progress({"processed": processed, "max_id": next_max})
|
||||
await ctx.report_progress(
|
||||
{"processed": processed, "errors": errors, "max_id": next_max}
|
||||
)
|
||||
if await ctx.is_canceled():
|
||||
return
|
||||
except PeerIdInvalid:
|
||||
await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"})
|
||||
return
|
||||
await ctx.report_progress({"processed": processed, "done": True})
|
||||
await ctx.report_progress({"processed": processed, "errors": errors, "done": True})
|
||||
@@ -8,6 +8,7 @@ from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.jobs.context import JobContext
|
||||
from userbot.modules.jobs.registry import register
|
||||
from userbot.modules.stories.service import save_story
|
||||
from utils.logging import logger
|
||||
|
||||
SAVE_EVERY = 10
|
||||
|
||||
@@ -38,6 +39,9 @@ async def _drain(
|
||||
raise
|
||||
except RPCError:
|
||||
continue
|
||||
except Exception:
|
||||
logger.exception(f"Failed to save story {story.id} from {name}")
|
||||
continue
|
||||
saved += 1
|
||||
if saved % SAVE_EVERY == 0:
|
||||
await ctx.report_progress({"saved": saved, "source": name})
|
||||
|
||||
@@ -11,6 +11,7 @@ from userbot.modules.jobs.context import JobContext
|
||||
from userbot.modules.jobs.registry import register
|
||||
from userbot.modules.profiles.parse import snapshot_from_high_level
|
||||
from userbot.modules.profiles.repository import write_profile
|
||||
from userbot.modules.raw_json import raw_json
|
||||
|
||||
MEMBER_CAP = 200
|
||||
|
||||
@@ -26,7 +27,7 @@ LIMIT 100
|
||||
|
||||
async def _save_user(ctx: CaptureContext, user: User) -> None:
|
||||
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
|
||||
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user))
|
||||
await write_profile(ctx.pool, ctx.account_id, user.id, fields, raw_json(user))
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
ctx.pool, ctx.account_id, user.id, "peer", photo_unique_id, photo_file_id
|
||||
@@ -35,6 +36,8 @@ async def _save_user(ctx: CaptureContext, user: User) -> None:
|
||||
|
||||
async def _enrich_chat_meta(client: Client, ctx: CaptureContext, chat_id: int) -> None:
|
||||
chat = await client.get_chat(chat_id)
|
||||
if chat is None:
|
||||
return
|
||||
photo = chat.photo
|
||||
photo_unique_id = photo.big_photo_unique_id if photo else None
|
||||
photo_file_id = photo.big_file_id if photo else None
|
||||
@@ -48,7 +51,7 @@ async def _enrich_chat_meta(client: Client, ctx: CaptureContext, chat_id: int) -
|
||||
photo_unique_id,
|
||||
None,
|
||||
datetime.now(UTC),
|
||||
str(chat),
|
||||
raw_json(chat),
|
||||
)
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
@@ -65,16 +68,28 @@ async def _enrich_members(client: Client, ctx: CaptureContext, chat_id: int) ->
|
||||
return
|
||||
|
||||
|
||||
async def _fetch_users(client: Client, ids: list[int]) -> list[User]:
|
||||
try:
|
||||
users = await client.get_users(ids)
|
||||
except BadRequest:
|
||||
users = []
|
||||
for sender_id in ids:
|
||||
try:
|
||||
users.append(await client.get_users(sender_id))
|
||||
except BadRequest:
|
||||
continue
|
||||
if isinstance(users, User):
|
||||
return [users]
|
||||
return [user for user in users if isinstance(user, User)]
|
||||
|
||||
|
||||
async def _enrich_senders(client: Client, ctx: CaptureContext, chat_id: int) -> None:
|
||||
rows = await ctx.pool.fetch(_MISSING_SENDERS, ctx.account_id, chat_id)
|
||||
ids = [row["sender_id"] for row in rows]
|
||||
for sender_id in ids:
|
||||
try:
|
||||
user = await client.get_users(sender_id)
|
||||
except BadRequest:
|
||||
continue
|
||||
if isinstance(user, User):
|
||||
await _save_user(ctx, user)
|
||||
if not ids:
|
||||
return
|
||||
for user in await _fetch_users(client, ids):
|
||||
await _save_user(ctx, user)
|
||||
|
||||
|
||||
@register("enrich_chat")
|
||||
|
||||
@@ -1,9 +1,48 @@
|
||||
from io import BytesIO
|
||||
from pyrogram import Client
|
||||
from pyrogram.errors import (
|
||||
ChannelInvalid,
|
||||
ChannelPrivate,
|
||||
FileIdInvalid,
|
||||
FileReferenceExpired,
|
||||
FileReferenceInvalid,
|
||||
PeerIdInvalid,
|
||||
)
|
||||
from pyrogram.types import ChatPhoto
|
||||
|
||||
from userbot.modules.avatars.repository import get_avatar_file, mark_avatar_downloaded
|
||||
from userbot.modules.download import download_bytes
|
||||
from userbot.modules.jobs.context import JobContext
|
||||
from userbot.modules.jobs.registry import register
|
||||
|
||||
STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)
|
||||
UNKNOWN_OWNER = (PeerIdInvalid, ChannelInvalid, ChannelPrivate)
|
||||
|
||||
|
||||
async def _fresh_file_id(client: Client, owner_id: int, unique_id: str) -> str | None:
|
||||
try:
|
||||
chat = await client.get_chat(owner_id)
|
||||
except UNKNOWN_OWNER:
|
||||
return None
|
||||
photo = chat.photo if chat else None
|
||||
if photo and photo.big_photo_unique_id == unique_id:
|
||||
return photo.big_file_id
|
||||
async for item in client.get_chat_photos(owner_id):
|
||||
if isinstance(item, ChatPhoto) and item.big_photo_unique_id == unique_id:
|
||||
return item.big_file_id
|
||||
return None
|
||||
|
||||
|
||||
async def _download(
|
||||
client: Client, owner_id: int, unique_id: str, file_id: str
|
||||
) -> bytes | None:
|
||||
try:
|
||||
return await download_bytes(client, file_id)
|
||||
except STALE_FILE_ID:
|
||||
fresh = await _fresh_file_id(client, owner_id, unique_id)
|
||||
if fresh is None:
|
||||
return None
|
||||
return await download_bytes(client, fresh)
|
||||
|
||||
|
||||
@register("fetch_avatar")
|
||||
async def fetch_avatar(ctx: JobContext) -> None:
|
||||
@@ -24,10 +63,9 @@ async def fetch_avatar(ctx: JobContext) -> None:
|
||||
file_id, downloaded = found
|
||||
if downloaded or file_id is None:
|
||||
return
|
||||
buffer = await client.download_media(file_id, in_memory=True)
|
||||
if not isinstance(buffer, BytesIO):
|
||||
data = await _download(client, owner_id, unique_id, file_id)
|
||||
if data is None:
|
||||
return
|
||||
data = buffer.getvalue()
|
||||
storage_key = capture.storage.put(data)
|
||||
await mark_avatar_downloaded(
|
||||
ctx.pool,
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
from io import BytesIO
|
||||
|
||||
from pyrogram.types import Sticker
|
||||
|
||||
from userbot.modules.custom_emoji.repository import is_downloaded, upsert_downloaded
|
||||
from userbot.modules.download import download_bytes
|
||||
from userbot.modules.jobs.context import JobContext
|
||||
from userbot.modules.jobs.registry import register
|
||||
|
||||
@@ -30,10 +29,9 @@ async def fetch_custom_emoji(ctx: JobContext) -> None:
|
||||
if not stickers:
|
||||
return
|
||||
sticker = stickers[0]
|
||||
buffer = await client.download_media(sticker.file_id, in_memory=True)
|
||||
if not isinstance(buffer, BytesIO):
|
||||
data = await download_bytes(client, sticker.file_id)
|
||||
if data is None:
|
||||
return
|
||||
data = buffer.getvalue()
|
||||
storage_key = capture.storage.put(data)
|
||||
await upsert_downloaded(
|
||||
ctx.pool,
|
||||
|
||||
@@ -18,6 +18,12 @@ async def fetch_media(ctx: JobContext) -> None:
|
||||
if isinstance(message, list):
|
||||
message = message[0] if message else None
|
||||
if message is None or message.empty:
|
||||
return
|
||||
msg = f"message {chat_id}/{message_id} is unavailable"
|
||||
raise RuntimeError(msg)
|
||||
toggles = CaptureToggles(media=True, self_destruct_media=True)
|
||||
await capture_media(client, message, capture, chat_id, message_id, toggles)
|
||||
ok = await capture_media(
|
||||
client, message, capture, chat_id, message_id, toggles, force=True
|
||||
)
|
||||
if not ok:
|
||||
msg = f"media download failed for {chat_id}/{message_id}"
|
||||
raise RuntimeError(msg)
|
||||
@@ -46,7 +46,7 @@ async def _save_found(
|
||||
source = _source(peer, users, chats)
|
||||
if source is None:
|
||||
return None
|
||||
chat = Chat._parse_chat(client, source) # noqa: SLF001
|
||||
chat = await Chat._parse_chat(client, source) # noqa: SLF001
|
||||
if chat is None or chat.id is None:
|
||||
return None
|
||||
await save_chat(ctx, chat)
|
||||
|
||||
@@ -5,6 +5,7 @@ from userbot.modules.jobs.context import JobContext
|
||||
from userbot.modules.jobs.registry import register
|
||||
from userbot.modules.profiles.parse import snapshot_from_high_level
|
||||
from userbot.modules.profiles.repository import write_profile
|
||||
from userbot.modules.raw_json import raw_json
|
||||
|
||||
|
||||
@register("sync_contacts")
|
||||
@@ -21,7 +22,7 @@ async def sync_contacts(ctx: JobContext) -> None:
|
||||
if not isinstance(user, User):
|
||||
continue
|
||||
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
|
||||
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user))
|
||||
await write_profile(ctx.pool, ctx.account_id, user.id, fields, raw_json(user))
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
ctx.pool,
|
||||
|
||||
@@ -9,6 +9,8 @@ from userbot.modules.jobs.registry import register
|
||||
from userbot.modules.profiles.parse import snapshot_from_high_level
|
||||
from userbot.modules.profiles.repository import write_profile
|
||||
from userbot.modules.profiles.snapshots import save_group, save_private
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from utils.logging import logger
|
||||
|
||||
SAVE_EVERY = 100
|
||||
USERS_BATCH = 200
|
||||
@@ -31,7 +33,9 @@ async def _enrich_users(client: Client, ctx: CaptureContext, ids: list[int]) ->
|
||||
if not isinstance(user, User):
|
||||
continue
|
||||
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
|
||||
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user))
|
||||
await write_profile(
|
||||
ctx.pool, ctx.account_id, user.id, fields, raw_json(user)
|
||||
)
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
ctx.pool,
|
||||
@@ -66,6 +70,8 @@ async def sync_dialogs(ctx: JobContext) -> None:
|
||||
await save_group(capture, chat)
|
||||
except (BadRequest, Forbidden):
|
||||
pass
|
||||
except Exception:
|
||||
logger.exception(f"Failed to snapshot dialog {chat_id}")
|
||||
await ctx.pool.execute(_UPSERT_DIALOG, ctx.account_id, chat_id)
|
||||
processed += 1
|
||||
if processed % SAVE_EVERY == 0:
|
||||
|
||||
@@ -48,6 +48,16 @@ async def requeue_running(pool: asyncpg.Pool, account_id: int) -> int:
|
||||
return int(result.split()[-1])
|
||||
|
||||
|
||||
async def prune_finished(pool: asyncpg.Pool, account_id: int) -> int:
|
||||
result = await pool.execute(
|
||||
"DELETE FROM jobs WHERE account_id = $1 "
|
||||
"AND status IN ('done', 'canceled', 'failed') "
|
||||
"AND finished_at < now() - interval '30 days'",
|
||||
account_id,
|
||||
)
|
||||
return int(result.split()[-1])
|
||||
|
||||
|
||||
async def save_cursor(pool: asyncpg.Pool, job_id: int, cursor: dict[str, Any]) -> None:
|
||||
await pool.execute(
|
||||
"UPDATE jobs SET cursor = $2::jsonb, updated_at = now() WHERE id = $1",
|
||||
|
||||
@@ -1,13 +1,23 @@
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
from pyrogram import Client
|
||||
from pyrogram.errors import (
|
||||
FileIdInvalid,
|
||||
FileReferenceExpired,
|
||||
FileReferenceInvalid,
|
||||
FloodPremiumWait,
|
||||
FloodWait,
|
||||
)
|
||||
from pyrogram.types import Message
|
||||
|
||||
from userbot.modules.capture import repository
|
||||
from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.download import Downloadable, download_bytes
|
||||
from utils.logging import logger
|
||||
from utils.policy.models import CaptureToggles
|
||||
|
||||
STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)
|
||||
|
||||
_MEDIA_ATTRS = (
|
||||
"photo",
|
||||
"video",
|
||||
@@ -47,6 +57,48 @@ def media_unique_id(message: Message) -> str | None:
|
||||
return getattr(obj, "file_unique_id", None) if obj is not None else None
|
||||
|
||||
|
||||
def _download_target(message: Message) -> Downloadable | None:
|
||||
kind, obj = media_object(message)
|
||||
if obj is None:
|
||||
return None
|
||||
return message if getattr(message, kind or "", None) is obj else obj
|
||||
|
||||
|
||||
async def _fresh_target(
|
||||
client: Client, chat_id: int, message_id: int
|
||||
) -> Downloadable | None:
|
||||
message = await client.get_messages(chat_id, message_id)
|
||||
if not isinstance(message, Message) or message.empty:
|
||||
return None
|
||||
return _download_target(message)
|
||||
|
||||
|
||||
async def _download(
|
||||
client: Client, target: Downloadable | None, chat_id: int, message_id: int
|
||||
) -> bytes | None:
|
||||
if target is None:
|
||||
return None
|
||||
try:
|
||||
return await download_bytes(client, target)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
f"[yellow]Media download timed out for {chat_id}/{message_id}.[/]"
|
||||
)
|
||||
return None
|
||||
except STALE_FILE_ID:
|
||||
fresh = await _fresh_target(client, chat_id, message_id)
|
||||
if fresh is None:
|
||||
return None
|
||||
try:
|
||||
return await download_bytes(client, fresh)
|
||||
except (TimeoutError, *STALE_FILE_ID):
|
||||
logger.warning(
|
||||
f"[yellow]Media download failed after refetch "
|
||||
f"for {chat_id}/{message_id}.[/]"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def capture_media( # noqa: PLR0913
|
||||
client: Client,
|
||||
message: Message,
|
||||
@@ -54,10 +106,12 @@ async def capture_media( # noqa: PLR0913
|
||||
chat_id: int,
|
||||
message_id: int,
|
||||
toggles: CaptureToggles,
|
||||
) -> None:
|
||||
*,
|
||||
force: bool = False,
|
||||
) -> bool:
|
||||
kind, obj = media_object(message)
|
||||
if obj is None:
|
||||
return
|
||||
return True
|
||||
unique_id = getattr(obj, "file_unique_id", None)
|
||||
ttl = getattr(obj, "ttl_seconds", None)
|
||||
want = toggles.self_destruct_media if ttl else toggles.media
|
||||
@@ -65,6 +119,8 @@ async def capture_media( # noqa: PLR0913
|
||||
mime = getattr(obj, "mime_type", None)
|
||||
storage_key: str | None = None
|
||||
downloaded = False
|
||||
evicted = False
|
||||
flood: FloodWait | FloodPremiumWait | None = None
|
||||
if want:
|
||||
existing = await repository.current_media(
|
||||
ctx.pool, ctx.account_id, chat_id, message_id
|
||||
@@ -78,14 +134,19 @@ async def capture_media( # noqa: PLR0913
|
||||
storage_key = existing["storage_key"]
|
||||
file_size = existing["file_size"]
|
||||
downloaded = True
|
||||
else:
|
||||
target = message if getattr(message, kind or "", None) is obj else obj
|
||||
buffer = await client.download_media(target, in_memory=True)
|
||||
if isinstance(buffer, BytesIO):
|
||||
data = buffer.getvalue()
|
||||
elif force or not await repository.is_evicted(ctx.pool, unique_id):
|
||||
target = _download_target(message)
|
||||
try:
|
||||
data = await _download(client, target, chat_id, message_id)
|
||||
except (FloodWait, FloodPremiumWait) as exc:
|
||||
flood = exc
|
||||
data = None
|
||||
if data is not None:
|
||||
storage_key = ctx.storage.put(data)
|
||||
file_size = len(data)
|
||||
downloaded = True
|
||||
else:
|
||||
evicted = True
|
||||
await repository.insert_media(
|
||||
ctx.pool,
|
||||
ctx.account_id,
|
||||
@@ -98,4 +159,14 @@ async def capture_media( # noqa: PLR0913
|
||||
ttl,
|
||||
unique_id,
|
||||
downloaded=downloaded,
|
||||
evicted=evicted,
|
||||
)
|
||||
if (
|
||||
force
|
||||
and storage_key is not None
|
||||
and await repository.unevict(ctx.pool, unique_id, storage_key, file_size)
|
||||
):
|
||||
ctx.storage.purge(storage_key)
|
||||
if flood is not None:
|
||||
raise flood
|
||||
return downloaded or not want
|
||||
@@ -21,10 +21,10 @@ def active_username(usernames: list[raw.types.Username]) -> str | None:
|
||||
return usernames[0].username if usernames else None
|
||||
|
||||
|
||||
def snapshot_from_user(
|
||||
async def snapshot_from_user(
|
||||
client: Client, raw_user: raw.types.User
|
||||
) -> tuple[ProfileFields, str | None, str | None]:
|
||||
user = User._parse(client, raw_user) # noqa: SLF001
|
||||
user = await User._parse(client, raw_user) # noqa: SLF001
|
||||
photo = user.photo if user else None
|
||||
photo_unique_id = photo.big_photo_unique_id if photo else None
|
||||
photo_file_id = photo.big_file_id if photo else None
|
||||
|
||||
@@ -7,12 +7,13 @@ from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.groups.repository import insert_chat_history
|
||||
from userbot.modules.profiles.parse import snapshot_from_chat
|
||||
from userbot.modules.profiles.repository import write_profile
|
||||
from userbot.modules.raw_json import raw_json
|
||||
|
||||
|
||||
async def save_private(ctx: CaptureContext, chat: Chat) -> bool:
|
||||
chat_id = chat.id or 0
|
||||
fields, photo_file_id, photo_unique_id = snapshot_from_chat(chat)
|
||||
await write_profile(ctx.pool, ctx.account_id, chat_id, fields, str(chat))
|
||||
await write_profile(ctx.pool, ctx.account_id, chat_id, fields, raw_json(chat))
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
ctx.pool, ctx.account_id, chat_id, "peer", photo_unique_id, photo_file_id
|
||||
@@ -35,7 +36,7 @@ async def save_group(ctx: CaptureContext, chat: Chat) -> None:
|
||||
photo_unique_id,
|
||||
None,
|
||||
datetime.now(UTC),
|
||||
str(chat),
|
||||
raw_json(chat),
|
||||
)
|
||||
if photo_file_id and photo_unique_id:
|
||||
await note_avatar(
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from pyrogram.raw.core import TLObject
|
||||
from pyrogram.types import Contact, Object
|
||||
|
||||
EMPTY_JSON = "{}"
|
||||
PHONE_KEY = "phone_number"
|
||||
|
||||
|
||||
def _default(obj: Any) -> Any: # noqa: ANN401
|
||||
if isinstance(obj, TLObject):
|
||||
return TLObject.default(obj)
|
||||
data = Object.default(obj)
|
||||
if isinstance(obj, Contact) and isinstance(data, dict) and PHONE_KEY in data:
|
||||
data[PHONE_KEY] = obj.phone_number
|
||||
return data
|
||||
|
||||
|
||||
def raw_json(obj: Object | TLObject | None) -> str:
|
||||
if obj is None:
|
||||
return EMPTY_JSON
|
||||
return json.dumps(obj, indent=4, default=_default, ensure_ascii=False)
|
||||
@@ -1,9 +1,9 @@
|
||||
from io import BytesIO
|
||||
|
||||
from pyrogram import Client
|
||||
from pyrogram.types import Story
|
||||
|
||||
from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.download import download_bytes
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from userbot.modules.stories import repository
|
||||
|
||||
|
||||
@@ -24,9 +24,8 @@ async def save_story(client: Client, capture: CaptureContext, story: Story) -> N
|
||||
capture.pool, capture.account_id, peer_id, story.id
|
||||
)
|
||||
if not (stored or story.deleted or story.media is None):
|
||||
buffer = await client.download_media(story, in_memory=True)
|
||||
if isinstance(buffer, BytesIO):
|
||||
data = buffer.getvalue()
|
||||
data = await download_bytes(client, story)
|
||||
if data is not None:
|
||||
storage_key = capture.storage.put(data)
|
||||
file_size = len(data)
|
||||
downloaded = True
|
||||
@@ -42,7 +41,7 @@ async def save_story(client: Client, capture: CaptureContext, story: Story) -> N
|
||||
storage_key,
|
||||
file_size,
|
||||
story.views,
|
||||
str(story.raw),
|
||||
raw_json(story.raw),
|
||||
pinned=bool(story.pinned),
|
||||
deleted=bool(story.deleted),
|
||||
downloaded=downloaded,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import signal
|
||||
from collections.abc import Callable, Coroutine
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
@@ -9,8 +10,8 @@ import asyncpg
|
||||
import uvloop
|
||||
|
||||
from dependencies.container import container
|
||||
from userbot import PyroClient
|
||||
from userbot.modules.capture import CaptureContext, build_capture_context
|
||||
from userbot.modules.client import PyroClient
|
||||
from userbot.modules.jobs import JobConsumer
|
||||
from utils.env import env
|
||||
from utils.jobs import enqueue
|
||||
@@ -26,6 +27,8 @@ from utils.storage import ContentAddressedStorage
|
||||
|
||||
setup_logging()
|
||||
|
||||
SESSION_FLUSH_INTERVAL = 60
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunningAccount:
|
||||
@@ -46,18 +49,6 @@ async def _cancel(task: asyncio.Task) -> None:
|
||||
await task
|
||||
|
||||
|
||||
async def _enqueue_once(pool: asyncpg.Pool, account_id: int, kind: str) -> None:
|
||||
existing = await pool.fetchval(
|
||||
"SELECT 1 FROM jobs WHERE account_id = $1 AND kind = $2 "
|
||||
"AND status IN ('pending', 'running') LIMIT 1",
|
||||
account_id,
|
||||
kind,
|
||||
)
|
||||
if existing is None:
|
||||
await enqueue(pool, account_id, kind, {})
|
||||
logger.info(f"[green]Queued {kind}.[/]")
|
||||
|
||||
|
||||
class AccountRegistry:
|
||||
def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
|
||||
self._pool = pool
|
||||
@@ -118,8 +109,8 @@ class AccountRegistry:
|
||||
client, asyncio.create_task(consumer.run()), device_model
|
||||
)
|
||||
logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
|
||||
await _enqueue_once(self._pool, account_id, "sync_dialogs")
|
||||
await _enqueue_once(self._pool, account_id, "sync_contacts")
|
||||
await enqueue(self._pool, account_id, "sync_dialogs", {})
|
||||
await enqueue(self._pool, account_id, "sync_contacts", {})
|
||||
|
||||
async def _stop(self, session_name: str) -> None:
|
||||
account = self._running.pop(session_name, None)
|
||||
@@ -149,6 +140,14 @@ class AccountRegistry:
|
||||
logger.info(f"[yellow]Account logged out:[/] {path.stem}")
|
||||
|
||||
|
||||
async def _flush_sessions(registry: AccountRegistry) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(SESSION_FLUSH_INTERVAL)
|
||||
for client in registry.clients:
|
||||
with contextlib.suppress(Exception):
|
||||
await client.storage.save()
|
||||
|
||||
|
||||
async def _listen_changes(
|
||||
registry: AccountRegistry, tasks: set[asyncio.Task]
|
||||
) -> asyncpg.Connection:
|
||||
@@ -193,14 +192,18 @@ async def runner() -> None:
|
||||
registry = AccountRegistry(pool, storage)
|
||||
tasks: set[asyncio.Task] = set()
|
||||
listen_conn: asyncpg.Connection | None = None
|
||||
stop = asyncio.Event()
|
||||
asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set)
|
||||
flush = asyncio.create_task(_flush_sessions(registry))
|
||||
try:
|
||||
await registry.sync()
|
||||
if not registry.clients:
|
||||
logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]")
|
||||
listen_conn = await _listen_changes(registry, tasks)
|
||||
logger.info("[green]Userbot running.[/]")
|
||||
await asyncio.Event().wait()
|
||||
await stop.wait()
|
||||
finally:
|
||||
await _cancel(flush)
|
||||
if listen_conn is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await listen_conn.close()
|
||||
|
||||
@@ -493,3 +493,53 @@ class Dialog(SQLModel, table=True):
|
||||
onupdate=func.now(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class FileShare(SQLModel, table=True):
|
||||
__tablename__ = "file_shares"
|
||||
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
account_id: int
|
||||
token: str = Field(unique=True)
|
||||
kind: str
|
||||
storage_key: str
|
||||
file_name: str
|
||||
mime: str | None = None
|
||||
file_size: int | None = Field(default=None, sa_column=Column(BigInteger))
|
||||
title: str | None = None
|
||||
chat_id: int | None = Field(default=None, sa_column=Column(BigInteger))
|
||||
message_id: int | None = Field(default=None, sa_column=Column(BigInteger))
|
||||
peer_id: int | None = Field(default=None, sa_column=Column(BigInteger))
|
||||
story_id: int | None = Field(default=None, sa_column=Column(BigInteger))
|
||||
expires_at: datetime | None = Field(
|
||||
default=None, sa_column=Column(DateTime(timezone=True))
|
||||
)
|
||||
max_downloads: int | None = None
|
||||
download_count: int = 0
|
||||
last_download_at: datetime | None = Field(
|
||||
default=None, sa_column=Column(DateTime(timezone=True))
|
||||
)
|
||||
revoked_at: datetime | None = Field(
|
||||
default=None, sa_column=Column(DateTime(timezone=True))
|
||||
)
|
||||
created_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class FileShareHit(SQLModel, table=True):
|
||||
__tablename__ = "file_share_hits"
|
||||
|
||||
id: int | None = Field(default=None, sa_column=Column(BigInteger, primary_key=True))
|
||||
share_id: int = Field(foreign_key="file_shares.id")
|
||||
ts: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
method: str
|
||||
ip: str | None = None
|
||||
user_agent: str | None = None
|
||||
counted: bool = False
|
||||
@@ -31,6 +31,12 @@ class TelegramSettings(BaseSettings):
|
||||
sessions_dir: str = "sessions"
|
||||
|
||||
|
||||
class BotsSettings(BaseSettings):
|
||||
sessions_dir: str = "sessions/bots"
|
||||
backfill_per_chat: int = 0
|
||||
backfill_chats: int = 0
|
||||
|
||||
|
||||
class ApiSettings(BaseSettings):
|
||||
host: str = "0.0.0.0" # noqa: S104
|
||||
port: int = 8080
|
||||
@@ -50,12 +56,14 @@ class LogSettings(BaseSettings):
|
||||
level: str = "INFO"
|
||||
level_external: str = "WARNING"
|
||||
show_time: bool = False
|
||||
show_locals: bool = False
|
||||
console_width: int = 150
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
db: DatabaseSettings = Field(default_factory=DatabaseSettings)
|
||||
tg: TelegramSettings = Field(default_factory=TelegramSettings)
|
||||
bots: BotsSettings = Field(default_factory=BotsSettings)
|
||||
api: ApiSettings = Field(default_factory=ApiSettings)
|
||||
auth: AuthSettings = Field(default_factory=AuthSettings)
|
||||
storage: StorageSettings = Field(default_factory=StorageSettings)
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
import re
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from urllib.parse import quote
|
||||
|
||||
_MIME_EXTENSIONS = {
|
||||
"application/pdf": ".pdf",
|
||||
"application/x-tgsticker": ".tgs",
|
||||
"application/zip": ".zip",
|
||||
"audio/mpeg": ".mp3",
|
||||
"audio/ogg": ".ogg",
|
||||
"image/gif": ".gif",
|
||||
"image/jpeg": ".jpg",
|
||||
"image/png": ".png",
|
||||
"image/webp": ".webp",
|
||||
"text/plain": ".txt",
|
||||
"video/mp4": ".mp4",
|
||||
"video/quicktime": ".mov",
|
||||
"video/webm": ".webm",
|
||||
}
|
||||
|
||||
_KIND_EXTENSIONS = {
|
||||
"animation": ".mp4",
|
||||
"audio": ".mp3",
|
||||
"gif": ".mp4",
|
||||
"photo": ".jpg",
|
||||
"sticker": ".webp",
|
||||
"video": ".mp4",
|
||||
"video_note": ".mp4",
|
||||
"voice": ".ogg",
|
||||
}
|
||||
|
||||
_EXTENSION_MIMES = {
|
||||
".flac": "audio/flac",
|
||||
".gif": "image/gif",
|
||||
".heic": "image/heic",
|
||||
".jpeg": "image/jpeg",
|
||||
".jpg": "image/jpeg",
|
||||
".m4a": "audio/mp4",
|
||||
".mov": "video/quicktime",
|
||||
".mp3": "audio/mpeg",
|
||||
".mp4": "video/mp4",
|
||||
".ogg": "audio/ogg",
|
||||
".pdf": "application/pdf",
|
||||
".png": "image/png",
|
||||
".svg": "image/svg+xml",
|
||||
".txt": "text/plain",
|
||||
".wav": "audio/wav",
|
||||
".webm": "video/webm",
|
||||
".webp": "image/webp",
|
||||
}
|
||||
|
||||
_KIND_MIMES = {
|
||||
"animation": "video/mp4",
|
||||
"gif": "video/mp4",
|
||||
"photo": "image/jpeg",
|
||||
"video": "video/mp4",
|
||||
"video_note": "video/mp4",
|
||||
"voice": "audio/ogg",
|
||||
}
|
||||
|
||||
_GENERIC_MIMES = {
|
||||
"application/octet-stream",
|
||||
"application/binary",
|
||||
"binary/octet-stream",
|
||||
"",
|
||||
}
|
||||
|
||||
_INLINE_MIME_PREFIXES = ("image/", "video/", "audio/", "text/")
|
||||
_INLINE_MIMES = {"application/pdf", "application/json"}
|
||||
_ACTIVE_MIMES = {
|
||||
"text/html",
|
||||
"application/xhtml+xml",
|
||||
"image/svg+xml",
|
||||
"text/xml",
|
||||
"application/xml",
|
||||
}
|
||||
|
||||
_PREVIEW_AGENTS = (
|
||||
"telegrambot",
|
||||
"twitterbot",
|
||||
"facebookexternalhit",
|
||||
"whatsapp",
|
||||
"discordbot",
|
||||
"slackbot",
|
||||
"skypeuripreview",
|
||||
"vkshare",
|
||||
"redditbot",
|
||||
"linkedinbot",
|
||||
"embedly",
|
||||
"quora link preview",
|
||||
"googlebot",
|
||||
"bingbot",
|
||||
"applebot",
|
||||
"yandexbot",
|
||||
"duckduckbot",
|
||||
"petalbot",
|
||||
"ahrefsbot",
|
||||
"semrushbot",
|
||||
"headlesschrome",
|
||||
"python-requests",
|
||||
)
|
||||
|
||||
_TRANSLIT = {
|
||||
"а": "a",
|
||||
"б": "b",
|
||||
"в": "v",
|
||||
"г": "g",
|
||||
"д": "d",
|
||||
"е": "e",
|
||||
"ё": "e",
|
||||
"ж": "zh",
|
||||
"з": "z",
|
||||
"и": "i",
|
||||
"й": "y",
|
||||
"к": "k",
|
||||
"л": "l",
|
||||
"м": "m",
|
||||
"н": "n",
|
||||
"о": "o",
|
||||
"п": "p",
|
||||
"р": "r",
|
||||
"с": "s",
|
||||
"т": "t",
|
||||
"у": "u",
|
||||
"ф": "f",
|
||||
"х": "h",
|
||||
"ц": "c",
|
||||
"ч": "ch",
|
||||
"ш": "sh",
|
||||
"щ": "sch",
|
||||
"ъ": "",
|
||||
"ы": "y",
|
||||
"ь": "",
|
||||
"э": "e",
|
||||
"ю": "yu",
|
||||
"я": "ya",
|
||||
}
|
||||
|
||||
_URL_UNSAFE = re.compile(r"[^A-Za-z0-9._-]+")
|
||||
_URL_REPEATS = re.compile(r"_{2,}")
|
||||
|
||||
_UNSAFE_CHARS = re.compile(r'[\\/:*?"<>|\x00-\x1f]+')
|
||||
_SPACES = re.compile(r"\s+")
|
||||
_NAME_LIMIT = 120
|
||||
|
||||
|
||||
def extension_for(kind: str | None, mime: str | None) -> str:
|
||||
if mime and mime in _MIME_EXTENSIONS:
|
||||
return _MIME_EXTENSIONS[mime]
|
||||
if kind and kind in _KIND_EXTENSIONS:
|
||||
return _KIND_EXTENSIONS[kind]
|
||||
if mime and "/" in mime:
|
||||
tail = mime.rsplit("/", 1)[1].split(";")[0].strip()
|
||||
if tail.isalnum():
|
||||
return f".{tail}"
|
||||
return ".bin"
|
||||
|
||||
|
||||
def sanitize_name(name: str) -> str:
|
||||
cleaned = _SPACES.sub(" ", _UNSAFE_CHARS.sub("_", name)).strip(" .")
|
||||
if len(cleaned) > _NAME_LIMIT:
|
||||
head, dot, tail = cleaned.rpartition(".")
|
||||
cleaned = (
|
||||
f"{head[: _NAME_LIMIT - len(tail) - 1]}{dot}{tail}"
|
||||
if dot and len(tail) < 12 # noqa: PLR2004
|
||||
else cleaned[:_NAME_LIMIT]
|
||||
)
|
||||
return cleaned or "file"
|
||||
|
||||
|
||||
def url_slug(file_name: str) -> str:
|
||||
lowered = "".join(
|
||||
_TRANSLIT.get(ch, _TRANSLIT.get(ch.lower(), ch)) for ch in file_name
|
||||
)
|
||||
slug = _URL_REPEATS.sub("_", _URL_UNSAFE.sub("_", lowered)).strip("_.")
|
||||
if not slug:
|
||||
return "file"
|
||||
return slug[:_NAME_LIMIT]
|
||||
|
||||
|
||||
def media_file_name(
|
||||
kind: str | None, mime: str | None, message_id: int, original: str | None = None
|
||||
) -> str:
|
||||
if original:
|
||||
name = sanitize_name(original)
|
||||
return name if "." in name else f"{name}{extension_for(kind, mime)}"
|
||||
return f"{kind or 'media'}_{message_id}{extension_for(kind, mime)}"
|
||||
|
||||
|
||||
def story_file_name(peer_id: int, story_id: int, media_kind: str | None) -> str:
|
||||
return f"story_{peer_id}_{story_id}{extension_for(media_kind or 'photo', None)}"
|
||||
|
||||
|
||||
def resolve_mime(
|
||||
kind: str | None, mime: str | None, file_name: str | None = None
|
||||
) -> str:
|
||||
if mime and mime.lower() not in _GENERIC_MIMES:
|
||||
return mime
|
||||
if file_name and "." in file_name:
|
||||
by_extension = _EXTENSION_MIMES.get(f".{file_name.rsplit('.', 1)[1].lower()}")
|
||||
if by_extension:
|
||||
return by_extension
|
||||
return _KIND_MIMES.get(kind or "", mime or "application/octet-stream")
|
||||
|
||||
|
||||
def is_inline_mime(mime: str | None) -> bool:
|
||||
if not mime:
|
||||
return False
|
||||
return mime in _INLINE_MIMES or mime.startswith(_INLINE_MIME_PREFIXES)
|
||||
|
||||
|
||||
def is_active_mime(mime: str | None) -> bool:
|
||||
return bool(mime) and mime.split(";", 1)[0].strip().lower() in _ACTIVE_MIMES
|
||||
|
||||
|
||||
def content_disposition(file_name: str, *, attachment: bool) -> str:
|
||||
kind = "attachment" if attachment else "inline"
|
||||
ascii_name = file_name.encode("ascii", "replace").decode("ascii").replace('"', "_")
|
||||
return f"{kind}; filename=\"{ascii_name}\"; filename*=UTF-8''{quote(file_name)}"
|
||||
|
||||
|
||||
def is_preview_agent(user_agent: str | None) -> bool:
|
||||
if not user_agent:
|
||||
return False
|
||||
lowered = user_agent.lower()
|
||||
return any(marker in lowered for marker in _PREVIEW_AGENTS)
|
||||
|
||||
|
||||
def counts_as_download(
|
||||
method: str, user_agent: str | None, range_header: str | None
|
||||
) -> bool:
|
||||
if method.upper() != "GET" or is_preview_agent(user_agent):
|
||||
return False
|
||||
if range_header is None:
|
||||
return True
|
||||
return range_header.replace(" ", "").startswith("bytes=0-")
|
||||
|
||||
|
||||
def expiry_from_seconds(seconds: int | None) -> datetime | None:
|
||||
if seconds is None or seconds <= 0:
|
||||
return None
|
||||
return datetime.now(UTC) + timedelta(seconds=seconds)
|
||||
@@ -0,0 +1,6 @@
|
||||
import asyncio
|
||||
from importlib import import_module
|
||||
|
||||
|
||||
async def preload(module: str) -> None:
|
||||
await asyncio.to_thread(import_module, module)
|
||||
@@ -9,12 +9,22 @@ JOBS_CHANGED_CHANNEL = "jobs_changed"
|
||||
async def enqueue(
|
||||
pool: asyncpg.Pool, account_id: int, kind: str, params: dict[str, Any]
|
||||
) -> int:
|
||||
params_json = json.dumps(params)
|
||||
existing = await pool.fetchval(
|
||||
"SELECT id FROM jobs WHERE account_id = $1 AND kind = $2 "
|
||||
"AND params = $3::jsonb AND status IN ('pending', 'running') LIMIT 1",
|
||||
account_id,
|
||||
kind,
|
||||
params_json,
|
||||
)
|
||||
if existing is not None:
|
||||
return existing
|
||||
job_id = await pool.fetchval(
|
||||
"INSERT INTO jobs (account_id, kind, params) "
|
||||
"VALUES ($1, $2, $3::jsonb) RETURNING id",
|
||||
account_id,
|
||||
kind,
|
||||
json.dumps(params),
|
||||
params_json,
|
||||
)
|
||||
await pool.execute(f"NOTIFY {JOBS_CHANGED_CHANNEL}")
|
||||
return job_id
|
||||
@@ -20,13 +20,13 @@ def setup_logging() -> None:
|
||||
markup=True,
|
||||
rich_tracebacks=True,
|
||||
enable_link_path=False,
|
||||
tracebacks_show_locals=True,
|
||||
tracebacks_show_locals=env.log.show_locals,
|
||||
omit_repeated_times=False,
|
||||
show_time=env.log.show_time,
|
||||
)
|
||||
],
|
||||
)
|
||||
install(console=console, show_locals=True)
|
||||
install(console=console, show_locals=env.log.show_locals)
|
||||
|
||||
|
||||
logger = logging.getLogger("beavergram")
|
||||
|
||||
@@ -79,6 +79,24 @@ class FolderSpec:
|
||||
is_chatlist: bool = False
|
||||
|
||||
|
||||
class PolicyParent(BaseModel):
|
||||
scope_type: ScopeType
|
||||
scope_id: int | None
|
||||
title: str | None = None
|
||||
toggles: CaptureToggles
|
||||
|
||||
|
||||
class PolicyRuleView(PolicyRecord):
|
||||
title: str | None = None
|
||||
parent: PolicyParent | None = None
|
||||
|
||||
|
||||
class PolicyScopeView(BaseModel):
|
||||
title: str | None
|
||||
rule: PolicyRecord | None
|
||||
parent: PolicyParent | None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PolicySet:
|
||||
chat: dict[int, CaptureToggles] = field(default_factory=dict)
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.policy import repository
|
||||
from utils.policy.models import (
|
||||
ChatKind,
|
||||
ChatMeta,
|
||||
FolderSpec,
|
||||
PolicyParent,
|
||||
PolicyRecord,
|
||||
PolicyRuleView,
|
||||
PolicyScopeView,
|
||||
PolicySet,
|
||||
ScopeType,
|
||||
)
|
||||
from utils.policy.resolver import resolve_parent
|
||||
from utils.read import chats
|
||||
from utils.read.models import ChatListItem
|
||||
|
||||
|
||||
def _meta(chat_id: int, chat: ChatListItem | None) -> ChatMeta:
|
||||
if chat_id > 0:
|
||||
kind = ChatKind.DM
|
||||
elif chat is not None and chat.is_broadcast:
|
||||
kind = ChatKind.CHANNEL
|
||||
else:
|
||||
kind = ChatKind.GROUP
|
||||
return ChatMeta(
|
||||
chat_id=chat_id,
|
||||
kind=kind,
|
||||
is_bot=bool(chat and chat.is_bot),
|
||||
is_contact=chat.is_contact if chat else None,
|
||||
)
|
||||
|
||||
|
||||
def _without_chats(policies: PolicySet) -> PolicySet:
|
||||
return PolicySet(folder=policies.folder, defaults=policies.defaults)
|
||||
|
||||
|
||||
def _chat_parent(
|
||||
chat_id: int,
|
||||
chat: ChatListItem | None,
|
||||
folders: list[FolderSpec],
|
||||
policies: PolicySet,
|
||||
) -> PolicyParent:
|
||||
return resolve_parent(_meta(chat_id, chat), folders, _without_chats(policies))
|
||||
|
||||
|
||||
async def list_rules(pool: asyncpg.Pool, account_id: int) -> list[PolicyRuleView]:
|
||||
records = await repository.list_policies(pool, account_id)
|
||||
folders = await repository.list_folders(pool, account_id)
|
||||
policies = await repository.load_policy_set(pool, account_id)
|
||||
chat_ids = [
|
||||
record.scope_id
|
||||
for record in records
|
||||
if record.scope_type is ScopeType.CHAT and record.scope_id is not None
|
||||
]
|
||||
known = {
|
||||
chat.chat_id: chat for chat in await chats.get_chats(pool, account_id, chat_ids)
|
||||
}
|
||||
titles = {folder.folder_id: folder.title for folder in folders}
|
||||
rules: list[PolicyRuleView] = []
|
||||
for record in records:
|
||||
rule = PolicyRuleView(**record.model_dump())
|
||||
if record.scope_type is ScopeType.FOLDER and record.scope_id is not None:
|
||||
rule.title = titles.get(record.scope_id)
|
||||
if record.scope_type is ScopeType.CHAT and record.scope_id is not None:
|
||||
chat = known.get(record.scope_id)
|
||||
rule.title = chat.title if chat else None
|
||||
rule.parent = _chat_parent(record.scope_id, chat, folders, policies)
|
||||
rules.append(rule)
|
||||
return rules
|
||||
|
||||
|
||||
async def _rule(
|
||||
pool: asyncpg.Pool, account_id: int, scope_type: ScopeType, scope_id: int
|
||||
) -> PolicyRecord | None:
|
||||
return await repository.find_policy(pool, account_id, scope_type, scope_id)
|
||||
|
||||
|
||||
async def chat_scope(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
) -> PolicyScopeView:
|
||||
folders = await repository.list_folders(pool, account_id)
|
||||
policies = await repository.load_policy_set(pool, account_id)
|
||||
chat = await chats.get_chat(pool, account_id, chat_id)
|
||||
return PolicyScopeView(
|
||||
title=chat.title if chat else None,
|
||||
rule=await _rule(pool, account_id, ScopeType.CHAT, chat_id),
|
||||
parent=_chat_parent(chat_id, chat, folders, policies),
|
||||
)
|
||||
|
||||
|
||||
async def folder_scope(
|
||||
pool: asyncpg.Pool, account_id: int, folder_id: int
|
||||
) -> PolicyScopeView | None:
|
||||
folder = await repository.get_folder(pool, account_id, folder_id)
|
||||
if folder is None:
|
||||
return None
|
||||
return PolicyScopeView(
|
||||
title=folder.title,
|
||||
rule=await _rule(pool, account_id, ScopeType.FOLDER, folder_id),
|
||||
parent=None,
|
||||
)
|
||||
@@ -1,11 +1,14 @@
|
||||
from collections.abc import Iterable
|
||||
|
||||
from utils.policy.models import (
|
||||
DEFAULT_SCOPE_BY_KIND,
|
||||
CaptureToggles,
|
||||
ChatKind,
|
||||
ChatMeta,
|
||||
FolderSpec,
|
||||
PolicyParent,
|
||||
PolicySet,
|
||||
ScopeType,
|
||||
)
|
||||
|
||||
|
||||
@@ -33,14 +36,29 @@ def folder_contains(folder: FolderSpec, chat: ChatMeta) -> bool:
|
||||
return _category_match(folder, chat)
|
||||
|
||||
|
||||
def resolve_parent(
|
||||
chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet
|
||||
) -> PolicyParent:
|
||||
for folder in sorted(folders, key=lambda f: f.order_index):
|
||||
toggles = policies.folder.get(folder.folder_id)
|
||||
if toggles is not None and folder_contains(folder, chat):
|
||||
return PolicyParent(
|
||||
scope_type=ScopeType.FOLDER,
|
||||
scope_id=folder.folder_id,
|
||||
title=folder.title,
|
||||
toggles=toggles,
|
||||
)
|
||||
return PolicyParent(
|
||||
scope_type=DEFAULT_SCOPE_BY_KIND[chat.kind],
|
||||
scope_id=None,
|
||||
toggles=policies.defaults.get(chat.kind, CaptureToggles()),
|
||||
)
|
||||
|
||||
|
||||
def resolve(
|
||||
chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet
|
||||
) -> CaptureToggles:
|
||||
chat_override = policies.chat.get(chat.chat_id)
|
||||
if chat_override is not None:
|
||||
return chat_override
|
||||
for folder in sorted(folders, key=lambda f: f.order_index):
|
||||
toggles = policies.folder.get(folder.folder_id)
|
||||
if toggles is not None and folder_contains(folder, chat):
|
||||
return toggles
|
||||
return policies.defaults.get(chat.kind, CaptureToggles())
|
||||
return resolve_parent(chat, folders, policies).toggles
|
||||
@@ -7,17 +7,18 @@ from utils.read.models import AccountView
|
||||
|
||||
ACCOUNTS_CHANGED_CHANNEL = "accounts_changed"
|
||||
|
||||
_ACCOUNT_COLS = "account_id, label, phone, tg_user_id, is_active, device_model"
|
||||
_ACCOUNT_COLS = "account_id, label, phone, tg_user_id, is_active, is_bot, device_model"
|
||||
|
||||
_UPSERT_ACCOUNT = """
|
||||
INSERT INTO accounts
|
||||
(tg_user_id, label, phone, session_name, is_active, raw, updated_at)
|
||||
VALUES ($1, $2, $3, $4, TRUE, $5::jsonb, now())
|
||||
(tg_user_id, label, phone, session_name, is_active, is_bot, raw, updated_at)
|
||||
VALUES ($1, $2, $3, $4, TRUE, $6, $5::jsonb, now())
|
||||
ON CONFLICT (tg_user_id) DO UPDATE SET
|
||||
label = EXCLUDED.label,
|
||||
phone = EXCLUDED.phone,
|
||||
session_name = EXCLUDED.session_name,
|
||||
is_active = TRUE,
|
||||
is_bot = EXCLUDED.is_bot,
|
||||
raw = EXCLUDED.raw,
|
||||
updated_at = now()
|
||||
RETURNING account_id
|
||||
@@ -30,10 +31,48 @@ RETURNING {_ACCOUNT_COLS}
|
||||
""" # noqa: S608
|
||||
|
||||
|
||||
_self_ids: dict[int, int] = {}
|
||||
|
||||
|
||||
async def self_user_id(pool: asyncpg.Pool, account_id: int) -> int | None:
|
||||
return await pool.fetchval(
|
||||
cached = _self_ids.get(account_id)
|
||||
if cached is not None:
|
||||
return cached
|
||||
tg_user_id = await pool.fetchval(
|
||||
"SELECT tg_user_id FROM accounts WHERE account_id = $1", account_id
|
||||
)
|
||||
if tg_user_id is not None:
|
||||
_self_ids[account_id] = tg_user_id
|
||||
return tg_user_id
|
||||
|
||||
|
||||
_known_ids: set[int] = set()
|
||||
|
||||
|
||||
def _account_label(row: asyncpg.Record) -> str:
|
||||
named = f", {row['label']}" if row["label"] else ""
|
||||
return f"account_id={row['account_id']} (Telegram id {row['tg_user_id']}{named})"
|
||||
|
||||
|
||||
async def unknown_account_hint(pool: asyncpg.Pool, account_id: int) -> str | None:
|
||||
"""Explain an account_id that matches no account, or None if it is valid."""
|
||||
if account_id in _known_ids:
|
||||
return None
|
||||
rows = await pool.fetch(
|
||||
"SELECT account_id, tg_user_id, label FROM accounts ORDER BY account_id"
|
||||
)
|
||||
_known_ids.update(row["account_id"] for row in rows)
|
||||
if account_id in _known_ids:
|
||||
return None
|
||||
known = "; ".join(_account_label(row) for row in rows) or "none"
|
||||
mistaken = next((row for row in rows if row["tg_user_id"] == account_id), None)
|
||||
if mistaken is not None:
|
||||
return (
|
||||
f"{account_id} is a Telegram user id, not an account_id. In beavergram "
|
||||
f"that account is account_id={mistaken['account_id']}. "
|
||||
f"Retry with account_id={mistaken['account_id']}."
|
||||
)
|
||||
return f"No account with account_id={account_id}. Archived accounts: {known}."
|
||||
|
||||
|
||||
async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]:
|
||||
@@ -51,7 +90,9 @@ async def get_account(pool: asyncpg.Pool, account_id: int) -> AccountView | None
|
||||
return AccountView(**dict(row)) if row else None
|
||||
|
||||
|
||||
async def sync_account(pool: asyncpg.Pool, me: User, session_name: str) -> int:
|
||||
async def sync_account(
|
||||
pool: asyncpg.Pool, me: User, session_name: str, *, is_bot: bool = False
|
||||
) -> int:
|
||||
raw = json.dumps(
|
||||
{
|
||||
"id": me.id,
|
||||
@@ -63,7 +104,7 @@ async def sync_account(pool: asyncpg.Pool, me: User, session_name: str) -> int:
|
||||
)
|
||||
label = " ".join(filter(None, [me.first_name, me.last_name])) or me.username
|
||||
return await pool.fetchval(
|
||||
_UPSERT_ACCOUNT, me.id, label, me.phone_number, session_name, raw
|
||||
_UPSERT_ACCOUNT, me.id, label, me.phone_number, session_name, raw, is_bot
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -33,7 +33,7 @@ async def message_volume(
|
||||
|
||||
|
||||
async def response_stats(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, *, days: int = 365
|
||||
) -> ResponseStats:
|
||||
self_id = await self_user_id(pool, account_id)
|
||||
rows = await pool.fetch(
|
||||
@@ -43,6 +43,7 @@ async def response_stats(
|
||||
"lag(date) OVER w AS prev_date "
|
||||
"FROM messages "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND sender_id IS NOT NULL "
|
||||
"AND date >= $4 "
|
||||
"WINDOW w AS (ORDER BY date, message_id)), "
|
||||
"resp AS ("
|
||||
"SELECT (sender_id = $3) AS is_mine, "
|
||||
@@ -55,6 +56,7 @@ async def response_stats(
|
||||
account_id,
|
||||
chat_id,
|
||||
self_id,
|
||||
datetime.now(UTC) - timedelta(days=days),
|
||||
)
|
||||
stats = ResponseStats(
|
||||
mine_median_seconds=None, mine_count=0, their_median_seconds=None, their_count=0
|
||||
|
||||
@@ -21,7 +21,7 @@ WHERE account_id = $1 AND owner_id = $2 AND unique_id = $3
|
||||
_AVATAR_HISTORY = """
|
||||
SELECT unique_id, first_seen_at, downloaded FROM avatars
|
||||
WHERE account_id = $1 AND owner_id = $2
|
||||
ORDER BY first_seen_at DESC
|
||||
ORDER BY first_seen_at DESC LIMIT 200
|
||||
"""
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Registry of bot tokens the watcher holds.
|
||||
|
||||
Only session bookkeeping lives here: what the archive itself keeps about a
|
||||
bot is an ordinary account row and ordinary messages, read through
|
||||
``utils.read.accounts`` and ``utils.read.chats`` like any other account.
|
||||
"""
|
||||
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import BotView
|
||||
|
||||
BOTS_CHANGED_CHANNEL = "watched_bots_changed"
|
||||
|
||||
STATUS_PENDING = "pending"
|
||||
STATUS_OK = "ok"
|
||||
STATUS_UNAUTHORIZED = "unauthorized"
|
||||
STATUS_ERROR = "error"
|
||||
STATUS_REMOVED = "removed"
|
||||
|
||||
_BOT_COLS = """
|
||||
SELECT bot_id, account_id, username, title, enabled, status, status_detail,
|
||||
backfilled_at, checked_at, added_at
|
||||
FROM watched_bots
|
||||
"""
|
||||
|
||||
_REGISTER = """
|
||||
INSERT INTO watched_bots (bot_id, status) VALUES ($1, 'pending')
|
||||
ON CONFLICT (bot_id) DO UPDATE SET
|
||||
enabled = true, status = 'pending', status_detail = NULL, updated_at = now()
|
||||
"""
|
||||
|
||||
_REMOVE = """
|
||||
UPDATE watched_bots
|
||||
SET enabled = false, status = 'removed', status_detail = NULL, updated_at = now()
|
||||
WHERE bot_id = $1 AND enabled
|
||||
"""
|
||||
|
||||
_BACKLOG_CURSOR = """
|
||||
SELECT backfill_pts, backfill_qts, backfill_date
|
||||
FROM watched_bots
|
||||
WHERE bot_id = $1 AND backfill_pts IS NOT NULL
|
||||
"""
|
||||
|
||||
_SAVE_BACKLOG_CURSOR = """
|
||||
UPDATE watched_bots SET
|
||||
backfill_pts = $2, backfill_qts = $3, backfill_date = $4, updated_at = now()
|
||||
WHERE bot_id = $1
|
||||
"""
|
||||
|
||||
_MARK_BACKFILLED = """
|
||||
UPDATE watched_bots SET
|
||||
backfilled_at = now(),
|
||||
backfill_pts = NULL,
|
||||
backfill_qts = NULL,
|
||||
backfill_date = NULL,
|
||||
updated_at = now()
|
||||
WHERE bot_id = $1
|
||||
"""
|
||||
|
||||
_SET_STATUS = """
|
||||
UPDATE watched_bots SET
|
||||
status = $2,
|
||||
status_detail = $3,
|
||||
username = coalesce($4, username),
|
||||
title = coalesce($5, title),
|
||||
checked_at = now(),
|
||||
updated_at = now()
|
||||
WHERE bot_id = $1
|
||||
"""
|
||||
|
||||
|
||||
async def notify_bots_changed(pool: asyncpg.Pool) -> None:
|
||||
await pool.execute("SELECT pg_notify($1, '')", BOTS_CHANGED_CHANNEL)
|
||||
|
||||
|
||||
async def list_bots(pool: asyncpg.Pool) -> list[BotView]:
|
||||
rows = await pool.fetch(f"{_BOT_COLS} ORDER BY added_at")
|
||||
return [BotView(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
async def get_bot(pool: asyncpg.Pool, bot_id: int) -> BotView | None:
|
||||
row = await pool.fetchrow(f"{_BOT_COLS} WHERE bot_id = $1", bot_id)
|
||||
return BotView(**dict(row)) if row else None
|
||||
|
||||
|
||||
async def enabled_bot_ids(pool: asyncpg.Pool) -> set[int]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT bot_id FROM watched_bots WHERE enabled AND status <> $1",
|
||||
STATUS_UNAUTHORIZED,
|
||||
)
|
||||
return {row["bot_id"] for row in rows}
|
||||
|
||||
|
||||
async def removed_bot_ids(pool: asyncpg.Pool) -> set[int]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT bot_id FROM watched_bots WHERE NOT enabled AND status = $1",
|
||||
STATUS_REMOVED,
|
||||
)
|
||||
return {row["bot_id"] for row in rows}
|
||||
|
||||
|
||||
async def known_bot_ids(pool: asyncpg.Pool) -> set[int]:
|
||||
rows = await pool.fetch("SELECT bot_id FROM watched_bots")
|
||||
return {row["bot_id"] for row in rows}
|
||||
|
||||
|
||||
async def register_bot(pool: asyncpg.Pool, bot_id: int) -> None:
|
||||
await pool.execute(_REGISTER, bot_id)
|
||||
|
||||
|
||||
async def remove_bot(pool: asyncpg.Pool, bot_id: int) -> bool:
|
||||
return await pool.execute(_REMOVE, bot_id) != "UPDATE 0"
|
||||
|
||||
|
||||
async def link_account(pool: asyncpg.Pool, bot_id: int, account_id: int) -> None:
|
||||
await pool.execute(
|
||||
"UPDATE watched_bots SET account_id = $2, updated_at = now() WHERE bot_id = $1",
|
||||
bot_id,
|
||||
account_id,
|
||||
)
|
||||
|
||||
|
||||
async def set_status( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
bot_id: int,
|
||||
status: str,
|
||||
*,
|
||||
detail: str | None = None,
|
||||
username: str | None = None,
|
||||
title: str | None = None,
|
||||
) -> None:
|
||||
await pool.execute(_SET_STATUS, bot_id, status, detail, username, title)
|
||||
|
||||
|
||||
async def needs_backfill(pool: asyncpg.Pool, bot_id: int) -> bool:
|
||||
return not await pool.fetchval(
|
||||
"SELECT backfilled_at IS NOT NULL FROM watched_bots WHERE bot_id = $1", bot_id
|
||||
)
|
||||
|
||||
|
||||
async def backlog_cursor(
|
||||
pool: asyncpg.Pool, bot_id: int
|
||||
) -> tuple[int, int, int] | None:
|
||||
row = await pool.fetchrow(_BACKLOG_CURSOR, bot_id)
|
||||
if row is None:
|
||||
return None
|
||||
return row["backfill_pts"], row["backfill_qts"], row["backfill_date"]
|
||||
|
||||
|
||||
async def save_backlog_cursor(
|
||||
pool: asyncpg.Pool, bot_id: int, pts: int, qts: int, date: int
|
||||
) -> None:
|
||||
await pool.execute(_SAVE_BACKLOG_CURSOR, bot_id, pts, qts, date)
|
||||
|
||||
|
||||
async def mark_backfilled(pool: asyncpg.Pool, bot_id: int) -> None:
|
||||
await pool.execute(_MARK_BACKFILLED, bot_id)
|
||||
@@ -62,6 +62,18 @@ _ONE_ID = """
|
||||
WHERE account_id = $1 AND scope_type = 'chat' AND scope_id = $2
|
||||
"""
|
||||
|
||||
_MANY_IDS = """
|
||||
SELECT chat_id FROM chat_stats
|
||||
WHERE account_id = $1 AND chat_id = ANY($2::bigint[])
|
||||
UNION
|
||||
SELECT chat_id FROM dialogs
|
||||
WHERE account_id = $1 AND chat_id = ANY($2::bigint[])
|
||||
UNION
|
||||
SELECT scope_id AS chat_id FROM capture_policy
|
||||
WHERE account_id = $1 AND scope_type = 'chat'
|
||||
AND scope_id = ANY($2::bigint[])
|
||||
"""
|
||||
|
||||
_CHAT_ROWS = """
|
||||
WITH ids AS ({ids})
|
||||
SELECT ids.chat_id,
|
||||
@@ -164,6 +176,18 @@ async def list_chats(
|
||||
return [_chat_item(row) for row in rows]
|
||||
|
||||
|
||||
async def folder_chat_ids(
|
||||
pool: asyncpg.Pool, account_id: int, folder: FolderSpec
|
||||
) -> list[int]:
|
||||
rows_sql = _CHAT_ROWS.format(ids=_ALL_IDS)
|
||||
rows = await pool.fetch(
|
||||
f"SELECT chat.chat_id FROM ({rows_sql}) chat WHERE {_folder_filter(1)}", # noqa: S608
|
||||
account_id,
|
||||
*_folder_params(folder),
|
||||
)
|
||||
return [row["chat_id"] for row in rows]
|
||||
|
||||
|
||||
async def get_chat(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
) -> ChatListItem | None:
|
||||
@@ -171,6 +195,15 @@ async def get_chat(
|
||||
return _chat_item(row) if row is not None else None
|
||||
|
||||
|
||||
async def get_chats(
|
||||
pool: asyncpg.Pool, account_id: int, ids: list[int]
|
||||
) -> list[ChatListItem]:
|
||||
if not ids:
|
||||
return []
|
||||
rows = await pool.fetch(_CHAT_ROWS.format(ids=_MANY_IDS), account_id, ids)
|
||||
return [_chat_item(row) for row in rows]
|
||||
|
||||
|
||||
async def get_chat_history( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
@@ -308,7 +341,7 @@ async def get_message_versions(
|
||||
rows = await pool.fetch(
|
||||
"SELECT observed_at, edit_date, text FROM message_versions "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
|
||||
"ORDER BY observed_at",
|
||||
"ORDER BY observed_at LIMIT 500",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
|
||||
@@ -54,9 +54,8 @@ WITH chat_meta AS (
|
||||
bool_or(is_broadcast) AS is_broadcast
|
||||
FROM hits GROUP BY chat_id
|
||||
), counts AS (
|
||||
SELECT chat_id, count(*) AS message_count FROM messages
|
||||
SELECT chat_id, message_count FROM chat_stats
|
||||
WHERE account_id = $1 AND chat_id IN (SELECT chat_id FROM merged)
|
||||
GROUP BY chat_id
|
||||
)
|
||||
SELECT m.chat_id, m.title, m.username, m.is_bot, m.is_contact, m.is_broadcast,
|
||||
COALESCE(c.message_count, 0) AS message_count,
|
||||
|
||||
Loaded 100 of 283 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user