Coverage for faststream / _internal / fastapi / router.py: 96%

1import json (empty)

2import warnings (empty)

3from abc import abstractmethod (empty)

4from collections.abc import ( (empty)

5 AsyncIterator, 

6 Awaitable, 

7 Callable, 

8 Iterable, 

9 Mapping, 

10 Sequence, 

11) 

12from contextlib import asynccontextmanager (empty)

13from enum import Enum (empty)

14from typing import ( (empty)

15 TYPE_CHECKING, 

16 Any, 

17 Generic, 

18 Optional, 

19 Union, 

20 cast, 

21 overload, 

22) 

23 

24from fastapi.datastructures import Default (empty)

25from fastapi.responses import HTMLResponse (empty)

26from fastapi.routing import APIRoute, APIRouter (empty)

27from fastapi.utils import generate_unique_id (empty)

28from starlette.responses import JSONResponse, Response (empty)

29from starlette.routing import BaseRoute, _DefaultLifespan (empty)

30 

31from faststream._internal.application import StartAbleApplication (empty)

32from faststream._internal.broker import BrokerRouter (empty)

33from faststream._internal.context import ContextRepo (empty)

34from faststream._internal.di.config import FastDependsConfig (empty)

35from faststream._internal.types import ( (empty)

36 MsgType, 

37 P_HandlerParams, 

38 T_HandlerReturn, 

39) 

40from faststream._internal.utils.functions import to_async (empty)

41from faststream.middlewares import BaseMiddleware (empty)

42from faststream.specification.asyncapi.site import get_asyncapi_html (empty)

43 

44from .config import FastAPIConfig (empty)

45from .get_dependant import get_fastapi_dependant (empty)

46from .route import wrap_callable_to_fastapi_compatible (empty)

47 

48if TYPE_CHECKING: (empty)

49 from types import TracebackType 

50 

51 from fastapi import FastAPI, params 

52 from fastapi.background import BackgroundTasks 

53 from fastapi.types import IncEx 

54 from starlette import routing 

55 from starlette.types import ASGIApp, AppType, Lifespan 

56 

57 from faststream._internal.broker import BrokerUsecase 

58 from faststream._internal.endpoint.call_wrapper import HandlerCallWrapper 

59 from faststream._internal.endpoint.publisher import PublisherUsecase 

60 from faststream._internal.proto import NameRequired 

61 from faststream._internal.types import BrokerMiddleware 

62 from faststream.message import StreamMessage 

63 from faststream.specification.base import SpecificationFactory 

64 from faststream.specification.schema import Tag, TagDict 

65 

66 

67class _BackgroundMiddleware(BaseMiddleware): (empty)

68 async def __aexit__( (empty)

69 self, 

70 exc_type: type[BaseException] | None = None, 

71 exc_val: BaseException | None = None, 

72 exc_tb: Optional["TracebackType"] = None, 

73 ) -> bool | None: 

74 if not exc_type and ( (empty)

75 background := cast( 

76 "BackgroundTasks | None", 

77 getattr(self.context.get_local("message"), "background", None), 

78 ) 

79 ): 

80 await background() (empty)

81 

82 return await super().after_processed(exc_type, exc_val, exc_tb) (empty)

83 

84 

85class StreamRouter(APIRouter, StartAbleApplication, Generic[MsgType]): (empty)

86 """A class to route streams.""" 

87 

88 broker_class: type["BrokerUsecase[MsgType, Any]"] (empty)

89 broker: "BrokerUsecase[MsgType, Any]" (empty)

90 docs_router: APIRouter | None (empty)

91 _after_startup_hooks: list[Callable[[Any], Awaitable[Mapping[str, Any] | None]]] (empty)

92 _on_shutdown_hooks: list[Callable[[Any], Awaitable[None]]] (empty)

93 

94 title: str (empty)

95 description: str (empty)

96 version: str (empty)

97 license: dict[str, Any] | None (empty)

