Merge pull request #7 from modern-python/2-feature-add-healthchecks

lesnik512 · web-flow · commit 8283ab7ce46e · 2025-03-17T21:48:04.000+03:00
add health-checks instrument and refactor bootstrap configuration
diff --git a/lite_bootstrap/bootstraps/base.py b/lite_bootstrap/bootstraps/base.py
@@ -2,17 +2,21 @@
 import typing
 
 from lite_bootstrap.instruments.base import BaseInstrument
+from lite_bootstrap.service_config import ServiceConfig
+from lite_bootstrap.types import ApplicationT
 
 
-class BaseBootstrap(abc.ABC):
+class BaseBootstrap(abc.ABC, typing.Generic[ApplicationT]):
+    application: ApplicationT
     instruments: typing.Sequence[BaseInstrument]
+    service_config: ServiceConfig
 
     def bootstrap(self) -> None:
         for one_instrument in self.instruments:
             if one_instrument.is_ready():
-                one_instrument.bootstrap()
+                one_instrument.bootstrap(self.service_config, self.application)
 
     def teardown(self) -> None:
         for one_instrument in self.instruments:
             if one_instrument.is_ready():
-                one_instrument.teardown()
+                one_instrument.teardown(self.application)
diff --git a/lite_bootstrap/bootstraps/fastapi_bootstrap/__init__.py b/lite_bootstrap/bootstraps/fastapi_bootstrap/__init__.py
@@ -4,22 +4,25 @@
 import fastapi
 
 from lite_bootstrap.bootstraps.base import BaseBootstrap
+from lite_bootstrap.bootstraps.fastapi_bootstrap.healthchecks_instrument import FastAPIHealthChecksInstrument
 from lite_bootstrap.bootstraps.fastapi_bootstrap.opentelemetry_instrument import FastAPIOpenTelemetryInstrument
 from lite_bootstrap.bootstraps.fastapi_bootstrap.sentry_instrument import FastAPISentryInstrument
 
 
 __all__ = [
     "FastAPIBootstrap",
+    "FastAPIHealthChecksInstrument",
     "FastAPIOpenTelemetryInstrument",
     "FastAPISentryInstrument",
 ]
 
+from lite_bootstrap.service_config import ServiceConfig
 
-@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
-class FastAPIBootstrap(BaseBootstrap):
-    app: fastapi.FastAPI
-    instruments: typing.Sequence[FastAPIOpenTelemetryInstrument | FastAPISentryInstrument]
 
-    def __post_init__(self) -> None:
-        for one_instrument in self.instruments:
-            one_instrument.app = self.app
+@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
+class FastAPIBootstrap(BaseBootstrap[fastapi.FastAPI]):
+    application: fastapi.FastAPI
+    instruments: typing.Sequence[
+        FastAPIOpenTelemetryInstrument | FastAPISentryInstrument | FastAPIHealthChecksInstrument
+    ]
+    service_config: ServiceConfig
diff --git a/lite_bootstrap/bootstraps/fastapi_bootstrap/healthchecks_instrument.py b/lite_bootstrap/bootstraps/fastapi_bootstrap/healthchecks_instrument.py
@@ -0,0 +1,30 @@
+import dataclasses
+import typing
+
+import fastapi
+
+from lite_bootstrap.instruments.healthchecks_instrument import HealthChecksInstrument, HealthCheckTypedDict
+from lite_bootstrap.service_config import ServiceConfig
+
+
+@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
+class FastAPIHealthChecksInstrument(HealthChecksInstrument):
+    enabled: bool = True
+    path: str = "/health/"
+    include_in_schema: bool = False
+
+    def build_fastapi_health_check_router(self, service_config: ServiceConfig) -> fastapi.APIRouter:
+        fastapi_router: typing.Final = fastapi.APIRouter(
+            tags=["probes"],
+            include_in_schema=self.include_in_schema,
+        )
+
+        @fastapi_router.get(self.path)
+        async def health_check_handler() -> HealthCheckTypedDict:
+            return self.render_health_check_data(service_config)
+
+        return fastapi_router
+
+    def bootstrap(self, service_config: ServiceConfig, application: fastapi.FastAPI | None = None) -> None:
+        if application:
+            application.include_router(self.build_fastapi_health_check_router(service_config))
diff --git a/lite_bootstrap/bootstraps/fastapi_bootstrap/opentelemetry_instrument.py b/lite_bootstrap/bootstraps/fastapi_bootstrap/opentelemetry_instrument.py
@@ -4,6 +4,7 @@
 import fastapi
 
 from lite_bootstrap.instruments.opentelemetry_instrument import OpenTelemetryInstrument
