dataops-sdk 0.3.0__py3-none-any.whl

dataops_sdk/runner/schemas/nacos/naming_param.py

@@ -0,0 +1,62 @@
+from collections.abc import Callable
+from typing import Annotated
+
+from pydantic import BaseModel, Field
+
+from .constants import Constants
+
+
+class RegisterInstanceParam(BaseModel):
+    ip: Annotated[str, Field(description="实例IP地址")]
+    port: Annotated[int, Field(description="实例端口号")]
+    weight: Annotated[float, Field(description="实例权重")] = 1.0
+    enabled: Annotated[bool, Field(description="实例是否启用")] = True
+    healthy: Annotated[bool, Field(description="实例健康状态")] = True
+    metadata: Annotated[dict[str, str], Field(description="实例元数据")] = {}
+    cluster_name: Annotated[str, Field(description="实例所在集群名称")] = ""
+    service_name: Annotated[str, Field(description="实例所在服务名称")]
+    group_name: Annotated[str, Field(description="实例所在分组")] = Constants.DEFAULT_GROUP
+    ephemeral: Annotated[bool, Field(description="实例是否临时实例")] = True
+
+
+class BatchRegisterInstanceParam(BaseModel):
+    service_name: Annotated[str, Field(description="服务名称")]
+    group_name: Annotated[str, Field(description="分组名称")] = Constants.DEFAULT_GROUP
+    instances: Annotated[list[RegisterInstanceParam], Field(description="实例列表")] = []
+
+
+class DeregisterInstanceParam(BaseModel):
+    ip: Annotated[str, Field(description="实例IP地址")]
+    port: Annotated[int, Field(description="实例端口号")]
+    cluster_name: Annotated[str, Field(description="实例所在集群名称")] = ""
+    service_name: Annotated[str, Field(description="实例所在服务名称")]
+    group_name: Annotated[str, Field(description="实例所在分组")] = Constants.DEFAULT_GROUP
+    ephemeral: Annotated[bool, Field(description="实例是否临时实例")] = True
+
+
+class ListInstanceParam(BaseModel):
+    service_name: Annotated[str, Field(description="服务名称")]
+    group_name: Annotated[str, Field(description="分组名称")] = Constants.DEFAULT_GROUP
+    clusters: Annotated[list[str], Field(description="集群列表")] = []
+    subscribe: Annotated[bool, Field(description="是否订阅")] = True
+    healthy_only: Annotated[bool | None, Field(description="仅返回健康实例")] = None
+
+
+class SubscribeServiceParam(BaseModel):
+    service_name: Annotated[str, Field(description="服务名称")]
+    group_name: Annotated[str, Field(description="分组名称")] = Constants.DEFAULT_GROUP
+    clusters: Annotated[list[str], Field(description="集群列表")] = []
+    subscribe_callback: Annotated[Callable | None, Field(description="订阅回调函数")] = None
+
+
+class GetServiceParam(BaseModel):
+    service_name: Annotated[str, Field(description="服务名称")]
+    group_name: Annotated[str, Field(description="分组名称")] = Constants.DEFAULT_GROUP
+    clusters: Annotated[list[str], Field(description="集群列表")] = []
+
+
+class ListServiceParam(BaseModel):
+    namespace_id: Annotated[str, Field(description="命名空间ID")] = Constants.DEFAULT_NAMESPACE_ID
+    group_name: Annotated[str, Field(description="分组名称")] = Constants.DEFAULT_GROUP
+    page_no: Annotated[int, Field(description="页码")] = 1
+    page_size: Annotated[int, Field(description="每页数量")] = 10

