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

sqlspec/extensions/adk/converters.py

@@ -0,0 +1,172 @@
+"""Conversion functions between ADK models and database records."""
+
+import json
+import pickle
+from datetime import datetime, timezone
+from typing import Any
+
+from google.adk.events.event import Event
+from google.adk.sessions import Session
+from google.genai import types
+
+from sqlspec.extensions.adk._types import EventRecord, SessionRecord
+from sqlspec.utils.logging import get_logger
+
+logger = get_logger("extensions.adk.converters")
+
+__all__ = ("event_to_record", "record_to_event", "record_to_session", "session_to_record")
+
+
+def session_to_record(session: "Session") -> SessionRecord:
+    """Convert ADK Session to database record.
+
+    Args:
+        session: ADK Session object.
+
+    Returns:
+        SessionRecord for database storage.
+    """
+    return SessionRecord(
+        id=session.id,
+        app_name=session.app_name,
+        user_id=session.user_id,
+        state=session.state,
+        create_time=datetime.now(timezone.utc),
+        update_time=datetime.fromtimestamp(session.last_update_time, tz=timezone.utc),
+    )
+
+
+def record_to_session(record: SessionRecord, events: "list[EventRecord]") -> "Session":
+    """Convert database record to ADK Session.
+
+    Args:
+        record: Session database record.
+        events: List of event records for this session.
+
+    Returns:
+        ADK Session object.
+    """
+    event_objects = [record_to_event(event_record) for event_record in events]
+
+    return Session(
+        id=record["id"],
+        app_name=record["app_name"],
+        user_id=record["user_id"],
+        state=record["state"],
+        events=event_objects,
+        last_update_time=record["update_time"].timestamp(),
+    )
+
+
+def event_to_record(event: "Event", session_id: str, app_name: str, user_id: str) -> EventRecord:
+    """Convert ADK Event to database record.
+
+    Args:
+        event: ADK Event object.
+        session_id: ID of the parent session.
+        app_name: Name of the application.
+        user_id: ID of the user.
+
+    Returns:
+        EventRecord for database storage.
+    """
+    actions_bytes = pickle.dumps(event.actions)
+
+    long_running_tool_ids_json = None
+    if event.long_running_tool_ids:
+        long_running_tool_ids_json = json.dumps(list(event.long_running_tool_ids))
+
+    content_dict = None
+    if event.content:
+        content_dict = event.content.model_dump(exclude_none=True, mode="json")
+
+    grounding_metadata_dict = None
+    if event.grounding_metadata:
+        grounding_metadata_dict = event.grounding_metadata.model_dump(exclude_none=True, mode="json")
+
+    custom_metadata_dict = event.custom_metadata
+
+    return EventRecord(
+        id=event.id,
+        app_name=app_name,
+        user_id=user_id,
+        session_id=session_id,
+        invocation_id=event.invocation_id,
+        author=event.author,
+        branch=event.branch,
+        actions=actions_bytes,
+        long_running_tool_ids_json=long_running_tool_ids_json,
+        timestamp=datetime.fromtimestamp(event.timestamp, tz=timezone.utc),
+        content=content_dict,
+        grounding_metadata=grounding_metadata_dict,
+        custom_metadata=custom_metadata_dict,
+        partial=event.partial,
+        turn_complete=event.turn_complete,
+        interrupted=event.interrupted,
+        error_code=event.error_code,
+        error_message=event.error_message,
+    )
+
+
+def record_to_event(record: "EventRecord") -> "Event":
+    """Convert database record to ADK Event.
+
+    Args:
+        record: Event database record.
+
+    Returns:
+        ADK Event object.
+    """
+    actions = pickle.loads(record["actions"])  # noqa: S301
+
+    long_running_tool_ids = None
+    if record["long_running_tool_ids_json"]:
+        long_running_tool_ids = set(json.loads(record["long_running_tool_ids_json"]))
+
+    return Event(
+        id=record["id"],
+        invocation_id=record["invocation_id"],
+        author=record["author"],
+        branch=record["branch"],
+        actions=actions,
+        timestamp=record["timestamp"].timestamp(),
+        content=_decode_content(record["content"]),
+        long_running_tool_ids=long_running_tool_ids,
+        partial=record["partial"],
+        turn_complete=record["turn_complete"],
+        error_code=record["error_code"],
+        error_message=record["error_message"],
+        interrupted=record["interrupted"],
+        grounding_metadata=_decode_grounding_metadata(record["grounding_metadata"]),
+        custom_metadata=record["custom_metadata"],
+    )
+
+
+def _decode_content(content_dict: "dict[str, Any] | None") -> Any:
+    """Decode content dictionary from database to ADK Content object.
+
+    Args:
+        content_dict: Content dictionary from database.
+
+    Returns:
+        ADK Content object or None.
+    """
+    if not content_dict:
+        return None
+
+    return types.Content.model_validate(content_dict)
+
+
+def _decode_grounding_metadata(grounding_dict: "dict[str, Any] | None") -> Any:
+    """Decode grounding metadata dictionary from database to ADK object.
+
+    Args:
+        grounding_dict: Grounding metadata dictionary from database.
+
+    Returns:
+        ADK GroundingMetadata object or None.
+    """
+    if not grounding_dict:
+        return None
+
+    return types.GroundingMetadata.model_validate(grounding_dict)