+from lite_bootstrap.service_config import ServiceConfig
 
 
 with contextlib.suppress(ImportError):
@@ -13,16 +14,16 @@
 @dataclasses.dataclass(kw_only=True)
 class FastAPIOpenTelemetryInstrument(OpenTelemetryInstrument):
     excluded_urls: list[str] = dataclasses.field(default_factory=list)
-    app: fastapi.FastAPI = dataclasses.field(init=False)
 
-    def bootstrap(self) -> None:
-        super().bootstrap()
+    def bootstrap(self, service_config: ServiceConfig, application: fastapi.FastAPI | None = None) -> None:
+        super().bootstrap(service_config, application)
         FastAPIInstrumentor.instrument_app(
-            app=self.app,
+            app=application,
             tracer_provider=self.tracer_provider,
             excluded_urls=",".join(self.excluded_urls),
         )
 
-    def teardown(self) -> None:
-        FastAPIInstrumentor.uninstrument_app(self.app)
+    def teardown(self, application: fastapi.FastAPI | None = None) -> None:
+        if application:
+            FastAPIInstrumentor.uninstrument_app(application)
         super().teardown()
diff --git a/lite_bootstrap/bootstraps/fastapi_bootstrap/sentry_instrument.py b/lite_bootstrap/bootstraps/fastapi_bootstrap/sentry_instrument.py
@@ -1,19 +1,7 @@
-import contextlib
 import dataclasses
 
-import fastapi
-
 from lite_bootstrap.instruments.sentry_instrument import SentryInstrument
 
 
-with contextlib.suppress(ImportError):
-    from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
-
-
-@dataclasses.dataclass(kw_only=True)
-class FastAPISentryInstrument(SentryInstrument):
-    app: fastapi.FastAPI = dataclasses.field(init=False)
-
-    def bootstrap(self) -> None:
-        super().bootstrap()
-        self.app.add_middleware(SentryAsgiMiddleware)  # type: ignore[arg-type]
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class FastAPISentryInstrument(SentryInstrument): ...
diff --git a/lite_bootstrap/instruments/base.py b/lite_bootstrap/instruments/base.py
@@ -1,12 +1,13 @@
 import abc
 
+from lite_bootstrap.service_config import ServiceConfig
+from lite_bootstrap.types import ApplicationT
+
 
 class BaseInstrument(abc.ABC):
-    @abc.abstractmethod
-    def bootstrap(self) -> None: ...
+    def bootstrap(self, service_config: ServiceConfig, application: ApplicationT | None = None) -> None: ...  # noqa: B027
 
-    @abc.abstractmethod
-    def teardown(self) -> None: ...
+    def teardown(self, application: ApplicationT | None = None) -> None: ...  # noqa: B027
 
     @abc.abstractmethod
     def is_ready(self) -> bool: ...
diff --git a/lite_bootstrap/instruments/healthchecks_instrument.py b/lite_bootstrap/instruments/healthchecks_instrument.py
@@ -0,0 +1,30 @@
+import dataclasses
+
+import typing_extensions
+
+from lite_bootstrap.instruments.base import BaseInstrument
+from lite_bootstrap.service_config import ServiceConfig
+
+
+class HealthCheckTypedDict(typing_extensions.TypedDict, total=False):
+    service_version: str | None
+    service_name: str | None
+    health_status: bool
+
+
+@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
+class HealthChecksInstrument(BaseInstrument):
+    enabled: bool = True
+    path: str = "/health/"
+    include_in_schema: bool = False
+
+    def is_ready(self) -> bool:
+        return self.enabled
+
+    @staticmethod
+    def render_health_check_data(service_config: ServiceConfig) -> HealthCheckTypedDict:
+        return {
+            "service_version": service_config.service_version,
+            "service_name": service_config.service_name,
+            "health_status": True,
+        }
diff --git a/lite_bootstrap/instruments/opentelemetry_instrument.py b/lite_bootstrap/instruments/opentelemetry_instrument.py
@@ -3,14 +3,16 @@
 import typing
 
 from lite_bootstrap.instruments.base import BaseInstrument
