gloria/feishu/client.py

#!/usr/bin/env python3
"""
飞书表格 API 客户端模块
统一版本，支持月度表格和年度表格
支持自动获取和刷新 tenant_access_token
"""

import os
import requests
import time
from typing import Dict, List, Optional, Tuple
import logging

# 加载环境变量
from dotenv import load_dotenv

load_dotenv()

# 配置日志
logger = logging.getLogger(__name__)


class FeishuClientError(Exception):
    """飞书客户端异常基类"""

    pass


class FeishuSheetsClient:
    """飞书表格 API 客户端"""

    def __init__(
        self,
        base_url: Optional[str] = None,
        token: Optional[str] = None,
        spreadsheet_token: Optional[str] = None,
        app_id: Optional[str] = None,
        app_secret: Optional[str] = None,
    ):
        """
        初始化客户端

        参数:
            base_url: 飞书 API 基础URL，如果为None则使用配置
            token: Bearer 认证令牌，如果为None则使用配置或自动获取
            spreadsheet_token: 表格 token，如果为None则使用配置
            app_id: 飞书应用ID，用于获取tenant_access_token
            app_secret: 飞书应用密钥，用于获取tenant_access_token
        """
        self.base_url = (base_url or os.getenv("FEISHU_BASE_URL")).rstrip("/")
        self.spreadsheet_token = spreadsheet_token or os.getenv(
            "FEISHU_SPREADSHEET_TOKEN"
        )
        self.app_id = app_id or os.getenv("FEISHU_APP_ID")
        self.app_secret = app_secret or os.getenv("FEISHU_APP_SECRET")

        # Token管理相关属性
        self._token = token or os.getenv("FEISHU_TOKEN")
        self._token_expire_time = 0  # token过期时间戳
        self._token_obtained_time = 0  # token获取时间戳

        # 使用 Session 重用连接
        self.session = requests.Session()
        self.session.timeout = int(os.getenv("REQUEST_TIMEOUT", 30))

        # 初始化headers
        self._update_session_headers()

        logger.debug(f"飞书客户端初始化完成，基础URL: {self.base_url}")
        logger.debug(
            f"使用应用ID: {self.app_id[:8]}... 如果配置"
            if self.app_id
            else "未配置应用ID"
        )

    def _update_session_headers(self):
        """更新session的headers"""
        if self._token:
            self.session.headers.update(
                {
                    "Authorization": f"Bearer {self._token}",
                    "Content-Type": "application/json",
                    "Accept": "application/json",
                }
            )
        else:
            # 如果没有token，移除Authorization头
            if "Authorization" in self.session.headers:
                del self.session.headers["Authorization"]

    def _get_tenant_access_token(self) -> Tuple[str, int]:
        """
        获取tenant_access_token

        返回:
            (token, expire_time): token字符串和过期时间（秒）

        异常:
            requests.exceptions.RequestException: 网络请求失败
            ValueError: API返回错误
        """
        if not self.app_id or not self.app_secret:
            raise ValueError("未配置飞书应用ID和密钥，无法获取tenant_access_token")

        token_url = (
            "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal"
        )

        payload = {"app_id": self.app_id, "app_secret": self.app_secret}

        headers = {"Content-Type": "application/json; charset=utf-8"}

        try:
            logger.info(f"正在获取tenant_access_token，应用ID: {self.app_id[:8]}...")
            response = requests.post(
                token_url, json=payload, headers=headers, timeout=int(os.getenv("REQUEST_TIMEOUT", 30))
            )
            response.raise_for_status()
            data = response.json()

            if data.get("code") != 0:
                error_msg = f"获取tenant_access_token失败: {data.get('msg')}"
                logger.error(error_msg)
                raise ValueError(error_msg)

            token = data.get("tenant_access_token")
            expire = data.get("expire", 7200)  # 默认2小时

            if not token:
                raise ValueError("API返回的token为空")

            logger.info(f"成功获取tenant_access_token，有效期: {expire}秒")
            return token, expire

        except requests.exceptions.RequestException as e:
            logger.error(f"获取tenant_access_token网络请求失败: {e}")
            raise
        except Exception as e:
            logger.error(f"获取tenant_access_token失败: {e}")
            raise

    def _ensure_valid_token(self):
        """
        确保当前token有效，如果无效则重新获取

        返回:
            bool: token是否有效
        """
        current_time = time.time()

        # 如果有手动配置的token，直接使用
        if os.getenv("FEISHU_TOKEN") and self._token == os.getenv("FEISHU_TOKEN"):
            logger.debug("使用手动配置的FEISHU_TOKEN")
            return True

        # 检查token是否过期（提前30分钟刷新）
        if self._token and self._token_expire_time > 0:
            time_remaining = self._token_expire_time - current_time
            if time_remaining > 1800:  # 剩余时间大于30分钟
                logger.debug(f"token仍然有效，剩余时间: {int(time_remaining)}秒")
                return True
            elif time_remaining > 0:  # 剩余时间小于30分钟但大于0
                logger.info(
                    f"token即将过期，剩余时间: {int(time_remaining)}秒，重新获取"
                )
            else:  # 已过期
                logger.info("token已过期，重新获取")

        # 需要获取新token
        try:
            token, expire = self._get_tenant_access_token()
            self._token = token
            self._token_obtained_time = current_time
            self._token_expire_time = current_time + expire
            self._update_session_headers()
            logger.info(f"token获取成功，将在 {expire} 秒后过期")
            return True
        except Exception as e:
            logger.error(f"获取token失败: {e}")
            # 如果配置了备用token，尝试使用
            if os.getenv("FEISHU_TOKEN") and os.getenv("FEISHU_TOKEN") != self._token:
                logger.warning("使用备用FEISHU_TOKEN")
                self._token = os.getenv("FEISHU_TOKEN")
                self._update_session_headers()
                return True
            return False

    def get_sheets_info(self) -> List[Dict[str, str]]:
        """
        获取所有表格信息（sheet_id 和 title）

        返回:
            表格信息列表 [{'sheet_id': 'xxx', 'title': 'xxx'}, ...]

        异常:
            requests.exceptions.RequestException: 网络请求失败
            ValueError: API返回错误
        """
        # 确保token有效
        if not self._ensure_valid_token():
            raise FeishuClientError("无法获取有效的飞书token")

        url = f"{self.base_url}/spreadsheets/{self.spreadsheet_token}/sheets/query"

        try:
            response = self.session.get(url, timeout=int(os.getenv("REQUEST_TIMEOUT", 30)))
            response.raise_for_status()
            data = response.json()

            if data.get("code") != 0:
                error_msg = f"飞书API错误: {data.get('msg')}"
                logger.error(error_msg)
                raise ValueError(error_msg)

            sheets = data.get("data", {}).get("sheets", [])
            result = []
            for sheet in sheets:
                result.append(
                    {"sheet_id": sheet.get("sheet_id"), "title": sheet.get("title")}
                )

            logger.info(f"获取到 {len(result)} 个表格")
            return result

        except requests.exceptions.RequestException as e:
            logger.error(f"获取表格信息失败: {e}")
            raise
        except Exception as e:
            logger.error(f"解析表格信息失败: {e}")
            raise

    def get_sheet_data(self, sheet_id: str, range_: Optional[str] = None) -> Dict:
        """
        获取指定表格的数据

        参数:
            sheet_id: 表格ID
            range_: 数据范围，如果为None则使用配置

        返回:
            飞书API返回的原始数据，包含revision版本号

        异常:
            requests.exceptions.RequestException: 网络请求失败
            ValueError: API返回错误
        """
        # 确保token有效
        if not self._ensure_valid_token():
            raise FeishuClientError("无法获取有效的飞书token")

        if range_ is None:
            range_ = os.getenv("SHEET_RANGE", "A1:Z100")

        # 注意：获取表格数据使用 v2 API，而不是 v3
        url = f"{self.base_url.replace('/v3', '/v2')}/spreadsheets/{self.spreadsheet_token}/values/{sheet_id}!{range_}"
        params = {
            "valueRenderOption": "ToString",
            "dateTimeRenderOption": "FormattedString",
        }

        try:
            response = self.session.get(
                url, params=params, timeout=int(os.getenv("REQUEST_TIMEOUT", 30))
            )
            response.raise_for_status()
            data = response.json()

            if data.get("code") != 0:
                error_msg = f"飞书API错误: {data.get('msg')}"
                logger.error(error_msg)
                raise ValueError(error_msg)

            logger.debug(f"获取表格数据成功: {sheet_id}, 范围: {range_}")
            return data.get("data", {})

        except requests.exceptions.RequestException as e:
            logger.error(f"获取表格数据失败: {e}, sheet_id: {sheet_id}")
            raise
        except Exception as e:
            logger.error(f"解析表格数据失败: {e}, sheet_id: {sheet_id}")
            raise

    def get_token_info(self) -> Dict[str, any]:
        """
        获取当前token信息

        返回:
            token信息字典
        """
        current_time = time.time()
        time_remaining = (
            self._token_expire_time - current_time if self._token_expire_time > 0 else 0
        )

        return {
            "has_token": bool(self._token),
            "token_preview": self._token[:20] + "..."
            if self._token and len(self._token) > 20
            else self._token,
            "token_obtained_time": self._token_obtained_time,
            "token_expire_time": self._token_expire_time,
            "time_remaining": max(0, time_remaining),
            "using_app_credentials": bool(self.app_id and self.app_secret),
            "using_manual_token": bool(
                os.getenv("FEISHU_TOKEN") and self._token == os.getenv("FEISHU_TOKEN")
            ),
        }

    def test_connection(self) -> bool:
        """
        测试飞书连接是否正常

        返回:
            连接是否正常
        """
        try:
            # 首先测试token获取
            if not self._ensure_valid_token():
                logger.error("无法获取有效的飞书token")
                return False

            # 然后测试表格访问
            sheets = self.get_sheets_info()
            if sheets:
                logger.info(f"飞书连接测试成功，找到 {len(sheets)} 个表格")
                return True
            else:
                logger.warning("飞书连接测试成功，但未找到表格")
                return False
        except Exception as e:
            logger.error(f"飞书连接测试失败: {e}")
            return False

    def refresh_token(self) -> bool:
        """
        强制刷新token

        返回:
            刷新是否成功
        """
        try:
            logger.info("强制刷新token...")
            current_time = time.time()
            token, expire = self._get_tenant_access_token()
            self._token = token
            self._token_obtained_time = current_time
            self._token_expire_time = current_time + expire
            self._update_session_headers()
            logger.info(f"token刷新成功，将在 {expire} 秒后过期")
            return True
        except Exception as e:
            logger.error(f"强制刷新token失败: {e}")
            return False


