Replace core database utcnow defaults (#1457)
Co-authored-by: ghreprimand <203024559+ghreprimand@users.noreply.github.com>
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
import os
|
||||
import logging
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from sqlalchemy import event, create_engine, Column, String, Text, Boolean, DateTime, Integer, ForeignKey, JSON, Index, func, text
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.types import TypeDecorator
|
||||
@@ -13,15 +13,21 @@ logger = logging.getLogger(__name__)
|
||||
# Create base class for declarative models
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
def utcnow_naive() -> datetime:
|
||||
"""Return naive UTC for existing DateTime columns."""
|
||||
return datetime.now(timezone.utc).replace(tzinfo=None)
|
||||
|
||||
|
||||
class TimestampMixin:
|
||||
"""Mixin that adds timestamp fields to models"""
|
||||
@declared_attr
|
||||
def created_at(cls):
|
||||
return Column(DateTime, default=datetime.utcnow, nullable=False)
|
||||
return Column(DateTime, default=utcnow_naive, nullable=False)
|
||||
|
||||
@declared_attr
|
||||
def updated_at(cls):
|
||||
return Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
|
||||
return Column(DateTime, default=utcnow_naive, onupdate=utcnow_naive, nullable=False)
|
||||
|
||||
# Get database URL from environment, default to SQLite
|
||||
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./data/app.db")
|
||||
@@ -171,7 +177,7 @@ class ChatMessage(Base):
|
||||
meta_data = Column("metadata", Text, nullable=True) # JSON string for metrics etc.
|
||||
|
||||
# Timestamp
|
||||
timestamp = Column(DateTime, default=datetime.utcnow)
|
||||
timestamp = Column(DateTime, default=utcnow_naive)
|
||||
|
||||
# Relationship to Session
|
||||
session = relationship("Session", back_populates="messages")
|
||||
@@ -224,7 +230,7 @@ class DocumentVersion(Base):
|
||||
content = Column(Text, nullable=False)
|
||||
summary = Column(String, nullable=True) # Edit description
|
||||
source = Column(String, default="ai") # "ai" or "user"
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
created_at = Column(DateTime, default=utcnow_naive)
|
||||
|
||||
document = relationship("Document", back_populates="versions")
|
||||
|
||||
@@ -472,8 +478,8 @@ class UserToolData(Base):
|
||||
tool_id = Column(String, ForeignKey("user_tools.id", ondelete="CASCADE"), nullable=False)
|
||||
key = Column(String, nullable=False)
|
||||
value = Column(Text, nullable=True)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
created_at = Column(DateTime, default=utcnow_naive)
|
||||
updated_at = Column(DateTime, default=utcnow_naive, onupdate=utcnow_naive)
|
||||
|
||||
tool = relationship("UserTool", backref=backref("data_entries", cascade="all, delete-orphan"))
|
||||
|
||||
@@ -592,7 +598,7 @@ class TaskRun(Base):
|
||||
|
||||
id = Column(String, primary_key=True, index=True)
|
||||
task_id = Column(String, ForeignKey("scheduled_tasks.id", ondelete="CASCADE"), nullable=False)
|
||||
started_at = Column(DateTime, nullable=False, default=datetime.utcnow)
|
||||
started_at = Column(DateTime, nullable=False, default=utcnow_naive)
|
||||
finished_at = Column(DateTime, nullable=True)
|
||||
status = Column(String, default="running") # "running", "success", "error"
|
||||
result = Column(Text, nullable=True)
|
||||
@@ -633,7 +639,7 @@ class Memory(Base):
|
||||
session_id = Column(String, ForeignKey("sessions.id", ondelete="SET NULL"), nullable=True, index=True)
|
||||
|
||||
# Timestamp as Unix timestamp
|
||||
timestamp = Column(Integer, default=lambda: int(datetime.utcnow().timestamp()))
|
||||
timestamp = Column(Integer, default=lambda: int(utcnow_naive().timestamp()))
|
||||
|
||||
# Relationship to Session
|
||||
session = relationship("Session", backref="memories")
|
||||
@@ -1480,7 +1486,7 @@ def _migrate_seed_email_account():
|
||||
if not imap_host and not smtp_host:
|
||||
return # nothing to migrate
|
||||
|
||||
now = datetime.utcnow()
|
||||
now = utcnow_naive()
|
||||
with engine.begin() as conn:
|
||||
conn.execute(text("""
|
||||
INSERT INTO email_accounts
|
||||
@@ -1760,7 +1766,7 @@ def bulk_insert_messages(session_id: str, messages: list):
|
||||
'session_id': session_id,
|
||||
'role': msg['role'],
|
||||
'content': msg['content'],
|
||||
'timestamp': datetime.utcnow()
|
||||
'timestamp': utcnow_naive()
|
||||
}
|
||||
for msg in messages
|
||||
]
|
||||
@@ -1771,7 +1777,7 @@ def cleanup_old_sessions(days: int = 30):
|
||||
from datetime import timedelta
|
||||
|
||||
with get_db_session() as db:
|
||||
cutoff_date = datetime.utcnow() - timedelta(days=days)
|
||||
cutoff_date = utcnow_naive() - timedelta(days=days)
|
||||
|
||||
deleted_count = db.query(Session).filter(
|
||||
Session.archived == True,
|
||||
@@ -1816,7 +1822,7 @@ def update_session_last_accessed(session_id: str):
|
||||
with get_db_session() as db:
|
||||
db_session = db.query(Session).filter(Session.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.last_accessed = datetime.utcnow()
|
||||
db_session.last_accessed = utcnow_naive()
|
||||
db.commit()
|
||||
return True
|
||||
return False
|
||||
@@ -1861,7 +1867,7 @@ def get_upcoming_events(owner, horizon_days: int = 60, limit: int = 40):
|
||||
The autonomous email->calendar pass relies on this to avoid disclosing (and
|
||||
acting on) other users' calendars."""
|
||||
from datetime import timedelta
|
||||
now = datetime.utcnow()
|
||||
now = utcnow_naive()
|
||||
with get_db_session() as db:
|
||||
q = db.query(CalendarEvent).join(CalendarCal).filter(
|
||||
CalendarEvent.dtstart >= now,
|
||||
|
||||
33
tests/test_database_utcnow.py
Normal file
33
tests/test_database_utcnow.py
Normal file
@@ -0,0 +1,33 @@
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
sqlalchemy = pytest.importorskip("sqlalchemy")
|
||||
if not isinstance(sqlalchemy, types.ModuleType):
|
||||
pytest.skip("sqlalchemy is stubbed in this environment", allow_module_level=True)
|
||||
|
||||
from core.database import ChatMessage, DocumentVersion, Session, TaskRun, UserToolData, utcnow_naive
|
||||
|
||||
|
||||
def test_utcnow_naive_returns_naive_utc_datetime():
|
||||
now = utcnow_naive()
|
||||
|
||||
assert now.tzinfo is None
|
||||
assert abs((now - utcnow_naive()).total_seconds()) < 2
|
||||
|
||||
|
||||
def test_database_timestamp_defaults_use_utcnow_naive():
|
||||
defaults = (
|
||||
Session.created_at.default.arg,
|
||||
Session.updated_at.default.arg,
|
||||
Session.updated_at.onupdate.arg,
|
||||
ChatMessage.timestamp.default.arg,
|
||||
DocumentVersion.created_at.default.arg,
|
||||
UserToolData.created_at.default.arg,
|
||||
UserToolData.updated_at.default.arg,
|
||||
UserToolData.updated_at.onupdate.arg,
|
||||
TaskRun.started_at.default.arg,
|
||||
)
|
||||
|
||||
for fn in defaults:
|
||||
assert fn.__name__ == "utcnow_naive"
|
||||
Reference in New Issue
Block a user