Skip to content

Commit

Permalink
Updateted lifespan setup.
Browse files Browse the repository at this point in the history 
Signed-off-by: Pavel Kirilin <[email protected]>
  • Loading branch information
@s3rius
s3rius committed Jul 13, 2024
1 parent ec5f725 commit 6a6271a
Show file tree
Hide file tree
Showing 7 changed files with 69 additions and 88 deletions.
2 changes: 1 addition & 1 deletion fastapi_template/template/{{cookiecutter.project_name}}/README.md
View file
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
4 changes: 2 additions & 2 deletions fastapi_template/template/{{cookiecutter.project_name}}/tests/conftest.py
View file
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
0 ...project_name}}/services/kafka/lifetime.py → ...project_name}}/services/kafka/lifespan.py
View file
File renamed without changes.
0 ...roject_name}}/services/rabbit/lifetime.py → ...roject_name}}/services/rabbit/lifespan.py
View file
File renamed without changes.
0 ...project_name}}/services/redis/lifetime.py → ...project_name}}/services/redis/lifespan.py
View file
File renamed without changes.
8 changes: 2 additions & 6 deletions ...e/template/{{cookiecutter.project_name}}/{{cookiecutter.project_name}}/web/application.py
View file
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
143 changes: 64 additions & 79 deletions ...okiecutter.project_name}}/web/lifetime.py → ...okiecutter.project_name}}/web/lifespan.py
View file
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 %}

0 comments on commit 6a6271a

Please sign in to comment.