dts-dance 0.1.8__py3-none-any.whl

dtsdance/dsyncer.py

@@ -0,0 +1,301 @@
+from .bytecloud import ByteCloudHelper
+from typing import Any, NamedTuple, Tuple, cast, Optional
+from loguru import logger
+import requests
+import re
+from datetime import datetime
+
+DSyncer_Task_Detail_URL = "{endpoint}/dsyncer/tasks/all?task_id_list={task_id}"
+DSyncer_Backlog_Dashboard_URL = "https://grafana.sretools.bytedance.net/d/b1cad70c-8ce3-490c-be48-fbcab32c7a2e/dsyncer-backlog?orgId=1&from=now-1h&to=now&timezone=browser&refresh=1m"
+DSyncer_N8N_OPT_URL = "https://n8n.sretools.bytedance.net/webhook/Xt73CZPpKjMu?opt={opt}&env={env}&task_id={task_id}"
+
+
+class DSyncerEnvInfo(NamedTuple):
+    name: str
+    token: str
+
+
+class DSyncerHelper:
+
+    def __init__(self, envs: dict[str, DSyncerEnvInfo], bytecloud_helper: ByteCloudHelper) -> None:
+        self.envs = envs
+        self.bytecloud_helper = bytecloud_helper
+
+    def _build_headers(self, env: str, secret_api: bool = False) -> dict[str, str]:
+        """
+        构建请求头
+
+        Args:
+            env: 环境名称
+            include_auth: 是否包含 Authorization 头
+
+        Returns:
+            dict[str, str]: 请求头字典
+        """
+        jwt_token = self.bytecloud_helper.get_jwt_token(env)
+        headers = {"X-Jwt-Token": jwt_token, "x-bcgw-vregion": env}
+
+        if secret_api:
+            token = self.envs[env].token
+            headers["Authorization"] = f"Token {token}"
+
+        return headers
+
+    def _make_request(self, method: str, url: str, headers: dict[str, str], json_data: Optional[dict] = None) -> dict[str, Any]:
+        """
+        发送 HTTP 请求的通用方法
+
+        Args:
+            method: HTTP 方法 (GET/POST)
+            url: 请求 URL
+            headers: 请求头
+            json_data: POST 请求的 JSON 数据
+            operation_name: 操作名称，用于日志记录
+
+        Returns:
+            dict[str, Any]: 解析后的 JSON 响应
+        """
+        response = None
+        try:
+            if method.upper() == "GET":
+                response = requests.get(url, headers=headers)
+            elif method.upper() == "POST":
+                response = requests.post(url, json=json_data, headers=headers)
+            else:
+                raise ValueError(f"不支持的 HTTP 方法: {method}")
+
+            # 检查响应状态码
+            # response.raise_for_status()
+
+            # 解析 JSON 响应
+            return response.json()
+
+        except Exception as e:
+            error_msg = f"_make_request occur error, error: {e}"
+            if response is not None:
+                error_msg += f", response.text: {response.text}"
+            logger.warning(error_msg)
+            raise
+
+    def _acquire_task_info(self, env: str, task_id: str) -> dict[str, str]:
+        """
+        获取 DSyncer 任务信息
+
+        Args:
+            env: 环境名称
+            task_id: DSyncer 任务 ID
+
+        Returns:
+            dict[str, str]: DSyncer 任务的 rocket_mq_connection 信息，只包含 cluster、topic 和 group 字段
+        """
+        # 构建 API URL
+        env_info = self.bytecloud_helper.get_env_info(env)
+        url = f"{env_info.endpoint}/api/v1/dsyncer/openapi/taskinfo/{task_id}/"
+
+        # 准备请求头
+        headers = self._build_headers(env)
+
+        return self._make_request("GET", url, headers)
+
+    def get_dflow_task_info(self, env: str, task_id: str) -> tuple[str, str]:
+        """
+        获取迁移后的 DFlow 任务信息
+        """
+        # 构建 API URL
+        env_info = self.bytecloud_helper.get_env_info(env)
+        url = f"{env_info.endpoint}/api/v1/dsyncer/openapi/taskinfo/{task_id}/migrate/"
+
+        # 准备请求头
+        headers = self._build_headers(env)
+
+        json_data = self._make_request("GET", url, headers)
+
+        message = json_data.get("message", "")
+        logger.debug(f"get task migrate info {env} {task_id}, message: {message}")
+
+        # 从消息中提取 DFlow 任务 URL 和 ID
+        # 消息格式: "任务已迁移至ByteDTS平台，请前往[https://cloud.bytedance.net/bytedts/datasync/detail/93127366537986?scope=China-North&tabKey=DetailInfo]查看详情"
+
+        # 提取完整的 URL
+        url_pattern = r"\[(https://[^\]]+)\]"
+        url_match = re.search(url_pattern, message)
+
+        # 提取任务 ID
+        id_pattern = r"/bytedts/datasync/detail/(\d+)"
+        id_match = re.search(id_pattern, message)
+
+        if url_match and id_match:
+            dflow_task_url = url_match.group(1)
+            dflow_task_id = id_match.group(1)
+            logger.info(f"extracted dflow task_url: {dflow_task_url}, task_id: {dflow_task_id}")
+            return (dflow_task_url, dflow_task_id)
+        else:
+            logger.warning(f"could not extract dflow task info from message: {message}")
+            return ("", "")
+
+    def generate_task_url(self, env: str, task_id: str) -> str:
+        """
+        获取 DSyncer 任务详情页面的 URL
+
+        Args:
+            env: 环境名称
+            task_id: DSyncer 任务 ID
+
+        Returns:
+            str: DSyncer 任务详情页面的 URL
+        """
+        env_info = self.bytecloud_helper.get_env_info(env)
+        return DSyncer_Task_Detail_URL.format(endpoint=env_info.endpoint, task_id=task_id)
+
+    def generate_task_grafana_url(self, task_id: str, change_time: str) -> str:
+        """
+        获取 DSyncer 任务的 Grafana 监控页面 URL
+        """
+        try:
+            # 将 change_time 字符串转换为毫秒时间戳
+            dt = datetime.strptime(change_time, "%Y-%m-%d %H:%M:%S")
+            change_time_timestamp = int(dt.timestamp() * 1000)
+        except ValueError as e:
+            logger.warning(f"无法解析时间格式 '{change_time}': {e}，使用当前时间戳")
+            change_time_timestamp = int(datetime.now().timestamp() * 1000)
+
+        return f"{DSyncer_Backlog_Dashboard_URL}&var-change_time={change_time_timestamp}&var-task_id={task_id}"
+
+    def get_task_info(self, env: str, task_id: str) -> dict[str, Any]:
+        """
+        获取 DSyncer 任务状态
+
+        Args:
+            env: 环境名称
+            task_id: DSyncer 任务 ID
+
+        Returns:
+            dict[str, str]: DSyncer 任务状态
+        """
+        json_data = self._acquire_task_info(env, task_id)
+        if json_data.get("code") == 400 and "message" in json_data:
+            raise Exception(json_data.get("message", "未知错误"))
+
+        logger.info(f"get_task_info success {env} {task_id}")
+
+        try:
+            data = cast(dict, json_data.get("data", {}))
+            increment_task = cast(dict, data.get("increment_task", {}))
+            rocket_mq_connection = cast(dict, increment_task.get("rocket_mq_connection", {}))
+            filtered_data = {
+                "task_id": data.get("task_id", ""),
+                "status": data.get("status", ""),
+                "desc": data.get("desc", ""),
+                "scene": data.get("scene", ""),
+                "psm": data.get("psm", ""),
+                "increment_task_id": increment_task.get("task_id", ""),
+                "mq_info": {
+                    "cluster": rocket_mq_connection.get("cluster", ""),
+                    "topic": rocket_mq_connection.get("topic", ""),
+                    "group": rocket_mq_connection.get("group", ""),
+                },
+            }
+
+            return filtered_data
+        except (KeyError, AttributeError, Exception) as e:
+            raise Exception(f"无法从响应中提取 DSyncer 任务状态数据: {str(e)}")
+
+    def is_task_migrate_running(self, env: str, task_id: str) -> Tuple[bool, str]:
+        """
+        检查 DSyncer 任务是否正在迁移中
+
+        Returns:
+            bool: 如果任务正在迁移中，返回 True；否则返回 False
+        """
+        # 构建 API URL
+        env_info = self.bytecloud_helper.get_env_info(env)
+        url = f"{env_info.endpoint}/api/v1/dsyncer/secret_api/task/migrate/check"
+
+        # 准备请求头
+        headers = self._build_headers(env, secret_api=True)
+
+        json_data = {"task_id": task_id}
+
+        response_data = self._make_request("POST", url, headers, json_data)
+
+        message = response_data.get("message", "")
+        logger.debug(f"get task migrate status {env} {task_id}, message: {message}")
+        if "task migrate is running" in message:
+            return True, ""
+        else:
+            return False, response_data.get("data", {}).get("msg", {})
+
+    def migration_rollback(self, env: str, task_id: str) -> bool:
+        """
+        执行回滚
+        """
+        # 构建 API URL
+        env_info = self.bytecloud_helper.get_env_info(env)
+        url = f"{env_info.endpoint}/api/v1/dsyncer/secret_api/task/rollback_migrate2dsyncer/"
+
+        # 准备请求头
+        headers = self._build_headers(env, secret_api=True)
+
+        json_data = {"task_id_list": [task_id]}
+
+        response_data = self._make_request("POST", url, headers, json_data)
+
+        logger.debug(f"migration_mark_rollback return {env} {task_id}, json_data: {response_data}")
+        success_task = response_data.get("data", {}).get("success_task", [])
+        return task_id in success_task
+
+    def migration_mark_success(self, env: str, task_id: str) -> bool:
+        """
+        标记迁移成功
+        """
+        # 构建 API URL
+        env_info = self.bytecloud_helper.get_env_info(env)
+        url = f"{env_info.endpoint}/api/v1/dsyncer/secret_api/task/mark_migrate_success/"
+
+        # 准备请求头
+        headers = self._build_headers(env, secret_api=True)
+
+        json_data = {"task_id": task_id}
+
+        response_data = self._make_request("POST", url, headers, json_data)
+
+        logger.debug(f"migration_mark_success return {env} {task_id}, json_data: {response_data}")
+        success_task = response_data.get("data", {}).get("success_task", [])
+        return task_id in success_task
+
+    def migrate_task(self, env: str, task_id: str, app_parallel: int) -> Optional[str]:
+        """
+        迁移任务到DFlow
+
+        Returns:
+            Optional[str]: 错误信息，成功时返回None，失败时返回错误信息
+        """
+        # 构建 API URL
+        env_info = self.bytecloud_helper.get_env_info(env)
+        url = f"{env_info.endpoint}/api/v1/dsyncer/secret_api/task/migrate2dflow/single"
+
+        # 准备请求头
+        headers = self._build_headers(env, secret_api=True)
+
+        json_data = {
+            "task_id": task_id,
+            "delay_threshold": "20s",
+            "dsyncer_task_pause_threshold": "180s",
+            "dflow_delay_threshold": "1m",
+            "dflow_diff_threshold": 100,
+            "check_minutes": 10,
+            "enable_psm": True,
+            "disable_check": True,
+            "app_parallel": app_parallel,
+        }
+
+        response_data = self._make_request("POST", url, headers, json_data)
+
+        logger.debug(f"migrate_task return {env} {task_id}, json_data: {response_data}")
+        message = response_data.get("message")
+        if response_data.get("code") == 0 and message == "ok":
+            err_message = None
+        else:
+            err_message = message
+
+        return err_message

