feat: 领取 Admin 采集任务并保存本地 (#30)

client/src/http_admin_gateway.py

@@ -1,28 +1,31 @@
-"""使用 Python 标准库调用 Admin 登记接口。"""
+"""使用 Python 标准库调用 Admin 登记和任务领取接口。"""
 
 import json
 import socket
 from http.client import RemoteDisconnected
-from typing import Callable, Optional
+from typing import Callable, Mapping, Optional
 from urllib.error import HTTPError, URLError
 from urllib.parse import urlparse
 from urllib.request import ProxyHandler, Request, build_opener
 from uuid import uuid4
 
 from .admin_gateway import (
+    AdminTask,
     AdminGatewayError,
     ClaimCapabilities,
     ClientInfo,
     ClientRegistrationGateway,
     RegistrationReceipt,
+    TaskClaimGateway,
 )
+from .task_models import TaskType
 
 
 DEFAULT_ADMIN_BASE_URL = "http://127.0.0.1:8080"
 
 
-class HttpAdminGateway(ClientRegistrationGateway):
-    """通过 HTTP 登记 Client；访问令牌只保存在内存。"""
+class HttpAdminGateway(ClientRegistrationGateway, TaskClaimGateway):
+    """通过 HTTP 登记 Client 和领取采集任务；令牌只保存在内存。"""
 
     def __init__(
         self,
@@ -135,6 +138,151 @@ class HttpAdminGateway(ClientRegistrationGateway):
             registered_at=data["registered_at"],
         )
 
+    def claim_next(
+        self,
+        client: ClientInfo,
+        capabilities: ClaimCapabilities,
+    ) -> Optional[AdminTask]:
+        """领取一个采集任务；Admin 返回 204 时返回 ``None``。"""
+
+        request_id = str(uuid4())
+        payload = {
+            "client": {"name": client.name.strip()},
+            # #30 只允许领取采集任务。采购能力必须由安全门禁工单开启。
+            "supported_types": [TaskType.COLLECT.value],
+            "capabilities": {
+                "purchase_mode": "dry_run",
+                "schema_versions": list(capabilities.schema_versions),
+            },
+        }
+        if capabilities.device is not None:
+            payload["device"] = {
+                "address": capabilities.device.address.strip(),
+                "platform": capabilities.device.platform,
+                "pdd_package": capabilities.device.pdd_package.strip(),
+            }
+
+        headers = {
+            "Content-Type": "application/json; charset=utf-8",
+            "Accept": "application/json",
+            "X-Client-Id": client.client_id.strip(),
+            "X-Request-Id": request_id,
+        }
+        if self._token:
+            headers["Authorization"] = f"Bearer {self._token}"
+
+        request = Request(
+            f"{self._base_url}/api/v1/client/tasks/claim",
+            data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
+            headers=headers,
+            method="POST",
+        )
+
+        try:
+            with self._opener(request, timeout=self._timeout_seconds) as response:
+                status = getattr(response, "status", None)
+                if status is None:
+                    status = response.getcode()
+                body = response.read()
+        except HTTPError as exc:
+            self._raise_http_error(exc, request_id, "领取任务")
+        except (
+            URLError,
+            RemoteDisconnected,
+            ConnectionError,
+            socket.timeout,
+            TimeoutError,
+        ) as exc:
+            reason = getattr(exc, "reason", exc)
+            if isinstance(reason, (socket.timeout, TimeoutError)):
+                raise AdminGatewayError(
+                    "ADMIN_TIMEOUT",
+                    "Admin 领取请求超时，请稍后重试",
+                    True,
+                    request_id,
+                ) from exc
+            raise AdminGatewayError(
+                "ADMIN_UNAVAILABLE",
+                "无法连接 Admin，请检查服务地址",
+                True,
+                request_id,
+            ) from exc
+
+        if status == 204:
+            return None
+        if status != 200:
+            raise AdminGatewayError(
+                "ADMIN_UNEXPECTED_RESPONSE",
+                f"Admin 返回了未预期的状态码 {status}",
+                status >= 500,
+                request_id,
+            )
+        return self._parse_claim_response(body, request_id)
+
+    @classmethod
+    def _parse_claim_response(
+        cls,
+        body: bytes,
+        request_id: str,
+    ) -> AdminTask:
+        data = cls._decode_json(body, request_id)
+        task = data.get("task")
+        if not isinstance(task, Mapping):
+            raise AdminGatewayError(
+                "ADMIN_INVALID_RESPONSE",
+                "Admin 领取响应缺少 task 对象",
+                False,
+                request_id,
+            )
+
+        task_id = task.get("id")
+        raw_type = task.get("type")
+        version = task.get("version")
+        priority = task.get("priority")
+        task_payload = task.get("payload")
+        created_at = task.get("created_at")
+        updated_at = task.get("updated_at")
+        valid = (
+            isinstance(task_id, str)
+            and bool(task_id.strip())
+            and isinstance(raw_type, str)
+            and isinstance(version, int)
+            and not isinstance(version, bool)
+            and version > 0
+            and isinstance(priority, int)
+            and not isinstance(priority, bool)
+            and isinstance(task_payload, Mapping)
+            and isinstance(created_at, str)
+            and isinstance(updated_at, str)
+        )
+        if not valid:
+            raise AdminGatewayError(
+                "ADMIN_INVALID_RESPONSE",
+                "Admin 领取响应的任务字段不完整",
+                False,
+                request_id,
+            )
+
+        try:
+            task_type = TaskType(raw_type)
+        except ValueError as exc:
+            raise AdminGatewayError(
+                "ADMIN_INVALID_RESPONSE",
+                f"Admin 返回了不支持的任务类型 {raw_type}",
+                False,
+                request_id,
+            ) from exc
+
+        return AdminTask(
+            task_id=task_id,
+            task_type=task_type,
+            version=version,
+            priority=priority,
+            payload=dict(task_payload),
+            created_at=created_at,
+            updated_at=updated_at,
+        )
+
     @staticmethod
     def _decode_json(body: bytes, request_id: str) -> dict:
         try:
@@ -156,7 +304,12 @@ class HttpAdminGateway(ClientRegistrationGateway):
         return data
 
     @classmethod
-    def _raise_http_error(cls, error: HTTPError, request_id: str) -> None:
+    def _raise_http_error(
+        cls,
+        error: HTTPError,
+        request_id: str,
+        operation: str = "登记",
+    ) -> None:
         try:
             body = error.read()
         except OSError:
@@ -168,7 +321,7 @@ class HttpAdminGateway(ClientRegistrationGateway):
             if not isinstance(details, dict):
                 raise ValueError
             code = str(details.get("code") or "ADMIN_HTTP_ERROR")
-            message = str(details.get("message") or "Admin 拒绝了登记请求")
+            message = str(details.get("message") or f"Admin 拒绝了{operation}请求")
             retryable = bool(details.get("retryable", error.code >= 500))
             response_request_id = str(details.get("request_id") or request_id)
             extra = details.get("details")
@@ -176,7 +329,7 @@ class HttpAdminGateway(ClientRegistrationGateway):
                 extra = {}
         except (AdminGatewayError, ValueError):
             code = "ADMIN_HTTP_ERROR"
-            message = f"Admin 登记失败，状态码 {error.code}"
+            message = f"Admin {operation}失败，状态码 {error.code}"
             retryable = error.code >= 500
             response_request_id = request_id
             extra = {}