mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
dd8cd27b2c
commit
666cebe369
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user