dtsdance/feishu_base.py

@@ -0,0 +1,107 @@
+from datetime import datetime, timedelta
+import threading
+from typing import Optional
+from loguru import logger
+import requests
+
+URL_FEISHU_OPEN_API: str = "https://fsopen.bytedance.net/open-apis"
+
+
+class FeishuBase:
+    """飞书基础类，传入机器人认证信息，支持自动token续期管理"""
+
+    def __init__(self, bot_app_id: str, bot_app_secret: str):
+        """
+        初始化飞书基础类
+        """
+        self.tenant_access_token: Optional[str] = None
+        self.token_expire_time: Optional[datetime] = None
+        self._token_lock = threading.Lock()
+
+        self.bot_app_id = bot_app_id
+        self.bot_app_secret = bot_app_secret
+
+    def _get_tenant_access_token(self) -> str:
+        """
+        获取tenant_access_token
+
+        Returns:
+            tenant_access_token字符串
+
+        Raises:
+            Exception: 获取token失败时抛出异常
+        """
+        url = f"{URL_FEISHU_OPEN_API}/auth/v3/tenant_access_token/internal"
+        payload = {"app_id": self.bot_app_id, "app_secret": self.bot_app_secret}
+        headers = {"Content-Type": "application/json"}
+
+        try:
+            response = requests.post(url, json=payload, headers=headers)
+            response.raise_for_status()
+            result = response.json()
+            if result.get("code") != 0:
+                raise Exception(f"获取tenant_access_token失败: {result.get('msg')}")
+
+            token = result.get("tenant_access_token")
+            expire_in = result.get("expire", 7200)  # 默认2小时过期
+
+            # 设置过期时间，提前5分钟刷新
+            self.token_expire_time = datetime.now() + timedelta(seconds=expire_in - 300)
+
+            logger.info(f"成功获取tenant_access_token，过期时间: {self.token_expire_time}")
+            return token
+
+        except requests.RequestException as e:
+            logger.warning(f"请求tenant_access_token失败: {e}")
+            raise Exception(f"网络请求失败: {e}")
+        except Exception as e:
+            logger.warning(f"获取tenant_access_token失败: {e}")
+            raise
+
+    def get_access_token(self) -> str:
+        """
+        获取有效的access token，自动处理续期
+
+        Returns:
+            有效的tenant_access_token
+        """
+        with self._token_lock:
+            # 检查token是否存在或已过期
+            if self.tenant_access_token is None or self.token_expire_time is None or datetime.now() >= self.token_expire_time:
+
+                logger.info("Token不存在或已过期，重新获取")
+                self.tenant_access_token = self._get_tenant_access_token()
+
+            return self.tenant_access_token
+
+    def make_request(self, method: str, endpoint: str, **kwargs) -> requests.Response:
+        """
+        发起API请求的通用方法
+
+        Args:
+            method: HTTP方法
+            endpoint: API端点
+            **kwargs: 其他请求参数
+
+        Returns:
+            Response对象
+        """
+        url = f"{URL_FEISHU_OPEN_API}{endpoint}"
+
+        # 获取有效token
+        token = self.get_access_token()
+        # logger.debug(f"使用token: {token} 发送请求: {method} {url}")
+
+        # 设置认证头
+        headers = kwargs.get("headers", {"Content-Type": "application/json"})
+        headers["Authorization"] = f"Bearer {token}"
+        kwargs["headers"] = headers
+
+        try:
+            response = requests.request(method, url, **kwargs)
+            # logger.debug(f"response_json: {response.json()}")
+            response.raise_for_status()
+            return response
+        except requests.RequestException as e:
+            logger.warning(f"API请求失败: {method} {url}, error: {e}")
+            raise

