quantdb-sdk 0.1.0__py3-none-any.whl

quantdb_sdk/__init__.py

@@ -0,0 +1,33 @@
+"""QuantDB 量化数据平台官方 Python SDK。"""
+
+try:
+    from importlib.metadata import version
+    # 注意：此处用 PyPI 包名（连字符）查找版本号，不是模块名（下划线）
+    __version__ = version("quantdb-sdk")
+except ImportError:
+    __version__ = "0.1.0"
+
+from .async_client import AsyncQuantDBClient
+from .client import DuckDBWarehouse, QuantDBClient
+from .errors import (
+    AuthError,
+    InsufficientTrafficError,
+    NotFoundError,
+    QuantDBError,
+    RateLimitError,
+    ServerError,
+    ValidationError,
+)
+
+__all__ = [
+    "QuantDBClient",
+    "AsyncQuantDBClient",
+    "DuckDBWarehouse",
+    "QuantDBError",
+    "AuthError",
+    "InsufficientTrafficError",
+    "NotFoundError",
+    "RateLimitError",
+    "ServerError",
+    "ValidationError",
+]

quantdb_sdk/__main__.py

@@ -0,0 +1,77 @@
+"""QuantDB SDK 命令行入口。
+
+提供快速验证和诊断功能：
+  python -m quantdb --version
+  python -m quantdb --help
+  python -m quantdb check <api_key>  # 验证 API Key 是否有效
+"""
+
+import argparse
+import sys
+
+from . import __version__
+from .client import QuantDBClient
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        prog="quantdb",
+        description="QuantDB 量化数据平台官方 Python SDK",
+    )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"%(prog)s {__version__}",
+    )
+    parser.add_argument(
+        "--api-host",
+        default="https://quantdb.quantmind.cloud",
+        help="API 服务地址 (默认: https://quantdb.quantmind.cloud)",
+    )
+
+    subparsers = parser.add_subparsers(dest="command", help="可用命令")
+
+    # check 命令：验证 API Key
+    check_parser = subparsers.add_parser("check", help="验证 API Key 是否有效")
+    check_parser.add_argument("api_key", help="要验证的 API Key")
+
+    # usage 命令：查询用量
+    usage_parser = subparsers.add_parser("usage", help="查询账户用量和订阅状态")
+    usage_parser.add_argument("api_key", help="API Key")
+
+    args = parser.parse_args()
+
+    if not args.command:
+        parser.print_help()
+        return 0
+
+    try:
+        client = QuantDBClient(api_host=args.api_host, api_key=args.api_key)
+
+        if args.command == "check":
+            me = client.get_me()
+            print(f"✓ API Key 有效")
+            print(f"  用户: {me.get('username', 'N/A')}")
+            print(f"  ID: {me.get('id', 'N/A')}")
+            return 0
+
+        if args.command == "usage":
+            usage = client.get_usage()
+            print(f"账户用量:")
+            print(f"  已用: {usage['used_gb']:.2f} GB")
+            print(f"  限额: {usage['limit_gb']:.1f} GB")
+            print(f"  剩余: {usage['remaining_gb']:.2f} GB")
+            print(f"  余额: ¥{usage['balance_yuan']:.2f}")
+            sub = usage.get("subscription", {})
+            print(f"  订阅: {sub.get('status', 'none')}")
+            return 0
+
+    except Exception as e:
+        print(f"✗ 错误: {e}", file=sys.stderr)
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

quantdb_sdk/_utils.py