dataops_sdk/runner/schemas/nacos/naming_request.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from abc import ABC
+from typing import Annotated
+
+from pydantic import Field
+from v2.nacos.naming.model.instance import Instance
+from v2.nacos.transport.model.rpc_request import Request
+
+from .service import Service
+
+
+class AbstractNamingRequest(Request, ABC):
+    namespace: Annotated[str | None, Field(description="命名空间")] = ""
+    serviceName: Annotated[str | None, Field(description="服务名称")] = ""
+    groupName: Annotated[str | None, Field(description="分组名称")] = ""
+
+    def get_module(self):
+        return "naming"
+
+    def get_request_type(self) -> str:
+        """
+        提供一个默认实现或抛出NotImplementedError，明确指示子类需要覆盖此方法。
+        """
+        raise NotImplementedError("Subclasses should implement this method.")
+
+
+NOTIFY_SUBSCRIBER_REQUEST_TYPE = "NotifySubscriberRequest"
+
+
+class ServiceQueryRequest(AbstractNamingRequest):
+    cluster: Annotated[str | None, Field(description="集群名称")]
+    healthOnly: Annotated[bool | None, Field(description="仅健康实例")]
+    udpPort: Annotated[int | None, Field(description="UDP端口")] = None
+
+    def get_request_type(self) -> str:
+        return "ServiceQueryRequest"
+
+
+class InstanceRequest(AbstractNamingRequest):
+    type: Annotated[str | None, Field(description="请求类型")]
+    instance: Annotated[Instance | None, Field(description="实例")]
+
+    def get_request_type(self) -> str:
+        return "InstanceRequest"
+
+
+class PersistentInstanceRequest(AbstractNamingRequest):
+    type: Annotated[str | None, Field(description="请求类型")]
+    instance: Annotated[Instance | None, Field(description="实例")]
+
+    def get_request_type(self) -> str:
+        return "PersistentInstanceRequest"
+
+
+class BatchInstanceRequest(AbstractNamingRequest):
+    type: Annotated[str | None, Field(description="请求类型")]
+    instances: Annotated[list[Instance] | None, Field(description="实例列表")]
+
+    def get_request_type(self) -> str:
+        return "BatchInstanceRequest"
+
+
+class NotifySubscriberRequest(AbstractNamingRequest):
+    serviceInfo: Annotated[Service | None, Field(description="服务信息")]
+
+    def get_request_type(self) -> str:
+        return "NotifySubscriberRequest"
+
+
+class ServiceListRequest(AbstractNamingRequest):
+    pageNo: Annotated[int | None, Field(description="页码")]
+    pageSize: Annotated[int | None, Field(description="每页数量")]
+
+    def get_request_type(self) -> str:
+        return "ServiceListRequest"
+
+
+class SubscribeServiceRequest(AbstractNamingRequest):
+    subscribe: Annotated[bool | None, Field(description="是否订阅")]
+    clusters: Annotated[str | None, Field(description="集群名称")]
+
+    def get_request_type(self) -> str:
+        return "SubscribeServiceRequest"

dataops_sdk/runner/schemas/nacos/naming_response.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from pydantic import Field
+from v2.nacos.naming.model.service import Service
+from v2.nacos.transport.model.rpc_response import Response
+
+
+class NotifySubscriberResponse(Response):
+    def get_response_type(self) -> str:
+        return "NotifySubscriberResponse"
+
+
+class QueryServiceResponse(Response):
+    serviceInfo: Annotated[Service | None, Field(description="服务信息")] = None
+
+    def get_service_info(self) -> Service:
+        return self.serviceInfo
+
+    def get_response_type(self) -> str:
+        return "QueryServiceResponse"
+
+
+class SubscribeServiceResponse(Response):
+    serviceInfo: Annotated[Service | None, Field(description="服务信息")] = None
+
+    def get_response_type(self) -> str:
+        return "SubscribeServiceResponse"
+
+    def get_service_info(self) -> Service:
+        return self.serviceInfo
+
+
+class InstanceResponse(Response):
+    def get_response_type(self) -> str:
+        return "InstanceResponse"
+
+
+class BatchInstanceResponse(Response):
+    def get_response_type(self) -> str:
+        return "BatchInstanceResponse"
+
+
+class ServiceListResponse(Response):
+    count: Annotated[int, Field(description="服务数量")]
+    serviceNames: Annotated[list[str], Field(description="服务名称列表")]
+
+    def get_response_type(self) -> str:
+        return "ServiceListResponse"

dataops_sdk/runner/schemas/nacos/preserved_metadata_key.py

@@ -0,0 +1,10 @@
+class PreservedMetadataKeys:
+    REGISTER_SOURCE = "preserved.register.source"
+
+    HEART_BEAT_TIMEOUT = "preserved.heart.beat.timeout"
+
+    IP_DELETE_TIMEOUT = "preserved.ip.delete.timeout"
+
+    HEART_BEAT_INTERVAL = "preserved.heart.beat.interval"
+
+    INSTANCE_ID_GENERATOR = "preserved.instance.id.generator"

