metrics with aioinflux

apps/tgbot/pyproject.toml

@@ -10,6 +10,8 @@ aiogram = "^2.11.2"
 httpx = "^0.16.1"
 python-whois = "^0.7.3"
 core = {path = "../core"}
+aioinflux = "^0.9.0"
+loguru = "^0.5.3"
 
 [tool.poetry.dev-dependencies]
 

apps/tgbot/tgbot/bot.py

@@ -1,6 +1,6 @@
 from aiogram import Bot, Dispatcher, executor
 from aiogram.contrib.fsm_storage.memory import MemoryStorage
-
+from tgbot.middlewares import WriteCommandMetric
 import config
 import handlers
 
@@ -11,6 +11,7 @@ dp = Dispatcher(telegram_bot, storage=storage)
 
 def on_startup():
 	handlers.default.setup(dp)
+	dp.middleware.setup(WriteCommandMetric())
 
 
 if __name__ == '__main__':

apps/tgbot/tgbot/config.py

@@ -2,3 +2,10 @@ import os
 
 # Loading token from .env
 TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
+
+# Influx for metrics
+INFLUX_HOST = os.getenv("INFLUX_HOST", None)
+INFLUX_PORT = os.getenv("INFLUX_PORT", None)
+INFLUX_USERNAME = os.getenv("INFLUX_USERNAME", None)
+INFLUX_PASSWORD = os.getenv("INFLUX_PASSWORD", None)
+INFLUX_DB = os.getenv("INFLUX_DB", None)

apps/tgbot/tgbot/handlers/base.py

@@ -7,6 +7,7 @@ from aiogram.bot import Bot
 from datetime import datetime
 from core.coretypes import APINodeInfo
 from .helpers import send_api_requests, check_int, validate_local
+from .metrics import push_metric
 
 header = "Отчет о проверке хоста:" \
          "\n\n— Хост: {target_fq}"\

apps/tgbot/tgbot/handlers/default/icmp.py

@@ -2,6 +2,7 @@ from aiogram.types import Message
 from httpx import Response
 from core.coretypes import ErrorPayload, ICMPCheckerResponse, ResponseStatus
 from ..base import CheckerBaseHandler, NotEnoughArgs, LocalhostForbidden
+from ..metrics import push_status_metric
 
 icmp_help_message = """
 ❓ Производит проверку хоста по протоколу ICMP.
@@ -45,4 +46,5 @@ class ICMPCheckerHandler(CheckerBaseHandler):
         if status == ResponseStatus.ERROR:
             payload = ErrorPayload(**res.json().get("payload"))
             message += f"❌️ {payload.message}"
+        await push_status_metric(status, self.api_endpoint)
         return message

apps/tgbot/tgbot/handlers/default/minecraft.py

@@ -2,7 +2,8 @@ from aiogram.types import Message
 from core.coretypes import ResponseStatus, ErrorPayload, MinecraftResponse
 from httpx import Response
 
-from tgbot.handlers.base import CheckerBaseHandler, NotEnoughArgs, InvalidPort, process_args_for_host_port
+from tgbot.handlers.base import CheckerBaseHandler, process_args_for_host_port
+from tgbot.handlers.metrics import push_status_metric
 
 minecraft_help_message = """
 ❓ Получает статистику о Minecraft сервере
@@ -37,4 +38,5 @@ class MinecraftCheckerHandler(CheckerBaseHandler):
         if status == ResponseStatus.ERROR:
             payload = ErrorPayload(**res.json().get("payload"))
             message += f"❌️ {payload.message}"
+        await push_status_metric(status, self.api_endpoint)
         return message

apps/tgbot/tgbot/handlers/default/tcp.py

@@ -4,6 +4,7 @@ from httpx import Response
 
 from tgbot.handlers.base import CheckerBaseHandler, NotEnoughArgs, InvalidPort
 from tgbot.handlers.helpers import check_int
