sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/extensions/adk/memory/store.py

@@ -0,0 +1,525 @@
+"""Base store classes for ADK memory backend (sync and async)."""
+
+import logging
+import re
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar
+
+from sqlspec.extensions.adk._config_utils import _ADKMemoryStoreConfig, _get_adk_memory_store_config
+from sqlspec.observability import resolve_db_system
+from sqlspec.utils.logging import get_logger, log_with_context
+
+if TYPE_CHECKING:
+    from sqlspec.config import DatabaseConfigProtocol
+    from sqlspec.extensions.adk.memory._types import MemoryRecord
+
+ConfigT = TypeVar("ConfigT", bound="DatabaseConfigProtocol[Any, Any, Any]")
+
+logger = get_logger("sqlspec.extensions.adk.memory.store")
+
+__all__ = ("BaseAsyncADKMemoryStore", "BaseSyncADKMemoryStore")
+
+VALID_TABLE_NAME_PATTERN: Final = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
+COLUMN_NAME_PATTERN: Final = re.compile(r"^(\w+)")
+MAX_TABLE_NAME_LENGTH: Final = 63
+
+
+def _parse_owner_id_column(owner_id_column_ddl: str) -> str:
+    """Extract column name from owner ID column DDL definition.
+
+    Args:
+        owner_id_column_ddl: Full column DDL string.
+
+    Returns:
+        Column name only (first word).
+
+    Raises:
+        ValueError: If DDL format is invalid.
+    """
+    match = COLUMN_NAME_PATTERN.match(owner_id_column_ddl.strip())
+    if not match:
+        msg = f"Invalid owner_id_column DDL: {owner_id_column_ddl!r}. Must start with column name."
+        raise ValueError(msg)
+
+    return match.group(1)
+
+
+def _validate_table_name(table_name: str) -> None:
+    """Validate table name for SQL safety.
+
+    Args:
+        table_name: Table name to validate.
+
+    Raises:
+        ValueError: If table name is invalid.
+    """
+    if not table_name:
+        msg = "Table name cannot be empty"
+        raise ValueError(msg)
+
+    if len(table_name) > MAX_TABLE_NAME_LENGTH:
+        msg = f"Table name too long: {len(table_name)} chars (max {MAX_TABLE_NAME_LENGTH})"
+        raise ValueError(msg)
+
+    if not VALID_TABLE_NAME_PATTERN.match(table_name):
+        msg = (
+            f"Invalid table name: {table_name!r}. "
+            "Must start with letter/underscore and contain only alphanumeric characters and underscores"
+        )
+        raise ValueError(msg)
+
+
+class BaseAsyncADKMemoryStore(ABC, Generic[ConfigT]):
+    """Base class for async SQLSpec-backed ADK memory stores.
+
+    Implements storage operations for Google ADK memory entries using
+    SQLSpec database adapters with async/await.
+
+    This abstract base class provides common functionality for all database-specific
+    memory store implementations including:
+    - Connection management via SQLSpec configs
+    - Table name validation
+    - Memory entry CRUD operations
+    - Text search with optional full-text search support
+
+    Subclasses must implement dialect-specific SQL queries and will be created
+    in each adapter directory (e.g., sqlspec/adapters/asyncpg/adk/store.py).
+
+    Args:
+        config: SQLSpec database configuration with extension_config["adk"] settings.
+
+    Notes:
+        Configuration is read from config.extension_config["adk"]:
+        - memory_table: Memory table name (default: "adk_memory_entries")
+        - memory_use_fts: Enable full-text search when supported (default: False)
+        - memory_max_results: Max search results (default: 20)
+        - owner_id_column: Optional owner FK column DDL (default: None)
+        - enable_memory: Whether memory is enabled (default: True)
+    """
+
+    __slots__ = (
+        "_config",
+        "_enabled",
+        "_max_results",
+        "_memory_table",
+        "_owner_id_column_ddl",
+        "_owner_id_column_name",
+        "_use_fts",
+    )
+
+    def __init__(self, config: ConfigT) -> None:
+        """Initialize the ADK memory store.
+
+        Args:
+            config: SQLSpec database configuration.
+
+        Notes:
+            Reads configuration from config.extension_config["adk"]:
+            - memory_table: Memory table name (default: "adk_memory_entries")
+            - memory_use_fts: Enable full-text search when supported (default: False)
+            - memory_max_results: Max search results (default: 20)
+            - owner_id_column: Optional owner FK column DDL (default: None)
+            - enable_memory: Whether memory is enabled (default: True)
+        """
+        self._config = config
+        store_config = self._get_store_config_from_extension()
+        self._enabled: bool = store_config.get("enable_memory", True)
+        self._memory_table: str = str(store_config["memory_table"])
+        self._use_fts: bool = bool(store_config.get("use_fts", False))
+        self._max_results: int = store_config.get("max_results", 20)
+        self._owner_id_column_ddl: str | None = store_config.get("owner_id_column")
+        self._owner_id_column_name: str | None = (
+            _parse_owner_id_column(self._owner_id_column_ddl) if self._owner_id_column_ddl else None
+        )
+        _validate_table_name(self._memory_table)
+
+    def _get_store_config_from_extension(self) -> "_ADKMemoryStoreConfig":
+        """Extract ADK memory configuration from config.extension_config.
+
+        Returns:
+            Dict with memory_table, use_fts, max_results, and optionally owner_id_column.
+        """
+        return _get_adk_memory_store_config(self._config)
+
+    @property
+    def config(self) -> ConfigT:
+        """Return the database configuration."""
+        return self._config
+
+    @property
+    def memory_table(self) -> str:
+        """Return the memory table name."""
+        return self._memory_table
+
+    @property
+    def enabled(self) -> bool:
+        """Return whether memory store is enabled."""
+        return self._enabled
+
+    @property
+    def use_fts(self) -> bool:
+        """Return whether full-text search is enabled."""
+        return self._use_fts
+
+    @property
+    def max_results(self) -> int:
+        """Return the max search results limit."""
+        return self._max_results
+
+    @property
+    def owner_id_column_ddl(self) -> "str | None":
+        """Return the full owner ID column DDL (or None if not configured)."""
+        return self._owner_id_column_ddl
+
+    @property
+    def owner_id_column_name(self) -> "str | None":
+        """Return the owner ID column name only (or None if not configured)."""
+        return self._owner_id_column_name
+
+    @abstractmethod
+    async def create_tables(self) -> None:
+        """Create the memory table and indexes if they don't exist.
+
+        Should check self._enabled and skip table creation if False.
+        """
+        raise NotImplementedError
+
+    async def ensure_tables(self) -> None:
+        """Create tables when enabled and emit a standardized log entry."""
+
+        if not self._enabled:
+            self._log_memory_table_skipped()
+            return
+        await self.create_tables()
+        self._log_memory_table_created()
+
+    @abstractmethod
+    async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
+        """Bulk insert memory entries with deduplication.
+
+        Uses UPSERT pattern to skip duplicates based on event_id.
+
+        Args:
+            entries: List of memory records to insert.
+            owner_id: Optional owner ID value for owner_id_column (if configured).
+
+        Returns:
+            Number of entries actually inserted (excludes duplicates).
+
+        Raises:
+            RuntimeError: If memory store is disabled.
+        """
+        raise NotImplementedError
+
+    def _log_memory_table_created(self) -> None:
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "adk.memory.table.ready",
+            db_system=resolve_db_system(type(self).__name__),
+            memory_table=self._memory_table,
+        )
+
+    def _log_memory_table_skipped(self) -> None:
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "adk.memory.table.skipped",
+            db_system=resolve_db_system(type(self).__name__),
+            memory_table=self._memory_table,
+            reason="disabled",
+        )
+
+    @abstractmethod
+    async def search_entries(
+        self, query: str, app_name: str, user_id: str, limit: "int | None" = None
+    ) -> "list[MemoryRecord]":
+        """Search memory entries by text query.
+
+        Uses the configured search strategy (simple ILIKE or FTS).
+
+        Args:
+            query: Text query to search for.
+            app_name: Application name to filter by.
+            user_id: User ID to filter by.
+            limit: Maximum number of results (defaults to max_results config).
+
+        Returns:
+            List of matching memory records ordered by relevance/timestamp.
+
+        Raises:
+            RuntimeError: If memory store is disabled.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete_entries_by_session(self, session_id: str) -> int:
+        """Delete all memory entries for a specific session.
+
+        Args:
+            session_id: Session ID to delete entries for.
+
+        Returns:
+            Number of entries deleted.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete_entries_older_than(self, days: int) -> int:
+        """Delete memory entries older than specified days.
+
+        Used for TTL cleanup operations.
+
+        Args:
+            days: Number of days to retain entries.
+
+        Returns:
+            Number of entries deleted.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def _get_create_memory_table_sql(self) -> "str | list[str]":
+        """Get the CREATE TABLE SQL for the memory table.
+
+        Returns:
+            SQL statement(s) to create the memory table with indexes.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_drop_memory_table_sql(self) -> "list[str]":
+        """Get the DROP TABLE SQL statements for this database dialect.
+
+        Returns:
+            List of SQL statements to drop the memory table and indexes.
+        """
+        raise NotImplementedError
+
+
+class BaseSyncADKMemoryStore(ABC, Generic[ConfigT]):
+    """Base class for sync SQLSpec-backed ADK memory stores.
+
+    Implements storage operations for Google ADK memory entries using
+    SQLSpec database adapters with synchronous execution.
+
+    This abstract base class provides common functionality for sync database-specific
+    memory store implementations including:
+    - Connection management via SQLSpec configs
+    - Table name validation
+    - Memory entry CRUD operations
+    - Text search with optional full-text search support
+
+    Subclasses must implement dialect-specific SQL queries and will be created
+    in each adapter directory (e.g., sqlspec/adapters/sqlite/adk/store.py).
+
+    Args:
+        config: SQLSpec database configuration with extension_config["adk"] settings.
+
+    Notes:
+        Configuration is read from config.extension_config["adk"]:
+        - memory_table: Memory table name (default: "adk_memory_entries")
+        - memory_use_fts: Enable full-text search when supported (default: False)
+        - memory_max_results: Max search results (default: 20)
+        - owner_id_column: Optional owner FK column DDL (default: None)
+        - enable_memory: Whether memory is enabled (default: True)
+    """
+
+    __slots__ = (
+        "_config",
+        "_enabled",
+        "_max_results",
+        "_memory_table",
+        "_owner_id_column_ddl",
+        "_owner_id_column_name",
+        "_use_fts",
+    )
+
+    def __init__(self, config: ConfigT) -> None:
+        """Initialize the sync ADK memory store.
+
+        Args:
+            config: SQLSpec database configuration.
+
+        Notes:
+            Reads configuration from config.extension_config["adk"]:
+            - memory_table: Memory table name (default: "adk_memory_entries")
+            - memory_use_fts: Enable full-text search when supported (default: False)
+            - memory_max_results: Max search results (default: 20)
+            - owner_id_column: Optional owner FK column DDL (default: None)
+            - enable_memory: Whether memory is enabled (default: True)
+        """
+        self._config = config
+        store_config = self._get_store_config_from_extension()
+        self._enabled: bool = store_config.get("enable_memory", True)
+        self._memory_table: str = str(store_config["memory_table"])
+        self._use_fts: bool = bool(store_config.get("use_fts", False))
+        self._max_results: int = store_config.get("max_results", 20)
+        self._owner_id_column_ddl: str | None = store_config.get("owner_id_column")
+        self._owner_id_column_name: str | None = (
+            _parse_owner_id_column(self._owner_id_column_ddl) if self._owner_id_column_ddl else None
+        )
+        _validate_table_name(self._memory_table)
+
+    def _get_store_config_from_extension(self) -> "_ADKMemoryStoreConfig":
+        """Extract ADK memory configuration from config.extension_config.
+
+        Returns:
+            Dict with memory_table, use_fts, max_results, and optionally owner_id_column.
+        """
+        return _get_adk_memory_store_config(self._config)
+
+    @property
+    def config(self) -> ConfigT:
+        """Return the database configuration."""
+        return self._config
+
+    @property
+    def memory_table(self) -> str:
+        """Return the memory table name."""
+        return self._memory_table
+
+    @property
+    def enabled(self) -> bool:
+        """Return whether memory store is enabled."""
+        return self._enabled
+
+    @property
+    def use_fts(self) -> bool:
+        """Return whether full-text search is enabled."""
+        return self._use_fts
+
+    @property
+    def max_results(self) -> int:
+        """Return the max search results limit."""
+        return self._max_results
+
+    @property
+    def owner_id_column_ddl(self) -> "str | None":
+        """Return the full owner ID column DDL (or None if not configured)."""
+        return self._owner_id_column_ddl
+
+    @property
+    def owner_id_column_name(self) -> "str | None":
+        """Return the owner ID column name only (or None if not configured)."""
+        return self._owner_id_column_name
+
+    @abstractmethod
+    def create_tables(self) -> None:
+        """Create the memory table and indexes if they don't exist.
+
+        Should check self._enabled and skip table creation if False.
+        """
+        raise NotImplementedError
+
+    def ensure_tables(self) -> None:
+        """Create tables when enabled and emit a standardized log entry."""
+
+        if not self._enabled:
+            self._log_memory_table_skipped()
+            return
+        self.create_tables()
+        self._log_memory_table_created()
+
+    @abstractmethod
+    def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
+        """Bulk insert memory entries with deduplication.
+
+        Uses UPSERT pattern to skip duplicates based on event_id.
+
+        Args:
+            entries: List of memory records to insert.
+            owner_id: Optional owner ID value for owner_id_column (if configured).
+
+        Returns:
+            Number of entries actually inserted (excludes duplicates).
+
+        Raises:
+            RuntimeError: If memory store is disabled.
+        """
+        raise NotImplementedError
+
+    def _log_memory_table_created(self) -> None:
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "adk.memory.table.ready",
+            db_system=resolve_db_system(type(self).__name__),
+            memory_table=self._memory_table,
+        )
+
+    def _log_memory_table_skipped(self) -> None:
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "adk.memory.table.skipped",
+            db_system=resolve_db_system(type(self).__name__),
+            memory_table=self._memory_table,
+            reason="disabled",
+        )
+
+    @abstractmethod
+    def search_entries(
+        self, query: str, app_name: str, user_id: str, limit: "int | None" = None
+    ) -> "list[MemoryRecord]":
+        """Search memory entries by text query.
+
+        Uses the configured search strategy (simple ILIKE or FTS).
+
+        Args:
+            query: Text query to search for.
+            app_name: Application name to filter by.
+            user_id: User ID to filter by.
+            limit: Maximum number of results (defaults to max_results config).
+
+        Returns:
+            List of matching memory records ordered by relevance/timestamp.
+
+        Raises:
+            RuntimeError: If memory store is disabled.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_entries_by_session(self, session_id: str) -> int:
+        """Delete all memory entries for a specific session.
+
+        Args:
+            session_id: Session ID to delete entries for.
+
+        Returns:
+            Number of entries deleted.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_entries_older_than(self, days: int) -> int:
+        """Delete memory entries older than specified days.
+
+        Used for TTL cleanup operations.
+
+        Args:
+            days: Number of days to retain entries.
+
+        Returns:
+            Number of entries deleted.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_create_memory_table_sql(self) -> "str | list[str]":
+        """Get the CREATE TABLE SQL for the memory table.
+
+        Returns:
+            SQL statement(s) to create the memory table with indexes.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_drop_memory_table_sql(self) -> "list[str]":
+        """Get the DROP TABLE SQL statements for this database dialect.
+
+        Returns:
+            List of SQL statements to drop the memory table and indexes.
+        """
+        raise NotImplementedError

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

@@ -0,0 +1,184 @@
+"""Create ADK session, events, and memory tables migration using store DDL definitions."""
+
+import inspect
+import logging
+from typing import TYPE_CHECKING, NoReturn, cast
+
+from sqlspec.exceptions import SQLSpecError
+from sqlspec.extensions.adk._config_utils import (
+    _get_adk_adapter_store_class,
+    _get_adk_memory_migration_store_class,
+    _is_adk_memory_migration_enabled,
+)
+from sqlspec.utils.logging import get_logger, log_with_context
+
+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.AsyncpgADKStore'.
+    """
+    if not context or not context.config:
+        _raise_missing_config()
+
+    return cast("type[BaseAsyncADKStore]", _get_adk_adapter_store_class(context.config, "ADKStore"))
+
+
+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.AsyncpgADKMemoryStore'.
+    """
+    if not context or not context.config:
+        return None
+
+    store_class = _get_adk_memory_migration_store_class(context.config)
+    if store_class is None:
+        log_with_context(logger, logging.DEBUG, "adk.migration.memory_store.missing")
+        return None
+    return cast("type[BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore]", 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
+
+    return _is_adk_memory_migration_enabled(context.config)
+
+
+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)
+
+
+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