PyObservability 0.1.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

pyobservability/config/enums.py

@@ -4,3 +4,8 @@ from enum import StrEnum
 class APIEndpoints(StrEnum):
     root = "/"
     ws = "/ws"
+
+
+class Log(StrEnum):
+    file = "file"
+    stdout = "stdout"

pyobservability/config/settings.py

@@ -2,13 +2,49 @@ import json
 import os
 import pathlib
 import socket
-from typing import List
+from typing import Any, Dict, List
 
 import yaml
-from pydantic import BaseModel, Field, HttpUrl, PositiveInt
+from pydantic import BaseModel, Field, FilePath, HttpUrl, PositiveInt
 from pydantic.aliases import AliasChoices
 from pydantic_settings import BaseSettings
 
+from pyobservability.config import enums
+
+
+def detailed_log_config(filename: str | None = None, debug: bool = False) -> Dict[str, Any]:
+    if filename:
+        log_handler = {
+            "class": "logging.FileHandler",
+            "formatter": "default",
+            "filename": filename,
+            "mode": "a",
+        }
+    else:
+        log_handler = {
+            "class": "logging.StreamHandler",
+            "formatter": "default",
+            "stream": "ext://sys.stdout",
+        }
+    level = "DEBUG" if debug else "INFO"
+    return {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "default": {
+                "format": "%(asctime)s - %(levelname)s - [%(module)s:%(lineno)d] - %(funcName)s - %(message)s",
+                "datefmt": "%b-%d-%Y %I:%M:%S %p",
+            }
+        },
+        "handlers": {"default": log_handler},
+        "loggers": {
+            "uvicorn": {"handlers": ["default"], "level": level},
+            "uvicorn.error": {"handlers": ["default"], "level": level, "propagate": False},
+            "uvicorn.access": {"handlers": ["default"], "level": level, "propagate": False},
+        },
+        "root": {"handlers": ["default"], "level": level},
+    }
+
 
 class PydanticEnvConfig(BaseSettings):
     """Pydantic BaseSettings with custom order for loading environment variables.
@@ -63,6 +99,10 @@ class EnvConfig(PydanticEnvConfig):
     targets: List[MonitorTarget] = Field(..., validation_alias=alias_choices("TARGETS"))
     interval: PositiveInt = Field(3, validation_alias=alias_choices("INTERVAL"))
 
+    log: enums.Log | None = None
+    debug: bool = False
+    log_config: Dict[str, Any] | FilePath | None = None
+
     username: str | None = Field(None, validation_alias=alias_choices("USERNAME"))
     password: str | None = Field(None, validation_alias=alias_choices("PASSWORD"))
 
@@ -124,3 +164,4 @@ def env_loader(**kwargs) -> EnvConfig:
 
 
 env: EnvConfig
+targets_by_url: Dict[str, Dict[str, str]]

pyobservability/main.py

@@ -1,17 +1,19 @@
 import logging
+import os
 import pathlib
 import warnings
+from datetime import datetime
 
-import time
 import uiauth
 import uvicorn
-from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, Request
 from fastapi.routing import APIRoute, APIWebSocketRoute
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 
 from pyobservability.config import enums, settings
-from pyobservability.monitor import Monitor
+from pyobservability.transport import websocket_endpoint
+from pyobservability.version import __version__
 
 LOGGER = logging.getLogger("uvicorn.default")
 
@@ -33,37 +35,9 @@ async def index(request: Request):
     Args:
         request: FastAPI request object.
     """
-    return templates.TemplateResponse("index.html", {"request": request, "targets": settings.env.targets})
-
-
-async def websocket_endpoint(websocket: WebSocket):
-    """Websocket endpoint to render the metrics.
-
-    Args:
-        websocket: FastAPI websocket object.
-    """
-    monitor = Monitor(targets=settings.env.targets, poll_interval=settings.env.interval)
-    await websocket.accept()
-    await monitor.start()
-    q = monitor.subscribe()
-    try:
-        while True:
-            start = time.time()
-            payload = await q.get()
-            end = time.time()
-            nodes = [d['name'] for d in payload["data"]]
-            LOGGER.debug("Payload generated in %s - %d %s", end - start, len(nodes), nodes)
-            # send as JSON text
-            await websocket.send_json(payload)
-    except WebSocketDisconnect:
-        monitor.unsubscribe(q)
-    except Exception:
-        monitor.unsubscribe(q)
-        try:
-            await websocket.close()
-        except Exception as err:
-            LOGGER.warning(err)
-    await monitor.stop()
+    return templates.TemplateResponse(
+        "index.html", {"request": request, "targets": settings.env.targets, "version": __version__}
+    )
 
 
 def include_routes():
@@ -72,6 +46,7 @@ def include_routes():
             app=PyObservability,
             username=settings.env.username,
             password=settings.env.password,
