fix: Add update_timestamp_tz property to StorageSession

This property is a compatibility alias that returns the update timestamp as a POSIX timestamp. It infers whether the database is SQLite using sqlalchemy.inspect to call get_update_timestamp correctly

Close #4334

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 864595914
This commit is contained in:
George Weale
2026-02-02 17:39:42 -08:00
committed by Copybara-Service
parent dd8cd27b2c
commit 666cebe369
3 changed files with 66 additions and 1 deletions
+15
View File
@@ -38,6 +38,7 @@ from google.genai import types
from sqlalchemy import Boolean
from sqlalchemy import ForeignKeyConstraint
from sqlalchemy import func
from sqlalchemy import inspect
from sqlalchemy import Text
from sqlalchemy.dialects import mysql
from sqlalchemy.ext.mutable import MutableDict
@@ -130,6 +131,20 @@ class StorageSession(Base):
def __repr__(self):
return f"<StorageSession(id={self.id}, update_time={self.update_time})>"
@property
def update_timestamp_tz(self) -> float:
"""Returns the update timestamp as a POSIX timestamp.
This is a compatibility alias for callers that used the pre-`main` API.
"""
sqlalchemy_session = inspect(self).session
is_sqlite = bool(
sqlalchemy_session
and sqlalchemy_session.bind
and sqlalchemy_session.bind.dialect.name == "sqlite"
)
return self.get_update_timestamp(is_sqlite=is_sqlite)
def get_update_timestamp(self, is_sqlite: bool) -> float:
"""Returns the time zone aware update timestamp."""
if is_sqlite:
+15 -1
View File
@@ -26,11 +26,11 @@ from __future__ import annotations
from datetime import datetime
from datetime import timezone
from typing import Any
from typing import Optional
import uuid
from sqlalchemy import ForeignKeyConstraint
from sqlalchemy import func
from sqlalchemy import inspect
from sqlalchemy.ext.mutable import MutableDict
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped
@@ -105,6 +105,20 @@ class StorageSession(Base):
def __repr__(self):
return f"<StorageSession(id={self.id}, update_time={self.update_time})>"
@property
def update_timestamp_tz(self) -> float:
"""Returns the update timestamp as a POSIX timestamp.
This is a compatibility alias for callers that used the pre-`main` API.
"""
sqlalchemy_session = inspect(self).session
is_sqlite = bool(
sqlalchemy_session
and sqlalchemy_session.bind
and sqlalchemy_session.bind.dialect.name == "sqlite"
)
return self.get_update_timestamp(is_sqlite=is_sqlite)
def get_update_timestamp(self, is_sqlite: bool) -> float:
"""Returns the time zone aware update timestamp."""
if is_sqlite:
@@ -245,3 +245,39 @@ def test_migrate_from_sqlalchemy_pickle_with_async_driver_urls(tmp_path):
assert session_res.id == "async_session"
dest_session.close()
def _assert_update_timestamp_tz_is_utc_timestamp(schema_module) -> None:
engine = create_engine("sqlite:///:memory:")
schema_module.Base.metadata.create_all(engine)
SessionLocal = sessionmaker(bind=engine)
update_time = datetime(2026, 1, 1, 0, 0, 0)
storage_session = schema_module.StorageSession(
app_name="app",
user_id="user",
id="sid",
state={},
create_time=update_time,
update_time=update_time,
)
with SessionLocal() as db:
db.add(storage_session)
db.commit()
fetched = db.get(schema_module.StorageSession, ("app", "user", "sid"))
assert fetched is not None
assert isinstance(fetched.update_timestamp_tz, float)
assert (
fetched.update_timestamp_tz
== update_time.replace(tzinfo=timezone.utc).timestamp()
)
def test_v1_storage_session_update_timestamp_tz() -> None:
_assert_update_timestamp_tz_is_utc_timestamp(v1)
def test_v0_storage_session_update_timestamp_tz() -> None:
_assert_update_timestamp_tz_is_utc_timestamp(v0)