Source code for sqlspec.extensions.adk.maintenance

"""Retention helpers for ADK stores.

Standalone functions to prune aged sessions, purge expired events, sweep scoped
memories, and prune aged artifact versions, decoupled from any framework worker
so retention can run from a cron job, a CLI, or a task queue.

Storage-level upkeep such as vacuuming, checkpointing, or refreshing optimizer
statistics is deliberately out of scope. Those are operator decisions that need
elevated privileges and their own scheduling, and several cannot run inside the
transaction a store session provides.
"""

import time
from datetime import datetime, timedelta, timezone
from typing import Any, Literal

from typing_extensions import TypedDict

from sqlspec.extensions.adk._config_utils import _adk_adapter_store_class
from sqlspec.utils.sync_tools import await_

__all__ = (
    "PruneReport",
    "prune_artifacts",
    "prune_artifacts_sync",
    "prune_events",
    "prune_events_sync",
    "prune_memory",
    "prune_memory_sync",
    "prune_sessions",
    "prune_sessions_sync",
    "prune_user_state",
    "prune_user_state_sync",
)




class PruneReport(TypedDict):
    """Performance telemetry for a pruning operation."""

    deleted_count: int
    elapsed_ms: float
    table: str

def _resolve_session_store(target: Any) -> Any: """Resolve an ADK session store instance from target.""" if hasattr(target, "delete_idle_sessions") or hasattr(target, "delete_expired_events"): return target if hasattr(target, "extension_config") or hasattr(target, "provide_session"): store_cls = _adk_adapter_store_class(target, "ADKStore") return store_cls(target) msg = f"Cannot resolve ADK session store from target of type {type(target).__name__}" raise TypeError(msg) def _resolve_memory_store(target: Any) -> Any: """Resolve an ADK memory store instance from target.""" if hasattr(target, "delete_entries_older_than"): return target if hasattr(target, "extension_config") or hasattr(target, "provide_session"): store_cls = _adk_adapter_store_class(target, "ADKMemoryStore") return store_cls(target) msg = f"Cannot resolve ADK memory store from target of type {type(target).__name__}" raise TypeError(msg) def _resolve_artifact_service(target: Any) -> Any: """Resolve a storage-aware ADK artifact service instance from target. Args: target: Candidate artifact service. Returns: The artifact service. Raises: TypeError: If target does not expose both the retention method and its metadata store. """ if hasattr(target, "delete_artifacts_older_than") and hasattr(target, "store"): return target msg = ( f"Cannot resolve ADK artifact service from target of type {type(target).__name__}. " "Artifact pruning removes content objects as well as metadata rows, so it requires a " "storage-aware artifact service rather than a database config or a bare metadata store." ) raise TypeError(msg) def _ensure_positive_days(value: Any, parameter: str) -> int: """Validate a retention age expressed in whole days. Args: value: Candidate age value. parameter: Name of the keyword argument being validated. Returns: The validated age in days. Raises: ValueError: If the value is not a positive integer. """ if not isinstance(value, int) or isinstance(value, bool) or value <= 0: msg = f"{parameter} must be a positive integer, got {value!r}" raise ValueError(msg) return value async def _call_store_method(store: Any, method_name: str, *args: Any, **kwargs: Any) -> int: """Call a store method handling both async and sync implementations.""" method = getattr(store, method_name) result = method(*args, **kwargs) if hasattr(result, "__await__"): deleted = await result else: deleted = result return int(deleted) if deleted is not None else 0

async def prune_sessions(target: Any, *, idle_days: int = 30, app_name: str | None = None) -> PruneReport: """Prune sessions that have been idle longer than specified days. Args: target: ADKStore instance, DatabaseConfig, or DriverAdapter. idle_days: Number of days of inactivity before a session is pruned. app_name: Optional application name to limit pruning. Returns: PruneReport containing deleted row count and timing. Raises: ValueError: If ``idle_days`` is not a positive integer. """ retention_days = _ensure_positive_days(idle_days, "idle_days") start = time.perf_counter() store = _resolve_session_store(target) table_name = getattr(store, "session_table", "adk_session") cutoff = datetime.now(timezone.utc) - timedelta(days=retention_days) deleted = await _call_store_method(store, "delete_idle_sessions", cutoff, app_name=app_name) elapsed_ms = (time.perf_counter() - start) * 1000.0 return PruneReport(deleted_count=deleted, elapsed_ms=elapsed_ms, table=str(table_name))