72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
"""shift all naive timestamps to beijing and set db timezone
|
|
|
|
Revision ID: shift_all_timestamps_to_beijing
|
|
Revises: add_conversation_sessions_table
|
|
Create Date: 2026-03-11
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
|
|
revision = "shift_all_timestamps_to_beijing"
|
|
down_revision = "add_conversation_sessions_table"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def _is_postgresql() -> bool:
|
|
bind = op.get_bind()
|
|
return str(getattr(getattr(bind, "dialect", None), "name", "")).lower().startswith("postgres")
|
|
|
|
|
|
def _list_timestamp_columns() -> list[tuple[str, str]]:
|
|
bind = op.get_bind()
|
|
rows = bind.execute(
|
|
sa.text(
|
|
"""
|
|
SELECT table_name, column_name
|
|
FROM information_schema.columns
|
|
WHERE table_schema = 'public'
|
|
AND data_type = 'timestamp without time zone'
|
|
ORDER BY table_name, column_name
|
|
"""
|
|
)
|
|
).fetchall()
|
|
return [(str(row[0]), str(row[1])) for row in rows]
|
|
|
|
|
|
def _shift_all_timestamp_columns(*, hours: int) -> None:
|
|
delta_expr = f"INTERVAL '{int(hours)} hours'" if int(hours) >= 0 else f"- INTERVAL '{abs(int(hours))} hours'"
|
|
for table_name, column_name in _list_timestamp_columns():
|
|
if int(hours) >= 0:
|
|
sql = f'UPDATE "{table_name}" SET "{column_name}" = "{column_name}" + {delta_expr} WHERE "{column_name}" IS NOT NULL'
|
|
else:
|
|
sql = f'UPDATE "{table_name}" SET "{column_name}" = "{column_name}" {delta_expr} WHERE "{column_name}" IS NOT NULL'
|
|
op.execute(sa.text(sql))
|
|
|
|
|
|
def _set_database_timezone(timezone_name: str) -> None:
|
|
bind = op.get_bind()
|
|
database_name = str(bind.execute(sa.text("SELECT current_database()")).scalar_one())
|
|
role_name = str(bind.execute(sa.text("SELECT current_user")).scalar_one())
|
|
op.execute(sa.text(f'ALTER DATABASE "{database_name}" SET timezone TO \'{timezone_name}\''))
|
|
op.execute(sa.text(f'ALTER ROLE "{role_name}" SET timezone TO \'{timezone_name}\''))
|
|
|
|
|
|
def upgrade() -> None:
|
|
if not _is_postgresql():
|
|
return
|
|
_set_database_timezone("Asia/Shanghai")
|
|
_shift_all_timestamp_columns(hours=8)
|
|
|
|
|
|
def downgrade() -> None:
|
|
if not _is_postgresql():
|
|
return
|
|
_shift_all_timestamp_columns(hours=-8)
|
|
# 统一策略：即使降级历史数据，也保持数据库会话时区为北京时间。
|
|
_set_database_timezone("Asia/Shanghai")
|