sqlspec/extensions/adk/migrations/0001_create_adk_tables.py

@@ -0,0 +1,144 @@
+"""Create ADK session and events tables migration using store DDL definitions."""
+
+from typing import TYPE_CHECKING, NoReturn
+
+from sqlspec.exceptions import SQLSpecError
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.module_loader import import_string
+
+if TYPE_CHECKING:
+    from sqlspec.extensions.adk.store import BaseAsyncADKStore
+    from sqlspec.migrations.context import MigrationContext
+
+logger = get_logger("migrations.adk.tables")
+
+__all__ = ("down", "up")
+
+
+def _get_store_class(context: "MigrationContext | None") -> "type[BaseAsyncADKStore]":
+    """Get the appropriate store class based on the config's module path.
+
+    Args:
+        context: Migration context containing config.
+
+    Returns:
+        Store class matching the config's adapter.
+
+    Notes:
+        Dynamically imports the store class from the config's module path.
+        For example, AsyncpgConfig at 'sqlspec.adapters.asyncpg.config'
+        maps to AsyncpgADKStore at 'sqlspec.adapters.asyncpg.adk.store.AsyncpgADKStore'.
+    """
+    if not context or not context.config:
+        _raise_missing_config()
+
+    config_class = type(context.config)
+    config_module = config_class.__module__
+    config_name = config_class.__name__
+
+    if not config_module.startswith("sqlspec.adapters."):
+        _raise_unsupported_config(f"{config_module}.{config_name}")
+
+    adapter_name = config_module.split(".")[2]
+    store_class_name = config_name.replace("Config", "ADKStore")
+
+    store_path = f"sqlspec.adapters.{adapter_name}.adk.store.{store_class_name}"
+
+    try:
+        store_class: type[BaseAsyncADKStore] = import_string(store_path)
+    except ImportError as e:
+        _raise_store_import_failed(store_path, e)
+
+    return store_class
+
+
+def _raise_missing_config() -> NoReturn:
+    """Raise error when migration context has no config.
+
+    Raises:
+        SQLSpecError: Always raised.
+    """
+    msg = "Migration context must have a config to determine store class"
+    raise SQLSpecError(msg)
+
+
+def _raise_unsupported_config(config_type: str) -> NoReturn:
+    """Raise error for unsupported config type.
+
+    Args:
+        config_type: The unsupported config type name.
+
+    Raises:
+        SQLSpecError: Always raised with config type info.
+    """
+    msg = f"Unsupported config type for ADK migration: {config_type}"
+    raise SQLSpecError(msg)
+
+
+def _raise_store_import_failed(store_path: str, error: ImportError) -> NoReturn:
+    """Raise error when store class import fails.
+
+    Args:
+        store_path: The import path that failed.
+        error: The original import error.
+
+    Raises:
+        SQLSpecError: Always raised with import details.
+    """
+    msg = f"Failed to import ADK store class from {store_path}: {error}"
+    raise SQLSpecError(msg) from error
+
+
+async def up(context: "MigrationContext | None" = None) -> "list[str]":
+    """Create the ADK session and events tables using store DDL definitions.
+
+    This migration delegates to the appropriate store class to generate
+    dialect-specific DDL. The store classes contain the single source of
+    truth for table schemas.
+
+    Args:
+        context: Migration context containing config.
+
+    Returns:
+        List of SQL statements to execute for upgrade.
+
+    Notes:
+        Configuration is read from context.config.extension_config["adk"].
+        Supports custom table names and optional owner_id_column for linking
+        sessions to owner tables (users, tenants, teams, etc.).
+    """
+    if context is None or context.config is None:
+        _raise_missing_config()
+
+    store_class = _get_store_class(context)
+    store_instance = store_class(config=context.config)
+
+    return [
+        await store_instance._get_create_sessions_table_sql(),  # pyright: ignore[reportPrivateUsage]
+        await store_instance._get_create_events_table_sql(),  # pyright: ignore[reportPrivateUsage]
+    ]
+
+
+async def down(context: "MigrationContext | None" = None) -> "list[str]":
+    """Drop the ADK session and events tables using store DDL definitions.
+
+    This migration delegates to the appropriate store class to generate
+    dialect-specific DROP statements. The store classes contain the single
+    source of truth for table schemas.
+
+    Args:
+        context: Migration context containing config.
+
+    Returns:
+        List of SQL statements to execute for downgrade.
+
+    Notes:
+        Configuration is read from context.config.extension_config["adk"].
+    """
+    if context is None or context.config is None:
+        _raise_missing_config()
+
+    store_class = _get_store_class(context)
+    store_instance = store_class(config=context.config)
+
+    return store_instance._get_drop_tables_sql()  # pyright: ignore[reportPrivateUsage]

sqlspec/extensions/adk/service.py