+from tgbot.handlers.metrics import push_status_metric
 
 tcp_help_message = """
 ❓ Производит проверку TCP порта, открыт ли он или нет
@@ -48,4 +49,5 @@ class TCPCheckerHandler(CheckerBaseHandler):
         if status == ResponseStatus.ERROR:
             payload = ErrorPayload(**res.json().get("payload"))
             message += f"❌️ {payload.message}"
+        await push_status_metric(status, self.api_endpoint)
         return message

apps/tgbot/tgbot/handlers/default/web.py

@@ -1,7 +1,8 @@
 from aiogram.types import Message
 from httpx import Response
 from core.coretypes import ResponseStatus, HTTP_EMOJI, HttpCheckerResponse, ErrorPayload
-from ..base import CheckerBaseHandler, NotEnoughArgs, InvalidPort, process_args_for_host_port
+from ..base import CheckerBaseHandler, process_args_for_host_port
+from ..metrics import push_status_metric
 
 web_help_message = """
 ❓ Производит проверку хоста по протоколу HTTP.
@@ -36,4 +37,5 @@ class WebCheckerHandler(CheckerBaseHandler):
         if status == ResponseStatus.ERROR:
             payload = ErrorPayload(**res.json().get("payload"))
             message += f"❌️ {payload.message}"
+        await push_status_metric(status, self.api_endpoint)
         return message

apps/tgbot/tgbot/handlers/helpers.py

@@ -4,6 +4,8 @@ from core.coretypes import APINode
 from ipaddress import ip_address
 from contextlib import suppress
 
+from tgbot.handlers.metrics import push_api_request_status
+
 
 def check_int(value) -> bool:
     try:
@@ -47,4 +49,8 @@ async def send_api_requests(endpoint: str, data: dict, nodes: List[APINode]):
             # TODO: Report problems to admins
             # We yield 500 response when backend is offline
             result = Response(500)
+        await push_api_request_status(
+            result.status_code,
+            endpoint
+        )
         yield result

apps/tgbot/tgbot/handlers/metrics.py

@@ -0,0 +1,50 @@
+from aioinflux import InfluxDBClient
+from typing import Dict
+from tgbot.config import INFLUX_DB, INFLUX_HOST, INFLUX_PORT, INFLUX_PASSWORD, INFLUX_USERNAME
+
+
+async def push_metric(measurement, tags: Dict, fields: Dict):
+    try:
+        point = {
+            'measurement': measurement,
+            'tags': tags,
+            'fields': fields
+        }
+        async with InfluxDBClient(
+            host=INFLUX_HOST,
+            port=INFLUX_PORT,
+            username=INFLUX_USERNAME,
+            password=INFLUX_PASSWORD,
+            db=INFLUX_DB,
+            mode='async'
+        ) as client:
+            await client.write(point)
+    except Exception as e:
+        print(e)
+        pass
+
+
+async def push_api_request_status(status_code: int, endpoint: str):
+    await push_metric(
+        measurement="bot_api_request",
+        fields=dict(
+            value=1,
+        ),
+        tags=dict(
+            status=status_code,
+            endpoint=endpoint
+        )
+    )
+
+
+async def push_status_metric(status, api_endpoint):
+    await push_metric(
+        measurement="bot_prepared_messages",
+        fields=dict(
+            value=1,
+        ),
+        tags=dict(
+            rsp_status=status,
+            api_endpoint=api_endpoint
+        )
+    )

apps/tgbot/tgbot/middlewares/__init__.py

@@ -0,0 +1 @@
+from tgbot.middlewares.write_command_metric import WriteCommandMetric

apps/tgbot/tgbot/middlewares/write_command_metric.py

@@ -0,0 +1,23 @@
+from aiogram.dispatcher.middlewares import BaseMiddleware
+from aiogram.types import Message
+from tgbot.handlers.metrics import push_metric
+
+
+class WriteCommandMetric(BaseMiddleware):
+
+    def __init__(self):
+        super().__init__()
+
+    async def on_process_message(self, message: Message, data: dict):
+        await push_metric(
+            measurement="bot_processed_messages",
+            fields=dict(
+                telegram_id=message.from_user.id,
+                full_command=message.text,
+                value=1,
+            ),
+            tags=dict(
+                command=message.text.split(" ")[0],
+                type="command"
+            )
+        )