Source code for cow.monitoring.manager

"""
Module implement base class for monitoring
"""

from collections.abc import Iterable
from typing import Literal, Protocol

from luna_plugins.base.manager import PluginManager

from ..utils import mixins
from ..utils.check_connection import checkConnectionToClickhouse, checkConnectionToInflux
from ..utils.log import Logger
from .base_adapter import BaseMonitoringAdapter
from .clickhouse_adapter import ClickhouseMonitoringAdaptor, ClickhouseSettings
from .influx_adapter import InfluxMonitoringAdapter, InfluxSettings
from .points import BaseMonitoringPoint

logger = Logger("luna.monitoring")




class MonitoringSettings(Protocol):
    """
    Monitoring settings protocol
    """



    class InfluxCredentials:
        """
        Monitoring credentials
        """

        organization: str
        token: str
        bucket: str




    class ClickhouseCredentials:
        """
        Monitoring credentials
        """

        dbUser: str
        dbPassword: str
        dbName: str


    storageType: Literal["influx", "clickhouse"]
    sendData: Literal[0, 1]
    useSsl: Literal[0, 1]
    flushingPeriod: int
    host: str
    port: int
    credentials: InfluxCredentials | ClickhouseCredentials





class LunaMonitoringManager[TAdaptor: BaseMonitoringAdapter, TPoint: BaseMonitoringPoint](mixins.Initializable):
    """
    Monitoring manager. Sends data to the monitoring storage and monitoring plugins.
    Attributes:
        settings: monitoring storage settings
    """

    adapter: TAdaptor | None = None

    def __init__(self, settings: MonitoringSettings, name: str, pluginManager: PluginManager):
        self.settings = settings
        self.serviceName = name
        self.pluginManager = pluginManager
        # storageType monitoring setting might not be supported for services that either
        # use their own configs e.g. agents, or for some reason
        # use older version of configs service (with no `storageType` setting) with newer version of cow
        self.storageType = getattr(self.settings, "storageType", "influx")



    async def initialize(self) -> None:
        """
        Initialize monitoring
        """
        if not self.settings.sendData:
            return
        url = f"{'https' if self.settings.useSsl else 'http'}://{self.settings.host}"
        if self.storageType == "influx":
            url = f"{url}:{self.settings.port}"
            adapterSettings = InfluxSettings(
                url=url,
                bucket=self.settings.credentials.bucket,
                organization=self.settings.credentials.organization,
                token=self.settings.credentials.token,
            )
            self.adapter = InfluxMonitoringAdapter(
                settings=adapterSettings, flushingPeriod=self.settings.flushingPeriod
            )
        elif self.storageType == "clickhouse":
            url = f"{url}:{self.settings.port}"
            self.adapter = ClickhouseMonitoringAdaptor(
                ClickhouseSettings(
                    url=url,
                    user=self.settings.credentials.dbUser,
                    password=self.settings.credentials.dbPassword,
                    database=self.settings.credentials.dbName,
                ),
                flushingPeriod=self.settings.flushingPeriod,
            )
        else:
            raise RuntimeError(f"Bad monitoring storage type {self.storageType}")
        await self.adapter.initializeMonitoring()




    async def close(self) -> None:
        """
        Stop monitoring.
        """
        if self.adapter:
            await self.adapter.stopMonitoring()
            self.adapter = None




    def flushPoints(self, points: Iterable[TPoint]) -> None:
        """
        Flush point to monitoring.

        Args:
            points: point
        """
        if self.pluginManager:
            self.pluginManager.sendEventToPlugins("monitoring_event", points, logger)
        if self.adapter:
            self.adapter.addPointsToBuffer(points)




    async def probe(self):
        """
        Check that we can connect to the monitoring service with current configuration.
        Can be used without initialization
        """
        if self.settings.sendData:
            if self.storageType == "influx":
                return await checkConnectionToInflux(
                    host=self.settings.host,
                    port=self.settings.port,
                    bucket=self.settings.credentials.bucket,
                    organization=self.settings.credentials.organization,
                    token=self.settings.credentials.token,
                    ssl=bool(self.settings.useSsl),
                    asyncCheck=True,
                )
            if self.storageType == "clickhouse":
                return await checkConnectionToClickhouse(
                    host=self.settings.host,
                    port=self.settings.port,
                    database=self.settings.credentials.dbName,
                    user=self.settings.credentials.dbUser,
                    password=self.settings.credentials.dbPassword,
                    ssl=bool(self.settings.useSsl),
                    asyncCheck=True,
                )
            raise RuntimeError(f"Bad monitoring storage type {self.storageType}")
        return True