diff --git a/docs/low_api/connection/index.md b/docs/low_api/connection/index.md new file mode 100644 index 0000000..10bdbaa --- /dev/null +++ b/docs/low_api/connection/index.md @@ -0,0 +1,35 @@ +# Connection + +## Startup setup +You can configure the importance of connecting to brokers using the `safe_connection` parameter. + +If `safe_connection=True`, the application will not start if the connection to the broker fails. + +If `safe_connection=False`, the application will start even if there is a connection error, followed by an error in the logs. + +!!! tip + By default `safe_connection` is `False` + +## Connection hook + +You can also set up your own hook for the connection. + +Example with tenacity: +```python +import uvicorn +from fastapi import FastAPI +from faststream.kafka import KafkaBroker +from faststream_fastapi import FastStreamAPI +from tenacity import retry + +@retry +async def my_connection_hook(broker: KafkaBroker) -> None: + await broker.start() + +application = FastStreamAPI( + KafkaBroker(), + application=FastAPI(), + connection_hook=my_connection_hook, +) +uvicorn.run(application) +``` diff --git a/src/faststream_fastapi/faststream_api.py b/src/faststream_fastapi/faststream_api.py index bfca824..5746aff 100644 --- a/src/faststream_fastapi/faststream_api.py +++ b/src/faststream_fastapi/faststream_api.py @@ -1,6 +1,7 @@ import traceback from collections.abc import AsyncIterator, Awaitable, Callable, Iterable from contextlib import asynccontextmanager +from logging import getLogger from typing import Any from fastapi import FastAPI @@ -8,6 +9,7 @@ from faststream.message import StreamMessage from faststream.specification.base import SpecificationFactory from starlette.types import Receive, Scope, Send +from typing_extensions import Self from faststream_fastapi._internal.asyncapi_router import AsyncAPIRouter from faststream_fastapi._internal.background_middleware import _BackgroundMiddleware @@ -22,6 +24,7 @@ ) from faststream_fastapi.asyncapi_config import AsyncAPIConfig +_logger = getLogger(__name__) def _subscriber_compatibility_wrapper( config: Config, @@ -41,6 +44,10 @@ def subscriber_compatibility_wrapper( return subscriber_compatibility_wrapper +async def _default_connection_hook(broker: BrokerUsecase[Any, Any]) -> None: + await broker.start() + + class FastStreamAPI: def __init__( self, @@ -50,6 +57,9 @@ def __init__( # AsyncAPI specification: SpecificationFactory | None = None, asyncapi_path: str | AsyncAPIConfig | None = None, + # Low API + safe_connection: bool = False, + connection_hook: Callable[[Any], Awaitable[None]] | None = None, ) -> None: self._application = application @@ -97,6 +107,13 @@ def __init__( *subscriber._call_decorators, ) + self._safe_connection = safe_connection + + if connection_hook is None: + connection_hook = _default_connection_hook + + self._connection_hook = connection_hook + # For FastStream docs gen @property def schema(self) -> SpecificationFactory: @@ -145,7 +162,7 @@ async def lifespan(self, scope: Scope, receive: Receive, send: Send) -> None: await send({"type": "lifespan.shutdown.complete"}) @asynccontextmanager - async def _lifespan_context(self, application: Any) -> AsyncIterator[None]: + async def _lifespan_context(self, application: Self) -> AsyncIterator[None]: if self._asyncapi_config is not None: asyncapi_router = AsyncAPIRouter( brokers=self._brokers, @@ -158,10 +175,18 @@ async def _lifespan_context(self, application: Any) -> AsyncIterator[None]: try: for broker in self._brokers: - await broker.start() + await self._connect_broker(broker) started_brokers.append(broker) yield None finally: for started_broker in started_brokers: await started_broker.stop() + + async def _connect_broker(self, broker: BrokerUsecase[Any, Any]) -> None: + try: + await self._connection_hook(broker) + except Exception as error: + if self._safe_connection: + raise + _logger.warning("The start to broker %r was unsuccessful.", broker, exc_info=error) diff --git a/tests/base/abstract.py b/tests/base/abstract.py index 31f3d86..aa85973 100644 --- a/tests/base/abstract.py +++ b/tests/base/abstract.py @@ -22,7 +22,7 @@ class AbstractTestCaseConfig(ABC, Generic[_BrokerT]): @abstractmethod def get_broker( self, - *, + *args: Any, apply_types: bool = False, **kwargs: Any, ) -> _BrokerT: diff --git a/tests/base/real.py b/tests/base/real.py index 83634e6..d9d9796 100644 --- a/tests/base/real.py +++ b/tests/base/real.py @@ -1,11 +1,12 @@ +import asyncio from asyncio import Event, create_task, wait from typing import Annotated, Any, Generic, TypeVar -from unittest.mock import Mock +from unittest.mock import MagicMock, Mock import pytest from fastapi import BackgroundTasks -from faststream_fastapi import Context, StreamMessage +from faststream_fastapi import Context, FastStreamAPI, StreamMessage from faststream_fastapi._internal.fs_re_exports.broker import BrokerUsecase from faststream_fastapi._internal.fs_re_exports.context import ContextRepo from tests.base.abstract import AbstractTestCaseConfig @@ -263,3 +264,12 @@ async def subscriber_handler(message: StreamMessage) -> None: ) mock.assert_called_once_with(True) + + async def test_conneciton_safe_true(self, mock: MagicMock) -> None: + broker = self.get_broker() + app = FastStreamAPI(broker, application=mock, safe_connection=True) + + await asyncio.wait( + [asyncio.create_task(app._connect_broker(broker))], + timeout=self.timeout, + ) diff --git a/tests/integrations/confluent/abstract.py b/tests/integrations/confluent/abstract.py index 84fa667..dbd9363 100644 --- a/tests/integrations/confluent/abstract.py +++ b/tests/integrations/confluent/abstract.py @@ -38,10 +38,11 @@ def get_subscriber_params( @override def get_broker( self, + *args: Any, apply_types: bool = False, **kwargs: Any, ) -> KafkaBroker: - return KafkaBroker(apply_types=apply_types, **kwargs) + return KafkaBroker(*args, apply_types=apply_types, **kwargs) class ConfluentAbstractInMemoryTestCaseConfig(ConfluentAbstractTestCaseConfig): diff --git a/tests/integrations/confluent/test_real.py b/tests/integrations/confluent/test_real.py index 2af759a..f0b6e6e 100644 --- a/tests/integrations/confluent/test_real.py +++ b/tests/integrations/confluent/test_real.py @@ -5,6 +5,7 @@ import pytest from faststream.confluent import KafkaBroker +from faststream_fastapi import FastStreamAPI from tests.base.real import BaseRealTestCaseConfig from tests.integrations.confluent.abstract import ConfluentAbstractTestCaseConfig @@ -38,3 +39,14 @@ async def subscriber_handler(msg: list[str]) -> None: assert event.is_set() mock.assert_called_with(["hi"]) + + async def test_conneciton_safe_false(self, mock: MagicMock) -> None: + broker = self.get_broker("localhost:6767") + + app = FastStreamAPI(broker, application=mock) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait( + [asyncio.create_task(app._connect_broker(broker))], + timeout=self.timeout, + ) diff --git a/tests/integrations/kafka/abstract.py b/tests/integrations/kafka/abstract.py index 615299c..74e9093 100644 --- a/tests/integrations/kafka/abstract.py +++ b/tests/integrations/kafka/abstract.py @@ -11,6 +11,7 @@ class KafkaAbstractTestCaseConfig(AbstractTestCaseConfig[KafkaBroker]): @override def get_broker( self, + *args: Any, apply_types: bool = False, **kwargs: Any, ) -> KafkaBroker: diff --git a/tests/integrations/kafka/test_real.py b/tests/integrations/kafka/test_real.py index deabb58..80f4762 100644 --- a/tests/integrations/kafka/test_real.py +++ b/tests/integrations/kafka/test_real.py @@ -5,6 +5,7 @@ import pytest from faststream.kafka import KafkaBroker +from faststream_fastapi import FastStreamAPI from tests.base.real import BaseRealTestCaseConfig from tests.integrations.kafka.abstract import KafkaAbstractTestCaseConfig @@ -38,3 +39,14 @@ async def subscriber_handler(msg: list[str]) -> None: assert event.is_set() mock.assert_called_with(["hi"]) + + async def test_conneciton_safe_false(self, mock: MagicMock) -> None: + broker = self.get_broker("localhost:6767") + + app = FastStreamAPI(broker, application=mock) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait( + [asyncio.create_task(app._connect_broker(broker))], + timeout=self.timeout, + ) diff --git a/tests/integrations/mqtt/abstract.py b/tests/integrations/mqtt/abstract.py index 4623403..4e240f7 100644 --- a/tests/integrations/mqtt/abstract.py +++ b/tests/integrations/mqtt/abstract.py @@ -12,6 +12,7 @@ class MQTTAbstractTestCaseConfig(AbstractTestCaseConfig[MQTTBroker]): @override def get_broker( self, + *args: Any, apply_types: bool = False, **kwargs: Any, ) -> MQTTBroker: diff --git a/tests/integrations/mqtt/test_real.py b/tests/integrations/mqtt/test_real.py index 3fa57a8..fe4cc71 100644 --- a/tests/integrations/mqtt/test_real.py +++ b/tests/integrations/mqtt/test_real.py @@ -5,6 +5,7 @@ import pytest from faststream.mqtt import MQTTBroker +from faststream_fastapi import FastStreamAPI from tests.base.real import BaseRealTestCaseConfig from tests.integrations.mqtt.abstract import MQTTAbstractTestCaseConfig @@ -40,3 +41,14 @@ def subscriber_handler(msg: str, name: str) -> None: assert event.is_set() mock.assert_called_once_with(msg="hello", name="john") + + async def test_conneciton_safe_false(self, mock: MagicMock) -> None: + broker = self.get_broker("mqtt://localhost:6767") + + app = FastStreamAPI(broker, application=mock) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait( + [asyncio.create_task(app._connect_broker(broker))], + timeout=self.timeout, + ) diff --git a/tests/integrations/nats/abstract.py b/tests/integrations/nats/abstract.py index 8abc1dc..2216639 100644 --- a/tests/integrations/nats/abstract.py +++ b/tests/integrations/nats/abstract.py @@ -11,6 +11,7 @@ class NatsAbstractTestCaseConfig(AbstractTestCaseConfig[NatsBroker]): @override def get_broker( self, + *args: Any, apply_types: bool = False, **kwargs: Any, ) -> NatsBroker: diff --git a/tests/integrations/nats/test_real.py b/tests/integrations/nats/test_real.py index 392b315..a8a7b73 100644 --- a/tests/integrations/nats/test_real.py +++ b/tests/integrations/nats/test_real.py @@ -5,6 +5,7 @@ import pytest from faststream.nats import JStream, NatsBroker, PullSub +from faststream_fastapi import FastStreamAPI from tests.base.real import BaseRealTestCaseConfig from tests.integrations.nats.abstract import NatsAbstractTestCaseConfig @@ -72,3 +73,14 @@ def subscriber_handler(messages: list[str]) -> None: assert event.is_set() mock.assert_called_once_with(["hello"]) + + async def test_conneciton_safe_false(self, mock: MagicMock) -> None: + broker = self.get_broker("nats://localhost:6767") + + app = FastStreamAPI(broker, application=mock) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait( + [asyncio.create_task(app._connect_broker(broker))], + timeout=self.timeout, + ) diff --git a/tests/integrations/rabbit/abstract.py b/tests/integrations/rabbit/abstract.py index 27634c9..1e9fb37 100644 --- a/tests/integrations/rabbit/abstract.py +++ b/tests/integrations/rabbit/abstract.py @@ -11,6 +11,7 @@ class RabbitAbstractTestCaseConfig(AbstractTestCaseConfig[RabbitBroker]): @override def get_broker( self, + *args: Any, apply_types: bool = False, **kwargs: Any, ) -> RabbitBroker: diff --git a/tests/integrations/rabbit/test_real.py b/tests/integrations/rabbit/test_real.py index 7b3f474..9cb0c1c 100644 --- a/tests/integrations/rabbit/test_real.py +++ b/tests/integrations/rabbit/test_real.py @@ -5,6 +5,7 @@ import pytest from faststream.rabbit import ExchangeType, RabbitBroker, RabbitExchange, RabbitQueue +from faststream_fastapi import FastStreamAPI from tests.base.real import BaseRealTestCaseConfig from tests.integrations.rabbit.abstract import RabbitAbstractTestCaseConfig @@ -50,3 +51,14 @@ def subscriber_handler(msg: str, name: str) -> None: assert event.is_set() mock.assert_called_once_with(msg="hello", name="john") + + async def test_conneciton_safe_false(self, mock: MagicMock) -> None: + broker = self.get_broker("amqp://guest:guest@localhost:6767/") + + app = FastStreamAPI(broker, application=mock) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait( + [asyncio.create_task(app._connect_broker(broker))], + timeout=self.timeout, + ) diff --git a/tests/integrations/redis/abstract.py b/tests/integrations/redis/abstract.py index 9e94450..3f093d5 100644 --- a/tests/integrations/redis/abstract.py +++ b/tests/integrations/redis/abstract.py @@ -11,6 +11,7 @@ class RedisAbstractTestCaseConfig(AbstractTestCaseConfig[RedisBroker]): @override def get_broker( self, + *args: Any, apply_types: bool = False, **kwargs: Any, ) -> RedisBroker: diff --git a/tests/integrations/redis/test_real.py b/tests/integrations/redis/test_real.py index a250dc8..905eec8 100644 --- a/tests/integrations/redis/test_real.py +++ b/tests/integrations/redis/test_real.py @@ -1,9 +1,11 @@ +import asyncio from asyncio import Event, create_task, sleep, wait from unittest.mock import MagicMock import pytest from faststream.redis import ListSub, RedisBroker, StreamSub +from faststream_fastapi import FastStreamAPI from tests.base.real import BaseRealTestCaseConfig from tests.integrations.redis.abstract import RedisAbstractTestCaseConfig @@ -134,3 +136,14 @@ async def subscriber_handler(msg: list[str]) -> None: ) mock.assert_called_once_with(["hello"]) + + async def test_conneciton_safe_false(self, mock: MagicMock) -> None: + broker = self.get_broker("redis://localhost:6767") + + app = FastStreamAPI(broker, application=mock) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait( + [asyncio.create_task(app._connect_broker(broker))], + timeout=self.timeout, + ) diff --git a/zensical.toml b/zensical.toml index 26a42e4..1b6afcc 100644 --- a/zensical.toml +++ b/zensical.toml @@ -29,6 +29,9 @@ nav = [ { "Testing" = [ { "Testing" = "testing/index.md" }, ] }, + { "Low API" = [ + { "Connection" = "low_api/connection/index.md"} + ] } ] [project.theme]