+from lite_bootstrap.service_config import ServiceConfig
+from lite_bootstrap.types import ApplicationT
 
 
 with contextlib.suppress(ImportError):
     from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
     from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore[attr-defined]
     from opentelemetry.sdk import resources
     from opentelemetry.sdk.trace import TracerProvider
-    from opentelemetry.sdk.trace.export import BatchSpanProcessor
+    from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
 
 
 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
@@ -21,37 +23,24 @@ class InstrumentorWithParams:
 
 @dataclasses.dataclass(kw_only=True, slots=True)
 class OpenTelemetryInstrument(BaseInstrument):
-    service_version: str = "1.0.0"
-    service_name: str | None = None
     container_name: str | None = None
     endpoint: str | None = None
     namespace: str | None = None
     insecure: bool = True
     instrumentors: list[InstrumentorWithParams | BaseInstrumentor] = dataclasses.field(default_factory=list)
+    span_exporter: SpanExporter | None = None
 
     tracer_provider: TracerProvider = dataclasses.field(init=False)
 
     def is_ready(self) -> bool:
-        return all(
-            (
-                self.endpoint,
-                self.service_name,
-            ),
-        )
-
-    def teardown(self) -> None:
-        for one_instrumentor in self.instrumentors:
-            if isinstance(one_instrumentor, InstrumentorWithParams):
-                one_instrumentor.instrumentor.uninstrument(**one_instrumentor.additional_params)
-            else:
-                one_instrumentor.uninstrument()
+        return bool(self.endpoint)
 
