289 lines
10 KiB
Python
289 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
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 NbClient:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
env_name: str = "nb-mcp",
|
|
data_source: str = "main",
|
|
nb_bin: str = "nb",
|
|
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_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) -> dict[str, Any]:
|
|
result = self._run_raw(["env", "status", self.env_name])
|
|
output = result.stdout.strip()
|
|
return {
|
|
"ok": result.returncode == 0 and "unreachable" not in output.lower(),
|
|
"env": self.env_name,
|
|
"status_output": output,
|
|
"stderr": self._redact(result.stderr.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 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.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.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:
|
|
raise NbCommandError(self._format_error(command, result))
|
|
try:
|
|
parsed = json.loads(result.stdout)
|
|
except json.JSONDecodeError as error:
|
|
raise NbCommandError(
|
|
f"nb returned invalid JSON: {error.msg}\nstdout={self._redact(result.stdout.strip())}"
|
|
) from error
|
|
if not isinstance(parsed, dict):
|
|
raise NbCommandError("nb returned JSON that is not an object")
|
|
return parsed
|
|
|
|
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())}"
|
|
)
|
|
|
|
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)])
|