@@ -0,0 +1,68 @@
+"""QuantDB SDK 内部工具函数。
+
+包含下载目录管理、文件名解析、字节换算等通用工具。
+"""
+
+import os
+import re
+from urllib.parse import unquote
+
+
+def default_download_dir() -> str:
+    """返回默认下载目录。
+
+    优先读取环境变量 ``QUANTDB_DOWNLOAD_DIR``，未设置时按平台回退：
+
+    - Windows: ``D:\\QuantDB\\downloads\\``
+    - Linux/macOS: ``~/QuantDB/downloads/``
+
+    Returns:
+        绝对路径字符串。
+    """
+    env = os.getenv("QUANTDB_DOWNLOAD_DIR")
+    if env:
+        return env.strip()
+    if os.name == "nt":
+        return r"D:\QuantDB\downloads"
+    return os.path.join(os.path.expanduser("~"), "QuantDB", "downloads")
+
+
+def parse_filename_from_content_disposition(value: str, fallback: str) -> str:
+    """从 HTTP Content-Disposition 响应头解析文件名。
+
+    支持 RFC 5987 (filename*=utf-8''name) 和常规 filename="name" 两种格式。
+    解析失败时返回 *fallback*。
+
+    Args:
+        value: Content-Disposition 头的值。
+        fallback: 解析失败时的默认文件名。
+
+    Returns:
+        解析后的文件名或 fallback。
+    """
+    if not value:
+        return fallback
+
+    # RFC 5987 filename*=utf-8''name
+    m = re.search(r"filename\*\s*=\s*(?:[^']*)'" r"(?:[^']*)'([^;]+)", value, re.IGNORECASE)
+    if m:
+        return unquote(m.group(1))
+
+    # filename="name" 或 filename=name
+    m = re.search(r'filename\s*=\s*"?([^";]+)"?', value, re.IGNORECASE)
+    if m:
+        return unquote(m.group(1).strip('"'))
+
+    return fallback
+
+
+def bytes_to_gb(n: int) -> float:
+    """将字节数转换为 GB（1024 进制）。
+
+    Args:
+        n: 字节数。
+
+    Returns:
+        GB 数值（浮点数）。
+    """
+    return n / (1024 ** 3)

quantdb_sdk/async_client.py

