diff --git a/src/basic_memory/models.py b/src/basic_memory/models.py index 59b956d1..a692ab95 100644 --- a/src/basic_memory/models.py +++ b/src/basic_memory/models.py @@ -1,34 +1,11 @@ """Database models for basic-memory.""" -from datetime import datetime, UTC +from datetime import datetime from typing import List, Optional -from sqlalchemy import String, DateTime, ForeignKey, Text, TypeDecorator, Integer, text, UniqueConstraint +from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, UniqueConstraint from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase from sqlalchemy.ext.asyncio import AsyncAttrs -class UTCDateTime(TypeDecorator): - """Automatically handle UTC timezone for datetime columns""" - impl = DateTime - cache_ok = True - - def process_bind_param(self, value: Optional[datetime], dialect): - if value is not None: - if value.tzinfo is None: - return value.replace(tzinfo=UTC) - return value - return value - - def process_result_value(self, value: Optional[datetime], dialect): - if value is not None: - return value.replace(tzinfo=UTC) - return value - - -def utc_now() -> datetime: - """Helper to get current UTC time""" - return datetime.now(UTC) - - class Base(AsyncAttrs, DeclarativeBase): """Base class for all models""" pass @@ -55,13 +32,13 @@ class Entity(Base): entity_type: Mapped[str] = mapped_column(String) description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) created_at: Mapped[datetime] = mapped_column( - UTCDateTime, + DateTime, server_default=text('CURRENT_TIMESTAMP') ) updated_at: Mapped[datetime] = mapped_column( - UTCDateTime, + DateTime, server_default=text('CURRENT_TIMESTAMP'), - onupdate=utc_now + onupdate=text('CURRENT_TIMESTAMP') ) # Relationships @@ -152,7 +129,7 @@ class Relation(Base): ) relation_type: Mapped[str] = mapped_column(String) created_at: Mapped[datetime] = mapped_column( - UTCDateTime, + DateTime, server_default=text('CURRENT_TIMESTAMP') ) context: Mapped[str | None] = mapped_column(String, nullable=True) diff --git a/tests/test_entity_repository.py b/tests/test_entity_repository.py index 8758b948..d8987a8e 100644 --- a/tests/test_entity_repository.py +++ b/tests/test_entity_repository.py @@ -25,7 +25,6 @@ class TestEntityRepository: assert entity.name == 'Test' assert entity.description == 'Test description' assert isinstance(entity.created_at, datetime) - assert entity.created_at.tzinfo == UTC # Verify in database stmt = select(Entity).where(Entity.id == entity.id)