From c8337cf08085d0c8117d47a0c7c237ff0c14b211 Mon Sep 17 00:00:00 2001 From: Abeautifulsnow <1491444340@qq.com> Date: Tue, 11 Jun 2024 22:37:58 +0800 Subject: [PATCH 1/5] add use_tz and timezone in RegisterTortoise --- tortoise/contrib/fastapi/__init__.py | 32 +++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/tortoise/contrib/fastapi/__init__.py b/tortoise/contrib/fastapi/__init__.py index aa9df8874..b33af6b62 100644 --- a/tortoise/contrib/fastapi/__init__.py +++ b/tortoise/contrib/fastapi/__init__.py @@ -96,6 +96,8 @@ def __init__( modules: Optional[Dict[str, Iterable[Union[str, ModuleType]]]] = None, generate_schemas: bool = False, add_exception_handlers: bool = False, + use_tz: bool = False, + timezone: str = "UTC", ) -> None: self.app = app self.config = config @@ -103,24 +105,44 @@ def __init__( self.db_url = db_url self.modules = modules self.generate_schemas = generate_schemas + self.use_tz = use_tz + self.timezone = timezone + if add_exception_handlers: @app.exception_handler(DoesNotExist) - async def doesnotexist_exception_handler(request: "Request", exc: DoesNotExist): + async def doesnotexist_exception_handler( + request: "Request", exc: DoesNotExist + ): return JSONResponse(status_code=404, content={"detail": str(exc)}) @app.exception_handler(IntegrityError) - async def integrityerror_exception_handler(request: "Request", exc: IntegrityError): + async def integrityerror_exception_handler( + request: "Request", exc: IntegrityError + ): return JSONResponse( status_code=422, - content={"detail": [{"loc": [], "msg": str(exc), "type": "IntegrityError"}]}, + content={ + "detail": [ + {"loc": [], "msg": str(exc), "type": "IntegrityError"} + ] + }, ) async def init_orm(self) -> None: # pylint: disable=W0612 config, config_file = self.config, self.config_file db_url, modules = self.db_url, self.modules - await Tortoise.init(config=config, config_file=config_file, db_url=db_url, modules=modules) - logger.info("Tortoise-ORM started, %s, %s", connections._get_storage(), Tortoise.apps) + await Tortoise.init( + config=config, + config_file=config_file, + db_url=db_url, + modules=modules, + use_tz=self.use_tz, + timezone=self.timezone, + ) + logger.info( + "Tortoise-ORM started, %s, %s", connections._get_storage(), Tortoise.apps + ) if self.generate_schemas: logger.info("Tortoise-ORM generating schema") await Tortoise.generate_schemas() From 720059ab172c927949cb46e1f2f57fda8314aa98 Mon Sep 17 00:00:00 2001 From: 15639936570 Date: Fri, 14 Jun 2024 10:49:53 +0800 Subject: [PATCH 2/5] add tests for east-8 timezone --- examples/fastapi/_tests.py | 32 +++++++++++++++++++++++- examples/fastapi/main.py | 50 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 1 deletion(-) diff --git a/examples/fastapi/_tests.py b/examples/fastapi/_tests.py index f5a4443d0..454511a9d 100644 --- a/examples/fastapi/_tests.py +++ b/examples/fastapi/_tests.py @@ -1,11 +1,12 @@ # mypy: no-disallow-untyped-decorators # pylint: disable=E0611,E0401 +import datetime from typing import AsyncGenerator import pytest from asgi_lifespan import LifespanManager from httpx import ASGITransport, AsyncClient -from main import app +from main import app, app_east from models import Users @@ -33,3 +34,32 @@ async def test_create_user(client: AsyncClient) -> None: # nosec user_obj = await Users.get(id=user_id) assert user_obj.id == user_id + + # UTC timezone. + created_at = user_obj.created_at + assert created_at.hour - datetime.datetime.now().hour == -8 + + +@pytest.fixture(scope="module") +async def client_east() -> AsyncGenerator[AsyncClient, None]: + async with LifespanManager(app_east): + transport = ASGITransport(app=app_east) + async with AsyncClient(transport=transport, base_url="http://test") as c: + yield c + + +@pytest.mark.anyio +async def test_create_user_east(client_east: AsyncClient) -> None: # nosec + response = await client_east.post("/users_east", json={"username": "admin"}) + assert response.status_code == 200, response.text + data = response.json() + assert data["username"] == "admin" + assert "id" in data + user_id = data["id"] + + user_obj = await Users.get(id=user_id) + assert user_obj.id == user_id + + # Verify that the time zone is East 8. + created_at = user_obj.created_at + assert created_at.hour - datetime.datetime.now().hour == 0 diff --git a/examples/fastapi/main.py b/examples/fastapi/main.py index 0c6a28c49..9c994dbb0 100644 --- a/examples/fastapi/main.py +++ b/examples/fastapi/main.py @@ -37,7 +37,26 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # db connections closed +@asynccontextmanager +async def lifespan_east(app: FastAPI) -> AsyncGenerator[None, None]: + # app startup + async with RegisterTortoise( + app, + db_url="sqlite://:memory:", + modules={"models": ["models"]}, + generate_schemas=True, + add_exception_handlers=True, + use_tz=False, + timezone="Asia/Shanghai", + ): + # db connected + yield + # app teardown + # db connections closed + + app = FastAPI(title="Tortoise ORM FastAPI example", lifespan=lifespan) +app_east = FastAPI(title="Tortoise ORM FastAPI example", lifespan=lifespan_east) class Status(BaseModel): @@ -72,3 +91,34 @@ async def delete_user(user_id: int): if not deleted_count: raise HTTPException(status_code=404, detail=f"User {user_id} not found") return Status(message=f"Deleted user {user_id}") + + +############################ East 8 ############################ +@app_east.get("/users_east", response_model=List[User_Pydantic]) +async def get_users_east(): + return await User_Pydantic.from_queryset(Users.all()) + + +@app_east.post("/users_east", response_model=User_Pydantic) +async def create_user_east(user: UserIn_Pydantic): + user_obj = await Users.create(**user.model_dump(exclude_unset=True)) + return await User_Pydantic.from_tortoise_orm(user_obj) + + +@app_east.get("/user_east/{user_id}", response_model=User_Pydantic) +async def get_user_east(user_id: int): + return await User_Pydantic.from_queryset_single(Users.get(id=user_id)) + + +@app_east.put("/user_east/{user_id}", response_model=User_Pydantic) +async def update_user_east(user_id: int, user: UserIn_Pydantic): + await Users.filter(id=user_id).update(**user.model_dump(exclude_unset=True)) + return await User_Pydantic.from_queryset_single(Users.get(id=user_id)) + + +@app_east.delete("/user_east/{user_id}", response_model=Status) +async def delete_user_east(user_id: int): + deleted_count = await Users.filter(id=user_id).delete() + if not deleted_count: + raise HTTPException(status_code=404, detail=f"User {user_id} not found") + return Status(message=f"Deleted user {user_id}") From 1eb61701c2f11aca8d783bc4c407aa21cc56d338 Mon Sep 17 00:00:00 2001 From: Abeautifulsnow <1491444340@qq.com> Date: Fri, 14 Jun 2024 10:53:02 +0800 Subject: [PATCH 3/5] Update main.py --- examples/fastapi/main.py | 50 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/examples/fastapi/main.py b/examples/fastapi/main.py index 0c6a28c49..9c994dbb0 100644 --- a/examples/fastapi/main.py +++ b/examples/fastapi/main.py @@ -37,7 +37,26 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # db connections closed +@asynccontextmanager +async def lifespan_east(app: FastAPI) -> AsyncGenerator[None, None]: + # app startup + async with RegisterTortoise( + app, + db_url="sqlite://:memory:", + modules={"models": ["models"]}, + generate_schemas=True, + add_exception_handlers=True, + use_tz=False, + timezone="Asia/Shanghai", + ): + # db connected + yield + # app teardown + # db connections closed + + app = FastAPI(title="Tortoise ORM FastAPI example", lifespan=lifespan) +app_east = FastAPI(title="Tortoise ORM FastAPI example", lifespan=lifespan_east) class Status(BaseModel): @@ -72,3 +91,34 @@ async def delete_user(user_id: int): if not deleted_count: raise HTTPException(status_code=404, detail=f"User {user_id} not found") return Status(message=f"Deleted user {user_id}") + + +############################ East 8 ############################ +@app_east.get("/users_east", response_model=List[User_Pydantic]) +async def get_users_east(): + return await User_Pydantic.from_queryset(Users.all()) + + +@app_east.post("/users_east", response_model=User_Pydantic) +async def create_user_east(user: UserIn_Pydantic): + user_obj = await Users.create(**user.model_dump(exclude_unset=True)) + return await User_Pydantic.from_tortoise_orm(user_obj) + + +@app_east.get("/user_east/{user_id}", response_model=User_Pydantic) +async def get_user_east(user_id: int): + return await User_Pydantic.from_queryset_single(Users.get(id=user_id)) + + +@app_east.put("/user_east/{user_id}", response_model=User_Pydantic) +async def update_user_east(user_id: int, user: UserIn_Pydantic): + await Users.filter(id=user_id).update(**user.model_dump(exclude_unset=True)) + return await User_Pydantic.from_queryset_single(Users.get(id=user_id)) + + +@app_east.delete("/user_east/{user_id}", response_model=Status) +async def delete_user_east(user_id: int): + deleted_count = await Users.filter(id=user_id).delete() + if not deleted_count: + raise HTTPException(status_code=404, detail=f"User {user_id} not found") + return Status(message=f"Deleted user {user_id}") From ffaa45d6119fe1ff7a4487e062371a932dbb8e13 Mon Sep 17 00:00:00 2001 From: Abeautifulsnow <1491444340@qq.com> Date: Fri, 14 Jun 2024 10:53:19 +0800 Subject: [PATCH 4/5] Update _tests.py --- examples/fastapi/_tests.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/examples/fastapi/_tests.py b/examples/fastapi/_tests.py index f5a4443d0..454511a9d 100644 --- a/examples/fastapi/_tests.py +++ b/examples/fastapi/_tests.py @@ -1,11 +1,12 @@ # mypy: no-disallow-untyped-decorators # pylint: disable=E0611,E0401 +import datetime from typing import AsyncGenerator import pytest from asgi_lifespan import LifespanManager from httpx import ASGITransport, AsyncClient -from main import app +from main import app, app_east from models import Users @@ -33,3 +34,32 @@ async def test_create_user(client: AsyncClient) -> None: # nosec user_obj = await Users.get(id=user_id) assert user_obj.id == user_id + + # UTC timezone. + created_at = user_obj.created_at + assert created_at.hour - datetime.datetime.now().hour == -8 + + +@pytest.fixture(scope="module") +async def client_east() -> AsyncGenerator[AsyncClient, None]: + async with LifespanManager(app_east): + transport = ASGITransport(app=app_east) + async with AsyncClient(transport=transport, base_url="http://test") as c: + yield c + + +@pytest.mark.anyio +async def test_create_user_east(client_east: AsyncClient) -> None: # nosec + response = await client_east.post("/users_east", json={"username": "admin"}) + assert response.status_code == 200, response.text + data = response.json() + assert data["username"] == "admin" + assert "id" in data + user_id = data["id"] + + user_obj = await Users.get(id=user_id) + assert user_obj.id == user_id + + # Verify that the time zone is East 8. + created_at = user_obj.created_at + assert created_at.hour - datetime.datetime.now().hour == 0 From 46e88470166237707ec52796998882278ff19348 Mon Sep 17 00:00:00 2001 From: 15639936570 Date: Tue, 18 Jun 2024 14:06:38 +0800 Subject: [PATCH 5/5] add code to adjust the timezone --- examples/fastapi/_tests.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/examples/fastapi/_tests.py b/examples/fastapi/_tests.py index 454511a9d..dd9f5203e 100644 --- a/examples/fastapi/_tests.py +++ b/examples/fastapi/_tests.py @@ -4,6 +4,7 @@ from typing import AsyncGenerator import pytest +import pytz from asgi_lifespan import LifespanManager from httpx import ASGITransport, AsyncClient from main import app, app_east @@ -35,10 +36,6 @@ async def test_create_user(client: AsyncClient) -> None: # nosec user_obj = await Users.get(id=user_id) assert user_obj.id == user_id - # UTC timezone. - created_at = user_obj.created_at - assert created_at.hour - datetime.datetime.now().hour == -8 - @pytest.fixture(scope="module") async def client_east() -> AsyncGenerator[AsyncClient, None]: @@ -62,4 +59,13 @@ async def test_create_user_east(client_east: AsyncClient) -> None: # nosec # Verify that the time zone is East 8. created_at = user_obj.created_at - assert created_at.hour - datetime.datetime.now().hour == 0 + + # Asia/Shanghai timezone + asia_tz = pytz.timezone("Asia/Shanghai") + asia_now = datetime.datetime.now(pytz.utc).astimezone(asia_tz) + assert created_at.hour - asia_now.hour == 0 + + # UTC timezone + utc_tz = pytz.timezone("UTC") + utc_now = datetime.datetime.now(pytz.utc).astimezone(utc_tz) + assert created_at.hour - utc_now.hour == 8