117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
import asyncio
|
|
import json
|
|
from typing import Annotated
|
|
|
|
import asyncpg
|
|
from dishka.integrations.fastapi import DishkaRoute, FromDishka
|
|
from fastapi import APIRouter, Query
|
|
from pydantic import BaseModel
|
|
|
|
from api.routers.policy import POLICY_CHANGED_CHANNEL
|
|
from utils.jobs import enqueue
|
|
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
|
|
|
|
router = APIRouter(prefix="/api", tags=["discover"], route_class=DishkaRoute)
|
|
|
|
DEFAULT_LIMIT = 30
|
|
REMOTE_TIMEOUT_SECONDS = 20.0
|
|
POLL_INTERVAL_SECONDS = 0.2
|
|
FINISHED = ("done", "failed", "canceled")
|
|
|
|
_CHAT_POLICY_ID = """
|
|
SELECT id FROM capture_policy
|
|
WHERE account_id = $1 AND scope_type = 'chat' AND scope_id = $2
|
|
"""
|
|
|
|
AccountId = Annotated[int, Query()]
|
|
|
|
|
|
class TrackRequest(BaseModel):
|
|
account_id: int
|
|
backfill: bool = True
|
|
|
|
|
|
class SyncContactsRequest(BaseModel):
|
|
account_id: int
|
|
|
|
|
|
async def _remote_ids(
|
|
pool: asyncpg.Pool, account_id: int, query: str, limit: int
|
|
) -> list[int]:
|
|
job_id = await enqueue(
|
|
pool, account_id, "search_peers", {"query": query, "limit": limit}
|
|
)
|
|
loop = asyncio.get_running_loop()
|
|
deadline = loop.time() + REMOTE_TIMEOUT_SECONDS
|
|
while loop.time() < deadline:
|
|
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
|
row = await pool.fetchrow(
|
|
"SELECT status, progress FROM jobs WHERE id = $1", job_id
|
|
)
|
|
if row is not None and row["status"] in FINISHED:
|
|
await pool.execute("DELETE FROM jobs WHERE id = $1", job_id)
|
|
return json.loads(row["progress"]).get("ids", [])
|
|
return []
|
|
|
|
|
|
@router.get("/discover")
|
|
async def discover_peers(
|
|
pool: FromDishka[asyncpg.Pool],
|
|
account_id: AccountId,
|
|
query: Annotated[str, Query()] = "",
|
|
remote: Annotated[bool, Query()] = False,
|
|
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
|
) -> list[DiscoverItem]:
|
|
if not query.strip():
|
|
return []
|
|
items = await discover.search(pool, account_id, query, limit)
|
|
if not remote:
|
|
return items
|
|
known = {item.chat_id for item in items}
|
|
ids = await _remote_ids(pool, account_id, query, limit)
|
|
extra = await discover.by_ids(
|
|
pool, account_id, [chat_id for chat_id in ids if chat_id not in known]
|
|
)
|
|
return [*items, *extra]
|
|
|
|
|
|
@router.get("/discover/{chat_id}")
|
|
async def discover_chat(
|
|
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
|
|
) -> DiscoverItem:
|
|
return await discover.get_item(pool, account_id, chat_id)
|
|
|
|
|
|
@router.post("/chats/{chat_id}/track", status_code=201)
|
|
async def track_chat(
|
|
pool: FromDishka[asyncpg.Pool], chat_id: int, body: TrackRequest
|
|
) -> DiscoverItem:
|
|
kind = await discover.chat_kind(pool, body.account_id, chat_id)
|
|
toggles = TRACKING[kind]
|
|
policy_id = await pool.fetchval(_CHAT_POLICY_ID, body.account_id, chat_id)
|
|
if policy_id is None:
|
|
await policy_repository.create_policy(
|
|
pool, body.account_id, ScopeType.CHAT, chat_id, toggles
|
|
)
|
|
else:
|
|
await policy_repository.update_policy(pool, policy_id, toggles)
|
|
await pool.execute(f"NOTIFY {POLICY_CHANGED_CHANNEL}")
|
|
await enqueue(pool, body.account_id, "enrich_chat", {"chat_id": chat_id})
|
|
if body.backfill:
|
|
await enqueue(
|
|
pool, body.account_id, "backfill", {"chat_id": chat_id, "media": True}
|
|
)
|
|
return await discover.get_item(pool, body.account_id, chat_id)
|
|
|
|
|
|
@router.post("/contacts/sync", status_code=201)
|
|
async def sync_contacts(
|
|
pool: FromDishka[asyncpg.Pool], body: SyncContactsRequest
|
|
) -> dict[str, int]:
|
|
job_id = await enqueue(pool, body.account_id, "sync_contacts", {})
|
|
return {"job_id": job_id}
|