@@ -0,0 +1,181 @@
+"""SQLSpec-backed session service for Google ADK."""
+
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from google.adk.sessions.base_session_service import BaseSessionService, GetSessionConfig, ListSessionsResponse
+
+from sqlspec.extensions.adk.converters import event_to_record, record_to_session
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from google.adk.events.event import Event
+    from google.adk.sessions import Session
+
+    from sqlspec.extensions.adk.store import BaseAsyncADKStore
+
+logger = get_logger("extensions.adk.service")
+
+__all__ = ("SQLSpecSessionService",)
+
+
+class SQLSpecSessionService(BaseSessionService):
+    """SQLSpec-backed implementation of BaseSessionService.
+
+    Provides session and event storage using SQLSpec database adapters.
+    Delegates all database operations to a store implementation.
+
+    Args:
+        store: Database store implementation (e.g., AsyncpgADKStore).
+
+    Example:
+        from sqlspec.adapters.asyncpg import AsyncpgConfig
+        from sqlspec.adapters.asyncpg.adk.store import AsyncpgADKStore
+        from sqlspec.extensions.adk.service import SQLSpecSessionService
+
+        config = AsyncpgConfig(pool_config={"dsn": "postgresql://..."})
+        store = AsyncpgADKStore(config)
+        await store.create_tables()
+
+        service = SQLSpecSessionService(store)
+        session = await service.create_session(
+            app_name="my_app",
+            user_id="user123",
+            state={"key": "value"}
+        )
+    """
+
+    def __init__(self, store: "BaseAsyncADKStore") -> None:
+        """Initialize the session service.
+
+        Args:
+            store: Database store implementation.
+        """
+        self._store = store
+
+    @property
+    def store(self) -> "BaseAsyncADKStore":
+        """Return the database store."""
+        return self._store
+
+    async def create_session(
+        self, *, app_name: str, user_id: str, state: "dict[str, Any] | None" = None, session_id: "str | None" = None
+    ) -> "Session":
+        """Create a new session.
+
+        Args:
+            app_name: Name of the application.
+            user_id: ID of the user.
+            state: Initial state of the session.
+            session_id: Client-provided session ID. If None, generates a UUID.
+
+        Returns:
+            The newly created session.
+        """
+        if session_id is None:
+            session_id = str(uuid.uuid4())
+
+        if state is None:
+            state = {}
+
+        record = await self._store.create_session(
+            session_id=session_id, app_name=app_name, user_id=user_id, state=state
+        )
+
+        return record_to_session(record, events=[])
+
+    async def get_session(
+        self, *, app_name: str, user_id: str, session_id: str, config: "GetSessionConfig | None" = None
+    ) -> "Session | None":
+        """Get a session by ID.
+
+        Args:
+            app_name: Name of the application.
+            user_id: ID of the user.
+            session_id: Session identifier.
+            config: Configuration for retrieving events.
+
+        Returns:
+            Session object if found, None otherwise.
+        """
+        record = await self._store.get_session(session_id)
+
+        if not record:
+            return None
+
+        if record["app_name"] != app_name or record["user_id"] != user_id:
+            return None
+
+        after_timestamp = None
+        limit = None
+
+        if config:
+            if config.after_timestamp:
+                after_timestamp = datetime.fromtimestamp(config.after_timestamp, tz=timezone.utc)
+            limit = config.num_recent_events
+
+        events = await self._store.get_events(session_id=session_id, after_timestamp=after_timestamp, limit=limit)
+
+        return record_to_session(record, events)
+
+    async def list_sessions(self, *, app_name: str, user_id: str | None = None) -> "ListSessionsResponse":
+        """List all sessions for an app, optionally filtered by user.
+
+        Args:
+            app_name: Name of the application.
+            user_id: ID of the user. If None, all sessions for the app are listed.
+
+        Returns:
+            Response containing list of sessions (without events).
+        """
+        records = await self._store.list_sessions(app_name=app_name, user_id=user_id)
+
+        sessions = [record_to_session(record, events=[]) for record in records]
+
+        return ListSessionsResponse(sessions=sessions)
+
+    async def delete_session(self, *, app_name: str, user_id: str, session_id: str) -> None:
+        """Delete a session and all its events.
+
+        Args:
+            app_name: Name of the application.
+            user_id: ID of the user.
+            session_id: Session identifier.
+        """
+        record = await self._store.get_session(session_id)
+
+        if not record:
+            return
+
+        if record["app_name"] != app_name or record["user_id"] != user_id:
+            return
+
+        await self._store.delete_session(session_id)
+
+    async def append_event(self, session: "Session", event: "Event") -> "Event":
+        """Append an event to a session.
+
+        Args:
+            session: Session to append to.
+            event: Event to append.
+
+        Returns:
+            The appended event.
+        """
+        event = await super().append_event(session, event)  # pyright: ignore
+
+        if event.partial:
+            return event
+
+        event_record = event_to_record(
+            event=event, session_id=session.id, app_name=session.app_name, user_id=session.user_id
+        )
+
+        await self._store.append_event(event_record)
+
+        session_record = await self._store.get_session(session.id)
+        if session_record:
+            session.last_update_time = session_record["update_time"].timestamp()
+
+        return event