98 lines
3.3 KiB
Python
98 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Callable
|
|
|
|
|
|
def legacy_headers(*, bridge_auth_user: str, bridge_auth_role: str) -> dict[str, str]:
|
|
return {"X-User": bridge_auth_user, "X-Role": bridge_auth_role}
|
|
|
|
|
|
def bearer_headers(*, issue_token: Callable[[], str]) -> dict[str, str]:
|
|
return {"Authorization": f"Bearer {issue_token()}"}
|
|
|
|
|
|
def request_with_bridge_auth(
|
|
*,
|
|
client: Any,
|
|
httpx_module: Any,
|
|
method: str,
|
|
url: str,
|
|
bridge_auth_mode: str,
|
|
bridge_auth_fallback_legacy: bool,
|
|
legacy_headers_payload: dict[str, str],
|
|
bearer_headers_payload: dict[str, str],
|
|
retry_reset: Callable[[], None] | None = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
mode = bridge_auth_mode
|
|
|
|
if mode == "legacy_headers":
|
|
response = client.request(method, url, headers=legacy_headers_payload, **kwargs)
|
|
response.raise_for_status()
|
|
return response
|
|
|
|
if mode == "bearer":
|
|
response = client.request(method, url, headers=bearer_headers_payload, **kwargs)
|
|
response.raise_for_status()
|
|
return response
|
|
|
|
response = client.request(method, url, headers=bearer_headers_payload, **kwargs)
|
|
if (
|
|
mode == "bearer_first"
|
|
and bridge_auth_fallback_legacy
|
|
and response.status_code in {401, 403}
|
|
):
|
|
if retry_reset:
|
|
retry_reset()
|
|
response = client.request(method, url, headers=legacy_headers_payload, **kwargs)
|
|
response.raise_for_status()
|
|
return response
|
|
|
|
|
|
def request_json(
|
|
*,
|
|
httpx_module: Any,
|
|
send_request: Callable[..., Any],
|
|
method: str,
|
|
url: str,
|
|
payload: dict[str, Any],
|
|
timeout_seconds: float | None,
|
|
max_attempts: int | None,
|
|
retry_backoff_seconds: float | None,
|
|
forward_timeout_seconds: Callable[[float], float],
|
|
forward_max_attempts: Callable[[], int],
|
|
forward_retry_backoff_seconds: Callable[[], float],
|
|
sleep: Callable[[float], None],
|
|
) -> dict[str, Any]:
|
|
last_exc: Exception | None = None
|
|
resolved_timeout = timeout_seconds if timeout_seconds is not None else forward_timeout_seconds(45.0)
|
|
resolved_attempts = max(max_attempts if max_attempts is not None else forward_max_attempts(), 1)
|
|
resolved_backoff = max(
|
|
retry_backoff_seconds if retry_backoff_seconds is not None else forward_retry_backoff_seconds(),
|
|
0.0,
|
|
)
|
|
with httpx_module.Client(timeout=resolved_timeout) as client:
|
|
for attempt in range(1, resolved_attempts + 1):
|
|
try:
|
|
response = send_request(
|
|
client,
|
|
method=method,
|
|
url=url,
|
|
json=payload,
|
|
)
|
|
return response.json()
|
|
except httpx_module.HTTPStatusError as exc:
|
|
last_exc = exc
|
|
status = exc.response.status_code if exc.response is not None else 0
|
|
if status < 500 or attempt >= resolved_attempts:
|
|
raise
|
|
except httpx_module.RequestError as exc:
|
|
last_exc = exc
|
|
if attempt >= resolved_attempts:
|
|
raise
|
|
if resolved_backoff > 0:
|
|
sleep(resolved_backoff * attempt)
|
|
if last_exc is not None:
|
|
raise last_exc
|
|
raise RuntimeError(f"Failed to {method} JSON payload")
|