98 contact: dict[str, Any] | None (empty)

99 

100 def __init__( (empty)

101 self, 

102 *connection_args: Any, 

103 context: ContextRepo | None, 

104 middlewares: Sequence["BrokerMiddleware[MsgType]"] = (), 

105 prefix: str = "", 

106 tags: list[str | Enum] | None = None, 

107 dependencies: Sequence["params.Depends"] | None = None, 

108 default_response_class: type["Response"] = Default(JSONResponse), 

109 responses: dict[int | str, dict[str, Any]] | None = None, 

110 callbacks: list["routing.BaseRoute"] | None = None, 

111 routes: list["routing.BaseRoute"] | None = None, 

112 redirect_slashes: bool = True, 

113 default: Optional["ASGIApp"] = None, 

114 dependency_overrides_provider: Any | None = None, 

115 route_class: type["APIRoute"] = APIRoute, 

116 on_startup: Sequence[Callable[[], Any]] | None = None, 

117 on_shutdown: Sequence[Callable[[], Any]] | None = None, 

118 deprecated: bool | None = None, 

119 include_in_schema: bool = True, 

120 setup_state: bool = True, 

121 lifespan: Optional["Lifespan[Any]"] = None, 

122 generate_unique_id_function: Callable[["APIRoute"], str] = Default( 

123 generate_unique_id, 

124 ), 

125 # Specification information 

126 specification: Optional["SpecificationFactory"] = None, 

127 specification_tags: Iterable[Union["Tag", "TagDict"]] = (), 

128 schema_url: str | None = "/asyncapi", 

129 **connection_kwars: Any, 

130 ) -> None: 

131 broker = self.broker_class( (empty)

132 *connection_args, 

133 middlewares=( 

134 *middlewares, 

135 # allow to catch background exceptions in user middlewares 

136 _BackgroundMiddleware, 

137 ), 

138 tags=specification_tags, 

139 apply_types=False, 

140 **connection_kwars, 

141 ) 

142 

143 self._init_setupable_( (empty)

144 broker, 

145 config=FastDependsConfig( 

146 get_dependent=get_fastapi_dependant, 

147 context=context or ContextRepo(), 

148 ), 

149 specification=specification, 

150 ) 

151 

152 self.setup_state = setup_state (empty)

153 

154 super().__init__( (empty)

155 prefix=prefix, 

156 tags=tags, 

157 dependencies=dependencies, 

158 default_response_class=default_response_class, 

159 responses=responses, 

160 callbacks=callbacks, 

161 routes=routes, 

162 redirect_slashes=redirect_slashes, 

163 default=default, 

164 dependency_overrides_provider=dependency_overrides_provider, 

165 route_class=route_class, 

166 deprecated=deprecated, 

167 include_in_schema=include_in_schema, 

168 generate_unique_id_function=generate_unique_id_function, 

169 lifespan=self._wrap_lifespan(lifespan), 

170 on_startup=on_startup, 

171 on_shutdown=on_shutdown, 

172 ) 

173 

174 self.fastapi_config = FastAPIConfig( (empty)

175 dependency_overrides_provider=dependency_overrides_provider, 

176 ) 

177 

178 if self.include_in_schema: (empty)

179 self.docs_router = self._asyncapi_router(schema_url) (empty)

180 else: 

181 self.docs_router = None (empty)

182 

183 self._after_startup_hooks = [] (empty)

184 self._on_shutdown_hooks = [] (empty)

185 

186 self._lifespan_started = False (empty)

187 

188 def _subscriber_compatibility_wrapper( (empty)

189 self, 

190 dependencies: Iterable["params.Depends"] = (), (empty)

191 response_model: Any = Default(None), (empty)

192 response_model_include: Optional["IncEx"] = None, (empty)

193 response_model_exclude: Optional["IncEx"] = None, (empty)

194 response_model_by_alias: bool = True, (empty)

195 response_model_exclude_unset: bool = False, (empty)

196 response_model_exclude_defaults: bool = False, (empty)

197 response_model_exclude_none: bool = False, (empty)

198 ) -> Callable[ (empty)

199 [Callable[..., Any]], (empty)

200 Callable[["StreamMessage[Any]"], Awaitable[Any]], (empty)

201 ]: 

