Update the model datetime column type to custom (#786)

* Update the model datetime column type to custom

* Update the schema datetime filed json encoder

Author: wu-clan

backend/app/admin/model/login_log.py

@@ -2,12 +2,12 @@
 # -*- coding: utf-8 -*-
 from datetime import datetime
 
-from sqlalchemy import DateTime, String
+from sqlalchemy import String
 from sqlalchemy.dialects.mysql import LONGTEXT
 from sqlalchemy.dialects.postgresql import TEXT
 from sqlalchemy.orm import Mapped, mapped_column
 
-from backend.common.model import DataClassBase, id_key
+from backend.common.model import DataClassBase, TimeZone, id_key
 from backend.utils.timezone import timezone
 
 
@@ -29,7 +29,7 @@ class LoginLog(DataClassBase):
     browser: Mapped[str | None] = mapped_column(String(50), comment='浏览器')
     device: Mapped[str | None] = mapped_column(String(50), comment='设备')
     msg: Mapped[str] = mapped_column(LONGTEXT().with_variant(TEXT, 'postgresql'), comment='提示消息')
-    login_time: Mapped[datetime] = mapped_column(DateTime(timezone=True), comment='登录时间')
+    login_time: Mapped[datetime] = mapped_column(TimeZone, comment='登录时间')
     created_time: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), init=False, default_factory=timezone.now, comment='创建时间'
+        TimeZone, init=False, default_factory=timezone.now, comment='创建时间'
     )

backend/app/admin/model/opera_log.py

@@ -2,12 +2,12 @@
 # -*- coding: utf-8 -*-
 from datetime import datetime
 
-from sqlalchemy import DateTime, String
+from sqlalchemy import String
 from sqlalchemy.dialects.mysql import JSON, LONGTEXT
 from sqlalchemy.dialects.postgresql import TEXT
 from sqlalchemy.orm import Mapped, mapped_column
 
-from backend.common.model import DataClassBase, id_key
+from backend.common.model import DataClassBase, TimeZone, id_key
 from backend.utils.timezone import timezone
 
 
@@ -35,7 +35,7 @@ class OperaLog(DataClassBase):
     code: Mapped[str] = mapped_column(String(20), insert_default='200', comment='操作状态码')
     msg: Mapped[str | None] = mapped_column(LONGTEXT().with_variant(TEXT, 'postgresql'), comment='提示消息')
     cost_time: Mapped[float] = mapped_column(insert_default=0.0, comment='请求耗时（ms）')
-    opera_time: Mapped[datetime] = mapped_column(DateTime(timezone=True), comment='操作时间')
+    opera_time: Mapped[datetime] = mapped_column(TimeZone, comment='操作时间')
     created_time: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), init=False, default_factory=timezone.now, comment='创建时间'
+        TimeZone, init=False, default_factory=timezone.now, comment='创建时间'
     )

backend/app/admin/model/user.py

@@ -5,12 +5,12 @@
 from datetime import datetime
 from typing import TYPE_CHECKING
 
-from sqlalchemy import VARBINARY, Boolean, DateTime, ForeignKey, String
+from sqlalchemy import VARBINARY, Boolean, ForeignKey, String
 from sqlalchemy.dialects.postgresql import BYTEA, INTEGER
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from backend.app.admin.model.m2m import sys_user_role
-from backend.common.model import Base, id_key
+from backend.common.model import Base, TimeZone, id_key
 from backend.database.db import uuid4_str
 from backend.utils.timezone import timezone
 
@@ -42,11 +42,9 @@ class User(Base):
     is_multi_login: Mapped[bool] = mapped_column(
         Boolean().with_variant(INTEGER, 'postgresql'), default=False, comment='是否重复登陆(0否 1是)'
     )
-    join_time: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), init=False, default_factory=timezone.now, comment='注册时间'
-    )
+    join_time: Mapped[datetime] = mapped_column(TimeZone, init=False, default_factory=timezone.now, comment='注册时间')
     last_login_time: Mapped[datetime | None] = mapped_column(
-        DateTime(timezone=True), init=False, onupdate=timezone.now, comment='上次登录'
+        TimeZone, init=False, onupdate=timezone.now, comment='上次登录'
     )
 
     # 部门用户一对多

backend/app/task/model/scheduler.py

@@ -7,7 +7,6 @@
 from sqlalchemy import (
     JSON,
     Boolean,
-    DateTime,
     String,
     event,
 )