+            custom_logger=LOGGER,
             params=[
                 uiauth.Parameters(
                     path=enums.APIEndpoints.root,
@@ -106,10 +81,23 @@ def include_routes():
 def start(**kwargs):
     settings.env = settings.env_loader(**kwargs)
     settings.env.targets = [{k: str(v) for k, v in target.model_dump().items()} for target in settings.env.targets]
+    settings.targets_by_url = {t["base_url"]: t for t in settings.env.targets}
     include_routes()
     uvicorn_args = dict(
         host=settings.env.host,
         port=settings.env.port,
         app=PyObservability,
     )
+    if log := settings.env.log:
+        if log == enums.Log.stdout:
+            uvicorn_args["log_config"] = settings.detailed_log_config(debug=settings.env.debug)
+        else:
+            log_file = datetime.now().strftime(os.path.join("logs", "pyobservability_%d-%m-%Y.log"))
+            os.makedirs("logs", exist_ok=True)
+            uvicorn_args["log_config"] = settings.detailed_log_config(filename=log_file, debug=settings.env.debug)
+    # log_config will take precedence if both log and log_config are set
+    if settings.env.log_config:
+        uvicorn_args["log_config"] = (
+            settings.env.log_config if isinstance(settings.env.log_config, dict) else str(settings.env.log_config)
+        )
     uvicorn.run(**uvicorn_args)

pyobservability/monitor.py

@@ -2,120 +2,175 @@ import asyncio
 import json
 import logging
 from asyncio import CancelledError
-from typing import Dict, List
+from collections.abc import Generator
+from typing import Any, Dict, List
 
 import aiohttp
 
+from pyobservability.config import settings
+
 LOGGER = logging.getLogger("uvicorn.default")
 OBS_PATH = "/observability"
 
 
-class Monitor:
-    def __init__(self, targets: List[Dict[str, str]], poll_interval: float):
-        self.targets = targets
-        self.poll_interval = poll_interval
-        self.sessions: Dict[str, aiohttp.ClientSession] = {}
-        self._ws_subscribers: List[asyncio.Queue] = []
-        self._task: asyncio.Task | None = None
-        self._stop = asyncio.Event()
-
-    ############################################################################
-    # LIFECYCLE
-    ############################################################################
-    async def start(self):
-        self._tasks = []
-
-        for target in self.targets:
-            base = target["base_url"]
-            name = target.get("name")
-            apikey = target.get("apikey")
-
-            session = aiohttp.ClientSession()
-            self.sessions[base] = session
+def refine_service(service_list: List[Dict[str, Any]]) -> Generator[Dict[str, Dict[str, str]]]:
+    for service in service_list:
+        service["memory"] = dict(rss=service.get("memory", {}).get("rss"), vms=service.get("memory", {}).get("vms"))
+        service["cpu"] = dict(
+            user=round(service.get("cpu", {}).get("user", 0), 2),
+            system=round(service.get("cpu", {}).get("system", 0), 2),
+        )
+        yield service
 
-            task = asyncio.create_task(self._stream_target(name, base, session, apikey))
-            self._tasks.append(task)
 
-    async def stop(self):
-        self._stop.set()
+class Monitor:
+    def __init__(self, target: Dict[str, str]):
+        self.name = target["name"]
+        self.base_url = target["base_url"]
+        self.apikey = target["apikey"]
+        self.flags = {"all_services": False}
 
-        for task in self._tasks:
-            task.cancel()
-            try:
-                await task
-            except CancelledError:
-                pass
+        self.session: aiohttp.ClientSession | None = None
+        self._task: asyncio.Task | None = None
+        self._stop = asyncio.Event()
 
-        for session in self.sessions.values():
-            await session.close()
+        self._ws_subscribers = []  # list of asyncio.Queue
 
-    ############################################################################
+    # ------------------------------
     # SUBSCRIBE / UNSUBSCRIBE
-    ############################################################################
+    # ------------------------------
     def subscribe(self) -> asyncio.Queue:
         q = asyncio.Queue(maxsize=10)
         self._ws_subscribers.append(q)
         return q
 
     def unsubscribe(self, q: asyncio.Queue):
-        try:
+        if q in self._ws_subscribers:
             self._ws_subscribers.remove(q)
-        except ValueError:
-            pass
-
-    ############################################################################
-    # FETCH OBSERVABILITY FOR SINGLE TARGET
-    ############################################################################
-    async def _fetch_observability(self, session, base_url, apikey):
-        url = base_url.rstrip("/") + OBS_PATH + f"?interval={self.poll_interval}"
-        headers = {"accept": "application/json"}
-        if apikey:
-            headers["Authorization"] = f"Bearer {apikey}"
-
-        try:
-            async with session.get(url, headers=headers, timeout=None) as resp:
-                if resp.status != 200:
-                    LOGGER.debug("Exception - [%d]: %s", resp.status, await resp.text())
-                    return
-                async for line in resp.content:
-                    line = line.decode().strip()
-                    if not line:
-                        continue
+
+    # ------------------------------
+    # START / STOP
+    # ------------------------------
+    async def start(self):
+        if self._task:
+            return  # already running
+
+        self._stop.clear()
+        self.session = aiohttp.ClientSession()
+
+        self._task = asyncio.create_task(self._stream_target())
+
+    async def stop(self):
+        self._stop.set()
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except CancelledError:
+                pass
+            self._task = None
+
+        if self.session:
+            await self.session.close()
+            self.session = None
+
+    async def update_flags(self, **kwargs):
+        for k, v in kwargs.items():
+            if k in self.flags:
+                self.flags[k] = v
+
+        # restart stream with new params
+        await self._restart_stream()
+
+    async def _restart_stream(self):
+        await self.stop()
+        await self.start()
+
+    # ------------------------------
+    # FETCH STREAM
+    # ------------------------------
+    async def _fetch_stream(self):
+        query = f"?interval={settings.env.interval}"
+        if self.flags["all_services"]:
+            query += "&all_services=true"
+        url = self.base_url.rstrip("/") + OBS_PATH + query
+        headers = {"Accept": "application/json", "Authorization": f"Bearer {self.apikey}"}
+
+        async with self.session.get(
+            url, headers=headers, timeout=aiohttp.ClientTimeout(total=None, connect=3, sock_read=None, sock_connect=3)
+        ) as resp:
+            if resp.status != 200:
+                LOGGER.error("Bad response [%d] from %s", resp.status, url)
+                return
+
+            async for raw in resp.content:
+                line = raw.decode().strip()
+                if not line:
+                    continue
+
+                try:
+                    parsed = json.loads(line)
                     try:
-                        payload = json.loads(line)
-                        # yield each record instead of returning
-                        yield payload
-                    except json.JSONDecodeError as err:
-                        LOGGER.debug("JSON decode error: %s | line=%s", err, line)
-        except Exception as err:
-            LOGGER.debug("Exception: %s", err)
-            return
-
-    ############################################################################
-    # STREAM A SPECIFIC TARGET - SEQUENTIAL OVER ALL TARGETS
-    ############################################################################
-    async def _stream_target(self, name, base, session, apikey):
+                        if service_stats := parsed.get("service_stats"):
+                            parsed["service_stats"] = list(refine_service(service_stats))
+                    except Exception as error:
+                        LOGGER.error("Received [%s: %s] when parsing services for %s", type(error), error, self.name)
+                    yield parsed
+                except json.JSONDecodeError:
+                    LOGGER.debug("Bad JSON from %s: %s", self.base_url, line)
+
+    # ------------------------------
+    # STREAM LOOP
+    # ------------------------------
+    async def _stream_target(self):
+        errors = {}
         while not self._stop.is_set():
             try:
-                async for payload in self._fetch_observability(session, base, apikey):
+                async for payload in self._fetch_stream():
                     result = {
                         "type": "metrics",
                         "ts": asyncio.get_event_loop().time(),
-                        "data": [{"name": name, "base_url": base, "metrics": payload}],
+                        "data": [
+                            {
+                                "name": self.base_url,
+                                "base_url": self.base_url,
+                                "metrics": payload,
+                            }
+                        ],
                     }
 
-                    # The queue gets full when subscribers aren’t reading fast enough
-                    # [OR]
-                    # The monitor produces new metrics faster than the queue’s fixed maxsize of 10 can be consumed
                     for q in list(self._ws_subscribers):
                         try:
                             q.put_nowait(result)
                         except asyncio.QueueFull:
+                            _ = q.get_nowait()
+                            q.put_nowait(result)
+            except Exception as err:
+                LOGGER.debug("Stream error for %s: %s", self.base_url, err)
+                if errors.get(self.base_url):
+                    if errors[self.base_url] < 10:
+                        errors[self.base_url] += 1
+                    else:
+                        LOGGER.error(err.with_traceback())
+                        LOGGER.error("%s exceeded error threshold.", self.base_url)
+
+                        # notify subscribers before stopping
+                        error_msg = {
+                            "type": "error",
+                            "base_url": self.base_url,
+                            "message": f"{self.name!r} is unreachable.",
+                        }
+
+                        for q in list(self._ws_subscribers):
                             try:
+                                q.put_nowait(error_msg)
+                            except asyncio.QueueFull as warn:
+                                LOGGER.warning(warn)
                                 _ = q.get_nowait()
                                 q.put_nowait(result)
-                            except Exception as debug:
-                                LOGGER.debug(f"Queue full: {debug}")
-            except Exception as debug:
-                LOGGER.debug(f"Stream failed for {base}, retrying: {debug}")
+                        await self.stop()
+                        return
+                else:
+                    errors[self.base_url] = 1
+
             await asyncio.sleep(1)