202 """Decorator before `broker.subscriber`, that wraps function to FastAPI-compatible one.""" 

203 

204 def wrapper( (empty)

205 endpoint: Callable[..., Any], (empty)

206 ) -> Callable[["StreamMessage[Any]"], Awaitable[Any]]: (empty)

207 """Patch user function to make it FastAPI-compatible.""" 

208 return wrap_callable_to_fastapi_compatible( (empty)

209 user_callable=endpoint, (empty)

210 dependencies=dependencies, (empty)

211 response_model=response_model, (empty)

212 response_model_include=response_model_include, (empty)

213 response_model_exclude=response_model_exclude, (empty)

214 response_model_by_alias=response_model_by_alias, (empty)

215 response_model_exclude_unset=response_model_exclude_unset, (empty)

216 response_model_exclude_defaults=response_model_exclude_defaults, (empty)

217 response_model_exclude_none=response_model_exclude_none, (empty)

218 context=self.context, (empty)

219 fastapi_config=self.fastapi_config, (empty)

220 ) 

221 

222 return wrapper (empty)

223 

224 def subscriber( (empty)

225 self, 

226 *extra: Union["NameRequired", str], 

227 dependencies: Iterable["params.Depends"], 

228 response_model: Any, 

229 response_model_include: Optional["IncEx"], 

230 response_model_exclude: Optional["IncEx"], 

231 response_model_by_alias: bool, 

232 response_model_exclude_unset: bool, 

233 response_model_exclude_defaults: bool, 

234 response_model_exclude_none: bool, 

235 **broker_kwargs: Any, 

236 ) -> Callable[ 

237 [Callable[P_HandlerParams, T_HandlerReturn]], 

238 "HandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", 

239 ]: 

240 """A function decorator for subscribing to a message queue.""" 

241 dependencies = (*self.dependencies, *dependencies) (empty)

242 

243 sub = self.broker.subscriber( # type: ignore[call-arg] (empty)

244 *extra, # type: ignore[arg-type] 

245 dependencies=dependencies, 

246 **broker_kwargs, 

247 ) 

248 

249 sub._call_decorators = ( (empty)

250 self._subscriber_compatibility_wrapper( 

251 dependencies=dependencies, 

252 response_model=response_model, 

253 response_model_include=response_model_include, 

254 response_model_exclude=response_model_exclude, 

255 response_model_by_alias=response_model_by_alias, 

256 response_model_exclude_unset=response_model_exclude_unset, 

257 response_model_exclude_defaults=response_model_exclude_defaults, 

258 response_model_exclude_none=response_model_exclude_none, 

259 ), 

260 *sub._call_decorators, 

261 ) 

262 

263 return sub (empty)

264 

265 def _wrap_lifespan( (empty)

266 self, 

267 lifespan: Optional["Lifespan[Any]"] = None, 

268 ) -> "Lifespan[Any]": 

269 lifespan_context = lifespan if lifespan is not None else _DefaultLifespan(self) (empty)

270 

271 @asynccontextmanager (empty)

272 async def start_broker_lifespan( (empty)

273 app: "FastAPI", 

274 ) -> AsyncIterator[Mapping[str, Any] | None]: 

275 """Starts the lifespan of a broker.""" 

276 self.fastapi_config.set_application(app) (empty)

277 

278 if self.docs_router: (empty)

279 self.schema.title = app.title (empty)

280 self.schema.description = app.description (empty)

281 self.schema.version = app.version (empty)

282 self.schema.contact = app.contact (empty)

283 self.schema.license = app.license_info (empty)

284 app.include_router(self.docs_router) (empty)

285 