@@ -16,7 +15,7 @@
 from sqlalchemy.orm import Mapped, mapped_column
 
 from backend.common.exception import errors
-from backend.common.model import Base, id_key
+from backend.common.model import Base, TimeZone, id_key
 from backend.core.conf import settings
 from backend.database.redis import redis_client
 from backend.utils.timezone import timezone
@@ -35,8 +34,8 @@ class TaskScheduler(Base):
     queue: Mapped[str | None] = mapped_column(String(255), comment='CELERY_TASK_QUEUES 中定义的队列')
     exchange: Mapped[str | None] = mapped_column(String(255), comment='低级别 AMQP 路由的交换机')
     routing_key: Mapped[str | None] = mapped_column(String(255), comment='低级别 AMQP 路由的路由密钥')
-    start_time: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), comment='任务开始触发的时间')
-    expire_time: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), comment='任务不再触发的截止时间')
+    start_time: Mapped[datetime | None] = mapped_column(TimeZone, comment='任务开始触发的时间')
+    expire_time: Mapped[datetime | None] = mapped_column(TimeZone, comment='任务不再触发的截止时间')
     expire_seconds: Mapped[int | None] = mapped_column(comment='任务不再触发的秒数时间差')
     type: Mapped[int] = mapped_column(comment='调度类型（0间隔 1定时）')
     interval_every: Mapped[int | None] = mapped_column(comment='任务再次运行前的间隔周期数')
@@ -49,9 +48,7 @@ class TaskScheduler(Base):
         Boolean().with_variant(INTEGER, 'postgresql'), default=True, comment='是否启用任务'
     )
     total_run_count: Mapped[int] = mapped_column(default=0, comment='任务触发的总次数')
-    last_run_time: Mapped[datetime | None] = mapped_column(
-        DateTime(timezone=True), default=None, comment='任务最后触发的时间'
-    )
+    last_run_time: Mapped[datetime | None] = mapped_column(TimeZone, default=None, comment='任务最后触发的时间')
     remark: Mapped[str | None] = mapped_column(
         LONGTEXT().with_variant(TEXT, 'postgresql'), default=None, comment='备注'
     )

backend/common/model.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from typing import Annotated
 
-from sqlalchemy import BigInteger, DateTime
+from sqlalchemy import BigInteger, DateTime, TypeDecorator
 from sqlalchemy.ext.asyncio import AsyncAttrs
 from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, declared_attr, mapped_column
 
@@ -51,14 +51,37 @@ class UserMixin(MappedAsDataclass):
     updated_by: Mapped[int | None] = mapped_column(init=False, default=None, sort_order=998, comment='修改者')
 
 
+class TimeZone(TypeDecorator[datetime]):
+    """时区感知 DateTime"""
+
+    impl = DateTime(timezone=True)
+    cache_ok = True
+
+    @property
+    def python_type(self) -> type[datetime]:
+        return datetime
+
+    def process_bind_param(self, value: datetime | None, dialect) -> datetime | None:
+        if value is not None:
+            if value.tzinfo != timezone.tz_info:
+                value = timezone.from_datetime(value)
+        return value
+
+    def process_result_value(self, value: datetime | None, dialect) -> datetime | None:
+        if value is not None:
+            if value.tzinfo is None:
+                value = value.replace(tzinfo=timezone.tz_info)
+        return value
+
+
 class DateTimeMixin(MappedAsDataclass):
     """日期时间 Mixin 数据类"""
 
     created_time: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), init=False, default_factory=timezone.now, sort_order=999, comment='创建时间'
+        TimeZone, init=False, default_factory=timezone.now, sort_order=999, comment='创建时间'
     )
     updated_time: Mapped[datetime | None] = mapped_column(
-        DateTime(timezone=True), init=False, onupdate=timezone.now, sort_order=999, comment='更新时间'
+        TimeZone, init=False, onupdate=timezone.now, sort_order=999, comment='更新时间'
     )
 
 

backend/common/schema.py

@@ -25,7 +25,7 @@ class SchemaBase(BaseModel):
         use_enum_values=True,
         json_encoders={
             datetime: lambda x: timezone.to_str(timezone.from_datetime(x))
-            if x.tzinfo is not None
+            if x.tzinfo is not None and x.tzinfo != timezone.tz_info
             else timezone.to_str(x)
         },
     )