dataops_sdk/runner/schemas/nacos/service.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import time
+import urllib.parse
+from typing import Annotated
+
+from pydantic import BaseModel, Field
+
+from .constants import Constants
+from .instance import Instance
+from .nacos_exception import NacosException
+
+EMPTY = ""
+
+ALL_IPS = "000--00-ALL_IPS--00--000"
+
+SPLITER = "@@"
+
+DEFAULT_CHARSET = "UTF-8"
+
+
+class Service(BaseModel):
+    name: Annotated[str, Field(description="服务名称")]
+    groupName: Annotated[str, Field(description="分组名称")]
+    clusters: Annotated[str | None, Field(description="集群名称")] = ""
+    cacheMillis: Annotated[int, Field(description="缓存时间(毫秒)")] = 1000
+    hosts: Annotated[list[Instance], Field(description="实例列表")] = []
+    lastRefTime: Annotated[int, Field(description="最后引用时间")] = 0
+    checksum: Annotated[str, Field(description="校验和")] = ""
+    allIps: Annotated[bool, Field(description="是否所有IP")] = False
+    reachProtectionThreshold: Annotated[bool, Field(description="是否达到保护阈值")] = False
+    jsonFromServer: Annotated[str, Field(description="服务器返回的JSON")] = ""
+
+    def init_from_key(self, key=None):
+        if key:
+            max_index = 2
+            cluster_index = 2
+            service_name_index = 1
+            group_index = 0
+            keys = key.split(Constants.SERVICE_INFO_SPLITER)
+            if len(keys) >= max_index + 1:
+                self.groupName = keys[group_index]
+                self.name = keys[service_name_index]
+                self.clusters = keys[cluster_index]
+            elif len(keys) == max_index:
+                self.groupName = keys[group_index]
+                self.name = keys[service_name_index]
+            else:
+                raise NacosException("Can't parse out 'group_name', but it must not None!")
+
+    def get_ip_count(self):
+        return len(self.hosts)
+
+    def is_expired(self):
+        return int(round(time.time() * 1000)) - self.lastRefTime > self.cacheMillis
+
+    def add_host(self, host):
+        self.hosts.append(host)
+
+    def add_all_hosts(self, hosts):
+        self.hosts.extend(hosts)
+
+    def is_valid(self):
+        return self.hosts != []
+
+    def validate(self):
+        if self.allIps:
+            return True
+
+        if not self.hosts:
+            return False
+
+        valid_hosts = []
+        for host in self.hosts:
+            if not host.is_healthy():
+                continue
+
+            for i in range(host.get_weight()):
+                valid_hosts.append(i)
+
+        return len(valid_hosts) > 0
+
+    def get_key_default(self):
+        service_name = self.get_grouped_service_name()
+        return self.get_key(service_name, self.clusters)
+
+    def get_key_encoded(self):
+        service_name = self.get_grouped_service_name().encode("utf-8")
+        service_name = urllib.parse.quote(service_name)
+        return self.get_key(service_name, self.clusters)
+
+    def get_grouped_service_name(self):
+        service_name = self.name
+        if self.groupName and Constants.SERVICE_INFO_SPLITER not in service_name:
+            service_name = self.groupName + Constants.SERVICE_INFO_SPLITER + service_name
+        return service_name
+
+    @staticmethod
+    def from_key(key: str):
+        info = key.split(Constants.SERVICE_INFO_SPLITER)
+        if len(info) > 2:
+            service = Service(name=info[1], groupName=info[0], clusters=info[2])
+        else:
+            service = Service(name=info[1], groupName=info[0], clusters="")
+        return service
+
+    def get_hosts_str(self):
+        hosts_str = ""
+        for host in self.hosts:
+            hosts_str += host.json() + ";"
+        return hosts_str
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class ServiceList(BaseModel):
+    count: Annotated[int, Field(description="服务数量")]
+    services: Annotated[list[str], Field(description="服务名称列表")]

dataops_sdk/runner/utils/__init__.py

