Skip to content

Commit 7406976

Browse files
authored
Merge pull request #23 from modern-python/16-feature-add-faststream-integration
add faststream bootstrapper
2 parents e2d24ac + c89f3f1 commit 7406976

File tree

5 files changed

+241
-4
lines changed

5 files changed

+241
-4
lines changed

lite_bootstrap/__init__.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
from lite_bootstrap.bootstrappers.fastapi_bootstrapper import FastAPIBootstrapper, FastAPIConfig
2+
from lite_bootstrap.bootstrappers.faststream_bootstrapper import FastStreamBootstrapper, FastStreamConfig
23
from lite_bootstrap.bootstrappers.free_bootstrapper import FreeBootstrapper, FreeBootstrapperConfig
34
from lite_bootstrap.bootstrappers.litestar_bootstrapper import LitestarBootstrapper, LitestarConfig
45

56

67
__all__ = [
78
"FastAPIBootstrapper",
89
"FastAPIConfig",
10+
"FastStreamBootstrapper",
11+
"FastStreamConfig",
912
"FreeBootstrapper",
1013
"FreeBootstrapperConfig",
1114
"LitestarBootstrapper",

lite_bootstrap/bootstrappers/fastapi_bootstrapper.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -112,8 +112,5 @@ class FastAPIBootstrapper(BaseBootstrapper[fastapi.FastAPI]):
112112
bootstrap_config: FastAPIConfig
113113
__slots__ = "bootstrap_config", "instruments"
114114

115-
def __init__(self, bootstrap_config: FastAPIConfig) -> None:
116-
super().__init__(bootstrap_config)
117-
118115
def _prepare_application(self) -> fastapi.FastAPI:
119116
return self.bootstrap_config.application
Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
from __future__ import annotations
2+
import contextlib
3+
import dataclasses
4+
import json
5+
import typing
6+
7+
from lite_bootstrap.bootstrappers.base import BaseBootstrapper
8+
from lite_bootstrap.instruments.healthchecks_instrument import HealthChecksConfig, HealthChecksInstrument
9+
from lite_bootstrap.instruments.logging_instrument import LoggingConfig, LoggingInstrument
10+
from lite_bootstrap.instruments.opentelemetry_instrument import OpentelemetryConfig, OpenTelemetryInstrument
11+
from lite_bootstrap.instruments.prometheus_instrument import PrometheusConfig, PrometheusInstrument
12+
from lite_bootstrap.instruments.sentry_instrument import SentryConfig, SentryInstrument
13+
14+
15+
with contextlib.suppress(ImportError):
16+
import faststream
17+
import prometheus_client
18+
from faststream.asgi import AsgiFastStream, AsgiResponse
19+
from faststream.asgi import get as handle_get
20+
from faststream.broker.core.usecase import BrokerUsecase
21+
from opentelemetry.metrics import Meter, MeterProvider
22+
from opentelemetry.trace import TracerProvider, get_tracer_provider
23+
24+
25+
@typing.runtime_checkable
26+
class FastStreamTelemetryMiddlewareProtocol(typing.Protocol):
27+
def __init__(
28+
self,
29+
*,
30+
tracer_provider: TracerProvider | None = None,
31+
meter_provider: MeterProvider | None = None,
32+
meter: Meter | None = None,
33+
) -> None: ...
34+
def __call__(self, msg: typing.Any | None) -> faststream.BaseMiddleware: ... # noqa: ANN401
35+
36+
37+
@typing.runtime_checkable
38+
class FastStreamPrometheusMiddlewareProtocol(typing.Protocol):
39+
def __init__(
40+
self,
41+
*,
42+
registry: prometheus_client.CollectorRegistry,
43+
app_name: str = ...,
44+
metrics_prefix: str = "faststream",
45+
received_messages_size_buckets: typing.Sequence[float] | None = None,
46+
) -> None: ...
47+
def __call__(self, msg: typing.Any | None) -> faststream.BaseMiddleware: ... # noqa: ANN401
48+
49+
50+
@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
51+
class FastStreamConfig(HealthChecksConfig, LoggingConfig, OpentelemetryConfig, PrometheusConfig, SentryConfig):
52+
application: AsgiFastStream = dataclasses.field(default_factory=AsgiFastStream)
53+
broker: BrokerUsecase[typing.Any, typing.Any] | None = None
54+
opentelemetry_middleware_cls: type[FastStreamTelemetryMiddlewareProtocol] | None = None
55+
prometheus_middleware_cls: type[FastStreamPrometheusMiddlewareProtocol] | None = None
56+
57+
58+
@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
59+
class FastStreamHealthChecksInstrument(HealthChecksInstrument):
60+
bootstrap_config: FastStreamConfig
61+
62+
def bootstrap(self) -> None:
63+
@handle_get
64+
async def check_health(_: object) -> AsgiResponse:
65+
return (
66+
AsgiResponse(
67+
json.dumps(self.render_health_check_data()).encode(), 200, headers={"content-type": "text/plain"}
68+
)
69+
if await self._define_health_status()
70+
else AsgiResponse(b"Service is unhealthy", 500, headers={"content-type": "application/json"})
71+
)
72+
73+
self.bootstrap_config.application.mount(self.bootstrap_config.health_checks_path, check_health)
74+
75+
async def _define_health_status(self) -> bool:
76+
if not self.bootstrap_config.application or not self.bootstrap_config.application.broker:
77+
return False
78+
79+
return await self.bootstrap_config.application.broker.ping(timeout=5)
80+
81+
82+
@dataclasses.dataclass(kw_only=True, frozen=True)
83+
class FastStreamLoggingInstrument(LoggingInstrument):
84+
bootstrap_config: FastStreamConfig
85+
86+
87+
@dataclasses.dataclass(kw_only=True, frozen=True)
88+
class FastStreamOpenTelemetryInstrument(OpenTelemetryInstrument):
89+
bootstrap_config: FastStreamConfig
90+
91+
def is_ready(self) -> bool:
92+
return bool(self.bootstrap_config.opentelemetry_middleware_cls and super().is_ready())
93+
94+
def bootstrap(self) -> None:
95+
if self.bootstrap_config.opentelemetry_middleware_cls and self.bootstrap_config.application.broker:
96+
self.bootstrap_config.application.broker.add_middleware(
97+
self.bootstrap_config.opentelemetry_middleware_cls(tracer_provider=get_tracer_provider())
98+
)
99+
100+
101+
@dataclasses.dataclass(kw_only=True, frozen=True)
102+
class FastStreamSentryInstrument(SentryInstrument):
103+
bootstrap_config: FastStreamConfig
104+
105+
106+
@dataclasses.dataclass(kw_only=True, frozen=True)
107+
class FastStreamPrometheusInstrument(PrometheusInstrument):
108+
bootstrap_config: FastStreamConfig
109+
collector_registry: prometheus_client.CollectorRegistry = dataclasses.field(
110+
default_factory=prometheus_client.CollectorRegistry, init=False
111+
)
112+
113+
def is_ready(self) -> bool:
114+
return bool(self.bootstrap_config.prometheus_middleware_cls and super().is_ready())
115+
116+
def bootstrap(self) -> None:
117+
self.bootstrap_config.application.mount(
118+
self.bootstrap_config.prometheus_metrics_path, prometheus_client.make_asgi_app(self.collector_registry)
119+
)
120+
if self.bootstrap_config.prometheus_middleware_cls and self.bootstrap_config.application.broker:
121+
self.bootstrap_config.application.broker.add_middleware(
122+
self.bootstrap_config.prometheus_middleware_cls(registry=self.collector_registry)
123+
)
124+
125+
126+
class FastStreamBootstrapper(BaseBootstrapper[AsgiFastStream]):
127+
instruments_types: typing.ClassVar = [
128+
FastStreamOpenTelemetryInstrument,
129+
FastStreamSentryInstrument,
130+
FastStreamHealthChecksInstrument,
131+
FastStreamLoggingInstrument,
132+
FastStreamPrometheusInstrument,
133+
]
134+
bootstrap_config: FastStreamConfig
135+
__slots__ = "bootstrap_config", "instruments"
136+
137+
def __init__(self, bootstrap_config: FastStreamConfig) -> None:
138+
super().__init__(bootstrap_config)
139+
if self.bootstrap_config.broker:
140+
self.bootstrap_config.application.broker = self.bootstrap_config.broker
141+
142+
def _prepare_application(self) -> AsgiFastStream:
143+
return self.bootstrap_config.application

pyproject.toml

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,15 +67,29 @@ litestar = [
6767
litestar-otl = [
6868
"opentelemetry-instrumentation-asgi>=0.46b0",
6969
]
70+
litestar-metrics = [
71+
"prometheus-client>=0.20",
72+
]
7073
litestar-all = [
71-
"lite-bootstrap[sentry,otl,logging,litestar,litestar-otl]"
74+
"lite-bootstrap[sentry,otl,logging,litestar,litestar-otl,litestar-metrics]"
75+
]
76+
faststream = [
77+
"faststream",
78+
]
79+
faststream-metrics = [
80+
"prometheus-client>=0.20",
81+
]
82+
faststream-all = [
83+
"lite-bootstrap[sentry,otl,logging,faststream,faststream-metrics]"
7284
]
7385

7486
[dependency-groups]
7587
dev = [
7688
"pytest",
7789
"pytest-cov",
90+
"pytest-asyncio",
7891
"httpx", # for test client
92+
"redis>=5.2.1",
7993
"mypy",
8094
"ruff",
8195
]
@@ -117,6 +131,8 @@ isort.no-lines-before = ["standard-library", "local-folder"]
117131

118132
[tool.pytest.ini_options]
119133
addopts = "--cov=. --cov-report term-missing"
134+
asyncio_mode = "auto"
135+
asyncio_default_fixture_loop_scope = "function"
120136

121137
[tool.coverage.report]
122138
exclude_also = ["if typing.TYPE_CHECKING:"]

tests/test_faststream_bootstrap.py

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
import pytest
2+
import structlog
3+
from faststream.redis import RedisBroker, TestRedisBroker
4+
from faststream.redis.opentelemetry import RedisTelemetryMiddleware
5+
from faststream.redis.prometheus import RedisPrometheusMiddleware
6+
from opentelemetry.sdk.trace.export import ConsoleSpanExporter
7+
from starlette import status
8+
from starlette.testclient import TestClient
9+
10+
from lite_bootstrap import FastStreamBootstrapper, FastStreamConfig
11+
from tests.conftest import CustomInstrumentor
12+
13+
14+
logger = structlog.getLogger(__name__)
15+
16+
17+
@pytest.fixture
18+
def broker() -> RedisBroker:
19+
return RedisBroker()
20+
21+
22+
async def test_faststream_bootstrap(broker: RedisBroker) -> None:
23+
prometheus_metrics_path = "/test-metrics-path"
24+
health_check_path = "/custom-health-check-path"
25+
bootstrapper = FastStreamBootstrapper(
26+
bootstrap_config=FastStreamConfig(
27+
broker=broker,
28+
service_name="microservice",
29+
service_version="2.0.0",
30+
service_environment="test",
31+
service_debug=False,
32+
opentelemetry_endpoint="otl",
33+
opentelemetry_instrumentors=[CustomInstrumentor()],
34+
opentelemetry_span_exporter=ConsoleSpanExporter(),
35+
opentelemetry_middleware_cls=RedisTelemetryMiddleware,
36+
prometheus_metrics_path=prometheus_metrics_path,
37+
prometheus_middleware_cls=RedisPrometheusMiddleware,
38+
sentry_dsn="https://testdsn@localhost/1",
39+
health_checks_path=health_check_path,
40+
logging_buffer_capacity=0,
41+
),
42+
)
43+
application = bootstrapper.bootstrap()
44+
logger.info("testing logging", key="value")
45+
test_client = TestClient(app=application)
46+
47+
async with TestRedisBroker(broker):
48+
response = test_client.get(prometheus_metrics_path)
49+
assert response.status_code == status.HTTP_200_OK
50+
51+
response = test_client.get(health_check_path)
52+
assert response.status_code == status.HTTP_200_OK
53+
assert response.json() == {"health_status": True, "service_name": "microservice", "service_version": "2.0.0"}
54+
55+
56+
async def test_faststream_bootstrap_health_check_wo_broker() -> None:
57+
health_check_path = "/custom-health-check-path"
58+
bootstrapper = FastStreamBootstrapper(
59+
bootstrap_config=FastStreamConfig(
60+
service_name="microservice",
61+
service_version="2.0.0",
62+
service_environment="test",
63+
service_debug=False,
64+
opentelemetry_endpoint="otl",
65+
opentelemetry_instrumentors=[CustomInstrumentor()],
66+
opentelemetry_span_exporter=ConsoleSpanExporter(),
67+
sentry_dsn="https://testdsn@localhost/1",
68+
health_checks_path=health_check_path,
69+
logging_buffer_capacity=0,
70+
),
71+
)
72+
application = bootstrapper.bootstrap()
73+
logger.info("testing logging", key="value")
74+
test_client = TestClient(app=application)
75+
76+
response = test_client.get(health_check_path)
77+
assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
78+
assert response.text == "Service is unhealthy"

0 commit comments

Comments
 (0)