from __future__ import annotations import json import os import re import subprocess from collections.abc import Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from pathlib import Path from typing import Any Runner = Callable[..., subprocess.CompletedProcess[str]] class NbCommandError(RuntimeError): """Raised when the NocoBase CLI command fails or returns invalid output.""" class NbAuthenticationError(NbCommandError): """Raised when NocoBase rejects the configured API credential.""" class NbClient: def __init__( self, *, env_name: str = "nb-mcp", data_source: str = "main", nb_bin: str = "nb", api_base_url: str | None = None, api_token: str | None = None, role: str | None = None, cwd: str | Path | None = None, timeout: int = 30, runner: Runner | None = None, ) -> None: self.env_name = env_name self.data_source = data_source expanded_nb_bin = Path(nb_bin).expanduser() self.nb_bin = str(expanded_nb_bin) if str(nb_bin).startswith("~") or expanded_nb_bin.is_absolute() else nb_bin self.api_base_url = api_base_url.strip().rstrip("/") if api_base_url else None self.api_token = api_token self.role = role self.cwd = Path(cwd).resolve() if cwd else None self.timeout = timeout self.runner = runner or subprocess.run def health(self, resource: str, id_field: str) -> dict[str, Any]: self.list_records(resource, fields=[id_field], page=1, page_size=1) return { "ok": True, "env": self.env_name, "status_output": "authenticated", "stderr": "", } def environment_credentials(self) -> tuple[str, str]: """Return the selected CLI environment's API URL and current user token. The token is read from ``nb env info`` output and is never added to a subprocess command line. This is intended for the few multipart or relation endpoints that the CLI does not expose as first-class commands. """ command = [ self.nb_bin, "env", "info", self.env_name, "--json", "--show-secrets", ] result = self._run_raw(command[1:], full_command=command) if result.returncode != 0: raise NbCommandError( f"Could not read nb environment {self.env_name!r} (exit code {result.returncode}). " f"Run `nb env info {self.env_name}` to verify it." ) try: payload = self._parse_json_object(result.stdout) except (json.JSONDecodeError, NbCommandError) as error: raise NbCommandError(f"nb env info returned invalid JSON for environment {self.env_name!r}") from error api = payload.get("api") auth = api.get("auth") if isinstance(api, Mapping) else None api_base_url = api.get("apiBaseUrl") if isinstance(api, Mapping) else None access_token = auth.get("accessToken") if isinstance(auth, Mapping) else None if not isinstance(api_base_url, str) or not api_base_url.strip(): raise NbCommandError(f"nb environment {self.env_name!r} has no API base URL") if not isinstance(access_token, str) or not access_token.strip() or access_token in {"-", "********"}: raise NbAuthenticationError( f"nb environment {self.env_name!r} is not authenticated. " "Ask the local installation Agent to reconnect it with your personal NocoBase API Key." ) return api_base_url.strip().rstrip("/"), access_token.strip() def list_records( self, resource: str, *, filter: Mapping[str, Any] | None = None, fields: Sequence[str] | None = None, appends: Sequence[str] | None = None, except_fields: Sequence[str] | None = None, sort: Sequence[str] | None = None, page: int | None = None, page_size: int | None = None, paginate: bool | None = None, ) -> dict[str, Any]: command = self._resource_command("list", resource) self._append_json(command, "--filter", filter) self._append_repeated(command, "--fields", fields) self._append_repeated(command, "--appends", appends) self._append_repeated(command, "--except", except_fields) self._append_repeated(command, "--sort", sort) self._append_value(command, "--page", page) self._append_value(command, "--page-size", page_size) if paginate is True: command.append("--paginate") elif paginate is False: command.append("--no-paginate") return self._run_json(command) def get_record( self, resource: str, filter_by_tk: Any, *, fields: Sequence[str] | None = None, appends: Sequence[str] | None = None, except_fields: Sequence[str] | None = None, ) -> dict[str, Any]: command = self._resource_command("get", resource) self._append_json_or_scalar(command, "--filter-by-tk", filter_by_tk) self._append_repeated(command, "--fields", fields) self._append_repeated(command, "--appends", appends) self._append_repeated(command, "--except", except_fields) return self._run_json(command) def create_record( self, resource: str, values: Mapping[str, Any], whitelist: Sequence[str], ) -> dict[str, Any]: self._validate_scalar_values(values) command = self._resource_command("create", resource) self._append_json(command, "--values", values) self._append_repeated(command, "--whitelist", whitelist) return self._run_json(command) def update_record( self, resource: str, record_id: Any, values: Mapping[str, Any], whitelist: Sequence[str], ) -> dict[str, Any]: self._validate_scalar_values(values) command = self._resource_command("update", resource) self._append_json_or_scalar(command, "--filter-by-tk", record_id) self._append_json(command, "--values", values) self._append_repeated(command, "--whitelist", whitelist) return self._run_json(command) def query_records( self, resource: str, *, measures: Sequence[Mapping[str, Any]] | None = None, dimensions: Sequence[Mapping[str, Any]] | None = None, orders: Sequence[Mapping[str, Any]] | None = None, filter: Mapping[str, Any] | None = None, having: Mapping[str, Any] | None = None, limit: int | None = None, offset: int | None = None, timezone: str | None = None, ) -> dict[str, Any]: command = self._resource_command("query", resource) self._append_json(command, "--measures", measures) self._append_json(command, "--dimensions", dimensions) self._append_json(command, "--orders", orders) self._append_json(command, "--filter", filter) self._append_json(command, "--having", having) self._append_value(command, "--limit", limit) self._append_value(command, "--offset", offset) self._append_value(command, "--timezone", timezone) return self._run_json(command) def get_collection( self, collection_name: str, *, appends: Sequence[str] | None = None, except_fields: Sequence[str] | None = None, ) -> dict[str, Any]: command = self._data_modeling_command("collections", "get") command.extend(["--filter-by-tk", collection_name]) self._append_repeated(command, "--appends", appends) self._append_repeated(command, "--except", except_fields) return self._run_json(command) def list_collection_fields( self, collection_name: str, *, fields: Sequence[str] | None = None, appends: Sequence[str] | None = None, except_fields: Sequence[str] | None = None, page: int | None = None, page_size: int | None = None, paginate: bool | None = None, ) -> dict[str, Any]: command = self._data_modeling_command("collections", "fields", "list") command.extend(["--collection-name", collection_name]) self._append_repeated(command, "--fields", fields) self._append_repeated(command, "--appends", appends) self._append_repeated(command, "--except", except_fields) self._append_value(command, "--page", page) self._append_value(command, "--page-size", page_size) if paginate is True: command.append("--paginate") elif paginate is False: command.append("--no-paginate") return self._run_json(command) def _resource_command(self, action: str, resource: str) -> list[str]: command = [ self.nb_bin, "api", "resource", action, "--resource", resource, "--env", self.env_name, "--yes", "--json-output", ] if self.data_source: command.extend(["--data-source", self.data_source]) if self.api_base_url: command.extend(["--api-base-url", self.api_base_url]) if self.role: command.extend(["--role", self.role]) if self.api_token: command.extend(["--token", self.api_token]) return command def _data_modeling_command(self, *parts: str) -> list[str]: command = [ self.nb_bin, "api", "data-modeling", *parts, "--env", self.env_name, "--yes", "--json-output", ] if self.api_base_url: command.extend(["--api-base-url", self.api_base_url]) if self.role: command.extend(["--role", self.role]) if self.api_token: command.extend(["--token", self.api_token]) return command def _run_json(self, command: list[str]) -> dict[str, Any]: result = self._run_raw(command[1:], full_command=command) if result.returncode != 0: error_type = NbAuthenticationError if self._is_authentication_failure(result) else NbCommandError raise error_type(self._format_error(command, result)) try: parsed = self._parse_json_object(result.stdout) except json.JSONDecodeError as error: raise NbCommandError( f"nb returned invalid JSON: {error.msg}\nstdout={self._redact(result.stdout.strip())}" ) from error return parsed @staticmethod def _parse_json_object(output: str) -> dict[str, Any]: stripped = output.strip() try: parsed = json.loads(stripped) except json.JSONDecodeError: parsed = None if parsed is not None: if not isinstance(parsed, dict): raise NbCommandError("nb returned JSON that is not an object") return parsed decoder = json.JSONDecoder() candidates: list[tuple[int, int, dict[str, Any]]] = [] for start, character in enumerate(output): if character != "{": continue try: value, end = decoder.raw_decode(output, start) except json.JSONDecodeError: continue if isinstance(value, dict): candidates.append((start, end, value)) outer_candidates = [ candidate for candidate in candidates if not any( other_start <= candidate[0] and other_end >= candidate[1] and (other_start, other_end) != (candidate[0], candidate[1]) for other_start, other_end, _ in candidates ) ] if len(outer_candidates) != 1: raise json.JSONDecodeError("could not locate one complete JSON object", output, 0) return outer_candidates[0][2] def _run_raw( self, args: list[str], *, full_command: list[str] | None = None, ) -> subprocess.CompletedProcess[str]: command = full_command or [self.nb_bin, *args] env = os.environ.copy() env["NB_CLI_LOG_DISABLED"] = "1" env["NB_SKIP_STARTUP_UPDATE"] = "1" self._prepend_nb_bin_dir_to_path(env) return self.runner( command, cwd=str(self.cwd) if self.cwd else None, env=env, timeout=self.timeout, capture_output=True, text=True, check=False, ) def _prepend_nb_bin_dir_to_path(self, env: dict[str, str]) -> None: nb_path = Path(self.nb_bin).expanduser() if not nb_path.is_absolute(): return nb_dir = str(nb_path.parent) existing_path = env.get("PATH", "") parts = [part for part in existing_path.split(os.pathsep) if part] if nb_dir in parts: return env["PATH"] = nb_dir + (os.pathsep + existing_path if existing_path else "") def _format_error(self, command: list[str], result: subprocess.CompletedProcess[str]) -> str: return ( f"nb command failed with exit code {result.returncode}: {self._redact_command(command)}\n" f"stdout={self._redact(result.stdout.strip())}\n" f"stderr={self._redact(result.stderr.strip())}" ) @staticmethod def _is_authentication_failure(result: subprocess.CompletedProcess[str]) -> bool: output = f"{result.stdout}\n{result.stderr}" if re.search(r"\bstatus\s+401\b", output, flags=re.IGNORECASE): return True upper_output = output.upper() return any( code in upper_output for code in ( "EMPTY_TOKEN", "INVALID_TOKEN", "EXPIRED_TOKEN", "BLOCKED_TOKEN", "EXPIRED_SESSION", "NOT_EXIST_USER", ) ) def _redact_command(self, command: Sequence[str]) -> str: redacted: list[str] = [] skip_next = False for token in command: if skip_next: redacted.append("") skip_next = False continue redacted.append(token) if token == "--token": skip_next = True return " ".join(redacted) def _redact(self, value: str) -> str: redacted = value if self.api_token: redacted = redacted.replace(self.api_token, "") return redacted @staticmethod def _append_json(command: list[str], flag: str, value: Any) -> None: if value is None or value == {}: return command.extend([flag, json.dumps(value, ensure_ascii=False, separators=(",", ":"))]) @staticmethod def _append_json_or_scalar(command: list[str], flag: str, value: Any) -> None: if value is None: return if isinstance(value, (list, dict)): command.extend([flag, json.dumps(value, ensure_ascii=False, separators=(",", ":"))]) else: command.extend([flag, str(value)]) @staticmethod def _append_repeated(command: list[str], flag: str, values: Sequence[str] | None) -> None: for value in values or []: command.extend([flag, str(value)]) @staticmethod def _append_value(command: list[str], flag: str, value: Any) -> None: if value is None or value == "": return command.extend([flag, str(value)]) @staticmethod def _validate_scalar_values(values: Mapping[str, Any]) -> None: if not values: raise ValueError("values must not be empty") for field, value in values.items(): if isinstance(value, (Mapping, list, AbstractSet, tuple)): raise ValueError(f"values[{field!r}] must be scalar")