@@ -0,0 +1,42 @@
+"""Runner 工具模块。
+
+提供数据库连接、数据转换、日志等工具函数。
+
+Functions:
+    encrypt_body: 加密请求体
+    decrypt_body: 解密请求体
+    generate_uuid: 生成 UUID
+    print_json: 打印 JSON 格式数据
+
+Classes:
+    DatabaseConn: 数据库连接管理类
+    DuckDBManager: DuckDB 连接管理类
+    Neo4jConn: Neo4j 图数据库连接类
+
+Example:
+    >>> from runner.utils import DatabaseConn, encrypt_body
+"""
+
+from .common import decrypt_body, encrypt_body, generate_uuid, parse_args, print_json
+from .dbConn_sync import DatabaseConn
+from .ftp_file_transfer import FTPFileTransfer
+from .log import getLogger
+from .nacos import get_service_instance
+from .neo4j_sync import Neo4jConn
+from .queue_kafka import KafkaQueue
+from .sftp_file_transfer import SFTPFileTransfer
+
+__all__ = [
+    "decrypt_body",
+    "encrypt_body",
+    "generate_uuid",
+    "print_json",
+    "DatabaseConn",
+    "getLogger",
+    "Neo4jConn",
+    "get_service_instance",
+    "KafkaQueue",
+    "FTPFileTransfer",
+    "SFTPFileTransfer",
+    "parse_args",
+]

dataops_sdk/runner/utils/common.py

