sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

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

@@ -0,0 +1,246 @@
+"""Create ADK session, events, and memory tables migration using store DDL definitions."""
+
+import inspect
+import logging
+from typing import TYPE_CHECKING, Any, NoReturn, cast
+
+from sqlspec.exceptions import SQLSpecError
+from sqlspec.utils.logging import get_logger, log_with_context
+from sqlspec.utils.module_loader import import_string
+
+if TYPE_CHECKING:
+    from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore
+    from sqlspec.extensions.adk.store import BaseAsyncADKStore
+    from sqlspec.migrations.context import MigrationContext
+
+logger = get_logger("sqlspec.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 _get_memory_store_class(
+    context: "MigrationContext | None",
+) -> "type[BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore] | None":
+    """Get the appropriate memory store class based on the config's module path.
+
+    Args:
+        context: Migration context containing config.
+
+    Returns:
+        Memory store class matching the config's adapter, or None if not available.
+
+    Notes:
+        Dynamically imports the memory store class from the config's module path.
+        For example, AsyncpgConfig at 'sqlspec.adapters.asyncpg.config'
+        maps to AsyncpgADKMemoryStore at 'sqlspec.adapters.asyncpg.adk.store.AsyncpgADKMemoryStore'.
+    """
+    if not context or not context.config:
+        return None
+
+    config_class = type(context.config)
+    config_module = config_class.__module__
+    config_name = config_class.__name__
+
+    if not config_module.startswith("sqlspec.adapters."):
+        return None
+
+    adapter_name = config_module.split(".")[2]
+    store_class_name = config_name.replace("Config", "ADKMemoryStore")
+
+    store_path = f"sqlspec.adapters.{adapter_name}.adk.store.{store_class_name}"
+
+    try:
+        store_class: type[BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore] = import_string(store_path)
+    except ImportError:
+        log_with_context(logger, logging.DEBUG, "adk.migration.memory_store.missing", store_path=store_path)
+        return None
+    else:
+        return store_class
+
+
+def _is_memory_enabled(context: "MigrationContext | None") -> bool:
+    """Check if memory migration is enabled in the config.
+
+    Args:
+        context: Migration context containing config.
+
+    Returns:
+        True if memory migration should be included, False otherwise.
+
+    Notes:
+        Checks config.extension_config["adk"]["include_memory_migration"].
+        Defaults to True if not specified and enable_memory is True.
+    """
+    if not context or not context.config:
+        return False
+
+    config = context.config
+    extension_config = cast("dict[str, dict[str, Any]]", config.extension_config)
+    adk_config: dict[str, Any] = extension_config.get("adk", {})
+
+    include_memory = adk_config.get("include_memory_migration")
+    if include_memory is not None:
+        return bool(include_memory)
+
+    return bool(adk_config.get("enable_memory", True))
+
+
+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, events, and memory 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.).
+        Memory table is included if enable_memory or include_memory_migration is True.
+    """
+    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)
+
+    statements = [
+        await store_instance._get_create_sessions_table_sql(),  # pyright: ignore[reportPrivateUsage]
+        await store_instance._get_create_events_table_sql(),  # pyright: ignore[reportPrivateUsage]
+    ]
+
+    if _is_memory_enabled(context):
+        memory_store_class = _get_memory_store_class(context)
+        if memory_store_class is not None:
+            memory_store = memory_store_class(config=context.config)
+            memory_sql = memory_store._get_create_memory_table_sql()  # pyright: ignore[reportPrivateUsage]
+            if inspect.isawaitable(memory_sql):
+                memory_sql = await memory_sql
+            if isinstance(memory_sql, list):
+                statements.extend(memory_sql)
+            else:
+                statements.append(memory_sql)
+            log_with_context(
+                logger, logging.DEBUG, "adk.migration.memory.include", table_name=memory_store.memory_table
+            )
+
+    return statements
+
+
+async def down(context: "MigrationContext | None" = None) -> "list[str]":
+    """Drop the ADK session, events, and memory 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"].
+        Memory table is included if enable_memory or include_memory_migration is True.
+    """
+    if context is None or context.config is None:
+        _raise_missing_config()
+
+    statements: list[str] = []
+
+    if _is_memory_enabled(context):
+        memory_store_class = _get_memory_store_class(context)
+        if memory_store_class is not None:
+            memory_store = memory_store_class(config=context.config)
+            memory_drop_stmts = memory_store._get_drop_memory_table_sql()  # pyright: ignore[reportPrivateUsage]
+            statements.extend(memory_drop_stmts)
+            log_with_context(
+                logger, logging.DEBUG, "adk.migration.memory.drop.include", table_name=memory_store.memory_table
+            )
+
+    store_class = _get_store_class(context)
+    store_instance = store_class(config=context.config)
+    statements.extend(store_instance._get_drop_tables_sql())  # pyright: ignore[reportPrivateUsage]
+
+    return statements

sqlspec/extensions/adk/service.py

@@ -0,0 +1,225 @@
+"""SQLSpec-backed session service for Google ADK."""
+
+import logging
+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, log_with_context
+
+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("sqlspec.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(connection_config={"dsn": "postgresql://..."})
+        store = AsyncpgADKStore(config)
+        await store.ensure_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
+        )
+        log_with_context(
+            logger, logging.DEBUG, "adk.session.create", app_name=app_name, session_id=session_id, has_state=bool(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:
+            log_with_context(
+                logger, logging.DEBUG, "adk.session.get", app_name=app_name, session_id=session_id, found=False
+            )
+            return None
+
+        if record["app_name"] != app_name or record["user_id"] != user_id:
+            log_with_context(
+                logger, logging.DEBUG, "adk.session.get", app_name=app_name, session_id=session_id, found=False
+            )
+            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)
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "adk.session.get",
+            app_name=app_name,
+            session_id=session_id,
+            found=True,
+            event_count=len(events),
+        )
+
+        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]
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "adk.session.list",
+            app_name=app_name,
+            has_user_id=user_id is not None,
+            count=len(sessions),
+        )
+
+        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:
+            log_with_context(
+                logger, logging.DEBUG, "adk.session.delete", app_name=app_name, session_id=session_id, deleted=False
+            )
+            return
+
+        if record["app_name"] != app_name or record["user_id"] != user_id:
+            log_with_context(
+                logger, logging.DEBUG, "adk.session.delete", app_name=app_name, session_id=session_id, deleted=False
+            )
+            return
+
+        await self._store.delete_session(session_id)
+        log_with_context(
+            logger, logging.DEBUG, "adk.session.delete", app_name=app_name, session_id=session_id, deleted=True
+        )
+
+    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)
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "adk.session.event.append",
+            app_name=session.app_name,
+            session_id=session.id,
+            partial=bool(event.partial),
+        )
+
+        session_record = await self._store.get_session(session.id)
+        if session_record:
+            session.last_update_time = session_record["update_time"].timestamp()
+
+        return event