installers/charm/pol/src/charm.py

#!/usr/bin/env python3
# Copyright 2020 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 pydantic import ValidationError
from typing import Any, Dict, NoReturn

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

from pod_spec import make_pod_spec

logger = logging.getLogger(__name__)


class ConfigurePodEvent(EventBase):
    """Configure Pod event"""

    pass


class PolEvents(CharmEvents):
    """POL Events"""

    configure_pod = EventSource(ConfigurePodEvent)


class PolCharm(CharmBase):
    """POL Charm."""

    state = StoredState()
    on = PolEvents()

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

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

        # Message bus data initialization
        self.state.set_default(message_host=None)
        self.state.set_default(message_port=None)

        # Database data initialization
        self.state.set_default(database_uri=None)

        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)
        self.framework.observe(self.on.upgrade_charm, self.configure_pod)

        # Registering custom internal events
        self.framework.observe(self.on.configure_pod, self.configure_pod)

        # Registering required relation events
        self.framework.observe(
            self.on.kafka_relation_changed, self._on_kafka_relation_changed
        )
        self.framework.observe(
            self.on.mongodb_relation_changed, self._on_mongodb_relation_changed
        )

        # Registering required relation departed events
        self.framework.observe(
            self.on.kafka_relation_departed, self._on_kafka_relation_departed
        )
        self.framework.observe(
            self.on.mongodb_relation_departed, self._on_mongodb_relation_departed
        )

    def _on_kafka_relation_changed(self, event: EventBase) -> NoReturn:
        """Reads information about the kafka relation.

        Args:
            event (EventBase): Kafka relation event.
        """
        data_loc = event.unit if event.unit else event.app

        message_host = event.relation.data[data_loc].get("host")
        message_port = event.relation.data[data_loc].get("port")

        if (
            message_host
            and message_port
            and (
                self.state.message_host != message_host
                or self.state.message_port != message_port
            )
        ):
            self.state.message_host = message_host
            self.state.message_port = message_port
            self.on.configure_pod.emit()

    def _on_kafka_relation_departed(self, event: EventBase) -> NoReturn:
        """Clear kafka relation data.

        Args:
            event (EventBase): Kafka relation event.
        """
        self.state.message_host = None
        self.state.message_port = None
        self.on.configure_pod.emit()

    def _on_mongodb_relation_changed(self, event: EventBase) -> NoReturn:
        """Reads information about the DB relation.

        Args:
            event (EventBase): DB relation event.
        """
        data_loc = event.unit if event.unit else event.app

        database_uri = event.relation.data[data_loc].get("connection_string")

        if database_uri and self.state.database_uri != database_uri:
            self.state.database_uri = database_uri
            self.on.configure_pod.emit()

    def _on_mongodb_relation_departed(self, event: EventBase) -> NoReturn:
        """Clear mongodb relation data.

        Args:
            event (EventBase): DB relation event.
        """
        self.state.database_uri = None
        self.on.configure_pod.emit()

    def _missing_relations(self) -> str:
        """Checks if there missing relations.

        Returns:
            str: string with missing relations
        """
        data_status = {
            "kafka": self.state.message_host,
            "mongodb": self.state.database_uri,
        }

        missing_relations = [k for k, v in data_status.items() if not v]

        return ", ".join(missing_relations)

    @property
    def relation_state(self) -> Dict[str, Any]:
        """Collects relation state configuration for pod spec assembly.

        Returns:
            Dict[str, Any]: relation state information.
        """
        relation_state = {
            "message_host": self.state.message_host,
            "message_port": self.state.message_port,
            "database_uri": self.state.database_uri,
        }

        return relation_state

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

        Args:
            event (EventBase): Hook or Relation event that started the
                               function.
        """
        if missing := self._missing_relations():
            self.unit.status = BlockedStatus(
                "Waiting for {0} relation{1}".format(
                    missing, "s" if "," in missing else ""
                )
            )
            return

        if not self.unit.is_leader():
            self.unit.status = ActiveStatus("ready")
            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,
                self.relation_state,
                self.model.app.name,
            )
        except ValidationError 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(PolCharm)