dtsdance/feishu_table.py

@@ -0,0 +1,224 @@
+from datetime import datetime
+import threading
+from typing import Callable, Optional, Dict, Any
+from loguru import logger
+import requests
+
+from .feishu_base import FeishuBase
+
+URL_FEISHU_OPEN_API: str = "https://fsopen.bytedance.net/open-apis"
+
+
+class FeishuTable:
+    """飞书表格类，支持自动token续期管理"""
+
+    def __init__(self, feishu_base: FeishuBase, table_app_token: str, table_id: str):
+        """
+        初始化飞书表格
+        """
+        self.tenant_access_token: Optional[str] = None
+        self.token_expire_time: Optional[datetime] = None
+        self._token_lock = threading.Lock()
+
+        self.feishu_base = feishu_base
+        self.table_app_token = table_app_token
+        self.table_id = table_id
+
+    def get_app_table_record_id(self, table_view_id: str, task_id: str) -> Optional[str]:
+        """
+        根据任务ID查询飞书多维表格记录，返回记录ID
+        """
+        # 构建请求体
+        payload = {
+            "view_id": table_view_id,
+            "filter": {"conditions": [{"field_name": "任务id", "operator": "is", "value": [task_id]}], "conjunction": "and"},
+        }
+
+        try:
+            endpoint = f"/bitable/v1/apps/{self.table_app_token}/tables/{self.table_id}/records/search"
+            response = self.feishu_base.make_request("POST", endpoint, json=payload)
+            result = response.json()
+            if result.get("code") != 0:
+                logger.warning(f"get_app_table_record_id {task_id} 失败: {result}")
+                return None
+
+            # 提取记录数据
+            data = result.get("data", {})
+            items = data.get("items", [])
+            if not items:
+                logger.warning(f"get_app_table_record_id {task_id} 失败: items为空")
+                return None
+
+            # 返回第一个匹配记录的record_id
+            record_id = items[0].get("record_id")
+            logger.info(f"成功查询到任务ID {task_id} 对应的记录ID: {record_id}")
+
+            return record_id
+
+        except Exception as e:
+            logger.warning(f"查询 {task_id} 表格记录失败: {e}")
+            return None
+
+    def update_app_table_record(self, record_id: str, fields: Dict[str, Any]) -> bool:
+        """
+        更新飞书多维表格记录
+
+        Args:
+            record_id: 记录ID
+            fields: 要更新的字段，格式为 {"字段名": "字段值"}
+
+        Returns:
+            是否更新成功
+        """
+        # 构建请求体
+        payload = {"fields": fields}
+
+        try:
+            endpoint = f"/bitable/v1/apps/{self.table_app_token}/tables/{self.table_id}/records/{record_id}"
+            response = self.feishu_base.make_request("PUT", endpoint, json=payload)
+            result = response.json()
+            if result.get("code") != 0:
+                logger.warning(f"更新表格记录失败: {result.get('msg')}")
+                return False
+
+            logger.info(f"成功更新记录ID {record_id}，字段: {fields}")
+            return True
+
+        except Exception as e:
+            logger.warning(f"更新表格记录失败: {e}")
+            return False
+
+    def fetch_records(self, table_view_id: str, field_names: list[str], page_size: int = 100, page_token: Optional[str] = None) -> Dict[str, Any]:
+        """
+        获取表格记录
+
+        Args:
+            page_size: 每页记录数，最大500
+            page_token: 分页标记
+
+        Returns:
+            API响应数据
+        """
+        endpoint = f"/bitable/v1/apps/{self.table_app_token}/tables/{self.table_id}/records/search"
+        params = {"page_size": page_size, "user_id_type": "open_id"}
+        if page_token:
+            params["page_token"] = page_token
+
+        payload = {"field_names": field_names, "view_id": table_view_id}
+
+        try:
+            response = self.feishu_base.make_request("POST", endpoint, json=payload, params=params)
+            result = response.json()
+            if result.get("code") != 0:
+                raise Exception(f"API请求失败: {result.get('msg')}")
+
+            return result.get("data", {})
+
+        except requests.RequestException as e:
+            logger.error(f"请求失败: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"获取记录失败: {e}")
+            raise
+
+    def parse_record(self, record: Dict[str, Any], field_names: Optional[list[str]] = None) -> Dict[str, str]:
+        """
+        解析单条记录，动态提取指定字段
+
+        Args:
+            record: 飞书表格记录
+            field_names: 需要解析的字段名列表，如果为None则解析所有字段
+
+        Returns:
+            解析后的记录信息，包含record_id和指定的字段
+        """
+        fields = record.get("fields", {})
+        result = {"record_id": record.get("record_id", "")}
+
+        # 如果没有指定字段名，则解析所有字段
+        target_fields = field_names if field_names else list(fields.keys())
+
+        # 动态解析每个字段
+        for field_name in target_fields:
+            field_value = fields.get(field_name, [])
+
+            # 处理文本类型字段（列表格式）
+            if field_value and isinstance(field_value, list) and len(field_value) > 0:
+                # 如果是字典且包含text字段
+                if isinstance(field_value[0], dict) and "text" in field_value[0]:
+                    result[field_name] = field_value[0].get("text", "")
+                else:
+                    # 其他类型直接取第一个元素
+                    result[field_name] = str(field_value[0])
+            else:
+                # 空值或非列表类型
+                result[field_name] = field_value
+
+        return result
+
+    def loop_all(
+        self, table_view_id: str, field_names: list[str], callback: Optional[Callable[[Dict[str, str]], None]] = None, limit: Optional[int] = None
+    ):
+        """
+        遍历视图中的所有记录
+
+        Args:
+            callback: 可选的回调函数，对每条记录进行处理
+            limit: 可选的记录处理限制数量
+        """
+        logger.info("开始获取飞书多维表格数据...")
+
+        page_count = 0
+        total_count = 0
+        processed_count = 0
+        page_token = None
+
+        while True:
+            page_count += 1
+            logger.info(f"正在获取第 {page_count} 页数据...")
+
+            try:
+                data = self.fetch_records(table_view_id, field_names, page_size=200, page_token=page_token)
+                items = data.get("items", [])
+                size = len(items)
+                total_count += size
+                page_token = data.get("page_token")
+                logger.info(f"第 {page_count} 页获取到 {size} 条记录，next page_token: {page_token}")
+
+                # 处理每条记录
+                for item in items:
+                    logger.info(f"正在处理第 {processed_count + 1} 条数据...")
+                    # 检查是否达到处理限制
+                    if limit and processed_count >= limit:
+                        logger.info(f"已达到处理限制 {limit}，停止处理")
+                        break
+
+                    try:
+                        parsed_record = self.parse_record(item)
+                        processed_count += 1
+
+                        # 如果提供了回调函数，则调用它
+                        if callback:
+                            callback(parsed_record)
+                        else:
+                            # 默认行为：打印记录信息
+                            logger.info(f"任务ID: {parsed_record['task_id']}")
+
+                    except Exception as e:
+                        logger.error(f"处理记录失败: {e}, 记录: {item}")
+                        continue
+
+                # 如果达到处理限制，退出外层循环
+                if limit and processed_count >= limit:
+                    break
+
+                # 检查是否还有更多数据
+                if not data.get("has_more", False) or not page_token:
+                    break
+
+            except Exception as e:
+                logger.error(f"获取第 {page_count} 页数据失败: {e}")
+                break
+
+        logger.info(f"总共获取到 {total_count} 条记录，成功处理 {processed_count} 条")
+        logger.info("数据遍历完成")