Add timezone and use_utc parameters into RegisterTortoise class. (

#1649)

* add use_tz and timezone in RegisterTortoise

* add tests for east-8 timezone

* Update main.py

* Update _tests.py

* add code to adjust the timezone

examples/fastapi/_tests.py

@@ -1,11 +1,13 @@
 # mypy: no-disallow-untyped-decorators
 # pylint: disable=E0611,E0401
+import datetime
 from typing import AsyncGenerator
 
 import pytest
+import pytz
 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 +35,37 @@ async def test_create_user(client: AsyncClient) -> None:  # nosec
 
     user_obj = await Users.get(id=user_id)
     assert user_obj.id == user_id
+
+
+@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
+
+    # 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

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}")

tortoise/contrib/fastapi/__init__.py

@@ -96,31 +96,53 @@ 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
         self.config_file = config_file
         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()