sqlspec 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/adapters/aiosqlite/adk/store.py

@@ -0,0 +1,536 @@
+"""Aiosqlite async ADK store for Google Agent Development Kit session/event storage."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.serializers import from_json, to_json
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.aiosqlite.config import AiosqliteConfig
+
+logger = get_logger("adapters.aiosqlite.adk.store")
+
+SECONDS_PER_DAY = 86400.0
+JULIAN_EPOCH = 2440587.5
+
+__all__ = ("AiosqliteADKStore",)
+
+
+def _datetime_to_julian(dt: datetime) -> float:
+    """Convert datetime to Julian Day number for SQLite storage.
+
+    Args:
+        dt: Datetime to convert (must be UTC-aware).
+
+    Returns:
+        Julian Day number as REAL.
+
+    Notes:
+        Julian Day number is days since November 24, 4714 BCE (proleptic Gregorian).
+        This enables direct comparison with julianday('now') in SQL queries.
+    """
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    epoch = datetime(1970, 1, 1, tzinfo=timezone.utc)
+    delta_days = (dt - epoch).total_seconds() / SECONDS_PER_DAY
+    return JULIAN_EPOCH + delta_days
+
+
+def _julian_to_datetime(julian: float) -> datetime:
+    """Convert Julian Day number back to datetime.
+
+    Args:
+        julian: Julian Day number.
+
+    Returns:
+        UTC-aware datetime.
+    """
+    days_since_epoch = julian - JULIAN_EPOCH
+    timestamp = days_since_epoch * SECONDS_PER_DAY
+    return datetime.fromtimestamp(timestamp, tz=timezone.utc)
+
+
+def _to_sqlite_bool(value: "bool | None") -> "int | None":
+    """Convert Python bool to SQLite INTEGER.
+
+    Args:
+        value: Boolean value or None.
+
+    Returns:
+        1 for True, 0 for False, None for None.
+    """
+    if value is None:
+        return None
+    return 1 if value else 0
+
+
+def _from_sqlite_bool(value: "int | None") -> "bool | None":
+    """Convert SQLite INTEGER to Python bool.
+
+    Args:
+        value: Integer value (0/1) or None.
+
+    Returns:
+        True for 1, False for 0, None for None.
+    """
+    if value is None:
+        return None
+    return bool(value)
+
+
+class AiosqliteADKStore(BaseAsyncADKStore["AiosqliteConfig"]):
+    """Aiosqlite ADK store using asynchronous SQLite driver.
+
+    Implements session and event storage for Google Agent Development Kit
+    using SQLite via the asynchronous aiosqlite driver.
+
+    Provides:
+    - Session state management with JSON storage (as TEXT)
+    - Event history tracking with BLOB-serialized actions
+    - Julian Day timestamps (REAL) for efficient date operations
+    - Foreign key constraints with cascade delete
+    - Efficient upserts using INSERT OR REPLACE
+
+    Args:
+        config: AiosqliteConfig with extension_config["adk"] settings.
+
+    Example:
+        from sqlspec.adapters.aiosqlite import AiosqliteConfig
+        from sqlspec.adapters.aiosqlite.adk import AiosqliteADKStore
+
+        config = AiosqliteConfig(
+            pool_config={"database": ":memory:"},
+            extension_config={
+                "adk": {
+                    "session_table": "my_sessions",
+                    "events_table": "my_events"
+                }
+            }
+        )
+        store = AiosqliteADKStore(config)
+        await store.create_tables()
+
+    Notes:
+        - JSON stored as TEXT with SQLSpec serializers (msgspec/orjson/stdlib)
+        - BOOLEAN as INTEGER (0/1, with None for NULL)
+        - Timestamps as REAL (Julian day: julianday('now'))
+        - BLOB for pre-serialized actions from Google ADK
+        - PRAGMA foreign_keys = ON (enable per connection)
+        - Configuration is read from config.extension_config["adk"]
+    """
+
+    __slots__ = ()
+
+    def __init__(self, config: "AiosqliteConfig") -> None:
+        """Initialize Aiosqlite ADK store.
+
+        Args:
+            config: AiosqliteConfig instance.
+
+        Notes:
+            Configuration is read from config.extension_config["adk"]:
+            - session_table: Sessions table name (default: "adk_sessions")
+            - events_table: Events table name (default: "adk_events")
+        """
+        super().__init__(config)
+
+    async def _get_create_sessions_table_sql(self) -> str:
+        """Get SQLite CREATE TABLE SQL for sessions.
+
+        Returns:
+            SQL statement to create adk_sessions table with indexes.
+
+        Notes:
+            - TEXT for IDs, names, and JSON state
+            - REAL for Julian Day timestamps
+            - Composite index on (app_name, user_id)
+            - Index on update_time DESC for recent session queries
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._session_table} (
+            id TEXT PRIMARY KEY,
+            app_name TEXT NOT NULL,
+            user_id TEXT NOT NULL,
+            state TEXT NOT NULL DEFAULT '{{}}',
+            create_time REAL NOT NULL,
+            update_time REAL NOT NULL
+        );
+        CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user
+            ON {self._session_table}(app_name, user_id);
+        CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time
+            ON {self._session_table}(update_time DESC);
+        """
+
+    async def _get_create_events_table_sql(self) -> str:
+        """Get SQLite CREATE TABLE SQL for events.
+
+        Returns:
+            SQL statement to create adk_events table with indexes.
+
+        Notes:
+            - TEXT for IDs, strings, and JSON content
+            - BLOB for pickled actions
+            - INTEGER for booleans (0/1/NULL)
+            - REAL for Julian Day timestamps
+            - Foreign key to sessions with CASCADE delete
+            - Index on (session_id, timestamp ASC)
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._events_table} (
+            id TEXT PRIMARY KEY,
+            session_id TEXT NOT NULL,
+            app_name TEXT NOT NULL,
+            user_id TEXT NOT NULL,
+            invocation_id TEXT NOT NULL,
+            author TEXT NOT NULL,
+            actions BLOB NOT NULL,
+            long_running_tool_ids_json TEXT,
+            branch TEXT,
+            timestamp REAL NOT NULL,
+            content TEXT,
+            grounding_metadata TEXT,
+            custom_metadata TEXT,
+            partial INTEGER,
+            turn_complete INTEGER,
+            interrupted INTEGER,
+            error_code TEXT,
+            error_message TEXT,
+            FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
+        );
+        CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session
+            ON {self._events_table}(session_id, timestamp ASC);
+        """
+
+    def _get_drop_tables_sql(self) -> "list[str]":
+        """Get SQLite DROP TABLE SQL statements.
+
+        Returns:
+            List of SQL statements to drop tables and indexes.
+
+        Notes:
+            Order matters: drop events table (child) before sessions (parent).
+            SQLite automatically drops indexes when dropping tables.
+        """
+        return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
+
+    async def _enable_foreign_keys(self, connection: Any) -> None:
+        """Enable foreign key constraints for this connection.
+
+        Args:
+            connection: Aiosqlite connection.
+
+        Notes:
+            SQLite requires PRAGMA foreign_keys = ON per connection.
+        """
+        await connection.execute("PRAGMA foreign_keys = ON")
+
+    async def create_tables(self) -> None:
+        """Create both sessions and events tables if they don't exist."""
+        async with self._config.provide_session() as driver:
+            await self._enable_foreign_keys(driver.connection)
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
+        logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
+
+    async def create_session(
+        self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
+    ) -> SessionRecord:
+        """Create a new session.
+
+        Args:
+            session_id: Unique session identifier.
+            app_name: Application name.
+            user_id: User identifier.
+            state: Initial session state.
+            owner_id: Optional owner ID value for owner_id_column.
+
+        Returns:
+            Created session record.
+
+        Notes:
+            Uses Julian Day for create_time and update_time.
+            State is JSON-serialized before insertion.
+        """
+        now = datetime.now(timezone.utc)
+        now_julian = _datetime_to_julian(now)
+        state_json = to_json(state) if state else None
+
+        params: tuple[Any, ...]
+        if self._owner_id_column_name:
+            sql = f"""
+            INSERT INTO {self._session_table}
+            (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
+            VALUES (?, ?, ?, ?, ?, ?, ?)
+            """
+            params = (session_id, app_name, user_id, owner_id, state_json, now_julian, now_julian)
+        else:
+            sql = f"""
+            INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time)
+            VALUES (?, ?, ?, ?, ?, ?)
+            """
+            params = (session_id, app_name, user_id, state_json, now_julian, now_julian)
+
+        async with self._config.provide_connection() as conn:
+            await self._enable_foreign_keys(conn)
+            await conn.execute(sql, params)
+            await conn.commit()
+
+        return SessionRecord(
+            id=session_id, app_name=app_name, user_id=user_id, state=state, create_time=now, update_time=now
+        )
+
+    async def get_session(self, session_id: str) -> "SessionRecord | None":
+        """Get session by ID.
+
+        Args:
+            session_id: Session identifier.
+
+        Returns:
+            Session record or None if not found.
+
+        Notes:
+            SQLite returns Julian Day (REAL) for timestamps.
+            JSON is parsed from TEXT storage.
+        """
+        sql = f"""
+        SELECT id, app_name, user_id, state, create_time, update_time
+        FROM {self._session_table}
+        WHERE id = ?
+        """
+
+        async with self._config.provide_connection() as conn:
+            await self._enable_foreign_keys(conn)
+            cursor = await conn.execute(sql, (session_id,))
+            row = await cursor.fetchone()
+
+            if row is None:
+                return None
+
+            return SessionRecord(
+                id=row[0],
+                app_name=row[1],
+                user_id=row[2],
+                state=from_json(row[3]) if row[3] else {},
+                create_time=_julian_to_datetime(row[4]),
+                update_time=_julian_to_datetime(row[5]),
+            )
+
+    async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
+        """Update session state.
+
+        Args:
+            session_id: Session identifier.
+            state: New state dictionary (replaces existing state).
+
+        Notes:
+            This replaces the entire state dictionary.
+            Updates update_time to current Julian Day.
+        """
+        now_julian = _datetime_to_julian(datetime.now(timezone.utc))
+        state_json = to_json(state) if state else None
+
+        sql = f"""
+        UPDATE {self._session_table}
+        SET state = ?, update_time = ?
+        WHERE id = ?
+        """
+
+        async with self._config.provide_connection() as conn:
+            await self._enable_foreign_keys(conn)
+            await conn.execute(sql, (state_json, now_julian, session_id))
+            await conn.commit()
+
+    async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
+        """List sessions for an app, optionally filtered by user.
+
+        Args:
+            app_name: Application name.
+            user_id: User identifier. If None, lists all sessions for the app.
+
+        Returns:
+            List of session records ordered by update_time DESC.
+
+        Notes:
+            Uses composite index on (app_name, user_id) when user_id is provided.
+        """
+        if user_id is None:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = ?
+            ORDER BY update_time DESC
+            """
+            params: tuple[str, ...] = (app_name,)
+        else:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = ? AND user_id = ?
+            ORDER BY update_time DESC
+            """
+            params = (app_name, user_id)
+
+        async with self._config.provide_connection() as conn:
+            await self._enable_foreign_keys(conn)
+            cursor = await conn.execute(sql, params)
+            rows = await cursor.fetchall()
+
+            return [
+                SessionRecord(
+                    id=row[0],
+                    app_name=row[1],
+                    user_id=row[2],
+                    state=from_json(row[3]) if row[3] else {},
+                    create_time=_julian_to_datetime(row[4]),
+                    update_time=_julian_to_datetime(row[5]),
+                )
+                for row in rows
+            ]
+
+    async def delete_session(self, session_id: str) -> None:
+        """Delete session and all associated events (cascade).
+
+        Args:
+            session_id: Session identifier.
+
+        Notes:
+            Foreign key constraint ensures events are cascade-deleted.
+        """
+        sql = f"DELETE FROM {self._session_table} WHERE id = ?"
+
+        async with self._config.provide_connection() as conn:
+            await self._enable_foreign_keys(conn)
+            await conn.execute(sql, (session_id,))
+            await conn.commit()
+
+    async def append_event(self, event_record: EventRecord) -> None:
+        """Append an event to a session.
+
+        Args:
+            event_record: Event record to store.
+
+        Notes:
+            Uses Julian Day for timestamp.
+            JSON fields are serialized to TEXT.
+            Boolean fields converted to INTEGER (0/1/NULL).
+        """
+        timestamp_julian = _datetime_to_julian(event_record["timestamp"])
+
+        content_json = to_json(event_record.get("content")) if event_record.get("content") else None
+        grounding_metadata_json = (
+            to_json(event_record.get("grounding_metadata")) if event_record.get("grounding_metadata") else None
+        )
+        custom_metadata_json = (
+            to_json(event_record.get("custom_metadata")) if event_record.get("custom_metadata") else None
+        )
+
+        partial_int = _to_sqlite_bool(event_record.get("partial"))
+        turn_complete_int = _to_sqlite_bool(event_record.get("turn_complete"))
+        interrupted_int = _to_sqlite_bool(event_record.get("interrupted"))
+
+        sql = f"""
+        INSERT INTO {self._events_table} (
+            id, session_id, app_name, user_id, invocation_id, author, actions,
+            long_running_tool_ids_json, branch, timestamp, content,
+            grounding_metadata, custom_metadata, partial, turn_complete,
+            interrupted, error_code, error_message
+        ) VALUES (
+            ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
+        )
+        """
+
+        async with self._config.provide_connection() as conn:
+            await self._enable_foreign_keys(conn)
+            await conn.execute(
+                sql,
+                (
+                    event_record["id"],
+                    event_record["session_id"],
+                    event_record["app_name"],
+                    event_record["user_id"],
+                    event_record["invocation_id"],
+                    event_record["author"],
+                    event_record["actions"],
+                    event_record.get("long_running_tool_ids_json"),
+                    event_record.get("branch"),
+                    timestamp_julian,
+                    content_json,
+                    grounding_metadata_json,
+                    custom_metadata_json,
+                    partial_int,
+                    turn_complete_int,
+                    interrupted_int,
+                    event_record.get("error_code"),
+                    event_record.get("error_message"),
+                ),
+            )
+            await conn.commit()
+
+    async def get_events(
+        self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
+    ) -> "list[EventRecord]":
+        """Get events for a session.
+
+        Args:
+            session_id: Session identifier.
+            after_timestamp: Only return events after this time.
+            limit: Maximum number of events to return.
+
+        Returns:
+            List of event records ordered by timestamp ASC.
+
+        Notes:
+            Uses index on (session_id, timestamp ASC).
+            Parses JSON fields and converts BLOB actions to bytes.
+            Converts INTEGER booleans back to bool/None.
+        """
+        where_clauses = ["session_id = ?"]
+        params: list[Any] = [session_id]
+
+        if after_timestamp is not None:
+            where_clauses.append("timestamp > ?")
+            params.append(_datetime_to_julian(after_timestamp))
+
+        where_clause = " AND ".join(where_clauses)
+        limit_clause = f" LIMIT {limit}" if limit else ""
+
+        sql = f"""
+        SELECT id, session_id, app_name, user_id, invocation_id, author, actions,
+               long_running_tool_ids_json, branch, timestamp, content,
+               grounding_metadata, custom_metadata, partial, turn_complete,
+               interrupted, error_code, error_message
+        FROM {self._events_table}
+        WHERE {where_clause}
+        ORDER BY timestamp ASC{limit_clause}
+        """
+
+        async with self._config.provide_connection() as conn:
+            await self._enable_foreign_keys(conn)
+            cursor = await conn.execute(sql, params)
+            rows = await cursor.fetchall()
+
+            return [
+                EventRecord(
+                    id=row[0],
+                    session_id=row[1],
+                    app_name=row[2],
+                    user_id=row[3],
+                    invocation_id=row[4],
+                    author=row[5],
+                    actions=bytes(row[6]),
+                    long_running_tool_ids_json=row[7],
+                    branch=row[8],
+                    timestamp=_julian_to_datetime(row[9]),
+                    content=from_json(row[10]) if row[10] else None,
+                    grounding_metadata=from_json(row[11]) if row[11] else None,
+                    custom_metadata=from_json(row[12]) if row[12] else None,
+                    partial=_from_sqlite_bool(row[13]),
+                    turn_complete=_from_sqlite_bool(row[14]),
+                    interrupted=_from_sqlite_bool(row[15]),
+                    error_code=row[16],
+                    error_message=row[17],
+                )
+                for row in rows
+            ]