@@ -0,0 +1,209 @@
+"""通用工具模块。
+
+提供加解密、时间解析、UUID 生成和 JSON 打印等功能。
+
+Functions:
+    encrypt_body: 加密请求体
+    decrypt_body: 解密请求体
+    generate_uuid: 生成 UUID
+    print_json: 打印 JSON 格式数据
+    parse_trigger_time: 解析时间字符串为 datetime 对象
+    build_trigger_time_vars: 根据触发时间构建日期变量
+
+Example:
+    >>> encrypt_body({"key": "value"})
+    {'__base64': '...'}
+"""
+
+from base64 import b64decode, b64encode
+from calendar import monthrange
+from datetime import date, datetime, timedelta
+from email.utils import parsedate_to_datetime
+from typing import Any, Literal, Mapping
+from uuid import uuid4
+
+from pydantic import BaseModel, JsonValue
+from pydantic_core import PydanticSerializationError, from_json, to_json
+
+
+def encrypt_body(body: JsonValue | BaseModel):
+    if body is None:
+        raise ValueError("body is None")
+    return {"__base64": f"FMIND-{b64encode(to_json(body)).decode('utf-8')}"}
+
+
+def decrypt_body(body: dict[Literal["__base64"], str]) -> Mapping[str, JsonValue]:
+    """
+    解密请求体
+
+    Args:
+        body (dict[Literal["__base64"], str]): 加密后的请求体，包含一个 "__base64" 键，其值为原始请求体的 Base64 编码字符串
+
+    Returns:
+        JsonValue: 原始请求体
+    """
+    if "__base64" not in body:
+        raise ValueError("Invalid body format: missing '__base64' key")
+    base64_str = body["__base64"]
+    if base64_str.startswith("FMIND-"):
+        base64_str = base64_str[6:]
+    bytes_body = b64decode(base64_str)
+    return cast(Mapping[str, JsonValue], from_json(bytes_body))
+
+
+def generate_uuid() -> int:
+    """
+    生成 UUID
+
+    Returns:
+        int: 64位无符号整数对应的字符串形式
+    """
+    return uuid4().int
+
+
+def print_json(data: Any, indent: int | None = None) -> None:
+    """打印 Pydantic 兼容的 JSON 格式数据。
+
+    将数据序列化为 JSON 字符串并打印到标准输出。
+    仅支持字典或 Pydantic 模型等可序列化对象。
+
+    Args:
+        data: 要打印的数据，必须为字典或可序列化为 JSON 的对象。
+
+    Example:
+        >>> print_json({"name": "test", "value": 123})
+        #
+        # {"name": "test", "value": 123}
+    """
+    if not isinstance(data, dict):
+        print("ERROR: data must be a dict or json serializable object.")
+        return
+    try:
+        json_str = to_json(data, indent=indent).decode("utf-8")
+        print()
+        print(json_str, end="")
+    except PydanticSerializationError as e:
+        print(f"ERROR in Serialization: {e}")
+    except Exception as e:
+        print(f"ERROR: {e}")
+
+
+def parse_trigger_time(time_str: str, strict: bool = True) -> datetime:
+    """解析时间字符串为 datetime 对象
+
+    支持格式:
+    - 秒级时间戳: 1704067200
+    - 毫秒级时间戳: 1704067200000
+    - 标准格式: 2024-01-01 12:00:00
+    - ISO 格式: 2024-01-01T12:00:00, 2024-01-01T12:00:00Z
+    - 其他常见格式: 2024/01/01 12:00:00, 2024-01-01
+    """
+    if not time_str:
+        return datetime.now()
+
+    time_str = time_str.strip()
+
+    # 尝试解析为数字（时间戳）
+    if time_str.isdigit():
+        ts = int(time_str)
+        # 判断是秒级还是毫秒级
+        if ts > 1e11:  # 毫秒级时间戳
+            return datetime.fromtimestamp(ts / 1000)
+        else:  # 秒级时间戳
+            return datetime.fromtimestamp(ts)
+
+    # 尝试 ISO 格式
+    try:
+        # 处理带 Z 的 ISO 格式
+        return datetime.fromisoformat(time_str)
+    except ValueError:
+        pass
+
+    # 尝试常见格式
+    formats = [
+        "%Y-%m-%d %H:%M:%S",
+        "%Y-%m-%d %H:%M",
+        "%Y-%m-%d",
+        "%Y/%m/%d %H:%M:%S",
+        "%Y/%m/%d %H:%M",
+        "%Y/%m/%d",
+        "%Y%m%d",
+    ]
+
+    for fmt in formats:
+        try:
+            return datetime.strptime(time_str, fmt)
+        except ValueError:
+            continue
+
+    # 尝试 RFC 2822 格式
+    try:
+        return parsedate_to_datetime(time_str)
+    except Exception:
+        pass
+
+    # 无法解析，抛出异常
+    if strict:
+        raise ValueError(f"无法解析时间字符串: {time_str}")
+    return datetime.now()
+
+
+def build_trigger_time_vars(trigger_time_raw: str) -> dict[str, datetime | date | int]:
+    """根据触发时间构建日期变量"""
+    trigger_time = parse_trigger_time(trigger_time_raw)
+    today = trigger_time.date()
+    yesterday = today - timedelta(days=1)
+    before_yesterday = today - timedelta(days=2)
+    tomorrow = today + timedelta(days=1)
+
+    # 月首月末
+    _, month_last_day_num = monthrange(today.year, today.month)
+    month_first_day = today.replace(day=1)
+    month_last_day = today.replace(day=month_last_day_num)
+
+    # 上月
+    last_month_num = today.month - 1 if today.month != 1 else 12
+    last_month_year_num = today.year if today.month != 1 else today.year - 1
+    _, last_month_last_num = monthrange(last_month_year_num, last_month_num)
+    if today.month == 1:
+        last_month = today.replace(year=today.year - 1, month=12)
+    else:
+        last_month = today.replace(month=today.month - 1, day=min(today.day, last_month_last_num))
+    last_month_first_day = last_month.replace(day=1)
+    last_month_last_day = last_month.replace(day=last_month_last_num)
+
+    # 季度首日
+    quarter = (today.month - 1) // 3
+    quarter_first_month = quarter * 3 + 1
+    quarter_first_day = today.replace(month=quarter_first_month, day=1)
+
+    # 周一
+    weekday = today.weekday()  # 0=Monday, 6=Sunday
+    monday = today - timedelta(days=weekday)
+
+    # 设置全局变量
+    return {
+        "trigger_time": trigger_time,
+        "today": today,
+        "tomorrow": tomorrow,
+        "yesterday": yesterday,
+        "before_yesterday": before_yesterday,
+        "month_first_day": month_first_day,
+        "month_last_day": month_last_day,
+        "last_month": last_month,
+        "last_month_first_day": last_month_first_day,
+        "last_month_last_day": last_month_last_day,
+        "quarter_first_day": quarter_first_day,
+        "monday": monday,
+        # 拓展
+        "cur_year": today.year,
+        "cur_month": today.month,
+    }
+
+
+def parse_args() -> tuple[str, ...]:
+    """解析命令行参数"""
+    import sys
+
+    args = sys.argv[1:]
+    return args