Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updateted lifespan setup. #215

Merged
merged 1 commit into from
Jul 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ $ tree "{{cookiecutter.project_name}}"
├── api # Package with all handlers.
│   └── router.py # Main router.
├── application.py # FastAPI application configuration.
└── lifetime.py # Contains actions to perform on startup and shutdown.
└── lifespan.py # Contains actions to perform on startup and shutdown.
```

## Configuration
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,14 +22,14 @@
from aio_pika.pool import Pool
from {{cookiecutter.project_name}}.services.rabbit.dependencies import \
get_rmq_channel_pool
from {{cookiecutter.project_name}}.services.rabbit.lifetime import (init_rabbit,
from {{cookiecutter.project_name}}.services.rabbit.lifespan import (init_rabbit,
shutdown_rabbit)

{%- endif %}
{%- if cookiecutter.enable_kafka == "True" %}
from aiokafka import AIOKafkaProducer
from {{cookiecutter.project_name}}.services.kafka.dependencies import get_kafka_producer
from {{cookiecutter.project_name}}.services.kafka.lifetime import (init_kafka,
from {{cookiecutter.project_name}}.services.kafka.lifespan import (init_kafka,
shutdown_kafka)

{%- endif %}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,7 @@
{%- endif %}
from importlib import metadata

from {{cookiecutter.project_name}}.web.lifetime import (register_shutdown_event,
register_startup_event)
from {{cookiecutter.project_name}}.web.lifespan import lifespan_setup

{%- if cookiecutter.orm == 'tortoise' %}
from tortoise.contrib.fastapi import register_tortoise
Expand Down Expand Up @@ -80,6 +79,7 @@ def get_app() -> FastAPI:
app = FastAPI(
title="{{cookiecutter.project_name}}",
version=metadata.version("{{cookiecutter.project_name}}"),
lifespan=lifespan_setup,
{%- if cookiecutter.self_hosted_swagger == 'True' %}
docs_url=None,
redoc_url=None,
Expand All @@ -91,10 +91,6 @@ def get_app() -> FastAPI:
default_response_class=UJSONResponse,
)

# Adds startup and shutdown events.
register_startup_event(app)
register_shutdown_event(app)

# Main router for the API.
app.include_router(router=api_router, prefix="/api")
{%- if cookiecutter.api_type == 'graphql' %}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import logging
from typing import Awaitable, Callable
from typing import AsyncGenerator, Awaitable, Callable
from contextlib import asynccontextmanager

from fastapi import FastAPI
from {{cookiecutter.project_name}}.settings import settings
Expand All @@ -11,19 +12,19 @@
{%- endif %}

{%- if cookiecutter.enable_redis == "True" %}
from {{cookiecutter.project_name}}.services.redis.lifetime import (init_redis,
from {{cookiecutter.project_name}}.services.redis.lifespan import (init_redis,
shutdown_redis)

{%- endif %}

{%- if cookiecutter.enable_rmq == "True" %}
from {{cookiecutter.project_name}}.services.rabbit.lifetime import (init_rabbit,
from {{cookiecutter.project_name}}.services.rabbit.lifespan import (init_rabbit,
shutdown_rabbit)

{%- endif %}

{%- if cookiecutter.enable_kafka == "True" %}
from {{cookiecutter.project_name}}.services.kafka.lifetime import (init_kafka,
from {{cookiecutter.project_name}}.services.kafka.lifespan import (init_kafka,
shutdown_kafka)

{%- endif %}
Expand Down Expand Up @@ -267,7 +268,8 @@ def setup_prometheus(app: FastAPI) -> None: # pragma: no cover
{%- endif %}


def register_startup_event(app: FastAPI) -> Callable[[], Awaitable[None]]: # pragma: no cover
@asynccontextmanager
async def lifespan_setup(app: FastAPI) -> AsyncGenerator[None, None]: # pragma: no cover
"""
Actions to run on application startup.

Expand All @@ -278,79 +280,62 @@ def register_startup_event(app: FastAPI) -> Callable[[], Awaitable[None]]: # pr
:return: function that actually performs actions.
"""

@app.on_event("startup")
async def _startup() -> None: # noqa: WPS430
app.middleware_stack = None
{%- if cookiecutter.enable_taskiq == "True" %}
if not broker.is_worker_process:
await broker.startup()
{%- endif %}
{%- if cookiecutter.orm == "sqlalchemy" %}
_setup_db(app)
{%- elif cookiecutter.orm == "ormar" %}
await database.connect()
{%- elif cookiecutter.orm in ["beanie", "psycopg"] %}
await _setup_db(app)
{%- endif %}
{%- if cookiecutter.db_info.name != "none" and cookiecutter.enable_migrations != "True" %}
{%- if cookiecutter.orm in ["ormar", "sqlalchemy"] %}
await _create_tables()
{%- endif %}
{%- endif %}
{%- if cookiecutter.otlp_enabled == "True" %}
setup_opentelemetry(app)
{%- endif %}
{%- if cookiecutter.enable_redis == "True" %}
init_redis(app)
{%- endif %}
{%- if cookiecutter.enable_rmq == "True" %}
init_rabbit(app)
{%- endif %}
{%- if cookiecutter.enable_kafka == "True" %}
await init_kafka(app)
{%- endif %}
{%- if cookiecutter.prometheus_enabled == "True" %}
setup_prometheus(app)
{%- endif %}
app.middleware_stack = app.build_middleware_stack()
pass # noqa: WPS420

return _startup


def register_shutdown_event(app: FastAPI) -> Callable[[], Awaitable[None]]: # pragma: no cover
"""
Actions to run on application's shutdown.

:param app: fastAPI application.
:return: function that actually performs actions.
"""
app.middleware_stack = None
{%- if cookiecutter.enable_taskiq == "True" %}
if not broker.is_worker_process:
await broker.startup()
{%- endif %}
{%- if cookiecutter.orm == "sqlalchemy" %}
_setup_db(app)
{%- elif cookiecutter.orm == "ormar" %}
await database.connect()
{%- elif cookiecutter.orm in ["beanie", "psycopg"] %}
await _setup_db(app)
{%- endif %}
{%- if cookiecutter.db_info.name != "none" and cookiecutter.enable_migrations != "True" %}
{%- if cookiecutter.orm in ["ormar", "sqlalchemy"] %}
await _create_tables()
{%- endif %}
{%- endif %}
{%- if cookiecutter.otlp_enabled == "True" %}
setup_opentelemetry(app)
{%- endif %}
{%- if cookiecutter.enable_redis == "True" %}
init_redis(app)
{%- endif %}
{%- if cookiecutter.enable_rmq == "True" %}
init_rabbit(app)
{%- endif %}
{%- if cookiecutter.enable_kafka == "True" %}
await init_kafka(app)
{%- endif %}
{%- if cookiecutter.prometheus_enabled == "True" %}
setup_prometheus(app)
{%- endif %}
app.middleware_stack = app.build_middleware_stack()

@app.on_event("shutdown")
async def _shutdown() -> None: # noqa: WPS430
{%- if cookiecutter.enable_taskiq == "True" %}
if not broker.is_worker_process:
await broker.shutdown()
{%- endif %}
{%- if cookiecutter.orm == "sqlalchemy" %}
await app.state.db_engine.dispose()
{% elif cookiecutter.orm == "ormar" %}
await database.disconnect()
{%- elif cookiecutter.orm == "psycopg" %}
await app.state.db_pool.close()
{%- endif %}
{%- if cookiecutter.enable_redis == "True" %}
await shutdown_redis(app)
{%- endif %}
{%- if cookiecutter.enable_rmq == "True" %}
await shutdown_rabbit(app)
{%- endif %}
{%- if cookiecutter.enable_kafka == "True" %}
await shutdown_kafka(app)
{%- endif %}
{%- if cookiecutter.otlp_enabled == "True" %}
stop_opentelemetry(app)
{%- endif %}
pass # noqa: WPS420
yield

return _shutdown
{%- if cookiecutter.enable_taskiq == "True" %}
if not broker.is_worker_process:
await broker.shutdown()
{%- endif %}
{%- if cookiecutter.orm == "sqlalchemy" %}
await app.state.db_engine.dispose()
{% elif cookiecutter.orm == "ormar" %}
await database.disconnect()
{%- elif cookiecutter.orm == "psycopg" %}
await app.state.db_pool.close()
{%- endif %}
{%- if cookiecutter.enable_redis == "True" %}
await shutdown_redis(app)
{%- endif %}
{%- if cookiecutter.enable_rmq == "True" %}
await shutdown_rabbit(app)
{%- endif %}
{%- if cookiecutter.enable_kafka == "True" %}
await shutdown_kafka(app)
{%- endif %}
{%- if cookiecutter.otlp_enabled == "True" %}
stop_opentelemetry(app)
{%- endif %}
Loading