87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
"""Асинхронный клиент API для генератора: общий семафор, повторный вход при 401, повтор при сбоях сети."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
API = os.environ.get("PERF_API", "http://localhost:8000")
|
|
PREFIX = "/api/v1"
|
|
|
|
|
|
class ApiError(Exception):
|
|
def __init__(self, status: int, text: str, method: str, path: str):
|
|
super().__init__(f"{method} {path} -> {status}: {text[:400]}")
|
|
self.status = status
|
|
self.text = text
|
|
self.method = method
|
|
self.path = path
|
|
|
|
|
|
class Api:
|
|
def __init__(self, username: str, password: str, sem: asyncio.Semaphore, base: str = API):
|
|
self.username = username
|
|
self.password = password
|
|
self.sem = sem
|
|
self.base = base
|
|
self.client = httpx.AsyncClient(
|
|
base_url=base + PREFIX,
|
|
timeout=httpx.Timeout(600.0, connect=30.0),
|
|
limits=httpx.Limits(max_connections=64, max_keepalive_connections=64),
|
|
)
|
|
self.token: str | None = None
|
|
self._login_lock = asyncio.Lock()
|
|
self.requests = 0
|
|
|
|
async def aclose(self) -> None:
|
|
await self.client.aclose()
|
|
|
|
async def login(self, stale: str | None = None) -> str:
|
|
async with self._login_lock:
|
|
if self.token is not None and self.token != stale:
|
|
return self.token
|
|
r = await self.client.post("/auth/login", json={"username": self.username, "password": self.password})
|
|
if r.status_code != 200:
|
|
raise ApiError(r.status_code, r.text, "POST", "/auth/login")
|
|
self.token = r.json()["access_token"]
|
|
return self.token
|
|
|
|
async def req(self, method: str, path: str, *, ok_empty: bool = False, **kw: Any) -> Any:
|
|
if self.token is None:
|
|
await self.login()
|
|
last: Exception | None = None
|
|
for attempt in range(4):
|
|
try:
|
|
async with self.sem:
|
|
self.requests += 1
|
|
token = self.token
|
|
r = await self.client.request(method, path, headers={"Authorization": f"Bearer {token}"}, **kw)
|
|
if r.status_code == 401:
|
|
await self.login(stale=token)
|
|
continue
|
|
if r.status_code in (502, 503, 504):
|
|
await asyncio.sleep(1 + attempt)
|
|
last = ApiError(r.status_code, r.text, method, path)
|
|
continue
|
|
if r.status_code >= 400:
|
|
raise ApiError(r.status_code, r.text, method, path)
|
|
return r.json() if r.content else None
|
|
except (httpx.ConnectError, httpx.RemoteProtocolError) as e:
|
|
last = e
|
|
await asyncio.sleep(1 + attempt)
|
|
raise last if last else RuntimeError(f"{method} {path}: исчерпаны попытки")
|
|
|
|
async def get(self, path: str, **kw: Any) -> Any:
|
|
return await self.req("GET", path, **kw)
|
|
|
|
async def post(self, path: str, **kw: Any) -> Any:
|
|
return await self.req("POST", path, **kw)
|
|
|
|
async def patch(self, path: str, **kw: Any) -> Any:
|
|
return await self.req("PATCH", path, **kw)
|
|
|
|
async def put(self, path: str, **kw: Any) -> Any:
|
|
return await self.req("PUT", path, **kw)
|