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

sqlspec/adapters/psycopg/adk/store.py

@@ -0,0 +1,944 @@
+"""Psycopg ADK store for Google Agent Development Kit session/event storage."""
+
+from typing import TYPE_CHECKING, Any, cast
+
+from psycopg import errors
+from psycopg import sql as pg_sql
+from psycopg.types.json import Jsonb
+
+from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, EventRecord, SessionRecord
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from datetime import datetime
+
+    from psycopg.abc import Query
+
+    from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig
+
+logger = get_logger("adapters.psycopg.adk.store")
+
+__all__ = ("PsycopgAsyncADKStore", "PsycopgSyncADKStore")
+
+
+class PsycopgAsyncADKStore(BaseAsyncADKStore["PsycopgAsyncConfig"]):
+    """PostgreSQL ADK store using Psycopg3 driver.
+
+    Implements session and event storage for Google Agent Development Kit
+    using PostgreSQL via psycopg3 with native async/await support.
+
+    Provides:
+    - Session state management with JSONB storage and merge operations
+    - Event history tracking with BYTEA-serialized actions
+    - Microsecond-precision timestamps with TIMESTAMPTZ
+    - Foreign key constraints with cascade delete
+    - Efficient upserts using ON CONFLICT
+    - GIN indexes for JSONB queries
+    - HOT updates with FILLFACTOR 80
+
+    Args:
+        config: PsycopgAsyncConfig with extension_config["adk"] settings.
+
+    Example:
+        from sqlspec.adapters.psycopg import PsycopgAsyncConfig
+        from sqlspec.adapters.psycopg.adk import PsycopgAsyncADKStore
+
+        config = PsycopgAsyncConfig(
+            pool_config={"conninfo": "postgresql://..."},
+            extension_config={
+                "adk": {
+                    "session_table": "my_sessions",
+                    "events_table": "my_events",
+                    "owner_id_column": "tenant_id INTEGER NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
+                }
+            }
+        )
+        store = PsycopgAsyncADKStore(config)
+        await store.create_tables()
+
+    Notes:
+        - PostgreSQL JSONB type used for state (more efficient than JSON)
+        - Psycopg requires wrapping dicts with Jsonb() for type safety
+        - TIMESTAMPTZ provides timezone-aware microsecond precision
+        - State merging uses `state || $1::jsonb` operator for efficiency
+        - BYTEA for pre-serialized actions from Google ADK
+        - GIN index on state for JSONB queries (partial index)
+        - FILLFACTOR 80 leaves space for HOT updates
+        - Parameter style: $1, $2, $3 (PostgreSQL numeric placeholders)
+        - Configuration is read from config.extension_config["adk"]
+    """
+
+    __slots__ = ()
+
+    def __init__(self, config: "PsycopgAsyncConfig") -> None:
+        """Initialize Psycopg ADK store.
+
+        Args:
+            config: PsycopgAsyncConfig 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")
+            - owner_id_column: Optional owner FK column DDL (default: None)
+        """
+        super().__init__(config)
+
+    def _get_create_sessions_table_sql(self) -> str:
+        """Get PostgreSQL CREATE TABLE SQL for sessions.
+
+        Returns:
+            SQL statement to create adk_sessions table with indexes.
+
+        Notes:
+            - VARCHAR(128) for IDs and names (sufficient for UUIDs and app names)
+            - JSONB type for state storage with default empty object
+            - TIMESTAMPTZ with microsecond precision
+            - FILLFACTOR 80 for HOT updates (reduces table bloat)
+            - Composite index on (app_name, user_id) for listing
+            - Index on update_time DESC for recent session queries
+            - Partial GIN index on state for JSONB queries (only non-empty)
+            - Optional owner ID column for multi-tenancy or user references
+        """
+        owner_id_line = ""
+        if self._owner_id_column_ddl:
+            owner_id_line = f",\n            {self._owner_id_column_ddl}"
+
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._session_table} (
+            id VARCHAR(128) PRIMARY KEY,
+            app_name VARCHAR(128) NOT NULL,
+            user_id VARCHAR(128) NOT NULL{owner_id_line},
+            state JSONB NOT NULL DEFAULT '{{}}'::jsonb,
+            create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+        ) WITH (fillfactor = 80);
+
+        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);
+
+        CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state
+            ON {self._session_table} USING GIN (state)
+            WHERE state != '{{}}'::jsonb;
+        """
+
+    def _get_create_events_table_sql(self) -> str:
+        """Get PostgreSQL CREATE TABLE SQL for events.
+
+        Returns:
+            SQL statement to create adk_events table with indexes.
+
+        Notes:
+            - VARCHAR sizes: id(128), session_id(128), invocation_id(256), author(256),
+              branch(256), error_code(256), error_message(1024)
+            - BYTEA for pickled actions (no size limit)
+            - JSONB for content, grounding_metadata, custom_metadata, long_running_tool_ids_json
+            - BOOLEAN for partial, turn_complete, interrupted
+            - Foreign key to sessions with CASCADE delete
+            - Index on (session_id, timestamp ASC) for ordered event retrieval
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._events_table} (
+            id VARCHAR(128) PRIMARY KEY,
+            session_id VARCHAR(128) NOT NULL,
+            app_name VARCHAR(128) NOT NULL,
+            user_id VARCHAR(128) NOT NULL,
+            invocation_id VARCHAR(256),
+            author VARCHAR(256),
+            actions BYTEA,
+            long_running_tool_ids_json JSONB,
+            branch VARCHAR(256),
+            timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            content JSONB,
+            grounding_metadata JSONB,
+            custom_metadata JSONB,
+            partial BOOLEAN,
+            turn_complete BOOLEAN,
+            interrupted BOOLEAN,
+            error_code VARCHAR(256),
+            error_message VARCHAR(1024),
+            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 PostgreSQL DROP TABLE SQL statements.
+
+        Returns:
+            List of SQL statements to drop tables and indexes.
+
+        Notes:
+            Order matters: drop events table (child) before sessions (parent).
+            PostgreSQL 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 create_tables(self) -> None:
+        """Create both sessions and events tables if they don't exist."""
+        async with self._config.provide_connection() as conn, conn.cursor() as cur:
+            await cur.execute(cast("Query", self._get_create_sessions_table_sql()))
+            await cur.execute(cast("Query", 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 (if configured).
+
+        Returns:
+            Created session record.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP for create_time and update_time.
+            State is wrapped with Jsonb() for PostgreSQL type safety.
+            If owner_id_column is configured, owner_id value must be provided.
+        """
+        params: tuple[Any, ...]
+        if self._owner_id_column_name:
+            query = pg_sql.SQL("""
+            INSERT INTO {table} (id, app_name, user_id, {owner_id_col}, state, create_time, update_time)
+            VALUES (%s, %s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+            """).format(
+                table=pg_sql.Identifier(self._session_table), owner_id_col=pg_sql.Identifier(self._owner_id_column_name)
+            )
+            params = (session_id, app_name, user_id, owner_id, Jsonb(state))
+        else:
+            query = pg_sql.SQL("""
+            INSERT INTO {table} (id, app_name, user_id, state, create_time, update_time)
+            VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+            """).format(table=pg_sql.Identifier(self._session_table))
+            params = (session_id, app_name, user_id, Jsonb(state))
+
+        async with self._config.provide_connection() as conn, conn.cursor() as cur:
+            await cur.execute(query, params)
+
+        return await self.get_session(session_id)  # type: ignore[return-value]
+
+    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:
+            PostgreSQL returns datetime objects for TIMESTAMPTZ columns.
+            JSONB is automatically deserialized by psycopg to Python dict.
+        """
+        query = pg_sql.SQL("""
+        SELECT id, app_name, user_id, state, create_time, update_time
+        FROM {table}
+        WHERE id = %s
+        """).format(table=pg_sql.Identifier(self._session_table))
+
+        try:
+            async with self._config.provide_connection() as conn, conn.cursor() as cur:
+                await cur.execute(query, (session_id,))
+                row = await cur.fetchone()
+
+                if row is None:
+                    return None
+
+                return SessionRecord(
+                    id=row["id"],
+                    app_name=row["app_name"],
+                    user_id=row["user_id"],
+                    state=row["state"],
+                    create_time=row["create_time"],
+                    update_time=row["update_time"],
+                )
+        except errors.UndefinedTable:
+            return None
+
+    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.
+            Uses CURRENT_TIMESTAMP for update_time.
+            State is wrapped with Jsonb() for PostgreSQL type safety.
+        """
+        query = pg_sql.SQL("""
+        UPDATE {table}
+        SET state = %s, update_time = CURRENT_TIMESTAMP
+        WHERE id = %s
+        """).format(table=pg_sql.Identifier(self._session_table))
+
+        async with self._config.provide_connection() as conn, conn.cursor() as cur:
+            await cur.execute(query, (Jsonb(state), session_id))
+
+    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.
+        """
+        query = pg_sql.SQL("DELETE FROM {table} WHERE id = %s").format(table=pg_sql.Identifier(self._session_table))
+
+        async with self._config.provide_connection() as conn, conn.cursor() as cur:
+            await cur.execute(query, (session_id,))
+
+    async def list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
+        """List all sessions for a user in an app.
+
+        Args:
+            app_name: Application name.
+            user_id: User identifier.
+
+        Returns:
+            List of session records ordered by update_time DESC.
+
+        Notes:
+            Uses composite index on (app_name, user_id).
+        """
+        query = pg_sql.SQL("""
+        SELECT id, app_name, user_id, state, create_time, update_time
+        FROM {table}
+        WHERE app_name = %s AND user_id = %s
+        ORDER BY update_time DESC
+        """).format(table=pg_sql.Identifier(self._session_table))
+
+        try:
+            async with self._config.provide_connection() as conn, conn.cursor() as cur:
+                await cur.execute(query, (app_name, user_id))
+                rows = await cur.fetchall()
+
+                return [
+                    SessionRecord(
+                        id=row["id"],
+                        app_name=row["app_name"],
+                        user_id=row["user_id"],
+                        state=row["state"],
+                        create_time=row["create_time"],
+                        update_time=row["update_time"],
+                    )
+                    for row in rows
+                ]
+        except errors.UndefinedTable:
+            return []
+
+    async def append_event(self, event_record: EventRecord) -> None:
+        """Append an event to a session.
+
+        Args:
+            event_record: Event record to store.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP for timestamp if not provided.
+            JSONB fields are wrapped with Jsonb() for PostgreSQL type safety.
+        """
+        content_json = event_record.get("content")
+        grounding_metadata_json = event_record.get("grounding_metadata")
+        custom_metadata_json = event_record.get("custom_metadata")
+
+        query = pg_sql.SQL("""
+        INSERT INTO {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 (
+            %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
+        )
+        """).format(table=pg_sql.Identifier(self._events_table))
+
+        async with self._config.provide_connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                query,
+                (
+                    event_record["id"],
+                    event_record["session_id"],
+                    event_record["app_name"],
+                    event_record["user_id"],
+                    event_record.get("invocation_id"),
+                    event_record.get("author"),
+                    event_record.get("actions"),
+                    event_record.get("long_running_tool_ids_json"),
+                    event_record.get("branch"),
+                    event_record["timestamp"],
+                    Jsonb(content_json) if content_json is not None else None,
+                    Jsonb(grounding_metadata_json) if grounding_metadata_json is not None else None,
+                    Jsonb(custom_metadata_json) if custom_metadata_json is not None else None,
+                    event_record.get("partial"),
+                    event_record.get("turn_complete"),
+                    event_record.get("interrupted"),
+                    event_record.get("error_code"),
+                    event_record.get("error_message"),
+                ),
+            )
+
+    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).
+            JSONB fields are automatically deserialized by psycopg.
+            BYTEA actions are converted to bytes.
+        """
+        where_clauses = ["session_id = %s"]
+        params: list[Any] = [session_id]
+
+        if after_timestamp is not None:
+            where_clauses.append("timestamp > %s")
+            params.append(after_timestamp)
+
+        where_clause = " AND ".join(where_clauses)
+        if limit:
+            params.append(limit)
+
+        query = pg_sql.SQL(
+            """
+        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 {table}
+        WHERE {where_clause}
+        ORDER BY timestamp ASC{limit_clause}
+        """
+        ).format(
+            table=pg_sql.Identifier(self._events_table),
+            where_clause=pg_sql.SQL(where_clause),  # pyright: ignore[reportArgumentType]
+            limit_clause=pg_sql.SQL(" LIMIT %s" if limit else ""),  # pyright: ignore[reportArgumentType]
+        )
+
+        try:
+            async with self._config.provide_connection() as conn, conn.cursor() as cur:
+                await cur.execute(query, tuple(params))
+                rows = await cur.fetchall()
+
+                return [
+                    EventRecord(
+                        id=row["id"],
+                        session_id=row["session_id"],
+                        app_name=row["app_name"],
+                        user_id=row["user_id"],
+                        invocation_id=row["invocation_id"],
+                        author=row["author"],
+                        actions=bytes(row["actions"]) if row["actions"] else b"",
+                        long_running_tool_ids_json=row["long_running_tool_ids_json"],
+                        branch=row["branch"],
+                        timestamp=row["timestamp"],
+                        content=row["content"],
+                        grounding_metadata=row["grounding_metadata"],
+                        custom_metadata=row["custom_metadata"],
+                        partial=row["partial"],
+                        turn_complete=row["turn_complete"],
+                        interrupted=row["interrupted"],
+                        error_code=row["error_code"],
+                        error_message=row["error_message"],
+                    )
+                    for row in rows
+                ]
+        except errors.UndefinedTable:
+            return []
+
+
+class PsycopgSyncADKStore(BaseSyncADKStore["PsycopgSyncConfig"]):
+    """PostgreSQL synchronous ADK store using Psycopg3 driver.
+
+    Implements session and event storage for Google Agent Development Kit
+    using PostgreSQL via psycopg3 with synchronous execution.
+
+    Provides:
+    - Session state management with JSONB storage and merge operations
+    - Event history tracking with BYTEA-serialized actions
+    - Microsecond-precision timestamps with TIMESTAMPTZ
+    - Foreign key constraints with cascade delete
+    - Efficient upserts using ON CONFLICT
+    - GIN indexes for JSONB queries
+    - HOT updates with FILLFACTOR 80
+
+    Args:
+        config: PsycopgSyncConfig with extension_config["adk"] settings.
+
+    Example:
+        from sqlspec.adapters.psycopg import PsycopgSyncConfig
+        from sqlspec.adapters.psycopg.adk import PsycopgSyncADKStore
+
+        config = PsycopgSyncConfig(
+            pool_config={"conninfo": "postgresql://..."},
+            extension_config={
+                "adk": {
+                    "session_table": "my_sessions",
+                    "events_table": "my_events",
+                    "owner_id_column": "tenant_id INTEGER NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
+                }
+            }
+        )
+        store = PsycopgSyncADKStore(config)
+        store.create_tables()
+
+    Notes:
+        - PostgreSQL JSONB type used for state (more efficient than JSON)
+        - Psycopg requires wrapping dicts with Jsonb() for type safety
+        - TIMESTAMPTZ provides timezone-aware microsecond precision
+        - State merging uses `state || $1::jsonb` operator for efficiency
+        - BYTEA for pre-serialized actions from Google ADK
+        - GIN index on state for JSONB queries (partial index)
+        - FILLFACTOR 80 leaves space for HOT updates
+        - Parameter style: $1, $2, $3 (PostgreSQL numeric placeholders)
+        - Configuration is read from config.extension_config["adk"]
+    """
+
+    __slots__ = ()
+
+    def __init__(self, config: "PsycopgSyncConfig") -> None:
+        """Initialize Psycopg synchronous ADK store.
+
+        Args:
+            config: PsycopgSyncConfig 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")
+            - owner_id_column: Optional owner FK column DDL (default: None)
+        """
+        super().__init__(config)
+
+    def _get_create_sessions_table_sql(self) -> str:
+        """Get PostgreSQL CREATE TABLE SQL for sessions.
+
+        Returns:
+            SQL statement to create adk_sessions table with indexes.
+
+        Notes:
+            - VARCHAR(128) for IDs and names (sufficient for UUIDs and app names)
+            - JSONB type for state storage with default empty object
+            - TIMESTAMPTZ with microsecond precision
+            - FILLFACTOR 80 for HOT updates (reduces table bloat)
+            - Composite index on (app_name, user_id) for listing
+            - Index on update_time DESC for recent session queries
+            - Partial GIN index on state for JSONB queries (only non-empty)
+            - Optional owner ID column for multi-tenancy or user references
+        """
+        owner_id_line = ""
+        if self._owner_id_column_ddl:
+            owner_id_line = f",\n            {self._owner_id_column_ddl}"
+
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._session_table} (
+            id VARCHAR(128) PRIMARY KEY,
+            app_name VARCHAR(128) NOT NULL,
+            user_id VARCHAR(128) NOT NULL{owner_id_line},
+            state JSONB NOT NULL DEFAULT '{{}}'::jsonb,
+            create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+        ) WITH (fillfactor = 80);
+
+        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);
+
+        CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state
+            ON {self._session_table} USING GIN (state)
+            WHERE state != '{{}}'::jsonb;
+        """
+
+    def _get_create_events_table_sql(self) -> str:
+        """Get PostgreSQL CREATE TABLE SQL for events.
+
+        Returns:
+            SQL statement to create adk_events table with indexes.
+
+        Notes:
+            - VARCHAR sizes: id(128), session_id(128), invocation_id(256), author(256),
+              branch(256), error_code(256), error_message(1024)
+            - BYTEA for pickled actions (no size limit)
+            - JSONB for content, grounding_metadata, custom_metadata, long_running_tool_ids_json
+            - BOOLEAN for partial, turn_complete, interrupted
+            - Foreign key to sessions with CASCADE delete
+            - Index on (session_id, timestamp ASC) for ordered event retrieval
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._events_table} (
+            id VARCHAR(128) PRIMARY KEY,
+            session_id VARCHAR(128) NOT NULL,
+            app_name VARCHAR(128) NOT NULL,
+            user_id VARCHAR(128) NOT NULL,
+            invocation_id VARCHAR(256),
+            author VARCHAR(256),
+            actions BYTEA,
+            long_running_tool_ids_json JSONB,
+            branch VARCHAR(256),
+            timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            content JSONB,
+            grounding_metadata JSONB,
+            custom_metadata JSONB,
+            partial BOOLEAN,
+            turn_complete BOOLEAN,
+            interrupted BOOLEAN,
+            error_code VARCHAR(256),
+            error_message VARCHAR(1024),
+            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 PostgreSQL DROP TABLE SQL statements.
+
+        Returns:
+            List of SQL statements to drop tables and indexes.
+
+        Notes:
+            Order matters: drop events table (child) before sessions (parent).
+            PostgreSQL automatically drops indexes when dropping tables.
+        """
+        return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
+
+    def create_tables(self) -> None:
+        """Create both sessions and events tables if they don't exist."""
+        with self._config.provide_connection() as conn, conn.cursor() as cur:
+            cur.execute(cast("Query", self._get_create_sessions_table_sql()))
+            cur.execute(cast("Query", self._get_create_events_table_sql()))
+        logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
+
+    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 (if configured).
+
+        Returns:
+            Created session record.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP for create_time and update_time.
+            State is wrapped with Jsonb() for PostgreSQL type safety.
+            If owner_id_column is configured, owner_id value must be provided.
+        """
+        params: tuple[Any, ...]
+        if self._owner_id_column_name:
+            query = pg_sql.SQL("""
+            INSERT INTO {table} (id, app_name, user_id, {owner_id_col}, state, create_time, update_time)
+            VALUES (%s, %s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+            """).format(
+                table=pg_sql.Identifier(self._session_table), owner_id_col=pg_sql.Identifier(self._owner_id_column_name)
+            )
+            params = (session_id, app_name, user_id, owner_id, Jsonb(state))
+        else:
+            query = pg_sql.SQL("""
+            INSERT INTO {table} (id, app_name, user_id, state, create_time, update_time)
+            VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+            """).format(table=pg_sql.Identifier(self._session_table))
+            params = (session_id, app_name, user_id, Jsonb(state))
+
+        with self._config.provide_connection() as conn, conn.cursor() as cur:
+            cur.execute(query, params)
+
+        return self.get_session(session_id)  # type: ignore[return-value]
+
+    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:
+            PostgreSQL returns datetime objects for TIMESTAMPTZ columns.
+            JSONB is automatically deserialized by psycopg to Python dict.
+        """
+        query = pg_sql.SQL("""
+        SELECT id, app_name, user_id, state, create_time, update_time
+        FROM {table}
+        WHERE id = %s
+        """).format(table=pg_sql.Identifier(self._session_table))
+
+        try:
+            with self._config.provide_connection() as conn, conn.cursor() as cur:
+                cur.execute(query, (session_id,))
+                row = cur.fetchone()
+
+                if row is None:
+                    return None
+
+                return SessionRecord(
+                    id=row["id"],
+                    app_name=row["app_name"],
+                    user_id=row["user_id"],
+                    state=row["state"],
+                    create_time=row["create_time"],
+                    update_time=row["update_time"],
+                )
+        except errors.UndefinedTable:
+            return None
+
+    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.
+            Uses CURRENT_TIMESTAMP for update_time.
+            State is wrapped with Jsonb() for PostgreSQL type safety.
+        """
+        query = pg_sql.SQL("""
+        UPDATE {table}
+        SET state = %s, update_time = CURRENT_TIMESTAMP
+        WHERE id = %s
+        """).format(table=pg_sql.Identifier(self._session_table))
+
+        with self._config.provide_connection() as conn, conn.cursor() as cur:
+            cur.execute(query, (Jsonb(state), session_id))
+
+    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.
+        """
+        query = pg_sql.SQL("DELETE FROM {table} WHERE id = %s").format(table=pg_sql.Identifier(self._session_table))
+
+        with self._config.provide_connection() as conn, conn.cursor() as cur:
+            cur.execute(query, (session_id,))
+
+    def list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
+        """List all sessions for a user in an app.
+
+        Args:
+            app_name: Application name.
+            user_id: User identifier.
+
+        Returns:
+            List of session records ordered by update_time DESC.
+
+        Notes:
+            Uses composite index on (app_name, user_id).
+        """
+        query = pg_sql.SQL("""
+        SELECT id, app_name, user_id, state, create_time, update_time
+        FROM {table}
+        WHERE app_name = %s AND user_id = %s
+        ORDER BY update_time DESC
+        """).format(table=pg_sql.Identifier(self._session_table))
+
+        try:
+            with self._config.provide_connection() as conn, conn.cursor() as cur:
+                cur.execute(query, (app_name, user_id))
+                rows = cur.fetchall()
+
+                return [
+                    SessionRecord(
+                        id=row["id"],
+                        app_name=row["app_name"],
+                        user_id=row["user_id"],
+                        state=row["state"],
+                        create_time=row["create_time"],
+                        update_time=row["update_time"],
+                    )
+                    for row in rows
+                ]
+        except errors.UndefinedTable:
+            return []
+
+    def create_event(
+        self,
+        event_id: str,
+        session_id: str,
+        app_name: str,
+        user_id: str,
+        author: "str | None" = None,
+        actions: "bytes | None" = None,
+        content: "dict[str, Any] | None" = None,
+        **kwargs: Any,
+    ) -> EventRecord:
+        """Create a new event.
+
+        Args:
+            event_id: Unique event identifier.
+            session_id: Session identifier.
+            app_name: Application name.
+            user_id: User identifier.
+            author: Event author (user/assistant/system).
+            actions: Pickled actions object.
+            content: Event content (JSONB).
+            **kwargs: Additional optional fields (invocation_id, branch, timestamp,
+                     grounding_metadata, custom_metadata, partial, turn_complete,
+                     interrupted, error_code, error_message, long_running_tool_ids_json).
+
+        Returns:
+            Created event record.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP for timestamp if not provided in kwargs.
+            JSONB fields are wrapped with Jsonb() for PostgreSQL type safety.
+        """
+        content_json = Jsonb(content) if content is not None else None
+        grounding_metadata = kwargs.get("grounding_metadata")
+        grounding_metadata_json = Jsonb(grounding_metadata) if grounding_metadata is not None else None
+        custom_metadata = kwargs.get("custom_metadata")
+        custom_metadata_json = Jsonb(custom_metadata) if custom_metadata is not None else None
+
+        query = pg_sql.SQL("""
+        INSERT INTO {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 (
+            %s, %s, %s, %s, %s, %s, %s, %s, %s, COALESCE(%s, CURRENT_TIMESTAMP), %s, %s, %s, %s, %s, %s, %s, %s
+        )
+        RETURNING 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
+        """).format(table=pg_sql.Identifier(self._events_table))
+
+        with self._config.provide_connection() as conn, conn.cursor() as cur:
+            cur.execute(
+                query,
+                (
+                    event_id,
+                    session_id,
+                    app_name,
+                    user_id,
+                    kwargs.get("invocation_id"),
+                    author,
+                    actions,
+                    kwargs.get("long_running_tool_ids_json"),
+                    kwargs.get("branch"),
+                    kwargs.get("timestamp"),
+                    content_json,
+                    grounding_metadata_json,
+                    custom_metadata_json,
+                    kwargs.get("partial"),
+                    kwargs.get("turn_complete"),
+                    kwargs.get("interrupted"),
+                    kwargs.get("error_code"),
+                    kwargs.get("error_message"),
+                ),
+            )
+            row = cur.fetchone()
+
+            if row is None:
+                msg = f"Failed to create event {event_id}"
+                raise RuntimeError(msg)
+
+            return EventRecord(
+                id=row["id"],
+                session_id=row["session_id"],
+                app_name=row["app_name"],
+                user_id=row["user_id"],
+                invocation_id=row["invocation_id"],
+                author=row["author"],
+                actions=bytes(row["actions"]) if row["actions"] else b"",
+                long_running_tool_ids_json=row["long_running_tool_ids_json"],
+                branch=row["branch"],
+                timestamp=row["timestamp"],
+                content=row["content"],
+                grounding_metadata=row["grounding_metadata"],
+                custom_metadata=row["custom_metadata"],
+                partial=row["partial"],
+                turn_complete=row["turn_complete"],
+                interrupted=row["interrupted"],
+                error_code=row["error_code"],
+                error_message=row["error_message"],
+            )
+
+    def list_events(self, session_id: str) -> "list[EventRecord]":
+        """List events for a session ordered by timestamp.
+
+        Args:
+            session_id: Session identifier.
+
+        Returns:
+            List of event records ordered by timestamp ASC.
+
+        Notes:
+            Uses index on (session_id, timestamp ASC).
+            JSONB fields are automatically deserialized by psycopg.
+            BYTEA actions are converted to bytes.
+        """
+        query = pg_sql.SQL("""
+        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 {table}
+        WHERE session_id = %s
+        ORDER BY timestamp ASC
+        """).format(table=pg_sql.Identifier(self._events_table))
+
+        try:
+            with self._config.provide_connection() as conn, conn.cursor() as cur:
+                cur.execute(query, (session_id,))
+                rows = cur.fetchall()
+
+                return [
+                    EventRecord(
+                        id=row["id"],
+                        session_id=row["session_id"],
+                        app_name=row["app_name"],
+                        user_id=row["user_id"],
+                        invocation_id=row["invocation_id"],
+                        author=row["author"],
+                        actions=bytes(row["actions"]) if row["actions"] else b"",
+                        long_running_tool_ids_json=row["long_running_tool_ids_json"],
+                        branch=row["branch"],
+                        timestamp=row["timestamp"],
+                        content=row["content"],
+                        grounding_metadata=row["grounding_metadata"],
+                        custom_metadata=row["custom_metadata"],
+                        partial=row["partial"],
+                        turn_complete=row["turn_complete"],
+                        interrupted=row["interrupted"],
+                        error_code=row["error_code"],
+                        error_message=row["error_message"],
+                    )
+                    for row in rows
+                ]
+        except errors.UndefinedTable:
+            return []