Feature/bytes openmetrics

.pre-commit-config.yaml

@@ -50,7 +50,7 @@ repos:
     rev: v1.2.0
     hooks:
       - id: mypy
-        additional_dependencies: ['types-PyYAML', 'types-requests', 'pydantic', 'pynacl']
+        additional_dependencies: ['types-PyYAML', 'types-requests', 'types-cachetools', 'pydantic', 'pynacl']
         exclude: |
           (?x)(
           ^boefjes/ |

bytes/.env-dist

@@ -29,3 +29,4 @@ QUEUE_URI=
 # Optional environment variables
 BYTES_LOG_FILE=  # Optional file with Bytes logs.
 BYTES_DATA_DIR=  # Root for all the data. A change means that you no longer have access to old data unless you move it!
+BYTES_METRICS_TTL_SECONDS=0  # The time to cache slow queries performed in the metrics endpoint.

bytes/README.md

@@ -85,6 +85,7 @@ QUEUE_URI=
 # Optional environment variables
 BYTES_LOG_FILE=  # Optional file with Bytes logs.
 BYTES_DATA_DIR=  # Root for all the data. A change means that you no longer have access to old data unless you move it!
+BYTES_METRICS_TTL_SECONDS=0  # The time to cache slow queries performed in the metrics endpoint.
 ```
 
 Most of these are self-explanatory, but a few sets of variables require more explanation.
@@ -130,6 +131,14 @@ KAT_PRIVATE_KEY_B64=""
 VWS_PUBLIC_KEY_B64=""
 ```
 
+### Observability
+
+Bytes exposes a `/metrics` endpoint for basic application level observability,
+such as the amount of organizations and the amount of raw files per organization.
+Another important component to monitor is the disk usage of Bytes.
+It is recommended to install [node exporter](https://prometheus.io/docs/guides/node-exporter/) to keep track of this.
+
+
 ## Design
 
 We now include two levels of design, according to the [C4 model](https://c4model.com/).
@@ -147,7 +156,6 @@ graph
     Scheduler["Scheduler<br/><i>Software System"]
     Boefjes["Boefjes<br/><i>Python App"]
 
-
     Boefjes -- GET/POST Raw/Meta --> Bytes
     User -- Interacts with --> Rocky
     Rocky -- GET/POST Raw/Meta --> Bytes
@@ -238,3 +246,13 @@ To export raw SQL from the SQLAlchemy migration files, run the following target
 ```shell
 $ make sql rev1=0003 rev2=0004 > sql_migrations/0004_change_x_to_y_add_column_z.sql
 ```
+
+
+## Production
+
+
+### Performance tuning
+
+Bytes caches some metrics for performance, but the default is not to cache these queries.
+It is recommended to tune the `BYTES_METRICS_TTL_SECONDS` variable to on the amount of calls to the `/metrics` endpoint.
+As a guideline, add at least 10 seconds to the cache for every million of raw files in the database.

bytes/bytes/api/metrics.py

@@ -0,0 +1,46 @@
+import logging
+from typing import Dict
+
+from cachetools import TTLCache, cached
+from prometheus_client import CollectorRegistry, Gauge
+
+from bytes.config import get_settings
+from bytes.repositories.meta_repository import MetaDataRepository
+
+collector_registry = CollectorRegistry()
+
+
+bytes_database_organizations_total = Gauge(
+    name="bytes_database_organizations_total",
+    documentation="Total amount of organizations in the bytes database.",
+    registry=collector_registry,
+)
+bytes_database_raw_files_total = Gauge(
+    name="bytes_database_raw_files_total",
+    documentation="Total amount of raw files in the bytes database.",
+    registry=collector_registry,
+    labelnames=["organization_id"],
+)
+
+logger = logging.getLogger(__name__)
+
+
+def ignore_arguments_key(meta_repository: MetaDataRepository):
+    return ""
+
+
+@cached(cache=TTLCache(maxsize=1, ttl=get_settings().bytes_metrics_ttl_seconds), key=ignore_arguments_key)
+def cached_counts_per_organization(meta_repository: MetaDataRepository) -> Dict[str, int]:
+    logger.debug("Metrics cache miss, ttl set to %s seconds", get_settings().bytes_metrics_ttl_seconds)
+
+    return meta_repository.get_raw_file_count_per_organization()
+
+
+def get_registry(meta_repository: MetaDataRepository) -> CollectorRegistry:
+    counts_per_organization = cached_counts_per_organization(meta_repository)
+    bytes_database_organizations_total.set(len(counts_per_organization))
+
+    for organization_id, count in counts_per_organization.items():
+        bytes_database_raw_files_total.labels(organization_id).set(count)
+
+    return collector_registry

bytes/bytes/api/root.py

@@ -1,16 +1,20 @@
 import logging
 from typing import Any, List, Optional, Union
 
+import prometheus_client
 from fastapi import APIRouter, Depends
 from fastapi.exceptions import RequestValidationError
-from fastapi.responses import RedirectResponse
+from fastapi.responses import RedirectResponse, Response
 from fastapi.security import OAuth2PasswordRequestForm
 from pydantic import BaseModel, ValidationError
 from starlette.requests import Request
 from starlette.responses import JSONResponse
 from starlette.status import HTTP_422_UNPROCESSABLE_ENTITY
 
-from bytes.auth import TokenResponse, get_access_token
+from bytes.api.metrics import get_registry
+from bytes.auth import TokenResponse, authenticate_token, get_access_token
+from bytes.database.sql_meta_repository import create_meta_data_repository
+from bytes.repositories.meta_repository import MetaDataRepository
 from bytes.version import __version__
 
 router = APIRouter()
@@ -46,6 +50,14 @@ def root() -> ServiceHealth:
     return bytes_health
 
 
+@router.get("/metrics", dependencies=[Depends(authenticate_token)])
+def metrics(meta_repository: MetaDataRepository = Depends(create_meta_data_repository)):
+    collector_registry = get_registry(meta_repository)
+    data = prometheus_client.generate_latest(collector_registry)
+
+    return Response(media_type="text/plain", content=data)
+
+
 @router.post("/token", response_model=TokenResponse)
 def login(form_data: OAuth2PasswordRequestForm = Depends()) -> TokenResponse:
     access_token, expire_time = get_access_token(form_data)

bytes/bytes/config.py

@@ -35,6 +35,8 @@ class Settings(BaseSettings):
     kat_private_key_b64: str = ""
     vws_public_key_b64: str = ""
 
+    bytes_metrics_ttl_seconds: int = 0
+
 
 @lru_cache()
 def get_settings() -> Settings:

bytes/bytes/database/sql_meta_repository.py

@@ -1,7 +1,8 @@
 import logging
 import uuid
-from typing import Iterator, List, Optional, Type
+from typing import Dict, Iterator, List, Optional, Type
 
+from sqlalchemy import func
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session, sessionmaker
 
@@ -157,6 +158,15 @@ def has_raw(self, boefje_meta: BoefjeMeta, mime_types: List[MimeType]) -> bool:
 
         return count > 0
 
+    def get_raw_file_count_per_organization(self) -> Dict[str, int]:
+        query = (
+            self.session.query(BoefjeMetaInDB.organization, func.count(RawFileInDB.id))
+            .join(BoefjeMetaInDB)
+            .group_by(BoefjeMetaInDB.organization)
+        )
+
+        return {organization_id: count for organization_id, count in query}
+
     def _to_raw(self, raw_file_in_db: RawFileInDB) -> RawData:
         boefje_meta = to_boefje_meta(raw_file_in_db.boefje_meta)
         data = self.raw_repository.get_raw(raw_file_in_db.id, boefje_meta)

bytes/bytes/repositories/meta_repository.py

@@ -45,10 +45,7 @@ def save_boefje_meta(self, boefje_meta: BoefjeMeta) -> None:
     def get_boefje_meta_by_id(self, boefje_meta_id: str) -> BoefjeMeta:
         raise NotImplementedError()
 
-    def get_boefje_meta(
-        self,
-        query_filter: BoefjeMetaFilter,
-    ) -> List[BoefjeMeta]:
+    def get_boefje_meta(self, query_filter: BoefjeMetaFilter) -> List[BoefjeMeta]:
         raise NotImplementedError()
 
     def save_normalizer_meta(self, normalizer_meta: NormalizerMeta) -> None:
@@ -68,3 +65,6 @@ def has_raw(self, boefje_meta: BoefjeMeta, mime_types: List[MimeType]) -> bool:
 
     def get_raw(self, raw_id: str) -> RawData:
         raise NotImplementedError()
+
+    def get_raw_file_count_per_organization(self) -> Dict[str, int]:
+        raise NotImplementedError()

bytes/requirements-dev.txt

@@ -5,6 +5,7 @@ astroid==2.12.12
 attrs==23.1.0
 bcrypt==4.0.1
 black==23.1.0
+cachetools==5.3.0
 certifi==2022.12.7
 cffi==1.15.1
 charset-normalizer==3.1.0
@@ -32,6 +33,7 @@ pika==1.3.1
 platformdirs==3.0.0
 pluggy==1.0.0
 pre-commit==3.2.2
+prometheus-client==0.16.0
 psycopg2==2.9.6
 py==1.11.0
 pyasn1==0.4.8

bytes/requirements.txt

@@ -2,6 +2,7 @@ alembic==1.8.1
 anyio==3.6.2
 asgiref==3.5.2
 bcrypt==4.0.1
+cachetools==5.3.0
 certifi==2022.12.7
 cffi==1.15.1
 charset-normalizer==3.1.0
@@ -18,6 +19,7 @@ mako==1.2.4
 markupsafe==2.1.2
 passlib[bcrypt]==1.7.4
 pika==1.3.1
+prometheus-client==0.16.0
 psycopg2==2.9.6
 pyasn1==0.4.8
 pycparser==2.21

bytes/tests/client.py

@@ -3,7 +3,7 @@
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import requests
-from requests.models import HTTPError
+from requests.exceptions import HTTPError
 
 from bytes.models import BoefjeMeta, NormalizerMeta
 from bytes.repositories.meta_repository import BoefjeMetaFilter, RawDataFilter
@@ -70,6 +70,14 @@ def _get_token(self) -> str:
 
         return str(response.json()["access_token"])
 
+    @retry_with_login
+    def get_metrics(self) -> bytes:
+        response = self._session.get("/metrics", headers=self.headers)
+
+        self._verify_response(response)
+
+        return response.content
+
     @retry_with_login
     def save_boefje_meta(self, boefje_meta: BoefjeMeta) -> None:
         response = self._session.post("/bytes/boefje_meta", data=boefje_meta.json(), headers=self.headers)

bytes/tests/integration/test_bytes_api_client.py → bytes/tests/integration/test_bytes_api.py

@@ -2,6 +2,7 @@
 
 import pytest
 import requests
+from prometheus_client.parser import text_string_to_metric_families
 from requests import HTTPError
 from tests.client import BytesAPIClient
 from tests.loading import get_boefje_meta, get_normalizer_meta
@@ -17,6 +18,57 @@ def test_login(bytes_api_client: BytesAPIClient) -> None:
     assert "bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" in bytes_api_client.headers["Authorization"]
 
 
+def test_metrics(bytes_api_client: BytesAPIClient) -> None:
+    metrics = bytes_api_client.get_metrics()
+
+    metrics = list(text_string_to_metric_families(metrics.decode()))
+    assert len(metrics) == 2
+    assert metrics[0].name == "bytes_database_organizations_total"
+    assert len(metrics[0].samples) == 1
+    assert metrics[1].name == "bytes_database_raw_files_total"
+    assert len(metrics[1].samples) == 0
+
+    database_organizations, database_files = metrics
+
+    assert database_organizations.samples[0].value == 0.0
+
+    boefje_meta = get_boefje_meta()
+    bytes_api_client.save_boefje_meta(boefje_meta)
+    bytes_api_client.save_raw(boefje_meta.id, b"test 123")
+    bytes_api_client.save_raw(boefje_meta.id, b"test 12334", ["text/boefje"])
+
+    metrics = bytes_api_client.get_metrics()
+    metrics = list(text_string_to_metric_families(metrics.decode()))
+    assert len(metrics[0].samples) == 1
+    assert len(metrics[1].samples) == 1
+    database_organizations, database_files = metrics
+
+    assert database_organizations.samples[0].value == 1.0
+
+    assert database_files.samples[0].labels["organization_id"] == "test"
+    assert database_files.samples[0].value == 2.0
+
+    boefje_meta = get_boefje_meta()
+    boefje_meta.id = str(uuid.uuid4())
+    boefje_meta.organization = "test2"
+    bytes_api_client.save_boefje_meta(boefje_meta)
+    bytes_api_client.save_raw(boefje_meta.id, b"test 123")
+
+    metrics = bytes_api_client.get_metrics()
+    metrics = list(text_string_to_metric_families(metrics.decode()))
+    assert len(metrics[0].samples) == 1
+    assert len(metrics[1].samples) == 2
+    database_organizations, database_files = metrics
+
+    assert database_organizations.samples[0].value == 2.0
+
+    assert len(database_files.samples) == 2
+    assert database_files.samples[0].labels["organization_id"] == "test"
+    assert database_files.samples[0].value == 2.0
+    assert database_files.samples[1].labels["organization_id"] == "test2"
+    assert database_files.samples[1].value == 1.0
+
+
 def test_boefje_meta(bytes_api_client: BytesAPIClient) -> None:
     boefje_meta = get_boefje_meta()
     bytes_api_client.save_boefje_meta(boefje_meta)

bytes/tests/integration/test_meta_repository.py

@@ -146,6 +146,8 @@ def test_save_boefje_meta_hash(meta_repository: SQLMetaDataRepository) -> None:
     non_empty_raws = meta_repository.get_raws(query_filter)
     assert len(non_empty_raws) == 1
 
+    assert meta_repository.get_raw_file_count_per_organization() == {"test": 1}
+
 
 def test_save_normalizer_meta(meta_repository: SQLMetaDataRepository) -> None:
     with meta_repository:

bytes/tests/seed_database.py

@@ -0,0 +1,29 @@
+import uuid
+
+from bytes.database.sql_meta_repository import create_meta_data_repository
+from bytes.models import RawData
+from tests.loading import get_boefje_meta
+
+
+def seed():
+    repository = next(create_meta_data_repository())
+    number_of_raw_files, chunk_size = int(2e6), 1000
+
+    for i in range(int(number_of_raw_files / chunk_size)):
+        with repository:
+            for j in range(chunk_size):
+                raw = b"asdf           ---                \n\n testdata" + str(i).encode()
+
+                boefje_meta = get_boefje_meta(meta_id=str(uuid.uuid4()))
+                boefje_meta.organization = ["a", "b", "c", "d", "e", "f", "g"][i % 7]
+
+                repository.save_boefje_meta(boefje_meta)
+                repository.save_raw(RawData(value=raw, boefje_meta=boefje_meta))
+
+            print(f"Committing chunk {i}")
+
+
+if __name__ == "__main__":
+    """This script is just a helper to generate a large set of objects to test performance with."""
+
+    seed()