From 28c6aef1e84d90a3c8355852175b04574893a43e Mon Sep 17 00:00:00 2001 From: Ruochen Pan Date: Thu, 24 Sep 2026 11:28:59 +0800 Subject: [PATCH] fix: preserve datetime storage across SQLModel upgrades (#10211) --- astrbot/core/db/po.py | 29 ++++-- .../db/vec_db/faiss_impl/document_storage.py | 7 +- astrbot/core/knowledge_base/models.py | 19 +++- tests/unit/test_datetime_storage.py | 97 +++++++++++++++++++ 4 files changed, 137 insertions(+), 15 deletions(-) create mode 100644 tests/unit/test_datetime_storage.py diff --git a/astrbot/core/db/po.py b/astrbot/core/db/po.py index 366a07292..9d7c59121 100644 --- a/astrbot/core/db/po.py +++ b/astrbot/core/db/po.py @@ -4,14 +4,19 @@ from datetime import datetime, timezone from typing import TypedDict from deprecated import deprecated -from sqlalchemy import Index, desc +from sqlalchemy import DateTime, Index, desc from sqlmodel import JSON, Field, SQLModel, Text, UniqueConstraint class TimestampMixin(SQLModel): - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + # Preserve legacy storage instead of SQLModel's inferred UTCDateTime type. + created_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_type=DateTime(timezone=False), + ) updated_at: datetime = Field( default_factory=lambda: datetime.now(timezone.utc), + sa_type=DateTime(timezone=False), sa_column_kwargs={"onupdate": lambda: datetime.now(timezone.utc)}, ) @@ -25,7 +30,7 @@ class PlatformStat(SQLModel, table=True): __tablename__: str = "platform_stats" id: int = Field(primary_key=True, sa_column_kwargs={"autoincrement": True}) - timestamp: datetime = Field(nullable=False) + timestamp: datetime = Field(nullable=False, sa_type=DateTime(timezone=False)) platform_id: str = Field(nullable=False) platform_type: str = Field(nullable=False) # such as "aiocqhttp", "slack", etc. count: int = Field(default=0, nullable=False) @@ -204,8 +209,10 @@ class CronJob(TimestampMixin, SQLModel, table=True): persistent: bool = Field(default=True) run_once: bool = Field(default=False) status: str = Field(default="scheduled", max_length=32) - last_run_at: datetime | None = Field(default=None) - next_run_time: datetime | None = Field(default=None) + last_run_at: datetime | None = Field(default=None, sa_type=DateTime(timezone=False)) + next_run_time: datetime | None = Field( + default=None, sa_type=DateTime(timezone=False) + ) last_error: str | None = Field(default=None, sa_type=Text) @@ -411,9 +418,11 @@ class ApiKey(TimestampMixin, SQLModel, table=True): key_prefix: str = Field(max_length=24, nullable=False) scopes: list | None = Field(default=None, sa_type=JSON) created_by: str = Field(max_length=255, nullable=False) - last_used_at: datetime | None = Field(default=None) - expires_at: datetime | None = Field(default=None) - revoked_at: datetime | None = Field(default=None) + last_used_at: datetime | None = Field( + default=None, sa_type=DateTime(timezone=False) + ) + expires_at: datetime | None = Field(default=None, sa_type=DateTime(timezone=False)) + revoked_at: datetime | None = Field(default=None, sa_type=DateTime(timezone=False)) __table_args__ = ( UniqueConstraint( @@ -439,7 +448,9 @@ class DashboardTrustedDevice(TimestampMixin, SQLModel, table=True): ) token_hash: str = Field(max_length=64, nullable=False, unique=True, index=True) totp_secret_hash: str = Field(max_length=64, nullable=False, index=True) - expires_at: datetime = Field(nullable=False, index=True) + expires_at: datetime = Field( + nullable=False, index=True, sa_type=DateTime(timezone=False) + ) class ChatUIProject(TimestampMixin, SQLModel, table=True): diff --git a/astrbot/core/db/vec_db/faiss_impl/document_storage.py b/astrbot/core/db/vec_db/faiss_impl/document_storage.py index f45184796..1b2d9ceea 100644 --- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py @@ -4,7 +4,7 @@ from contextlib import asynccontextmanager from datetime import datetime from pathlib import Path -from sqlalchemy import Column, Text, bindparam +from sqlalchemy import Column, DateTime, Text, bindparam from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker from sqlmodel import Field, MetaData, SQLModel, col, func, select, text @@ -37,8 +37,9 @@ class Document(BaseDocModel, table=True): doc_id: str = Field(nullable=False, unique=True) text: str = Field(nullable=False) metadata_: str | None = Field(default=None, sa_column=Column("metadata", Text)) - created_at: datetime | None = Field(default=None) - updated_at: datetime | None = Field(default=None) + # Retain the legacy local timestamps without reinterpreting them as UTC. + created_at: datetime | None = Field(default=None, sa_type=DateTime(timezone=False)) + updated_at: datetime | None = Field(default=None, sa_type=DateTime(timezone=False)) class DocumentStorage: diff --git a/astrbot/core/knowledge_base/models.py b/astrbot/core/knowledge_base/models.py index da919a384..578d1f5cb 100644 --- a/astrbot/core/knowledge_base/models.py +++ b/astrbot/core/knowledge_base/models.py @@ -1,6 +1,7 @@ import uuid from datetime import datetime, timezone +from sqlalchemy import DateTime from sqlmodel import Field, MetaData, SQLModel, Text, UniqueConstraint @@ -40,9 +41,14 @@ class KnowledgeBase(BaseKBModel, table=True): top_k_dense: int | None = Field(default=50, nullable=True) top_k_sparse: int | None = Field(default=50, nullable=True) top_m_final: int | None = Field(default=5, nullable=True) - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + # Keep existing database timestamps unchanged across SQLModel upgrades. + created_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_type=DateTime(timezone=False), + ) updated_at: datetime = Field( default_factory=lambda: datetime.now(timezone.utc), + sa_type=DateTime(timezone=False), sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, ) doc_count: int = Field(default=0, nullable=False) @@ -83,9 +89,13 @@ class KBDocument(BaseKBModel, table=True): file_path: str = Field(max_length=512, nullable=False) chunk_count: int = Field(default=0, nullable=False) media_count: int = Field(default=0, nullable=False) - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + created_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_type=DateTime(timezone=False), + ) updated_at: datetime = Field( default_factory=lambda: datetime.now(timezone.utc), + sa_type=DateTime(timezone=False), sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, ) @@ -117,4 +127,7 @@ class KBMedia(BaseKBModel, table=True): file_path: str = Field(max_length=512, nullable=False) file_size: int = Field(nullable=False) mime_type: str = Field(max_length=100, nullable=False) - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + created_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_type=DateTime(timezone=False), + ) diff --git a/tests/unit/test_datetime_storage.py b/tests/unit/test_datetime_storage.py new file mode 100644 index 000000000..9684d064d --- /dev/null +++ b/tests/unit/test_datetime_storage.py @@ -0,0 +1,97 @@ +"""Regression coverage for the existing SQLite datetime representation.""" + +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest +from sqlalchemy import Column, DateTime, MetaData, create_engine, literal, select +from sqlmodel import SQLModel + +from astrbot.core.db.po import ConversationV2, PlatformStat +from astrbot.core.db.vec_db.faiss_impl.document_storage import BaseDocModel, Document +from astrbot.core.knowledge_base.models import BaseKBModel, KnowledgeBase + + +@pytest.mark.parametrize( + "metadata", + [SQLModel.metadata, BaseKBModel.metadata, BaseDocModel.metadata], + ids=["main", "knowledge-base", "documents"], +) +@pytest.mark.parametrize( + "value", + [ + None, + datetime(2024, 1, 2, 3, 4, 5, 123456), + datetime(2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc), + datetime(2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone(timedelta(hours=8))), + ], + ids=["null", "naive", "utc", "non-utc"], +) +def test_datetime_columns_preserve_wall_time( + metadata: MetaData, value: datetime | None +): + """Exercise each datetime column's SQLite bind and result processing. + + Args: + metadata: One of AstrBot's three database model registries. + value: An input accepted by the existing datetime storage contract. + """ + columns = [ + column + for table in metadata.tables.values() + for column in table.columns + if isinstance(column.type, DateTime) + or isinstance(getattr(column.type, "impl", None), DateTime) + ] + assert columns + expected = value.replace(tzinfo=None) if value is not None else None + engine = create_engine("sqlite://") + try: + with engine.connect() as connection: + for column in columns: + # SQLite historically stores the wall time without an offset. + result = connection.scalar(select(literal(value, type_=column.type))) + assert result == expected, str(column) + finally: + engine.dispose() + + +@pytest.mark.parametrize( + "column", + [ + PlatformStat.__table__.c.timestamp, + ConversationV2.__table__.c.created_at, + KnowledgeBase.__table__.c.created_at, + Document.__table__.c.created_at, + ], + ids=["statistics", "inherited-timestamp", "knowledge-base", "documents"], +) +def test_existing_datetime_rows_keep_naive_reads_and_filters( + tmp_path: Path, column: Column +): + """Read legacy DATETIME values without changing their timezone or filtering. + + Args: + tmp_path: Temporary directory for the pre-existing database file. + column: A production model column used to read the legacy values. + """ + engine = create_engine(f"sqlite:///{tmp_path / 'legacy.db'}") + timestamp = datetime(2024, 1, 2, 3, 4, 5, 123456) + try: + with engine.begin() as connection: + # Create the old schema independently of the current model metadata. + connection.exec_driver_sql( + f'CREATE TABLE "{column.table.name}" ("{column.name}" DATETIME)' + ) + connection.exec_driver_sql( + f'INSERT INTO "{column.table.name}" ("{column.name}") VALUES (?)', + (timestamp.isoformat(sep=" "),), + ) + with engine.connect() as connection: + result = connection.execute( + select(column).where(column >= timestamp) + ).scalar_one() + assert result == timestamp + assert result.isoformat() == "2024-01-02T03:04:05.123456" + finally: + engine.dispose()