quantdb-sdk 0.1.0__py3-none-any.whl

quantdb_sdk/client.py

@@ -0,0 +1,610 @@
+"""QuantDB 同步 Python SDK。"""
+
+import io
+import os
+import re
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+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 QuantDBClient:
+    """QuantDB 同步客户端。
+
+    支持 API Key 或用户名/密码两种鉴权方式。
+    数据查询：K线/Tick 走下载（消耗流量），列表/日历/元数据走网关 JSON（不计流量）。
+    """
+
+    def __init__(
+        self,
+        api_host: str = "https://quantdb.quantmind.cloud",
+        api_key: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        timeout: tuple = (5, 60),
+        max_retries: int = 2,
+    ):
+        self.api_host = api_host.rstrip("/")
+        self.timeout = timeout
+        self.token: Optional[str] = None
+        self.headers: Dict[str, str] = {}
+
+        if api_key:
+            self.headers = {"X-API-Key": api_key}
+        elif username and password:
+            self._login(username, password)
+        else:
+            raise ValueError("必须提供 api_key 或 username+password")
+
+        self.session = requests.Session()
+        self.session.headers.update({"User-Agent": "QuantDB-Python-SDK/0.1.0"})
+        if self.headers:
+            self.session.headers.update(self.headers)
+
+        # 进程内 parquet 缓存：cache_key -> {"etag": str, "df": DataFrame}
+        # 用 ETag 判断对象是否变化：同对象（同 ETag）不重复下载，避免重复计费。
+        # 典型场景：量化脚本内对同一 symbol 多次查询不同字段/日期范围，底层 parquet
+        # 只需下载一次。进程退出即失效（无跨进程副作用）。
+        self._cache: Dict[str, Dict[str, Any]] = {}
+
+        retries = Retry(
+            total=max_retries,
+            backoff_factor=0.5,
+            status_forcelist=[429, 502, 503, 504],
+            allowed_methods=["GET", "POST"],
+        )
+        self.session.mount("http://", HTTPAdapter(max_retries=retries))
+        self.session.mount("https://", HTTPAdapter(max_retries=retries))
+
+    def _login(self, username: str, password: str) -> None:
+        """用户名密码登录，获取 JWT Token。"""
+        resp = self.session.post(
+            f"{self.api_host}/api/v1/auth/login",
+            json={"username": username, "password": password},
+            timeout=self.timeout,
+        )
+        data = self._check_response(resp)
+        self.token = data["access_token"]
+        self.session.headers.update({"Authorization": f"Bearer {self.token}"})
+
+    def _request(
+        self,
+        method: str,
+        path: str,
+        params: Optional[dict] = None,
+        json: Optional[dict] = None,
+        stream: bool = False,
+        headers: Optional[dict] = None,
+    ) -> requests.Response:
+        url = f"{self.api_host}{path}"
+        resp = self.session.request(
+            method=method,
+            url=url,
+            params=params,
+            json=json,
+            stream=stream,
+            timeout=self.timeout,
+            headers=headers,
+        )
+        return resp
+
+    def _get(self, path: str, params: Optional[dict] = None) -> dict:
+        resp = self._request("GET", path, params=params)
+        return self._check_response(resp)
+
+    def _post(self, path: str, json: Optional[dict] = None) -> dict:
+        resp = self._request("POST", path, json=json)
+        return self._check_response(resp)
+
+    def _check_response(self, response: requests.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}")
+
+    def clear_cache(self) -> None:
+        """清空进程内 Parquet 缓存。
+
+        load_as_df 会按 ETag 缓存已下载的 parquet，避免重复查询同 symbol 被重复计费。
+        当确认需要强制重新下载最新数据时调用此方法。
+        """
+        self._cache.clear()
+
+    # ========== 账户信息 ==========
+
+    def get_me(self) -> Dict[str, Any]:
+        """获取当前登录用户信息。"""
+        return self._get("/api/v1/auth/me")
+
+    def get_usage(self) -> Dict[str, Any]:
+        """获取当前账户流量与订阅使用情况。"""
+        data = 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,
+        }
+
+    def register(self, username: str, email: str, password: str) -> Dict[str, Any]:
+        """用户注册。"""
+        return self._post(
+            "/api/v1/auth/register",
+            json={"username": username, "email": email, "password": password},
+        )
+
+    def forgot_password(self, email: str) -> Dict[str, Any]:
+        """发送密码重置邮件。"""
+        return self._post("/api/v1/auth/forgot-password", json={"email": email})
+
+    def reset_password(self, token: str, new_password: str) -> Dict[str, Any]:
+        """使用邮件中的重置令牌设置新密码。"""
+        return self._post(
+            "/api/v1/auth/reset-password",
+            json={"token": token, "new_password": new_password},
+        )
+
+    # ========== API Key 管理 ==========
+
+    def list_api_keys(self) -> List[Dict[str, Any]]:
+        """列出当前用户的 API Key。"""
+        data = self._get("/api/v1/auth/api-keys")
+        return data.get("keys", data.get("data", []))
+
+    def create_api_key(self, description: str = "") -> Dict[str, Any]:
+        """创建新的 API Key（每个账号限制 1 个有效 Key）。"""
+        return self._post("/api/v1/auth/api-keys", json={"description": description})
+
+    def delete_api_key(self, key_token: str) -> Dict[str, Any]:
+        """删除指定的 API Key。
+
+        Args:
+            key_token: 要删除的 API Key 令牌。
+
+        Returns:
+            删除结果信息。
+        """
+        return self._post(f"/api/v1/auth/api-keys/{key_token}/delete")
+
+    # ========== 订阅与支付 ==========
+
+    def list_plans(self) -> List[Dict[str, Any]]:
+        """获取可购买套餐列表。"""
+        data = self._get("/api/v1/subscription/plans")
+        return data.get("plans", data.get("data", []))
+
+    def create_order(self, plan_id: str) -> Dict[str, Any]:
+        """创建订单，返回支付跳转 URL。"""
+        return self._post("/api/v1/subscription/orders", json={"plan_id": plan_id})
+
+    def get_order(self, order_id: int) -> Dict[str, Any]:
+        """查询订单状态。"""
+        return self._get(f"/api/v1/subscription/orders/{order_id}")
+
+    def get_version(self) -> Dict[str, Any]:
+        """获取服务端版本信息。
+
+        Returns:
+            包含版本号、最新版本号、下载地址等信息的字典。
+        """
+        return self._get("/api/v1/version")
+
+    # ========== 数据查询：K线/Tick 走下载（消耗流量），列表/日历走网关 JSON（不计流量） ==========
+
+    def 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 切片后客户端解析，消耗下载流量）。
+
+        下载整个 symbol 的日线 parquet，按 start_date/end_date 过滤日期、按 fields 选列。
+        limit 非 None 时只返回尾部 limit 行。
+        """
+        sub_category = f"daily_{adj_type}"
+        df = self.load_as_df("1", sub_category, symbol)
+        # 日期过滤：K线 parquet 的日期列为 'time'（采集端 01_kline.py 统一命名）。
+        # 对外仍以 'trade_date' 暴露（YYYY-MM-DD），便于用户按日期理解。
+        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"])
+        # 字段过滤：保留 trade_date + 选中字段
+        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)
+
+    def 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 切片后客户端解析，消耗下载流量）。
+
+        下载 trade_date 当日该 symbol 的 tick parquet，按 start_ts/end_ts 过滤时间、按 fields 选列。
+        start_ts/end_ts 可传完整时间戳或 "HH:MM:SS"（自动补 trade_date 日期）。
+        """
+        df = self.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)
+
+    def query_stock_list(
+        self, keyword: Optional[str] = None, limit: int = 200
+    ) -> pd.DataFrame:
+        """查询 A 股基础列表。"""
+        params = {"limit": limit}
+        if keyword:
+            params["keyword"] = keyword
+        data = self._get("/api/v1/data/stock-list", params)
+        return pd.DataFrame(data.get("data", []))
+
+    def query_calendar(
+        self,
+        start_date: Optional[str] = None,
+        end_date: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """查询 A 股交易日历。"""
+        params: Dict[str, Any] = {}
+        if start_date:
+            params["start_date"] = start_date
+        if end_date:
+            params["end_date"] = end_date
+        data = self._get("/api/v1/data/calendar", params)
+        return pd.DataFrame(data.get("data", []))
+
+    def query_manifest(
+        self,
+        category_id: str,
+        sub_category: str,
+        trade_date: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        """查询 COS 可下载文件清单。"""
+        params: Dict[str, Any] = {
+            "category_id": category_id,
+            "sub_category": sub_category,
+        }
+        if trade_date:
+            params["trade_date"] = trade_date
+        data = self._get("/api/v1/data/download/manifest", params)
+        return data.get("files", [])
+
+    def query_meta(
+        self,
+        dataset: Optional[str] = None,
+        source: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """查询数据元数据（各数据集日期范围/行数/可用性）。
+
+        数据来自 PostgreSQL data_meta 表（由每日流水线从 COS 重建）。
+        可按 dataset（kline/financial/basic/calendar/sector/tick）与
+        source（cos）过滤。
+        """
+        params: Dict[str, Any] = {}
+        if dataset:
+            params["dataset"] = dataset
+        if source:
+            params["source"] = source
+        data = self._get("/api/v1/data/meta", params)
+        return pd.DataFrame(data.get("data", []))
+
+    # ========== 数据预览与下载（消耗下载流量） ==========
+
+    def preview_as_df(
+        self,
+        category_id: str,
+        sub_category: str,
+        symbol: Optional[str] = None,
+        limit: int = 30,
+    ) -> pd.DataFrame:
+        """预览 Parquet 尾部 N 条数据（不消耗流量）。"""
+        params: Dict[str, Any] = {
+            "category_id": category_id,
+            "sub_category": sub_category,
+            "limit": limit,
+        }
+        if symbol:
+            params["symbol"] = symbol
+        data = self._get("/api/v1/data/preview", params)
+        return pd.DataFrame(data.get("data", []), columns=data.get("columns"))
+
+    def download_file(
+        self,
+        category_id: str,
+        sub_category: str,
+        symbol: Optional[str] = None,
+        save_dir: Optional[str] = None,
+        trade_date: Optional[str] = None,
+    ) -> str:
+        """流式下载原始 Parquet 切片到本地，返回保存路径（消耗下载流量）。
+
+        trade_date 仅对 tick_data 子分类有效（按交易日下载 {SYM}_{YYYYMMDD}.parquet）。
+        """
+        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
+
+        resp = self._request(
+            "GET", "/api/v1/data/download", params=params, stream=True
+        )
+        if resp.status_code != 200:
+            self._check_response(resp)
+
+        fallback = f"{sub_category}.parquet"
+        if symbol:
+            fallback = f"{sub_category}_{symbol}.parquet"
+        filename = parse_filename_from_content_disposition(
+            resp.headers.get("Content-Disposition", ""), fallback
+        )
+
+        save_path = os.path.join(save_dir, filename)
+        with open(save_path, "wb") as f:
+            for chunk in resp.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+        return os.path.abspath(save_path)
+
+    def load_as_df(
+        self,
+        category_id: str,
+        sub_category: str,
+        symbol: Optional[str] = None,
+        trade_date: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """将远端 Parquet 切片直接加载到内存 DataFrame（不落盘，消耗下载流量）。
+
+        trade_date 仅对 tick_data 子分类有效（按交易日下载 {SYM}_{YYYYMMDD}.parquet）。
+
+        带进程内 ETag 缓存：同一对象（ETag 未变）不重复下载，避免对同一 symbol
+        多次查询时重复消耗流量。
+        """
+        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
+        # 若已有缓存，带 If-None-Match 让服务端在对象未变时返回 304（不计流量下载 body）。
+        # 注：当前网关透传 COS ETag，但未实现 304；若服务端不支持 304，仍会返回 200
+        # 全量 body，此时用响应 ETag 命中本地缓存跳过重复解析。
+        cached = self._cache.get(cache_key)
+        headers = {}
+        if cached and cached.get("etag"):
+            headers["If-None-Match"] = cached["etag"]
+        resp = self._request("GET", "/api/v1/data/download", params=params, headers=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 ""
+        # 命中缓存（ETag 未变）：复用已解析的 df，不重复消耗解析
+        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
+
+    def 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 = self.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}'，已被拒绝。"
+                    "query_local 仅支持 SELECT 查询，不支持数据修改或联合查询。"
+                )
+
+        if "FROM" in sql_upper:
+            # 完整 SELECT：替换第一个 FROM 后的表名为 parquet 路径
+            sql = re.sub(
+                r"FROM\s+[`\"\']?\w+[`\"\']?",
+                f"FROM '{clean_path}'",
+                sql,
+                count=1,
+                flags=re.IGNORECASE,
+            )
+        else:
+            # 纯 WHERE 条件
+            sql = f"SELECT * FROM '{clean_path}' WHERE {sql}"
+        return duckdb.query(sql).df()
+
+    def get_local_warehouse(self, save_dir: Optional[str] = None) -> "DuckDBWarehouse":
+        """获取 DuckDB 本地数据仓库实例，用于离线多表 SQL JOIN 查询。"""
+        return DuckDBWarehouse(client=self, save_dir=save_dir)
+
+
+class DuckDBWarehouse:
+    """DuckDB 本地仓库管理，支持自动下载数据并注册为 DuckDB 视图，方便多表复杂 SQL 查询。"""
+
+    def __init__(self, client: QuantDBClient, save_dir: Optional[str] = None):
+        try:
+            import duckdb  # type: ignore
+        except ImportError as exc:
+            raise RuntimeError("DuckDBWarehouse 需要 duckdb：pip install duckdb") from exc
+
+        self.client = client
+        self.save_dir = save_dir or default_download_dir()
+        self.conn = duckdb.connect(database=":memory:")
+        self._views: Dict[str, str] = {}
+
+    def register_file(self, view_name: str, file_path: str) -> None:
+        """把指定的本地 Parquet 文件注册为 DuckDB 视图。"""
+        clean_path = file_path.replace("\\", "/")
+        self.conn.execute(f"CREATE OR REPLACE VIEW {view_name} AS SELECT * FROM '{clean_path}'")
+        self._views[view_name] = clean_path
+
+    def mount_dataset(
+        self,
+        category_id: str,
+        sub_category: str,
+        symbol: Optional[str] = None,
+        trade_date: Optional[str] = None,
+        view_name: Optional[str] = None,
+    ) -> str:
+        """自动下载或指定 Parquet 文件，并将其挂载为 DuckDB View。
+
+        Returns:
+            注册的视图名称。
+        """
+        file_path = self.client.download_file(
+            category_id=category_id,
+            sub_category=sub_category,
+            symbol=symbol,
+            trade_date=trade_date,
+            save_dir=self.save_dir,
+        )
+        if not view_name:
+            v_sym = f"_{symbol.replace('.', '_')}" if symbol else ""
+            v_td = f"_{trade_date}" if trade_date else ""
+            view_name = f"{sub_category}{v_sym}{v_td}"
+            # 清理 view_name 中的非法字符
+            view_name = re.sub(r"[^a-zA-Z0-9_]", "_", view_name)
+
+        self.register_file(view_name, file_path)
+        return view_name
+
+    def query(self, sql: str) -> pd.DataFrame:
+        """在所有已注册的视图上执行 SQL 查询，返回 DataFrame。"""
+        return self.conn.query(sql).df()
+
+    def list_views(self) -> Dict[str, str]:
+        """列出当前已注册的视图及其对应本地路径。"""
+        return dict(self._views)
+

quantdb_sdk/errors.py

@@ -0,0 +1,43 @@
+"""QuantDB SDK 异常类型。"""
+
+
+class QuantDBError(RuntimeError):
+    """QuantDB SDK 通用异常基类。"""
+
+    pass
+
+
+class AuthError(QuantDBError):
+    """认证失败：API Key 或 JWT Token 无效。"""
+
+    pass
+
+
+class InsufficientTrafficError(QuantDBError):
+    """无有效订阅或本月下载流量已用完。"""
+
+    pass
+
+
+class NotFoundError(QuantDBError):
+    """请求的数据或资源不存在。"""
+
+    pass
+
+
+class RateLimitError(QuantDBError):
+    """请求过于频繁，触发限流。"""
+
+    pass
+
+
+class ServerError(QuantDBError):
+    """服务端错误。"""
+
+    pass
+
+
+class ValidationError(QuantDBError):
+    """请求参数校验失败。"""
+
+    pass