@@ -0,0 +1,475 @@
+"""QuantDB 异步 Python SDK（基于 httpx）。"""
+
+import io
+import os
+import re
+from typing import Any, Dict, List, Optional
+
+import httpx
+import pandas as pd
+
+from ._utils import bytes_to_gb, default_download_dir, parse_filename_from_content_disposition
+from .errors import (
+    AuthError,
+    InsufficientTrafficError,
+    NotFoundError,
+    QuantDBError,
+    RateLimitError,
+    ServerError,
+    ValidationError,
+)
+
+
+class AsyncQuantDBClient:
+    """QuantDB 异步客户端。
+
+    基于 httpx，适用于 asyncio 量化框架。所有异步方法名以 a_ 开头。
+    """
+
+    def __init__(
+        self,
+        api_host: str = "https://quantdb.quantmind.cloud",
+        api_key: Optional[str] = None,
+        token: Optional[str] = None,
+        timeout: float = 60.0,
+    ):
+        self.api_host = api_host.rstrip("/")
+        self.timeout = timeout
+        headers = {"User-Agent": "QuantDB-Python-SDK/0.1.0"}
+        if api_key:
+            headers["X-API-Key"] = api_key
+        elif token:
+            headers["Authorization"] = f"Bearer {token}"
+        else:
+            raise ValueError("必须提供 api_key 或 token")
+
+        self.client = httpx.AsyncClient(
+            headers=headers,
+            timeout=httpx.Timeout(timeout, connect=5.0),
+        )
+        # 进程内 parquet ETag 缓存：同对象（ETag 未变）不重复下载，避免重复计费。
+        self._cache: Dict[str, Dict[str, Any]] = {}
+
+    async def close(self) -> None:
+        """关闭底层 httpx 客户端。"""
+        await self.client.aclose()
+
+    def clear_cache(self) -> None:
+        """清空进程内 Parquet 缓存（强制下次重新下载最新数据）。"""
+        self._cache.clear()
+
+    async def __aenter__(self) -> "AsyncQuantDBClient":
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        await self.close()
+
+    async def _request(
+        self,
+        method: str,
+        path: str,
+        params: Optional[dict] = None,
+        json: Optional[dict] = None,
+    ) -> httpx.Response:
+        url = f"{self.api_host}{path}"
+        return await self.client.request(method=method, url=url, params=params, json=json)
+
+    async def _get(self, path: str, params: Optional[dict] = None) -> dict:
+        resp = await self._request("GET", path, params=params)
+        return self._check_response(resp)
+
+    async def _post(self, path: str, json: Optional[dict] = None) -> dict:
+        resp = await self._request("POST", path, json=json)
+        return self._check_response(resp)
+
+    def _check_response(self, response: httpx.Response) -> dict:
+        if response.status_code == 200:
+            try:
+                return response.json()
+            except Exception as exc:
+                raise QuantDBError(f"响应不是合法 JSON: {exc}") from exc
+
+        try:
+            payload = response.json()
+            detail = payload.get("detail", "")
+        except Exception:
+            detail = response.text[:200]
+
+        msg = detail or f"HTTP {response.status_code}"
+
+        if response.status_code == 401:
+            raise AuthError(f"认证失败：{msg}")
+        if response.status_code == 402:
+            raise InsufficientTrafficError(f"无有效订阅或本月流量已用完：{msg}")
+        if response.status_code == 403:
+            raise AuthError(f"无权访问：{msg}")
+        if response.status_code == 404:
+            raise NotFoundError(f"资源不存在：{msg}")
+        if response.status_code == 422:
+            raise ValidationError(f"参数错误：{msg}")
+        if response.status_code == 429:
+            raise RateLimitError(f"请求过于频繁：{msg}")
+        if response.status_code >= 500:
+            raise ServerError(f"服务端错误：{msg}")
+        raise QuantDBError(f"请求失败：{msg}")
+
+    # ========== 账户信息 ==========
+
+    async def a_get_me(self) -> Dict[str, Any]:
+        return await self._get("/api/v1/auth/me")
+
+    async def a_get_usage(self) -> Dict[str, Any]:
+        data = await self._get("/api/v1/auth/usage")
+        used = data.get("used_traffic", 0)
+        limit = data.get("traffic_limit", 0)
+        credit = data.get("credit_limit", 0)
+        subscription = data.get("subscription", {}) or {}
+        return {
+            "used_gb": bytes_to_gb(used),
+            "limit_gb": bytes_to_gb(limit),
+            "credit_gb": bytes_to_gb(credit),
+            "remaining_gb": max(0.0, bytes_to_gb(limit + credit - used)),
+            "balance_yuan": data.get("balance_yuan", 0.0),
+            "subscription": subscription,
+            "is_active": subscription.get("status") == "active" and used < limit + credit,
+        }
+
+    async def a_register(self, username: str, email: str, password: str) -> Dict[str, Any]:
+        return await self._post(
+            "/api/v1/auth/register",
+            json={"username": username, "email": email, "password": password},
+        )
+
+    async def a_forgot_password(self, email: str) -> Dict[str, Any]:
+        return await self._post("/api/v1/auth/forgot-password", json={"email": email})
+
+    async def a_reset_password(self, token: str, new_password: str) -> Dict[str, Any]:
+        return await self._post(
+            "/api/v1/auth/reset-password",
+            json={"token": token, "new_password": new_password},
+        )
+
+    # ========== API Key 管理 ==========
+
+    async def a_list_api_keys(self) -> List[Dict[str, Any]]:
+        data = await self._get("/api/v1/auth/api-keys")
+        return data.get("keys", data.get("data", []))
+
+    async def a_create_api_key(self, description: str = "") -> Dict[str, Any]:
+        """创建新的 API Key（每个账号限制 1 个有效 Key）。"""
+        return await self._post("/api/v1/auth/api-keys", json={"description": description})
+
+    async def a_delete_api_key(self, key_token: str) -> Dict[str, Any]:
+        """删除指定的 API Key。
+
+        Args:
+            key_token: 要删除的 API Key 令牌。
+
+        Returns:
+            删除结果信息。
+        """
+        return await self._post(f"/api/v1/auth/api-keys/{key_token}/delete")
+
+    # ========== 订阅与支付 ==========
+
+    async def a_list_plans(self) -> List[Dict[str, Any]]:
+        data = await self._get("/api/v1/subscription/plans")
+        return data.get("plans", data.get("data", []))
+
+    async def a_create_order(self, plan_id: str) -> Dict[str, Any]:
+        return await self._post("/api/v1/subscription/orders", json={"plan_id": plan_id})
+
+    async def a_get_order(self, order_id: int) -> Dict[str, Any]:
+        return await self._get(f"/api/v1/subscription/orders/{order_id}")
+
+    async def a_get_version(self) -> Dict[str, Any]:
+        """获取服务端版本信息。
+
+        Returns:
+            包含版本号、最新版本号、下载地址等信息的字典。
+        """
+        return await self._get("/api/v1/version")
+
+    # ========== 数据查询：K线/Tick 走下载（消耗流量），列表/日历走网关 JSON（不计流量） ==========
+
+    async def a_query_kline(
+        self,
+        symbol: str,
+        adj_type: str = "unadjusted",
+        start_date: Optional[str] = None,
+        end_date: Optional[str] = None,
+        fields: str = "open,high,low,close,volume,amount",
+        limit: Optional[int] = None,
+    ) -> pd.DataFrame:
+        """查询 K 线数据（下载 COS parquet 切片后客户端解析，消耗下载流量，异步）。"""
+        sub_category = f"daily_{adj_type}"
+        df = await self.a_load_as_df("1", sub_category, symbol)
+        # 日期过滤：K线 parquet 的日期列为 'time'，对外以 'trade_date' 暴露。
+        if "time" in df.columns:
+            dt = pd.to_datetime(df["time"], errors="coerce")
+            if dt.dt.tz is not None:
+                dt = dt.dt.tz_convert(None)
+            mask = pd.Series(True, index=df.index)
+            if start_date:
+                mask &= dt >= pd.to_datetime(start_date)
+            if end_date:
+                mask &= dt <= pd.to_datetime(end_date)
+            df = df[mask].copy()
+            df["trade_date"] = dt[mask].dt.strftime("%Y-%m-%d")
+            df = df.drop(columns=["time"])
+        field_list = [f.strip() for f in fields.split(",") if f.strip()]
+        cols = [c for c in field_list if c in df.columns]
+        keep = (["trade_date"] if "trade_date" in df.columns else []) + cols
+        keep = list(dict.fromkeys(keep))
+        df = df[keep]
+        if limit is not None:
+            df = df.tail(limit)
+        return df.reset_index(drop=True)
+
+    async def a_query_tick(
+        self,
+        symbol: str,
+        trade_date: str,
+        start_ts: Optional[str] = None,
+        end_ts: Optional[str] = None,
+        fields: str = "last_price,open,high,low,last_close,volume,amount",
+        limit: Optional[int] = None,
+    ) -> pd.DataFrame:
+        """查询 Tick 分笔数据（下载 COS parquet 切片后客户端解析，消耗下载流量，异步）。"""
+        df = await self.a_load_as_df("1", "tick_data", symbol, trade_date=trade_date)
+        ts_col = "ts" if "ts" in df.columns else ("time" if "time" in df.columns else None)
+        if ts_col and (start_ts or end_ts):
+            ts = pd.to_datetime(df[ts_col], errors="coerce")
+            if ts.dt.tz is not None:
+                ts = ts.dt.tz_convert(None)
+            mask = pd.Series(True, index=df.index)
+            if start_ts:
+                s = start_ts if (len(start_ts) > 8 or "-" in start_ts) else f"{trade_date} {start_ts}"
+                mask &= ts >= pd.to_datetime(s)
+            if end_ts:
+                e = end_ts if (len(end_ts) > 8 or "-" in end_ts) else f"{trade_date} {end_ts}"
+                mask &= ts <= pd.to_datetime(e)
+            df = df[mask].copy()
+        field_list = [f.strip() for f in fields.split(",") if f.strip()]
+        cols = [c for c in field_list if c in df.columns]
+        keep = ([ts_col] if ts_col else []) + cols
+        keep = list(dict.fromkeys(keep))
+        df = df[keep]
+        if limit is not None:
+            df = df.head(limit)
+        return df.reset_index(drop=True)
+
+    async def a_query_stock_list(
+        self, keyword: Optional[str] = None, limit: int = 200
+    ) -> pd.DataFrame:
+        params = {"limit": limit}
+        if keyword:
+            params["keyword"] = keyword
+        data = await self._get("/api/v1/data/stock-list", params)
+        return pd.DataFrame(data.get("data", []))
+
+    async def a_query_calendar(
+        self,
+        start_date: Optional[str] = None,
+        end_date: Optional[str] = None,
+    ) -> pd.DataFrame:
+        params: Dict[str, Any] = {}
+        if start_date:
+            params["start_date"] = start_date
+        if end_date:
+            params["end_date"] = end_date
+        data = await self._get("/api/v1/data/calendar", params)
+        return pd.DataFrame(data.get("data", []))
+
+    async def a_query_manifest(
+        self,
+        category_id: str,
+        sub_category: str,
+        trade_date: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        params: Dict[str, Any] = {
+            "category_id": category_id,
+            "sub_category": sub_category,
+        }
+        if trade_date:
+            params["trade_date"] = trade_date
+        data = await self._get("/api/v1/data/download/manifest", params)
+        return data.get("files", [])
+
+    async def a_query_meta(
+        self,
+        dataset: Optional[str] = None,
+        source: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """查询数据元数据（各数据集日期范围/行数/可用性，异步）。"""
+        params: Dict[str, Any] = {}
+        if dataset:
+            params["dataset"] = dataset
+        if source:
+            params["source"] = source
+        data = await self._get("/api/v1/data/meta", params)
+        return pd.DataFrame(data.get("data", []))
+
+    # ========== 数据预览与下载 ==========
+
+    async def a_preview_as_df(
+        self,
+        category_id: str,
+        sub_category: str,
+        symbol: Optional[str] = None,
+        limit: int = 30,
+    ) -> pd.DataFrame:
+        params: Dict[str, Any] = {
+            "category_id": category_id,
+            "sub_category": sub_category,
+            "limit": limit,
+        }
+        if symbol:
+            params["symbol"] = symbol
+        data = await self._get("/api/v1/data/preview", params)
+        return pd.DataFrame(data.get("data", []), columns=data.get("columns"))
+
+    async def a_download_file(
+        self,
+        category_id: str,
+        sub_category: str,
+        symbol: Optional[str] = None,
+        save_dir: Optional[str] = None,
+        trade_date: Optional[str] = None,
+    ) -> str:
+        if save_dir is None:
+            save_dir = default_download_dir()
+        os.makedirs(save_dir, exist_ok=True)
+
+        params: Dict[str, Any] = {
+            "category_id": category_id,
+            "sub_category": sub_category,
+        }
+        if symbol:
+            params["symbol"] = symbol
+        if trade_date:
+            params["trade_date"] = trade_date
+
+        async with self.client.stream(
+            "GET", f"{self.api_host}/api/v1/data/download", params=params
+        ) as resp:
+            if resp.status_code != 200:
+                body = await resp.aread()
+                # 构造一个完整的 httpx.Response 用于错误解析，然后立即抛出
+                err_resp = httpx.Response(resp.status_code, content=body)
+                self._check_response(err_resp)
+                # 如果 _check_response 没有抛出（理论上不应该），安全退出
+                raise QuantDBError(f"下载失败，HTTP {resp.status_code}")
+
+            cd = ""
+            if "Content-Disposition" in resp.headers:
+                cd = resp.headers["Content-Disposition"]
+            fallback = f"{sub_category}.parquet"
+            if symbol:
+                fallback = f"{sub_category}_{symbol}.parquet"
+            filename = parse_filename_from_content_disposition(cd, fallback)
+
+            save_path = os.path.join(save_dir, filename)
+            with open(save_path, "wb") as f:
+                async for chunk in resp.aiter_bytes(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+        return os.path.abspath(save_path)
+
+    async def a_load_as_df(
+        self,
+        category_id: str,
+        sub_category: str,
+        symbol: Optional[str] = None,
+        trade_date: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """远端 Parquet 切片加载到内存 DataFrame（消耗下载流量，异步）。
+
+        带进程内 ETag 缓存：同对象（ETag 未变）不重复下载，避免重复计费。
+        """
+        cache_key = f"{category_id}/{sub_category}/{symbol}/{trade_date}"
+        params: Dict[str, Any] = {
+            "category_id": category_id,
+            "sub_category": sub_category,
+        }
+        if symbol:
+            params["symbol"] = symbol
+        if trade_date:
+            params["trade_date"] = trade_date
+        cached = self._cache.get(cache_key)
+        req_headers = {}
+        if cached and cached.get("etag"):
+            req_headers["If-None-Match"] = cached["etag"]
+        resp = await self.client.get(
+            f"{self.api_host}/api/v1/data/download", params=params, headers=req_headers
+        )
+        if resp.status_code == 304 and cached:
+            return cached["df"]
+        if resp.status_code != 200:
+            self._check_response(resp)
+        etag = resp.headers.get("ETag") or ""
+        if cached and etag and cached.get("etag") == etag:
+            return cached["df"]
+        df = pd.read_parquet(io.BytesIO(resp.content))
+        if etag:
+            self._cache[cache_key] = {"etag": etag, "df": df}
+        return df
+
+    async def a_query_local(
+        self,
+        sql: str,
+        category_id: str,
+        sub_category: str,
+        symbol: Optional[str] = None,
+        save_dir: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """用 DuckDB 查询本地已下载的 Parquet 文件（零流量，异步）。
+
+        若文件不存在会先自动下载。
+
+        sql 支持两种写法：
+            1. 完整 SELECT 语句（必须包含 FROM 子句，表名会被替换为 parquet 路径）
+            2. 纯 WHERE 条件字符串（不含 FROM 时自动拼装为 SELECT * FROM <path> WHERE <条件>）
+
+        安全限制：
+            - 不允许分号（;），防止多语句注入
+            - 不允许注释（-- 或 /* */），防止注释注入
+            - 不允许 UNION / DROP / DELETE / INSERT / UPDATE / ALTER / CREATE / EXEC / EXECUTE
+        """
+        file_path = await self.a_download_file(
+            category_id=category_id,
+            sub_category=sub_category,
+            symbol=symbol,
+            save_dir=save_dir,
+        )
+        clean_path = file_path.replace("\\", "/")
+
+        try:
+            import duckdb  # type: ignore
+        except ImportError as exc:
+            raise RuntimeError("本地 SQL 查询需要 duckdb：pip install duckdb") from exc
+
+        # 安全检查：拒绝危险字符和关键字
+        dangerous_keywords = [
+            ";", "--", "/*", "*/", "union", "drop", "delete",
+            "insert", "update", "alter", "create", "exec", "execute",
+        ]
+        sql_upper = sql.upper()
+        for kw in dangerous_keywords:
+            if kw in sql_upper:
+                raise ValidationError(
+                    f"SQL 包含危险关键字 '{kw}'，已被拒绝。"
+                    "a_query_local 仅支持 SELECT 查询，不支持数据修改或联合查询。"
+                )
+
+        if "FROM" in sql_upper:
+            sql = re.sub(
+                r"FROM\s+[`\"\']?\w+[`\"\']?",
+                f"FROM '{clean_path}'",
+                sql,
+                count=1,
+                flags=re.IGNORECASE,
+            )
+        else:
+            sql = f"SELECT * FROM '{clean_path}' WHERE {sql}"
+        return duckdb.query(sql).df()