Simple SwarmPit API adapter example

swarmpit_api.py

from dataclasses import dataclass, field
from typing import Optional
from urllib.parse import urljoin

import httpx


@dataclass
class Swarmpit:
    """
    A class for interacting with the Swarmpit API.

    Args:
        api_url (str): The URL of the Swarmpit API. For example: "http://host"
        api_key (str): The Swarmpit API token.

    Example:
        ```
        swarmpit = Swarmpit(api_url="https://swarmpit.example.com", api_key="your_api_token")
        node_id = swarmpit.get_node_id_by_ip("192.168.1.2")
        swarmpit.set_node_label(node_id, "my_label_key", "my_label_value")
        ```
    """

    # public args:
    api_url: str
    api_key: str
    # ---
    # private attrs:
    _client: httpx.Client = field(init=False, default=None)

    def __post_init__(self):
        """Initialize the Swarmpit instance after dataclass fields are set."""
        self.api_url = self.api_url.rstrip("/")
        self._client = self._create_client()

    def _create_client(self) -> httpx.Client:
        """Create an HTTP client with the appropriate headers for Swarmpit API."""
        return httpx.Client(
            base_url=urljoin(self.api_url, "/api"),
            headers={"Authorization": f"Bearer {self.api_key}"},
        )

    def get_node_id_by_ip(self, node_ip: str) -> Optional[str]:
        """
        Get the node ID by IP address.

        Args:
            node_ip (str): The IP address of the node.

        Returns:
            Optional[str]: The node ID if found, otherwise None.

        Raises:
            httpx.HTTPError: If there's a problem with the API request.

        Example:
            ```
            node_id = swarmpit.get_node_id_by_ip("192.168.1.2")
            ```
        """
        # Retrieve the list of nodes
        response = self._client.get("/nodes")
        response.raise_for_status()
        nodes = response.json()

        # Search for the node with the given IP address
        for node in nodes:
            if node["address"] == node_ip:
                return node["id"]

        return None

    def set_node_label(self, node_id: str, label_key: str, label_value: str) -> bool:
        """
        Set a label for a node.

        Args:
            node_id (str): The ID of the node.
            label_key (str): The key for the label.
            label_value (str): The value for the label.

        Returns:
            bool: True if successful, otherwise False.

        Raises:
            httpx.HTTPError: If there's a problem with the API request.

        Example:
            ```
            swarmpit.set_node_label(node_id, "my_label_key", "my_label_value")
            ```
        """
        # Retrieve the node information
        response = self._client.get(f"/nodes/{node_id}")
        response.raise_for_status()
        payload = response.json().copy()

        # Update the label in the node's specification
        for label_data in payload["labels"]:
            if label_data["name"] == label_key:
                label_data["value"] = label_value
                break
        else:
            payload["labels"].append(
                {
                    "name": label_key,
                    "value": label_value,
                }
            )

        # Send the update request
        response = self._client.post(
            f"/nodes/{node_id}",
            json=payload,
        )
        response.raise_for_status()

        return True