"""Command line interface for the MXAccess Gateway Python client.""" from __future__ import annotations import asyncio import contextlib import io import json import logging import os import sys import time from collections.abc import Awaitable, Callable from datetime import datetime, timezone from typing import Any import click from google.protobuf.json_format import MessageToDict from zb_mom_ww_mxgateway import __version__ from zb_mom_ww_mxgateway.auth import redact_secret from zb_mom_ww_mxgateway.client import GatewayClient from zb_mom_ww_mxgateway.errors import MxGatewayError from zb_mom_ww_mxgateway.galaxy import GalaxyRepositoryClient from zb_mom_ww_mxgateway.generated import galaxy_repository_pb2 as galaxy_pb from zb_mom_ww_mxgateway.generated import mxaccess_gateway_pb2 as pb from zb_mom_ww_mxgateway.options import ClientOptions from zb_mom_ww_mxgateway.values import MxValueInput, to_mx_value logger = logging.getLogger(__name__) MAX_AGGREGATE_EVENTS = 10_000 _BATCH_EOR = "__MXGW_BATCH_EOR__" @click.group() def main() -> None: """MXAccess Gateway Python test CLI.""" @main.command() @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def version(output_json: bool) -> None: """Print client package version information.""" payload = { "client": "mxgw-py", "package": "mxaccess-gateway-client", "version": __version__, } _emit(payload, output_json=output_json, text=f"mxgw-py {__version__}") @main.command() def batch() -> None: """Read commands from stdin and execute each, writing output + __MXGW_BATCH_EOR__ after each. Each non-empty line of stdin is a complete argument string (no quoting support — the harness never passes whitespace-containing arguments). Lines are split on runs of ASCII whitespace and dispatched through the normal CLI parser. On EOF or an empty line, exit 0. Errors do NOT terminate the loop. Each command's output (including any error JSON) is written to stdout followed by a line containing exactly ``__MXGW_BATCH_EOR__``, then stdout is flushed. Error output is formatted as ``{"error": "...", "type": "..."}``. Recursive ``batch`` lines are rejected (Client.Python-024) — re-entering the batch dispatcher would silently spawn a nested loop reading from the same exhausted stdin. """ for raw_line in sys.stdin: line = raw_line.rstrip("\n").rstrip("\r") if not line: # Empty line signals clean exit (matches the spec and .NET behaviour). break args = line.split() # Reject a recursive `batch` line outright: the nested invocation would # read from the already-exhausted stdin (or, depending on harness, the # same stream the outer batch is consuming line-by-line) and silently # exit. Surface it as an explicit error block so callers can audit the # mis-routed line. if args and args[0] == "batch": _batch_write_error( "RecursiveBatchError", "nested 'batch' invocation is not allowed inside batch mode", ) _batch_flush_eor() continue _dispatch_batch_line(args) def _dispatch_batch_line(args: list[str]) -> None: """Run a single batch line through the Click parser directly (no CliRunner). Captures the subcommand's stdout via :func:`contextlib.redirect_stdout` and synthesises the standard ``{"error": ..., "type": ...}`` shape on failure. Click exceptions (`ClickException`, `UsageError`) are caught and rendered; `SystemExit(0)` from a Click command is treated as a clean exit, while a non-zero `SystemExit` is rendered as a CLI error. All other exceptions are captured and rendered as `{"error": str(exc), "type": exc.__class__.__name__}` so the loop never dies. """ buffer = io.StringIO() exit_code = 0 exc: BaseException | None = None try: with contextlib.redirect_stdout(buffer): try: # `standalone_mode=False` makes Click raise instead of calling # `sys.exit`; we still need to handle SystemExit because some # commands explicitly raise it (or `click.UsageError` converts # to a SystemExit under some entry-point paths). main.main(args=args, standalone_mode=False, prog_name="mxgw-py") except click.exceptions.Exit as click_exit: exit_code = click_exit.exit_code except click.ClickException as click_exc: exit_code = click_exc.exit_code exc = click_exc click.echo(f"Error: {click_exc.format_message()}", err=False) except SystemExit as sys_exit: code = sys_exit.code exit_code = int(code) if isinstance(code, int) else (0 if code is None else 1) except Exception as captured: # noqa: BLE001 — never let batch loop die exc = captured exit_code = 1 output = buffer.getvalue() if exit_code == 0 and exc is None: sys.stdout.write(output) else: if output.lstrip().startswith("{"): # Inner command already emitted JSON (e.g. a structured error) — # relay verbatim. sys.stdout.write(output) if output and not output.endswith("\n"): sys.stdout.write("\n") elif exc is not None: _batch_write_error(type(exc).__name__, str(exc)) else: msg = output.strip() if msg.startswith("Error: "): msg = msg[len("Error: "):] _batch_write_error("CliError", msg) _batch_flush_eor() def _batch_write_error(exc_type: str, message: str) -> None: """Write a JSON error record to stdout in the standard batch error shape.""" sys.stdout.write(json.dumps({"error": message, "type": exc_type}) + "\n") def _batch_flush_eor() -> None: """Write the end-of-record sentinel and flush stdout.""" sys.stdout.write(_BATCH_EOR + "\n") sys.stdout.flush() def gateway_options(command: Callable[..., Any]) -> Callable[..., Any]: """Apply the shared gateway connection options to a Click command.""" command = click.option("--endpoint", default="localhost:5000", show_default=True)(command) command = click.option("--api-key", default=None, help="Gateway API key.")(command) command = click.option( "--api-key-env", default=None, help="Environment variable containing the gateway API key.", )(command) command = click.option("--plaintext", is_flag=True, help="Use plaintext gRPC.")(command) command = click.option("--tls", "use_tls", is_flag=True, help="Use TLS gRPC.")(command) command = click.option("--ca-file", default=None, help="Custom root certificate file.")(command) command = click.option( "--require-certificate-validation", "require_certificate_validation", is_flag=True, help="Verify the TLS certificate against the system trust store " "instead of the lenient trust-on-first-use default.", )(command) command = click.option( "--server-name-override", default=None, help="TLS server name override for test environments.", )(command) command = click.option("--call-timeout", default=30.0, type=float, show_default=True)(command) command = click.option("--stream-timeout", default=None, type=float)(command) return command @main.command("open-session") @gateway_options @click.option("--client-name", default="", help="Client session name.") @click.option("--requested-backend", default="", help="Requested backend name.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def open_session(**kwargs: Any) -> None: """Open a gateway session.""" _run( _open_session(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("close-session") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def close_session(**kwargs: Any) -> None: """Close a gateway session.""" _run( _close_session(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command() @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--message", default="ping", show_default=True, help="Ping payload.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def ping(**kwargs: Any) -> None: """Send a diagnostic ping command.""" _run(_ping(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command() @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--client-name", required=True, help="MXAccess client name.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def register(**kwargs: Any) -> None: """Invoke MXAccess Register.""" _run( _register(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("add-item") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item", required=True, help="MXAccess item definition.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def add_item(**kwargs: Any) -> None: """Invoke MXAccess AddItem.""" _run( _add_item(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command() @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item-handle", required=True, type=int, help="MXAccess item handle.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def advise(**kwargs: Any) -> None: """Invoke MXAccess Advise.""" _run(_advise(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command("subscribe-bulk") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--items", required=True, help="Comma-separated MXAccess item definitions.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def subscribe_bulk(**kwargs: Any) -> None: """Invoke MXAccess SubscribeBulk.""" _run( _subscribe_bulk(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("unsubscribe-bulk") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item-handles", required=True, help="Comma-separated MXAccess item handles.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def unsubscribe_bulk(**kwargs: Any) -> None: """Invoke MXAccess UnsubscribeBulk.""" _run( _unsubscribe_bulk(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("read-bulk") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--items", required=True, help="Comma-separated MXAccess tag addresses.") @click.option("--timeout-ms", default=0, type=int, show_default=True, help="Per-tag snapshot timeout in milliseconds. 0 = worker default.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def read_bulk(**kwargs: Any) -> None: """Invoke MXAccess ReadBulk — cached value when advised, snapshot otherwise.""" _run(_read_bulk(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command("write-bulk") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item-handles", required=True, help="Comma-separated MXAccess item handles.") @click.option("--type", "value_type", default="string", show_default=True) @click.option("--values", required=True, help="Comma-separated values, one per item handle.") @click.option("--user-id", default=0, type=int, show_default=True) @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def write_bulk(**kwargs: Any) -> None: """Invoke MXAccess WriteBulk — sequential Write per entry.""" _run(_write_bulk(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command("write2-bulk") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item-handles", required=True, help="Comma-separated MXAccess item handles.") @click.option("--type", "value_type", default="string", show_default=True) @click.option("--values", required=True, help="Comma-separated values, one per item handle.") @click.option("--timestamp", required=True, help="ISO-8601 timestamp shared across all entries.") @click.option("--user-id", default=0, type=int, show_default=True) @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def write2_bulk(**kwargs: Any) -> None: """Invoke MXAccess Write2Bulk — timestamped sequential Write2 per entry.""" _run(_write2_bulk(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command("write-secured-bulk") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item-handles", required=True, help="Comma-separated MXAccess item handles.") @click.option("--type", "value_type", default="string", show_default=True) @click.option("--values", required=True, help="Comma-separated values, one per item handle.") @click.option("--current-user-id", default=0, type=int, show_default=True) @click.option("--verifier-user-id", default=0, type=int, show_default=True) @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def write_secured_bulk(**kwargs: Any) -> None: """Invoke MXAccess WriteSecuredBulk — credential-sensitive.""" _run(_write_secured_bulk(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command("write-secured2-bulk") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item-handles", required=True, help="Comma-separated MXAccess item handles.") @click.option("--type", "value_type", default="string", show_default=True) @click.option("--values", required=True, help="Comma-separated values, one per item handle.") @click.option("--timestamp", required=True, help="ISO-8601 timestamp shared across all entries.") @click.option("--current-user-id", default=0, type=int, show_default=True) @click.option("--verifier-user-id", default=0, type=int, show_default=True) @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def write_secured2_bulk(**kwargs: Any) -> None: """Invoke MXAccess WriteSecured2Bulk — timestamped + credential-sensitive.""" _run(_write_secured2_bulk(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command("bench-read-bulk") @gateway_options @click.option("--client-name", default="mxgw-python-bench", show_default=True) @click.option("--duration-seconds", default=30, type=int, show_default=True) @click.option("--warmup-seconds", default=3, type=int, show_default=True) @click.option("--bulk-size", default=6, type=int, show_default=True) @click.option("--tag-start", default=1, type=int, show_default=True) @click.option("--tag-prefix", default="TestMachine_", show_default=True) @click.option("--tag-attribute", default="TestChangingInt", show_default=True) @click.option("--timeout-ms", default=1500, type=int, show_default=True) @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def bench_read_bulk(**kwargs: Any) -> None: """Cross-language ReadBulk stress benchmark. Opens its own session, subscribes to bulk-size tags so the worker value cache populates from real OnDataChange events, runs ReadBulk in a tight loop for duration-seconds, and emits the shared JSON stats schema the scripts/bench-read-bulk.ps1 driver collates across all five clients. """ _run(_bench_read_bulk(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command("stream-events") @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--after-worker-sequence", default=0, type=int, show_default=True) @click.option("--max-events", default=1, type=int, show_default=True) @click.option("--timeout", default=5.0, type=float, show_default=True) @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def stream_events(**kwargs: Any) -> None: """Stream a bounded number of events.""" _run( _stream_events(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("stream-alarms") @gateway_options @click.option("--filter-prefix", default="", help="Alarm-reference prefix filter.") @click.option("--max-messages", default=1, type=int, show_default=True) @click.option("--timeout", default=5.0, type=float, show_default=True) @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def stream_alarms(**kwargs: Any) -> None: """Stream a bounded number of messages from the gateway's central alarm feed.""" _run( _stream_alarms(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("acknowledge-alarm") @gateway_options @click.option("--reference", required=True, help="Alarm full reference to acknowledge.") @click.option("--comment", default="", help="Acknowledgement comment.") @click.option("--operator", default="", help="Operator user name.") @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def acknowledge_alarm(**kwargs: Any) -> None: """Acknowledge an active MXAccess alarm condition (session-less).""" _run( _acknowledge_alarm(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command() @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item-handle", required=True, type=int, help="MXAccess item handle.") @click.option("--type", "value_type", default="string", show_default=True) @click.option("--value", required=True, help="Value to write.") @click.option("--user-id", default=0, type=int, show_default=True) @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def write(**kwargs: Any) -> None: """Invoke MXAccess Write.""" _run(_write(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command() @gateway_options @click.option("--session-id", required=True, help="Gateway session id.") @click.option("--server-handle", required=True, type=int, help="MXAccess server handle.") @click.option("--item-handle", required=True, type=int, help="MXAccess item handle.") @click.option("--type", "value_type", default="string", show_default=True) @click.option("--value", required=True, help="Value to write.") @click.option("--timestamp", required=True, help="ISO-8601 timestamp value.") @click.option("--user-id", default=0, type=int, show_default=True) @click.option("--correlation-id", default="", help="Client correlation id.") @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def write2(**kwargs: Any) -> None: """Invoke MXAccess Write2.""" _run(_write2(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command() @gateway_options @click.option("--client-name", default="mxgw-py-smoke", show_default=True) @click.option("--item", required=True, help="MXAccess item definition.") @click.option("--max-events", default=1, type=int, show_default=True) @click.option("--timeout", default=5.0, type=float, show_default=True) @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def smoke(**kwargs: Any) -> None: """Run a bounded open/register/add/advise/stream/close smoke flow.""" _run(_smoke(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs)) @main.command("galaxy-test-connection") @gateway_options @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def galaxy_test_connection(**kwargs: Any) -> None: """Test whether the gateway can reach the Galaxy Repository DB.""" _run( _galaxy_test_connection(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("galaxy-last-deploy") @gateway_options @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def galaxy_last_deploy(**kwargs: Any) -> None: """Read the last Galaxy deploy timestamp.""" _run( _galaxy_last_deploy(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("galaxy-discover") @gateway_options @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def galaxy_discover(**kwargs: Any) -> None: """Enumerate the deployed Galaxy object hierarchy.""" _run( _galaxy_discover(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) @main.command("galaxy-watch") @gateway_options @click.option( "--last-seen-deploy-time", "last_seen_deploy_time", default=None, help="ISO-8601 timestamp; when it matches the current cached deploy time the " "bootstrap event is suppressed.", ) @click.option( "--max-events", default=1, type=int, show_default=True, help="Stop after collecting this many deploy events.", ) @click.option( "--timeout", default=5.0, type=float, show_default=True, help="Seconds to wait for each event before stopping.", ) @click.option("--json", "output_json", is_flag=True, help="Emit JSON output.") def galaxy_watch(**kwargs: Any) -> None: """Stream a bounded number of Galaxy deploy events.""" _run( _galaxy_watch(**kwargs), output_json=kwargs["output_json"], secrets=_secrets(kwargs), ) async def _open_session(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: reply = await client.open_session_raw( pb.OpenSessionRequest( requested_backend=kwargs["requested_backend"], client_session_name=kwargs["client_name"], client_correlation_id=kwargs["correlation_id"], ), ) return {"sessionId": reply.session_id, "rawReply": _message_dict(reply)} async def _close_session(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: reply = await client.close_session_raw( pb.CloseSessionRequest( session_id=kwargs["session_id"], client_correlation_id=kwargs["correlation_id"], ), ) return {"sessionId": reply.session_id, "rawReply": _message_dict(reply)} async def _ping(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: reply = await client.invoke_raw( pb.MxCommandRequest( session_id=kwargs["session_id"], command=pb.MxCommand( kind=pb.MX_COMMAND_KIND_PING, ping=pb.PingCommand(message=kwargs["message"]), ), ), ) return {"kind": "ping", "rawReply": _message_dict(reply)} async def _register(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) server_handle = await session.register( kwargs["client_name"], correlation_id=kwargs["correlation_id"], ) return {"serverHandle": server_handle} async def _add_item(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) item_handle = await session.add_item( kwargs["server_handle"], kwargs["item"], correlation_id=kwargs["correlation_id"], ) return {"itemHandle": item_handle} async def _advise(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) await session.advise( kwargs["server_handle"], kwargs["item_handle"], correlation_id=kwargs["correlation_id"], ) return {"ok": True} async def _subscribe_bulk(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) results = await session.subscribe_bulk( kwargs["server_handle"], _parse_string_list(kwargs["items"]), correlation_id=kwargs["correlation_id"], ) return {"results": [_message_dict(result) for result in results]} async def _unsubscribe_bulk(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) results = await session.unsubscribe_bulk( kwargs["server_handle"], _parse_int_list(kwargs["item_handles"]), correlation_id=kwargs["correlation_id"], ) return {"results": [_message_dict(result) for result in results]} async def _read_bulk(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) results = await session.read_bulk( kwargs["server_handle"], _parse_string_list(kwargs["items"]), timeout_ms=kwargs["timeout_ms"], correlation_id=kwargs["correlation_id"], ) return {"results": [_message_dict(result) for result in results]} def _build_write_bulk_entries(kwargs: dict[str, Any]): """Build (item_handle, MxValue) pairs for the bulk-write families. The CLI accepts a single ``--type`` plus ``--values`` (comma-separated string-encoded values, one per ``--item-handles`` entry). Returns the parsed item-handle list and the per-entry MxValue protobuf instances — callers wrap these into the appropriate per-entry message type. """ handles = _parse_int_list(kwargs["item_handles"]) value_texts = _parse_string_list(kwargs["values"]) if len(handles) != len(value_texts): raise click.UsageError( f"item-handles count ({len(handles)}) does not match values count ({len(value_texts)})", ) parsed = [_parse_value(text, kwargs["value_type"]) for text in value_texts] values = [to_mx_value(v) for v in parsed] return handles, values async def _write_bulk(**kwargs: Any) -> dict[str, Any]: handles, values = _build_write_bulk_entries(kwargs) entries = [ pb.WriteBulkEntry(item_handle=handle, user_id=kwargs["user_id"], value=value) for handle, value in zip(handles, values) ] async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) results = await session.write_bulk( kwargs["server_handle"], entries, correlation_id=kwargs["correlation_id"], ) return {"results": [_message_dict(result) for result in results]} async def _write2_bulk(**kwargs: Any) -> dict[str, Any]: handles, values = _build_write_bulk_entries(kwargs) timestamp_value = to_mx_value(_parse_datetime(kwargs["timestamp"])) entries = [ pb.Write2BulkEntry( item_handle=handle, user_id=kwargs["user_id"], value=value, timestamp_value=timestamp_value, ) for handle, value in zip(handles, values) ] async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) results = await session.write2_bulk( kwargs["server_handle"], entries, correlation_id=kwargs["correlation_id"], ) return {"results": [_message_dict(result) for result in results]} async def _write_secured_bulk(**kwargs: Any) -> dict[str, Any]: handles, values = _build_write_bulk_entries(kwargs) entries = [ pb.WriteSecuredBulkEntry( item_handle=handle, current_user_id=kwargs["current_user_id"], verifier_user_id=kwargs["verifier_user_id"], value=value, ) for handle, value in zip(handles, values) ] async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) results = await session.write_secured_bulk( kwargs["server_handle"], entries, correlation_id=kwargs["correlation_id"], ) return {"results": [_message_dict(result) for result in results]} async def _write_secured2_bulk(**kwargs: Any) -> dict[str, Any]: handles, values = _build_write_bulk_entries(kwargs) timestamp_value = to_mx_value(_parse_datetime(kwargs["timestamp"])) entries = [ pb.WriteSecured2BulkEntry( item_handle=handle, current_user_id=kwargs["current_user_id"], verifier_user_id=kwargs["verifier_user_id"], value=value, timestamp_value=timestamp_value, ) for handle, value in zip(handles, values) ] async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) results = await session.write_secured2_bulk( kwargs["server_handle"], entries, correlation_id=kwargs["correlation_id"], ) return {"results": [_message_dict(result) for result in results]} async def _bench_read_bulk(**kwargs: Any) -> dict[str, Any]: """ReadBulk stress benchmark — matches the .NET / Go / Rust / Java schema.""" bulk_size = int(kwargs["bulk_size"]) if bulk_size < 1: raise click.UsageError("bulk-size must be positive") duration_seconds = int(kwargs["duration_seconds"]) warmup_seconds = int(kwargs["warmup_seconds"]) tag_start = int(kwargs["tag_start"]) tag_prefix = kwargs["tag_prefix"] tag_attribute = kwargs["tag_attribute"] timeout_ms = int(kwargs["timeout_ms"]) client_name = kwargs["client_name"] tags = [f"{tag_prefix}{i:03d}.{tag_attribute}" for i in range(tag_start, tag_start + bulk_size)] async with await _connect(kwargs) as client: session = await client.open_session(client_session_name=client_name) server_handle = 0 item_handles: list[int] = [] try: server_handle = await session.register(client_name) subscribe_results = await session.subscribe_bulk(server_handle, tags) item_handles = [r.item_handle for r in subscribe_results if r.was_successful] # Warm-up window so JIT / connection pool / first-call costs are # amortised before the measurement window opens. warmup_deadline = time.perf_counter() + warmup_seconds while time.perf_counter() < warmup_deadline: await session.read_bulk(server_handle, tags, timeout_ms=timeout_ms) latencies_ms: list[float] = [] total_results = 0 cached_results = 0 successful = 0 failed = 0 steady_start = time.perf_counter() steady_deadline = steady_start + duration_seconds while time.perf_counter() < steady_deadline: call_start = time.perf_counter() try: results = await session.read_bulk(server_handle, tags, timeout_ms=timeout_ms) except Exception: failed += 1 latencies_ms.append((time.perf_counter() - call_start) * 1000.0) continue latencies_ms.append((time.perf_counter() - call_start) * 1000.0) successful += 1 for r in results: total_results += 1 if r.was_cached: cached_results += 1 steady_elapsed = time.perf_counter() - steady_start total_calls = successful + failed calls_per_second = total_calls / steady_elapsed if steady_elapsed > 0 else 0.0 finally: if item_handles: try: await session.unsubscribe_bulk(server_handle, item_handles) except Exception as exc: # noqa: BLE001 — bench is best-effort logger.warning("bench-read-bulk: unsubscribe_bulk cleanup failed: %s", exc) try: await session.close() except Exception as exc: # noqa: BLE001 — bench is best-effort logger.warning("bench-read-bulk: session.close cleanup failed: %s", exc) return { "language": "python", "command": "bench-read-bulk", "endpoint": kwargs.get("endpoint"), "clientName": client_name, "bulkSize": bulk_size, "durationSeconds": duration_seconds, "warmupSeconds": warmup_seconds, "durationMs": int(steady_elapsed * 1000), "tags": tags, "totalCalls": total_calls, "successfulCalls": successful, "failedCalls": failed, "totalReadResults": total_results, "cachedReadResults": cached_results, "callsPerSecond": round(calls_per_second, 2), "latencyMs": _percentile_summary(latencies_ms), } def _percentile_summary(sample: list[float]) -> dict[str, float]: if not sample: return {"p50": 0.0, "p95": 0.0, "p99": 0.0, "max": 0.0, "mean": 0.0} sorted_sample = sorted(sample) return { "p50": round(_percentile(sorted_sample, 0.50), 3), "p95": round(_percentile(sorted_sample, 0.95), 3), "p99": round(_percentile(sorted_sample, 0.99), 3), "max": round(sorted_sample[-1], 3), "mean": round(sum(sample) / len(sample), 3), } def _percentile(sorted_sample: list[float], quantile: float) -> float: """Nearest-rank with linear interpolation; matches every other client.""" n = len(sorted_sample) if n == 0: return 0.0 if n == 1: return sorted_sample[0] rank = quantile * (n - 1) lower = int(rank) upper = min(lower + 1, n - 1) fraction = rank - lower return sorted_sample[lower] + (sorted_sample[upper] - sorted_sample[lower]) * fraction async def _stream_events(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) events = await _collect_events( session.stream_events(after_worker_sequence=kwargs["after_worker_sequence"]), max_events=kwargs["max_events"], timeout=kwargs["timeout"], ) return {"events": [_message_dict(event) for event in events]} async def _stream_alarms(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: messages = await _collect_alarm_messages( client.stream_alarms( pb.StreamAlarmsRequest( client_correlation_id=kwargs["correlation_id"], alarm_filter_prefix=kwargs["filter_prefix"], ), ), max_messages=kwargs["max_messages"], timeout=kwargs["timeout"], ) return {"messages": [_message_dict(message) for message in messages]} async def _acknowledge_alarm(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: reply = await client.acknowledge_alarm( pb.AcknowledgeAlarmRequest( client_correlation_id=kwargs["correlation_id"], alarm_full_reference=kwargs["reference"], comment=kwargs["comment"], operator_user=kwargs["operator"], ), ) return _message_dict(reply) async def _write(**kwargs: Any) -> dict[str, Any]: value = _parse_value(kwargs["value"], kwargs["value_type"]) async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) await session.write( kwargs["server_handle"], kwargs["item_handle"], value, user_id=kwargs["user_id"], correlation_id=kwargs["correlation_id"], ) return {"ok": True} async def _write2(**kwargs: Any) -> dict[str, Any]: value = _parse_value(kwargs["value"], kwargs["value_type"]) timestamp = _parse_datetime(kwargs["timestamp"]) async with await _connect(kwargs) as client: session = _session(client, kwargs["session_id"]) await session.write2( kwargs["server_handle"], kwargs["item_handle"], value, timestamp, user_id=kwargs["user_id"], correlation_id=kwargs["correlation_id"], ) return {"ok": True} async def _smoke(**kwargs: Any) -> dict[str, Any]: async with await _connect(kwargs) as client: session = await client.open_session(client_session_name=kwargs["client_name"]) closed = False try: server_handle = await session.register(kwargs["client_name"]) item_handle = await session.add_item(server_handle, kwargs["item"]) await session.advise(server_handle, item_handle) events = await _collect_events( session.stream_events(), max_events=kwargs["max_events"], timeout=kwargs["timeout"], ) return { "sessionId": session.session_id, "serverHandle": server_handle, "itemHandle": item_handle, "events": [_message_dict(event) for event in events], } finally: if not closed: await session.close() async def _galaxy_test_connection(**kwargs: Any) -> dict[str, Any]: async with await _connect_galaxy(kwargs) as galaxy: ok = await galaxy.test_connection() return {"command": "galaxy-test-connection", "ok": ok} async def _galaxy_last_deploy(**kwargs: Any) -> dict[str, Any]: async with await _connect_galaxy(kwargs) as galaxy: last_deploy = await galaxy.get_last_deploy_time() payload: dict[str, Any] = { "command": "galaxy-last-deploy", "present": last_deploy is not None, } if last_deploy is not None: # galaxy.py returns a timezone-NAIVE UTC datetime (protobuf ToDatetime()). # Stamp it as UTC so the emitted ISO-8601 carries an unambiguous offset, # matching the Go client's "...Z" output. payload["timeOfLastDeploy"] = last_deploy.replace(tzinfo=timezone.utc).isoformat() return payload async def _galaxy_discover(**kwargs: Any) -> dict[str, Any]: async with await _connect_galaxy(kwargs) as galaxy: objects = await galaxy.discover_hierarchy() return { "command": "galaxy-discover", "objects": [_message_dict(obj) for obj in objects], } async def _galaxy_watch(**kwargs: Any) -> dict[str, Any]: last_seen = kwargs.get("last_seen_deploy_time") last_seen_dt = _parse_datetime(last_seen) if last_seen else None async with await _connect_galaxy(kwargs) as galaxy: events = await _collect_deploy_events( galaxy.watch_deploy_events(last_seen_dt), max_events=kwargs["max_events"], timeout=kwargs["timeout"], ) return { "command": "galaxy-watch", "events": [_message_dict(event) for event in events], } async def _connect(kwargs: dict[str, Any]) -> GatewayClient: api_key = kwargs.get("api_key") or _api_key_from_env(kwargs.get("api_key_env")) return await GatewayClient.connect( ClientOptions( endpoint=kwargs["endpoint"], api_key=api_key, plaintext=_use_plaintext(kwargs), ca_file=kwargs.get("ca_file"), require_certificate_validation=bool(kwargs.get("require_certificate_validation")), server_name_override=kwargs.get("server_name_override"), call_timeout=kwargs.get("call_timeout"), stream_timeout=kwargs.get("stream_timeout"), ), ) async def _connect_galaxy(kwargs: dict[str, Any]) -> GalaxyRepositoryClient: api_key = kwargs.get("api_key") or _api_key_from_env(kwargs.get("api_key_env")) return await GalaxyRepositoryClient.connect( ClientOptions( endpoint=kwargs["endpoint"], api_key=api_key, plaintext=_use_plaintext(kwargs), ca_file=kwargs.get("ca_file"), require_certificate_validation=bool(kwargs.get("require_certificate_validation")), server_name_override=kwargs.get("server_name_override"), call_timeout=kwargs.get("call_timeout"), stream_timeout=kwargs.get("stream_timeout"), ), ) def _session(client: GatewayClient, session_id: str): from zb_mom_ww_mxgateway.session import Session return Session(client=client, session_id=session_id) def _use_plaintext(kwargs: dict[str, Any]) -> bool: """Resolve the plaintext / TLS contract from the CLI flags. TLS is the default. ``--plaintext`` is the only way to opt in to an unencrypted channel; ``--tls`` is accepted as a redundant explicit affirmation. Combining the two is a usage error (regression-guarded by Client.Python-023 — the previous silent ``localhost:`` / ``127.0.0.1:`` auto-plaintext branch leaked the API-key bearer over a plaintext channel when a user ran the gateway behind TLS on loopback). """ plaintext = bool(kwargs.get("plaintext")) use_tls = bool(kwargs.get("use_tls")) if plaintext and use_tls: raise click.UsageError("--plaintext and --tls are mutually exclusive") return plaintext def _api_key_from_env(name: str | None) -> str | None: if not name: return None return os.environ.get(name) def _secrets(kwargs: dict[str, Any]) -> list[str | None]: return [ kwargs.get("api_key"), _api_key_from_env(kwargs.get("api_key_env")), ] def _run( awaitable: Awaitable[dict[str, Any]], *, output_json: bool, secrets: list[str | None], ) -> None: try: payload = asyncio.run(awaitable) except MxGatewayError as error: raise click.ClickException(redact_secret(str(error), secrets)) from error _emit(payload, output_json=output_json) def _emit( payload: dict[str, Any], *, output_json: bool, text: str | None = None, ) -> None: if output_json: click.echo(json.dumps(payload, sort_keys=True)) return click.echo(text or json.dumps(payload, sort_keys=True)) async def _collect_events( events: Any, *, max_events: int, timeout: float, ) -> list[pb.MxEvent]: if max_events > MAX_AGGREGATE_EVENTS: raise click.BadParameter( f"must be less than or equal to {MAX_AGGREGATE_EVENTS}", param_hint="--max-events", ) collected: list[pb.MxEvent] = [] iterator = events.__aiter__() try: while len(collected) < max_events: collected.append(await asyncio.wait_for(iterator.__anext__(), timeout=timeout)) except StopAsyncIteration: pass except asyncio.TimeoutError: pass finally: close = getattr(iterator, "aclose", None) if close is not None: await close() return collected async def _collect_alarm_messages( messages: Any, *, max_messages: int, timeout: float, ) -> list[pb.AlarmFeedMessage]: if max_messages > MAX_AGGREGATE_EVENTS: raise click.BadParameter( f"must be less than or equal to {MAX_AGGREGATE_EVENTS}", param_hint="--max-messages", ) collected: list[pb.AlarmFeedMessage] = [] iterator = messages.__aiter__() try: while len(collected) < max_messages: collected.append(await asyncio.wait_for(iterator.__anext__(), timeout=timeout)) except StopAsyncIteration: pass except asyncio.TimeoutError: pass finally: close = getattr(iterator, "aclose", None) if close is not None: await close() return collected async def _collect_deploy_events( events: Any, *, max_events: int, timeout: float, ) -> list[galaxy_pb.DeployEvent]: if max_events > MAX_AGGREGATE_EVENTS: raise click.BadParameter( f"must be less than or equal to {MAX_AGGREGATE_EVENTS}", param_hint="--max-events", ) collected: list[galaxy_pb.DeployEvent] = [] iterator = events.__aiter__() try: while len(collected) < max_events: collected.append(await asyncio.wait_for(iterator.__anext__(), timeout=timeout)) except StopAsyncIteration: pass except asyncio.TimeoutError: pass finally: close = getattr(iterator, "aclose", None) if close is not None: await close() return collected def _parse_value(raw_value: str, value_type: str) -> MxValueInput: normalized = value_type.lower() if normalized == "bool": return raw_value.lower() in ("1", "true", "yes", "on") if normalized in ("int", "int32", "int64"): return int(raw_value) if normalized in ("float", "double"): return float(raw_value) if normalized in ("time", "timestamp"): return _parse_datetime(raw_value) if normalized == "raw": return raw_value.encode("utf-8") if normalized == "string": return raw_value raise click.BadParameter(f"unsupported value type: {value_type}", param_hint="--type") def _parse_datetime(raw_value: str) -> datetime: if raw_value.endswith("Z"): raw_value = raw_value[:-1] + "+00:00" parsed = datetime.fromisoformat(raw_value) if parsed.tzinfo is None: parsed = parsed.replace(tzinfo=timezone.utc) return parsed def _parse_string_list(raw_value: str) -> list[str]: values = [item.strip() for item in raw_value.split(",") if item.strip()] if not values: raise click.BadParameter("at least one item is required", param_hint="--items") return values def _parse_int_list(raw_value: str) -> list[int]: values = [item.strip() for item in raw_value.split(",") if item.strip()] if not values: raise click.BadParameter("at least one item handle is required", param_hint="--item-handles") return [int(item) for item in values] def _message_dict(message: Any) -> dict[str, Any]: return MessageToDict( message, preserving_proto_field_name=False, use_integers_for_enums=False, )