DualFlow/scada_productor/base.py

from abc import ABC, abstractmethod
from typing import Optional, Dict, Union, final
import logging
logger = logging.getLogger("scada_productor_base")

from pathlib import Path
SCRIPT_DIR = Path(__file__).parent.absolute()

import os
import dotenv
dotenv.load_dotenv()

import requests


class ScadaProductorBase(ABC):
    def __init__(self):
        self._api_base_url: Optional[str] = os.getenv("DEFAULT_API_BASE_URL")
        self._login_user: Optional[str] = os.getenv("DEFAULT_LOGIN_USER")
        self._login_password: Optional[str] = os.getenv("DEFAULT_LOGIN_PASSWORD")
        self._login_dep_id: Optional[str] = os.getenv("DEFAULT_LOGIN_DEP_ID")
        self._scada_secret: Optional[str] = os.getenv("DEFAULT_SCADA_SECRET")

        self._token: Optional[str] = None
        self._headers: Dict[str, str] = {"Content-Type": "application/json"}

    # ======================== Token 管理 ========================

    @final
    def _update_token(self, max_retries: int = 3) -> Union[str, bool]:
        """
        向云平台请求登录，更新 JWT Token。

        Returns:
            str  — 成功时返回新 token
            False — 重试耗尽后仍失败
        """
        url = f"{self._api_base_url}/api/v2/user/login"
        payload = {
            "UserName": self._login_user,
            "Password": self._login_password,
            "type": "account",
            "DepId": self._login_dep_id,
        }

        for attempt in range(1, max_retries + 1):
            try:
                resp = requests.post(url, json=payload, timeout=5)
                resp.raise_for_status()
                data = resp.json()

                if data.get("code") != 200:
                    raise RuntimeError(f"登录失败: {data.get('msg', '未知错误')}")

                # 从响应中提取 token（兼容多种响应格式）
                resp_data = data.get("data")
                if isinstance(resp_data, dict):
                    new_token = resp_data.get("token") or resp_data.get("Token")
                elif isinstance(resp_data, str):
                    new_token = resp_data
                else:
                    new_token = data.get("token")

                if not new_token:
                    raise RuntimeError("登录响应中未找到 token")

                self._token = new_token
                self._headers["JWT-TOKEN"] = self._token
                logger.info(f"Token 更新成功 (第 {attempt} 次尝试)")
                return self._token

            except requests.Timeout:
                logger.warning(f"登录超时 (第 {attempt}/{max_retries} 次)", exc_info=True)
            except requests.RequestException as e:
                logger.warning(f"网络异常: {e} (第 {attempt}/{max_retries} 次)", exc_info=True)
            except (ValueError, KeyError) as e:
                logger.warning(f"响应解析异常: {e} (第 {attempt}/{max_retries} 次)", exc_info=True)
            except RuntimeError as e:
                logger.warning(f"{e} (第 {attempt}/{max_retries} 次)", exc_info=True)

        logger.error(f"Token 更新失败，已重试 {max_retries} 次")
        return False

    # ======================== 通用请求 ========================

    @final
    def _request(
        self,
        method: str,
        url: str,
        *,
        params: Optional[dict] = None,
        json: Optional[Union[dict, list]] = None,
        timeout: int = 30,
        max_retries: int = 3,
    ) -> Optional[dict]:
        """
        带自动 Token 刷新的通用请求方法。

        请求失败时（401 / 网络异常 / 响应异常）自动刷新 Token 并重试，
        最多重试 max_retries 次。

        Parameters:
            method:     "get" 或 "post"
            url:        完整请求 URL
            params:     GET 查询参数
            json:       POST 请求体
            timeout:    超时秒数
            max_retries: 最大重试次数

        Returns:
            dict   — 成功时返回解析后的 JSON 响应体
            None   — 重试耗尽仍失败
        """
        requester = getattr(requests, method.lower())
        kwargs = {"headers": self._headers, "timeout": timeout}
        if params is not None:
            kwargs["params"] = params
        if json is not None:
            kwargs["json"] = json

        for attempt in range(1, max_retries + 1):
            try:
                resp = requester(url, **kwargs)
                # Token 过期 → 刷新后重试（不消耗重试次数）
                if resp.status_code != 200:
                    if resp.status_code == 601: logger.warning(f"请求返回 601，尝试刷新 Token")
                    else : logger.warning(f"请求失败，尝试刷新 Token")
                    if self._update_token():
                        kwargs["headers"] = self._headers
                        # 重试，不要浪费一次循环
                        resp = requester(url, **kwargs)
                    else:
                        logger.error("Token 刷新失败，放弃重试")
                        return None

                resp.raise_for_status()
                if resp.status_code == 200:
                    data = resp.json()
                    return data

            except requests.Timeout:
                logger.warning(f"请求超时 {url} (第 {attempt}/{max_retries} 次)", exc_info=True)
            except requests.RequestException as e:
                logger.warning(f"请求异常 {url}: {e} (第 {attempt}/{max_retries} 次)", exc_info=True)
            except (ValueError, KeyError) as e:
                logger.warning(f"响应解析异常 {url}: {e} (第 {attempt}/{max_retries} 次)", exc_info=True)
            except Exception as e:
                logger.warning(f"请求异常 {url}: {e}", exc_info=True)

        logger.error(f"请求失败，已重试 {max_retries} 次: {method.upper()} {url}")
        return None