feat: add api and mcp
This commit is contained in:
+41
-3
@@ -4,18 +4,44 @@ from contextlib import asynccontextmanager
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka, setup_dishka
|
||||
from fastapi import FastAPI
|
||||
from fastmcp.utilities.lifespan import combine_lifespans
|
||||
from starlette.applications import Starlette
|
||||
|
||||
from api.routers import backfill, folders, policy
|
||||
from api.auth import BearerAuthMiddleware
|
||||
from api.mcp.server import mcp
|
||||
from api.routers import (
|
||||
annotations,
|
||||
backfill,
|
||||
chats,
|
||||
folders,
|
||||
media,
|
||||
peers,
|
||||
policy,
|
||||
presence,
|
||||
search,
|
||||
social,
|
||||
watches,
|
||||
)
|
||||
from dependencies.container import container
|
||||
from utils.env import env
|
||||
|
||||
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()
|
||||
|
||||
mcp_app = mcp.http_app(path="/")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app_: FastAPI) -> AsyncGenerator[None]:
|
||||
async def lifespan(app_: Starlette) -> AsyncGenerator[None]:
|
||||
yield
|
||||
await app_.state.dishka_container.close()
|
||||
|
||||
|
||||
app = FastAPI(title="beavergram API", lifespan=lifespan)
|
||||
app = FastAPI(
|
||||
title="beavergram API", lifespan=combine_lifespans(lifespan, mcp_app.lifespan)
|
||||
)
|
||||
app.router.route_class = DishkaRoute
|
||||
|
||||
|
||||
@@ -31,5 +57,17 @@ async def health(pool: FromDishka[asyncpg.Pool]) -> dict[str, bool]:
|
||||
app.include_router(policy.router)
|
||||
app.include_router(folders.router)
|
||||
app.include_router(backfill.router)
|
||||
app.include_router(search.router)
|
||||
app.include_router(chats.router)
|
||||
app.include_router(media.router)
|
||||
app.include_router(social.router)
|
||||
app.include_router(presence.router)
|
||||
app.include_router(peers.router)
|
||||
app.include_router(annotations.router)
|
||||
app.include_router(watches.router)
|
||||
|
||||
app.mount("/mcp", mcp_app)
|
||||
|
||||
app.add_middleware(BearerAuthMiddleware, token=_token)
|
||||
|
||||
setup_dishka(container, app)
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
PROTECTED_PREFIXES = ("/api", "/mcp")
|
||||
_UNAUTHORIZED = b'{"detail":"unauthorized"}'
|
||||
|
||||
|
||||
class BearerAuthMiddleware:
|
||||
def __init__(self, app: ASGIApp, token: str) -> None:
|
||||
self.app = app
|
||||
self.token = token
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
if scope["method"] == "OPTIONS" or not scope["path"].startswith(
|
||||
PROTECTED_PREFIXES
|
||||
):
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
headers = dict(scope["headers"])
|
||||
authorization = headers.get(b"authorization", b"").decode()
|
||||
if authorization == f"Bearer {self.token}":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 401,
|
||||
"headers": [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"www-authenticate", b"Bearer"),
|
||||
],
|
||||
}
|
||||
)
|
||||
await send({"type": "http.response.body", "body": _UNAUTHORIZED})
|
||||
@@ -0,0 +1,233 @@
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
from fastmcp import FastMCP
|
||||
from pydantic import BaseModel
|
||||
|
||||
from dependencies.container import container
|
||||
from utils.jobs import enqueue
|
||||
from utils.read import annotations, chats, media, peers, presence, social, watches
|
||||
from utils.read.models import DEFAULT_LIMIT, Page
|
||||
from utils.search.models import SearchFilters
|
||||
from utils.search.repository import search_messages
|
||||
|
||||
mcp: FastMCP = FastMCP("beavergram")
|
||||
|
||||
|
||||
async def _pool() -> asyncpg.Pool:
|
||||
return await container.get(asyncpg.Pool)
|
||||
|
||||
|
||||
def _dump(items: Sequence[BaseModel]) -> list[dict[str, Any]]:
|
||||
return [item.model_dump(mode="json") for item in items]
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def search_messages_tool(
|
||||
account_id: int,
|
||||
query: str | None = None,
|
||||
chat_id: int | None = None,
|
||||
sender_id: int | None = None,
|
||||
has_media: bool | None = None,
|
||||
date_from: datetime | None = None,
|
||||
date_to: datetime | None = None,
|
||||
regex: str | None = None,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Full-text search over message text and STT transcripts."""
|
||||
filters = SearchFilters(
|
||||
account_id=account_id,
|
||||
query=query,
|
||||
chat_id=chat_id,
|
||||
sender_id=sender_id,
|
||||
has_media=has_media,
|
||||
date_from=date_from,
|
||||
date_to=date_to,
|
||||
regex=regex,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
return _dump(await search_messages(await _pool(), filters))
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def list_chats(
|
||||
account_id: int, limit: int = DEFAULT_LIMIT, offset: int = 0
|
||||
) -> list[dict[str, Any]]:
|
||||
"""List archived chats with message counts and last activity."""
|
||||
page = Page(limit=limit, offset=offset)
|
||||
return _dump(await chats.list_chats(await _pool(), account_id, page))
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def get_chat_history(
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
offset: int = 0,
|
||||
include_deleted: bool = True,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Read archived messages of a chat, newest first."""
|
||||
return _dump(
|
||||
await chats.get_chat_history(
|
||||
await _pool(),
|
||||
account_id,
|
||||
chat_id,
|
||||
Page(limit=limit, offset=offset),
|
||||
include_deleted=include_deleted,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def get_deleted_messages(
|
||||
account_id: int,
|
||||
chat_id: int | None = None,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""List messages that were deleted in Telegram but kept in the archive."""
|
||||
return _dump(
|
||||
await chats.get_deleted_messages(
|
||||
await _pool(), account_id, Page(limit=limit, offset=offset), chat_id=chat_id
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def get_message_versions(
|
||||
account_id: int, chat_id: int, message_id: int
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Get the edit history of a message."""
|
||||
return _dump(
|
||||
await chats.get_message_versions(await _pool(), account_id, chat_id, message_id)
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def get_media(
|
||||
account_id: int, chat_id: int, message_id: int, fetch: bool = False
|
||||
) -> dict[str, Any] | None:
|
||||
"""Get media metadata for a message; set fetch=True to enqueue lazy download."""
|
||||
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")
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def get_callbacks(
|
||||
account_id: int, chat_id: int, message_id: int
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Get bot inline-button callback data (hex) for a message."""
|
||||
items = await social.get_callbacks(await _pool(), account_id, chat_id, message_id)
|
||||
return _dump(items)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def presence_history(
|
||||
account_id: int,
|
||||
peer_id: int,
|
||||
date_from: datetime | None = None,
|
||||
date_to: datetime | None = None,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Get online/offline status history of a peer."""
|
||||
return _dump(
|
||||
await presence.presence_history(
|
||||
await _pool(),
|
||||
account_id,
|
||||
peer_id,
|
||||
Page(limit=limit, offset=offset),
|
||||
date_from=date_from,
|
||||
date_to=date_to,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def get_peer_history(account_id: int, peer_id: int) -> list[dict[str, Any]]:
|
||||
"""Get name/username/avatar change history of a contact."""
|
||||
return _dump(await peers.get_peer_history(await _pool(), account_id, peer_id))
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def get_stories(
|
||||
account_id: int,
|
||||
peer_id: int | None = None,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""List archived stories of contacts."""
|
||||
return _dump(
|
||||
await peers.get_stories(
|
||||
await _pool(), account_id, Page(limit=limit, offset=offset), peer_id=peer_id
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def get_annotations(
|
||||
account_id: int,
|
||||
chat_id: int | None = None,
|
||||
message_id: int | None = None,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Read user annotations on messages (read-only via MCP)."""
|
||||
return _dump(
|
||||
await annotations.list_annotations(
|
||||
await _pool(),
|
||||
account_id,
|
||||
Page(limit=limit, offset=offset),
|
||||
chat_id=chat_id,
|
||||
message_id=message_id,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def set_watch(
|
||||
account_id: int,
|
||||
kind: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
enabled: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
"""Create a local watch rule (the only local write MCP is allowed)."""
|
||||
watch = await watches.create_watch(
|
||||
await _pool(), account_id, kind, params or {}, enabled=enabled
|
||||
)
|
||||
return watch.model_dump(mode="json")
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def list_watches(account_id: int) -> list[dict[str, Any]]:
|
||||
"""List local watch rules."""
|
||||
return _dump(await watches.list_watches(await _pool(), account_id))
|
||||
|
||||
|
||||
@mcp.tool
|
||||
async def list_alerts(
|
||||
account_id: int,
|
||||
seen: bool | None = None,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""List fired alerts from watch rules."""
|
||||
return _dump(
|
||||
await watches.list_alerts(
|
||||
await _pool(), account_id, Page(limit=limit, offset=offset), seen=seen
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,77 @@
|
||||
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.read import annotations
|
||||
from utils.read.models import DEFAULT_LIMIT, AnnotationView, Page
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/api/annotations", tags=["annotations"], route_class=DishkaRoute
|
||||
)
|
||||
|
||||
|
||||
class AnnotationCreate(BaseModel):
|
||||
account_id: int
|
||||
chat_id: int
|
||||
message_id: int
|
||||
text: str
|
||||
|
||||
|
||||
class AnnotationUpdate(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def list_annotations(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
chat_id: Annotated[int | None, Query()] = None,
|
||||
message_id: Annotated[int | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[AnnotationView]:
|
||||
return await annotations.list_annotations(
|
||||
pool,
|
||||
account_id,
|
||||
Page(limit=limit, offset=offset),
|
||||
chat_id=chat_id,
|
||||
message_id=message_id,
|
||||
)
|
||||
|
||||
|
||||
@router.post("", status_code=201)
|
||||
async def create_annotation(
|
||||
pool: FromDishka[asyncpg.Pool], body: AnnotationCreate
|
||||
) -> AnnotationView:
|
||||
return await annotations.create_annotation(
|
||||
pool, body.account_id, body.chat_id, body.message_id, body.text
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{annotation_id}")
|
||||
async def get_annotation(
|
||||
pool: FromDishka[asyncpg.Pool], annotation_id: int
|
||||
) -> AnnotationView:
|
||||
annotation = await annotations.get_annotation(pool, annotation_id)
|
||||
if annotation is None:
|
||||
raise HTTPException(status_code=404, detail="annotation not found")
|
||||
return annotation
|
||||
|
||||
|
||||
@router.put("/{annotation_id}")
|
||||
async def update_annotation(
|
||||
pool: FromDishka[asyncpg.Pool], annotation_id: int, body: AnnotationUpdate
|
||||
) -> AnnotationView:
|
||||
annotation = await annotations.update_annotation(pool, annotation_id, body.text)
|
||||
if annotation is None:
|
||||
raise HTTPException(status_code=404, detail="annotation not found")
|
||||
return annotation
|
||||
|
||||
|
||||
@router.delete("/{annotation_id}", status_code=204)
|
||||
async def delete_annotation(pool: FromDishka[asyncpg.Pool], annotation_id: int) -> None:
|
||||
if not await annotations.delete_annotation(pool, annotation_id):
|
||||
raise HTTPException(status_code=404, detail="annotation not found")
|
||||
@@ -7,9 +7,9 @@ from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from pydantic import BaseModel
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["backfill"], route_class=DishkaRoute)
|
||||
from utils.jobs import enqueue
|
||||
|
||||
JOBS_CHANGED_CHANNEL = "jobs_changed"
|
||||
router = APIRouter(prefix="/api", tags=["backfill"], route_class=DishkaRoute)
|
||||
|
||||
|
||||
class BackfillRequest(BaseModel):
|
||||
@@ -52,25 +52,11 @@ def _to_view(row: asyncpg.Record) -> JobView:
|
||||
return JobView(**data)
|
||||
|
||||
|
||||
async def _enqueue(
|
||||
pool: asyncpg.Pool, account_id: int, kind: str, params: dict[str, Any]
|
||||
) -> int:
|
||||
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),
|
||||
)
|
||||
await pool.execute(f"NOTIFY {JOBS_CHANGED_CHANNEL}")
|
||||
return job_id
|
||||
|
||||
|
||||
@router.post("/backfill", status_code=201)
|
||||
async def enqueue_backfill(
|
||||
pool: FromDishka[asyncpg.Pool], body: BackfillRequest
|
||||
) -> EnqueueResponse:
|
||||
job_id = await _enqueue(
|
||||
job_id = await enqueue(
|
||||
pool,
|
||||
body.account_id,
|
||||
"backfill",
|
||||
@@ -83,7 +69,7 @@ async def enqueue_backfill(
|
||||
async def enqueue_fetch_media(
|
||||
pool: FromDishka[asyncpg.Pool], body: FetchMediaRequest
|
||||
) -> EnqueueResponse:
|
||||
job_id = await _enqueue(
|
||||
job_id = await enqueue(
|
||||
pool,
|
||||
body.account_id,
|
||||
"fetch_media",
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, Query
|
||||
|
||||
from utils.read import chats
|
||||
from utils.read.models import (
|
||||
DEFAULT_LIMIT,
|
||||
ChatListItem,
|
||||
MessageVersionView,
|
||||
MessageView,
|
||||
Page,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["chats"], route_class=DishkaRoute)
|
||||
|
||||
AccountId = Annotated[int, Query()]
|
||||
Limit = Annotated[int, Query()]
|
||||
Offset = Annotated[int, Query()]
|
||||
|
||||
|
||||
@router.get("/chats")
|
||||
async def list_chats(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: AccountId,
|
||||
limit: Limit = DEFAULT_LIMIT,
|
||||
offset: Offset = 0,
|
||||
) -> list[ChatListItem]:
|
||||
return await chats.list_chats(pool, account_id, Page(limit=limit, offset=offset))
|
||||
|
||||
|
||||
@router.get("/chats/{chat_id}/messages")
|
||||
async def chat_history(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
chat_id: int,
|
||||
account_id: AccountId,
|
||||
limit: Limit = DEFAULT_LIMIT,
|
||||
offset: Offset = 0,
|
||||
include_deleted: Annotated[bool, Query()] = True,
|
||||
) -> list[MessageView]:
|
||||
return await chats.get_chat_history(
|
||||
pool,
|
||||
account_id,
|
||||
chat_id,
|
||||
Page(limit=limit, offset=offset),
|
||||
include_deleted=include_deleted,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/chats/{chat_id}/messages/{message_id}/versions")
|
||||
async def message_versions(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, message_id: int, account_id: AccountId
|
||||
) -> list[MessageVersionView]:
|
||||
return await chats.get_message_versions(pool, account_id, chat_id, message_id)
|
||||
|
||||
|
||||
@router.get("/deleted")
|
||||
async def deleted_messages(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: AccountId,
|
||||
chat_id: Annotated[int | None, Query()] = None,
|
||||
limit: Limit = DEFAULT_LIMIT,
|
||||
offset: Offset = 0,
|
||||
) -> list[MessageView]:
|
||||
return await chats.get_deleted_messages(
|
||||
pool, account_id, Page(limit=limit, offset=offset), chat_id=chat_id
|
||||
)
|
||||
@@ -0,0 +1,38 @@
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from utils.read.media import get_media
|
||||
from utils.read.models import MediaView
|
||||
from utils.storage import ContentAddressedStorage
|
||||
|
||||
router = APIRouter(prefix="/api/media", tags=["media"], route_class=DishkaRoute)
|
||||
|
||||
|
||||
@router.get("/{media_id}/meta")
|
||||
async def media_meta(pool: FromDishka[asyncpg.Pool], media_id: int) -> MediaView:
|
||||
media = await get_media(pool, media_id)
|
||||
if media is None:
|
||||
raise HTTPException(status_code=404, detail="media not found")
|
||||
return media
|
||||
|
||||
|
||||
@router.get("/{media_id}")
|
||||
async def serve_media(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
storage: FromDishka[ContentAddressedStorage],
|
||||
media_id: int,
|
||||
) -> FileResponse:
|
||||
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="media not downloaded; enqueue fetch via POST /api/media/fetch",
|
||||
)
|
||||
return FileResponse(
|
||||
storage.url(media.storage_key),
|
||||
media_type=media.mime or "application/octet-stream",
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
|
||||
from utils.read import peers
|
||||
from utils.read.models import DEFAULT_LIMIT, Page, PeerHistoryView, PeerView, StoryView
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["peers"], route_class=DishkaRoute)
|
||||
|
||||
AccountId = Annotated[int, Query()]
|
||||
|
||||
|
||||
@router.get("/peers/{peer_id}")
|
||||
async def get_peer(
|
||||
pool: FromDishka[asyncpg.Pool], peer_id: int, account_id: AccountId
|
||||
) -> PeerView:
|
||||
peer = await peers.get_peer(pool, account_id, peer_id)
|
||||
if peer is None:
|
||||
raise HTTPException(status_code=404, detail="peer not found")
|
||||
return peer
|
||||
|
||||
|
||||
@router.get("/peers/{peer_id}/history")
|
||||
async def peer_history(
|
||||
pool: FromDishka[asyncpg.Pool], peer_id: int, account_id: AccountId
|
||||
) -> list[PeerHistoryView]:
|
||||
return await peers.get_peer_history(pool, account_id, peer_id)
|
||||
|
||||
|
||||
@router.get("/stories")
|
||||
async def stories(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: AccountId,
|
||||
peer_id: Annotated[int | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[StoryView]:
|
||||
return await peers.get_stories(
|
||||
pool, account_id, Page(limit=limit, offset=offset), peer_id=peer_id
|
||||
)
|
||||
@@ -0,0 +1,49 @@
|
||||
from datetime import datetime
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, Query
|
||||
|
||||
from utils.read import presence
|
||||
from utils.read.models import DEFAULT_LIMIT, Page, PresenceHourly, PresenceSample
|
||||
|
||||
router = APIRouter(prefix="/api/presence", tags=["presence"], route_class=DishkaRoute)
|
||||
|
||||
AccountId = Annotated[int, Query()]
|
||||
PeerId = Annotated[int, Query()]
|
||||
DateFrom = Annotated[datetime | None, Query()]
|
||||
DateTo = Annotated[datetime | None, Query()]
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def presence_history(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: AccountId,
|
||||
peer_id: PeerId,
|
||||
date_from: DateFrom = None,
|
||||
date_to: DateTo = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[PresenceSample]:
|
||||
return await presence.presence_history(
|
||||
pool,
|
||||
account_id,
|
||||
peer_id,
|
||||
Page(limit=limit, offset=offset),
|
||||
date_from=date_from,
|
||||
date_to=date_to,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/hourly")
|
||||
async def presence_hourly(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: AccountId,
|
||||
peer_id: PeerId,
|
||||
date_from: DateFrom = None,
|
||||
date_to: DateTo = None,
|
||||
) -> list[PresenceHourly]:
|
||||
return await presence.presence_hourly(
|
||||
pool, account_id, peer_id, date_from=date_from, date_to=date_to
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, Query
|
||||
|
||||
from utils.search.models import SearchFilters, SearchHit
|
||||
from utils.search.repository import search_messages
|
||||
|
||||
router = APIRouter(prefix="/api/search", tags=["search"], route_class=DishkaRoute)
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def search(
|
||||
pool: FromDishka[asyncpg.Pool], filters: Annotated[SearchFilters, Query()]
|
||||
) -> list[SearchHit]:
|
||||
return await search_messages(pool, filters)
|
||||
@@ -0,0 +1,37 @@
|
||||
from typing import Annotated
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, Query
|
||||
|
||||
from utils.read import social
|
||||
from utils.read.models import CallbackView, LinkView, ReactionView
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/api/messages/{chat_id}/{message_id}",
|
||||
tags=["social"],
|
||||
route_class=DishkaRoute,
|
||||
)
|
||||
|
||||
AccountId = Annotated[int, Query()]
|
||||
|
||||
|
||||
@router.get("/callbacks")
|
||||
async def message_callbacks(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, message_id: int, account_id: AccountId
|
||||
) -> list[CallbackView]:
|
||||
return await social.get_callbacks(pool, account_id, chat_id, message_id)
|
||||
|
||||
|
||||
@router.get("/reactions")
|
||||
async def message_reactions(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, message_id: int, account_id: AccountId
|
||||
) -> list[ReactionView]:
|
||||
return await social.get_reactions(pool, account_id, chat_id, message_id)
|
||||
|
||||
|
||||
@router.get("/links")
|
||||
async def message_links(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, message_id: int, account_id: AccountId
|
||||
) -> list[LinkView]:
|
||||
return await social.get_links(pool, account_id, chat_id, message_id)
|
||||
@@ -0,0 +1,82 @@
|
||||
from typing import Annotated, Any
|
||||
|
||||
import asyncpg
|
||||
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils.read import watches
|
||||
from utils.read.models import DEFAULT_LIMIT, AlertView, Page, WatchView
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["watches"], route_class=DishkaRoute)
|
||||
|
||||
|
||||
class WatchCreate(BaseModel):
|
||||
account_id: int
|
||||
kind: str
|
||||
params: dict[str, Any] = {}
|
||||
enabled: bool = True
|
||||
|
||||
|
||||
class WatchUpdate(BaseModel):
|
||||
params: dict[str, Any] = {}
|
||||
enabled: bool = True
|
||||
|
||||
|
||||
@router.get("/watches")
|
||||
async def list_watches(
|
||||
pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()]
|
||||
) -> list[WatchView]:
|
||||
return await watches.list_watches(pool, account_id)
|
||||
|
||||
|
||||
@router.post("/watches", status_code=201)
|
||||
async def create_watch(pool: FromDishka[asyncpg.Pool], body: WatchCreate) -> WatchView:
|
||||
return await watches.create_watch(
|
||||
pool, body.account_id, body.kind, body.params, enabled=body.enabled
|
||||
)
|
||||
|
||||
|
||||
@router.get("/watches/{watch_id}")
|
||||
async def get_watch(pool: FromDishka[asyncpg.Pool], watch_id: int) -> WatchView:
|
||||
watch = await watches.get_watch(pool, watch_id)
|
||||
if watch is None:
|
||||
raise HTTPException(status_code=404, detail="watch not found")
|
||||
return watch
|
||||
|
||||
|
||||
@router.put("/watches/{watch_id}")
|
||||
async def update_watch(
|
||||
pool: FromDishka[asyncpg.Pool], watch_id: int, body: WatchUpdate
|
||||
) -> WatchView:
|
||||
watch = await watches.update_watch(
|
||||
pool, watch_id, body.params, enabled=body.enabled
|
||||
)
|
||||
if watch is None:
|
||||
raise HTTPException(status_code=404, detail="watch not found")
|
||||
return watch
|
||||
|
||||
|
||||
@router.delete("/watches/{watch_id}", status_code=204)
|
||||
async def delete_watch(pool: FromDishka[asyncpg.Pool], watch_id: int) -> None:
|
||||
if not await watches.delete_watch(pool, watch_id):
|
||||
raise HTTPException(status_code=404, detail="watch not found")
|
||||
|
||||
|
||||
@router.get("/alerts")
|
||||
async def list_alerts(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
seen: Annotated[bool | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[AlertView]:
|
||||
return await watches.list_alerts(
|
||||
pool, account_id, Page(limit=limit, offset=offset), seen=seen
|
||||
)
|
||||
|
||||
|
||||
@router.post("/alerts/{alert_id}/seen", status_code=204)
|
||||
async def mark_alert_seen(pool: FromDishka[asyncpg.Pool], alert_id: int) -> None:
|
||||
if not await watches.mark_alert_seen(pool, alert_id):
|
||||
raise HTTPException(status_code=404, detail="alert not found")
|
||||
@@ -410,3 +410,69 @@ class Link(SQLModel, table=True):
|
||||
raw: dict[str, Any] = Field(
|
||||
default_factory=dict, sa_column=Column(JSONB, nullable=False)
|
||||
)
|
||||
|
||||
|
||||
class Annotation(SQLModel, table=True):
|
||||
__tablename__ = "annotations"
|
||||
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
account_id: int
|
||||
chat_id: int = Field(sa_column=Column(BigInteger, nullable=False, index=True))
|
||||
message_id: int = Field(sa_column=Column(BigInteger, nullable=False))
|
||||
text: str
|
||||
created_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
updated_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class Watch(SQLModel, table=True):
|
||||
__tablename__ = "watches"
|
||||
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
account_id: int
|
||||
kind: str
|
||||
params: dict[str, Any] = Field(
|
||||
default_factory=dict, sa_column=Column(JSONB, nullable=False)
|
||||
)
|
||||
enabled: bool = True
|
||||
created_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
updated_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class Alert(SQLModel, table=True):
|
||||
__tablename__ = "alerts"
|
||||
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
account_id: int
|
||||
watch_id: int = Field(foreign_key="watches.id")
|
||||
ts: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))
|
||||
payload: dict[str, Any] = Field(
|
||||
default_factory=dict, sa_column=Column(JSONB, nullable=False)
|
||||
)
|
||||
seen: bool = False
|
||||
created_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
|
||||
@@ -36,6 +36,10 @@ class ApiSettings(BaseSettings):
|
||||
port: int = 8080
|
||||
|
||||
|
||||
class AuthSettings(BaseSettings):
|
||||
token: SecretStr | None = None
|
||||
|
||||
|
||||
class StorageSettings(BaseSettings):
|
||||
root: str = "storage"
|
||||
shard_depth: int = 2
|
||||
@@ -52,6 +56,7 @@ class Settings(BaseSettings):
|
||||
db: DatabaseSettings = Field(default_factory=DatabaseSettings)
|
||||
tg: TelegramSettings = Field(default_factory=TelegramSettings)
|
||||
api: ApiSettings = Field(default_factory=ApiSettings)
|
||||
auth: AuthSettings = Field(default_factory=AuthSettings)
|
||||
storage: StorageSettings = Field(default_factory=StorageSettings)
|
||||
log: LogSettings = Field(default_factory=LogSettings)
|
||||
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
JOBS_CHANGED_CHANNEL = "jobs_changed"
|
||||
|
||||
|
||||
async def enqueue(
|
||||
pool: asyncpg.Pool, account_id: int, kind: str, params: dict[str, Any]
|
||||
) -> int:
|
||||
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),
|
||||
)
|
||||
await pool.execute(f"NOTIFY {JOBS_CHANGED_CHANNEL}")
|
||||
return job_id
|
||||
@@ -0,0 +1,72 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import AnnotationView, Page
|
||||
|
||||
_COLS = "id, account_id, chat_id, message_id, text, created_at, updated_at"
|
||||
|
||||
|
||||
async def list_annotations(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
page: Page,
|
||||
*,
|
||||
chat_id: int | None = None,
|
||||
message_id: int | None = None,
|
||||
) -> list[AnnotationView]:
|
||||
params: list[object] = [account_id]
|
||||
where = "account_id = $1"
|
||||
if chat_id is not None:
|
||||
params.append(chat_id)
|
||||
where += f" AND chat_id = ${len(params)}"
|
||||
if message_id is not None:
|
||||
params.append(message_id)
|
||||
where += f" AND message_id = ${len(params)}"
|
||||
params.append(page.capped_limit)
|
||||
params.append(page.offset)
|
||||
rows = await pool.fetch(
|
||||
f"SELECT {_COLS} FROM annotations WHERE {where} " # noqa: S608
|
||||
f"ORDER BY created_at DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
|
||||
*params,
|
||||
)
|
||||
return [AnnotationView(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
async def get_annotation(
|
||||
pool: asyncpg.Pool, annotation_id: int
|
||||
) -> AnnotationView | None:
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT {_COLS} FROM annotations WHERE id = $1", # noqa: S608
|
||||
annotation_id,
|
||||
)
|
||||
return AnnotationView(**dict(row)) if row else None
|
||||
|
||||
|
||||
async def create_annotation(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int, text: str
|
||||
) -> AnnotationView:
|
||||
row = await pool.fetchrow(
|
||||
"INSERT INTO annotations (account_id, chat_id, message_id, text) " # noqa: S608
|
||||
f"VALUES ($1, $2, $3, $4) RETURNING {_COLS}",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
text,
|
||||
)
|
||||
return AnnotationView(**dict(row))
|
||||
|
||||
|
||||
async def update_annotation(
|
||||
pool: asyncpg.Pool, annotation_id: int, text: str
|
||||
) -> AnnotationView | None:
|
||||
row = await pool.fetchrow(
|
||||
"UPDATE annotations SET text = $2, updated_at = now() " # noqa: S608
|
||||
f"WHERE id = $1 RETURNING {_COLS}",
|
||||
annotation_id,
|
||||
text,
|
||||
)
|
||||
return AnnotationView(**dict(row)) if row else None
|
||||
|
||||
|
||||
async def delete_annotation(pool: asyncpg.Pool, annotation_id: int) -> bool:
|
||||
result = await pool.execute("DELETE FROM annotations WHERE id = $1", annotation_id)
|
||||
return result.endswith("1")
|
||||
@@ -0,0 +1,105 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import ChatListItem, MessageVersionView, MessageView, Page
|
||||
|
||||
_MESSAGE_COLS = (
|
||||
"chat_id, message_id, date, sender_id, text, "
|
||||
"has_media, is_self_destruct, edited_at, deleted_at"
|
||||
)
|
||||
|
||||
|
||||
def _peer_title(
|
||||
first: str | None, last: str | None, username: str | None
|
||||
) -> str | None:
|
||||
name = " ".join(part for part in (first, last) if part)
|
||||
return name or username
|
||||
|
||||
|
||||
async def list_chats(
|
||||
pool: asyncpg.Pool, account_id: int, page: Page
|
||||
) -> list[ChatListItem]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT m.chat_id, count(*) AS message_count, max(m.date) AS last_date, "
|
||||
"(SELECT p.first_name FROM peers p "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS first_name, "
|
||||
"(SELECT p.last_name FROM peers p "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS last_name, "
|
||||
"(SELECT p.username FROM peers p "
|
||||
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS username, "
|
||||
"(SELECT ch.title FROM chat_history ch "
|
||||
"WHERE ch.account_id = $1 AND ch.chat_id = m.chat_id "
|
||||
"AND ch.title IS NOT NULL ORDER BY ch.ts DESC LIMIT 1) AS group_title "
|
||||
"FROM messages m WHERE m.account_id = $1 "
|
||||
"GROUP BY m.chat_id ORDER BY last_date DESC LIMIT $2 OFFSET $3",
|
||||
account_id,
|
||||
page.capped_limit,
|
||||
page.offset,
|
||||
)
|
||||
items = []
|
||||
for row in rows:
|
||||
title = row["group_title"] or _peer_title(
|
||||
row["first_name"], row["last_name"], row["username"]
|
||||
)
|
||||
items.append(
|
||||
ChatListItem(
|
||||
chat_id=row["chat_id"],
|
||||
title=title,
|
||||
message_count=row["message_count"],
|
||||
last_date=row["last_date"],
|
||||
)
|
||||
)
|
||||
return items
|
||||
|
||||
|
||||
async def get_chat_history(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
page: Page,
|
||||
*,
|
||||
include_deleted: bool = True,
|
||||
) -> list[MessageView]:
|
||||
where = "account_id = $1 AND chat_id = $2"
|
||||
if not include_deleted:
|
||||
where += " AND deleted_at IS NULL"
|
||||
rows = await pool.fetch(
|
||||
f"SELECT {_MESSAGE_COLS} FROM messages WHERE {where} " # noqa: S608
|
||||
"ORDER BY date DESC, message_id DESC LIMIT $3 OFFSET $4",
|
||||
account_id,
|
||||
chat_id,
|
||||
page.capped_limit,
|
||||
page.offset,
|
||||
)
|
||||
return [MessageView(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
async def get_deleted_messages(
|
||||
pool: asyncpg.Pool, account_id: int, page: Page, *, chat_id: int | None = None
|
||||
) -> list[MessageView]:
|
||||
params: list[object] = [account_id]
|
||||
where = "account_id = $1 AND deleted_at IS NOT NULL"
|
||||
if chat_id is not None:
|
||||
params.append(chat_id)
|
||||
where += f" AND chat_id = ${len(params)}"
|
||||
params.append(page.capped_limit)
|
||||
params.append(page.offset)
|
||||
rows = await pool.fetch(
|
||||
f"SELECT {_MESSAGE_COLS} FROM messages WHERE {where} " # noqa: S608
|
||||
f"ORDER BY deleted_at DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
|
||||
*params,
|
||||
)
|
||||
return [MessageView(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
async def get_message_versions(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||
) -> list[MessageVersionView]:
|
||||
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",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
return [MessageVersionView(**dict(row)) for row in rows]
|
||||
@@ -0,0 +1,29 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import MediaView
|
||||
|
||||
_MEDIA_COLS = (
|
||||
"id, account_id, chat_id, message_id, kind, storage_key, file_size, "
|
||||
"mime, ttl_seconds, downloaded, extracted_text, created_at"
|
||||
)
|
||||
|
||||
|
||||
async def get_media(pool: asyncpg.Pool, media_id: int) -> MediaView | None:
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT {_MEDIA_COLS} FROM media WHERE id = $1", # noqa: S608
|
||||
media_id,
|
||||
)
|
||||
return MediaView(**dict(row)) if row else None
|
||||
|
||||
|
||||
async def get_message_media(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||
) -> MediaView | None:
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT {_MEDIA_COLS} FROM media " # noqa: S608
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
return MediaView(**dict(row)) if row else None
|
||||
@@ -0,0 +1,160 @@
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
DEFAULT_LIMIT = 50
|
||||
MAX_LIMIT = 500
|
||||
|
||||
|
||||
class Page(BaseModel):
|
||||
limit: int = DEFAULT_LIMIT
|
||||
offset: int = 0
|
||||
|
||||
@property
|
||||
def capped_limit(self) -> int:
|
||||
return min(self.limit, MAX_LIMIT)
|
||||
|
||||
|
||||
class ChatListItem(BaseModel):
|
||||
chat_id: int
|
||||
title: str | None
|
||||
message_count: int
|
||||
last_date: datetime | None
|
||||
|
||||
|
||||
class MessageView(BaseModel):
|
||||
chat_id: int
|
||||
message_id: int
|
||||
date: datetime
|
||||
sender_id: int | None
|
||||
text: str | None
|
||||
has_media: bool
|
||||
is_self_destruct: bool
|
||||
edited_at: datetime | None
|
||||
deleted_at: datetime | None
|
||||
|
||||
|
||||
class MessageVersionView(BaseModel):
|
||||
observed_at: datetime
|
||||
edit_date: datetime | None
|
||||
text: str | None
|
||||
|
||||
|
||||
class MediaView(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
chat_id: int
|
||||
message_id: int
|
||||
kind: str
|
||||
storage_key: str | None
|
||||
file_size: int | None
|
||||
mime: str | None
|
||||
ttl_seconds: int | None
|
||||
downloaded: bool
|
||||
extracted_text: str | None
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class CallbackView(BaseModel):
|
||||
position: int
|
||||
label: str | None
|
||||
data: str | None
|
||||
|
||||
|
||||
class ReactionView(BaseModel):
|
||||
peer_id: int
|
||||
reaction: str
|
||||
added_at: datetime
|
||||
removed_at: datetime | None
|
||||
|
||||
|
||||
class LinkView(BaseModel):
|
||||
position: int
|
||||
url: str
|
||||
kind: str
|
||||
web_url: str | None
|
||||
web_title: str | None
|
||||
web_description: str | None
|
||||
web_site_name: str | None
|
||||
|
||||
|
||||
class PresenceSample(BaseModel):
|
||||
peer_id: int
|
||||
ts: datetime
|
||||
status: str
|
||||
last_online_date: datetime | None
|
||||
next_offline_date: datetime | None
|
||||
|
||||
|
||||
class PresenceHourly(BaseModel):
|
||||
peer_id: int
|
||||
bucket: datetime
|
||||
samples: int
|
||||
online_samples: int
|
||||
last_seen: datetime | None
|
||||
|
||||
|
||||
class PeerView(BaseModel):
|
||||
peer_id: int
|
||||
first_name: str | None
|
||||
last_name: str | None
|
||||
username: str | None
|
||||
phone: str | None
|
||||
photo_unique_id: str | None
|
||||
is_deleted_account: bool
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class PeerHistoryView(BaseModel):
|
||||
observed_at: datetime
|
||||
first_name: str | None
|
||||
last_name: str | None
|
||||
username: str | None
|
||||
phone: str | None
|
||||
photo_unique_id: str | None
|
||||
is_deleted_account: bool
|
||||
|
||||
|
||||
class StoryView(BaseModel):
|
||||
peer_id: int
|
||||
story_id: int
|
||||
date: datetime | None
|
||||
expire_date: datetime | None
|
||||
caption: str | None
|
||||
media_kind: str | None
|
||||
storage_key: str | None
|
||||
downloaded: bool
|
||||
views: int | None
|
||||
pinned: bool
|
||||
deleted: bool
|
||||
|
||||
|
||||
class AnnotationView(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
chat_id: int
|
||||
message_id: int
|
||||
text: str
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class WatchView(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
kind: str
|
||||
params: dict[str, Any]
|
||||
enabled: bool
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class AlertView(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
watch_id: int
|
||||
ts: datetime
|
||||
payload: dict[str, Any]
|
||||
seen: bool
|
||||
created_at: datetime
|
||||
@@ -0,0 +1,50 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import Page, PeerHistoryView, PeerView, StoryView
|
||||
|
||||
|
||||
async def get_peer(
|
||||
pool: asyncpg.Pool, account_id: int, peer_id: int
|
||||
) -> PeerView | None:
|
||||
row = await pool.fetchrow(
|
||||
"SELECT peer_id, first_name, last_name, username, phone, "
|
||||
"photo_unique_id, is_deleted_account, updated_at FROM peers "
|
||||
"WHERE account_id = $1 AND peer_id = $2",
|
||||
account_id,
|
||||
peer_id,
|
||||
)
|
||||
return PeerView(**dict(row)) if row else None
|
||||
|
||||
|
||||
async def get_peer_history(
|
||||
pool: asyncpg.Pool, account_id: int, peer_id: int
|
||||
) -> list[PeerHistoryView]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT observed_at, first_name, last_name, username, phone, "
|
||||
"photo_unique_id, is_deleted_account FROM peer_history "
|
||||
"WHERE account_id = $1 AND peer_id = $2 ORDER BY observed_at DESC",
|
||||
account_id,
|
||||
peer_id,
|
||||
)
|
||||
return [PeerHistoryView(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
async def get_stories(
|
||||
pool: asyncpg.Pool, account_id: int, page: Page, *, peer_id: int | None = None
|
||||
) -> list[StoryView]:
|
||||
params: list[object] = [account_id]
|
||||
where = "account_id = $1"
|
||||
if peer_id is not None:
|
||||
params.append(peer_id)
|
||||
where += f" AND peer_id = ${len(params)}"
|
||||
params.append(page.capped_limit)
|
||||
params.append(page.offset)
|
||||
rows = await pool.fetch(
|
||||
"SELECT peer_id, story_id, date, expire_date, caption, media_kind, " # noqa: S608
|
||||
"storage_key, downloaded, views, pinned, deleted FROM stories "
|
||||
f"WHERE {where} "
|
||||
f"ORDER BY date DESC NULLS LAST LIMIT ${len(params) - 1} "
|
||||
f"OFFSET ${len(params)}",
|
||||
*params,
|
||||
)
|
||||
return [StoryView(**dict(row)) for row in rows]
|
||||
@@ -0,0 +1,57 @@
|
||||
from datetime import datetime
|
||||
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import Page, PresenceHourly, PresenceSample
|
||||
|
||||
|
||||
async def presence_history( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
peer_id: int,
|
||||
page: Page,
|
||||
*,
|
||||
date_from: datetime | None = None,
|
||||
date_to: datetime | None = None,
|
||||
) -> list[PresenceSample]:
|
||||
params: list[object] = [account_id, peer_id]
|
||||
where = "account_id = $1 AND peer_id = $2"
|
||||
if date_from is not None:
|
||||
params.append(date_from)
|
||||
where += f" AND ts >= ${len(params)}"
|
||||
if date_to is not None:
|
||||
params.append(date_to)
|
||||
where += f" AND ts <= ${len(params)}"
|
||||
params.append(page.capped_limit)
|
||||
params.append(page.offset)
|
||||
rows = await pool.fetch(
|
||||
"SELECT peer_id, ts, status, last_online_date, next_offline_date " # noqa: S608
|
||||
f"FROM presence WHERE {where} "
|
||||
f"ORDER BY ts DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
|
||||
*params,
|
||||
)
|
||||
return [PresenceSample(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
async def presence_hourly(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
peer_id: int,
|
||||
*,
|
||||
date_from: datetime | None = None,
|
||||
date_to: datetime | None = None,
|
||||
) -> list[PresenceHourly]:
|
||||
params: list[object] = [account_id, peer_id]
|
||||
where = "account_id = $1 AND peer_id = $2"
|
||||
if date_from is not None:
|
||||
params.append(date_from)
|
||||
where += f" AND bucket >= ${len(params)}"
|
||||
if date_to is not None:
|
||||
params.append(date_to)
|
||||
where += f" AND bucket <= ${len(params)}"
|
||||
rows = await pool.fetch(
|
||||
"SELECT peer_id, bucket, samples, online_samples, last_seen " # noqa: S608
|
||||
f"FROM presence_hourly WHERE {where} ORDER BY bucket DESC",
|
||||
*params,
|
||||
)
|
||||
return [PresenceHourly(**dict(row)) for row in rows]
|
||||
@@ -0,0 +1,53 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import CallbackView, LinkView, ReactionView
|
||||
|
||||
|
||||
async def get_callbacks(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||
) -> list[CallbackView]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT position, label, data FROM callbacks "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
|
||||
"ORDER BY position",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
return [
|
||||
CallbackView(
|
||||
position=row["position"],
|
||||
label=row["label"],
|
||||
data=row["data"].hex() if row["data"] is not None else None,
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
async def get_reactions(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||
) -> list[ReactionView]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT peer_id, reaction, added_at, removed_at FROM reactions "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
|
||||
"ORDER BY added_at",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
return [ReactionView(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
async def get_links(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||
) -> list[LinkView]:
|
||||
rows = await pool.fetch(
|
||||
"SELECT position, url, kind, web_url, web_title, web_description, "
|
||||
"web_site_name FROM links "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
|
||||
"ORDER BY position",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
return [LinkView(**dict(row)) for row in rows]
|
||||
@@ -0,0 +1,113 @@
|
||||
import json
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import AlertView, Page, WatchView
|
||||
|
||||
_WATCH_COLS = "id, account_id, kind, params, enabled, created_at, updated_at"
|
||||
_ALERT_COLS = "id, account_id, watch_id, ts, payload, seen, created_at"
|
||||
|
||||
|
||||
def _to_watch(row: asyncpg.Record) -> WatchView:
|
||||
data = dict(row)
|
||||
data["params"] = json.loads(data["params"])
|
||||
return WatchView(**data)
|
||||
|
||||
|
||||
def _to_alert(row: asyncpg.Record) -> AlertView:
|
||||
data = dict(row)
|
||||
data["payload"] = json.loads(data["payload"])
|
||||
return AlertView(**data)
|
||||
|
||||
|
||||
async def list_watches(pool: asyncpg.Pool, account_id: int) -> list[WatchView]:
|
||||
rows = await pool.fetch(
|
||||
f"SELECT {_WATCH_COLS} FROM watches WHERE account_id = $1 " # noqa: S608
|
||||
"ORDER BY id DESC",
|
||||
account_id,
|
||||
)
|
||||
return [_to_watch(row) for row in rows]
|
||||
|
||||
|
||||
async def get_watch(pool: asyncpg.Pool, watch_id: int) -> WatchView | None:
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT {_WATCH_COLS} FROM watches WHERE id = $1", # noqa: S608
|
||||
watch_id,
|
||||
)
|
||||
return _to_watch(row) if row else None
|
||||
|
||||
|
||||
async def create_watch(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
kind: str,
|
||||
params: dict[str, Any],
|
||||
*,
|
||||
enabled: bool = True,
|
||||
) -> WatchView:
|
||||
row = await pool.fetchrow(
|
||||
"INSERT INTO watches (account_id, kind, params, enabled) " # noqa: S608
|
||||
f"VALUES ($1, $2, $3::jsonb, $4) RETURNING {_WATCH_COLS}",
|
||||
account_id,
|
||||
kind,
|
||||
json.dumps(params),
|
||||
enabled,
|
||||
)
|
||||
return _to_watch(row)
|
||||
|
||||
|
||||
async def update_watch(
|
||||
pool: asyncpg.Pool, watch_id: int, params: dict[str, Any], *, enabled: bool
|
||||
) -> WatchView | None:
|
||||
row = await pool.fetchrow(
|
||||
"UPDATE watches SET params = $2::jsonb, enabled = $3, updated_at = now() " # noqa: S608
|
||||
f"WHERE id = $1 RETURNING {_WATCH_COLS}",
|
||||
watch_id,
|
||||
json.dumps(params),
|
||||
enabled,
|
||||
)
|
||||
return _to_watch(row) if row else None
|
||||
|
||||
|
||||
async def delete_watch(pool: asyncpg.Pool, watch_id: int) -> bool:
|
||||
result = await pool.execute("DELETE FROM watches WHERE id = $1", watch_id)
|
||||
return result.endswith("1")
|
||||
|
||||
|
||||
async def list_alerts(
|
||||
pool: asyncpg.Pool, account_id: int, page: Page, *, seen: bool | None = None
|
||||
) -> list[AlertView]:
|
||||
params: list[object] = [account_id]
|
||||
where = "account_id = $1"
|
||||
if seen is not None:
|
||||
params.append(seen)
|
||||
where += f" AND seen = ${len(params)}"
|
||||
params.append(page.capped_limit)
|
||||
params.append(page.offset)
|
||||
rows = await pool.fetch(
|
||||
f"SELECT {_ALERT_COLS} FROM alerts WHERE {where} " # noqa: S608
|
||||
f"ORDER BY ts DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
|
||||
*params,
|
||||
)
|
||||
return [_to_alert(row) for row in rows]
|
||||
|
||||
|
||||
async def insert_alert(
|
||||
pool: asyncpg.Pool, account_id: int, watch_id: int, payload: dict[str, Any]
|
||||
) -> AlertView:
|
||||
row = await pool.fetchrow(
|
||||
"INSERT INTO alerts (account_id, watch_id, ts, payload) " # noqa: S608
|
||||
f"VALUES ($1, $2, $3, $4::jsonb) RETURNING {_ALERT_COLS}",
|
||||
account_id,
|
||||
watch_id,
|
||||
datetime.now(UTC),
|
||||
json.dumps(payload),
|
||||
)
|
||||
return _to_alert(row)
|
||||
|
||||
|
||||
async def mark_alert_seen(pool: asyncpg.Pool, alert_id: int) -> bool:
|
||||
result = await pool.execute("UPDATE alerts SET seen = true WHERE id = $1", alert_id)
|
||||
return result.endswith("1")
|
||||
Reference in New Issue
Block a user