installers/charm/prometheus/src/charm.py

#!/usr/bin/env python3
# Copyright 2021 Canonical Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
#
# For those usages not covered by the Apache License, Version 2.0 please
# contact: legal@canonical.com
#
# To get in touch with the maintainers, please contact:
# osm-charmers@lists.launchpad.net
##

import logging
from typing import Dict, List, NoReturn

from ops.charm import CharmBase
from ops.framework import EventBase, StoredState
from ops.main import main
from ops.model import ActiveStatus, Application, BlockedStatus, MaintenanceStatus, Unit
from oci_image import OCIImageResource, OCIImageResourceError

from pod_spec import make_pod_spec

logger = logging.getLogger(__name__)

PROMETHEUS_PORT = 9090


class RelationsMissing(Exception):
    def __init__(self, missing_relations: List):
        self.message = ""
        if missing_relations and isinstance(missing_relations, list):
            self.message += f'Waiting for {", ".join(missing_relations)} relation'
            if "," in self.message:
                self.message += "s"


class RelationDefinition:
    def __init__(self, relation_name: str, keys: List, source_type):
        if source_type != Application and source_type != Unit:
            raise TypeError(
                "source_type should be ops.model.Application or ops.model.Unit"
            )
        self.relation_name = relation_name
        self.keys = keys
        self.source_type = source_type


def check_missing_relation_data(
    data: Dict,
    expected_relations_data: List[RelationDefinition],
):
    missing_relations = []
    for relation_data in expected_relations_data:
        if not all(
            f"{relation_data.relation_name}_{k}" in data for k in relation_data.keys
        ):
            missing_relations.append(relation_data.relation_name)
    if missing_relations:
        raise RelationsMissing(missing_relations)


def get_relation_data(
    charm: CharmBase,
    relation_data: RelationDefinition,
) -> Dict:
    data = {}
    relation = charm.model.get_relation(relation_data.relation_name)
    if relation:
        self_app_unit = (
            charm.app if relation_data.source_type == Application else charm.unit
        )
        expected_type = relation_data.source_type
        for app_unit in relation.data:
            if app_unit != self_app_unit and isinstance(app_unit, expected_type):
                if all(k in relation.data[app_unit] for k in relation_data.keys):
                    for k in relation_data.keys:
                        data[f"{relation_data.relation_name}_{k}"] = relation.data[
                            app_unit
                        ].get(k)
                    break
    return data


class PrometheusCharm(CharmBase):
    """Prometheus Charm."""

    state = StoredState()

    def __init__(self, *args) -> NoReturn:
        """Prometheus Charm constructor."""
        super().__init__(*args)

        # Internal state initialization
        self.state.set_default(pod_spec=None)

        self.port = PROMETHEUS_PORT
        self.image = OCIImageResource(self, "image")

        # Registering regular events
        self.framework.observe(self.on.start, self.configure_pod)
        self.framework.observe(self.on.config_changed, self.configure_pod)

        # Registering provided relation events
        self.framework.observe(
            self.on.prometheus_relation_joined, self._publish_prometheus_info
        )

    def _publish_prometheus_info(self, event: EventBase) -> NoReturn:
        """Publishes Prometheus information.

        Args:
            event (EventBase): Prometheus relation event.
        """
        if self.unit.is_leader():
            rel_data = {
                "host": self.model.app.name,
                "port": str(PROMETHEUS_PORT),
            }
            for k, v in rel_data.items():
                event.relation.data[self.app][k] = v

    @property
    def relations_requirements(self):
        return []

    def get_relation_state(self):
        relation_state = {}
        for relation_requirements in self.relations_requirements:
            data = get_relation_data(self, relation_requirements)
            relation_state = {**relation_state, **data}
        check_missing_relation_data(relation_state, self.relations_requirements)
        return relation_state

    def configure_pod(self, _=None) -> NoReturn:
        """Assemble the pod spec and apply it, if possible.

        Args:
            event (EventBase): Hook or Relation event that started the
                               function.
        """
        if not self.unit.is_leader():
            self.unit.status = ActiveStatus("ready")
            return

        relation_state = None
        try:
            relation_state = self.get_relation_state()
        except RelationsMissing as exc:
            logger.exception("Relation missing error")
            self.unit.status = BlockedStatus(exc.message)
            return

        self.unit.status = MaintenanceStatus("Assembling pod spec")

        # Fetch image information
        try:
            self.unit.status = MaintenanceStatus("Fetching image information")
            image_info = self.image.fetch()
        except OCIImageResourceError:
            self.unit.status = BlockedStatus("Error fetching image information")
            return

        try:
            pod_spec = make_pod_spec(
                image_info,
                self.model.config,
                relation_state,
                self.model.app.name,
                self.port,
            )
        except ValueError as exc:
            logger.exception("Config/Relation data validation error")
            self.unit.status = BlockedStatus(str(exc))
            return

        if self.state.pod_spec != pod_spec:
            self.model.pod.set_spec(pod_spec)
            self.state.pod_spec = pod_spec

        self.unit.status = ActiveStatus("ready")


if __name__ == "__main__":
    main(PrometheusCharm)