286 async with lifespan_context(app) as maybe_context: (empty)

287 lifespan_extra = {"broker": self.broker, **(maybe_context or {})} (empty)

288 

289 if not self._lifespan_started: 289 ↛ 293line 289 didn't jump to line 293 because the condition on line 289 was always true(empty)

290 await self._start_broker() (empty)

291 self._lifespan_started = True (empty)

292 else: 

293 warnings.warn( 

294 "Specifying 'lifespan_context' manually is no longer necessary with FastAPI >= 0.112.2.", 

295 category=RuntimeWarning, 

296 stacklevel=2, 

297 ) 

298 

299 for h in self._after_startup_hooks: (empty)

300 lifespan_extra.update(await h(app) or {}) (empty)

301 

302 try: (empty)

303 if self.setup_state: 303 ↛ 307line 303 didn't jump to line 307 because the condition on line 303 was always true(empty)

304 yield lifespan_extra (empty)

305 else: 

306 # NOTE: old asgi compatibility 

307 yield None 

308 

309 for h in self._on_shutdown_hooks: (empty)

310 await h(app) (empty)

311 

312 finally: 

313 await self.broker.stop() (empty)

314 

315 return start_broker_lifespan # type: ignore[return-value] (empty)

316 

317 @overload (empty)

318 def after_startup( (empty)

319 self, 

320 func: Callable[["AppType"], Mapping[str, Any]], (empty)

321 ) -> Callable[["AppType"], Mapping[str, Any]]: ... (empty)

322 

323 @overload (empty)

324 def after_startup( (empty)

325 self, 

326 func: Callable[["AppType"], Awaitable[Mapping[str, Any]]], (empty)

327 ) -> Callable[["AppType"], Awaitable[Mapping[str, Any]]]: ... (empty)

328 

329 @overload (empty)

330 def after_startup( (empty)

331 self, 

332 func: Callable[["AppType"], None], (empty)

333 ) -> Callable[["AppType"], None]: ... (empty)

334 

335 @overload (empty)

336 def after_startup( (empty)

337 self, 

338 func: Callable[["AppType"], Awaitable[None]], (empty)

339 ) -> Callable[["AppType"], Awaitable[None]]: ... (empty)

340 

341 def after_startup( (empty)

342 self, 

343 func: Callable[["AppType"], Mapping[str, Any]] 

344 | Callable[["AppType"], Awaitable[Mapping[str, Any]]] 

345 | Callable[["AppType"], None] 

346 | Callable[["AppType"], Awaitable[None]], 

347 ) -> ( 

348 Callable[["AppType"], Mapping[str, Any]] 

349 | Callable[["AppType"], Awaitable[Mapping[str, Any]]] 

350 | Callable[["AppType"], None] 

351 | Callable[["AppType"], Awaitable[None]] 

352 ): 

353 """Register a function to be executed after startup.""" 

354 self._after_startup_hooks.append(to_async(func)) (empty)

355 return func (empty)

356 

357 @overload (empty)

358 def on_broker_shutdown( (empty)

359 self, 

360 func: Callable[["AppType"], None], (empty)

361 ) -> Callable[["AppType"], None]: ... (empty)

362 

363 @overload (empty)

364 def on_broker_shutdown( (empty)

365 self, 

366 func: Callable[["AppType"], Awaitable[None]], (empty)

367 ) -> Callable[["AppType"], Awaitable[None]]: ... (empty)

368 

369 def on_broker_shutdown( (empty)

370 self, 

371 func: Callable[["AppType"], None] | Callable[["AppType"], Awaitable[None]], 

372 ) -> Callable[["AppType"], None] | Callable[["AppType"], Awaitable[None]]: 

373 """Register a function to be executed before broker stop.""" 

374 self._on_shutdown_hooks.append(to_async(func)) (empty)

375 return func (empty)

376 

377 @abstractmethod (empty)

378 def publisher(self) -> "PublisherUsecase": (empty)