-    def bootstrap(self) -> None:
+    def bootstrap(self, service_config: ServiceConfig, _: ApplicationT | None = None) -> None:
         attributes = {
-            resources.SERVICE_NAME: self.service_name,
+            resources.SERVICE_NAME: service_config.service_name,
             resources.TELEMETRY_SDK_LANGUAGE: "python",
             resources.SERVICE_NAMESPACE: self.namespace,
-            resources.SERVICE_VERSION: self.service_version,
+            resources.SERVICE_VERSION: service_config.service_version,
             resources.CONTAINER_NAME: self.container_name,
         }
         resource: typing.Final = resources.Resource.create(
@@ -60,7 +49,8 @@ def bootstrap(self) -> None:
         self.tracer_provider = TracerProvider(resource=resource)
         self.tracer_provider.add_span_processor(
             BatchSpanProcessor(
-                OTLPSpanExporter(
+                self.span_exporter
+                or OTLPSpanExporter(
                     endpoint=self.endpoint,
                     insecure=self.insecure,
                 ),
@@ -74,3 +64,10 @@ def bootstrap(self) -> None:
                 )
             else:
                 one_instrumentor.instrument(tracer_provider=self.tracer_provider)
+
+    def teardown(self, _: ApplicationT | None = None) -> None:
+        for one_instrumentor in self.instrumentors:
+            if isinstance(one_instrumentor, InstrumentorWithParams):
+                one_instrumentor.instrumentor.uninstrument(**one_instrumentor.additional_params)
+            else:
+                one_instrumentor.uninstrument()
diff --git a/lite_bootstrap/instruments/sentry_instrument.py b/lite_bootstrap/instruments/sentry_instrument.py
@@ -3,19 +3,20 @@
 import typing
 
 from lite_bootstrap.instruments.base import BaseInstrument
+from lite_bootstrap.service_config import ServiceConfig
+from lite_bootstrap.types import ApplicationT
 
 
 with contextlib.suppress(ImportError):
     import sentry_sdk
     from sentry_sdk.integrations import Integration
 
 
-@dataclasses.dataclass(kw_only=True, slots=True)
+@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
 class SentryInstrument(BaseInstrument):
     dsn: str | None = None
     sample_rate: float = dataclasses.field(default=1.0)
     traces_sample_rate: float | None = None
-    environment: str | None = None
     max_breadcrumbs: int = 15
     max_value_length: int = 16384
     attach_stacktrace: bool = True
@@ -26,12 +27,12 @@ class SentryInstrument(BaseInstrument):
     def is_ready(self) -> bool:
         return bool(self.dsn)
 
-    def bootstrap(self) -> None:
+    def bootstrap(self, service_config: ServiceConfig, _: ApplicationT | None = None) -> None:
         sentry_sdk.init(
             dsn=self.dsn,
             sample_rate=self.sample_rate,
             traces_sample_rate=self.traces_sample_rate,
-            environment=self.environment,
+            environment=service_config.service_environment,
             max_breadcrumbs=self.max_breadcrumbs,
             max_value_length=self.max_value_length,
             attach_stacktrace=self.attach_stacktrace,
@@ -41,4 +42,4 @@ def bootstrap(self) -> None:
         tags: dict[str, str] = self.tags or {}
         sentry_sdk.set_tags(tags)
 
-    def teardown(self) -> None: ...
+    def teardown(self, application: ApplicationT | None = None) -> None: ...
diff --git a/lite_bootstrap/service_config.py b/lite_bootstrap/service_config.py
@@ -0,0 +1,8 @@
+import dataclasses
+
+
+@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
+class ServiceConfig:
+    service_name: str = "micro-service"
+    service_version: str = "1.0.0"
+    service_environment: str | None = None
diff --git a/lite_bootstrap/types.py b/lite_bootstrap/types.py
@@ -0,0 +1,4 @@
+import typing
+
+
+ApplicationT = typing.TypeVar("ApplicationT", bound=typing.Any)
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,10 +1,11 @@
 import typing
 
 import pytest
-from fastapi import APIRouter, FastAPI
-from fastapi.responses import JSONResponse
+from fastapi import FastAPI
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore[attr-defined]
 
+from lite_bootstrap.service_config import ServiceConfig
+
 
 class CustomInstrumentor(BaseInstrumentor):  # type: ignore[misc]
     def instrumentation_dependencies(self) -> typing.Collection[str]:
@@ -16,12 +17,13 @@ def _uninstrument(self, **kwargs: typing.Mapping[str, typing.Any]) -> None:
 
 @pytest.fixture
 def fastapi_app() -> FastAPI:
-    app: typing.Final = FastAPI()
-    router: typing.Final = APIRouter()
+    return FastAPI()
 
-    @router.get("/test")
-    async def for_test_endpoint() -> JSONResponse:
-        return JSONResponse(content={"key": "value"})
 
-    app.include_router(router)
-    return app
+@pytest.fixture
+def service_config() -> ServiceConfig:
+    return ServiceConfig(
+        service_name="microservice",
+        service_version="2.0.0",
+        service_environment="test",
+    )
diff --git a/tests/instruments/__init__.py b/tests/instruments/__init__.py
diff --git a/tests/instruments/test_opentelemetry_instrument.py b/tests/instruments/test_opentelemetry_instrument.py
@@ -0,0 +1,31 @@
+from opentelemetry.sdk.trace.export import ConsoleSpanExporter
+
+from lite_bootstrap.instruments.opentelemetry_instrument import InstrumentorWithParams, OpenTelemetryInstrument
+from lite_bootstrap.service_config import ServiceConfig
+from tests.conftest import CustomInstrumentor
+
+
+def test_opentelemetry_instrument(service_config: ServiceConfig) -> None:
+    opentelemetry = OpenTelemetryInstrument(
+        endpoint="otl",
+        instrumentors=[
+            InstrumentorWithParams(instrumentor=CustomInstrumentor(), additional_params={"key": "value"}),
+            CustomInstrumentor(),
+        ],
+        span_exporter=ConsoleSpanExporter(),
+    )
+    try:
+        opentelemetry.bootstrap(service_config)
+    finally:
+        opentelemetry.teardown()
+
+
+def test_opentelemetry_instrument_empty_instruments(service_config: ServiceConfig) -> None:
+    opentelemetry = OpenTelemetryInstrument(
+        endpoint="otl",
+        span_exporter=ConsoleSpanExporter(),
+    )
+    try:
+        opentelemetry.bootstrap(service_config)
+    finally:
+        opentelemetry.teardown()
diff --git a/tests/instruments/test_sentry_instrument.py b/tests/instruments/test_sentry_instrument.py
diff --git a/tests/test_fastapi_bootstrap.py b/tests/test_fastapi_bootstrap.py
diff --git a/tests/test_opentelemetry_bootstrap.py b/tests/test_opentelemetry_bootstrap.py
diff --git a/tests/test_sentry_bootstrap.py b/tests/test_sentry_bootstrap.py
