lib-x17-fintech 2.1.3__py3-none-any.whl

xfintech/data/source/tushare/session/relay.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import gzip
+import hashlib
+import hmac
+import io
+import secrets
+import time
+from typing import Any, Callable, Dict, Optional
+
+import pandas as pd
+import requests
+
+from xfintech.data.relay.client import RelayClient
+
+
+class TushareRelayClient(RelayClient):
+    """
+    描述:
+    - Tushare 中继服务客户端类，用于通过中继服务器访问 Tushare 数据。
+    - 实现 HMAC-SHA256 签名认证机制，确保请求安全性。
+    - 支持数据压缩和 Parquet 格式传输，提高数据传输效率。
+    - 提供健康检查功能，确保服务可用性。
+
+    属性:
+    - url: str, 中继服务器 URL 地址。
+    - secret: str, 用于 HMAC 签名的密钥。
+    - timeout: int, 请求超时时间（秒），默认 180 秒。
+    - DEFAULT_TIMEOUT: int = 180, 默认超时时间常量。
+
+    方法:
+    - call(method, limit, offset, params): 调用 Tushare API 方法并返回 DataFrame。
+    - check_health(): 检查中继服务器健康状态。
+    - canonical_json(value): 生成规范化的 JSON 字节串用于签名。
+
+    例子:
+    ```python
+        from xfintech.data.source.tushare.session import TushareRelayClient
+
+        # 创建中继客户端
+        client = TushareRelayClient(
+            url="https://relay.example.com",
+            secret="your-secret-key",
+            timeout=120
+        )
+
+        # 调用 API 获取数据
+        df = client.call(
+            method="daily",
+            limit=100,
+            offset=0,
+            params={"ts_code": "000001.SZ", "start_date": "20240101"}
+        )
+        print(df.head())
+    ```
+    """
+
+    def __init__(
+        self,
+        url: str,
+        secret: str,
+        timeout: Optional[int] = None,
+    ) -> None:
+        super().__init__(
+            url=url,
+            secret=secret,
+            timeout=timeout,
+        )
+
+    def call(
+        self,
+        method: str,
+        limit: int | None,
+        offset: int | None,
+        params: Dict[str, Any],
+    ) -> pd.DataFrame:
+        payload: Dict[str, Any] = {
+            "path": f"{method}",
+            "params": params,
+            "limit": limit,
+            "offset": offset,
+        }
+        body = self.canonical_json(payload)
+        nonce = secrets.token_hex(16)
+        timestamp = str(int(time.time()))
+        msg = f"{nonce}.{timestamp}.".encode("utf-8") + body
+        sig = hmac.new(
+            self.secret.encode("utf-8"),
+            msg,
+            hashlib.sha256,
+        ).hexdigest()
+        url = f"{self.url}/v2/tushare/call"
+        response = requests.post(
+            url,
+            data=body,
+            headers={
+                "Content-Type": "application/json",
+                "X-YNONCE": nonce,
+                "X-YTS": timestamp,
+                "X-YSIGN": sig,
+                "X-Format": "parquet",
+                "X-Compression": "zstd+gzip",
+            },
+            timeout=self.timeout,
+        )
+        response.raise_for_status()
+        parquet_bytes = gzip.decompress(response.content)
+        return pd.read_parquet(io.BytesIO(parquet_bytes))
+
+    def refresh(
+        self,
+    ) -> bool:
+        try:
+            payload: Dict[str, Any] = {}
+            body = self.canonical_json(payload)
+            nonce = secrets.token_hex(16)
+            timestamp = str(int(time.time()))
+            msg = f"{nonce}.{timestamp}.".encode("utf-8") + body
+            sig = hmac.new(
+                self.secret.encode("utf-8"),
+                msg,
+                hashlib.sha256,
+            ).hexdigest()
+            url = f"{self.url}/v2/tushare/refresh"
+            response = requests.post(
+                url,
+                data=body,
+                headers={
+                    "Content-Type": "application/json",
+                    "X-YNONCE": nonce,
+                    "X-YTS": timestamp,
+                    "X-YSIGN": sig,
+                },
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+            data = response.json()
+            if data.get("status") != "ok":
+                msg = f"Refresh returned non-ok status: {data}"
+                raise RuntimeError(msg)
+            return True
+        except Exception as e:
+            msg = f"Tushare refresh failed: {e}"
+            raise RuntimeError(msg) from e
+
+    def check_health(
+        self,
+    ) -> bool:
+        try:
+            url = f"{self.url}/health"
+            response = requests.get(
+                url,
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+            data = response.json()
+            health = data.get("status") == "ok"
+            if not health:
+                raise RuntimeError("Health check returned non-ok status.")
+            return True
+        except Exception as e:
+            raise RuntimeError(f"Health check failed: {e}") from e
+
+
+class RelayConnection:
+    """
+    描述:
+    - Tushare 中继连接类，提供动态方法调用接口。
+    - 通过魔术方法 __getattr__ 实现类似 Tushare Pro API 的调用方式。
+    - 自动将方法调用转发到 TushareRelayClient 进行处理。
+
+    属性:
+    - client: TushareRelayClient, 中继客户端实例。
+
+    方法:
+    - 动态方法: 通过属性访问自动创建对应的 API 调用方法。
+
+    例子:
+    ```python
+        from xfintech.data.source.tushare.session import TushareRelayClient, RelayConnection
+
+        # 创建客户端和连接
+        client = TushareRelayClient(
+            url="https://relay.example.com",
+            secret="your-secret-key"
+        )
+        connection = RelayConnection(client=client)
+
+        # 使用类似 Tushare Pro API 的方式调用
+        # 等同于 pro.daily() 的调用方式
+        df = connection.daily(
+            ts_code="000001.SZ",
+            start_date="20240101",
+            end_date="20240131",
+            limit=100,
+            offset=0
+        )
+        print(df.head())
+
+        # 调用其他方法
+        df_basic = connection.stock_basic(
+            exchange="SSE",
+            list_status="L"
+        )
+    ```
+    """
+
+    def __init__(
+        self,
+        client: TushareRelayClient,
+    ) -> None:
+        self.client = client
+
+    def __getattr__(
+        self,
+        method: str,
+    ) -> Callable[..., pd.DataFrame]:
+        def _call(
+            *,
+            limit: int | None = None,
+            offset: int | None = None,
+            **params: Any,
+        ) -> pd.DataFrame:
+            return self.client.call(
+                method=method,
+                limit=limit,
+                offset=offset,
+                params=params,
+            )
+
+        return _call

xfintech/data/source/tushare/session/session.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+import uuid
+from datetime import datetime
+from typing import Any, Dict, Literal, Optional
+
+import pandas as pd
+import tushare as ts
+
+from xfintech.data.source.tushare.session.relay import (
+    RelayConnection,
+    TushareRelayClient,
+)
+
+
+class Session:
+    """
+    描述:
+    - Tushare 数据源会话管理类，支持直接连接和通过中继服务器连接两种模式。
+    - 提供连接管理、状态跟踪和描述功能。
+    - 支持会话持续时间跟踪，记录连接开始和结束时间。
+
+    属性:
+    - id: str, 会话唯一标识符（UUID 前 8 位）。
+    - mode: str, 连接模式，"direct" 或 "relay"。
+    - connection: object, Tushare 连接对象（pro_api 或 RelayConnection）。
+    - start_at: pd.Timestamp, 会话开始时间。
+    - finish_at: pd.Timestamp, 会话结束时间。
+    - relay_url: str | None, 中继服务器 URL（relay 模式下必填）。
+    - relay_secret: str | None, 中继服务器密钥（relay 模式下必填）。
+
+    方法:
+    - connect(): 建立连接，根据模式创建 direct 或 relay 连接。
+    - disconnect(): 断开连接并记录结束时间。
+    - start(): 记录会话开始时间。
+    - end(): 记录会话结束时间。
+    - get_start_iso(): 返回 ISO 格式的开始时间。
+    - get_finish_iso(): 返回 ISO 格式的结束时间。
+    - describe(): 返回会话描述信息（敏感信息已脱敏）。
+    - to_dict(): 返回会话的完整字典表示。
+
+    属性 (Property):
+    - duration: float, 会话持续时间（秒）。
+    - connected: bool, 是否已连接。
+
+    例子:
+    ```python
+        from xfintech.data.source.tushare.session import Session
+
+        # 直接连接模式
+        session = Session(
+            credential="your-tushare-token",
+            mode="direct"
+        )
+        print(f"会话 ID: {session.id}")
+        print(f"连接状态: {session.connected}")
+
+        # 使用 Tushare API
+        df = session.connection.daily(ts_code="000001.SZ", start_date="20240101")
+
+        # 查看会话信息
+        print(session.describe())
+        print(f"会话持续时间: {session.duration} 秒")
+
+        # 关闭连接
+        session.disconnect()
+
+        # 中继连接模式
+        relay_session = Session(
+            credential="your-tushare-token",
+            mode="relay",
+            relay_url="https://relay.example.com",
+            relay_secret="your-relay-secret"
+        )
+        df = relay_session.connection.daily(ts_code="000001.SZ")
+    ```
+    """
+
+    def __init__(
+        self,
+        credential: Optional[str] = None,
+        mode: Literal["direct", "relay"] = "direct",
+        relay_url: str | None = None,
+        relay_secret: str | None = None,
+    ) -> None:
+        self._credential = credential
+        self.id = str(uuid.uuid4())[:8]
+        self.connection = None
+        self.start_at = None
+        self.finish_at = None
+        self.mode = self._resolve_mode(mode)
+        self.relay_url = self._resolve_relay_url(
+            relay_url,
+            self.mode,
+        )
+        self.relay_secret = self._resolve_relay_secret(
+            relay_secret,
+            self.mode,
+        )
+        self.connect()
+
+    def _resolve_mode(
+        self,
+        mode: Literal["direct", "relay"],
+    ) -> str:
+        mode = mode.lower()
+        if not mode:
+            return "direct"
+        if mode not in ["direct", "relay"]:
+            msg = f"Unsupported mode: {mode}"
+            raise ValueError(msg)
+        return mode
+
+    def _resolve_relay_url(
+        self,
+        url: str | None,
+        mode: str,
+    ) -> str | None:
+        if mode == "relay":
+            if not url:
+                msg = "URL must be provided in relay mode."
+                raise ValueError(msg)
+            return url
+        else:
+            return None
+
+    def _resolve_relay_secret(
+        self,
+        secret: str | None,
+        mode: str,
+    ) -> str | None:
+        if mode == "relay":
+            if not secret:
+                msg = "Secret must be provided in relay mode."
+                raise ValueError(msg)
+            return secret
+        else:
+            return None
+
+    @property
+    def duration(self) -> float:
+        if not self.start_at:
+            return 0.0
+        if not self.finish_at:
+            now = datetime.now()
+            delta = now - self.start_at
+            return delta.total_seconds()
+        delta = self.finish_at - self.start_at
+        return delta.total_seconds()
+
+    @property
+    def connected(self) -> bool:
+        return self.connection is not None
+
+    def __str__(self) -> str:
+        return f"{self.id}"
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(connected={self.connected}, mode={self.mode})"
+
+    def start(self) -> None:
+        self.start_at = pd.Timestamp.now()
+
+    def get_start_iso(self) -> Optional[str]:
+        if self.start_at is None:
+            return None
+        return self.start_at.isoformat()
+
+    def end(self) -> None:
+        self.finish_at = pd.Timestamp.now()
+
+    def get_finish_iso(self) -> Optional[str]:
+        if self.finish_at is None:
+            return None
+        return self.finish_at.isoformat()
+
+    def connect(self) -> object:
+        if self.connected:
+            return self.connection
+        if self.mode == "direct":
+            ts.set_token(self._credential)
+            self.connection = ts.pro_api()
+        else:
+            client = TushareRelayClient(
+                url=self.relay_url,
+                secret=self.relay_secret,
+            )
+            client.check_health()
+            self.connection = RelayConnection(
+                client=client,
+            )
+        self.start()
+        return self.connection
+
+    def disconnect(self) -> None:
+        self.connection = None
+        self.end()
+
+    def refresh(self) -> None:
+        if self.mode == "relay" and self.connected:
+            self.connection.client.refresh()
+        else:
+            self.disconnect()
+            self.connect()
+
+    def describe(self) -> Dict[str, Any]:
+        start_at_iso = self.get_start_iso()
+        finish_at_iso = self.get_finish_iso()
+        result = {}
+        result["id"] = self.id
+        result["mode"] = self.mode
+        if self._credential:
+            result["credential"] = "******"
+        if self.mode == "relay":
+            result["relay"] = {
+                "url": self.relay_url,
+                "secret": "******",
+            }
+        result["connected"] = self.connected
+        if start_at_iso:
+            result["start_at"] = start_at_iso
+        if finish_at_iso:
+            result["finish_at"] = finish_at_iso
+        return result
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+            "connected": self.connected,
+            "credential": "******" if self._credential else None,
+            "mode": self.mode,
+            "relay": {
+                "url": self.relay_url,
+                "secret": "******" if self.relay_secret else None,
+            },
+            "start_at": self.get_start_iso(),
+            "finish_at": self.get_finish_iso(),
+            "duration": self.duration,
+        }