389 lines
14 KiB
Python
389 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
from collections.abc import Set as AbstractSet
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Mapping, Sequence
|
|
|
|
|
|
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 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("<redacted>")
|
|
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, "<redacted>")
|
|
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")
|