Parca

# Copyright 2022 Jon Seager
# See LICENSE file for licensing details.

"""Helpers for generating Parca configuration.

This library is used for generating YAML configuration files for Parca, the continuous profiling
tool. More information about Parca can be found at https://www.parca.dev/.

You can use this library as follows:

```python
from charms.parca_k8s.v0.parca_config import ParcaConfig, parca_command_line

# Generate a Parca config and get the dictionary representation
config = ParcaConfig().to_dict()

# Get the YAML representation of the config
yaml_config = str(ParcaConfig())

# Generate a command line to start Parca (pass the Parca charm config)
cmd = parca_command_line(app_config)
```
"""
from typing import Optional

import yaml

# The unique Charmhub library identifier, never change it
LIBID = "e02f282a42df4472b7b287fcb45c2991"

# Increment this major API version when introducing breaking changes
LIBAPI = 0

# Increment this PATCH version before using `charmcraft publish-lib` or reset
# to 0 if you are raising the major API version
LIBPATCH = 3

DEFAULT_BIN_PATH = "/parca"
DEFAULT_CONFIG_PATH = "/etc/parca/parca.yaml"
DEFAULT_PROFILE_PATH = "/var/lib/parca"


def parca_command_line(
        http_address: str = ":7070",
        app_config: dict = None,
        *,
        bin_path: str = DEFAULT_BIN_PATH,
        config_path: str = DEFAULT_CONFIG_PATH,
        profile_path: str = DEFAULT_PROFILE_PATH,
        path_prefix: Optional[str] = None,
        store_config: dict = None,
) -> str:
    """Generate a valid Parca command line.

    Args:
        app_config: Charm configuration dictionary.
        bin_path: Path to the Parca binary to be started.
        config_path: Path to the Parca YAML configuration file.
        profile_path: Path to profile storage directory.
        path_prefix: Path prefix to configure parca server with. Must start with a ``/``.
        store_config: Configuration to send profiles to a remote store
    """
    cmd = [
        str(bin_path),
        f"--config-path={config_path}",
        f"--http-address={http_address}"
    ]

    if path_prefix:
        if not path_prefix.startswith("/"):
            # parca will blow up if you try this
            raise ValueError("invalid path_prefix: should start with a slash.")
        # quote path_prefix so we don't have to escape the slashes
        path_prefix_option = f"--path-prefix='{path_prefix}'"
        cmd.append(path_prefix_option)

    # Render the template files with the correct values

    if app_config.get("enable-persistence", None):
        # Add the correct command line options for disk persistence
        cmd.append("--enable-persistence")
        cmd.append(f"--storage-path={profile_path}")
    else:
        limit = app_config["memory-storage-limit"] * 1048576
        cmd.append(f"--storage-active-memory={limit}")

    if store_config is not None:
        store_config_args = []

        if addr := store_config.get("remote-store-address", None):
            store_config_args.append(f"--store-address={addr}")

        if token := store_config.get("remote-store-bearer-token", None):
            store_config_args.append(f"--bearer-token={token}")

        if insecure := store_config.get("remote-store-insecure", None):
            store_config_args.append(f"--insecure={insecure}")

        if store_config_args:
            store_config_args.append("--mode=scraper-only")
            cmd += store_config_args

    return " ".join(cmd)


def parse_version(vstr: str) -> str:
    """Parse the output of 'parca --version' and return a representative string."""
    splits = vstr.split(" ")
    # If we're not on a 'proper' released version, include the first few digits of
    # the commit we're build from - e.g. 0.12.1-next+deadbeef
    if "-next" in splits[2]:
        return f"{splits[2]}+{splits[4][:6]}"
    return splits[2]


class ParcaConfig:
    """Class representing the Parca config file."""

    def __init__(self, scrape_configs=[], *, profile_path=DEFAULT_PROFILE_PATH):
        self._profile_path = str(profile_path)
        self._scrape_configs = scrape_configs

    @property
    def _config(self) -> dict:
        return {
            "object_storage": {
                "bucket": {"type": "FILESYSTEM", "config": {"directory": self._profile_path}}
            },
            "scrape_configs": self._scrape_configs,
        }

    def to_dict(self) -> dict:
        """Return the Parca config as a Python dictionary."""
        return self._config

    def __str__(self) -> str:
        """Return the Parca config as a YAML string."""
        return yaml.safe_dump(self._config)