379 """Create Publisher object.""" 

380 raise NotImplementedError 

381 

382 def _asyncapi_router(self, schema_url: str | None) -> APIRouter | None: (empty)

383 """Creates an API router for serving AsyncAPI documentation.""" 

384 if not self.include_in_schema or not schema_url: (empty)

385 return None (empty)

386 

387 def download_app_json_schema() -> Response: (empty)

388 return Response( (empty)

389 content=json.dumps( 

390 self.schema.to_specification().to_jsonable(), 

391 indent=2, 

392 ), 

393 headers={"Content-Type": "application/octet-stream"}, 

394 ) 

395 

396 def download_app_yaml_schema() -> Response: (empty)

397 return Response( (empty)

398 content=self.schema.to_specification().to_yaml(), 

399 headers={ 

400 "Content-Type": "application/octet-stream", 

401 }, 

402 ) 

403 

404 def serve_asyncapi_schema( (empty)

405 sidebar: bool = True, 

406 info: bool = True, 

407 servers: bool = True, 

408 operations: bool = True, 

409 messages: bool = True, 

410 schemas: bool = True, 

411 errors: bool = True, 

412 expandMessageExamples: bool = True, 

413 ) -> HTMLResponse: 

414 """Serve the AsyncAPI schema as an HTML response.""" 

415 return HTMLResponse( (empty)

416 content=get_asyncapi_html( 

417 self.schema.to_specification(), 

418 sidebar=sidebar, 

419 info=info, 

420 servers=servers, 

421 operations=operations, 

422 messages=messages, 

423 schemas=schemas, 

424 errors=errors, 

425 expand_message_examples=expandMessageExamples, 

426 ), 

427 ) 

428 

429 docs_router = APIRouter( (empty)

430 prefix=self.prefix, 

431 tags=["asyncapi"], 

432 redirect_slashes=self.redirect_slashes, 

433 default=self.default, 

434 deprecated=self.deprecated, 

435 ) 

436 docs_router.get(schema_url)(serve_asyncapi_schema) (empty)

437 docs_router.get(f"{schema_url}.json")(download_app_json_schema) (empty)

438 docs_router.get(f"{schema_url}.yaml")(download_app_yaml_schema) (empty)

439 return docs_router (empty)

440 

441 def include_router( # type: ignore[override] (empty)

442 self, 

443 router: Union["StreamRouter[MsgType]", "BrokerRouter[MsgType]"], 

444 *, 

445 prefix: str = "", 

446 tags: list[str | Enum] | None = None, 

447 dependencies: Sequence["params.Depends"] | None = None, 

448 default_response_class: type[Response] = Default(JSONResponse), 

449 responses: dict[int | str, dict[str, Any]] | None = None, 

450 callbacks: list["BaseRoute"] | None = None, 

451 deprecated: bool | None = None, 

452 include_in_schema: bool = True, 

453 generate_unique_id_function: Callable[["APIRoute"], str] = Default( 

454 generate_unique_id, 

455 ), 

456 ) -> None: 

457 """Includes a router in the API.""" 

458 if isinstance(router, BrokerRouter): (empty)

459 for sub in router.subscribers: (empty)

460 sub._call_decorators = ( (empty)

461 self._subscriber_compatibility_wrapper(), 

462 *sub._call_decorators, 

463 ) 

464 

465 self.broker.include_router(router) (empty)

466 return (empty)

467 

468 msg = ( (empty)

469 "Including a StreamRouter into another StreamRouter is not supported " 

470 "and may cause subtle context issues (e.g. message dependencies " 

471 "returning EmptyPlaceholder). " 

472 "Use a regular broker router (e.g. KafkaRouter, RabbitRouter, etc.) " 

473 "for grouping subscribers and include that into the StreamRouter instead. " 

474 "See: https://faststream.ag2.ai/latest/getting-started/integrations/fastapi/#multiple-routers" 

475 ) 

476 raise TypeError(msg) (empty)