yearning-cli 0.1.5__py3-none-any.whl

yearning_cli/client.py

@@ -0,0 +1,636 @@
+"""Yearning API client."""
+
+import json
+import os
+import sys
+import threading
+import time
+
+import msgpack
+import requests
+import websocket
+
+STATE_FILE = os.path.expanduser("~/.yearning_state")
+CONFIG_FILE = os.path.expanduser("~/.sqlrc")
+
+
+def _read_json_file(path):
+    try:
+        with open(path, encoding="utf-8") as f:
+            data = json.load(f)
+            return data if isinstance(data, dict) else {}
+    except (OSError, json.JSONDecodeError):
+        return {}
+
+
+def read_state():
+    """Load persisted CLI state."""
+    state = _read_json_file(STATE_FILE)
+    if not isinstance(state.get("query_orders"), dict):
+        state["query_orders"] = {}
+    return state
+
+
+def read_config():
+    """Load CLI config from ~/.sqlrc."""
+    cfg = {}
+    if not os.path.exists(CONFIG_FILE):
+        return cfg
+    try:
+        with open(CONFIG_FILE, encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith("#"):
+                    continue
+                if ":" in line:
+                    key, val = line.split(":", 1)
+                    cfg[key.strip()] = val.strip()
+    except OSError:
+        return {}
+    return cfg
+
+
+def write_state(state):
+    """Persist CLI state atomically."""
+    tmp_file = f"{STATE_FILE}.{os.getpid()}.{threading.get_ident()}.tmp"
+    with open(tmp_file, "w", encoding="utf-8") as f:
+        json.dump(state, f)
+    os.chmod(tmp_file, 0o600)
+    os.replace(tmp_file, STATE_FILE)
+
+
+def _log(msg, verbose=False):
+    """Print debug message to stderr if verbose mode is on."""
+    if verbose:
+        print(f"[DEBUG] {msg}", file=sys.stderr, flush=True)
+
+
+class TokenExpiredError(RuntimeError):
+    """Raised when the persisted Yearning token is no longer accepted."""
+
+
+AUTH_EXPIRED_MESSAGE = "本地 token 已过期或无效，请重新运行: sql login"
+AUTO_LOGIN_UNAVAILABLE_MESSAGE = (
+    "本地 token 已过期或无效，且缺少自动登录凭据。"
+    "请设置 YEARNING_USER/YEARNING_PASS，或在 ~/.sqlrc 配置 yearning_user/yearning_pass"
+)
+AUTH_EXPIRED_KEYWORDS = (
+    "token",
+    "jwt",
+    "expired",
+    "expire",
+    "unauthorized",
+    "forbidden",
+    "未登录",
+    "登录",
+    "过期",
+    "无效",
+    "认证",
+    "鉴权",
+    "授权",
+)
+
+
+def _looks_like_auth_error(value):
+    if isinstance(value, dict):
+        value = value.get("text") or value.get("error") or value.get("message") or value.get("msg") or ""
+    if not isinstance(value, str):
+        return False
+    text = str(value).lower()
+    return any(keyword.lower() in text for keyword in AUTH_EXPIRED_KEYWORDS)
+
+
+class YearningClient:
+    _login_lock = threading.Lock()
+
+    def __init__(
+        self,
+        base_url="http://192.168.1.135",
+        verbose=False,
+        username="",
+        password="",
+        auto_login=True,
+        max_auto_login_attempts=2,
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.api_base = f"{self.base_url}/api/v2"
+        self.verbose = verbose
+        self.session = requests.Session()
+        self.username = username
+        self.password = password
+        self.auto_login = auto_login
+        self.max_auto_login_attempts = max_auto_login_attempts
+        self.token = self._load_token()
+
+    def _log(self, msg):
+        _log(msg, self.verbose)
+
+    def _load_token(self):
+        data = read_state()
+        token = data.get("token", "")
+        if token:
+            self._log(f"Loaded token from {STATE_FILE}")
+            return token
+        return None
+
+    def _save_token(self, token, username):
+        state = read_state()
+        state["token"] = token
+        state["username"] = username
+        write_state(state)
+
+    def _auth_credentials(self):
+        cfg = read_config()
+        state = read_state()
+        username = (
+            self.username
+            or os.environ.get("YEARNING_USER", "")
+            or cfg.get("yearning_user", "")
+            or cfg.get("username", "")
+            or cfg.get("user", "")
+            or state.get("username", "")
+        )
+        password = (
+            self.password
+            or os.environ.get("YEARNING_PASS", "")
+            or cfg.get("yearning_pass", "")
+            or cfg.get("yearning_password", "")
+            or cfg.get("password", "")
+            or cfg.get("pass", "")
+        )
+        return username, password
+
+    def _adopt_newer_token_from_state(self):
+        token = self._load_token()
+        if token and token != self.token:
+            self._log("Adopted refreshed token from state")
+            self.token = token
+            return True
+        return False
+
+    def _login_for_retry(self):
+        if not self.auto_login:
+            raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+
+        username, password = self._auth_credentials()
+        if not username or not password:
+            raise TokenExpiredError(AUTO_LOGIN_UNAVAILABLE_MESSAGE)
+
+        with self._login_lock:
+            if self._adopt_newer_token_from_state():
+                return
+            self._log("Auto login with configured credentials")
+            self.login(username, password)
+
+    def _with_auto_login(self, label, operation):
+        login_attempts = 0
+        adopted_state_token = False
+        last_error = None
+
+        def login_until_success():
+            nonlocal login_attempts, last_error
+            while login_attempts < self.max_auto_login_attempts:
+                login_attempts += 1
+                self._log(f"{label} token invalid, auto login attempt {login_attempts}")
+                try:
+                    self._login_for_retry()
+                    return
+                except TokenExpiredError as login_error:
+                    last_error = login_error
+                    if str(login_error) == AUTO_LOGIN_UNAVAILABLE_MESSAGE:
+                        raise
+                except Exception as login_error:
+                    last_error = login_error
+
+            raise TokenExpiredError(
+                f"自动登录失败（已尝试 {login_attempts} 次）: {last_error}"
+            ) from last_error
+
+        while True:
+            if not self.token:
+                login_until_success()
+            try:
+                return operation()
+            except TokenExpiredError as e:
+                last_error = e
+                if not adopted_state_token and self._adopt_newer_token_from_state():
+                    adopted_state_token = True
+                    continue
+                login_until_success()
+
+    def ensure_authenticated(self):
+        """Ensure a token exists, using configured credentials when needed."""
+        return self._with_auto_login("认证检查", lambda: True)
+
+    def login(self, username, password):
+        """Login and save token."""
+        self._log(f"Logging in as {username} to {self.base_url}")
+        resp = self.session.post(
+            f"{self.base_url}/login",
+            json={"username": username, "password": password},
+            timeout=10,
+        )
+        self._log(f"Login response status: {resp.status_code}")
+        data = resp.json()
+        self._log(f"Login response: {json.dumps(data, ensure_ascii=False)[:200]}")
+        if data.get("code") != 1200:
+            raise ValueError(f"登录失败: {data.get('text', data.get('error', '未知错误'))}")
+        payload = data.get("payload", {})
+        token = payload.get("token", "")
+        if not token:
+            raise ValueError(f"登录响应中未找到 token: {data}")
+        self._save_token(token, username)
+        self.username = username
+        self.password = password
+        self.token = token
+        self._log("Token saved successfully")
+        return token
+
+    def _headers(self):
+        return {"Authorization": f"Bearer {self.token}"}
+
+    def _json_response(self, resp, label):
+        if resp.status_code in (401, 403):
+            raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+
+        try:
+            data = resp.json()
+        except ValueError as e:
+            body = resp.text.strip()
+            if _looks_like_auth_error(body):
+                raise TokenExpiredError(AUTH_EXPIRED_MESSAGE) from e
+            raise ValueError(f"{label} 响应不是有效 JSON: {body[:200]}") from e
+
+        if isinstance(data, dict):
+            if data.get("code") in (401, 403):
+                raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+            text = data.get("text") or data.get("error") or data.get("message") or data.get("msg")
+            if text and _looks_like_auth_error(text):
+                raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+            return data
+
+        if _looks_like_auth_error(data):
+            raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+
+        raise ValueError(f"{label} 响应格式异常: {data!r}")
+
+    def list_sources(self, idc=""):
+        return self._with_auto_login(
+            "数据源列表",
+            lambda: self._list_sources_once(idc),
+        )
+
+    def _list_sources_once(self, idc=""):
+        """List data sources. idc can be empty to list all."""
+        self._log(f"Fetching sources: tp=query, idc={idc}")
+        resp = self.session.get(
+            f"{self.api_base}/fetch/source",
+            headers=self._headers(),
+            params={"tp": "query", "idc": idc},
+            timeout=10,
+        )
+        data = self._json_response(resp, "数据源列表")
+        self._log(f"Sources response code: {data.get('code')}")
+        payload = data.get("payload", [])
+        if _looks_like_auth_error(payload):
+            raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+        if not isinstance(payload, list):
+            raise ValueError(f"数据源列表响应格式异常: {payload!r}")
+        return payload
+
+    def list_databases(self, source_id):
+        return self._with_auto_login(
+            "数据库列表",
+            lambda: self._list_databases_once(source_id),
+        )
+
+    def _list_databases_once(self, source_id):
+        """List databases for a source."""
+        self._log(f"Fetching databases for source_id={source_id}")
+        resp = self.session.get(
+            f"{self.api_base}/query/schema",
+            headers=self._headers(),
+            params={"source_id": source_id},
+            timeout=10,
+        )
+        data = self._json_response(resp, "数据库列表")
+        self._log(f"Databases response code: {data.get('code')}")
+        payload = data.get("payload") or {}
+        if not isinstance(payload, dict):
+            if _looks_like_auth_error(payload):
+                raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+            raise ValueError(f"数据库列表响应格式异常: {payload!r}")
+        info = payload.get("info", [])
+        if not isinstance(info, list):
+            raise ValueError(f"数据库列表响应格式异常: {info!r}")
+        return [
+            db.get("title", db.get("name", "")) if isinstance(db, dict) else str(db)
+            for db in info
+        ]
+
+    def list_tables(self, source_id, database):
+        return self._with_auto_login(
+            "表列表",
+            lambda: self._list_tables_once(source_id, database),
+        )
+
+    def _list_tables_once(self, source_id, database):
+        """List tables for a database."""
+        self._log(f"Fetching tables for source_id={source_id}, db={database}")
+        resp = self.session.get(
+            f"{self.api_base}/query/tables",
+            headers=self._headers(),
+            params={"source_id": source_id, "schema": database},
+            timeout=10,
+        )
+        data = self._json_response(resp, "表列表")
+        self._log(f"Tables response code: {data.get('code')}")
+        payload = data.get("payload") or {}
+        if not isinstance(payload, dict):
+            if _looks_like_auth_error(payload):
+                raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+            raise ValueError(f"表列表响应格式异常: {payload!r}")
+        tables = payload.get("table", [])
+        if not isinstance(tables, list):
+            raise ValueError(f"表列表响应格式异常: {tables!r}")
+        return [
+            t.get("title", t.get("name", "")) if isinstance(t, dict) else str(t)
+            for t in tables
+        ]
+
+    def _is_query_order_ready(self):
+        """Check whether the current query order is usable."""
+        try:
+            resp = self.session.get(
+                f"{self.api_base}/fetch/query_status",
+                headers=self._headers(),
+                timeout=10,
+            )
+            data = self._json_response(resp, "查询状态")
+            self._log(f"Query status: {data}")
+            if data.get("payload") is True:
+                return True
+        except TokenExpiredError:
+            raise
+        except Exception as e:
+            self._log(f"Query status check failed: {e}")
+
+        try:
+            resp = self.session.get(
+                f"{self.api_base}/fetch/is_query",
+                headers=self._headers(),
+                timeout=10,
+            )
+            data = self._json_response(resp, "查询权限状态")
+            self._log(f"Is query: {data}")
+            payload = data.get("payload", {})
+            if not isinstance(payload, dict):
+                if _looks_like_auth_error(payload):
+                    raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+                return False
+            if payload.get("status") is True:
+                return True
+        except TokenExpiredError:
+            raise
+        except Exception as e:
+            self._log(f"Is query check failed: {e}")
+
+        return False
+
+    def create_query_order(self, source_id, text="CLI query", verify_wait=0.5):
+        return self._with_auto_login(
+            "创建查询工单",
+            lambda: self._create_query_order_once(source_id, text, verify_wait),
+        )
+
+    def _create_query_order_once(self, source_id, text="CLI query", verify_wait=0.5):
+        """Create a query order. Returns True if order is ready for query."""
+        self._log(f"Creating query order for source_id={source_id}")
+        # First undo any existing orders
+        try:
+            resp = self.session.delete(
+                f"{self.api_base}/query/undo",
+                headers=self._headers(),
+                timeout=10,
+            )
+            self._log(f"Undo response: {resp.text[:200]}")
+            if resp.status_code in (401, 403) or _looks_like_auth_error(resp.text):
+                raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+        except TokenExpiredError:
+            raise
+        except Exception as e:
+            self._log(f"Undo failed (may be ok): {e}")
+
+        # Create new order
+        resp = self.session.post(
+            f"{self.api_base}/query/post",
+            headers=self._headers(),
+            json={"source_id": source_id, "text": text, "export": 0},
+            timeout=10,
+        )
+        self._log(f"Create order response: status={resp.status_code}, body='{resp.text[:200]}'")
+        if resp.status_code in (401, 403) or _looks_like_auth_error(resp.text):
+            raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+
+        # If response is empty with 200, likely auto-approve mode (ReferQueryOrder returns without JSON)
+        # If response contains ORDER_IS_CREATE, needs manual approval
+        body = resp.text.strip().strip('"')
+        if resp.status_code == 200 and body == "":
+            self._log("Empty response from create, assuming auto-approved")
+            return True
+
+        if resp.status_code == 200 and (body == "" or "已创建" in body):
+            if self._is_query_order_ready():
+                return True
+            time.sleep(verify_wait)
+            return self._is_query_order_ready()
+
+        return False
+
+    def execute_query(self, source_id, database, sql, timeout=30):
+        return self._with_auto_login(
+            "执行查询",
+            lambda: self._execute_query_once(source_id, database, sql, timeout),
+        )
+
+    def _execute_query_once(self, source_id, database, sql, timeout=30):
+        """Execute SQL query via WebSocket and return results.
+
+        Returns dict with:
+            - success: bool
+            - columns: list of column names
+            - rows: list of row dicts
+            - query_time: ms
+            - error: error message if failed
+        """
+        scheme = "wss" if self.base_url.startswith("https") else "ws"
+        host = self.base_url.replace("http://", "").replace("https://", "").rstrip("/")
+        ws_url = f"{scheme}://{host}{self.api_base.replace(self.base_url, '')}/query/results?source_id={source_id}"
+
+        self._log(f"WebSocket URL: {ws_url}")
+        self._log(f"Token: {self.token[:20]}... (len={len(self.token)})")
+        self._log(f"Database: {database}")
+        self._log(f"SQL: {sql}")
+
+        result = {"success": False, "columns": [], "rows": [], "query_time": 0, "error": ""}
+        received = threading.Event()
+        message_count = [0]
+
+        def on_message(ws, message):
+            message_count[0] += 1
+            try:
+                data = msgpack.unpackb(message, raw=False)
+                self._log(f"Message #{message_count[0]}: {json.dumps(data, ensure_ascii=False, default=str)[:500]}")
+
+                if not isinstance(data, dict):
+                    if _looks_like_auth_error(data):
+                        result["error"] = AUTH_EXPIRED_MESSAGE
+                    else:
+                        result["error"] = f"响应格式异常: {data!r}"
+                    result["success"] = False
+                    received.set()
+                    return
+
+                # Check for heartbeat first
+                heartbeat = data.get("heartbeat") or data.get("HeartBeat")
+                if heartbeat:
+                    # Respond to server heartbeat PING with PONG to keep connection alive
+                    if heartbeat == 1 or str(heartbeat).upper() == "PING":
+                        self._log("Received heartbeat PING, sending PONG")
+                        try:
+                            pong_msg = msgpack.packb({"heartbeat": "PONG"})
+                            ws.send(pong_msg, opcode=websocket.ABNF.OPCODE_BINARY)
+                        except Exception as e:
+                            self._log(f"Failed to send PONG: {e}")
+                    return
+
+                error = data.get("error", "")
+                if error:
+                    result["error"] = AUTH_EXPIRED_MESSAGE if _looks_like_auth_error(error) else error
+                    result["success"] = False
+                    received.set()
+                    return
+
+                # Yearning source code logic (from queryResults struct):
+                # - Success: {Export, Error:"", Results, QueryTime} — Status defaults to false (zero value)
+                # - No valid order: {Status: true} — sent when no approved order found or order expired
+                # - Heartbeat: {HeartBeat: PING, IsOnly} — Status defaults to false
+                # So status=False + results = SUCCESS, status=True + no results = no valid order
+                status = data.get("status")
+                results = data.get("results") or data.get("Results")
+
+                if results:
+                    # Query success (status=False is the default for data responses)
+                    first = results[0]
+                    fields = first.get("field") or first.get("Field", [])
+                    if fields and isinstance(fields[0], dict):
+                        result["columns"] = [
+                            f.get("title", f.get("dataIndex", "?")) for f in fields
+                        ]
+                    else:
+                        result["columns"] = [str(f) for f in fields]
+                    raw_data = first.get("data") or first.get("Data", [])
+                    if raw_data and isinstance(raw_data[0], dict):
+                        result["rows"] = [
+                            [row.get(col, "") for col in result["columns"]]
+                            for row in raw_data
+                        ]
+                    else:
+                        result["rows"] = raw_data
+                    result["query_time"] = data.get("query_time") or data.get("QueryTime", 0)
+                    result["export"] = data.get("export", False)
+                    result["success"] = True
+                    received.set()
+                    return
+
+                if status is True:
+                    # Informational: no valid order / order expired / server ready
+                    # Do NOT overwrite existing successful results
+                    if not result["success"]:
+                        result["error"] = "查询工单已过期或不存在，请先在 Web 界面提交查询工单并通过审核"
+                        result["success"] = False
+                        received.set()
+                    return
+
+                # Heartbeat or other unknown message (status=False, no results, no error)
+                self._log(f"Unhandled message, keys: {list(data.keys())}")
+
+            except Exception as e:
+                result["error"] = f"解析响应失败: {e}"
+                received.set()
+
+        def on_error(ws, error):
+            self._log(f"WebSocket error: {error}")
+            if _looks_like_auth_error(error):
+                result["error"] = AUTH_EXPIRED_MESSAGE
+            else:
+                result["error"] = f"WebSocket 错误: {error}"
+            received.set()
+
+        def on_close(ws, close_status_code, close_msg):
+            self._log(f"WebSocket closed: code={close_status_code}, msg={close_msg}")
+            if not received.is_set() and not result["success"]:
+                if close_status_code in (401, 403) or _looks_like_auth_error(close_msg):
+                    result["error"] = AUTH_EXPIRED_MESSAGE
+                elif close_status_code is None:
+                    # Server dropped TCP connection without close frame
+                    # Typical causes: server-side query timeout, reverse proxy timeout,
+                    # or heartbeat keepalive failure
+                    result["error"] = (
+                        "连接被服务端断开（无关闭帧）。"
+                        "可能原因：查询超时（Yearning 默认 60 秒）、反向代理超时、或心跳保活失败"
+                    )
+                else:
+                    result["error"] = f"连接关闭: code={close_status_code}, msg={close_msg}"
+                received.set()
+
+        def on_open(ws):
+            self._log("WebSocket connected")
+            # Send init (type=2) to establish DB connection
+            init_msg = msgpack.packb({"type": 2})
+            ws.send(init_msg, opcode=websocket.ABNF.OPCODE_BINARY)
+            self._log(f"Sent init (type=2): {init_msg.hex()}")
+
+            # Send query (type=4) immediately after init
+            # The server processes messages sequentially, so this will be handled after init
+            query_msg = msgpack.packb({
+                "type": 4,
+                "sql": sql,
+                "schema": database,
+            })
+            self._log(f"Sending query (type=4): {query_msg.hex()}")
+            ws.send(query_msg, opcode=websocket.ABNF.OPCODE_BINARY)
+            self._log("Query sent successfully")
+
+        self._log("Connecting to WebSocket...")
+        ws = websocket.WebSocketApp(
+            ws_url,
+            subprotocols=[self.token],
+            on_open=on_open,
+            on_message=on_message,
+            on_error=on_error,
+            on_close=on_close,
+        )
+
+        ws_thread = threading.Thread(
+            target=ws.run_forever,
+            kwargs={
+                # WebSocket-level ping keepalive: send PING every 30s
+                # to prevent reverse proxy / load balancer from dropping idle connections
+                "ping_interval": 30,
+                "ping_timeout": 10,
+            },
+            daemon=True,
+        )
+        ws_thread.start()
+
+        self._log(f"Waiting for response (timeout={timeout}s)...")
+        if not received.wait(timeout=timeout):
+            self._log("Timeout waiting for response")
+            result["error"] = f"查询超时 ({timeout}s)"
+
+        ws.close()
+        self._log(f"Final result: success={result['success']}, error={result['error']}, rows={len(result['rows'])}")
+        if not result["success"] and _looks_like_auth_error(result.get("error", "")):
+            raise TokenExpiredError(AUTH_EXPIRED_MESSAGE)
+        return result