diff --git a/lite_bootstrap/bootstraps/litestar_bootstrap.py b/lite_bootstrap/bootstraps/litestar_bootstrap.py new file mode 100644 index 0000000..8ce5cae --- /dev/null +++ b/lite_bootstrap/bootstraps/litestar_bootstrap.py @@ -0,0 +1,75 @@ +import contextlib +import dataclasses +import typing + +from lite_bootstrap.bootstraps.base import BaseBootstrap +from lite_bootstrap.instruments.healthchecks_instrument import HealthChecksInstrument, HealthCheckTypedDict +from lite_bootstrap.instruments.logging_instrument import LoggingInstrument +from lite_bootstrap.instruments.opentelemetry_instrument import OpenTelemetryInstrument +from lite_bootstrap.instruments.sentry_instrument import SentryInstrument +from lite_bootstrap.service_config import ServiceConfig + + +with contextlib.suppress(ImportError): + import litestar + from litestar.config.app import AppConfig + from litestar.contrib.opentelemetry import OpenTelemetryConfig + from opentelemetry.trace import get_tracer_provider + + +@dataclasses.dataclass(kw_only=True, slots=True, frozen=True) +class LitestarHealthChecksInstrument(HealthChecksInstrument): + enabled: bool = True + path: str = "/health/" + include_in_schema: bool = False + + def build_litestar_health_check_router(self, service_config: ServiceConfig) -> litestar.Router: + @litestar.get(media_type=litestar.MediaType.JSON) + async def health_check_handler() -> HealthCheckTypedDict: + return self.render_health_check_data(service_config) + + return litestar.Router( + path=self.path, + route_handlers=[health_check_handler], + tags=["probes"], + include_in_schema=self.include_in_schema, + ) + + def bootstrap(self, service_config: ServiceConfig, app_config: AppConfig | None = None) -> None: + if app_config: + app_config.route_handlers.append(self.build_litestar_health_check_router(service_config)) + + +@dataclasses.dataclass(kw_only=True, frozen=True) +class LitestarLoggingInstrument(LoggingInstrument): ... + + +@dataclasses.dataclass(kw_only=True, frozen=True) +class LitestarOpenTelemetryInstrument(OpenTelemetryInstrument): + excluded_urls: list[str] = dataclasses.field(default_factory=list) + + def bootstrap(self, service_config: ServiceConfig, app_config: AppConfig | None = None) -> None: + super().bootstrap(service_config, app_config) + if app_config: + app_config.middleware.append( + OpenTelemetryConfig( + tracer_provider=get_tracer_provider(), + exclude=self.excluded_urls, + ).middleware, + ) + + +@dataclasses.dataclass(kw_only=True, frozen=True) +class LitestarSentryInstrument(SentryInstrument): ... + + +@dataclasses.dataclass(kw_only=True, slots=True, frozen=True) +class LitestarBootstrap(BaseBootstrap[AppConfig]): + application: AppConfig + instruments: typing.Sequence[ + LitestarOpenTelemetryInstrument + | LitestarSentryInstrument + | LitestarHealthChecksInstrument + | LitestarLoggingInstrument + ] + service_config: ServiceConfig diff --git a/pyproject.toml b/pyproject.toml index 460cec1..6b1deb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,15 @@ fastapi-otl = [ fastapi-all = [ "lite-bootstrap[sentry,otl,logging,fastapi,fastapi-otl]" ] +litestar = [ + "litestar>=2.9", +] +litestar-otl = [ + "opentelemetry-instrumentation-asgi>=0.46b0", +] +litestar-all = [ + "lite-bootstrap[sentry,otl,logging,litestar,litestar-otl]" +] [dependency-groups] dev = [ diff --git a/tests/test_litestar_bootstrap.py b/tests/test_litestar_bootstrap.py new file mode 100644 index 0000000..701cd35 --- /dev/null +++ b/tests/test_litestar_bootstrap.py @@ -0,0 +1,56 @@ +import litestar +import structlog +from litestar import status_codes +from litestar.config.app import AppConfig +from litestar.testing import TestClient +from opentelemetry.sdk.trace.export import ConsoleSpanExporter + +from lite_bootstrap.bootstraps.litestar_bootstrap import ( + LitestarBootstrap, + LitestarHealthChecksInstrument, + LitestarLoggingInstrument, + LitestarOpenTelemetryInstrument, + LitestarSentryInstrument, +) +from lite_bootstrap.service_config import ServiceConfig +from tests.conftest import CustomInstrumentor + + +logger = structlog.getLogger(__name__) + + +def test_litestar_bootstrap(service_config: ServiceConfig) -> None: + app_config = AppConfig() + litestar_bootstrap = LitestarBootstrap( + application=app_config, + service_config=service_config, + instruments=[ + LitestarOpenTelemetryInstrument( + endpoint="otl", + instrumentors=[CustomInstrumentor()], + span_exporter=ConsoleSpanExporter(), + ), + LitestarSentryInstrument( + dsn="https://testdsn@localhost/1", + ), + LitestarHealthChecksInstrument( + path="/health/", + ), + LitestarLoggingInstrument(logging_buffer_capacity=0), + ], + ) + litestar_bootstrap.bootstrap() + application = litestar.Litestar.from_config(app_config) + logger.info("testing logging", key="value") + + try: + with TestClient(app=application) as async_client: + response = async_client.get("/health/") + assert response.status_code == status_codes.HTTP_200_OK + assert response.json() == { + "health_status": True, + "service_name": "microservice", + "service_version": "2.0.0", + } + finally: + litestar_bootstrap.teardown()