if __name__ == "__main__":
    # 测试代码
    import sys

    # 设置日志级别
    logging.basicConfig(level=logging.INFO)

    # 测试连接
    client = FeishuSheetsClient()

    # 显示token信息
    token_info = client.get_token_info()
    print("当前token信息:")
    print(f"  是否有token: {token_info['has_token']}")
    print(f"  token预览: {token_info['token_preview']}")
    print(f"  剩余时间: {int(token_info['time_remaining'])}秒")
    print(f"  使用应用凭证: {token_info['using_app_credentials']}")
    print(f"  使用手动token: {token_info['using_manual_token']}")

    if client.test_connection():
        print("\n飞书连接测试成功")

        # 获取表格信息
        sheets = client.get_sheets_info()
        for sheet in sheets[:3]:  # 只显示前3个
            print(f"表格: {sheet['title']} (ID: {sheet['sheet_id']})")

        if sheets:
            # 获取第一个表格的数据
            sheet_id = sheets[0]["sheet_id"]
            data = client.get_sheet_data(sheet_id, "A1:C5")
            print(f"获取到表格数据，版本: {data.get('revision', '未知')}")

        # 再次显示token信息
        token_info = client.get_token_info()
        print(f"\n测试后token剩余时间: {int(token_info['time_remaining'])}秒")
    else:
        print("\n飞书连接测试失败")
        sys.exit(1)