FastAPI结合SQLAlchemy处理时区感知对象的技术问题

问题背景 #

在FastAPI+PostgreSQL环境中，查询带时区的timestamp字段时，返回结果始终为UTC时间。即便通过SET TIMEZONE修改会话时区，输出仍保持UTC格式。需要根据终端用户位置动态调整会话时区，但不想为每个查询手动添加func.timezone()方法。

示例代码 #

async def query_database():
    """Query the database and display timezone information."""
    async with async_session_factory() as session:
        for timezone in ["America/Toronto", "Europe/London"]:
            await session.execute(text(f"SET TIME ZONE '{timezone}';"))
            timezone_result = await session.scalar(text("SHOW TIMEZONE;"))
            print(f"Current timezone: {timezone_result}")

            entity = await session.scalar(select(TimezoneAwareModel))
            print(f"Entity created at: {entity.created_at}")

输出结果 #

Current timezone: America/Toronto
Entity created at: 2024-07-15 10:49:07.182811+00:00
Current timezone: Europe/London
Entity created at: 2024-07-15 10:49:07.182811+00:00

可行解决方案 #

1. 自定义SQLAlchemy时区感知字段 #

通过自定义TypeDecorator，让字段在返回结果时自动转换为会话时区：

from sqlalchemy import DateTime, TypeDecorator
from sqlalchemy.ext.compiler import compiles
import pytz

class TZDateTime(TypeDecorator):
    impl = DateTime(timezone=True)

    def process_result_value(self, value, dialect):
        if value is not None:
            # 从会话信息中获取预设时区，默认UTC
            tz_name = self.session.info.get('timezone', 'UTC')
            tz = pytz.timezone(tz_name)
            return value.astimezone(tz)

@compiles(TZDateTime, 'postgresql')
def compile_tz_datetime(type_, compiler, **kw):
    return compiler.visit_TIMESTAMP(timezone=True, **kw)

• 模型中使用该字段：

class TimezoneAwareModel(Base):
    __tablename__ = 'timezone_aware'
    id = Column(Integer, primary_key=True)
    created_at = Column(TZDateTime)

• 会话初始化时注入用户时区：

async def get_session(user_timezone: str = 'UTC'):
    async with async_session_factory() as session:
        session.info['timezone'] = user_timezone
        await session.execute(text(f"SET TIME ZONE '{user_timezone}';"))
        yield session

2. 创建PostgreSQL时区转换视图 #

在数据库层面创建视图，自动将timestamp字段转换为当前会话时区：

CREATE OR REPLACE VIEW timezone_aware_view AS
SELECT 
    id,
    created_at AT TIME ZONE current_setting('TIMEZONE') AS created_at
FROM timezone_aware;

• 在SQLAlchemy中映射该视图：

class TimezoneAwareModelView(Base):
    __tablename__ = 'timezone_aware_view'
    id = Column(Integer, primary_key=True)
    created_at = Column(DateTime)

查询视图时，返回的created_at会自动适配会话时区，无需额外代码处理。

3. 会话级事件监听自动转换 #

利用SQLAlchemy的事件监听机制，在会话开始时设置时区，查询完成后自动转换时间字段：

from sqlalchemy import event
from sqlalchemy.ext.asyncio import AsyncSession
import pytz

@event.listens_for(AsyncSession, 'after_begin')
async def set_session_timezone(session, transaction, connection):
    timezone = session.info.get('timezone', 'UTC')
    await connection.execute(text(f"SET TIME ZONE '{timezone}';"))

@event.listens_for(AsyncSession, 'after_execute')
async def convert_result_timezone(session, result, execution_options):
    timezone = session.info.get('timezone', 'UTC')
    tz_obj = pytz.timezone(timezone)
    # 处理标量结果
    if hasattr(result, 'scalars'):
        for entity in result.scalars().all():
            if hasattr(entity, 'created_at') and entity.created_at:
                entity.created_at = entity.created_at.astimezone(tz_obj)

这种方式对现有代码侵入性极低，只需在会话初始化时传入用户时区即可。

内容的提问来源于stack exchange，提问作者Palaniyappan