yuyr
1e5e91b193
Reviewed-on: #17 Reviewed-by: sundapeng <sundp@mail.zgclab.edu.cn> Reviewed-by: xuxt <xuxt@zgclab.edu.cn>
61 lines
2.4 KiB
Python
61 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any, Dict, Optional
|
|
|
|
import requests
|
|
|
|
from .log import get_logger
|
|
|
|
LOGGER = get_logger("argus.agent.client")
|
|
|
|
|
|
class MasterAPIError(Exception):
|
|
def __init__(self, message: str, status_code: int, payload: Optional[Dict[str, Any]] = None) -> None:
|
|
super().__init__(message)
|
|
self.status_code = status_code
|
|
self.payload = payload or {}
|
|
|
|
|
|
class AgentClient:
|
|
def __init__(self, base_url: str, *, timeout: int = 10) -> None:
|
|
self._base_url = base_url.rstrip("/")
|
|
self._timeout = timeout
|
|
self._session = requests.Session()
|
|
|
|
def register_node(self, body: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""调用 master 注册接口，返回节点对象。"""
|
|
url = f"{self._base_url}/api/v1/master/nodes"
|
|
response = self._session.post(url, json=body, timeout=self._timeout)
|
|
return self._parse_response(response, "Failed to register node")
|
|
|
|
def update_status(self, node_id: str, body: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""上报健康信息，由 master 更新 last_report。"""
|
|
url = f"{self._base_url}/api/v1/master/nodes/{node_id}/status"
|
|
response = self._session.put(url, json=body, timeout=self._timeout)
|
|
return self._parse_response(response, "Failed to update node status")
|
|
|
|
def _parse_response(self, response: requests.Response, error_prefix: str) -> Dict[str, Any]:
|
|
content_type = response.headers.get("Content-Type", "")
|
|
payload: Dict[str, Any] | None = None
|
|
if "application/json" in content_type:
|
|
try:
|
|
payload = response.json()
|
|
except json.JSONDecodeError:
|
|
LOGGER.warning("Response contained invalid JSON", extra={"status": response.status_code})
|
|
|
|
if response.status_code >= 400:
|
|
message = payload.get("error") if isinstance(payload, dict) else response.text
|
|
raise MasterAPIError(
|
|
f"{error_prefix}: {message}",
|
|
status_code=response.status_code,
|
|
payload=payload if isinstance(payload, dict) else None,
|
|
)
|
|
|
|
if payload is None:
|
|
try:
|
|
payload = response.json()
|
|
except json.JSONDecodeError as exc:
|
|
raise MasterAPIError("Master returned non-JSON payload", response.status_code) from exc
|
|
return payload
|