Skip to content

Commit

Permalink
Merge pull request #15 from taskiq-python/feature/pass-app
Browse files Browse the repository at this point in the history
Added support for passing the FastaAPI instance.
  • Loading branch information
s3rius authored Dec 14, 2024
2 parents 9a43328 + 541da2d commit 49f38d4
Showing 1 changed file with 10 additions and 6 deletions.
16 changes: 10 additions & 6 deletions taskiq_fastapi/initializator.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Awaitable, Callable
from typing import Awaitable, Callable, Union

from fastapi import FastAPI, Request
from starlette.requests import HTTPConnection
Expand All @@ -8,7 +8,7 @@

def startup_event_generator(
broker: AsyncBroker,
app_path: str,
app_or_path: Union[str, FastAPI],
) -> Callable[[TaskiqState], Awaitable[None]]:
"""
Generate shutdown event.
Expand All @@ -24,12 +24,16 @@ def startup_event_generator(
async def startup(state: TaskiqState) -> None:
if not broker.is_worker_process:
return
app = import_object(app_path)
if isinstance(app_or_path, str):
app = import_object(app_or_path)
else:
app = app_or_path

if not isinstance(app, FastAPI):
app = app()

if not isinstance(app, FastAPI):
raise ValueError(f"'{app_path}' is not a FastAPI application.")
raise ValueError(f"'{app_or_path}' is not a FastAPI application.")

state.fastapi_app = app
await app.router.startup()
Expand Down Expand Up @@ -62,7 +66,7 @@ async def shutdown(state: TaskiqState) -> None:
return shutdown


def init(broker: AsyncBroker, app_path: str) -> None:
def init(broker: AsyncBroker, app_or_path: Union[str, FastAPI]) -> None:
"""
Add taskiq startup events.
Expand All @@ -78,7 +82,7 @@ def init(broker: AsyncBroker, app_path: str) -> None:
"""
broker.add_event_handler(
TaskiqEvents.WORKER_STARTUP,
startup_event_generator(broker, app_path),
startup_event_generator(broker, app_or_path),
)

broker.add_event_handler(
Expand Down

0 comments on commit 49f38d4

Please sign in to comment.