From 89b9b3c4af301eadfca9d1459052ef339388f2ce Mon Sep 17 00:00:00 2001 From: Logan Cary Date: Fri, 30 Aug 2024 08:57:27 -0400 Subject: [PATCH] NAS-129755 / 25.04 / Introduce docstrings and type annotations (#7) * add readme and pytype CI * manual merge prep * silence CI import error * fix CI * fix CI * import `ErrnoMixin` for middlewared.service_exception --- .github/workflows/pytype.yml | 27 ++ README.md | 96 +++++++ truenas_api_client/__init__.py | 480 ++++++++++++++++++++++++++++----- truenas_api_client/config.py | 1 + truenas_api_client/ejson.py | 62 ++++- truenas_api_client/exc.py | 59 +++- truenas_api_client/jsonrpc.py | 99 +++++++ truenas_api_client/legacy.py | 6 +- truenas_api_client/utils.py | 77 +++++- 9 files changed, 823 insertions(+), 84 deletions(-) create mode 100644 .github/workflows/pytype.yml create mode 100644 README.md diff --git a/.github/workflows/pytype.yml b/.github/workflows/pytype.yml new file mode 100644 index 0000000..5d84e6e --- /dev/null +++ b/.github/workflows/pytype.yml @@ -0,0 +1,27 @@ +name: "pytype" + +on: + pull_request: + types: + - 'synchronize' + - 'opened' + +jobs: + type-check: + + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v3 + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: 3.11 + + - name: Install Dependencies + run: | + python -m pip install --upgrade pip + pip install pytype websocket-client + + - name: Type-Check + run: pytype -j auto . \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..2b2e6e7 --- /dev/null +++ b/README.md @@ -0,0 +1,96 @@ +

+ Join Discord + Join Forums + File Issue +

+ +# TrueNAS Websocket Client + +*Found an issue? Please report it on our [Jira bugtracker](https://jira.ixsystems.com).* + +## About + +The TrueNAS websocket client provides the command line tool `midclt` and the means to easily communicate with [middleware](https://github.com/truenas/middleware) using Python by making calls through the [websocket API](https://www.truenas.com/docs/api/scale_websocket_api.html). The client can connect to a local TrueNAS instance by default or to a specified remote socket. This offers an alternative to going through the [web UI](https://github.com/truenas/webui) or connecting via ssh. + +By default, communication facilitated by the API between the client and middleware now uses the [JSON-RPC 2.0](https://www.jsonrpc.org/specification) protocol. However, it is still possible to use the legacy client by passing a legacy uri, e.g. `'ws://some.truenas.address/websocket'` as opposed to `'ws://some.truenas.address/api/current'`. + +## Getting Started + +TrueNAS comes with this client preinstalled, but it is also possible to use the TrueNAS websocket client from a non-TrueNAS host. + +On a non-TrueNAS host, ensure that Git is installed and run `pip install git+https://github.com/truenas/api_client.git` to automatically install dependencies. You may alternatively clone this repository and run `python setup.py install`. Using a Python venv is recommended. + +## Usage + +The `midclt` command (not to be confused with the [TrueNAS CLI](https://github.com/truenas/midcli)) provides a way to make direct API calls through the client. To view its syntax, enter `midclt -h`. The `-h` option can also be used with any of `midclt`'s subcommands. + +The client's default behavior is to connect to the localhost's middlewared socket. For a remote connection, e.g. from a Windows host, you must specify the `--uri` option and authenticate with either user credentials or an API key. For example: `midclt --uri ws:///api/current -K key ...` + +### Make local API calls + +``` +root@my_truenas[~]# midclt call user.create '{"full_name": "John Doe", "username": "user", "password": "pass", "group_create": true}' +``` + +### Login to a remote TrueNAS + +``` +root@my_truenas[~]# midclt --uri ws://some.other.truenas/api/current -U user -P password call system.info +``` + +### Start a job + +``` +root@my_truenas[~]# midclt call -j pool.dataset.lock mypool/mydataset +``` + +## Development + +The TrueNAS API client can also be used in Python scripts. + +### Make local API calls + +```python +from truenas_api_client import Client + +with Client() as c: # Local IPC + print(c.ping()) # pong + user = {"full_name": "John Doe", "username": "user", "password": "pass", "group_create": True} + entry_id = c.call("user.create", user) + user = c.call("user.get_instance", entry_id) + print(user["full_name"]) # John Doe +``` + +### Login with a user account or an API key + +```python +# User account +with Client(uri="ws://some.other.truenas/api/current") as c: + c.call("auth.login", username, password) + +# API key +with Client(uri="ws://some.other.truenas/api/current") as c: + c.call("auth.login_with_api_key", key) +``` + +### Start a job + +```python +with Client() as c: + is_locked = c.call("pool.dataset.lock", "mypool/mydataset", job=True) + if is_locked: + args = {"datasets": [{"name": "mypool/mydataset", "passphrase": "passphrase"}]} + c.call("pool.dataset.unlock", "mypool/mydataset", args, job=True) +``` + +## Helpful Links + + + + + +- [Websocket API docs](https://www.truenas.com/docs/api/scale_websocket_api.html) +- [Middleware repo](https://github.com/truenas/middleware) +- [Official TrueNAS Documentation Hub](https://www.truenas.com/docs/) +- [Get started building TrueNAS Scale](https://github.com/truenas/scale-build) +- [Forums](https://www.truenas.com/community/) diff --git a/truenas_api_client/__init__.py b/truenas_api_client/__init__.py index 462366b..520bef7 100644 --- a/truenas_api_client/__init__.py +++ b/truenas_api_client/__init__.py @@ -1,6 +1,37 @@ +"""Provides a simple way to call middleware API endpoints using a websocket connection. + +The full websocket API documentation can be found at https://www.truenas.com/docs/api/core_websocket_api.html. + +Example:: + + $ midclt ping && echo 'Connected' || echo 'Unable to ping' + Connected + $ midclt call user.create '{"full_name": "John Doe", "username": "user", "password": "pass", "group_create": true}' + 70 + $ midclt call user.get_instance 70 + {"id": 70, "uid": 3000, "username": "user", "unixhash": ... } + $ midclt call user.query '[["full_name", "=", "John Doe"]]' + {"id": 70, "uid": 3000, "username": "user", "unixhash": ... } + +Example:: + + with Client() as c: # Local IPC + print(c.ping()) # pong + user = {"full_name": "John Doe", "username": "user", "password": "pass", "group_create": True} + id = c.call("user.create", user) + user = c.call("user.get_instance", id) + print(user["full_name"]) # John Doe + +Example:: + + c = Client("ws://example.com/api/current") # Remote websocket connection + c.close() + +""" import argparse from base64 import b64decode from collections import defaultdict +from collections.abc import Callable, Iterable import errno import logging import pickle @@ -10,6 +41,7 @@ import sys from threading import Event, Lock, Thread import time +from typing import Any, Literal, NotRequired, Protocol, TypeAlias, TypedDict import urllib.parse import uuid @@ -24,15 +56,35 @@ from .config import CALL_TIMEOUT from .exc import ReserveFDException, ClientException, ErrnoMixin, ValidationErrors, CallTimeout from .legacy import LegacyClient -from .jsonrpc import JSONRPCError -from .utils import MIDDLEWARE_RUN_DIR, ProgressBar, undefined +from .jsonrpc import CollectionUpdateParams, ErrorObj, JobFields, JSONRPCError, JSONRPCMessage, TruenasError +from .utils import MIDDLEWARE_RUN_DIR, ProgressBar, undefined, UndefinedType logger = logging.getLogger(__name__) class Client: - def __init__(self, uri=None, reserved_ports=False, py_exceptions=False, log_py_exceptions=False, - call_timeout=undefined, verify_ssl=True): + """Implicit wrapper of either a `JSONRPCClient` or a `LegacyClient`.""" + + def __init__(self, uri: str | None=None, reserved_ports=False, py_exceptions=False, log_py_exceptions=False, + call_timeout: float | UndefinedType=undefined, verify_ssl=True): + """Initialize either a `JSONRPCClient` or a `LegacyClient`. + + Use `JSONRPCClient` unless `uri` ends with '/websocket'. + + Args: + uri: The address to connect to. Defaults to the local middlewared socket. + reserved_ports: `True` if the local socket should use a reserved port. + py_exceptions: `True` if the server should include exception objects in + `message['error']['data']['py_exception']`. + log_py_exceptions: `True` if exception tracebacks from API calls should be logged. + call_timeout: Number of seconds to allow an API call before timing out. Can be overridden on a per-call + basis. Defaults to `CALL_TIMEOUT`. + verify_ssl: `True` if SSL certificate should be verified before connecting. + + Raises: + ClientException: `WSClient` timed out or some other connection error occurred. + + """ if uri is not None and uri.endswith('/websocket'): client_class = LegacyClient else: @@ -51,16 +103,36 @@ def __exit__(self, exc_type, exc_val, exc_tb): class WSClient: - def __init__(self, url, *, client, reserved_ports=False, verify_ssl=True): + """A supporter class for `JSONRPCClient` that manages the `WebSocket` connection to the server. + + The object used by `JSONRPCClient` to send and receive data. + + """ + def __init__(self, url: str, *, client: 'JSONRPCClient', reserved_ports: bool=False, verify_ssl: bool=True): + """Initialize a `WSClient`. + + Args: + url: The websocket to connect to. `ws://` or `wss://` for secure connection. + client: Reference to the `JSONRPCClient` instance that uses this object. + reserved_ports: `True` if the `socket` should bind to a reserved port, i.e. 600-1024. + verify_ssl: `True` if SSL certificate should be verified before connecting. + + """ self.url = url self.client = client self.reserved_ports = reserved_ports self.verify_ssl = verify_ssl - self.socket = None - self.app = None + self.socket: socket.socket + self.app: WebSocketApp def connect(self): + """Connect a `socket` and start a `WebSocketApp` in a daemon `Thread`. + + Raises: + Exception: The `socket` failed to connect. + + """ unix_socket_prefix = "ws+unix://" if self.url.startswith(unix_socket_prefix): self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) @@ -93,14 +165,27 @@ def connect(self): ) Thread(daemon=True, target=self.app.run_forever).start() - def send(self, data): + def send(self, data: bytes | str): + """Send data to the server by calling `WebSocketApp.send()`. + + Args: + data: The serialized JSON-RPC v2.0-formatted request to send. + + """ return self.app.send(data) def close(self): + """Cleanly close the `WebSocket` connection to the server.""" self.app.close() self.client.on_close(STATUS_NORMAL) def _bind_to_reserved_port(self): + """Bind to a random port in the 600-1024 range. + + Raises: + ReserveFDException: Five failed attempts with different ports. + + """ # linux doesn't have a mechanism to allow the kernel to dynamically # assign ports in the "privileged" range (i.e. 600 - 1024) so we # loop through and call bind() on a privileged port explicitly since @@ -125,6 +210,11 @@ def _bind_to_reserved_port(self): raise ReserveFDException() def _on_open(self, app): + """Callback passed to the `WebSocketApp` to execute when `run_forever` is called. + + Configure the `socket` and call `client.on_open()`. + + """ # TCP keepalive settings don't apply to local unix sockets if 'ws+unix' not in self.url: # enable keepalives on the socket @@ -150,29 +240,77 @@ def _on_open(self, app): self.client.on_open() def _on_message(self, app, data): + """Callback passed to the `WebSocketApp` to execute when data is received. + + Pass the received data to the `JSONRPCClient`. + + """ self.client._recv(json.loads(data)) def _on_error(self, app, e): + """Callback passed to the `WebSocketApp` to execute when an error occurs. + + Log the error. + + """ logger.warning("Websocket client error: %r", e) self.client._ws_connection_error = e def _on_close(self, app, code, reason): + """Callback passed to the `WebSocketApp` to execute when it closes. + + Close the `JSONRPCClient`. + + """ self.client.on_close(code, reason) class Call: - def __init__(self, method, params): + """An encapsulation of the data from a single request-response pair.""" + + def __init__(self, method: str, params: tuple): + """Initialize a `Call` object with an automatically-assigned id. + + Args: + method: The API method being called. + params: Arguments passed to the method. + + """ self.id = str(uuid.uuid4()) self.method = method self.params = params self.returned = Event() - self.result = None - self.error = None - self.py_exception = None + self.result: Any = None + self.error: ClientException | None = None + self.py_exception: BaseException | None = None + + +class _JobDict(JobFields): + """Contains data received from the server for a particular running job.""" + __ready: Event + """Is set when the job returns or ends in error.""" + __callback: '_JobCallback | None' + """Procedure to execute each time a job update is received.""" + + +_JobCallback: TypeAlias = Callable[[_JobDict], None] class Job: - def __init__(self, client, job_id, callback=None): + """A long-running background process on the server initiated by an API call. + + Every `Job` is responsible for a corresponding `_JobDict` in the client's list of jobs. + + """ + def __init__(self, client: 'JSONRPCClient', job_id: str, callback: _JobCallback | None=None): + """Initialize `Job`. + + Args: + client: Reference to the client that created this `Job` and receives updates on its progress. + job_id: The job id returned by the server. Index of this `Job` in the `client._jobs` dictionary. + callback: A procedure to be called every time a job event is received. + + """ self.client = client self.job_id = job_id # If a job event has been received already then we must set an Event @@ -190,6 +328,16 @@ def __repr__(self): return f'' def result(self): + """Wait for the job to finish and return its result. + + Returns: + Any: The job's result. + + Raises: + ValidationErrors: The job failed due to one or more validation errors. + ClientException: No job event was received or it did not succeed. + + """ # Wait indefinitely for the job event with state SUCCESS/FAILED/ABORTED self.event.wait() job = self.client._jobs.pop(self.job_id, None) @@ -197,7 +345,7 @@ def result(self): raise ClientException('No job event was received.') if job['state'] != 'SUCCESS': if job['exc_info'] and job['exc_info']['type'] == 'VALIDATION': - raise ValidationErrors(job['exc_info']['extra']) + raise ValidationErrors(job['exc_info']['extra'] or []) raise ClientException( job['error'], trace={ @@ -210,12 +358,55 @@ def result(self): return job['result'] +class _EventCallbackProtocol(Protocol): + """Specifies how event callbacks should be defined.""" + def __call__(self, mtype: str, **message: Any) -> None: ... + + +class _Payload(TypedDict): + """Contains data for managing a subscription. + + Attributes: + callback: Procedure to call when the event is triggered. + sync: If `True`, main client thread blocks until `callback` finishes each time it is invoked. Otherwise, run + `callback` in the background as a daemon `Thread`. + event: `Event` that is set when the subscription should end. + error: Information included in the Notification if the subscription ended in error. + id: Random UUID assigned by `core.subscribe`. + ready: For backwards compatibility with `LegacyClient`. + + """ + callback: _EventCallbackProtocol | None + sync: bool + event: Event + error: NotRequired[str | TruenasError | None] + id: NotRequired[str] + ready: NotRequired[Event] + + class JSONRPCClient: - def __init__(self, uri=None, reserved_ports=False, py_exceptions=False, log_py_exceptions=False, - call_timeout=undefined, verify_ssl=True): - """ - Arguments: - :reserved_ports(bool): should the local socket used a reserved port + """The object used to interface with the TrueNAS API. + + Keeps track of the calls made, jobs submitted, and callbacks. Maintains a websocket connection using a `WSClient`. + + """ + def __init__(self, uri: str | None=None, reserved_ports=False, py_exceptions=False, log_py_exceptions=False, + call_timeout: float | UndefinedType=undefined, verify_ssl=True): + """Initialize a `JSONRPCClient`. + + Args: + uri: The address to connect to. Defaults to the local middlewared socket. + reserved_ports: `True` if the local socket should use a reserved port. + py_exceptions: `True` if the server should include exception objects in + `message['error']['data']['py_exception']`. + log_py_exceptions: `True` if exception tracebacks from API calls should be logged. + call_timeout: Number of seconds to allow an API call before timing out. Can be overridden on a per-call + basis. Defaults to `CALL_TIMEOUT`. + verify_ssl: `True` if SSL certificate should be verified before connecting. + + Raises: + ClientException: `WSClient` timed out or some other connection error occurred. + """ if uri is None: uri = f'ws+unix://{MIDDLEWARE_RUN_DIR}/middlewared.sock' @@ -223,19 +414,19 @@ def __init__(self, uri=None, reserved_ports=False, py_exceptions=False, log_py_e if call_timeout is undefined: call_timeout = CALL_TIMEOUT - self._calls = {} - self._jobs = defaultdict(dict) + self._calls: dict[str, Call] = {} + self._jobs: defaultdict[str, _JobDict] = defaultdict(dict) # type: ignore self._jobs_lock = Lock() self._jobs_watching = False self._py_exceptions = py_exceptions self._log_py_exceptions = log_py_exceptions self._call_timeout = call_timeout - self._event_callbacks = defaultdict(list) + self._event_callbacks: defaultdict[str, list[_Payload]] = defaultdict(list) self._set_options_call: Call | None = None self._closed = Event() self._connected = Event() - self._connection_error = None - self._ws_connection_error = None + self._connection_error: str | None = None + self._ws_connection_error: WebSocketException self._ws = WSClient( uri, client=self, @@ -246,7 +437,7 @@ def __init__(self, uri=None, reserved_ports=False, py_exceptions=False, log_py_e self._connected.wait(10) if not self._connected.is_set(): raise ClientException('Failed connection handshake') - if self._ws_connection_error is not None: + if hasattr(self, '_ws_connection_error'): if isinstance(self._ws_connection_error, WebSocketException): raise self._ws_connection_error if self._connection_error is not None: @@ -257,10 +448,17 @@ def __enter__(self): def __exit__(self, typ, value, traceback): self.close() - if typ is not None: - raise def _send(self, data): + """Send data to the server using `WSClient`. + + Args: + data: Object serializable with `ejson`. + + Raises: + ClientException: Connection to the server closed prematurely. + + """ try: self._ws.send(json.dumps(data)) except (AttributeError, WebSocketConnectionClosedException): @@ -268,13 +466,29 @@ def _send(self, data): # running tasks in the event loop (i.e. failover.call_remote failover.get_disks_local) raise ClientException('Unexpected closure of remote connection', errno.ECONNABORTED) - def _recv(self, message): + def _recv(self, message: JSONRPCMessage): + """Process a deserialized JSON-RPC v2.0 message from the server. + + The TrueNAS websocket `JSONRPCClient` receives data from the server in two standard forms: Notifications and + Responses. These are defined in the JSON-RPC v2.0 protocol at https://www.jsonrpc.org/specification. + + In the TrueNAS websocket client, Notifications are used to communicate subscription updates including when a + subscription is terminated. These subscription updates also include updates about jobs submitted by the client + via `core.get_jobs`. + + A Response is the server's answer to a Request sent by the client which may or may not come back immediately + depending on the Request sent. + + Args: + message: Deserialized JSON-RPC v2.0 data from the server. + + """ try: if 'method' in message: - params = message['params'] match message['method']: case 'collection_update': if self._event_callbacks: + params = message['params'] if '*' in self._event_callbacks: for event in self._event_callbacks['*']: self._run_callback(event, [params['msg'].upper()], params) @@ -282,6 +496,7 @@ def _recv(self, message): for event in self._event_callbacks[params['collection']]: self._run_callback(event, [params['msg'].upper()], params) case 'notify_unsubscribed': + params = message['params'] if params['collection'] in self._event_callbacks: for event in self._event_callbacks[params['collection']]: if 'error' in params: @@ -295,7 +510,7 @@ def _recv(self, message): try: self._parse_error(message['error'], self._set_options_call) except Exception: - logger.error('Unhandled exception in Client._parse_error', exc_info=True) + logger.error('Unhandled exception in JSONRPCClient._parse_error', exc_info=True) else: logger.error('Error setting client options: %r', self._set_options_call.error) self._connected.set() @@ -306,7 +521,7 @@ def _recv(self, message): try: self._parse_error(message['error'], call) except Exception: - logger.error('Unhandled exception in Client._parse_error', exc_info=True) + logger.error('Unhandled exception in JSONRPCClient._parse_error', exc_info=True) call.returned.set() self._unregister_call(call) else: @@ -314,14 +529,21 @@ def _recv(self, message): else: logger.error('Received unknown message %r', message) except Exception: - logger.error('Unhandled exception in Client._recv', exc_info=True) + logger.error('Unhandled exception in JSONRPCClient._recv', exc_info=True) - def _parse_error(self, error: dict, call: Call): + def _parse_error(self, error: ErrorObj, call: Call): + """Convert an error received from the server into a `ClientException` and store it. + + Args: + error: The JSON object received in an error Response. + call: The associated `Call` object with which to store the `ClientException`. + + """ code = JSONRPCError(error['code']) if self._py_exceptions and code in [JSONRPCError.INVALID_PARAMS, JSONRPCError.TRUENAS_CALL_ERROR]: data = error['data'] call.error = ClientException(data['reason'], data['error'], data['trace'], data['extra']) - if data.get('py_exception'): + if 'py_exception' in data: try: call.py_exception = pickle.loads(b64decode(data['py_exception'])) except Exception as e: @@ -334,16 +556,38 @@ def _parse_error(self, error: dict, call: Call): else: call.error = ClientException(code.name) - def _run_callback(self, event, args, kwargs): + def _run_callback(self, event: _Payload, args: Iterable[str], kwargs: CollectionUpdateParams): + """Call the passed `_Payload`'s callback function. + + Block until the callback returns if `event['sync']` is set. Otherwise, run in a separate daemon `Thread`. + + Args: + event: The `_Payload` whose callback to run. + args: Positional arguments to the callback. + kwargs: Keyword arguments to the callback. + + """ + if event['callback'] is None: + return if event['sync']: event['callback'](*args, **kwargs) else: Thread(target=event['callback'], args=args, kwargs=kwargs, daemon=True).start() def on_open(self): + """Make an API call to `core.set_options` to configure how middlewared sends its responses.""" self._set_options_call = self.call("core.set_options", {"py_exceptions": self._py_exceptions}, background=True) - def on_close(self, code, reason=None): + def on_close(self, code: int, reason: str | None=None): + """Close this `JSONRPCClient` in response to the `WebSocketApp` closing. + + End all unanswered calls and unreturned jobs with an error. + + Args: + code: One of several closing frame status codes defined in `websocket._abnf`. + reason: A message to accompany the closing code and provide more information. + + """ error = f'WebSocket connection closed with code={code!r}, reason={reason!r}' self._connection_error = error @@ -371,22 +615,33 @@ def on_close(self, code, reason=None): self._closed.set() - def _register_call(self, call): + def _register_call(self, call: Call): + """Save a `Call` and index it by its id.""" self._calls[call.id] = call - def _unregister_call(self, call): + def _unregister_call(self, call: Call): + """Remove a `Call` after it has returned.""" self._calls.pop(call.id, None) - def _jobs_callback(self, mtype, **message): - """ - Method to process the received job events. + def _jobs_callback(self, mtype: str, *, fields: JobFields, **message): + """Process a received job event. + + Update the saved job info, execute its saved callback in the background, and set its "__ready" flag if its + "state" is received. + + Args: + mtype: Indicates if the job state has changed. + **message: The members contained in `CollectionUpdateParams`. + + Keyword Args: + fields (JobFields): Contains job id and other information about the job from the server. + """ - fields = message.get('fields') job_id = fields['id'] with self._jobs_lock: if fields: job = self._jobs[job_id] - job.update(fields) + job.update(**fields) if callable(job.get('__callback')): Thread(target=job['__callback'], args=(job,), daemon=True).start() if mtype == 'CHANGED' and job['state'] in ('SUCCESS', 'FAILED', 'ABORTED'): @@ -400,13 +655,37 @@ def _jobs_callback(self, mtype, **message): event.set() def _jobs_subscribe(self): - """ - Subscribe to job updates, calling `_jobs_callback` on every new event. - """ + """Subscribe to job updates, calling `_jobs_callback` on every new event.""" self._jobs_watching = True self.subscribe('core.get_jobs', self._jobs_callback, sync=True) - def call(self, method, *params, background=False, callback=None, job=False, register_call=None, timeout=undefined): + def call(self, method: str, *params, background=False, callback: _JobCallback | None=None, + job: Literal['RETURN'] | bool=False, register_call: bool | None=None, + timeout: float | UndefinedType=undefined) -> Any: + """The primary way to send call requests to the API. + + Send a JSON-RPC v2.0 Request to the server. + + Args: + method: An API endpoint to call. + *params: Arguments to pass to the endpoint. + background: If `background=True`, send the request and return a `Call` object before receiving a response. + By default, wait for the call to return instead. + callback: The callback to pass to the job if `job` is set. + job: If set, subscribe to job updates and if `background=False`, create a `Job`. If `job='RETURN'`, return + the `Job` object rather than just its result. + timeout: Number of seconds to allow the call before timing out if `background=False`. + + Returns: + Call: If `background` is set, return an object representing the request-response pair. + Job: If `job='RETURN'`, return the `Job` object. + Any: Otherwise, return the result of the call. + + Raises: + ClientException: Connection to the server closed prematurely or the call ended in error. + CallTimeout: The call took longer than `timeout` seconds to return. + + """ if register_call is None: register_call = not background @@ -436,17 +715,36 @@ def call(self, method, *params, background=False, callback=None, job=False, regi if not background: self._unregister_call(c) - def wait(self, c, *, callback=None, job=False, timeout=undefined): + def wait(self, c: Call, *, callback: _JobCallback | None=None, job: Literal['RETURN'] | bool=False, + timeout: float | UndefinedType=undefined) -> Any: + """Wait for an API call to return and return its result. + + Args: + c: The `Call` object containing the data that was sent. + callback: The callback to pass to the job if `job` is set. + job: If set, create a `Job`. If `job='RETURN'`, return the `Job` object rather than just its result. + timeout: Override the default number of seconds until a timeout exception occurs. + + Returns: + Job: If `job='RETURN'`, return the `Job` object. + Any: If `job=True`, return the job's result. Otherwise, return the call's result. + + Raises: + CallTimeout: The call took longer than `timeout` seconds to return. + ClientException: The call ended in error and `py_exception` was not enabled for `c`. + BaseException: The call ended in error and `py_exception` was enabled for `c`. + + """ if timeout is undefined: timeout = self._call_timeout try: - if not c.returned.wait(timeout): + if not c.returned.wait(timeout): # type: ignore raise CallTimeout() if c.error: if c.py_exception: - if self._log_py_exceptions: + if self._log_py_exceptions and c.error.trace: logger.error(c.error.trace["formatted"]) raise c.py_exception else: @@ -463,14 +761,34 @@ def wait(self, c, *, callback=None, job=False, timeout=undefined): self._unregister_call(c) @staticmethod - def event_payload(): + def event_payload() -> _Payload: + """Create an empty payload. + + Returns: + _Payload: Empty `_Payload`. + + """ return { 'callback': None, 'sync': False, 'event': Event(), } - def subscribe(self, name, callback, payload=None, sync=False): + def subscribe(self, name: str, callback: _EventCallbackProtocol, payload: _Payload | None=None, + sync: bool=False) -> str: + """Subscribe to an event by calling `core.subscribe`. + + Args: + name: The name of the event to subscribe to. + callback: A procedure to call when an event is triggered. + payload: Dictionary containing subscription information. + sync: If `True`, main client thread blocks until `callback` finishes each time it is invoked. Otherwise, + run `callback` in the background as a daemon `Thread`. + + Returns: + str: The `_Payload` id assigned by `core.subscribe`. + + """ payload = payload or self.event_payload() payload.update({ 'callback': callback, @@ -480,27 +798,56 @@ def subscribe(self, name, callback, payload=None, sync=False): payload['id'] = self.call('core.subscribe', name, timeout=10) return payload['id'] - def unsubscribe(self, id_): + def unsubscribe(self, id_: str): + """Call `core.unsubscribe` and remove all associated `_Payload`s + + Args: + id_: `id` of the `_Payload` to remove. + + """ self.call('core.unsubscribe', id_) for k, events in list(self._event_callbacks.items()): - events = [v for v in events if v['id'] != id_] + events = [v for v in events if v.get('id') != id_] if events: self._event_callbacks[k] = events else: self._event_callbacks.pop(k) - def ping(self, timeout=10): + def ping(self, timeout: float=10) -> Literal['pong']: + """Call `core.ping` to verify connection to the server. + + Args: + timeout: Number of seconds to allow before raising `CallTimeout`. + + Raises: + ClientException: Connection to the server closed prematurely or the call ended in error. + CallTimeout: The call took longer than `timeout` seconds to return. + + """ c = self.call('core.ping', background=True, register_call=True) return self.wait(c, timeout=timeout) def close(self): + """Allow one second for the `WSClient` to close.""" self._ws.close() # Wait for websocketclient thread to close self._closed.wait(1) - self._ws = None + del self._ws def main(): + """The entry point for midclt. Run `midclt -h` to see usage. + + Sub-commands: + call, ping, subscribe + + Options: + -h, -q, -u URI, -U USERNAME, -P PASSWORD, -K API_KEY, -t TIMEOUT + + Raises: + ValueError: Login failed (`midclt call`) or a subscription terminated with an error (`midclt subscribe`). + + """ parser = argparse.ArgumentParser() parser.add_argument('-q', '--quiet', action='store_true') parser.add_argument('-u', '--uri') @@ -557,7 +904,8 @@ def from_json(args): if args.job_print == 'progressbar': # display the job progress and status message while we wait - def callback(progress_bar, job): + def pb_callback(progress_bar: ProgressBar, job: _JobDict): + """Update `progress_bar` with information in `job['progress']`.""" try: progress_bar.update( job['progress']['percent'], job['progress']['description'] @@ -566,16 +914,17 @@ def callback(progress_bar, job): print(f'Failed to update progress bar: {e!s}', file=sys.stderr) with ProgressBar() as progress_bar: - kwargs.update({ - 'job': True, - 'callback': lambda job: callback(progress_bar, job) - }) + kwargs.update( + job=True, + callback=lambda job: pb_callback(progress_bar, job) + ) rv = c.call(args.method[0], *list(from_json(args.method[1:])), **kwargs) progress_bar.finish() else: lastdesc = '' - def callback(job): + def callback(job: _JobDict): + """Print `job`'s description to `stderr` if it has changed.""" nonlocal lastdesc desc = job['progress']['description'] if desc is not None and desc != lastdesc: @@ -616,19 +965,20 @@ def callback(job): event = subscribe_payload['event'] number = 0 - def cb(mtype, **message): + def cb(mtype: str, **message): + """Print the event message and unsubscribe if the maximum number of events is reached.""" nonlocal number print(json.dumps(message)) number += 1 if args.number and number >= args.number: event.set() - c.subscribe(args.event, cb, subscribe_payload) + c.subscribe(args.event, cb, subscribe_payload) # type: ignore (`LegacyClient` does not return `_Payload`) if not event.wait(timeout=args.timeout): sys.exit(1) - if subscribe_payload['error']: + if 'error' in subscribe_payload and subscribe_payload['error']: raise ValueError(subscribe_payload['error']) sys.exit(0) diff --git a/truenas_api_client/config.py b/truenas_api_client/config.py index bcc066a..997f24b 100644 --- a/truenas_api_client/config.py +++ b/truenas_api_client/config.py @@ -1,3 +1,4 @@ import os CALL_TIMEOUT = int(os.environ.get("CALL_TIMEOUT", 60)) +"""Default number of seconds to allow an API call until timing out.""" diff --git a/truenas_api_client/ejson.py b/truenas_api_client/ejson.py index 2044c3b..721afa4 100644 --- a/truenas_api_client/ejson.py +++ b/truenas_api_client/ejson.py @@ -1,9 +1,42 @@ +"""Provides wrappers of the `json` module for handling Python sets and common objects of the `datetime` module. + +Specifically, this module allows `datetime.date`, `datetime.time`, +`datetime.datetime`, and `set` objects to be serialized and deserialized in +addition to the types handled by the `json` module (those types are listed +[here](https://docs.python.org/3.11/library/json.html#json.JSONDecoder)). + +Example:: + + >>> from ejson import dumps, loads + >>> obj = {'string', 4, date.today(), time(16, 22, 6)} + >>> serialized = dumps(obj) + >>> serialized + {"$set": [4, {"$type": "date", "$value": "2024-07-03"}, "string", {"$time": "16:22:06"}]} + >>> deserialized = loads(serialized) + +""" import calendar from datetime import date, datetime, time, timedelta, timezone import json class JSONEncoder(json.JSONEncoder): + """Custom JSON encoder that extends the default encoder to handle more types. + + In addition to the types already supported by `json.JSONEncoder`, this + encoder adds support for the following types: + + | Python | JSON | + | ----------------- | ------------------------------------------------- | + | datetime.date | {"$type": "date", "$value": string[YYYY-MM-DD]} | + | datetime.datetime | {"$date": number[Total milliseconds since EPOCH]} | + | datetime.time | {"$time": string[HH:MM:SS]} | + | set | {"$set": array[items...]} | + + Note: When serializing Python sets, the order that the elements appear in + the JSON array is undefined. + + """ def default(self, obj): if type(obj) is date: return {'$type': 'date', '$value': obj.isoformat()} @@ -19,13 +52,18 @@ def default(self, obj): return super(JSONEncoder, self).default(obj) -def object_hook(obj): +def object_hook(obj: dict): + """Used when deserializing `date`, `time`, `datetime`, and `set` objects. + + Passed as a kwarg to a JSON deserialization function like `json.dump()`. + + """ obj_len = len(obj) if obj_len == 1: if '$date' in obj: return datetime.fromtimestamp(obj['$date'] / 1000, tz=timezone.utc) + timedelta(milliseconds=obj['$date'] % 1000) if '$time' in obj: - return time(*[int(i) for i in obj['$time'].split(':')]) + return time(*[int(i) for i in obj['$time'].split(':')[:4]]) # type: ignore if '$set' in obj: return set(obj['$set']) if obj_len == 2 and '$type' in obj and '$value' in obj: @@ -35,12 +73,28 @@ def object_hook(obj): def dump(obj, fp, **kwargs): + """Wraps `json.dump()` and uses the custom `JSONEncoder`. + + Can serialize `date`, `time`, `datetime`, and `set` objects + to a file-like object. + + """ return json.dump(obj, fp, cls=JSONEncoder, **kwargs) -def dumps(obj, **kwargs): +def dumps(obj, **kwargs) -> str: + """Wraps `json.dumps()` and uses the custom `JSONEncoder`. + + Can serialize `date`, `time`, `datetime`, and `set` objects. + + """ return json.dumps(obj, cls=JSONEncoder, **kwargs) -def loads(obj, **kwargs): +def loads(obj: str | bytes | bytearray, **kwargs): + """Wraps `json.loads()` and uses a custom `object_hook` argument. + + Can deserialize `date`, `time`, `datetime`, and `set` objects. + + """ return json.loads(obj, object_hook=object_hook, **kwargs) diff --git a/truenas_api_client/exc.py b/truenas_api_client/exc.py index d6e3cf0..1bd27b0 100644 --- a/truenas_api_client/exc.py +++ b/truenas_api_client/exc.py @@ -1,8 +1,11 @@ -from collections import namedtuple +"""Defines general classes for handling exceptions which may be raised through the client.""" + +from collections.abc import Iterable import errno +from .jsonrpc import ErrorExtra, TruenasTraceback try: - from libzfs import Error as ZFSError + from libzfs import Error as ZFSError # pytype: disable=import-error except ImportError: # this happens on our CI/CD runners as they do not install the py-libzfs module to run our api integration tests LIBZFS = False @@ -11,21 +14,42 @@ class ReserveFDException(Exception): + """A `WSClient` instance failed to bind to a reserved port.""" pass class ErrnoMixin: + """Provides custom error codes and a function to get the name of an error code.""" + ENOMETHOD = 201 + """Service not found or method not found in service.""" ESERVICESTARTFAILURE = 202 + """Service failed to start.""" EALERTCHECKERUNAVAILABLE = 203 + """Alert checker unavailable.""" EREMOTENODEERROR = 204 + """Remote node responded with an error.""" EDATASETISLOCKED = 205 + """Locked datasets.""" EINVALIDRRDTIMESTAMP = 206 + """Invalid RRD timestamp.""" ENOTAUTHENTICATED = 207 + """Client not authenticated.""" ESSLCERTVERIFICATIONERROR = 208 + """SSL certificate/host key could not be verified.""" @classmethod - def _get_errname(cls, code): + def _get_errname(cls, code: int) -> str | None: + """Get the name of an error given its error code. + + Args: + code: An error code for either a ZFSError or a custom error defined in this class. + + Returns: + str: The name of the associated error. + None: `code` does not match any known errors. + + """ if LIBZFS and 2000 <= code <= 2100: return 'EZFS_' + ZFSError(code).name for k, v in cls.__dict__.items(): @@ -34,7 +58,19 @@ def _get_errname(cls, code): class ClientException(ErrnoMixin, Exception): - def __init__(self, error, errno=None, trace=None, extra=None): + """Represents any exception that might arise from a `Client`.""" + + def __init__(self, error: str, errno: int | None=None, trace: TruenasTraceback | None=None, + extra: list[ErrorExtra] | None=None): + """Initialize `ClientException`. + + Args: + error: An error message offering a reason for the exception. + errno: An error code to classify the error. + trace: Traceback information from the server. + extra: Any other errors pertaining to the exception. + + """ self.errno = errno self.error = error self.trace = trace @@ -44,14 +80,19 @@ def __str__(self): return self.error -Error = namedtuple('Error', ['attribute', 'errmsg', 'errcode']) +class ValidationErrors(ClientException): + """A raisable collection of `ErrorExtra`s that indicates a validation error occurred on the server.""" + def __init__(self, errors: Iterable[ErrorExtra]): + """Initialize `ValidationErrors`. -class ValidationErrors(ClientException): - def __init__(self, errors): + Args: + errors: List of error codes and messages from the server. + + """ self.errors = [] for e in errors: - self.errors.append(Error(e[0], e[1], e[2])) + self.errors.append(ErrorExtra(e[0], e[1], e[2])) super().__init__(str(self)) @@ -64,5 +105,7 @@ def __str__(self): class CallTimeout(ClientException): + """A special `ClientException` raised when a `Call` times out before it can return a result.""" def __init__(self): + """Initiate a `ClientException` with message `"Call timeout"`.""" super().__init__("Call timeout", errno.ETIMEDOUT) diff --git a/truenas_api_client/jsonrpc.py b/truenas_api_client/jsonrpc.py index c48f61f..cf11f87 100644 --- a/truenas_api_client/jsonrpc.py +++ b/truenas_api_client/jsonrpc.py @@ -1,4 +1,10 @@ +"""Collection of types used to reference the structure of JSONRPC-2.0 messages received from the server. + +https://www.jsonrpc.org/specification + +""" import enum +from typing import Any, Literal, NamedTuple, NotRequired, TypeAlias, TypedDict class JSONRPCError(enum.Enum): @@ -11,3 +17,96 @@ class JSONRPCError(enum.Enum): # Custom error codes from -32000 to -32099 as allowed by the specification above TRUENAS_TOO_MANY_CONCURRENT_CALLS = -32000 TRUENAS_CALL_ERROR = -32001 + + +class JobProgress(TypedDict): + percent: float + description: str + + +class ErrorExtra(NamedTuple): + attribute: str + errmsg: str + errcode: int + + +class ExcInfo(TypedDict): + type: str + extra: list[ErrorExtra] | None + repr: str + + +class JobFields(TypedDict): + id: str + state: str + progress: JobProgress + result: Any + exc_info: ExcInfo + error: str + exception: str + + +class CollectionUpdateParams(TypedDict): + msg: str + collection: str + id: NotRequired[Any] + fields: NotRequired[JobFields] + extra: NotRequired[dict] + + +class CollectionUpdate(TypedDict): + jsonrpc: Literal['2.0'] + method: Literal['collection_update'] + params: CollectionUpdateParams + + +TruenasTraceback = TypedDict('TruenasTraceback', { + 'class': str, + 'frames': NotRequired[list[dict[str, Any]]], + 'formatted': str, + 'repr': str, +}) +# Has to be defined this way because `class` is a keyword. + + +class TruenasError(TypedDict): + error: int + errname: str + reason: str + trace: TruenasTraceback | None + extra: list[ErrorExtra] + py_exception: NotRequired[str] + + +class NotifyUnsubscribedParams(TypedDict): + collection: str + error: TruenasError + + +class NotifyUnsubscribed(TypedDict): + jsonrpc: Literal['2.0'] + method: Literal['notify_unsubscribed'] + params: NotifyUnsubscribedParams + + +class SuccessResponse(TypedDict): + jsonrpc: Literal['2.0'] + result: Any + id: str + + +class ErrorObj(TypedDict): + code: int + message: str | None + data: TruenasError + + +class ErrorResponse(TypedDict): + jsonrpc: Literal['2.0'] + error: ErrorObj + id: str + + +Notification: TypeAlias = CollectionUpdate | NotifyUnsubscribed +Response: TypeAlias = SuccessResponse | ErrorResponse +JSONRPCMessage: TypeAlias = Notification | Response diff --git a/truenas_api_client/legacy.py b/truenas_api_client/legacy.py index 1377590..60ab72b 100644 --- a/truenas_api_client/legacy.py +++ b/truenas_api_client/legacy.py @@ -1,3 +1,5 @@ +"""The websocket client prior to implementing JSONRPC-2.0 protocol. Used for backwards compatibility.""" + from base64 import b64decode from collections import defaultdict import errno @@ -20,7 +22,7 @@ from . import ejson as json from .config import CALL_TIMEOUT from .exc import ReserveFDException, ClientException, ValidationErrors, CallTimeout -from .utils import MIDDLEWARE_RUN_DIR, undefined +from .utils import MIDDLEWARE_RUN_DIR, undefined, UndefinedType logger = logging.getLogger(__name__) @@ -190,7 +192,7 @@ def result(self): class LegacyClient: def __init__(self, uri=None, reserved_ports=False, py_exceptions=False, log_py_exceptions=False, - call_timeout=undefined, verify_ssl=True): + call_timeout: float | UndefinedType=undefined, verify_ssl=True): """ Arguments: :reserved_ports(bool): should the local socket used a reserved port diff --git a/truenas_api_client/utils.py b/truenas_api_client/utils.py index 55e18c6..3793bcf 100644 --- a/truenas_api_client/utils.py +++ b/truenas_api_client/utils.py @@ -1,18 +1,52 @@ +"""Utility classes for use in the TrueNAS API client. + +Includes `Struct` for creating regular objects out of `Mapping`s with string +keys and `ProgressBar` for displaying the progress of a task in the CLI. + +Attributes: + MIDDLEWARE_RUN_DIR: Directory containing the middlewared Unix domain socket. + undefined: A dummy object similar in purpose to `None` that indicates an unset variable. + +""" import sys +from typing import Any, final, Mapping + MIDDLEWARE_RUN_DIR = '/var/run/middleware' -undefined = object() + +@final +class UndefinedType: + def __new__(cls): + if not hasattr(cls, '_instance'): + cls._instance = super().__new__(cls) + return cls._instance +undefined = UndefinedType() class Struct: - """ - Simpler wrapper to access using object attributes instead of keys. + """Simpler wrapper to access using object attributes instead of keys. + This is meant for compatibility when switch scripts to use middleware client instead of django directly. + + Example:: + + >>> d = {'a':1, 'b':'2', 'c': [3, '4', {'d':5}], 'e':{'f':{'g':6}}} + >>> s = Struct(d) + >>> s.c + [3, '4', {'d':5}] + >>> s.e.f.g + 6 + """ + def __init__(self, mapping: Mapping[str, Any]): + """Initialize a `Struct` with a `Mapping`. - def __init__(self, mapping): + Args: + mapping: Contains string keys that will become the `Struct`'s attribute names. + + """ for k, v in mapping.items(): if isinstance(v, dict): setattr(self, k, Struct(v)) @@ -21,6 +55,26 @@ def __init__(self, mapping): class ProgressBar(object): + """A simple text-based progress bar that writes to `sys.stderr`. + + Status: (message) + Total Progress: [#####################___________________] 53.00% + + Example: + ``` + with ProgressBar() as pb: + for step in range(1, 101): + pb.update(step) + ``` + + Attributes: + message: String to display next to "Status". + percentage: A float from `0.0` to `100.0` representing the total progress. + write_stream: This is `sys.stderr` by default but can be any `TextIO`. + used_flag: Indicates whether `update()` has been called. + extra: A string or printable object to display after the status message. + + """ def __init__(self): self.message = None self.percentage = 0 @@ -32,6 +86,11 @@ def __enter__(self): return self def draw(self): + """Erase the previous progress bar and draw an updated one. + + If `self.extra` is set, will display "Status: (message) Extra: (extra)". + + """ progress_width = 40 filled_width = int(self.percentage * progress_width) self.write_stream.write('\033[2K\033[A\033[2K\r') @@ -47,7 +106,14 @@ def draw(self): ) self.write_stream.flush() - def update(self, percentage=None, message=None): + def update(self, percentage: float | None=None, message: str | None=None): + """Update the progress bar with a new percentage and/or message, redrawing it. + + Args: + percentage: The new percentage to display. A value of `100.0` represents full. + message: The "Status" message to display above the progress bar. + + """ if not self.used_flag: self.write_stream.write('\n') self.used_flag = True @@ -58,6 +124,7 @@ def update(self, percentage=None, message=None): self.draw() def finish(self): + """Fill the progress bar to 100%.""" self.percentage = 1 def __exit__(self, typ, value, traceback):