From dff22dc16e2403f6fcfa4beb39d72870197c22a5 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 11 Jul 2026 00:49:03 +0000 Subject: [PATCH 1/6] refactor(adk): consolidate service dispatch --- sqlspec/extensions/adk/_config_utils.py | 17 ++++--- sqlspec/extensions/adk/artifact/service.py | 26 +++++------ sqlspec/extensions/adk/memory/service.py | 28 +++++++---- .../test_adk/test_config_resolution.py | 15 +++++- .../test_adk/test_service_dispatch.py | 46 +++++++++++++++++++ 5 files changed, 104 insertions(+), 28 deletions(-) create mode 100644 tests/unit/extensions/test_adk/test_service_dispatch.py diff --git a/sqlspec/extensions/adk/_config_utils.py b/sqlspec/extensions/adk/_config_utils.py index ed6d347ac..eb1110bd4 100644 --- a/sqlspec/extensions/adk/_config_utils.py +++ b/sqlspec/extensions/adk/_config_utils.py @@ -76,9 +76,7 @@ def _adk_session_store_config(config: _ADKConfigSource) -> _ADKSessionStoreConfi "user_state_table": str(adk_config.get("user_state_table") or "adk_user_state"), "metadata_table": str(adk_config.get("metadata_table") or "adk_internal_metadata"), } - owner_id = adk_config.get("owner_id_column") - if owner_id is not None: - result["owner_id_column"] = cast("str", owner_id) + _apply_owner_id(result, adk_config) return result @@ -94,9 +92,7 @@ def _adk_memory_store_config(config: _ADKConfigSource) -> _ADKMemoryStoreConfig: "use_fts": bool(adk_config.get("memory_use_fts", False)), "max_results": int(max_results) if isinstance(max_results, int) else 20, } - owner_id = adk_config.get("owner_id_column") - if owner_id is not None: - result["owner_id_column"] = cast("str", owner_id) + _apply_owner_id(result, adk_config) return result @@ -107,6 +103,15 @@ def _adk_artifact_store_config(config: _ADKConfigSource) -> _ADKArtifactStoreCon return {"artifact_table": str(adk_config.get("artifact_table") or "adk_artifact")} +def _apply_owner_id( + result: "_ADKSessionStoreConfig | _ADKMemoryStoreConfig", adk_config: dict[str, Any] +) -> None: + """Copy the configured owner column into normalized store settings.""" + owner_id = adk_config.get("owner_id_column") + if owner_id is not None: + result["owner_id_column"] = cast("str", owner_id) + + def _adk_store_path(config: Any, store_suffix: str) -> str: """Return the adapter-specific ADK store import path.""" diff --git a/sqlspec/extensions/adk/artifact/service.py b/sqlspec/extensions/adk/artifact/service.py index d998143b3..99d792dd5 100644 --- a/sqlspec/extensions/adk/artifact/service.py +++ b/sqlspec/extensions/adk/artifact/service.py @@ -261,10 +261,7 @@ async def save_artifact( # Write content first (fail-fast before metadata) backend = self._registry.get(self._artifact_storage_uri) - if hasattr(backend, "write_bytes_async"): - await backend.write_bytes_async(content_path, content_bytes) - else: - backend.write_bytes_sync(content_path, content_bytes) + await _call_storage_backend(backend, "write_bytes_async", "write_bytes_sync", content_path, content_bytes) # Insert metadata row from datetime import datetime, timezone @@ -335,10 +332,7 @@ async def load_artifact( content_path = record["canonical_uri"].removeprefix(self._artifact_storage_uri + "/") backend = self._registry.get(self._artifact_storage_uri) - if hasattr(backend, "read_bytes_async"): - content_bytes = await backend.read_bytes_async(content_path) - else: - content_bytes = backend.read_bytes_sync(content_path) + content_bytes = await _call_storage_backend(backend, "read_bytes_async", "read_bytes_sync", content_path) log_with_context( logger, @@ -400,10 +394,7 @@ async def delete_artifact( for record in deleted_records: content_path = record["canonical_uri"].removeprefix(self._artifact_storage_uri + "/") try: - if hasattr(backend, "delete_async"): - await backend.delete_async(content_path) - else: - backend.delete_sync(content_path) + await _call_storage_backend(backend, "delete_async", "delete_sync", content_path) except Exception: log_with_context( logger, @@ -460,7 +451,6 @@ async def list_artifact_versions( app_name=app_name, user_id=user_id, filename=filename, session_id=session_id ) return [_record_to_artifact_version(r) for r in records] - async def get_artifact_version( self, *, @@ -488,3 +478,13 @@ async def get_artifact_version( if record is None: return None return _record_to_artifact_version(record) + + +async def _call_storage_backend( + backend: Any, async_method_name: str, sync_method_name: str, *args: Any, **kwargs: Any +) -> Any: + """Call the available async or sync storage-backend capability.""" + async_method = getattr(backend, async_method_name, None) + if async_method is not None: + return await async_method(*args, **kwargs) + return getattr(backend, sync_method_name)(*args, **kwargs) diff --git a/sqlspec/extensions/adk/memory/service.py b/sqlspec/extensions/adk/memory/service.py index 6d37eaf59..cf0972d3e 100644 --- a/sqlspec/extensions/adk/memory/service.py +++ b/sqlspec/extensions/adk/memory/service.py @@ -1,6 +1,7 @@ """SQLSpec-backed memory service for Google ADK.""" -from typing import TYPE_CHECKING +import inspect +from typing import TYPE_CHECKING, Any, cast from google.adk.memory.base_memory_service import BaseMemoryService, SearchMemoryResponse @@ -10,9 +11,10 @@ session_to_memory_records, ) from sqlspec.utils.logging import get_logger +from sqlspec.utils.sync_tools import async_ if TYPE_CHECKING: - from collections.abc import Mapping, Sequence + from collections.abc import Callable, Mapping, Sequence from google.adk.events.event import Event from google.adk.memory.memory_entry import MemoryEntry @@ -39,7 +41,7 @@ class SQLSpecMemoryService(BaseMemoryService): store: Database store implementation. """ - def __init__(self, store: "BaseAsyncADKMemoryStore") -> None: + def __init__(self, store: "BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore") -> None: """Initialize the memory service. Args: @@ -48,7 +50,7 @@ def __init__(self, store: "BaseAsyncADKMemoryStore") -> None: self._store = store @property - def store(self) -> "BaseAsyncADKMemoryStore": + def store(self) -> "BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore": """Return the database store.""" return self._store @@ -73,7 +75,7 @@ async def add_session_to_memory(self, session: "Session") -> None: ) return - inserted_count = await self._store.insert_memory_entries(records) + inserted_count = await self._call_store("insert_memory_entries", records) logger.debug( "Stored %d memory entries for session %s (total events: %d)", inserted_count, session.id, len(records) ) @@ -121,7 +123,7 @@ async def add_events_to_memory( ) return - inserted_count = await self._store.insert_memory_entries(records) + inserted_count = await self._call_store("insert_memory_entries", records) logger.debug( "Stored %d memory entries from %d events (app=%s, user=%s)", inserted_count, len(records), app_name, user_id ) @@ -161,7 +163,7 @@ async def add_memory( logger.debug("No content to store for memories (app=%s, user=%s)", app_name, user_id) return - inserted_count = await self._store.insert_memory_entries(records) + inserted_count = await self._call_store("insert_memory_entries", records) logger.debug( "Stored %d memory entries from %d memories (app=%s, user=%s)", inserted_count, @@ -183,7 +185,7 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str) -> "Se Returns: SearchMemoryResponse with memories: List[MemoryEntry]. """ - records = await self._store.search_entries(query=query, app_name=app_name, user_id=user_id) + records = await self._call_store("search_entries", query=query, app_name=app_name, user_id=user_id) memories = records_to_memory_entries(records) @@ -191,6 +193,16 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str) -> "Se return SearchMemoryResponse(memories=memories) + async def _call_store(self, method_name: str, *args: Any, **kwargs: Any) -> Any: + """Call an async store method or bridge a sync store method.""" + method = getattr(self._store, method_name) + if inspect.iscoroutinefunction(method): + return await method(*args, **kwargs) + sync_method = method + if TYPE_CHECKING: + sync_method = cast("Callable[..., Any]", method) + return await async_(sync_method)(*args, **kwargs) + class SQLSpecSyncMemoryService: """Synchronous SQLSpec-backed memory service. diff --git a/tests/unit/extensions/test_adk/test_config_resolution.py b/tests/unit/extensions/test_adk/test_config_resolution.py index 08a57b256..374b1f6ee 100644 --- a/tests/unit/extensions/test_adk/test_config_resolution.py +++ b/tests/unit/extensions/test_adk/test_config_resolution.py @@ -1,6 +1,6 @@ """Tests for ADK flat-config resolution.""" -from typing import Any +from typing import Any, cast from sqlspec.config import ADKConfig from sqlspec.extensions.adk._config_utils import ( @@ -8,6 +8,8 @@ _adk_memory_migration_enabled, _adk_memory_store_config, _adk_session_store_config, + _ADKSessionStoreConfig, + _apply_owner_id, ) @@ -18,6 +20,17 @@ def __init__(self, adk_config: dict[str, Any]) -> None: self.extension_config = {"adk": adk_config} +def test_apply_owner_id_adds_only_configured_values() -> None: + configured = cast("_ADKSessionStoreConfig", {}) + unconfigured = cast("_ADKSessionStoreConfig", {}) + + _apply_owner_id(configured, {"owner_id_column": "tenant_id UUID"}) + _apply_owner_id(unconfigured, {}) + + assert dict(configured) == {"owner_id_column": "tenant_id UUID"} + assert not unconfigured + + def test_adk_config_uses_flat_keys() -> None: """ADKConfig is a flat TypedDict; no per-adapter or nested negotiation blocks.""" annotations = set(ADKConfig.__annotations__) diff --git a/tests/unit/extensions/test_adk/test_service_dispatch.py b/tests/unit/extensions/test_adk/test_service_dispatch.py new file mode 100644 index 000000000..0a390d95a --- /dev/null +++ b/tests/unit/extensions/test_adk/test_service_dispatch.py @@ -0,0 +1,46 @@ +"""Regression tests for shared ADK service dispatch helpers.""" + +from typing import Any, cast + +import pytest + +from sqlspec.extensions.adk.artifact.service import _call_storage_backend +from sqlspec.extensions.adk.memory.service import SQLSpecMemoryService + + +class _SyncMemoryStore: + def insert_memory_entries(self, records: list[object]) -> int: + return len(records) + + +class _AsyncMemoryStore: + async def insert_memory_entries(self, records: list[object]) -> int: + return len(records) + + +@pytest.mark.anyio +@pytest.mark.parametrize("store", [_SyncMemoryStore(), _AsyncMemoryStore()]) +async def test_memory_service_call_store_dispatches_sync_and_async_methods(store: object) -> None: + service = SQLSpecMemoryService(cast("Any", store)) + + result = await service._call_store("insert_memory_entries", [object(), object()]) + + assert result == 2 + + +class _AsyncStorageBackend: + async def read_bytes_async(self, path: str) -> bytes: + return path.encode() + + +class _SyncStorageBackend: + def read_bytes_sync(self, path: str) -> bytes: + return path.encode() + + +@pytest.mark.anyio +@pytest.mark.parametrize("backend", [_AsyncStorageBackend(), _SyncStorageBackend()]) +async def test_artifact_backend_dispatch_uses_available_capability(backend: object) -> None: + result = await _call_storage_backend(cast("Any", backend), "read_bytes_async", "read_bytes_sync", "payload") + + assert result == b"payload" From 99ab2cb8c38c9f6f096db3c4c7b67327ba7e2ace Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 11 Jul 2026 00:52:37 +0000 Subject: [PATCH 2/6] refactor(adk): inline sync store operations --- .../adapters/cockroach_psycopg/adk/store.py | 553 +++---- sqlspec/adapters/mysqlconnector/adk/store.py | 245 +-- sqlspec/adapters/oracledb/adk/store.py | 1423 ++++++++--------- sqlspec/adapters/psycopg/adk/store.py | 555 +++---- sqlspec/adapters/sqlite/adk/store.py | 575 +++---- .../adapters/test_psycopg/test_adk_store.py | 8 +- .../test_adk/test_store_instantiation.py | 38 + 7 files changed, 1490 insertions(+), 1907 deletions(-) diff --git a/sqlspec/adapters/cockroach_psycopg/adk/store.py b/sqlspec/adapters/cockroach_psycopg/adk/store.py index 3141052ca..5211cea82 100644 --- a/sqlspec/adapters/cockroach_psycopg/adk/store.py +++ b/sqlspec/adapters/cockroach_psycopg/adk/store.py @@ -558,211 +558,6 @@ def __init__(self, config: "CockroachPsycopgSyncConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() - - 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.""" - return self._create_session(session_id, app_name, user_id, state, owner_id) - - def get_session( - self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": - """Get session by ID.""" - return self._get_session(app_name, user_id, session_id, renew_for=renew_for) - - def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: - """Update session state.""" - self._update_session_state(app_name, user_id, session_id, state) - - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": - """List sessions for an app.""" - return self._list_sessions(app_name, user_id) - - def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: - """Delete session and associated events.""" - self._delete_session(app_name, user_id, session_id) - - def append_event(self, event_record: EventRecord) -> None: - """Append an event to a session.""" - self._append_event(event_record) - - def append_event_and_update_state( - self, - event_record: EventRecord, - app_name: str, - user_id: str, - session_id: str, - state: "dict[str, Any]", - *, - app_state: "dict[str, Any] | None" = None, - user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: - """Atomically append an event and update session + scoped state.""" - return self._append_event_and_update_state( - event_record, app_name, user_id, session_id, state, app_state=app_state, user_state=user_state - ) - - def get_events( - self, - app_name: str, - user_id: str, - session_id: str, - after_timestamp: "datetime | None" = None, - limit: "int | None" = None, - ) -> "list[EventRecord]": - """Get events for a session.""" - return self._get_events(app_name, user_id, session_id, after_timestamp, limit) - - def delete_expired_events(self, before: "datetime") -> int: - """Delete events older than the given timestamp.""" - return self._delete_expired_events(before) - - def delete_idle_sessions(self, updated_before: "datetime") -> int: - """Delete sessions whose update_time predates the given threshold.""" - return self._delete_idle_sessions(updated_before) - - def get_app_state(self, app_name: str) -> "dict[str, Any] | None": - """Return app-scoped state for an application.""" - return self._get_app_state(app_name) - - def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": - """Return user-scoped state for an application user.""" - return self._get_user_state(app_name, user_id) - - def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: - """Insert or replace app-scoped state for an application.""" - self._upsert_app_state(app_name, state) - - def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: - """Insert or replace user-scoped state for an application user.""" - self._upsert_user_state(app_name, user_id, state) - - def get_metadata(self, key: str) -> "str | None": - """Return a value from the ADK internal metadata table.""" - return self._get_metadata(key) - - def set_metadata(self, key: str, value: str) -> None: - """Set a value in the ADK internal metadata table.""" - self._set_metadata(key, value) - - def _sessions_table_ddl(self) -> str: - adk_config = _adk_config(self._config) - owner_id_line = "" - if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" - session_locality = _cockroach_table_locality_clause(adk_config, "session_table_locality") - hash_shard_clause = _cockroach_hash_shard_clause(adk_config) - session_storing_clause = _cockroach_storing_clause(adk_config, ("state", "create_time", "update_time")) - - 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 - ){session_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user - ON {self._session_table}(app_name, user_id){session_storing_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time - ON {self._session_table}(update_time DESC){hash_shard_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state - ON {self._session_table} USING GIN (state) - WHERE state != '{{}}'::jsonb; - """ - - def _events_table_ddl(self) -> str: - adk_config = _adk_config(self._config) - events_locality = _cockroach_table_locality_clause(adk_config, "events_table_locality") - hash_shard_clause = _cockroach_hash_shard_clause(adk_config) - events_storing_clause = _cockroach_storing_clause(adk_config, ("invocation_id", "event_data")) - - return f""" - CREATE TABLE IF NOT EXISTS {self._events_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - invocation_id VARCHAR(256), - timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - event_data JSONB NOT NULL, - FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE - ){events_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session - ON {self._events_table}(session_id, timestamp ASC){hash_shard_clause}{events_storing_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_event_data - ON {self._events_table} USING GIN (event_data); - """ - - def _app_states_table_ddl(self) -> str: - adk_config = _adk_config(self._config) - app_state_locality = _cockroach_table_locality_clause(adk_config, "app_state_table_locality") - - return f""" - CREATE TABLE IF NOT EXISTS {self._app_state_table} ( - app_name VARCHAR(128) PRIMARY KEY, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ){app_state_locality}; - """ - - def _user_states_table_ddl(self) -> str: - adk_config = _adk_config(self._config) - user_state_locality = _cockroach_table_locality_clause(adk_config, "user_state_table_locality") - - return f""" - CREATE TABLE IF NOT EXISTS {self._user_state_table} ( - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY (app_name, user_id) - ){user_state_locality}; - """ - - def _metadata_table_ddl(self) -> str: - adk_config = _adk_config(self._config) - metadata_locality = _cockroach_table_locality_clause(adk_config, "metadata_table_locality") - - return f""" - CREATE TABLE IF NOT EXISTS {self._metadata_table} ( - key VARCHAR(128) PRIMARY KEY, - value VARCHAR(512) NOT NULL - ){metadata_locality}; - """ - - def _metadata_seed_sql(self) -> str: - return f""" - INSERT INTO {self._metadata_table} (key, value) - VALUES ('schema_version', '1') - ON CONFLICT (key) DO NOTHING - """ - - def _drop_app_states_table_sql(self) -> str: - return f"DROP TABLE IF EXISTS {self._app_state_table}" - - def _drop_user_states_table_sql(self) -> str: - return f"DROP TABLE IF EXISTS {self._user_state_table}" - - def _drop_metadata_table_sql(self) -> str: - return f"DROP TABLE IF EXISTS {self._metadata_table}" - - def _drop_tables_sql(self) -> "list[str]": - return [ - self._drop_metadata_table_sql(), - self._drop_user_states_table_sql(), - self._drop_app_states_table_sql(), - f"DROP TABLE IF EXISTS {self._events_table}", - f"DROP TABLE IF EXISTS {self._session_table}", - ] - - def _create_tables(self) -> None: with self._config.provide_session() as driver: driver.execute_script(self._sessions_table_ddl()) driver.execute_script(self._events_table_ddl()) @@ -771,9 +566,10 @@ def _create_tables(self) -> None: driver.execute_script(self._metadata_table_ddl()) driver.execute_script(self._metadata_seed_sql()) - def _create_session( + 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.""" state_json = Jsonb(state) params: tuple[Any, ...] if self._owner_id_column_name: @@ -793,15 +589,16 @@ def _create_session( cur.execute(sql.encode(), params) conn.commit() - result = self._get_session(app_name, user_id, session_id) + result = self.get_session(app_name, user_id, session_id) if result is None: msg = "Session creation failed" raise RuntimeError(msg) return result - def _get_session( - self, app_name: str, user_id: str, session_id: str, renew_for: "int | timedelta | None" = None + def get_session( + self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None ) -> "SessionRecord | None": + """Get session by ID.""" if renew_for is not None and self._calculate_expires_at(renew_for) is not None: sql = f""" UPDATE {self._session_table} @@ -835,7 +632,8 @@ def _get_session( except errors.UndefinedTable: return None - def _update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: + def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: + """Update session state.""" sql = f""" UPDATE {self._session_table} SET state = %s, update_time = CURRENT_TIMESTAMP @@ -846,14 +644,8 @@ def _update_session_state(self, app_name: str, user_id: str, session_id: str, st cur.execute(sql.encode(), (Jsonb(state), app_name, user_id, session_id)) conn.commit() - def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: - sql = f"DELETE FROM {self._session_table} WHERE app_name = %s AND user_id = %s AND id = %s" - - with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: - cur.execute(sql.encode(), (app_name, user_id, session_id)) - conn.commit() - - def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + """List sessions for an app.""" if user_id is None: sql = f""" SELECT id, app_name, user_id, state, create_time, update_time @@ -890,29 +682,20 @@ def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Ses except errors.UndefinedTable: return [] - def _insert_event(self, event_record: EventRecord) -> None: - sql = f""" - INSERT INTO {self._events_table} ( - id, session_id, invocation_id, timestamp, event_data - ) VALUES (%s, %s, %s, %s, %s) - """ - event_data_value = event_record["event_data"] - jsonb_value = Jsonb(event_data_value) if isinstance(event_data_value, dict) else event_data_value + def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: + """Delete session and associated events.""" + sql = f"DELETE FROM {self._session_table} WHERE app_name = %s AND user_id = %s AND id = %s" with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: - cur.execute( - sql.encode(), - ( - event_record["id"], - event_record["session_id"], - event_record["invocation_id"], - event_record["timestamp"], - jsonb_value, - ), - ) + cur.execute(sql.encode(), (app_name, user_id, session_id)) conn.commit() - def _append_event_and_update_state( + def append_event(self, event_record: EventRecord) -> None: + """Append an event to a session.""" + """Synchronous implementation of append_event.""" + self._insert_event(event_record) + + def append_event_and_update_state( self, event_record: EventRecord, app_name: str, @@ -923,6 +706,7 @@ def _append_event_and_update_state( app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, ) -> SessionRecord: + """Atomically append an event and update session + scoped state.""" insert_sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -979,7 +763,7 @@ def _append_event_and_update_state( update_time=row["update_time"], ) - def _get_events( + def get_events( self, app_name: str, user_id: str, @@ -987,6 +771,7 @@ def _get_events( after_timestamp: "datetime | None" = None, limit: "int | None" = None, ) -> "list[EventRecord]": + """Get events for a session.""" if limit == 0: return [] @@ -1030,7 +815,8 @@ def _get_events( except errors.UndefinedTable: return [] - def _delete_expired_events(self, before: "datetime") -> int: + def delete_expired_events(self, before: "datetime") -> int: + """Delete events older than the given timestamp.""" sql = f"DELETE FROM {self._events_table} WHERE timestamp < %s" try: @@ -1041,7 +827,8 @@ def _delete_expired_events(self, before: "datetime") -> int: except errors.UndefinedTable: return 0 - def _delete_idle_sessions(self, updated_before: "datetime") -> int: + def delete_idle_sessions(self, updated_before: "datetime") -> int: + """Delete sessions whose update_time predates the given threshold.""" sql = f"DELETE FROM {self._session_table} WHERE update_time < %s" try: @@ -1052,7 +839,8 @@ def _delete_idle_sessions(self, updated_before: "datetime") -> int: except errors.UndefinedTable: return 0 - def _get_app_state(self, app_name: str) -> "dict[str, Any] | None": + def get_app_state(self, app_name: str) -> "dict[str, Any] | None": + """Return app-scoped state for an application.""" sql = f"SELECT state FROM {self._app_state_table} WHERE app_name = %s" try: @@ -1063,7 +851,8 @@ def _get_app_state(self, app_name: str) -> "dict[str, Any] | None": except errors.UndefinedTable: return None - def _get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": + def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": + """Return user-scoped state for an application user.""" sql = f"SELECT state FROM {self._user_state_table} WHERE app_name = %s AND user_id = %s" try: @@ -1074,7 +863,8 @@ def _get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None except errors.UndefinedTable: return None - def _upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: + def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: + """Insert or replace app-scoped state for an application.""" sql = f""" UPSERT INTO {self._app_state_table} (app_name, state, update_time) VALUES (%s, %s, CURRENT_TIMESTAMP) @@ -1084,7 +874,8 @@ def _upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: cur.execute(sql.encode(), (app_name, Jsonb(state))) conn.commit() - def _upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: + def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: + """Insert or replace user-scoped state for an application user.""" sql = f""" UPSERT INTO {self._user_state_table} (app_name, user_id, state, update_time) VALUES (%s, %s, %s, CURRENT_TIMESTAMP) @@ -1094,7 +885,8 @@ def _upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any] cur.execute(sql.encode(), (app_name, user_id, Jsonb(state))) conn.commit() - def _get_metadata(self, key: str) -> "str | None": + def get_metadata(self, key: str) -> "str | None": + """Return a value from the ADK internal metadata table.""" sql = f"SELECT value FROM {self._metadata_table} WHERE key = %s" try: @@ -1105,7 +897,8 @@ def _get_metadata(self, key: str) -> "str | None": except errors.UndefinedTable: return None - def _set_metadata(self, key: str, value: str) -> None: + def set_metadata(self, key: str, value: str) -> None: + """Set a value in the ADK internal metadata table.""" sql = f""" UPSERT INTO {self._metadata_table} (key, value) VALUES (%s, %s) @@ -1115,9 +908,142 @@ def _set_metadata(self, key: str, value: str) -> None: cur.execute(sql.encode(), (key, value)) conn.commit() - def _append_event(self, event_record: EventRecord) -> None: - """Synchronous implementation of append_event.""" - self._insert_event(event_record) + def _sessions_table_ddl(self) -> str: + adk_config = _adk_config(self._config) + owner_id_line = "" + if self._owner_id_column_ddl: + owner_id_line = f",\n {self._owner_id_column_ddl}" + session_locality = _cockroach_table_locality_clause(adk_config, "session_table_locality") + hash_shard_clause = _cockroach_hash_shard_clause(adk_config) + session_storing_clause = _cockroach_storing_clause(adk_config, ("state", "create_time", "update_time")) + + 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 + ){session_locality}; + + CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user + ON {self._session_table}(app_name, user_id){session_storing_clause}; + + CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time + ON {self._session_table}(update_time DESC){hash_shard_clause}; + + CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state + ON {self._session_table} USING GIN (state) + WHERE state != '{{}}'::jsonb; + """ + + def _events_table_ddl(self) -> str: + adk_config = _adk_config(self._config) + events_locality = _cockroach_table_locality_clause(adk_config, "events_table_locality") + hash_shard_clause = _cockroach_hash_shard_clause(adk_config) + events_storing_clause = _cockroach_storing_clause(adk_config, ("invocation_id", "event_data")) + + return f""" + CREATE TABLE IF NOT EXISTS {self._events_table} ( + id VARCHAR(128) PRIMARY KEY, + session_id VARCHAR(128) NOT NULL, + invocation_id VARCHAR(256), + timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + event_data JSONB NOT NULL, + FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE + ){events_locality}; + + CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session + ON {self._events_table}(session_id, timestamp ASC){hash_shard_clause}{events_storing_clause}; + + CREATE INDEX IF NOT EXISTS idx_{self._events_table}_event_data + ON {self._events_table} USING GIN (event_data); + """ + + def _app_states_table_ddl(self) -> str: + adk_config = _adk_config(self._config) + app_state_locality = _cockroach_table_locality_clause(adk_config, "app_state_table_locality") + + return f""" + CREATE TABLE IF NOT EXISTS {self._app_state_table} ( + app_name VARCHAR(128) PRIMARY KEY, + state JSONB NOT NULL DEFAULT '{{}}'::jsonb, + update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP + ){app_state_locality}; + """ + + def _user_states_table_ddl(self) -> str: + adk_config = _adk_config(self._config) + user_state_locality = _cockroach_table_locality_clause(adk_config, "user_state_table_locality") + + return f""" + CREATE TABLE IF NOT EXISTS {self._user_state_table} ( + app_name VARCHAR(128) NOT NULL, + user_id VARCHAR(128) NOT NULL, + state JSONB NOT NULL DEFAULT '{{}}'::jsonb, + update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (app_name, user_id) + ){user_state_locality}; + """ + + def _metadata_table_ddl(self) -> str: + adk_config = _adk_config(self._config) + metadata_locality = _cockroach_table_locality_clause(adk_config, "metadata_table_locality") + + return f""" + CREATE TABLE IF NOT EXISTS {self._metadata_table} ( + key VARCHAR(128) PRIMARY KEY, + value VARCHAR(512) NOT NULL + ){metadata_locality}; + """ + + def _metadata_seed_sql(self) -> str: + return f""" + INSERT INTO {self._metadata_table} (key, value) + VALUES ('schema_version', '1') + ON CONFLICT (key) DO NOTHING + """ + + def _drop_app_states_table_sql(self) -> str: + return f"DROP TABLE IF EXISTS {self._app_state_table}" + + def _drop_user_states_table_sql(self) -> str: + return f"DROP TABLE IF EXISTS {self._user_state_table}" + + def _drop_metadata_table_sql(self) -> str: + return f"DROP TABLE IF EXISTS {self._metadata_table}" + + def _drop_tables_sql(self) -> "list[str]": + return [ + self._drop_metadata_table_sql(), + self._drop_user_states_table_sql(), + self._drop_app_states_table_sql(), + f"DROP TABLE IF EXISTS {self._events_table}", + f"DROP TABLE IF EXISTS {self._session_table}", + ] + + def _insert_event(self, event_record: EventRecord) -> None: + sql = f""" + INSERT INTO {self._events_table} ( + id, session_id, invocation_id, timestamp, event_data + ) VALUES (%s, %s, %s, %s, %s) + """ + event_data_value = event_record["event_data"] + jsonb_value = Jsonb(event_data_value) if isinstance(event_data_value, dict) else event_data_value + + with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: + cur.execute( + sql.encode(), + ( + event_record["id"], + event_record["session_id"], + event_record["invocation_id"], + event_record["timestamp"], + jsonb_value, + ), + ) + conn.commit() class CockroachPsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["CockroachPsycopgAsyncConfig"]): @@ -1305,82 +1231,14 @@ def __init__(self, config: "CockroachPsycopgSyncConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() - - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: - """Bulk insert memory entries with deduplication.""" - return self._insert_memory_entries(entries, owner_id) - - def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": - """Search memory entries by text query.""" - return self._search_entries(query, app_name, user_id, limit) - - def delete_entries_by_session(self, session_id: str) -> int: - """Delete all memory entries for a specific session.""" - return self._delete_entries_by_session(session_id) - - def delete_entries_older_than(self, days: int) -> int: - """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) - - def _memory_table_ddl(self) -> str: - adk_config = _adk_config(self._config) - owner_id_line = "" - if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" - memory_locality = _cockroach_table_locality_clause(adk_config, "memory_table_locality") - hash_shard_clause = _cockroach_hash_shard_clause(adk_config) - - fts_index = "" - if self._use_fts: - fts_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts - ON {self._memory_table} USING GIN (to_tsvector('english', content_text)); - """ - trigram_index = "" - if adk_config.get("enable_memory_trigram_index", False): - trigram_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_content_trgm - ON {self._memory_table} USING GIN (content_text gin_trgm_ops); - """ - - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMPTZ NOT NULL, - content_json JSONB NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSONB, - inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ){memory_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC){hash_shard_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session - ON {self._memory_table}(session_id); - {fts_index} - {trigram_index} - """ - - def _drop_memory_table_sql(self) -> "list[str]": - return [f"DROP TABLE IF EXISTS {self._memory_table}"] - - def _create_tables(self) -> None: if not self._enabled: return with self._config.provide_session() as driver: driver.execute_script(self._memory_table_ddl()) - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1423,9 +1281,10 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec inserted_count += cur.rowcount return inserted_count - def _search_entries( + def search_entries( self, query: str, app_name: str, user_id: str, limit: "int | None" = None ) -> "list[MemoryRecord]": + """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1464,7 +1323,8 @@ def _search_entries( return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] - def _delete_entries_by_session(self, session_id: str) -> int: + def delete_entries_by_session(self, session_id: str) -> int: + """Delete all memory entries for a specific session.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1475,7 +1335,8 @@ def _delete_entries_by_session(self, session_id: str) -> int: conn.commit() return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - def _delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int) -> int: + """Delete memory entries older than specified days.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1489,6 +1350,54 @@ def _delete_entries_older_than(self, days: int) -> int: conn.commit() return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 + def _memory_table_ddl(self) -> str: + adk_config = _adk_config(self._config) + owner_id_line = "" + if self._owner_id_column_ddl: + owner_id_line = f",\n {self._owner_id_column_ddl}" + memory_locality = _cockroach_table_locality_clause(adk_config, "memory_table_locality") + hash_shard_clause = _cockroach_hash_shard_clause(adk_config) + + fts_index = "" + if self._use_fts: + fts_index = f""" + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts + ON {self._memory_table} USING GIN (to_tsvector('english', content_text)); + """ + trigram_index = "" + if adk_config.get("enable_memory_trigram_index", False): + trigram_index = f""" + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_content_trgm + ON {self._memory_table} USING GIN (content_text gin_trgm_ops); + """ + + return f""" + CREATE TABLE IF NOT EXISTS {self._memory_table} ( + id VARCHAR(128) PRIMARY KEY, + session_id VARCHAR(128) NOT NULL, + app_name VARCHAR(128) NOT NULL, + user_id VARCHAR(128) NOT NULL, + event_id VARCHAR(128) NOT NULL UNIQUE, + author VARCHAR(256){owner_id_line}, + timestamp TIMESTAMPTZ NOT NULL, + content_json JSONB NOT NULL, + content_text TEXT NOT NULL, + metadata_json JSONB, + inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP + ){memory_locality}; + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time + ON {self._memory_table}(app_name, user_id, timestamp DESC){hash_shard_clause}; + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session + ON {self._memory_table}(session_id); + {fts_index} + {trigram_index} + """ + + def _drop_memory_table_sql(self) -> "list[str]": + return [f"DROP TABLE IF EXISTS {self._memory_table}"] + def _build_insert_params(entry: "MemoryRecord") -> "tuple[object, ...]": return ( diff --git a/sqlspec/adapters/mysqlconnector/adk/store.py b/sqlspec/adapters/mysqlconnector/adk/store.py index 1fe9646d8..46ef18acf 100644 --- a/sqlspec/adapters/mysqlconnector/adk/store.py +++ b/sqlspec/adapters/mysqlconnector/adk/store.py @@ -418,96 +418,6 @@ def __init__(self, config: "MysqlConnectorSyncConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() - - 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.""" - return self._create_session(session_id, app_name, user_id, state, owner_id) - - def get_session( - self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": - """Get session by ID.""" - return self._get_session(app_name, user_id, session_id, renew_for=renew_for) - - def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: - """Update session state.""" - self._update_session_state(app_name, user_id, session_id, state) - - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": - """List sessions for an app.""" - return self._list_sessions(app_name, user_id) - - def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: - """Delete session and associated events.""" - self._delete_session(app_name, user_id, session_id) - - def append_event(self, event_record: EventRecord) -> None: - """Append an event to a session.""" - self._append_event(event_record) - - def append_event_and_update_state( - self, - event_record: EventRecord, - app_name: str, - user_id: str, - session_id: str, - state: "dict[str, Any]", - *, - app_state: "dict[str, Any] | None" = None, - user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: - """Atomically append an event and update the session's durable state.""" - return self._append_event_and_update_state( - event_record, app_name, user_id, session_id, state, app_state=app_state, user_state=user_state - ) - - def get_events( - self, - app_name: str, - user_id: str, - session_id: str, - after_timestamp: "datetime | None" = None, - limit: "int | None" = None, - ) -> "list[EventRecord]": - """Get events for a session.""" - return self._get_events(app_name, user_id, session_id, after_timestamp, limit) - - def delete_expired_events(self, before: "datetime") -> int: - """Delete events older than the given timestamp.""" - return self._delete_expired_events(before) - - def delete_idle_sessions(self, updated_before: "datetime") -> int: - """Delete sessions whose update_time predates the threshold.""" - return self._delete_idle_sessions(updated_before) - - def get_app_state(self, app_name: str) -> "dict[str, Any] | None": - """Return app-scoped state for an application.""" - return self._get_app_state(app_name) - - def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": - """Return user-scoped state for an application user.""" - return self._get_user_state(app_name, user_id) - - def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: - """Insert or replace app-scoped state for an application.""" - self._upsert_app_state(app_name, state) - - def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: - """Insert or replace user-scoped state for an application user.""" - self._upsert_user_state(app_name, user_id, state) - - def get_metadata(self, key: str) -> "str | None": - """Return a value from the ADK internal metadata table.""" - return self._get_metadata(key) - - def set_metadata(self, key: str, value: str) -> None: - """Set a value in the ADK internal metadata table.""" - self._set_metadata(key, value) - - def _create_tables(self) -> None: with self._config.provide_session() as driver: driver.execute_script(self._sessions_table_ddl()) driver.execute_script(self._events_table_ddl()) @@ -516,9 +426,10 @@ def _create_tables(self) -> None: driver.execute_script(self._metadata_table_ddl()) driver.execute_script(self._metadata_seed_sql()) - def _create_session( + 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.""" params: tuple[Any, ...] if self._owner_id_column_name: sql = f""" @@ -541,15 +452,16 @@ def _create_session( cursor.close() conn.commit() - result = self._get_session(app_name, user_id, session_id) + result = self.get_session(app_name, user_id, session_id) if result is None: msg = "Failed to fetch created session" raise RuntimeError(msg) return result - def _get_session( + def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None ) -> "SessionRecord | None": + """Get session by ID.""" import mysql.connector try: @@ -584,7 +496,8 @@ def _get_session( return None raise - def _update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: + def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: + """Update session state.""" sql = f""" UPDATE {self._session_table} SET state = %s, update_time = UTC_TIMESTAMP(6) @@ -598,7 +511,8 @@ def _update_session_state(self, app_name: str, user_id: str, session_id: str, st cursor.close() conn.commit() - def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + """List sessions for an app.""" import mysql.connector if user_id is None: @@ -632,7 +546,8 @@ def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Ses return [] raise - def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: + def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: + """Delete session and associated events.""" sql = f"DELETE FROM {self._session_table} WHERE app_name = %s AND user_id = %s AND id = %s" with self._config.provide_connection() as conn: cursor = conn.cursor() @@ -642,7 +557,8 @@ def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: cursor.close() conn.commit() - def _append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: EventRecord) -> None: + """Append an event to a session.""" sql = f""" INSERT INTO {self._events_table} ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -656,7 +572,7 @@ def _append_event(self, event_record: EventRecord) -> None: cursor.close() conn.commit() - def _append_event_and_update_state( + def append_event_and_update_state( self, event_record: EventRecord, app_name: str, @@ -667,6 +583,7 @@ def _append_event_and_update_state( app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, ) -> SessionRecord: + """Atomically append an event and update the session's durable state.""" insert_sql = f""" INSERT INTO {self._events_table} ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -718,7 +635,7 @@ def _append_event_and_update_state( return _session_record_from_row(row) - def _get_events( + def get_events( self, app_name: str, user_id: str, @@ -726,6 +643,7 @@ def _get_events( after_timestamp: "datetime | None" = None, limit: "int | None" = None, ) -> "list[EventRecord]": + """Get events for a session.""" import mysql.connector if limit == 0: @@ -762,27 +680,34 @@ def _get_events( return [] raise - def _delete_expired_events(self, before: "datetime") -> int: + def delete_expired_events(self, before: "datetime") -> int: + """Delete events older than the given timestamp.""" return _sync_delete_before(self, self._events_table, "timestamp", before) - def _delete_idle_sessions(self, updated_before: "datetime") -> int: + def delete_idle_sessions(self, updated_before: "datetime") -> int: + """Delete sessions whose update_time predates the threshold.""" return _sync_delete_before(self, self._session_table, "update_time", updated_before) - def _get_app_state(self, app_name: str) -> "dict[str, Any] | None": + def get_app_state(self, app_name: str) -> "dict[str, Any] | None": + """Return app-scoped state for an application.""" return _sync_state(self, self._app_state_table, "app_name = %s", (app_name,)) - def _get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": + def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": + """Return user-scoped state for an application user.""" return _sync_state(self, self._user_state_table, "app_name = %s AND user_id = %s", (app_name, user_id)) - def _upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: + def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: + """Insert or replace app-scoped state for an application.""" _sync_execute_commit(self, _mysql_upsert_app_state_sql(self._app_state_table), (app_name, to_json(state))) - def _upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: + def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: + """Insert or replace user-scoped state for an application user.""" _sync_execute_commit( self, _mysql_upsert_user_state_sql(self._user_state_table), (app_name, user_id, to_json(state)) ) - def _get_metadata(self, key: str) -> "str | None": + def get_metadata(self, key: str) -> "str | None": + """Return a value from the ADK internal metadata table.""" import mysql.connector sql = f"SELECT value FROM {self._metadata_table} WHERE `key` = %s" @@ -800,7 +725,8 @@ def _get_metadata(self, key: str) -> "str | None": return None raise - def _set_metadata(self, key: str, value: str) -> None: + def set_metadata(self, key: str, value: str) -> None: + """Set a value in the ADK internal metadata table.""" _sync_execute_commit(self, _mysql_upsert_metadata_sql(self._metadata_table), (key, value)) def _sessions_table_ddl(self) -> str: @@ -1047,70 +973,14 @@ def __init__(self, config: "MysqlConnectorSyncConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() - - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: - """Bulk insert memory entries with deduplication.""" - return self._insert_memory_entries(entries, owner_id) - - def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": - """Search memory entries by text query.""" - return self._search_entries(query, app_name, user_id, limit) - - def delete_entries_by_session(self, session_id: str) -> int: - """Delete all memory entries for a specific session.""" - return self._delete_entries_by_session(session_id) - - def delete_entries_older_than(self, days: int) -> int: - """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) - - def _memory_table_ddl(self) -> str: - adk_config = _adk_config(self._config) - owner_id_line = "" - fk_constraint = "" - if self._owner_id_column_ddl: - col_def, fk_def = _mysql_owner_id_column_parts(self._owner_id_column_ddl) - owner_id_line = f",\n {col_def}" - if fk_def: - fk_constraint = f",\n {fk_def}" - - fts_index = "" - if self._use_fts: - fts_index = f",\n FULLTEXT INDEX idx_{self._memory_table}_fts (content_text)" - table_options = _mysql_table_options(adk_config, "memory_table_options") - - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - content_json JSON NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSON, - inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - INDEX idx_{self._memory_table}_app_user_time (app_name, user_id, timestamp), - INDEX idx_{self._memory_table}_session (session_id){fts_index}{fk_constraint} - ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} - """ - - def _drop_memory_table_sql(self) -> "list[str]": - return [f"DROP TABLE IF EXISTS {self._memory_table}"] - - def _create_tables(self) -> None: if not self._enabled: return with self._config.provide_session() as driver: driver.execute_script(self._memory_table_ddl()) - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1176,9 +1046,10 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec conn.commit() return inserted_count - def _search_entries( + def search_entries( self, query: str, app_name: str, user_id: str, limit: "int | None" = None ) -> "list[MemoryRecord]": + """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1216,7 +1087,8 @@ def _search_entries( return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] - def _delete_entries_by_session(self, session_id: str) -> int: + def delete_entries_by_session(self, session_id: str) -> int: + """Delete all memory entries for a specific session.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1231,7 +1103,8 @@ def _delete_entries_by_session(self, session_id: str) -> int: finally: cursor.close() - def _delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int) -> int: + """Delete memory entries older than specified days.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1249,6 +1122,42 @@ def _delete_entries_older_than(self, days: int) -> int: finally: cursor.close() + def _memory_table_ddl(self) -> str: + adk_config = _adk_config(self._config) + owner_id_line = "" + fk_constraint = "" + if self._owner_id_column_ddl: + col_def, fk_def = _mysql_owner_id_column_parts(self._owner_id_column_ddl) + owner_id_line = f",\n {col_def}" + if fk_def: + fk_constraint = f",\n {fk_def}" + + fts_index = "" + if self._use_fts: + fts_index = f",\n FULLTEXT INDEX idx_{self._memory_table}_fts (content_text)" + table_options = _mysql_table_options(adk_config, "memory_table_options") + + return f""" + CREATE TABLE IF NOT EXISTS {self._memory_table} ( + id VARCHAR(128) PRIMARY KEY, + session_id VARCHAR(128) NOT NULL, + app_name VARCHAR(128) NOT NULL, + user_id VARCHAR(128) NOT NULL, + event_id VARCHAR(128) NOT NULL UNIQUE, + author VARCHAR(256){owner_id_line}, + timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + content_json JSON NOT NULL, + content_text TEXT NOT NULL, + metadata_json JSON, + inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + INDEX idx_{self._memory_table}_app_user_time (app_name, user_id, timestamp), + INDEX idx_{self._memory_table}_session (session_id){fts_index}{fk_constraint} + ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} + """ + + def _drop_memory_table_sql(self) -> "list[str]": + return [f"DROP TABLE IF EXISTS {self._memory_table}"] + def _mysql_owner_id_column_parts(column_ddl: str) -> "tuple[str, str]": references_match = re.search(r"\s+REFERENCES\s+(.+)", column_ddl, re.IGNORECASE) diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index 72c6474ad..c088f51b7 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -1280,35 +1280,271 @@ def __init__(self, config: "OracleSyncConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() + """Create both sessions and events tables if they don't exist. + + Notes: + Detects Oracle version to determine optimal JSON storage type. + Uses version-appropriate table schema. + """ + storage_type = self._detect_json_storage_type() + logger.info("Creating ADK tables with storage type: %s", storage_type) + + with self._config.provide_session() as driver: + sessions_sql = SQL(self._sessions_table_ddl_for_type(storage_type)) + driver.execute_script(sessions_sql) + + events_sql = SQL(self._events_table_ddl_for_type(storage_type)) + driver.execute_script(events_sql) + driver.execute_script(SQL(self._app_states_table_ddl_for_type(storage_type))) + driver.execute_script(SQL(self._user_states_table_ddl_for_type(storage_type))) + driver.execute_script(SQL(self._metadata_table_ddl())) + driver.execute_script(SQL(self._metadata_seed_sql())) + driver.commit() 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.""" - return self._create_session(session_id, app_name, user_id, state, owner_id) + """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 SYSTIMESTAMP for create_time and update_time. + State is serialized using version-appropriate format. + owner_id is ignored if owner_id_column not configured. + """ + state_data = self._serialize_state(state) + + if self._owner_id_column_name: + sql = f""" + INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time, {self._owner_id_column_name}) + VALUES (:id, :app_name, :user_id, :state, SYSTIMESTAMP, SYSTIMESTAMP, :owner_id) + """ + params = { + "id": session_id, + "app_name": app_name, + "user_id": user_id, + "state": state_data, + "owner_id": owner_id, + } + else: + sql = f""" + INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time) + VALUES (:id, :app_name, :user_id, :state, SYSTIMESTAMP, SYSTIMESTAMP) + """ + params = {"id": session_id, "app_name": app_name, "user_id": user_id, "state": state_data} + + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, params) + conn.commit() + + result = self.get_session(app_name, user_id, session_id) + if result is None: + msg = "Failed to fetch created session" + raise RuntimeError(msg) + return result def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None ) -> "SessionRecord | None": """Get session by ID.""" - return self._get_session(app_name, user_id, session_id, renew_for=renew_for) + """Get session by ID. + + Args: + app_name: Application name. + user_id: User identifier. + session_id: Session identifier. + renew_for: If positive, touch update_time while reading. + + Returns: + Session record or None if not found. + + Notes: + Oracle returns datetime objects for TIMESTAMP columns. + State is deserialized using version-appropriate format. + """ + + sql = f""" + SELECT id, app_name, user_id, state, create_time, update_time + FROM {self._session_table} + WHERE app_name = :app_name AND user_id = :user_id AND id = :id + """ + + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + if renew_for is not None and self._calculate_expires_at(renew_for) is not None: + cursor.execute( + f"UPDATE {self._session_table} SET update_time = SYSTIMESTAMP WHERE app_name = :app_name AND user_id = :user_id AND id = :id", + {"app_name": app_name, "user_id": user_id, "id": session_id}, + ) + conn.commit() + + cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) + row = cursor.fetchone() + + if row is None: + return None + + session_id_val, app_name, user_id, state_data, create_time, update_time = row + + state = self._deserialize_state(state_data) + + return SessionRecord( + id=session_id_val, + app_name=app_name, + user_id=user_id, + state=state, + create_time=create_time, + update_time=update_time, + ) + except OracleDatabaseError as e: + error_obj = e.args[0] if e.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return None + raise def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: """Update session state.""" - self._update_session_state(app_name, user_id, session_id, state) + """Update session state. + + Args: + app_name: Application name. + user_id: User identifier. + session_id: Session identifier. + state: New state dictionary (replaces existing state). + + Notes: + This replaces the entire state dictionary. + Updates update_time to current timestamp. + State is serialized using version-appropriate format. + """ + state_data = self._serialize_state(state) + + sql = f""" + UPDATE {self._session_table} + SET state = :state, update_time = SYSTIMESTAMP + WHERE app_name = :app_name AND user_id = :user_id AND id = :id + """ + + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"state": state_data, "app_name": app_name, "user_id": user_id, "id": session_id}) + conn.commit() def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": """List sessions for an app.""" - return self._list_sessions(app_name, user_id) + """List sessions for an app, optionally filtered by user. + + Args: + app_name: Application name. + user_id: User identifier. If None, lists all sessions for the app. + + Returns: + List of session records ordered by update_time DESC. + + Notes: + Uses composite index on (app_name, user_id) when user_id is provided. + State is deserialized using version-appropriate format. + """ + + if user_id is None: + sql = f""" + SELECT id, app_name, user_id, state, create_time, update_time + FROM {self._session_table} + WHERE app_name = :app_name + ORDER BY update_time DESC + """ + params = {"app_name": app_name} + else: + sql = f""" + SELECT id, app_name, user_id, state, create_time, update_time + FROM {self._session_table} + WHERE app_name = :app_name AND user_id = :user_id + ORDER BY update_time DESC + """ + params = {"app_name": app_name, "user_id": user_id} + + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, params) + rows = cursor.fetchall() + + results = [] + for row in rows: + state = self._deserialize_state(row[3]) + + results.append( + SessionRecord( + id=row[0], + app_name=row[1], + user_id=row[2], + state=state, + create_time=row[4], + update_time=row[5], + ) + ) + return results + except OracleDatabaseError as e: + error_obj = e.args[0] if e.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return [] + raise def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: """Delete session and associated events.""" - self._delete_session(app_name, user_id, session_id) + """Delete session and all associated events (cascade). + + Args: + app_name: Application name. + user_id: User identifier. + session_id: Session identifier. + + Notes: + Foreign key constraint ensures events are cascade-deleted. + """ + sql = f"DELETE FROM {self._session_table} WHERE app_name = :app_name AND user_id = :user_id AND id = :id" + + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) + conn.commit() def append_event(self, event_record: EventRecord) -> None: """Append an event to a session.""" - self._append_event(event_record) + """Synchronous implementation of append_event.""" + sql = f""" + INSERT INTO {self._events_table} ( + id, session_id, invocation_id, timestamp, event_data + ) VALUES ( + :id, :session_id, :invocation_id, :timestamp, :event_data + ) + """ + + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute( + sql, + { + "id": event_record["id"], + "session_id": event_record["session_id"], + "invocation_id": event_record["invocation_id"], + "timestamp": event_record["timestamp"], + "event_data": self._serialize_event_data(event_record["event_data"]), + }, + ) + conn.commit() def append_event_and_update_state( self, @@ -1322,54 +1558,303 @@ def append_event_and_update_state( user_state: "dict[str, Any] | None" = None, ) -> SessionRecord: """Atomically append an event and update session + scoped state.""" - return self._append_event_and_update_state( - event_record, app_name, user_id, session_id, state, app_state=app_state, user_state=user_state + """Atomically create an event and update session + scoped state.""" + insert_sql = f""" + INSERT INTO {self._events_table} ( + id, session_id, invocation_id, timestamp, event_data + ) VALUES ( + :id, :session_id, :invocation_id, :timestamp, :event_data ) + """ - def get_events( - self, - app_name: str, - user_id: str, - session_id: str, - after_timestamp: "datetime | None" = None, - limit: "int | None" = None, + state_data = self._serialize_state(state) + update_sql = f""" + UPDATE {self._session_table} + SET state = :state, update_time = SYSTIMESTAMP + WHERE app_name = :app_name AND user_id = :user_id AND id = :id + """ + + select_sql = f""" + SELECT id, app_name, user_id, state, create_time, update_time + FROM {self._session_table} + WHERE app_name = :app_name AND user_id = :user_id AND id = :id + """ + + app_upsert_sql = f""" + MERGE INTO {self._app_state_table} target + USING (SELECT :app_name AS app_name, :state AS state FROM DUAL) source + ON (target.app_name = source.app_name) + WHEN MATCHED THEN + UPDATE SET target.state = source.state, target.update_time = SYSTIMESTAMP + WHEN NOT MATCHED THEN + INSERT (app_name, state, update_time) + VALUES (source.app_name, source.state, SYSTIMESTAMP) + """ + + user_upsert_sql = f""" + MERGE INTO {self._user_state_table} target + USING (SELECT :app_name AS app_name, :user_id AS user_id, :state AS state FROM DUAL) source + ON (target.app_name = source.app_name AND target.user_id = source.user_id) + WHEN MATCHED THEN + UPDATE SET target.state = source.state, target.update_time = SYSTIMESTAMP + WHEN NOT MATCHED THEN + INSERT (app_name, user_id, state, update_time) + VALUES (source.app_name, source.user_id, source.state, SYSTIMESTAMP) + """ + + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute( + update_sql, {"state": state_data, "app_name": app_name, "user_id": user_id, "id": session_id} + ) + cursor.execute(select_sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) + row = cursor.fetchone() + if row is None: + _raise_session_not_found(session_id) + cursor.execute( + insert_sql, + { + "id": event_record["id"], + "session_id": event_record["session_id"], + "invocation_id": event_record["invocation_id"], + "timestamp": event_record["timestamp"], + "event_data": self._serialize_event_data(event_record["event_data"]), + }, + ) + if app_state is not None: + cursor.execute(app_upsert_sql, {"app_name": app_name, "state": self._serialize_state(app_state)}) + if user_state is not None: + cursor.execute( + user_upsert_sql, + {"app_name": app_name, "user_id": user_id, "state": self._serialize_state(user_state)}, + ) + conn.commit() + except Exception: + conn.rollback() + raise + + session_id_val, row_app_name, row_user_id, state_data_row, create_time, update_time = row + return SessionRecord( + id=session_id_val, + app_name=row_app_name, + user_id=row_user_id, + state=self._deserialize_state(state_data_row), + create_time=create_time, + update_time=update_time, + ) + + def get_events( + self, + app_name: str, + user_id: str, + session_id: str, + after_timestamp: "datetime | None" = None, + limit: "int | None" = None, ) -> "list[EventRecord]": """Get events for a session.""" - return self._get_events(app_name, user_id, session_id, after_timestamp, limit) + """List events for a session ordered by timestamp. + + Args: + app_name: Application name. + user_id: User identifier. + 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. + """ + + if limit == 0: + return [] + + where_clauses = ["s.app_name = :app_name", "s.user_id = :user_id", "e.session_id = :session_id"] + params: dict[str, Any] = {"app_name": app_name, "user_id": user_id, "session_id": session_id} + + if after_timestamp is not None: + where_clauses.append("e.timestamp > :after_timestamp") + params["after_timestamp"] = after_timestamp + + where_clause = " AND ".join(where_clauses) + limit_clause = f" FETCH FIRST {limit} ROWS ONLY" if limit is not None else "" + sql = f""" + SELECT e.id, e.session_id, e.invocation_id, e.timestamp, e.event_data, s.app_name, s.user_id + FROM {self._events_table} e + JOIN {self._session_table} s ON e.session_id = s.id + WHERE {where_clause} + ORDER BY e.timestamp ASC{limit_clause} + """ + + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, params) + rows = cursor.fetchall() + + return [ + EventRecord( + id=row[0], + session_id=row[1], + invocation_id=_oracle_text_value(row[2]), + timestamp=row[3], + event_data=self._deserialize_json_field(row[4]) or {}, + app_name=row[5], + user_id=row[6], + ) + for row in rows + ] + except OracleDatabaseError as e: + error_obj = e.args[0] if e.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return [] + raise def delete_expired_events(self, before: "datetime") -> int: """Delete events older than the given timestamp.""" - return self._delete_expired_events(before) + sql = f"DELETE FROM {self._events_table} WHERE timestamp < :before" + + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"before": before}) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + except OracleDatabaseError as e: + error_obj = e.args[0] if e.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return 0 + raise def delete_idle_sessions(self, updated_before: "datetime") -> int: """Delete sessions whose update_time predates the given threshold.""" - return self._delete_idle_sessions(updated_before) + sql = f"DELETE FROM {self._session_table} WHERE update_time < :updated_before" + + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"updated_before": updated_before}) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + except OracleDatabaseError as e: + error_obj = e.args[0] if e.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return 0 + raise def get_app_state(self, app_name: str) -> "dict[str, Any] | None": """Return app-scoped state for an application.""" - return self._get_app_state(app_name) + """Synchronous implementation of get_app_state.""" + sql = f"SELECT state FROM {self._app_state_table} WHERE app_name = :app_name" + + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"app_name": app_name}) + row = cursor.fetchone() + return self._deserialize_state(row[0]) if row is not None else None + except OracleDatabaseError as e: + error_obj = e.args[0] if e.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return None + raise def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": """Return user-scoped state for an application user.""" - return self._get_user_state(app_name, user_id) + """Synchronous implementation of get_user_state.""" + sql = f""" + SELECT state + FROM {self._user_state_table} + WHERE app_name = :app_name AND user_id = :user_id + """ + + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"app_name": app_name, "user_id": user_id}) + row = cursor.fetchone() + return self._deserialize_state(row[0]) if row is not None else None + except OracleDatabaseError as e: + error_obj = e.args[0] if e.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return None + raise def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: """Insert or replace app-scoped state for an application.""" - self._upsert_app_state(app_name, state) + """Synchronous implementation of upsert_app_state.""" + sql = f""" + MERGE INTO {self._app_state_table} target + USING (SELECT :app_name AS app_name, :state AS state FROM DUAL) source + ON (target.app_name = source.app_name) + WHEN MATCHED THEN + UPDATE SET target.state = source.state, target.update_time = SYSTIMESTAMP + WHEN NOT MATCHED THEN + INSERT (app_name, state, update_time) + VALUES (source.app_name, source.state, SYSTIMESTAMP) + """ + + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"app_name": app_name, "state": self._serialize_state(state)}) + conn.commit() def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: """Insert or replace user-scoped state for an application user.""" - self._upsert_user_state(app_name, user_id, state) + """Synchronous implementation of upsert_user_state.""" + sql = f""" + MERGE INTO {self._user_state_table} target + USING (SELECT :app_name AS app_name, :user_id AS user_id, :state AS state FROM DUAL) source + ON (target.app_name = source.app_name AND target.user_id = source.user_id) + WHEN MATCHED THEN + UPDATE SET target.state = source.state, target.update_time = SYSTIMESTAMP + WHEN NOT MATCHED THEN + INSERT (app_name, user_id, state, update_time) + VALUES (source.app_name, source.user_id, source.state, SYSTIMESTAMP) + """ + + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "state": self._serialize_state(state)}) + conn.commit() def get_metadata(self, key: str) -> "str | None": """Return a value from the ADK internal metadata table.""" - return self._get_metadata(key) + """Synchronous implementation of get_metadata.""" + sql = f"SELECT value FROM {self._metadata_table} WHERE key = :key" + + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"key": key}) + row = cursor.fetchone() + return str(row[0]) if row is not None else None + except OracleDatabaseError as e: + error_obj = e.args[0] if e.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return None + raise def set_metadata(self, key: str, value: str) -> None: """Set a value in the ADK internal metadata table.""" - self._set_metadata(key, value) + """Synchronous implementation of set_metadata.""" + sql = f""" + MERGE INTO {self._metadata_table} target + USING (SELECT :key AS key, :value AS value FROM DUAL) source + ON (target.key = source.key) + WHEN MATCHED THEN + UPDATE SET target.value = source.value + WHEN NOT MATCHED THEN + INSERT (key, value) + VALUES (source.key, source.value) + """ - def _sessions_table_ddl(self) -> str: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"key": key, "value": value}) + conn.commit() + + def _sessions_table_ddl(self) -> str: """Get Oracle CREATE TABLE SQL for sessions table. Auto-detects optimal JSON storage type based on Oracle version. @@ -1736,642 +2221,83 @@ def _drop_user_states_table_sql(self) -> str: END; """ - def _drop_metadata_table_sql(self) -> str: - return f""" - BEGIN - EXECUTE IMMEDIATE 'DROP TABLE {self._metadata_table}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -942 THEN - RAISE; - END IF; - END; - """ - - def _drop_tables_sql(self) -> "list[str]": - """Get Oracle DROP TABLE SQL statements. - - Returns: - List of SQL statements to drop tables and indexes. - - Notes: - Order matters: drop events table (child) before sessions (parent). - Oracle automatically drops indexes when dropping tables. - """ - return [ - self._drop_metadata_table_sql(), - self._drop_user_states_table_sql(), - self._drop_app_states_table_sql(), - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP INDEX idx_{self._events_table}_session'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -1418 THEN - RAISE; - END IF; - END; - """, - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP INDEX idx_{self._session_table}_update_time'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -1418 THEN - RAISE; - END IF; - END; - """, - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP INDEX idx_{self._session_table}_app_user'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -1418 THEN - RAISE; - END IF; - END; - """, - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP TABLE {self._events_table}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -942 THEN - RAISE; - END IF; - END; - """, - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP TABLE {self._session_table}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -942 THEN - RAISE; - END IF; - END; - """, - ] - - def _create_tables(self) -> None: - """Create both sessions and events tables if they don't exist. - - Notes: - Detects Oracle version to determine optimal JSON storage type. - Uses version-appropriate table schema. - """ - storage_type = self._detect_json_storage_type() - logger.info("Creating ADK tables with storage type: %s", storage_type) - - with self._config.provide_session() as driver: - sessions_sql = SQL(self._sessions_table_ddl_for_type(storage_type)) - driver.execute_script(sessions_sql) - - events_sql = SQL(self._events_table_ddl_for_type(storage_type)) - driver.execute_script(events_sql) - driver.execute_script(SQL(self._app_states_table_ddl_for_type(storage_type))) - driver.execute_script(SQL(self._user_states_table_ddl_for_type(storage_type))) - driver.execute_script(SQL(self._metadata_table_ddl())) - driver.execute_script(SQL(self._metadata_seed_sql())) - driver.commit() - - 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 SYSTIMESTAMP for create_time and update_time. - State is serialized using version-appropriate format. - owner_id is ignored if owner_id_column not configured. - """ - state_data = self._serialize_state(state) - - if self._owner_id_column_name: - sql = f""" - INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time, {self._owner_id_column_name}) - VALUES (:id, :app_name, :user_id, :state, SYSTIMESTAMP, SYSTIMESTAMP, :owner_id) - """ - params = { - "id": session_id, - "app_name": app_name, - "user_id": user_id, - "state": state_data, - "owner_id": owner_id, - } - else: - sql = f""" - INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time) - VALUES (:id, :app_name, :user_id, :state, SYSTIMESTAMP, SYSTIMESTAMP) - """ - params = {"id": session_id, "app_name": app_name, "user_id": user_id, "state": state_data} - - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, params) - conn.commit() - - result = self._get_session(app_name, user_id, session_id) - if result is None: - msg = "Failed to fetch created session" - raise RuntimeError(msg) - return result - - def _get_session( - self, app_name: str, user_id: str, session_id: str, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": - """Get session by ID. - - Args: - app_name: Application name. - user_id: User identifier. - session_id: Session identifier. - renew_for: If positive, touch update_time while reading. - - Returns: - Session record or None if not found. - - Notes: - Oracle returns datetime objects for TIMESTAMP columns. - State is deserialized using version-appropriate format. - """ - - sql = f""" - SELECT id, app_name, user_id, state, create_time, update_time - FROM {self._session_table} - WHERE app_name = :app_name AND user_id = :user_id AND id = :id - """ - - try: - with self._config.provide_connection() as conn: - cursor = conn.cursor() - if renew_for is not None and self._calculate_expires_at(renew_for) is not None: - cursor.execute( - f"UPDATE {self._session_table} SET update_time = SYSTIMESTAMP WHERE app_name = :app_name AND user_id = :user_id AND id = :id", - {"app_name": app_name, "user_id": user_id, "id": session_id}, - ) - conn.commit() - - cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) - row = cursor.fetchone() - - if row is None: - return None - - session_id_val, app_name, user_id, state_data, create_time, update_time = row - - state = self._deserialize_state(state_data) - - return SessionRecord( - id=session_id_val, - app_name=app_name, - user_id=user_id, - state=state, - create_time=create_time, - update_time=update_time, - ) - except OracleDatabaseError as e: - error_obj = e.args[0] if e.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return None - raise - - def _update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: - """Update session state. - - Args: - app_name: Application name. - user_id: User identifier. - session_id: Session identifier. - state: New state dictionary (replaces existing state). - - Notes: - This replaces the entire state dictionary. - Updates update_time to current timestamp. - State is serialized using version-appropriate format. - """ - state_data = self._serialize_state(state) - - sql = f""" - UPDATE {self._session_table} - SET state = :state, update_time = SYSTIMESTAMP - WHERE app_name = :app_name AND user_id = :user_id AND id = :id - """ - - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"state": state_data, "app_name": app_name, "user_id": user_id, "id": session_id}) - conn.commit() - - def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": - """List sessions for an app, optionally filtered by user. - - Args: - app_name: Application name. - user_id: User identifier. If None, lists all sessions for the app. - - Returns: - List of session records ordered by update_time DESC. - - Notes: - Uses composite index on (app_name, user_id) when user_id is provided. - State is deserialized using version-appropriate format. - """ - - if user_id is None: - sql = f""" - SELECT id, app_name, user_id, state, create_time, update_time - FROM {self._session_table} - WHERE app_name = :app_name - ORDER BY update_time DESC - """ - params = {"app_name": app_name} - else: - sql = f""" - SELECT id, app_name, user_id, state, create_time, update_time - FROM {self._session_table} - WHERE app_name = :app_name AND user_id = :user_id - ORDER BY update_time DESC - """ - params = {"app_name": app_name, "user_id": user_id} - - try: - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, params) - rows = cursor.fetchall() - - results = [] - for row in rows: - state = self._deserialize_state(row[3]) - - results.append( - SessionRecord( - id=row[0], - app_name=row[1], - user_id=row[2], - state=state, - create_time=row[4], - update_time=row[5], - ) - ) - return results - except OracleDatabaseError as e: - error_obj = e.args[0] if e.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return [] - raise - - def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: - """Delete session and all associated events (cascade). - - Args: - app_name: Application name. - user_id: User identifier. - session_id: Session identifier. - - Notes: - Foreign key constraint ensures events are cascade-deleted. - """ - sql = f"DELETE FROM {self._session_table} WHERE app_name = :app_name AND user_id = :user_id AND id = :id" - - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) - conn.commit() - - def _append_event(self, event_record: EventRecord) -> None: - """Synchronous implementation of append_event.""" - sql = f""" - INSERT INTO {self._events_table} ( - id, session_id, invocation_id, timestamp, event_data - ) VALUES ( - :id, :session_id, :invocation_id, :timestamp, :event_data - ) - """ - - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute( - sql, - { - "id": event_record["id"], - "session_id": event_record["session_id"], - "invocation_id": event_record["invocation_id"], - "timestamp": event_record["timestamp"], - "event_data": self._serialize_event_data(event_record["event_data"]), - }, - ) - conn.commit() - - def _append_event_and_update_state( - self, - event_record: EventRecord, - app_name: str, - user_id: str, - session_id: str, - state: "dict[str, Any]", - *, - app_state: "dict[str, Any] | None" = None, - user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: - """Atomically create an event and update session + scoped state.""" - insert_sql = f""" - INSERT INTO {self._events_table} ( - id, session_id, invocation_id, timestamp, event_data - ) VALUES ( - :id, :session_id, :invocation_id, :timestamp, :event_data - ) - """ - - state_data = self._serialize_state(state) - update_sql = f""" - UPDATE {self._session_table} - SET state = :state, update_time = SYSTIMESTAMP - WHERE app_name = :app_name AND user_id = :user_id AND id = :id - """ - - select_sql = f""" - SELECT id, app_name, user_id, state, create_time, update_time - FROM {self._session_table} - WHERE app_name = :app_name AND user_id = :user_id AND id = :id - """ - - app_upsert_sql = f""" - MERGE INTO {self._app_state_table} target - USING (SELECT :app_name AS app_name, :state AS state FROM DUAL) source - ON (target.app_name = source.app_name) - WHEN MATCHED THEN - UPDATE SET target.state = source.state, target.update_time = SYSTIMESTAMP - WHEN NOT MATCHED THEN - INSERT (app_name, state, update_time) - VALUES (source.app_name, source.state, SYSTIMESTAMP) - """ - - user_upsert_sql = f""" - MERGE INTO {self._user_state_table} target - USING (SELECT :app_name AS app_name, :user_id AS user_id, :state AS state FROM DUAL) source - ON (target.app_name = source.app_name AND target.user_id = source.user_id) - WHEN MATCHED THEN - UPDATE SET target.state = source.state, target.update_time = SYSTIMESTAMP - WHEN NOT MATCHED THEN - INSERT (app_name, user_id, state, update_time) - VALUES (source.app_name, source.user_id, source.state, SYSTIMESTAMP) - """ - - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute( - update_sql, {"state": state_data, "app_name": app_name, "user_id": user_id, "id": session_id} - ) - cursor.execute(select_sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) - row = cursor.fetchone() - if row is None: - _raise_session_not_found(session_id) - cursor.execute( - insert_sql, - { - "id": event_record["id"], - "session_id": event_record["session_id"], - "invocation_id": event_record["invocation_id"], - "timestamp": event_record["timestamp"], - "event_data": self._serialize_event_data(event_record["event_data"]), - }, - ) - if app_state is not None: - cursor.execute(app_upsert_sql, {"app_name": app_name, "state": self._serialize_state(app_state)}) - if user_state is not None: - cursor.execute( - user_upsert_sql, - {"app_name": app_name, "user_id": user_id, "state": self._serialize_state(user_state)}, - ) - conn.commit() - except Exception: - conn.rollback() - raise - - session_id_val, row_app_name, row_user_id, state_data_row, create_time, update_time = row - return SessionRecord( - id=session_id_val, - app_name=row_app_name, - user_id=row_user_id, - state=self._deserialize_state(state_data_row), - create_time=create_time, - update_time=update_time, - ) - - def _get_events( - self, - app_name: str, - user_id: str, - session_id: str, - after_timestamp: "datetime | None" = None, - limit: "int | None" = None, - ) -> "list[EventRecord]": - """List events for a session ordered by timestamp. - - Args: - app_name: Application name. - user_id: User identifier. - 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. - """ - - if limit == 0: - return [] - - where_clauses = ["s.app_name = :app_name", "s.user_id = :user_id", "e.session_id = :session_id"] - params: dict[str, Any] = {"app_name": app_name, "user_id": user_id, "session_id": session_id} - - if after_timestamp is not None: - where_clauses.append("e.timestamp > :after_timestamp") - params["after_timestamp"] = after_timestamp - - where_clause = " AND ".join(where_clauses) - limit_clause = f" FETCH FIRST {limit} ROWS ONLY" if limit is not None else "" - sql = f""" - SELECT e.id, e.session_id, e.invocation_id, e.timestamp, e.event_data, s.app_name, s.user_id - FROM {self._events_table} e - JOIN {self._session_table} s ON e.session_id = s.id - WHERE {where_clause} - ORDER BY e.timestamp ASC{limit_clause} - """ - - try: - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, params) - rows = cursor.fetchall() - - return [ - EventRecord( - id=row[0], - session_id=row[1], - invocation_id=_oracle_text_value(row[2]), - timestamp=row[3], - event_data=self._deserialize_json_field(row[4]) or {}, - app_name=row[5], - user_id=row[6], - ) - for row in rows - ] - except OracleDatabaseError as e: - error_obj = e.args[0] if e.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return [] - raise - - def _delete_expired_events(self, before: "datetime") -> int: - sql = f"DELETE FROM {self._events_table} WHERE timestamp < :before" - - try: - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"before": before}) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - except OracleDatabaseError as e: - error_obj = e.args[0] if e.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return 0 - raise - - def _delete_idle_sessions(self, updated_before: "datetime") -> int: - sql = f"DELETE FROM {self._session_table} WHERE update_time < :updated_before" - - try: - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"updated_before": updated_before}) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - except OracleDatabaseError as e: - error_obj = e.args[0] if e.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return 0 - raise - - def _get_app_state(self, app_name: str) -> "dict[str, Any] | None": - """Synchronous implementation of get_app_state.""" - sql = f"SELECT state FROM {self._app_state_table} WHERE app_name = :app_name" - - try: - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"app_name": app_name}) - row = cursor.fetchone() - return self._deserialize_state(row[0]) if row is not None else None - except OracleDatabaseError as e: - error_obj = e.args[0] if e.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return None - raise - - def _get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": - """Synchronous implementation of get_user_state.""" - sql = f""" - SELECT state - FROM {self._user_state_table} - WHERE app_name = :app_name AND user_id = :user_id - """ - - try: - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"app_name": app_name, "user_id": user_id}) - row = cursor.fetchone() - return self._deserialize_state(row[0]) if row is not None else None - except OracleDatabaseError as e: - error_obj = e.args[0] if e.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return None - raise - - def _upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: - """Synchronous implementation of upsert_app_state.""" - sql = f""" - MERGE INTO {self._app_state_table} target - USING (SELECT :app_name AS app_name, :state AS state FROM DUAL) source - ON (target.app_name = source.app_name) - WHEN MATCHED THEN - UPDATE SET target.state = source.state, target.update_time = SYSTIMESTAMP - WHEN NOT MATCHED THEN - INSERT (app_name, state, update_time) - VALUES (source.app_name, source.state, SYSTIMESTAMP) - """ - - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"app_name": app_name, "state": self._serialize_state(state)}) - conn.commit() - - def _upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: - """Synchronous implementation of upsert_user_state.""" - sql = f""" - MERGE INTO {self._user_state_table} target - USING (SELECT :app_name AS app_name, :user_id AS user_id, :state AS state FROM DUAL) source - ON (target.app_name = source.app_name AND target.user_id = source.user_id) - WHEN MATCHED THEN - UPDATE SET target.state = source.state, target.update_time = SYSTIMESTAMP - WHEN NOT MATCHED THEN - INSERT (app_name, user_id, state, update_time) - VALUES (source.app_name, source.user_id, source.state, SYSTIMESTAMP) - """ - - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "state": self._serialize_state(state)}) - conn.commit() - - def _get_metadata(self, key: str) -> "str | None": - """Synchronous implementation of get_metadata.""" - sql = f"SELECT value FROM {self._metadata_table} WHERE key = :key" - - try: - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"key": key}) - row = cursor.fetchone() - return str(row[0]) if row is not None else None - except OracleDatabaseError as e: - error_obj = e.args[0] if e.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return None - raise - - def _set_metadata(self, key: str, value: str) -> None: - """Synchronous implementation of set_metadata.""" - sql = f""" - MERGE INTO {self._metadata_table} target - USING (SELECT :key AS key, :value AS value FROM DUAL) source - ON (target.key = source.key) - WHEN MATCHED THEN - UPDATE SET target.value = source.value - WHEN NOT MATCHED THEN - INSERT (key, value) - VALUES (source.key, source.value) + def _drop_metadata_table_sql(self) -> str: + return f""" + BEGIN + EXECUTE IMMEDIATE 'DROP TABLE {self._metadata_table}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -942 THEN + RAISE; + END IF; + END; """ - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"key": key, "value": value}) - conn.commit() + def _drop_tables_sql(self) -> "list[str]": + """Get Oracle DROP TABLE SQL statements. + + Returns: + List of SQL statements to drop tables and indexes. + + Notes: + Order matters: drop events table (child) before sessions (parent). + Oracle automatically drops indexes when dropping tables. + """ + return [ + self._drop_metadata_table_sql(), + self._drop_user_states_table_sql(), + self._drop_app_states_table_sql(), + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP INDEX idx_{self._events_table}_session'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -1418 THEN + RAISE; + END IF; + END; + """, + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP INDEX idx_{self._session_table}_update_time'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -1418 THEN + RAISE; + END IF; + END; + """, + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP INDEX idx_{self._session_table}_app_user'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -1418 THEN + RAISE; + END IF; + END; + """, + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP TABLE {self._events_table}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -942 THEN + RAISE; + END IF; + END; + """, + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP TABLE {self._session_table}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -942 THEN + RAISE; + END IF; + END; + """, + ] class OracleAsyncADKMemoryStore(BaseAsyncADKMemoryStore["OracleAsyncConfig"]): @@ -2733,25 +2659,100 @@ def __init__(self, config: "OracleSyncConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() + if not self._enabled: + return + + with self._config.provide_session() as driver: + driver.execute_script(self._memory_table_ddl()) def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" - return self._insert_memory_entries(entries, owner_id) + if not self._enabled: + msg = "Memory store is disabled" + raise RuntimeError(msg) + + if not entries: + return 0 + + owner_column = f", {self._owner_id_column_name}" if self._owner_id_column_name else "" + owner_param = ", :owner_id" if self._owner_id_column_name else "" + sql = f""" + INSERT INTO {self._memory_table} ( + id, session_id, app_name, user_id, event_id, author{owner_column}, + timestamp, content_json, content_text, metadata_json, inserted_at + ) VALUES ( + :id, :session_id, :app_name, :user_id, :event_id, :author{owner_param}, + :timestamp, :content_json, :content_text, :metadata_json, :inserted_at + ) + """ + + inserted_count = 0 + with self._config.provide_connection() as conn: + cursor = conn.cursor() + for entry in entries: + content_json = self._serialize_json_field(entry["content_json"]) + metadata_json = self._serialize_json_field(entry["metadata_json"]) + params = { + "id": entry["id"], + "session_id": entry["session_id"], + "app_name": entry["app_name"], + "user_id": entry["user_id"], + "event_id": entry["event_id"], + "author": entry["author"], + "timestamp": entry["timestamp"], + "content_json": content_json, + "content_text": entry["content_text"], + "metadata_json": metadata_json, + "inserted_at": entry["inserted_at"], + } + if self._owner_id_column_name: + params["owner_id"] = str(owner_id) if owner_id is not None else None + if self._execute_insert_entry(cursor, sql, params): + inserted_count += 1 + conn.commit() + + return inserted_count def search_entries( self, query: str, app_name: str, user_id: str, limit: "int | None" = None ) -> "list[MemoryRecord]": """Search memory entries by text query.""" - return self._search_entries(query, app_name, user_id, limit) + if not self._enabled: + msg = "Memory store is disabled" + raise RuntimeError(msg) + + effective_limit = limit if limit is not None else self._max_results + + try: + if self._use_fts: + return self._search_entries_fts(query, app_name, user_id, effective_limit) + return self._search_entries_simple(query, app_name, user_id, effective_limit) + except OracleDatabaseError as exc: + error_obj = exc.args[0] if exc.args else None + if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: + return [] + raise def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" - return self._delete_entries_by_session(session_id) + sql = f"DELETE FROM {self._memory_table} WHERE session_id = :session_id" + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"session_id": session_id}) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 def delete_entries_older_than(self, days: int) -> int: """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) + sql = f""" + DELETE FROM {self._memory_table} + WHERE inserted_at < SYSTIMESTAMP - NUMTODSINTERVAL(:days, 'DAY') + """ + with self._config.provide_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql, {"days": days}) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 def _detect_json_storage_type(self) -> "JSONStorageType": if self._json_storage_type is not None: @@ -2913,13 +2914,6 @@ def _drop_memory_table_sql(self) -> "list[str]": """, ] - def _create_tables(self) -> None: - if not self._enabled: - return - - with self._config.provide_session() as driver: - driver.execute_script(self._memory_table_ddl()) - def _execute_insert_entry(self, cursor: Any, sql: str, params: "dict[str, Any]") -> bool: """Execute an insert and skip duplicate key errors.""" try: @@ -2931,72 +2925,6 @@ def _execute_insert_entry(self, cursor: Any, sql: str, params: "dict[str, Any]") raise return True - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: - if not self._enabled: - msg = "Memory store is disabled" - raise RuntimeError(msg) - - if not entries: - return 0 - - owner_column = f", {self._owner_id_column_name}" if self._owner_id_column_name else "" - owner_param = ", :owner_id" if self._owner_id_column_name else "" - sql = f""" - INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author{owner_column}, - timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - :id, :session_id, :app_name, :user_id, :event_id, :author{owner_param}, - :timestamp, :content_json, :content_text, :metadata_json, :inserted_at - ) - """ - - inserted_count = 0 - with self._config.provide_connection() as conn: - cursor = conn.cursor() - for entry in entries: - content_json = self._serialize_json_field(entry["content_json"]) - metadata_json = self._serialize_json_field(entry["metadata_json"]) - params = { - "id": entry["id"], - "session_id": entry["session_id"], - "app_name": entry["app_name"], - "user_id": entry["user_id"], - "event_id": entry["event_id"], - "author": entry["author"], - "timestamp": entry["timestamp"], - "content_json": content_json, - "content_text": entry["content_text"], - "metadata_json": metadata_json, - "inserted_at": entry["inserted_at"], - } - if self._owner_id_column_name: - params["owner_id"] = str(owner_id) if owner_id is not None else None - if self._execute_insert_entry(cursor, sql, params): - inserted_count += 1 - conn.commit() - - return inserted_count - - def _search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": - if not self._enabled: - msg = "Memory store is disabled" - raise RuntimeError(msg) - - effective_limit = limit if limit is not None else self._max_results - - try: - if self._use_fts: - return self._search_entries_fts(query, app_name, user_id, effective_limit) - return self._search_entries_simple(query, app_name, user_id, effective_limit) - except OracleDatabaseError as exc: - error_obj = exc.args[0] if exc.args else None - if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: - return [] - raise - def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": sql = f""" SELECT id, session_id, app_name, user_id, event_id, author, @@ -3043,25 +2971,6 @@ def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: rows = cursor.fetchall() return self._rows_to_records(rows) - def _delete_entries_by_session(self, session_id: str) -> int: - sql = f"DELETE FROM {self._memory_table} WHERE session_id = :session_id" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"session_id": session_id}) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - - def _delete_entries_older_than(self, days: int) -> int: - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < SYSTIMESTAMP - NUMTODSINTERVAL(:days, 'DAY') - """ - with self._config.provide_connection() as conn: - cursor = conn.cursor() - cursor.execute(sql, {"days": days}) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]": records: list[MemoryRecord] = [] for row in rows: diff --git a/sqlspec/adapters/psycopg/adk/store.py b/sqlspec/adapters/psycopg/adk/store.py index 908b804e3..0173c14fb 100644 --- a/sqlspec/adapters/psycopg/adk/store.py +++ b/sqlspec/adapters/psycopg/adk/store.py @@ -566,196 +566,6 @@ def __init__(self, config: "PsycopgSyncConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() - - 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.""" - return self._create_session(session_id, app_name, user_id, state, owner_id) - - def get_session( - self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": - """Get session by ID.""" - return self._get_session(app_name, user_id, session_id, renew_for=renew_for) - - def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: - """Update session state.""" - self._update_session_state(app_name, user_id, session_id, state) - - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": - """List sessions for an app.""" - return self._list_sessions(app_name, user_id) - - def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: - """Delete session and associated events.""" - self._delete_session(app_name, user_id, session_id) - - def append_event(self, event_record: EventRecord) -> None: - """Append an event to a session.""" - self._append_event(event_record) - - def append_event_and_update_state( - self, - event_record: EventRecord, - app_name: str, - user_id: str, - session_id: str, - state: "dict[str, Any]", - *, - app_state: "dict[str, Any] | None" = None, - user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: - """Atomically append an event and update session + scoped state.""" - return self._append_event_and_update_state( - event_record, app_name, user_id, session_id, state, app_state=app_state, user_state=user_state - ) - - def get_events( - self, - app_name: str, - user_id: str, - session_id: str, - after_timestamp: "datetime | None" = None, - limit: "int | None" = None, - ) -> "list[EventRecord]": - """Get events for a session.""" - return self._get_events(app_name, user_id, session_id, after_timestamp, limit) - - def delete_expired_events(self, before: "datetime") -> int: - """Delete events older than the given timestamp.""" - return self._delete_expired_events(before) - - def delete_idle_sessions(self, updated_before: "datetime") -> int: - """Delete sessions whose update_time predates the given threshold.""" - return self._delete_idle_sessions(updated_before) - - def get_app_state(self, app_name: str) -> "dict[str, Any] | None": - """Return app-scoped state for an application.""" - return self._get_app_state(app_name) - - def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": - """Return user-scoped state for an application user.""" - return self._get_user_state(app_name, user_id) - - def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: - """Insert or replace app-scoped state for an application.""" - self._upsert_app_state(app_name, state) - - def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: - """Insert or replace user-scoped state for an application user.""" - self._upsert_user_state(app_name, user_id, state) - - def get_metadata(self, key: str) -> "str | None": - """Return a value from the ADK internal metadata table.""" - return self._get_metadata(key) - - def set_metadata(self, key: str, value: str) -> None: - """Set a value in the ADK internal metadata table.""" - self._set_metadata(key, value) - - def _sessions_table_ddl(self) -> str: - 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 _events_table_ddl(self) -> str: - adk_config = _adk_config(self._config) - generated_columns, generated_indexes, covering_columns = _postgres_event_ddl_options( - adk_config, self._events_table - ) - - return f""" - CREATE TABLE IF NOT EXISTS {self._events_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - invocation_id VARCHAR(256), - timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - event_data JSONB NOT NULL{generated_columns}, - FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE - ) WITH (fillfactor = 80); - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session - ON {self._events_table}(session_id, timestamp ASC){covering_columns}; - {generated_indexes} - """ - - def _app_states_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._app_state_table} ( - app_name VARCHAR(128) PRIMARY KEY, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = 80); - """ - - def _user_states_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._user_state_table} ( - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY (app_name, user_id) - ) WITH (fillfactor = 80); - """ - - def _metadata_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._metadata_table} ( - key VARCHAR(128) PRIMARY KEY, - value VARCHAR(512) NOT NULL - ); - """ - - def _metadata_seed_sql(self) -> str: - return f""" - INSERT INTO {self._metadata_table} (key, value) - VALUES ('schema_version', '1') - ON CONFLICT (key) DO NOTHING - """ - - def _drop_app_states_table_sql(self) -> str: - return f"DROP TABLE IF EXISTS {self._app_state_table}" - - def _drop_user_states_table_sql(self) -> str: - return f"DROP TABLE IF EXISTS {self._user_state_table}" - - def _drop_metadata_table_sql(self) -> str: - return f"DROP TABLE IF EXISTS {self._metadata_table}" - - def _drop_tables_sql(self) -> "list[str]": - return [ - self._drop_metadata_table_sql(), - self._drop_user_states_table_sql(), - self._drop_app_states_table_sql(), - f"DROP TABLE IF EXISTS {self._events_table}", - f"DROP TABLE IF EXISTS {self._session_table}", - ] - - def _create_tables(self) -> None: with self._config.provide_session() as driver: driver.execute_script(self._sessions_table_ddl()) driver.execute_script(self._events_table_ddl()) @@ -764,9 +574,10 @@ def _create_tables(self) -> None: driver.execute_script(self._metadata_table_ddl()) driver.execute_script(self._metadata_seed_sql()) - def _create_session( + 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.""" params: tuple[Any, ...] if self._owner_id_column_name: query = pg_sql.SQL(""" @@ -786,15 +597,16 @@ def _create_session( with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute(query, params) - result = self._get_session(app_name, user_id, session_id) + result = self.get_session(app_name, user_id, session_id) if result is None: msg = "Failed to fetch created session" raise RuntimeError(msg) return result - def _get_session( - self, app_name: str, user_id: str, session_id: str, renew_for: "int | timedelta | None" = None + def get_session( + self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None ) -> "SessionRecord | None": + """Get session by ID.""" if renew_for is not None and self._calculate_expires_at(renew_for) is not None: query = pg_sql.SQL(""" UPDATE {table} @@ -830,7 +642,8 @@ def _get_session( except errors.UndefinedTable: return None - def _update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: + def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: + """Update session state.""" query = pg_sql.SQL(""" UPDATE {table} SET state = %s, update_time = CURRENT_TIMESTAMP @@ -840,15 +653,8 @@ def _update_session_state(self, app_name: str, user_id: str, session_id: str, st with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute(query, (Jsonb(state), app_name, user_id, session_id)) - def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: - query = pg_sql.SQL("DELETE FROM {table} WHERE app_name = %s AND user_id = %s AND id = %s").format( - table=pg_sql.Identifier(self._session_table) - ) - - with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: - cur.execute(query, (app_name, user_id, session_id)) - - def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + """List sessions for an app.""" if user_id is None: query = pg_sql.SQL(""" SELECT id, app_name, user_id, state, create_time, update_time @@ -885,30 +691,21 @@ def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Ses except errors.UndefinedTable: return [] - def _insert_event(self, event_record: EventRecord) -> None: - insert_query = pg_sql.SQL(""" - INSERT INTO {table} ( - id, session_id, invocation_id, timestamp, event_data - ) VALUES (%s, %s, %s, %s, %s) - """).format(table=pg_sql.Identifier(self._events_table)) - - event_data_value = event_record["event_data"] - jsonb_value = Jsonb(event_data_value) if isinstance(event_data_value, dict) else event_data_value + def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: + """Delete session and associated events.""" + query = pg_sql.SQL("DELETE FROM {table} WHERE app_name = %s AND user_id = %s AND id = %s").format( + table=pg_sql.Identifier(self._session_table) + ) with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: - cur.execute( - insert_query, - ( - event_record["id"], - event_record["session_id"], - event_record["invocation_id"], - event_record["timestamp"], - jsonb_value, - ), - ) - conn.commit() + cur.execute(query, (app_name, user_id, session_id)) + + def append_event(self, event_record: EventRecord) -> None: + """Append an event to a session.""" + """Synchronous implementation of append_event.""" + self._insert_event(event_record) - def _append_event_and_update_state( + def append_event_and_update_state( self, event_record: EventRecord, app_name: str, @@ -919,6 +716,7 @@ def _append_event_and_update_state( app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, ) -> SessionRecord: + """Atomically append an event and update session + scoped state.""" insert_query = pg_sql.SQL(""" INSERT INTO {table} ( id, session_id, invocation_id, timestamp, event_data @@ -985,7 +783,7 @@ def _append_event_and_update_state( update_time=row["update_time"], ) - def _get_events( + def get_events( self, app_name: str, user_id: str, @@ -993,6 +791,7 @@ def _get_events( after_timestamp: "datetime | None" = None, limit: "int | None" = None, ) -> "list[EventRecord]": + """Get events for a session.""" if limit == 0: return [] @@ -1042,7 +841,8 @@ def _get_events( except errors.UndefinedTable: return [] - def _delete_expired_events(self, before: "datetime") -> int: + def delete_expired_events(self, before: "datetime") -> int: + """Delete events older than the given timestamp.""" query = pg_sql.SQL("DELETE FROM {table} WHERE timestamp < %s").format( table=pg_sql.Identifier(self._events_table) ) @@ -1055,7 +855,8 @@ def _delete_expired_events(self, before: "datetime") -> int: except errors.UndefinedTable: return 0 - def _delete_idle_sessions(self, updated_before: "datetime") -> int: + def delete_idle_sessions(self, updated_before: "datetime") -> int: + """Delete sessions whose update_time predates the given threshold.""" query = pg_sql.SQL("DELETE FROM {table} WHERE update_time < %s").format( table=pg_sql.Identifier(self._session_table) ) @@ -1068,7 +869,8 @@ def _delete_idle_sessions(self, updated_before: "datetime") -> int: except errors.UndefinedTable: return 0 - def _get_app_state(self, app_name: str) -> "dict[str, Any] | None": + def get_app_state(self, app_name: str) -> "dict[str, Any] | None": + """Return app-scoped state for an application.""" query = pg_sql.SQL("SELECT state FROM {table} WHERE app_name = %s").format( table=pg_sql.Identifier(self._app_state_table) ) @@ -1081,7 +883,8 @@ def _get_app_state(self, app_name: str) -> "dict[str, Any] | None": except errors.UndefinedTable: return None - def _get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": + def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": + """Return user-scoped state for an application user.""" query = pg_sql.SQL("SELECT state FROM {table} WHERE app_name = %s AND user_id = %s").format( table=pg_sql.Identifier(self._user_state_table) ) @@ -1094,7 +897,8 @@ def _get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None except errors.UndefinedTable: return None - def _upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: + def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: + """Insert or replace app-scoped state for an application.""" query = pg_sql.SQL(""" INSERT INTO {table} (app_name, state, update_time) VALUES (%s, %s, CURRENT_TIMESTAMP) @@ -1107,7 +911,8 @@ def _upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: cur.execute(query, (app_name, Jsonb(state))) conn.commit() - def _upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: + def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: + """Insert or replace user-scoped state for an application user.""" query = pg_sql.SQL(""" INSERT INTO {table} (app_name, user_id, state, update_time) VALUES (%s, %s, %s, CURRENT_TIMESTAMP) @@ -1120,7 +925,8 @@ def _upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any] cur.execute(query, (app_name, user_id, Jsonb(state))) conn.commit() - def _get_metadata(self, key: str) -> "str | None": + def get_metadata(self, key: str) -> "str | None": + """Return a value from the ADK internal metadata table.""" query = pg_sql.SQL("SELECT value FROM {table} WHERE key = %s").format( table=pg_sql.Identifier(self._metadata_table) ) @@ -1133,7 +939,8 @@ def _get_metadata(self, key: str) -> "str | None": except errors.UndefinedTable: return None - def _set_metadata(self, key: str, value: str) -> None: + def set_metadata(self, key: str, value: str) -> None: + """Set a value in the ADK internal metadata table.""" query = pg_sql.SQL(""" INSERT INTO {table} (key, value) VALUES (%s, %s) @@ -1144,9 +951,128 @@ def _set_metadata(self, key: str, value: str) -> None: cur.execute(query, (key, value)) conn.commit() - def _append_event(self, event_record: EventRecord) -> None: - """Synchronous implementation of append_event.""" - self._insert_event(event_record) + def _sessions_table_ddl(self) -> str: + 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 _events_table_ddl(self) -> str: + adk_config = _adk_config(self._config) + generated_columns, generated_indexes, covering_columns = _postgres_event_ddl_options( + adk_config, self._events_table + ) + + return f""" + CREATE TABLE IF NOT EXISTS {self._events_table} ( + id VARCHAR(128) PRIMARY KEY, + session_id VARCHAR(128) NOT NULL, + invocation_id VARCHAR(256), + timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + event_data JSONB NOT NULL{generated_columns}, + FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE + ) WITH (fillfactor = 80); + + CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session + ON {self._events_table}(session_id, timestamp ASC){covering_columns}; + {generated_indexes} + """ + + def _app_states_table_ddl(self) -> str: + return f""" + CREATE TABLE IF NOT EXISTS {self._app_state_table} ( + app_name VARCHAR(128) PRIMARY KEY, + state JSONB NOT NULL DEFAULT '{{}}'::jsonb, + update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP + ) WITH (fillfactor = 80); + """ + + def _user_states_table_ddl(self) -> str: + return f""" + CREATE TABLE IF NOT EXISTS {self._user_state_table} ( + app_name VARCHAR(128) NOT NULL, + user_id VARCHAR(128) NOT NULL, + state JSONB NOT NULL DEFAULT '{{}}'::jsonb, + update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (app_name, user_id) + ) WITH (fillfactor = 80); + """ + + def _metadata_table_ddl(self) -> str: + return f""" + CREATE TABLE IF NOT EXISTS {self._metadata_table} ( + key VARCHAR(128) PRIMARY KEY, + value VARCHAR(512) NOT NULL + ); + """ + + def _metadata_seed_sql(self) -> str: + return f""" + INSERT INTO {self._metadata_table} (key, value) + VALUES ('schema_version', '1') + ON CONFLICT (key) DO NOTHING + """ + + def _drop_app_states_table_sql(self) -> str: + return f"DROP TABLE IF EXISTS {self._app_state_table}" + + def _drop_user_states_table_sql(self) -> str: + return f"DROP TABLE IF EXISTS {self._user_state_table}" + + def _drop_metadata_table_sql(self) -> str: + return f"DROP TABLE IF EXISTS {self._metadata_table}" + + def _drop_tables_sql(self) -> "list[str]": + return [ + self._drop_metadata_table_sql(), + self._drop_user_states_table_sql(), + self._drop_app_states_table_sql(), + f"DROP TABLE IF EXISTS {self._events_table}", + f"DROP TABLE IF EXISTS {self._session_table}", + ] + + def _insert_event(self, event_record: EventRecord) -> None: + insert_query = pg_sql.SQL(""" + INSERT INTO {table} ( + id, session_id, invocation_id, timestamp, event_data + ) VALUES (%s, %s, %s, %s, %s) + """).format(table=pg_sql.Identifier(self._events_table)) + + event_data_value = event_record["event_data"] + jsonb_value = Jsonb(event_data_value) if isinstance(event_data_value, dict) else event_data_value + + with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: + cur.execute( + insert_query, + ( + event_record["id"], + event_record["session_id"], + event_record["invocation_id"], + event_record["timestamp"], + jsonb_value, + ), + ) + conn.commit() class PsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["PsycopgAsyncConfig"]): @@ -1346,67 +1272,6 @@ def __init__(self, config: "PsycopgSyncConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() - - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: - """Bulk insert memory entries with deduplication.""" - return self._insert_memory_entries(entries, owner_id) - - def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": - """Search memory entries by text query.""" - return self._search_entries(query, app_name, user_id, limit) - - def delete_entries_by_session(self, session_id: str) -> int: - """Delete all memory entries for a specific session.""" - return self._delete_entries_by_session(session_id) - - def delete_entries_older_than(self, days: int) -> int: - """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) - - def _memory_table_ddl(self) -> str: - """Get PostgreSQL CREATE TABLE SQL for memory entries.""" - owner_id_line = "" - if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" - - fts_index = "" - if self._use_fts: - fts_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts - ON {self._memory_table} USING GIN (to_tsvector('english', content_text)); - """ - - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMPTZ NOT NULL, - content_json JSONB NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSONB, - inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session - ON {self._memory_table}(session_id); - {fts_index} - """ - - def _drop_memory_table_sql(self) -> "list[str]": - """Get PostgreSQL DROP TABLE SQL statements.""" - return [f"DROP TABLE IF EXISTS {self._memory_table}"] - - def _create_tables(self) -> None: """Create the memory table and indexes if they don't exist.""" if not self._enabled: return @@ -1414,7 +1279,8 @@ def _create_tables(self) -> None: with self._config.provide_session() as driver: driver.execute_script(self._memory_table_ddl()) - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + """Bulk insert memory entries with deduplication.""" """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" @@ -1459,9 +1325,10 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec return inserted_count - def _search_entries( + def search_entries( self, query: str, app_name: str, user_id: str, limit: "int | None" = None ) -> "list[MemoryRecord]": + """Search memory entries by text query.""" """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -1479,6 +1346,71 @@ def _search_entries( except errors.UndefinedTable: return [] + def delete_entries_by_session(self, session_id: str) -> int: + """Delete all memory entries for a specific session.""" + """Delete all memory entries for a specific session.""" + sql = pg_sql.SQL("DELETE FROM {table} WHERE session_id = %s").format( + table=pg_sql.Identifier(self._memory_table) + ) + + with self._config.provide_connection() as conn, conn.cursor() as cur: + cur.execute(sql, (session_id,)) + return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 + + def delete_entries_older_than(self, days: int) -> int: + """Delete memory entries older than specified days.""" + """Delete memory entries older than specified days.""" + sql = pg_sql.SQL( + """ + DELETE FROM {table} + WHERE inserted_at < CURRENT_TIMESTAMP - {interval}::interval + """ + ).format(table=pg_sql.Identifier(self._memory_table), interval=pg_sql.Literal(f"{days} days")) + + with self._config.provide_connection() as conn, conn.cursor() as cur: + cur.execute(sql) + return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 + + def _memory_table_ddl(self) -> str: + """Get PostgreSQL CREATE TABLE SQL for memory entries.""" + owner_id_line = "" + if self._owner_id_column_ddl: + owner_id_line = f",\n {self._owner_id_column_ddl}" + + fts_index = "" + if self._use_fts: + fts_index = f""" + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts + ON {self._memory_table} USING GIN (to_tsvector('english', content_text)); + """ + + return f""" + CREATE TABLE IF NOT EXISTS {self._memory_table} ( + id VARCHAR(128) PRIMARY KEY, + session_id VARCHAR(128) NOT NULL, + app_name VARCHAR(128) NOT NULL, + user_id VARCHAR(128) NOT NULL, + event_id VARCHAR(128) NOT NULL UNIQUE, + author VARCHAR(256){owner_id_line}, + timestamp TIMESTAMPTZ NOT NULL, + content_json JSONB NOT NULL, + content_text TEXT NOT NULL, + metadata_json JSONB, + inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP + ); + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time + ON {self._memory_table}(app_name, user_id, timestamp DESC); + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session + ON {self._memory_table}(session_id); + {fts_index} + """ + + def _drop_memory_table_sql(self) -> "list[str]": + """Get PostgreSQL DROP TABLE SQL statements.""" + return [f"DROP TABLE IF EXISTS {self._memory_table}"] + def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": sql = pg_sql.SQL( """ @@ -1519,29 +1451,6 @@ def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: rows = cur.fetchall() return _rows_to_records(rows) - def _delete_entries_by_session(self, session_id: str) -> int: - """Delete all memory entries for a specific session.""" - sql = pg_sql.SQL("DELETE FROM {table} WHERE session_id = %s").format( - table=pg_sql.Identifier(self._memory_table) - ) - - with self._config.provide_connection() as conn, conn.cursor() as cur: - cur.execute(sql, (session_id,)) - return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - - def _delete_entries_older_than(self, days: int) -> int: - """Delete memory entries older than specified days.""" - sql = pg_sql.SQL( - """ - DELETE FROM {table} - WHERE inserted_at < CURRENT_TIMESTAMP - {interval}::interval - """ - ).format(table=pg_sql.Identifier(self._memory_table), interval=pg_sql.Literal(f"{days} days")) - - with self._config.provide_connection() as conn, conn.cursor() as cur: - cur.execute(sql) - return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - def _build_insert_params(entry: "MemoryRecord") -> "tuple[object, ...]": return ( diff --git a/sqlspec/adapters/sqlite/adk/store.py b/sqlspec/adapters/sqlite/adk/store.py index 4ea445cb0..c30e45c79 100644 --- a/sqlspec/adapters/sqlite/adk/store.py +++ b/sqlspec/adapters/sqlite/adk/store.py @@ -137,7 +137,15 @@ def __init__(self, config: "SqliteConfig") -> None: def create_tables(self) -> None: """Create both sessions and events tables if they don't exist.""" - self._create_tables() + """Synchronous implementation of create_tables.""" + with self._config.provide_session() as driver: + self._apply_pragmas(driver.connection) + driver.execute_script(self._sessions_table_ddl()) + driver.execute_script(self._events_table_ddl()) + driver.execute_script(self._app_states_table_ddl()) + driver.execute_script(self._user_states_table_ddl()) + driver.execute_script(self._metadata_table_ddl()) + driver.execute_script(self._metadata_seed_sql()) def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None @@ -154,176 +162,6 @@ def create_session( Returns: Created session record. """ - return self._create_session(session_id, app_name, user_id, state, owner_id) - - def get_session( - self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": - """Get session by ID. - - Args: - app_name: Application name. - user_id: User identifier. - session_id: Session identifier. - renew_for: If positive, touch the session update timestamp while reading. - - Returns: - Session record or None if not found. - """ - return self._get_session(app_name, user_id, session_id, renew_for=renew_for) - - def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: - """Update session state. - - Args: - app_name: Application name. - user_id: User identifier. - session_id: Session identifier. - state: New state dictionary (replaces existing state). - """ - self._update_session_state(app_name, user_id, session_id, state) - - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": - """List sessions for an app, optionally filtered by user. - - Args: - app_name: Application name. - user_id: User identifier. If None, lists all sessions for the app. - - Returns: - List of session records ordered by update_time DESC. - """ - return self._list_sessions(app_name, user_id) - - def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: - """Delete session and all associated events (cascade). - - Args: - app_name: Application name. - user_id: User identifier. - session_id: Session identifier. - """ - self._delete_session(app_name, user_id, session_id) - - def append_event(self, event_record: EventRecord) -> None: - """Append an event to a session. - - Args: - event_record: Event record to store. - """ - self._append_event(event_record) - - def append_event_and_update_state( - self, - event_record: EventRecord, - app_name: str, - user_id: str, - session_id: str, - state: "dict[str, Any]", - *, - app_state: "dict[str, Any] | None" = None, - user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: - """Atomically append an event and update the session's durable state. - - Inserts the event and updates the session state + update_time in a - single transaction, returning the updated SessionRecord via RETURNING. - - Args: - event_record: Event record to store. - app_name: Application name for scoped state. - user_id: User identifier for scoped state. - session_id: Session identifier whose state should be updated. - state: Post-append durable state snapshot (temp: keys already - stripped by the service layer). - app_state: App-scoped state snapshot to upsert when changed. - user_state: User-scoped state snapshot to upsert when changed. - """ - return self._append_event_and_update_state( - event_record, app_name, user_id, session_id, state, app_state=app_state, user_state=user_state - ) - - def get_events( - self, - app_name: str, - user_id: str, - session_id: str, - after_timestamp: "datetime | None" = None, - limit: "int | None" = None, - ) -> "list[EventRecord]": - """Get events for a session. - - Args: - app_name: Application name. - user_id: User identifier. - 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. - """ - return self._get_events(app_name, user_id, session_id, after_timestamp, limit) - - def delete_expired_events(self, before: datetime) -> int: - """Delete events older than the given timestamp.""" - return self._delete_expired_events(before) - - def delete_idle_sessions(self, updated_before: datetime) -> int: - """Delete sessions whose update_time predates the given threshold.""" - return self._delete_idle_sessions(updated_before) - - def get_app_state(self, app_name: str) -> "dict[str, Any] | None": - """Return app-scoped state for an application.""" - return self._get_app_state(app_name) - - def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": - """Return user-scoped state for an application user.""" - return self._get_user_state(app_name, user_id) - - def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: - """Insert or replace app-scoped state for an application.""" - self._upsert_app_state(app_name, state) - - def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: - """Insert or replace user-scoped state for an application user.""" - self._upsert_user_state(app_name, user_id, state) - - def get_metadata(self, key: str) -> "str | None": - """Return a value from the ADK internal metadata table.""" - return self._get_metadata(key) - - def set_metadata(self, key: str, value: str) -> None: - """Set a value in the ADK internal metadata table.""" - self._set_metadata(key, value) - - def _apply_pragmas(self, connection: Any) -> None: - """Apply PRAGMA optimization profile for this connection. - - Args: - connection: SQLite connection. - """ - connection.execute("PRAGMA foreign_keys = ON") - connection.execute("PRAGMA cache_size = -64000") - connection.execute("PRAGMA mmap_size = 30000000") - connection.execute("PRAGMA journal_size_limit = 67108864") - for pragma_name, pragma_value in self._pragma_overrides: - connection.execute(f"PRAGMA {pragma_name} = {pragma_value}") - - def _create_tables(self) -> None: - """Synchronous implementation of create_tables.""" - with self._config.provide_session() as driver: - self._apply_pragmas(driver.connection) - driver.execute_script(self._sessions_table_ddl()) - driver.execute_script(self._events_table_ddl()) - driver.execute_script(self._app_states_table_ddl()) - driver.execute_script(self._user_states_table_ddl()) - driver.execute_script(self._metadata_table_ddl()) - driver.execute_script(self._metadata_seed_sql()) - - def _create_session( - self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: """Synchronous implementation of create_session.""" now = datetime.now(timezone.utc) now_julian = _datetime_to_julian(now) @@ -353,9 +191,20 @@ def _create_session( id=session_id, app_name=app_name, user_id=user_id, state=state, create_time=now, update_time=now ) - def _get_session( - self, app_name: str, user_id: str, session_id: str, renew_for: "int | timedelta | None" = None + def get_session( + self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None ) -> "SessionRecord | None": + """Get session by ID. + + Args: + app_name: Application name. + user_id: User identifier. + session_id: Session identifier. + renew_for: If positive, touch the session update timestamp while reading. + + Returns: + Session record or None if not found. + """ """Synchronous implementation of get_session.""" params = (app_name, user_id, session_id) update_params: tuple[Any, ...] @@ -402,7 +251,15 @@ def _get_session( return None raise - def _update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: + def update_session_state(self, app_name: str, user_id: str, session_id: str, state: "dict[str, Any]") -> None: + """Update session state. + + Args: + app_name: Application name. + user_id: User identifier. + session_id: Session identifier. + state: New state dictionary (replaces existing state). + """ """Synchronous implementation of update_session_state.""" now_julian = _datetime_to_julian(datetime.now(timezone.utc)) state_json = to_json(state) @@ -418,7 +275,16 @@ def _update_session_state(self, app_name: str, user_id: str, session_id: str, st conn.execute(sql, (state_json, now_julian, app_name, user_id, session_id)) conn.commit() - def _list_sessions(self, app_name: str, user_id: "str | None") -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + """List sessions for an app, optionally filtered by user. + + Args: + app_name: Application name. + user_id: User identifier. If None, lists all sessions for the app. + + Returns: + List of session records ordered by update_time DESC. + """ """Synchronous implementation of list_sessions.""" if user_id is None: sql = f""" @@ -459,7 +325,14 @@ def _list_sessions(self, app_name: str, user_id: "str | None") -> "list[SessionR return [] raise - def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: + def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: + """Delete session and all associated events (cascade). + + Args: + app_name: Application name. + user_id: User identifier. + session_id: Session identifier. + """ """Synchronous implementation of delete_session.""" sql = f"DELETE FROM {self._session_table} WHERE app_name = ? AND user_id = ? AND id = ?" @@ -468,7 +341,12 @@ def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: conn.execute(sql, (app_name, user_id, session_id)) conn.commit() - def _append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: EventRecord) -> None: + """Append an event to a session. + + Args: + event_record: Event record to store. + """ """Synchronous implementation of append_event.""" timestamp_julian = _datetime_to_julian(event_record["timestamp"]) event_data_json = to_json(event_record["event_data"]) @@ -495,7 +373,7 @@ def _append_event(self, event_record: EventRecord) -> None: ) conn.commit() - def _append_event_and_update_state( + def append_event_and_update_state( self, event_record: EventRecord, app_name: str, @@ -506,6 +384,21 @@ def _append_event_and_update_state( app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, ) -> SessionRecord: + """Atomically append an event and update the session's durable state. + + Inserts the event and updates the session state + update_time in a + single transaction, returning the updated SessionRecord via RETURNING. + + Args: + event_record: Event record to store. + app_name: Application name for scoped state. + user_id: User identifier for scoped state. + session_id: Session identifier whose state should be updated. + state: Post-append durable state snapshot (temp: keys already + stripped by the service layer). + app_state: App-scoped state snapshot to upsert when changed. + user_state: User-scoped state snapshot to upsert when changed. + """ """Synchronous implementation of append_event_and_update_state.""" timestamp_julian = _datetime_to_julian(event_record["timestamp"]) event_data_json = to_json(event_record["event_data"]) @@ -585,7 +478,7 @@ def _append_event_and_update_state( update_time=_julian_to_datetime(row[5]), ) - def _get_events( + def get_events( self, app_name: str, user_id: str, @@ -593,6 +486,18 @@ def _get_events( after_timestamp: "datetime | None" = None, limit: "int | None" = None, ) -> "list[EventRecord]": + """Get events for a session. + + Args: + app_name: Application name. + user_id: User identifier. + 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. + """ """Synchronous implementation of get_events.""" if limit == 0: return [] @@ -637,7 +542,8 @@ def _get_events( return [] raise - def _delete_expired_events(self, before: datetime) -> int: + def delete_expired_events(self, before: datetime) -> int: + """Delete events older than the given timestamp.""" """Synchronous implementation of delete_expired_events.""" sql = f"DELETE FROM {self._events_table} WHERE timestamp < ?" @@ -653,7 +559,8 @@ def _delete_expired_events(self, before: datetime) -> int: return 0 raise - def _delete_idle_sessions(self, updated_before: datetime) -> int: + def delete_idle_sessions(self, updated_before: datetime) -> int: + """Delete sessions whose update_time predates the given threshold.""" """Synchronous implementation of delete_idle_sessions.""" sql = f"DELETE FROM {self._session_table} WHERE update_time < ?" @@ -669,7 +576,8 @@ def _delete_idle_sessions(self, updated_before: datetime) -> int: return 0 raise - def _get_app_state(self, app_name: str) -> "dict[str, Any] | None": + def get_app_state(self, app_name: str) -> "dict[str, Any] | None": + """Return app-scoped state for an application.""" """Synchronous implementation of get_app_state.""" sql = f"SELECT state FROM {self._app_state_table} WHERE app_name = ?" @@ -684,7 +592,8 @@ def _get_app_state(self, app_name: str) -> "dict[str, Any] | None": return None raise - def _get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": + def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None": + """Return user-scoped state for an application user.""" """Synchronous implementation of get_user_state.""" sql = f""" SELECT state @@ -703,7 +612,8 @@ def _get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None return None raise - def _upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: + def upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: + """Insert or replace app-scoped state for an application.""" """Synchronous implementation of upsert_app_state.""" sql = f""" INSERT INTO {self._app_state_table} (app_name, state, update_time) @@ -718,7 +628,8 @@ def _upsert_app_state(self, app_name: str, state: "dict[str, Any]") -> None: conn.execute(sql, (app_name, to_json(state), _datetime_to_julian(datetime.now(timezone.utc)))) conn.commit() - def _upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: + def upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any]") -> None: + """Insert or replace user-scoped state for an application user.""" """Synchronous implementation of upsert_user_state.""" sql = f""" INSERT INTO {self._user_state_table} (app_name, user_id, state, update_time) @@ -733,7 +644,8 @@ def _upsert_user_state(self, app_name: str, user_id: str, state: "dict[str, Any] conn.execute(sql, (app_name, user_id, to_json(state), _datetime_to_julian(datetime.now(timezone.utc)))) conn.commit() - def _get_metadata(self, key: str) -> "str | None": + def get_metadata(self, key: str) -> "str | None": + """Return a value from the ADK internal metadata table.""" """Synchronous implementation of get_metadata.""" sql = f"SELECT value FROM {self._metadata_table} WHERE key = ?" @@ -748,7 +660,8 @@ def _get_metadata(self, key: str) -> "str | None": return None raise - def _set_metadata(self, key: str, value: str) -> None: + def set_metadata(self, key: str, value: str) -> None: + """Set a value in the ADK internal metadata table.""" """Synchronous implementation of set_metadata.""" sql = f""" INSERT INTO {self._metadata_table} (key, value) @@ -761,10 +674,23 @@ def _set_metadata(self, key: str, value: str) -> None: conn.execute(sql, (key, value)) conn.commit() - def _sessions_table_ddl(self) -> str: - """Get SQLite CREATE TABLE SQL for sessions. + def _apply_pragmas(self, connection: Any) -> None: + """Apply PRAGMA optimization profile for this connection. - Returns: + Args: + connection: SQLite connection. + """ + connection.execute("PRAGMA foreign_keys = ON") + connection.execute("PRAGMA cache_size = -64000") + connection.execute("PRAGMA mmap_size = 30000000") + connection.execute("PRAGMA journal_size_limit = 67108864") + for pragma_name, pragma_value in self._pragma_overrides: + connection.execute(f"PRAGMA {pragma_name} = {pragma_value}") + + def _sessions_table_ddl(self) -> str: + """Get SQLite CREATE TABLE SQL for sessions. + + Returns: SQL statement to create adk_session table with indexes. """ owner_id_line = "" @@ -898,105 +824,6 @@ def __init__(self, config: "SqliteConfig") -> None: def create_tables(self) -> None: """Create tables if they don't exist.""" - self._create_tables() - - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: - """Bulk insert memory entries with deduplication.""" - return self._insert_memory_entries(entries, owner_id) - - def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": - """Search memory entries by text query.""" - return self._search_entries(query, app_name, user_id, limit) - - def delete_entries_by_session(self, session_id: str) -> int: - """Delete all memory entries for a specific session.""" - return self._delete_entries_by_session(session_id) - - def delete_entries_older_than(self, days: int) -> int: - """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) - - def _memory_table_ddl(self) -> str: - """Get SQLite CREATE TABLE SQL for memory entries. - - Returns: - SQL statement to create memory table with indexes. - """ - owner_id_line = "" - if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" - - fts_table = "" - if self._use_fts: - fts_options = _format_fts_options(self._fts_options) - fts_table = f""" - CREATE VIRTUAL TABLE IF NOT EXISTS {self._memory_table}_fts USING fts5( - content_text, - content={self._memory_table}, - content_rowid=rowid{fts_options} - ); - - CREATE TRIGGER IF NOT EXISTS {self._memory_table}_ai AFTER INSERT ON {self._memory_table} BEGIN - INSERT INTO {self._memory_table}_fts(rowid, content_text) VALUES (new.rowid, new.content_text); - END; - - CREATE TRIGGER IF NOT EXISTS {self._memory_table}_ad AFTER DELETE ON {self._memory_table} BEGIN - INSERT INTO {self._memory_table}_fts({self._memory_table}_fts, rowid, content_text) - VALUES('delete', old.rowid, old.content_text); - END; - - CREATE TRIGGER IF NOT EXISTS {self._memory_table}_au AFTER UPDATE ON {self._memory_table} BEGIN - INSERT INTO {self._memory_table}_fts({self._memory_table}_fts, rowid, content_text) - VALUES('delete', old.rowid, old.content_text); - INSERT INTO {self._memory_table}_fts(rowid, content_text) VALUES (new.rowid, new.content_text); - END; - """ - - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - app_name TEXT NOT NULL, - user_id TEXT NOT NULL, - event_id TEXT NOT NULL UNIQUE, - author TEXT{owner_id_line}, - timestamp REAL NOT NULL, - content_json TEXT NOT NULL, - content_text TEXT NOT NULL, - metadata_json TEXT, - inserted_at REAL NOT NULL - ); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session - ON {self._memory_table}(session_id); - {fts_table} - """ - - def _drop_memory_table_sql(self) -> "list[str]": - """Get SQLite DROP TABLE SQL statements. - - Returns: - List of SQL statements to drop the memory table and FTS table. - """ - statements = [f"DROP TABLE IF EXISTS {self._memory_table}"] - if self._use_fts: - statements.insert(0, f"DROP TABLE IF EXISTS {self._memory_table}_fts") - return statements - - def _enable_foreign_keys(self, connection: Any) -> None: - """Enable foreign key constraints for this connection. - - Args: - connection: SQLite connection. - """ - connection.execute("PRAGMA foreign_keys = ON") - - def _create_tables(self) -> None: """Create the memory table and indexes if they don't exist. Skips table creation if memory store is disabled. @@ -1008,7 +835,8 @@ def _create_tables(self) -> None: self._enable_foreign_keys(driver.connection) driver.execute_script(self._memory_table_ddl()) - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + """Bulk insert memory entries with deduplication.""" """Bulk insert memory entries with deduplication. Uses INSERT OR IGNORE to skip duplicates based on event_id @@ -1092,9 +920,10 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec return inserted_count - def _search_entries( + def search_entries( self, query: str, app_name: str, user_id: str, limit: "int | None" = None ) -> "list[MemoryRecord]": + """Search memory entries by text query.""" """Search memory entries by text query. Args: @@ -1122,6 +951,128 @@ def _search_entries( logger.warning("FTS search failed; falling back to simple search: %s", exc) return self._search_entries_simple(query, app_name, user_id, effective_limit) + def delete_entries_by_session(self, session_id: str) -> int: + """Delete all memory entries for a specific session.""" + """Delete all memory entries for a specific session. + + Args: + session_id: Session ID to delete entries for. + + Returns: + Number of entries deleted. + """ + sql = f"DELETE FROM {self._memory_table} WHERE session_id = ?" + + with self._config.provide_connection() as conn: + self._enable_foreign_keys(conn) + cursor = conn.execute(sql, (session_id,)) + deleted_count = cursor.rowcount + conn.commit() + + return deleted_count + + def delete_entries_older_than(self, days: int) -> int: + """Delete memory entries older than specified days.""" + """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. + """ + cutoff_julian = _datetime_to_julian(datetime.now(timezone.utc)) - days + + sql = f"DELETE FROM {self._memory_table} WHERE inserted_at < ?" + + with self._config.provide_connection() as conn: + self._enable_foreign_keys(conn) + cursor = conn.execute(sql, (cutoff_julian,)) + deleted_count = cursor.rowcount + conn.commit() + + return deleted_count + + def _memory_table_ddl(self) -> str: + """Get SQLite CREATE TABLE SQL for memory entries. + + Returns: + SQL statement to create memory table with indexes. + """ + owner_id_line = "" + if self._owner_id_column_ddl: + owner_id_line = f",\n {self._owner_id_column_ddl}" + + fts_table = "" + if self._use_fts: + fts_options = _format_fts_options(self._fts_options) + fts_table = f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {self._memory_table}_fts USING fts5( + content_text, + content={self._memory_table}, + content_rowid=rowid{fts_options} + ); + + CREATE TRIGGER IF NOT EXISTS {self._memory_table}_ai AFTER INSERT ON {self._memory_table} BEGIN + INSERT INTO {self._memory_table}_fts(rowid, content_text) VALUES (new.rowid, new.content_text); + END; + + CREATE TRIGGER IF NOT EXISTS {self._memory_table}_ad AFTER DELETE ON {self._memory_table} BEGIN + INSERT INTO {self._memory_table}_fts({self._memory_table}_fts, rowid, content_text) + VALUES('delete', old.rowid, old.content_text); + END; + + CREATE TRIGGER IF NOT EXISTS {self._memory_table}_au AFTER UPDATE ON {self._memory_table} BEGIN + INSERT INTO {self._memory_table}_fts({self._memory_table}_fts, rowid, content_text) + VALUES('delete', old.rowid, old.content_text); + INSERT INTO {self._memory_table}_fts(rowid, content_text) VALUES (new.rowid, new.content_text); + END; + """ + + return f""" + CREATE TABLE IF NOT EXISTS {self._memory_table} ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + app_name TEXT NOT NULL, + user_id TEXT NOT NULL, + event_id TEXT NOT NULL UNIQUE, + author TEXT{owner_id_line}, + timestamp REAL NOT NULL, + content_json TEXT NOT NULL, + content_text TEXT NOT NULL, + metadata_json TEXT, + inserted_at REAL NOT NULL + ); + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time + ON {self._memory_table}(app_name, user_id, timestamp DESC); + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session + ON {self._memory_table}(session_id); + {fts_table} + """ + + def _drop_memory_table_sql(self) -> "list[str]": + """Get SQLite DROP TABLE SQL statements. + + Returns: + List of SQL statements to drop the memory table and FTS table. + """ + statements = [f"DROP TABLE IF EXISTS {self._memory_table}"] + if self._use_fts: + statements.insert(0, f"DROP TABLE IF EXISTS {self._memory_table}_fts") + return statements + + def _enable_foreign_keys(self, connection: Any) -> None: + """Enable foreign key constraints for this connection. + + Args: + connection: SQLite connection. + """ + connection.execute("PRAGMA foreign_keys = ON") + def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": sql = f""" SELECT m.id, m.session_id, m.app_name, m.user_id, m.event_id, m.author, @@ -1174,48 +1125,6 @@ def _fetch_records(self, sql: str, params: "tuple[Any, ...]") -> "list[MemoryRec for row in rows ] - 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. - """ - sql = f"DELETE FROM {self._memory_table} WHERE session_id = ?" - - with self._config.provide_connection() as conn: - self._enable_foreign_keys(conn) - cursor = conn.execute(sql, (session_id,)) - deleted_count = cursor.rowcount - conn.commit() - - return deleted_count - - 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. - """ - cutoff_julian = _datetime_to_julian(datetime.now(timezone.utc)) - days - - sql = f"DELETE FROM {self._memory_table} WHERE inserted_at < ?" - - with self._config.provide_connection() as conn: - self._enable_foreign_keys(conn) - cursor = conn.execute(sql, (cutoff_julian,)) - deleted_count = cursor.rowcount - conn.commit() - - return deleted_count - def _datetime_to_julian(dt: datetime) -> float: """Convert datetime to Julian Day number for SQLite storage. diff --git a/tests/unit/adapters/test_psycopg/test_adk_store.py b/tests/unit/adapters/test_psycopg/test_adk_store.py index 6dae92130..4ddd8bfe3 100644 --- a/tests/unit/adapters/test_psycopg/test_adk_store.py +++ b/tests/unit/adapters/test_psycopg/test_adk_store.py @@ -164,7 +164,7 @@ def test_sync_append_event_inserts_without_session_update() -> None: "event_data": {"id": "event-1"}, } - store._append_event(event_record) # type: ignore[arg-type] + store.append_event(event_record) # type: ignore[arg-type] assert len(cursor.execute_calls) == 1 _, params = cursor.execute_calls[0] @@ -190,7 +190,7 @@ def test_sync_get_events_passes_after_timestamp_and_limit() -> None: ] store, cursor, _ = _build_store(rows) - result = store._get_events("app", "user", "session-1", after_timestamp=base_time, limit=1) + result = store.get_events("app", "user", "session-1", after_timestamp=base_time, limit=1) assert len(cursor.execute_calls) == 1 _, params = cursor.execute_calls[0] @@ -202,7 +202,7 @@ def test_sync_get_events_limit_zero_returns_empty_without_query() -> None: """get_events(limit=0) must return no events without querying.""" store, cursor, _ = _build_store() - result = store._get_events("app", "user", "session-1", limit=0) + result = store.get_events("app", "user", "session-1", limit=0) assert result == [] assert cursor.execute_calls == [] @@ -232,7 +232,7 @@ def test_sync_append_event_and_update_state_writes_scoped_state_in_one_unit() -> "event_data": {"id": "event-1"}, } - result = store._append_event_and_update_state( + result = store.append_event_and_update_state( event_record, # type: ignore[arg-type] "app", "user", diff --git a/tests/unit/extensions/test_adk/test_store_instantiation.py b/tests/unit/extensions/test_adk/test_store_instantiation.py index cf6ca6ece..58aa57965 100644 --- a/tests/unit/extensions/test_adk/test_store_instantiation.py +++ b/tests/unit/extensions/test_adk/test_store_instantiation.py @@ -9,6 +9,7 @@ when a concrete store no longer satisfies the base contract. """ +import ast import importlib import inspect from typing import cast @@ -168,6 +169,43 @@ def test_adk_store_registration_validator_resolves_sqlite_store_classes() -> Non _ensure_adk_store_registration(config) +@pytest.mark.parametrize( + "class_path", + [ + "sqlspec.adapters.psycopg.adk.PsycopgSyncADKStore", + "sqlspec.adapters.psycopg.adk.PsycopgSyncADKMemoryStore", + "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKStore", + "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKMemoryStore", + "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKStore", + "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKMemoryStore", + "sqlspec.adapters.oracledb.adk.OracleSyncADKStore", + "sqlspec.adapters.oracledb.adk.OracleSyncADKMemoryStore", + "sqlspec.adapters.sqlite.adk.SqliteADKStore", + "sqlspec.adapters.sqlite.adk.SqliteADKMemoryStore", + ], +) +def test_sync_store_public_methods_do_not_delegate_to_private_mirrors(class_path: str) -> None: + cls = _load_class(class_path) + tree = ast.parse(inspect.getsource(cls)) + + has_delegate = False + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef) or node.name.startswith("_"): + continue + for child in ast.walk(node): + if ( + isinstance(child, ast.Call) + and isinstance(child.func, ast.Attribute) + and isinstance(child.func.value, ast.Name) + and child.func.value.id == "self" + and child.func.attr.startswith("_") + and child.func.attr[1:] == node.name + ): + has_delegate = True + + assert not has_delegate + + def test_adk_store_registration_validator_resolves_duckdb_store_classes() -> None: """The migration registration validator handles DuckDB store export casing.""" from sqlspec.adapters.duckdb import DuckDBConfig From 4117813d53742df2f43eecbe53828001f1ef6791 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 11 Jul 2026 00:55:13 +0000 Subject: [PATCH 3/6] refactor(adk): share base store state --- sqlspec/extensions/adk/artifact/store.py | 55 +- sqlspec/extensions/adk/memory/store.py | 235 +++------ sqlspec/extensions/adk/store.py | 486 +++++++----------- .../test_adk/test_store_common_mixins.py | 30 ++ 4 files changed, 296 insertions(+), 510 deletions(-) create mode 100644 tests/unit/extensions/test_adk/test_store_common_mixins.py diff --git a/sqlspec/extensions/adk/artifact/store.py b/sqlspec/extensions/adk/artifact/store.py index ec857dc73..fcac2a28c 100644 --- a/sqlspec/extensions/adk/artifact/store.py +++ b/sqlspec/extensions/adk/artifact/store.py @@ -28,18 +28,8 @@ logger = get_logger("sqlspec.extensions.adk.artifact.store") -class BaseAsyncADKArtifactStore(ABC, Generic[ConfigT]): - """Base class for async SQLSpec-backed ADK artifact metadata stores. - - Manages artifact version metadata in a SQL table. Content bytes are - stored externally via ``sqlspec/storage/`` backends and referenced - by canonical URI in each metadata row. - - Subclasses must implement dialect-specific SQL queries. - - Args: - config: SQLSpec database configuration with extension_config["adk"] settings. - """ +class _ADKArtifactStoreCommon(Generic[ConfigT]): + """Shared non-async ADK store state and helpers.""" __slots__ = ("_artifact_table", "_config") @@ -64,6 +54,22 @@ def artifact_table(self) -> str: """Return the artifact versions table name.""" return self._artifact_table + +class BaseAsyncADKArtifactStore(_ADKArtifactStoreCommon[ConfigT], ABC): + """Base class for async SQLSpec-backed ADK artifact metadata stores. + + Manages artifact version metadata in a SQL table. Content bytes are + stored externally via ``sqlspec/storage/`` backends and referenced + by canonical URI in each metadata row. + + Subclasses must implement dialect-specific SQL queries. + + Args: + config: SQLSpec database configuration with extension_config["adk"] settings. + """ + + __slots__ = () + @abstractmethod async def insert_artifact(self, record: "ArtifactRecord") -> None: """Insert an artifact version metadata row. @@ -179,7 +185,7 @@ async def ensure_table(self) -> None: ) -class BaseSyncADKArtifactStore(ABC, Generic[ConfigT]): +class BaseSyncADKArtifactStore(_ADKArtifactStoreCommon[ConfigT], ABC): """Base class for sync SQLSpec-backed ADK artifact metadata stores. Synchronous counterpart of :class:`BaseAsyncADKArtifactStore`. @@ -188,28 +194,7 @@ class BaseSyncADKArtifactStore(ABC, Generic[ConfigT]): config: SQLSpec database configuration with extension_config["adk"] settings. """ - __slots__ = ("_artifact_table", "_config") - - def __init__(self, config: ConfigT) -> None: - """Initialize the sync ADK artifact store. - - Args: - config: SQLSpec database configuration. - """ - self._config = config - store_config = _adk_artifact_store_config(self._config) - self._artifact_table: str = store_config["artifact_table"] - ensure_table_name(self._artifact_table) - - @property - def config(self) -> ConfigT: - """Return the database configuration.""" - return self._config - - @property - def artifact_table(self) -> str: - """Return the artifact versions table name.""" - return self._artifact_table + __slots__ = () @abstractmethod def insert_artifact(self, record: "ArtifactRecord") -> None: diff --git a/sqlspec/extensions/adk/memory/store.py b/sqlspec/extensions/adk/memory/store.py index 609abfe3b..3c95c0450 100644 --- a/sqlspec/extensions/adk/memory/store.py +++ b/sqlspec/extensions/adk/memory/store.py @@ -10,6 +10,8 @@ from sqlspec.utils.logging import get_logger, log_with_context if TYPE_CHECKING: + from collections.abc import Callable + from sqlspec.config import DatabaseConfigProtocol from sqlspec.extensions.adk.memory._types import MemoryRecord @@ -23,25 +25,11 @@ ADK_RESET_MEMORY_TABLES: Final = ("adk_memory", "adk_memory_entries") -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. +class _ADKMemoryStoreCommon(Generic[ConfigT]): + """Shared non-async ADK store state and helpers.""" - Args: - config: SQLSpec database configuration with extension_config["adk"] settings. - """ + if TYPE_CHECKING: + _drop_memory_table_sql: "Callable[[], list[str]]" __slots__ = ( "_config", @@ -106,6 +94,70 @@ 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 + def _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 _adk_memory_store_config(self._config) + + def _reset_drop_memory_table_sql(self) -> "list[str]": + """Return memory drops needed before recreating the clean-break schema.""" + return reset_drop_sql( + list(self._drop_memory_table_sql()), ADK_RESET_MEMORY_TABLES, self._drop_memory_sql_for_table + ) + + def _drop_memory_sql_for_table(self, table_name: str) -> "list[str]": + current_table = self._memory_table + self._memory_table = table_name + try: + return list(self._drop_memory_table_sql()) + finally: + self._memory_table = current_table + + 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", + ) + + +class BaseAsyncADKMemoryStore(_ADKMemoryStoreCommon[ConfigT], ABC): + """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. + + Args: + config: SQLSpec database configuration with extension_config["adk"] settings. + """ + + __slots__ = () + @abstractmethod async def create_tables(self) -> None: """Create the memory table and indexes if they don't exist. @@ -189,14 +241,6 @@ async def delete_entries_older_than(self, days: int) -> int: """ raise NotImplementedError - def _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 _adk_memory_store_config(self._config) - @abstractmethod async def _memory_table_ddl(self) -> "str | list[str]": """Get the CREATE TABLE SQL for the memory table. @@ -215,41 +259,8 @@ def _drop_memory_table_sql(self) -> "list[str]": """ raise NotImplementedError - def _reset_drop_memory_table_sql(self) -> "list[str]": - """Return memory drops needed before recreating the clean-break schema.""" - return reset_drop_sql( - list(self._drop_memory_table_sql()), ADK_RESET_MEMORY_TABLES, self._drop_memory_sql_for_table - ) - - def _drop_memory_sql_for_table(self, table_name: str) -> "list[str]": - current_table = self._memory_table - self._memory_table = table_name - try: - return list(self._drop_memory_table_sql()) - finally: - self._memory_table = current_table - - 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", - ) - - -class BaseSyncADKMemoryStore(ABC, Generic[ConfigT]): +class BaseSyncADKMemoryStore(_ADKMemoryStoreCommon[ConfigT], ABC): """Base class for sync SQLSpec-backed ADK memory stores. Implements storage operations for Google ADK memory entries using @@ -269,68 +280,7 @@ class BaseSyncADKMemoryStore(ABC, Generic[ConfigT]): config: SQLSpec database configuration with extension_config["adk"] settings. """ - __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. - """ - self._config = config - store_config = self._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 = ( - owner_id_column_name(self._owner_id_column_ddl) if self._owner_id_column_ddl else None - ) - ensure_table_name(self._memory_table) - - @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 + __slots__ = () @abstractmethod def create_tables(self) -> None: @@ -415,14 +365,6 @@ def delete_entries_older_than(self, days: int) -> int: """ raise NotImplementedError - def _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 _adk_memory_store_config(self._config) - @abstractmethod def _memory_table_ddl(self) -> "str | list[str]": """Get the CREATE TABLE SQL for the memory table. @@ -432,39 +374,6 @@ def _memory_table_ddl(self) -> "str | list[str]": """ raise NotImplementedError - def _reset_drop_memory_table_sql(self) -> "list[str]": - """Return memory drops needed before recreating the clean-break schema.""" - return reset_drop_sql( - list(self._drop_memory_table_sql()), ADK_RESET_MEMORY_TABLES, self._drop_memory_sql_for_table - ) - - def _drop_memory_sql_for_table(self, table_name: str) -> "list[str]": - current_table = self._memory_table - self._memory_table = table_name - try: - return list(self._drop_memory_table_sql()) - finally: - self._memory_table = current_table - - 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 _drop_memory_table_sql(self) -> "list[str]": """Get the DROP TABLE SQL statements for this database dialect. diff --git a/sqlspec/extensions/adk/store.py b/sqlspec/extensions/adk/store.py index 2f40c7fe8..bc31d3784 100644 --- a/sqlspec/extensions/adk/store.py +++ b/sqlspec/extensions/adk/store.py @@ -13,6 +13,8 @@ from sqlspec.utils.sync_tools import async_ if TYPE_CHECKING: + from collections.abc import Callable + from sqlspec.config import DatabaseConfigProtocol from sqlspec.extensions.adk._types import EventRecord, SessionRecord @@ -30,33 +32,11 @@ ) -class BaseAsyncADKStore(ABC, Generic[ConfigT]): - """Base class for async SQLSpec-backed ADK session stores. +class _ADKStoreCommon(Generic[ConfigT]): + """Shared non-async ADK store state and helpers.""" - Implements storage operations for Google ADK sessions and events using - SQLSpec database adapters with async/await. - - This abstract base class provides common functionality for all database-specific - store implementations including: - - Connection management via SQLSpec configs - - Table name validation - - Session and event CRUD operations - - 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"]: - - session_table: Sessions table name (default: "adk_session") - - events_table: Events table name (default: "adk_event") - - app_state_table: App-scoped state table name (default: "adk_app_state") - - user_state_table: User-scoped state table name (default: "adk_user_state") - - metadata_table: Internal metadata table name (default: "adk_internal_metadata") - - owner_id_column: Optional owner FK column DDL (default: None) - """ + if TYPE_CHECKING: + _drop_tables_sql: "Callable[[], list[str]]" __slots__ = ( "_app_state_table", @@ -101,6 +81,162 @@ def __init__(self, config: ConfigT) -> None: ensure_table_name(self._user_state_table) ensure_table_name(self._metadata_table) + @property + def config(self) -> ConfigT: + """Return the database configuration.""" + return self._config + + @property + def session_table(self) -> str: + """Return the sessions table name.""" + return self._session_table + + @property + def events_table(self) -> str: + """Return the events table name.""" + return self._events_table + + @property + def app_state_table(self) -> str: + """Return the app-scoped state table name.""" + return self._app_state_table + + @property + def user_state_table(self) -> str: + """Return the user-scoped state table name.""" + return self._user_state_table + + @property + def metadata_table(self) -> str: + """Return the ADK metadata table name.""" + return self._metadata_table + + @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 + + def _reset_drop_tables_sql(self) -> "list[str]": + """Return all table drops needed before recreating the clean-break schema.""" + statements = list(self._drop_tables_sql()) + for table_profile in ADK_RESET_TABLE_PROFILES: + statements.extend(self._drop_sql_for_table_profile(table_profile)) + return unique_statements(statements) + + def _store_config_from_extension(self) -> "dict[str, Any]": + """Extract ADK store configuration from config.extension_config. + + Returns: + Dict with ADK table names and optionally owner_id_column. + """ + return dict(_adk_session_store_config(self._config)) + + def _calculate_expires_at(self, expires_in: "int | timedelta | None") -> "datetime | None": + """Calculate expiration timestamp from expires_in. + + Args: + expires_in: Seconds or timedelta until expiration. + + Returns: + UTC datetime of expiration, or None if no expiration. + """ + if expires_in is None: + return None + + expires_in_seconds = int(expires_in.total_seconds()) if isinstance(expires_in, timedelta) else expires_in + + if expires_in_seconds <= 0: + return None + + return datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds) + + def _drop_sql_for_table_profile(self, table_profile: "tuple[str, str, str, str, str]") -> "list[str]": + session_table, events_table, app_state_table, user_state_table, metadata_table = table_profile + current_session_table = self._session_table + current_events_table = self._events_table + current_app_state_table = self._app_state_table + current_user_state_table = self._user_state_table + current_table = self._metadata_table + self._session_table = session_table + self._events_table = events_table + self._app_state_table = app_state_table + self._user_state_table = user_state_table + self._metadata_table = metadata_table + try: + return list(self._drop_tables_sql()) + finally: + self._session_table = current_session_table + self._events_table = current_events_table + self._app_state_table = current_app_state_table + self._user_state_table = current_user_state_table + self._metadata_table = current_table + + def _log_tables_created(self) -> None: + log_with_context( + logger, + logging.DEBUG, + "adk.tables.ready", + db_system=resolve_db_system(type(self).__name__), + session_table=self._session_table, + events_table=self._events_table, + ) + + def _log_tables_dropped(self) -> None: + log_with_context( + logger, + logging.DEBUG, + "adk.tables.dropped", + db_system=resolve_db_system(type(self).__name__), + session_table=self._session_table, + events_table=self._events_table, + ) + + def _log_tables_recreated(self) -> None: + log_with_context( + logger, + logging.DEBUG, + "adk.tables.recreated", + db_system=resolve_db_system(type(self).__name__), + session_table=self._session_table, + events_table=self._events_table, + ) + + +class BaseAsyncADKStore(_ADKStoreCommon[ConfigT], ABC): + """Base class for async SQLSpec-backed ADK session stores. + + Implements storage operations for Google ADK sessions and events using + SQLSpec database adapters with async/await. + + This abstract base class provides common functionality for all database-specific + store implementations including: + - Connection management via SQLSpec configs + - Table name validation + - Session and event CRUD operations + + 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"]: + - session_table: Sessions table name (default: "adk_session") + - events_table: Events table name (default: "adk_event") + - app_state_table: App-scoped state table name (default: "adk_app_state") + - user_state_table: User-scoped state table name (default: "adk_user_state") + - metadata_table: Internal metadata table name (default: "adk_internal_metadata") + - owner_id_column: Optional owner FK column DDL (default: None) + """ + + __slots__ = () + async def create_tables(self) -> None: """Create the sessions and events tables if they don't exist.""" raise NotImplementedError @@ -346,46 +482,6 @@ async def set_metadata(self, key: str, value: str) -> None: """ raise NotImplementedError - @property - def config(self) -> ConfigT: - """Return the database configuration.""" - return self._config - - @property - def session_table(self) -> str: - """Return the sessions table name.""" - return self._session_table - - @property - def events_table(self) -> str: - """Return the events table name.""" - return self._events_table - - @property - def app_state_table(self) -> str: - """Return the app-scoped state table name.""" - return self._app_state_table - - @property - def user_state_table(self) -> str: - """Return the user-scoped state table name.""" - return self._user_state_table - - @property - def metadata_table(self) -> str: - """Return the ADK metadata table name.""" - return self._metadata_table - - @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 - async def ensure_tables(self) -> None: """Create tables and emit a standardized log entry.""" @@ -403,40 +499,6 @@ async def recreate_tables(self) -> None: await self.ensure_tables() self._log_tables_recreated() - def _reset_drop_tables_sql(self) -> "list[str]": - """Return all table drops needed before recreating the clean-break schema.""" - statements = list(self._drop_tables_sql()) - for table_profile in ADK_RESET_TABLE_PROFILES: - statements.extend(self._drop_sql_for_table_profile(table_profile)) - return unique_statements(statements) - - def _store_config_from_extension(self) -> "dict[str, Any]": - """Extract ADK store configuration from config.extension_config. - - Returns: - Dict with ADK table names and optionally owner_id_column. - """ - return dict(_adk_session_store_config(self._config)) - - def _calculate_expires_at(self, expires_in: "int | timedelta | None") -> "datetime | None": - """Calculate expiration timestamp from expires_in. - - Args: - expires_in: Seconds or timedelta until expiration. - - Returns: - UTC datetime of expiration, or None if no expiration. - """ - if expires_in is None: - return None - - expires_in_seconds = int(expires_in.total_seconds()) if isinstance(expires_in, timedelta) else expires_in - - if expires_in_seconds <= 0: - return None - - return datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds) - async def _execute_lifecycle_scripts(self, statements: list[str]) -> None: """Execute lifecycle DDL scripts for async and sync-backed configs.""" session_context = self._config.provide_session() @@ -453,15 +515,7 @@ async def _execute_lifecycle_scripts(self, statements: list[str]) -> None: await result return - def _execute_sync() -> None: - with cast("Any", self._config.provide_session()) as driver: - for statement in statements: - driver.execute_script(statement) - commit = getattr(driver, "commit", None) - if callable(commit): - commit() - - await async_(_execute_sync)() + await async_(_run_lifecycle_sync)(self._config, statements) @abstractmethod async def _sessions_table_ddl(self) -> str: @@ -558,59 +612,8 @@ def _drop_tables_sql(self) -> "list[str]": """ raise NotImplementedError - def _drop_sql_for_table_profile(self, table_profile: "tuple[str, str, str, str, str]") -> "list[str]": - session_table, events_table, app_state_table, user_state_table, metadata_table = table_profile - current_session_table = self._session_table - current_events_table = self._events_table - current_app_state_table = self._app_state_table - current_user_state_table = self._user_state_table - current_table = self._metadata_table - self._session_table = session_table - self._events_table = events_table - self._app_state_table = app_state_table - self._user_state_table = user_state_table - self._metadata_table = metadata_table - try: - return list(self._drop_tables_sql()) - finally: - self._session_table = current_session_table - self._events_table = current_events_table - self._app_state_table = current_app_state_table - self._user_state_table = current_user_state_table - self._metadata_table = current_table - def _log_tables_created(self) -> None: - log_with_context( - logger, - logging.DEBUG, - "adk.tables.ready", - db_system=resolve_db_system(type(self).__name__), - session_table=self._session_table, - events_table=self._events_table, - ) - - def _log_tables_dropped(self) -> None: - log_with_context( - logger, - logging.DEBUG, - "adk.tables.dropped", - db_system=resolve_db_system(type(self).__name__), - session_table=self._session_table, - events_table=self._events_table, - ) - - def _log_tables_recreated(self) -> None: - log_with_context( - logger, - logging.DEBUG, - "adk.tables.recreated", - db_system=resolve_db_system(type(self).__name__), - session_table=self._session_table, - events_table=self._events_table, - ) - - -class BaseSyncADKStore(ABC, Generic[ConfigT]): +class BaseSyncADKStore(_ADKStoreCommon[ConfigT], ABC): """Base class for sync SQLSpec-backed ADK session stores. Sync-backed adapters expose a real synchronous API for direct use in @@ -621,39 +624,7 @@ class BaseSyncADKStore(ABC, Generic[ConfigT]): config: SQLSpec database configuration with extension_config["adk"] settings. """ - __slots__ = ( - "_app_state_table", - "_config", - "_events_table", - "_metadata_table", - "_owner_id_column_ddl", - "_owner_id_column_name", - "_session_table", - "_user_state_table", - ) - - def __init__(self, config: ConfigT) -> None: - """Initialize the sync ADK store. - - Args: - config: SQLSpec database configuration. - """ - self._config = config - store_config = self._store_config_from_extension() - self._session_table: str = str(store_config["session_table"]) - self._events_table: str = str(store_config["events_table"]) - self._app_state_table: str = str(store_config["app_state_table"]) - self._user_state_table: str = str(store_config["user_state_table"]) - self._metadata_table: str = str(store_config["metadata_table"]) - self._owner_id_column_ddl: str | None = store_config.get("owner_id_column") - self._owner_id_column_name: str | None = ( - owner_id_column_name(self._owner_id_column_ddl) if self._owner_id_column_ddl else None - ) - ensure_table_name(self._session_table) - ensure_table_name(self._events_table) - ensure_table_name(self._app_state_table) - ensure_table_name(self._user_state_table) - ensure_table_name(self._metadata_table) + __slots__ = () @abstractmethod def create_tables(self) -> None: @@ -761,46 +732,6 @@ def set_metadata(self, key: str, value: str) -> None: """Set a value in the ADK internal metadata table.""" raise NotImplementedError - @property - def config(self) -> ConfigT: - """Return the database configuration.""" - return self._config - - @property - def session_table(self) -> str: - """Return the sessions table name.""" - return self._session_table - - @property - def events_table(self) -> str: - """Return the events table name.""" - return self._events_table - - @property - def app_state_table(self) -> str: - """Return the app-scoped state table name.""" - return self._app_state_table - - @property - def user_state_table(self) -> str: - """Return the user-scoped state table name.""" - return self._user_state_table - - @property - def metadata_table(self) -> str: - """Return the ADK metadata table name.""" - return self._metadata_table - - @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 - def ensure_tables(self) -> None: """Create tables and emit a standardized log entry.""" @@ -818,30 +749,9 @@ def recreate_tables(self) -> None: self.ensure_tables() self._log_tables_recreated() - def _store_config_from_extension(self) -> "dict[str, Any]": - """Extract ADK store configuration from config.extension_config.""" - return dict(_adk_session_store_config(self._config)) - - def _calculate_expires_at(self, expires_in: "int | timedelta | None") -> "datetime | None": - """Calculate expiration timestamp from expires_in.""" - if expires_in is None: - return None - - expires_in_seconds = int(expires_in.total_seconds()) if isinstance(expires_in, timedelta) else expires_in - - if expires_in_seconds <= 0: - return None - - return datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds) - def _execute_lifecycle_scripts(self, statements: list[str]) -> None: """Execute lifecycle DDL scripts using the sync driver session.""" - with cast("Any", self._config.provide_session()) as driver: - for statement in statements: - driver.execute_script(statement) - commit = getattr(driver, "commit", None) - if callable(commit): - commit() + _run_lifecycle_sync(self._config, statements) @abstractmethod def _sessions_table_ddl(self) -> str: @@ -893,60 +803,12 @@ def _drop_tables_sql(self) -> "list[str]": """Get the DROP TABLE SQL statements for this database dialect.""" raise NotImplementedError - def _reset_drop_tables_sql(self) -> "list[str]": - """Return all table drops needed before recreating the clean-break schema.""" - statements = list(self._drop_tables_sql()) - for table_profile in ADK_RESET_TABLE_PROFILES: - statements.extend(self._drop_sql_for_table_profile(table_profile)) - return unique_statements(statements) - - def _drop_sql_for_table_profile(self, table_profile: "tuple[str, str, str, str, str]") -> "list[str]": - session_table, events_table, app_state_table, user_state_table, metadata_table = table_profile - current_session_table = self._session_table - current_events_table = self._events_table - current_app_state_table = self._app_state_table - current_user_state_table = self._user_state_table - current_table = self._metadata_table - self._session_table = session_table - self._events_table = events_table - self._app_state_table = app_state_table - self._user_state_table = user_state_table - self._metadata_table = metadata_table - try: - return list(self._drop_tables_sql()) - finally: - self._session_table = current_session_table - self._events_table = current_events_table - self._app_state_table = current_app_state_table - self._user_state_table = current_user_state_table - self._metadata_table = current_table - - def _log_tables_created(self) -> None: - log_with_context( - logger, - logging.DEBUG, - "adk.tables.ready", - db_system=resolve_db_system(type(self).__name__), - session_table=self._session_table, - events_table=self._events_table, - ) - - def _log_tables_dropped(self) -> None: - log_with_context( - logger, - logging.DEBUG, - "adk.tables.dropped", - db_system=resolve_db_system(type(self).__name__), - session_table=self._session_table, - events_table=self._events_table, - ) - def _log_tables_recreated(self) -> None: - log_with_context( - logger, - logging.DEBUG, - "adk.tables.recreated", - db_system=resolve_db_system(type(self).__name__), - session_table=self._session_table, - events_table=self._events_table, - ) +def _run_lifecycle_sync(config: Any, statements: "list[str]") -> None: + """Execute lifecycle statements through a synchronous config session.""" + with cast("Any", config.provide_session()) as driver: + for statement in statements: + driver.execute_script(statement) + commit = getattr(driver, "commit", None) + if callable(commit): + commit() diff --git a/tests/unit/extensions/test_adk/test_store_common_mixins.py b/tests/unit/extensions/test_adk/test_store_common_mixins.py new file mode 100644 index 000000000..6870c76f9 --- /dev/null +++ b/tests/unit/extensions/test_adk/test_store_common_mixins.py @@ -0,0 +1,30 @@ +"""Regression tests for common ADK store state mixins.""" + +from sqlspec.extensions.adk.artifact.store import ( + BaseAsyncADKArtifactStore, + BaseSyncADKArtifactStore, + _ADKArtifactStoreCommon, +) +from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore, _ADKMemoryStoreCommon +from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore, _ADKStoreCommon + + +def test_session_store_bases_share_slotted_common_state() -> None: + assert issubclass(BaseAsyncADKStore, _ADKStoreCommon) + assert issubclass(BaseSyncADKStore, _ADKStoreCommon) + assert BaseAsyncADKStore.__slots__ == () + assert BaseSyncADKStore.__slots__ == () + + +def test_memory_store_bases_share_slotted_common_state() -> None: + assert issubclass(BaseAsyncADKMemoryStore, _ADKMemoryStoreCommon) + assert issubclass(BaseSyncADKMemoryStore, _ADKMemoryStoreCommon) + assert BaseAsyncADKMemoryStore.__slots__ == () + assert BaseSyncADKMemoryStore.__slots__ == () + + +def test_artifact_store_bases_share_slotted_common_state() -> None: + assert issubclass(BaseAsyncADKArtifactStore, _ADKArtifactStoreCommon) + assert issubclass(BaseSyncADKArtifactStore, _ADKArtifactStoreCommon) + assert BaseAsyncADKArtifactStore.__slots__ == () + assert BaseSyncADKArtifactStore.__slots__ == () From 2832b16001c3513af08611283933da5551e6e9fe Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 11 Jul 2026 00:58:25 +0000 Subject: [PATCH 4/6] refactor(adk): hoist adapter SQL templates --- .../adapters/cockroach_psycopg/adk/store.py | 408 +++++----- sqlspec/adapters/mysqlconnector/adk/store.py | 239 +++--- sqlspec/adapters/oracledb/adk/store.py | 720 ++++++++---------- sqlspec/adapters/psycopg/adk/store.py | 392 +++++----- .../test_adk/test_sql_template_hoist.py | 30 + 5 files changed, 919 insertions(+), 870 deletions(-) create mode 100644 tests/unit/extensions/test_adk/test_sql_template_hoist.py diff --git a/sqlspec/adapters/cockroach_psycopg/adk/store.py b/sqlspec/adapters/cockroach_psycopg/adk/store.py index 5211cea82..fc5ea2b43 100644 --- a/sqlspec/adapters/cockroach_psycopg/adk/store.py +++ b/sqlspec/adapters/cockroach_psycopg/adk/store.py @@ -31,6 +31,130 @@ logger = get_logger("sqlspec.adapters.cockroach_psycopg.adk.store") +_ADK_SESSIONS_TABLE_DDL_TEMPLATE = ",\n {0}" + +_ADK_SESSIONS_TABLE_DDL_TEMPLATE_2 = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL{1},\n" + " state JSONB NOT NULL DEFAULT '{{}}'::jsonb,\n" + " create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,\n" + " update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP\n" + " ){2};\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{3}_app_user\n" + " ON {4}(app_name, user_id){5};\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{6}_update_time\n" + " ON {7}(update_time DESC){8};\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{9}_state\n" + " ON {10} USING GIN (state)\n" + " WHERE state != '{{}}'::jsonb;\n" + " " +) + +_ADK_EVENTS_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " session_id VARCHAR(128) NOT NULL,\n" + " invocation_id VARCHAR(256),\n" + " timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,\n" + " event_data JSONB NOT NULL,\n" + " FOREIGN KEY (session_id) REFERENCES {1}(id) ON DELETE CASCADE\n" + " ){2};\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{3}_session\n" + " ON {4}(session_id, timestamp ASC){5}{6};\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{7}_event_data\n" + " ON {8} USING GIN (event_data);\n" + " " +) + +_ADK_APP_STATES_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " app_name VARCHAR(128) PRIMARY KEY,\n" + " state JSONB NOT NULL DEFAULT '{{}}'::jsonb,\n" + " update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP\n" + " ){1};\n" + " " +) + +_ADK_USER_STATES_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL,\n" + " state JSONB NOT NULL DEFAULT '{{}}'::jsonb,\n" + " update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,\n" + " PRIMARY KEY (app_name, user_id)\n" + " ){1};\n" + " " +) + +_ADK_METADATA_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " key VARCHAR(128) PRIMARY KEY,\n" + " value VARCHAR(512) NOT NULL\n" + " ){1};\n" + " " +) + +_ADK_METADATA_SEED_SQL_TEMPLATE = ( + "\n" + " INSERT INTO {0} (key, value)\n" + " VALUES ('schema_version', '1')\n" + " ON CONFLICT (key) DO NOTHING\n" + " " +) + +_ADK_MEMORY_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE INDEX IF NOT EXISTS idx_{0}_fts\n" + " ON {1} USING GIN (to_tsvector('english', content_text));\n" + " " +) + +_ADK_MEMORY_TABLE_DDL_TEMPLATE_2 = ( + "\n" + " CREATE INDEX IF NOT EXISTS idx_{0}_content_trgm\n" + " ON {1} USING GIN (content_text gin_trgm_ops);\n" + " " +) + +_ADK_MEMORY_TABLE_DDL_TEMPLATE_3 = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " session_id VARCHAR(128) NOT NULL,\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL,\n" + " event_id VARCHAR(128) NOT NULL UNIQUE,\n" + " author VARCHAR(256){1},\n" + " timestamp TIMESTAMPTZ NOT NULL,\n" + " content_json JSONB NOT NULL,\n" + " content_text TEXT NOT NULL,\n" + " metadata_json JSONB,\n" + " inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP\n" + " ){2};\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{3}_app_user_time\n" + " ON {4}(app_name, user_id, timestamp DESC){5};\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{6}_session\n" + " ON {7}(session_id);\n" + " {8}\n" + " {9}\n" + " " +) + + class CockroachPsycopgADKConfig(ADKConfig): """CockroachDB psycopg ADK extension settings. @@ -436,31 +560,24 @@ async def _sessions_table_ddl(self) -> str: adk_config = _adk_config(self._config) owner_id_line = "" if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" + owner_id_line = _ADK_SESSIONS_TABLE_DDL_TEMPLATE.format(self._owner_id_column_ddl) session_locality = _cockroach_table_locality_clause(adk_config, "session_table_locality") hash_shard_clause = _cockroach_hash_shard_clause(adk_config) session_storing_clause = _cockroach_storing_clause(adk_config, ("state", "create_time", "update_time")) - 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 - ){session_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user - ON {self._session_table}(app_name, user_id){session_storing_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time - ON {self._session_table}(update_time DESC){hash_shard_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state - ON {self._session_table} USING GIN (state) - WHERE state != '{{}}'::jsonb; - """ + return _ADK_SESSIONS_TABLE_DDL_TEMPLATE_2.format( + self._session_table, + owner_id_line, + session_locality, + self._session_table, + self._session_table, + session_storing_clause, + self._session_table, + self._session_table, + hash_shard_clause, + self._session_table, + self._session_table, + ) async def _events_table_ddl(self) -> str: adk_config = _adk_config(self._config) @@ -468,66 +585,38 @@ async def _events_table_ddl(self) -> str: hash_shard_clause = _cockroach_hash_shard_clause(adk_config) events_storing_clause = _cockroach_storing_clause(adk_config, ("invocation_id", "event_data")) - return f""" - CREATE TABLE IF NOT EXISTS {self._events_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - invocation_id VARCHAR(256), - timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - event_data JSONB NOT NULL, - FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE - ){events_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session - ON {self._events_table}(session_id, timestamp ASC){hash_shard_clause}{events_storing_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_event_data - ON {self._events_table} USING GIN (event_data); - """ + return _ADK_EVENTS_TABLE_DDL_TEMPLATE.format( + self._events_table, + self._session_table, + events_locality, + self._events_table, + self._events_table, + hash_shard_clause, + events_storing_clause, + self._events_table, + self._events_table, + ) async def _app_states_table_ddl(self) -> str: adk_config = _adk_config(self._config) app_state_locality = _cockroach_table_locality_clause(adk_config, "app_state_table_locality") - return f""" - CREATE TABLE IF NOT EXISTS {self._app_state_table} ( - app_name VARCHAR(128) PRIMARY KEY, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ){app_state_locality}; - """ + return _ADK_APP_STATES_TABLE_DDL_TEMPLATE.format(self._app_state_table, app_state_locality) async def _user_states_table_ddl(self) -> str: adk_config = _adk_config(self._config) user_state_locality = _cockroach_table_locality_clause(adk_config, "user_state_table_locality") - return f""" - CREATE TABLE IF NOT EXISTS {self._user_state_table} ( - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY (app_name, user_id) - ){user_state_locality}; - """ + return _ADK_USER_STATES_TABLE_DDL_TEMPLATE.format(self._user_state_table, user_state_locality) async def _metadata_table_ddl(self) -> str: adk_config = _adk_config(self._config) metadata_locality = _cockroach_table_locality_clause(adk_config, "metadata_table_locality") - return f""" - CREATE TABLE IF NOT EXISTS {self._metadata_table} ( - key VARCHAR(128) PRIMARY KEY, - value VARCHAR(512) NOT NULL - ){metadata_locality}; - """ + return _ADK_METADATA_TABLE_DDL_TEMPLATE.format(self._metadata_table, metadata_locality) async def _metadata_seed_sql(self) -> str: - return f""" - INSERT INTO {self._metadata_table} (key, value) - VALUES ('schema_version', '1') - ON CONFLICT (key) DO NOTHING - """ + return _ADK_METADATA_SEED_SQL_TEMPLATE.format(self._metadata_table) def _drop_app_states_table_sql(self) -> str: return f"DROP TABLE IF EXISTS {self._app_state_table}" @@ -912,31 +1001,24 @@ def _sessions_table_ddl(self) -> str: adk_config = _adk_config(self._config) owner_id_line = "" if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" + owner_id_line = _ADK_SESSIONS_TABLE_DDL_TEMPLATE.format(self._owner_id_column_ddl) session_locality = _cockroach_table_locality_clause(adk_config, "session_table_locality") hash_shard_clause = _cockroach_hash_shard_clause(adk_config) session_storing_clause = _cockroach_storing_clause(adk_config, ("state", "create_time", "update_time")) - 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 - ){session_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user - ON {self._session_table}(app_name, user_id){session_storing_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time - ON {self._session_table}(update_time DESC){hash_shard_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state - ON {self._session_table} USING GIN (state) - WHERE state != '{{}}'::jsonb; - """ + return _ADK_SESSIONS_TABLE_DDL_TEMPLATE_2.format( + self._session_table, + owner_id_line, + session_locality, + self._session_table, + self._session_table, + session_storing_clause, + self._session_table, + self._session_table, + hash_shard_clause, + self._session_table, + self._session_table, + ) def _events_table_ddl(self) -> str: adk_config = _adk_config(self._config) @@ -944,66 +1026,38 @@ def _events_table_ddl(self) -> str: hash_shard_clause = _cockroach_hash_shard_clause(adk_config) events_storing_clause = _cockroach_storing_clause(adk_config, ("invocation_id", "event_data")) - return f""" - CREATE TABLE IF NOT EXISTS {self._events_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - invocation_id VARCHAR(256), - timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - event_data JSONB NOT NULL, - FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE - ){events_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session - ON {self._events_table}(session_id, timestamp ASC){hash_shard_clause}{events_storing_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_event_data - ON {self._events_table} USING GIN (event_data); - """ + return _ADK_EVENTS_TABLE_DDL_TEMPLATE.format( + self._events_table, + self._session_table, + events_locality, + self._events_table, + self._events_table, + hash_shard_clause, + events_storing_clause, + self._events_table, + self._events_table, + ) def _app_states_table_ddl(self) -> str: adk_config = _adk_config(self._config) app_state_locality = _cockroach_table_locality_clause(adk_config, "app_state_table_locality") - return f""" - CREATE TABLE IF NOT EXISTS {self._app_state_table} ( - app_name VARCHAR(128) PRIMARY KEY, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ){app_state_locality}; - """ + return _ADK_APP_STATES_TABLE_DDL_TEMPLATE.format(self._app_state_table, app_state_locality) def _user_states_table_ddl(self) -> str: adk_config = _adk_config(self._config) user_state_locality = _cockroach_table_locality_clause(adk_config, "user_state_table_locality") - return f""" - CREATE TABLE IF NOT EXISTS {self._user_state_table} ( - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY (app_name, user_id) - ){user_state_locality}; - """ + return _ADK_USER_STATES_TABLE_DDL_TEMPLATE.format(self._user_state_table, user_state_locality) def _metadata_table_ddl(self) -> str: adk_config = _adk_config(self._config) metadata_locality = _cockroach_table_locality_clause(adk_config, "metadata_table_locality") - return f""" - CREATE TABLE IF NOT EXISTS {self._metadata_table} ( - key VARCHAR(128) PRIMARY KEY, - value VARCHAR(512) NOT NULL - ){metadata_locality}; - """ + return _ADK_METADATA_TABLE_DDL_TEMPLATE.format(self._metadata_table, metadata_locality) def _metadata_seed_sql(self) -> str: - return f""" - INSERT INTO {self._metadata_table} (key, value) - VALUES ('schema_version', '1') - ON CONFLICT (key) DO NOTHING - """ + return _ADK_METADATA_SEED_SQL_TEMPLATE.format(self._metadata_table) def _drop_app_states_table_sql(self) -> str: return f"DROP TABLE IF EXISTS {self._app_state_table}" @@ -1176,46 +1230,29 @@ async def _memory_table_ddl(self) -> str: adk_config = _adk_config(self._config) owner_id_line = "" if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" + owner_id_line = _ADK_SESSIONS_TABLE_DDL_TEMPLATE.format(self._owner_id_column_ddl) memory_locality = _cockroach_table_locality_clause(adk_config, "memory_table_locality") hash_shard_clause = _cockroach_hash_shard_clause(adk_config) fts_index = "" if self._use_fts: - fts_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts - ON {self._memory_table} USING GIN (to_tsvector('english', content_text)); - """ + fts_index = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(self._memory_table, self._memory_table) trigram_index = "" if adk_config.get("enable_memory_trigram_index", False): - trigram_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_content_trgm - ON {self._memory_table} USING GIN (content_text gin_trgm_ops); - """ - - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMPTZ NOT NULL, - content_json JSONB NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSONB, - inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ){memory_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC){hash_shard_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session - ON {self._memory_table}(session_id); - {fts_index} - {trigram_index} - """ + trigram_index = _ADK_MEMORY_TABLE_DDL_TEMPLATE_2.format(self._memory_table, self._memory_table) + + return _ADK_MEMORY_TABLE_DDL_TEMPLATE_3.format( + self._memory_table, + owner_id_line, + memory_locality, + self._memory_table, + self._memory_table, + hash_shard_clause, + self._memory_table, + self._memory_table, + fts_index, + trigram_index, + ) def _drop_memory_table_sql(self) -> "list[str]": return [f"DROP TABLE IF EXISTS {self._memory_table}"] @@ -1354,46 +1391,29 @@ def _memory_table_ddl(self) -> str: adk_config = _adk_config(self._config) owner_id_line = "" if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" + owner_id_line = _ADK_SESSIONS_TABLE_DDL_TEMPLATE.format(self._owner_id_column_ddl) memory_locality = _cockroach_table_locality_clause(adk_config, "memory_table_locality") hash_shard_clause = _cockroach_hash_shard_clause(adk_config) fts_index = "" if self._use_fts: - fts_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts - ON {self._memory_table} USING GIN (to_tsvector('english', content_text)); - """ + fts_index = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(self._memory_table, self._memory_table) trigram_index = "" if adk_config.get("enable_memory_trigram_index", False): - trigram_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_content_trgm - ON {self._memory_table} USING GIN (content_text gin_trgm_ops); - """ - - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMPTZ NOT NULL, - content_json JSONB NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSONB, - inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ){memory_locality}; - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC){hash_shard_clause}; - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session - ON {self._memory_table}(session_id); - {fts_index} - {trigram_index} - """ + trigram_index = _ADK_MEMORY_TABLE_DDL_TEMPLATE_2.format(self._memory_table, self._memory_table) + + return _ADK_MEMORY_TABLE_DDL_TEMPLATE_3.format( + self._memory_table, + owner_id_line, + memory_locality, + self._memory_table, + self._memory_table, + hash_shard_clause, + self._memory_table, + self._memory_table, + fts_index, + trigram_index, + ) def _drop_memory_table_sql(self) -> "list[str]": return [f"DROP TABLE IF EXISTS {self._memory_table}"] diff --git a/sqlspec/adapters/mysqlconnector/adk/store.py b/sqlspec/adapters/mysqlconnector/adk/store.py index 46ef18acf..0939abbd0 100644 --- a/sqlspec/adapters/mysqlconnector/adk/store.py +++ b/sqlspec/adapters/mysqlconnector/adk/store.py @@ -30,6 +30,104 @@ MYSQL_TABLE_NOT_FOUND_ERROR: Final = 1146 +_ADK_METADATA_SEED_SQL_TEMPLATE = "INSERT IGNORE INTO {0} (`key`, value) VALUES ('schema_version', '1')" + +_ADK_MEMORY_TABLE_DDL_TEMPLATE = ",\n {0}" + +_ADK_MEMORY_TABLE_DDL_TEMPLATE_2 = ",\n FULLTEXT INDEX idx_{0}_fts (content_text)" + +_ADK_MEMORY_TABLE_DDL_TEMPLATE_3 = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " session_id VARCHAR(128) NOT NULL,\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL,\n" + " event_id VARCHAR(128) NOT NULL UNIQUE,\n" + " author VARCHAR(256){1},\n" + " timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),\n" + " content_json JSON NOT NULL,\n" + " content_text TEXT NOT NULL,\n" + " metadata_json JSON,\n" + " inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),\n" + " INDEX idx_{2}_app_user_time (app_name, user_id, timestamp),\n" + " INDEX idx_{3}_session (session_id){4}{5}\n" + " ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{6}\n" + " " +) + +_ADK_MYSQL_SESSIONS_DDL_TEMPLATE = "\n {0}," + +_ADK_MYSQL_SESSIONS_DDL_TEMPLATE_2 = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL,{1}\n" + " state JSON NOT NULL,\n" + " create_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),\n" + " update_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),\n" + " INDEX idx_{2}_app_user (app_name, user_id),\n" + " INDEX idx_{3}_update_time (update_time DESC){4}\n" + " ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{5}\n" + " " +) + +_ADK_MYSQL_EVENTS_DDL_TEMPLATE = ( + ",\n" + " INDEX idx_{0}_author_gc (session_id, author_gc, timestamp ASC),\n" + " INDEX idx_{1}_node_path_gc (session_id, node_path_gc, timestamp ASC)" +) + +_ADK_MYSQL_EVENTS_DDL_TEMPLATE_2 = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL,\n" + " session_id VARCHAR(128) NOT NULL,\n" + " invocation_id VARCHAR(256) NOT NULL,\n" + " timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),\n" + " event_data JSON NOT NULL{1},\n" + " FOREIGN KEY (session_id) REFERENCES {2}(id) ON DELETE CASCADE,\n" + " INDEX idx_{3}_scope (app_name, user_id, session_id, timestamp ASC{4}),\n" + " INDEX idx_{5}_session (session_id, timestamp ASC{6}){7}\n" + " ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{8}\n" + " " +) + +_ADK_MYSQL_APP_STATE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " app_name VARCHAR(128) PRIMARY KEY,\n" + " state JSON NOT NULL,\n" + " update_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6)\n" + " ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{1}\n" + " " +) + +_ADK_MYSQL_USER_STATE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL,\n" + " state JSON NOT NULL,\n" + " update_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),\n" + " PRIMARY KEY (app_name, user_id)\n" + " ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{1}\n" + " " +) + +_ADK_MYSQL_METADATA_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " `key` VARCHAR(128) PRIMARY KEY,\n" + " value VARCHAR(512) NOT NULL\n" + " ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci\n" + " " +) + + class MysqlConnectorADKConfig(ADKConfig): """mysql-connector-specific ADK extension settings. @@ -387,7 +485,7 @@ async def _metadata_table_ddl(self) -> str: return _mysql_metadata_ddl(self._metadata_table) async def _metadata_seed_sql(self) -> str: - return f"INSERT IGNORE INTO {self._metadata_table} (`key`, value) VALUES ('schema_version', '1')" + return _ADK_METADATA_SEED_SQL_TEMPLATE.format(self._metadata_table) def _drop_app_states_table_sql(self) -> str: return f"DROP TABLE IF EXISTS {self._app_state_table}" @@ -751,7 +849,7 @@ def _metadata_table_ddl(self) -> str: return _mysql_metadata_ddl(self._metadata_table) def _metadata_seed_sql(self) -> str: - return f"INSERT IGNORE INTO {self._metadata_table} (`key`, value) VALUES ('schema_version', '1')" + return _ADK_METADATA_SEED_SQL_TEMPLATE.format(self._metadata_table) def _drop_app_states_table_sql(self) -> str: return f"DROP TABLE IF EXISTS {self._app_state_table}" @@ -932,32 +1030,24 @@ async def _memory_table_ddl(self) -> str: fk_constraint = "" if self._owner_id_column_ddl: col_def, fk_def = _mysql_owner_id_column_parts(self._owner_id_column_ddl) - owner_id_line = f",\n {col_def}" + owner_id_line = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(col_def) if fk_def: - fk_constraint = f",\n {fk_def}" + fk_constraint = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(fk_def) fts_index = "" if self._use_fts: - fts_index = f",\n FULLTEXT INDEX idx_{self._memory_table}_fts (content_text)" + fts_index = _ADK_MEMORY_TABLE_DDL_TEMPLATE_2.format(self._memory_table) table_options = _mysql_table_options(adk_config, "memory_table_options") - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - content_json JSON NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSON, - inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - INDEX idx_{self._memory_table}_app_user_time (app_name, user_id, timestamp), - INDEX idx_{self._memory_table}_session (session_id){fts_index}{fk_constraint} - ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} - """ + return _ADK_MEMORY_TABLE_DDL_TEMPLATE_3.format( + self._memory_table, + owner_id_line, + self._memory_table, + self._memory_table, + fts_index, + fk_constraint, + table_options, + ) def _drop_memory_table_sql(self) -> "list[str]": return [f"DROP TABLE IF EXISTS {self._memory_table}"] @@ -1128,32 +1218,24 @@ def _memory_table_ddl(self) -> str: fk_constraint = "" if self._owner_id_column_ddl: col_def, fk_def = _mysql_owner_id_column_parts(self._owner_id_column_ddl) - owner_id_line = f",\n {col_def}" + owner_id_line = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(col_def) if fk_def: - fk_constraint = f",\n {fk_def}" + fk_constraint = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(fk_def) fts_index = "" if self._use_fts: - fts_index = f",\n FULLTEXT INDEX idx_{self._memory_table}_fts (content_text)" + fts_index = _ADK_MEMORY_TABLE_DDL_TEMPLATE_2.format(self._memory_table) table_options = _mysql_table_options(adk_config, "memory_table_options") - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - content_json JSON NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSON, - inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - INDEX idx_{self._memory_table}_app_user_time (app_name, user_id, timestamp), - INDEX idx_{self._memory_table}_session (session_id){fts_index}{fk_constraint} - ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} - """ + return _ADK_MEMORY_TABLE_DDL_TEMPLATE_3.format( + self._memory_table, + owner_id_line, + self._memory_table, + self._memory_table, + fts_index, + fk_constraint, + table_options, + ) def _drop_memory_table_sql(self) -> "list[str]": return [f"DROP TABLE IF EXISTS {self._memory_table}"] @@ -1361,22 +1443,13 @@ def _mysql_sessions_ddl(session_table: str, owner_id_column_ddl: "str | None", t fk_constraint = "" if owner_id_column_ddl: col_def, fk_def = _mysql_owner_id_column_parts(owner_id_column_ddl) - owner_id_line = f"\n {col_def}," + owner_id_line = _ADK_MYSQL_SESSIONS_DDL_TEMPLATE.format(col_def) if fk_def: - fk_constraint = f",\n {fk_def}" + fk_constraint = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(fk_def) - return f""" - CREATE TABLE IF NOT EXISTS {session_table} ( - id VARCHAR(128) PRIMARY KEY, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL,{owner_id_line} - state JSON NOT NULL, - create_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - update_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6), - INDEX idx_{session_table}_app_user (app_name, user_id), - INDEX idx_{session_table}_update_time (update_time DESC){fk_constraint} - ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} - """ + return _ADK_MYSQL_SESSIONS_DDL_TEMPLATE_2.format( + session_table, owner_id_line, session_table, session_table, fk_constraint, table_options + ) def _mysql_events_ddl(events_table: str, session_table: str, adk_config: Mapping[str, Any] | None = None) -> str: @@ -1387,58 +1460,34 @@ def _mysql_events_ddl(events_table: str, session_table: str, adk_config: Mapping generated_columns = """, author_gc VARCHAR(256) GENERATED ALWAYS AS (JSON_UNQUOTE(JSON_EXTRACT(event_data, '$.author'))) STORED, node_path_gc VARCHAR(512) GENERATED ALWAYS AS (JSON_UNQUOTE(JSON_EXTRACT(event_data, '$.node_info.path'))) STORED""" - generated_indexes = f""", - INDEX idx_{events_table}_author_gc (session_id, author_gc, timestamp ASC), - INDEX idx_{events_table}_node_path_gc (session_id, node_path_gc, timestamp ASC)""" + generated_indexes = _ADK_MYSQL_EVENTS_DDL_TEMPLATE.format(events_table, events_table) covering_column = ", invocation_id" if adk_config.get("enable_covering_indexes", False) else "" table_options = _mysql_table_options(adk_config, "events_table_options") - return f""" - CREATE TABLE IF NOT EXISTS {events_table} ( - id VARCHAR(128) PRIMARY KEY, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - session_id VARCHAR(128) NOT NULL, - invocation_id VARCHAR(256) NOT NULL, - timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - event_data JSON NOT NULL{generated_columns}, - FOREIGN KEY (session_id) REFERENCES {session_table}(id) ON DELETE CASCADE, - INDEX idx_{events_table}_scope (app_name, user_id, session_id, timestamp ASC{covering_column}), - INDEX idx_{events_table}_session (session_id, timestamp ASC{covering_column}){generated_indexes} - ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} - """ + return _ADK_MYSQL_EVENTS_DDL_TEMPLATE_2.format( + events_table, + generated_columns, + session_table, + events_table, + covering_column, + events_table, + covering_column, + generated_indexes, + table_options, + ) def _mysql_app_state_ddl(app_state_table: str, table_options: str = "") -> str: - return f""" - CREATE TABLE IF NOT EXISTS {app_state_table} ( - app_name VARCHAR(128) PRIMARY KEY, - state JSON NOT NULL, - update_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6) - ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} - """ + return _ADK_MYSQL_APP_STATE_DDL_TEMPLATE.format(app_state_table, table_options) def _mysql_user_state_ddl(user_state_table: str, table_options: str = "") -> str: - return f""" - CREATE TABLE IF NOT EXISTS {user_state_table} ( - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - state JSON NOT NULL, - update_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6), - PRIMARY KEY (app_name, user_id) - ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} - """ + return _ADK_MYSQL_USER_STATE_DDL_TEMPLATE.format(user_state_table, table_options) def _mysql_metadata_ddl(metadata_table: str) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {metadata_table} ( - `key` VARCHAR(128) PRIMARY KEY, - value VARCHAR(512) NOT NULL - ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci - """ + return _ADK_MYSQL_METADATA_DDL_TEMPLATE.format(metadata_table) def _mysql_upsert_app_state_sql(app_state_table: str) -> str: diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index c088f51b7..128967057 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -72,6 +72,232 @@ OracleDatabaseError: Final[type[Exception]] = cast("type[Exception]", oracledb.DatabaseError) +_ADK_METADATA_TABLE_DDL_TEMPLATE = ( + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE TABLE {0} (\n" + " key VARCHAR2(128) PRIMARY KEY,\n" + " value VARCHAR2(512) NOT NULL\n" + " )';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + " " +) + +_ADK_METADATA_SEED_SQL_TEMPLATE = ( + "\n" + " BEGIN\n" + " INSERT INTO {0} (key, value)\n" + " SELECT 'schema_version', '1'\n" + " FROM DUAL\n" + " WHERE NOT EXISTS (\n" + " SELECT 1 FROM {1} WHERE key = 'schema_version'\n" + " );\n" + " END;\n" + " " +) + +_ADK_SESSIONS_TABLE_DDL_FOR_TYPE_TEMPLATE = ", {0}" + +_ADK_SESSIONS_TABLE_DDL_FOR_TYPE_TEMPLATE_2 = ( + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE TABLE {0} (\n" + " id VARCHAR2(128) PRIMARY KEY,\n" + " app_name VARCHAR2(128) NOT NULL,\n" + " user_id VARCHAR2(128) NOT NULL,\n" + " {1},\n" + " create_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,\n" + " update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL{2}\n" + " ){3}';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{4}_app_user\n" + " ON {5}(app_name, user_id)';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{6}_update_time\n" + " ON {7}(update_time DESC)';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + " " +) + +_ADK_EVENTS_TABLE_DDL_FOR_TYPE_TEMPLATE = ( + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE TABLE {0} (\n" + " id VARCHAR2(128) PRIMARY KEY,\n" + " session_id VARCHAR2(128) NOT NULL,\n" + " invocation_id VARCHAR2(256),\n" + " timestamp TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,\n" + " {1},\n" + " CONSTRAINT fk_{2}_session FOREIGN KEY (session_id)\n" + " REFERENCES {3}(id) ON DELETE CASCADE\n" + " ){4}';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{5}_session\n" + " ON {6}(session_id, timestamp ASC)';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{7}_invocation\n" + " ON {8}(invocation_id)';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{9}_timestamp\n" + " ON {10}(timestamp ASC)';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + " " +) + +_ADK_APP_STATES_TABLE_DDL_FOR_TYPE_TEMPLATE = ( + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE TABLE {0} (\n" + " app_name VARCHAR2(128) PRIMARY KEY,\n" + " {1},\n" + " update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL\n" + " )';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + " " +) + +_ADK_USER_STATES_TABLE_DDL_FOR_TYPE_TEMPLATE = ( + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE TABLE {0} (\n" + " app_name VARCHAR2(128) NOT NULL,\n" + " user_id VARCHAR2(128) NOT NULL,\n" + " {1},\n" + " update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,\n" + " PRIMARY KEY (app_name, user_id)\n" + " )';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + " " +) + +_ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE = ",\n {0}" + +_ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE_2 = ( + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{0}_fts\n" + " ON {1}(content_text) INDEXTYPE IS CTXSYS.CONTEXT';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + " " +) + +_ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE_3 = ( + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE TABLE {0} (\n" + " id VARCHAR2(128) PRIMARY KEY,\n" + " session_id VARCHAR2(128) NOT NULL,\n" + " app_name VARCHAR2(128) NOT NULL,\n" + " user_id VARCHAR2(128) NOT NULL,\n" + " event_id VARCHAR2(128) NOT NULL UNIQUE,\n" + " author VARCHAR2(256){1},\n" + " timestamp TIMESTAMP WITH TIME ZONE NOT NULL,\n" + " {2},\n" + " content_text CLOB NOT NULL,\n" + " inserted_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL\n" + " ){3}';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{4}_app_user_time\n" + " ON {5}(app_name, user_id, timestamp DESC)';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{6}_session\n" + " ON {7}(session_id)';\n" + " EXCEPTION\n" + " WHEN OTHERS THEN\n" + " IF SQLCODE != -955 THEN\n" + " RAISE;\n" + " END IF;\n" + " END;\n" + " {8}\n" + " " +) + +_ADK_JSON_COLUMN_DDL_TEMPLATE = "{0} JSON NOT NULL" + +_ADK_JSON_COLUMN_DDL_TEMPLATE_2 = "{0} BLOB CHECK ({1} IS JSON) NOT NULL" + +_ADK_JSON_COLUMN_DDL_TEMPLATE_3 = "{0} BLOB NOT NULL" + + class JSONStorageType(str, Enum): """JSON storage type based on Oracle version.""" @@ -810,32 +1036,11 @@ async def _user_states_table_ddl(self) -> str: async def _metadata_table_ddl(self) -> str: """Get Oracle CREATE TABLE SQL for ADK internal metadata.""" - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._metadata_table} ( - key VARCHAR2(128) PRIMARY KEY, - value VARCHAR2(512) NOT NULL - )'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_METADATA_TABLE_DDL_TEMPLATE.format(self._metadata_table) async def _metadata_seed_sql(self) -> str: """Get Oracle SQL to seed the ADK schema-version metadata row.""" - return f""" - BEGIN - INSERT INTO {self._metadata_table} (key, value) - SELECT 'schema_version', '1' - FROM DUAL - WHERE NOT EXISTS ( - SELECT 1 FROM {self._metadata_table} WHERE key = 'schema_version' - ); - END; - """ + return _ADK_METADATA_SEED_SQL_TEMPLATE.format(self._metadata_table, self._metadata_table) async def _detect_json_storage_type(self) -> JSONStorageType: """Detect the appropriate JSON storage type based on Oracle version. @@ -970,7 +1175,11 @@ def _sessions_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: else: state_column = "state BLOB NOT NULL" - owner_id_column_sql = f", {self._owner_id_column_ddl}" if self._owner_id_column_ddl else "" + owner_id_column_sql = ( + _ADK_SESSIONS_TABLE_DDL_FOR_TYPE_TEMPLATE.format(self._owner_id_column_ddl) + if self._owner_id_column_ddl + else "" + ) table_clauses = _oracle_table_feature_clauses( self._config, "session", @@ -979,43 +1188,16 @@ def _sessions_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: range_partition_key="create_time", ) - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._session_table} ( - id VARCHAR2(128) PRIMARY KEY, - app_name VARCHAR2(128) NOT NULL, - user_id VARCHAR2(128) NOT NULL, - {state_column}, - create_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL{owner_id_column_sql} - ){table_clauses}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._session_table}_app_user - ON {self._session_table}(app_name, user_id)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._session_table}_update_time - ON {self._session_table}(update_time DESC)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_SESSIONS_TABLE_DDL_FOR_TYPE_TEMPLATE_2.format( + self._session_table, + state_column, + owner_id_column_sql, + table_clauses, + self._session_table, + self._session_table, + self._session_table, + self._session_table, + ) def _events_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: """Get Oracle CREATE TABLE SQL for events with specified storage type. @@ -1038,94 +1220,31 @@ def _events_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: range_partition_key="timestamp", ) - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._events_table} ( - id VARCHAR2(128) PRIMARY KEY, - session_id VARCHAR2(128) NOT NULL, - invocation_id VARCHAR2(256), - timestamp TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - {event_data_col}, - CONSTRAINT fk_{self._events_table}_session FOREIGN KEY (session_id) - REFERENCES {self._session_table}(id) ON DELETE CASCADE - ){table_clauses}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._events_table}_session - ON {self._events_table}(session_id, timestamp ASC)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._events_table}_invocation - ON {self._events_table}(invocation_id)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._events_table}_timestamp - ON {self._events_table}(timestamp ASC)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_EVENTS_TABLE_DDL_FOR_TYPE_TEMPLATE.format( + self._events_table, + event_data_col, + self._events_table, + self._session_table, + table_clauses, + self._events_table, + self._events_table, + self._events_table, + self._events_table, + self._events_table, + self._events_table, + ) def _app_states_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: """Get Oracle CREATE TABLE SQL for app-scoped state with specified storage type.""" state_column = _json_column_ddl("state", storage_type) - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._app_state_table} ( - app_name VARCHAR2(128) PRIMARY KEY, - {state_column}, - update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL - )'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_APP_STATES_TABLE_DDL_FOR_TYPE_TEMPLATE.format(self._app_state_table, state_column) def _user_states_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: """Get Oracle CREATE TABLE SQL for user-scoped state with specified storage type.""" state_column = _json_column_ddl("state", storage_type) - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._user_state_table} ( - app_name VARCHAR2(128) NOT NULL, - user_id VARCHAR2(128) NOT NULL, - {state_column}, - update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - PRIMARY KEY (app_name, user_id) - )'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_USER_STATES_TABLE_DDL_FOR_TYPE_TEMPLATE.format(self._user_state_table, state_column) def _drop_app_states_table_sql(self) -> str: return f""" @@ -1884,32 +2003,11 @@ def _user_states_table_ddl(self) -> str: def _metadata_table_ddl(self) -> str: """Get Oracle CREATE TABLE SQL for ADK internal metadata.""" - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._metadata_table} ( - key VARCHAR2(128) PRIMARY KEY, - value VARCHAR2(512) NOT NULL - )'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_METADATA_TABLE_DDL_TEMPLATE.format(self._metadata_table) def _metadata_seed_sql(self) -> str: """Get Oracle SQL to seed the ADK schema-version metadata row.""" - return f""" - BEGIN - INSERT INTO {self._metadata_table} (key, value) - SELECT 'schema_version', '1' - FROM DUAL - WHERE NOT EXISTS ( - SELECT 1 FROM {self._metadata_table} WHERE key = 'schema_version' - ); - END; - """ + return _ADK_METADATA_SEED_SQL_TEMPLATE.format(self._metadata_table, self._metadata_table) def _detect_json_storage_type(self) -> JSONStorageType: """Detect the appropriate JSON storage type based on Oracle version. @@ -2040,7 +2138,11 @@ def _sessions_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: else: state_column = "state BLOB NOT NULL" - owner_id_column_sql = f", {self._owner_id_column_ddl}" if self._owner_id_column_ddl else "" + owner_id_column_sql = ( + _ADK_SESSIONS_TABLE_DDL_FOR_TYPE_TEMPLATE.format(self._owner_id_column_ddl) + if self._owner_id_column_ddl + else "" + ) table_clauses = _oracle_table_feature_clauses( self._config, "session", @@ -2049,43 +2151,16 @@ def _sessions_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: range_partition_key="create_time", ) - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._session_table} ( - id VARCHAR2(128) PRIMARY KEY, - app_name VARCHAR2(128) NOT NULL, - user_id VARCHAR2(128) NOT NULL, - {state_column}, - create_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL{owner_id_column_sql} - ){table_clauses}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._session_table}_app_user - ON {self._session_table}(app_name, user_id)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._session_table}_update_time - ON {self._session_table}(update_time DESC)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_SESSIONS_TABLE_DDL_FOR_TYPE_TEMPLATE_2.format( + self._session_table, + state_column, + owner_id_column_sql, + table_clauses, + self._session_table, + self._session_table, + self._session_table, + self._session_table, + ) def _events_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: """Get Oracle CREATE TABLE SQL for events with specified storage type. @@ -2108,94 +2183,31 @@ def _events_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: range_partition_key="timestamp", ) - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._events_table} ( - id VARCHAR2(128) PRIMARY KEY, - session_id VARCHAR2(128) NOT NULL, - invocation_id VARCHAR2(256), - timestamp TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - {event_data_col}, - CONSTRAINT fk_{self._events_table}_session FOREIGN KEY (session_id) - REFERENCES {self._session_table}(id) ON DELETE CASCADE - ){table_clauses}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._events_table}_session - ON {self._events_table}(session_id, timestamp ASC)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._events_table}_invocation - ON {self._events_table}(invocation_id)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._events_table}_timestamp - ON {self._events_table}(timestamp ASC)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_EVENTS_TABLE_DDL_FOR_TYPE_TEMPLATE.format( + self._events_table, + event_data_col, + self._events_table, + self._session_table, + table_clauses, + self._events_table, + self._events_table, + self._events_table, + self._events_table, + self._events_table, + self._events_table, + ) def _app_states_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: """Get Oracle CREATE TABLE SQL for app-scoped state with specified storage type.""" state_column = _json_column_ddl("state", storage_type) - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._app_state_table} ( - app_name VARCHAR2(128) PRIMARY KEY, - {state_column}, - update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL - )'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_APP_STATES_TABLE_DDL_FOR_TYPE_TEMPLATE.format(self._app_state_table, state_column) def _user_states_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: """Get Oracle CREATE TABLE SQL for user-scoped state with specified storage type.""" state_column = _json_column_ddl("state", storage_type) - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._user_state_table} ( - app_name VARCHAR2(128) NOT NULL, - user_id VARCHAR2(128) NOT NULL, - {state_column}, - update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - PRIMARY KEY (app_name, user_id) - )'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + return _ADK_USER_STATES_TABLE_DDL_FOR_TYPE_TEMPLATE.format(self._user_state_table, state_column) def _drop_app_states_table_sql(self) -> str: return f""" @@ -2464,7 +2476,11 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: metadata_json BLOB """ - owner_id_line = f",\n {self._owner_id_column_ddl}" if self._owner_id_column_ddl else "" + owner_id_line = ( + _ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE.format(self._owner_id_column_ddl) + if self._owner_id_column_ddl + else "" + ) table_clauses = _oracle_table_feature_clauses( self._config, "memory", @@ -2475,60 +2491,19 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: fts_index = "" if self._use_fts: - fts_index = f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._memory_table}_fts - ON {self._memory_table}(content_text) INDEXTYPE IS CTXSYS.CONTEXT'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ - - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._memory_table} ( - id VARCHAR2(128) PRIMARY KEY, - session_id VARCHAR2(128) NOT NULL, - app_name VARCHAR2(128) NOT NULL, - user_id VARCHAR2(128) NOT NULL, - event_id VARCHAR2(128) NOT NULL UNIQUE, - author VARCHAR2(256){owner_id_line}, - timestamp TIMESTAMP WITH TIME ZONE NOT NULL, - {json_columns}, - content_text CLOB NOT NULL, - inserted_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL - ){table_clauses}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._memory_table}_session - ON {self._memory_table}(session_id)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - {fts_index} - """ + fts_index = _ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE_2.format(self._memory_table, self._memory_table) + + return _ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE_3.format( + self._memory_table, + owner_id_line, + json_columns, + table_clauses, + self._memory_table, + self._memory_table, + self._memory_table, + self._memory_table, + fts_index, + ) def _drop_memory_table_sql(self) -> "list[str]": return [ @@ -2814,7 +2789,11 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: metadata_json BLOB """ - owner_id_line = f",\n {self._owner_id_column_ddl}" if self._owner_id_column_ddl else "" + owner_id_line = ( + _ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE.format(self._owner_id_column_ddl) + if self._owner_id_column_ddl + else "" + ) table_clauses = _oracle_table_feature_clauses( self._config, "memory", @@ -2825,60 +2804,19 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: fts_index = "" if self._use_fts: - fts_index = f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._memory_table}_fts - ON {self._memory_table}(content_text) INDEXTYPE IS CTXSYS.CONTEXT'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ - - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._memory_table} ( - id VARCHAR2(128) PRIMARY KEY, - session_id VARCHAR2(128) NOT NULL, - app_name VARCHAR2(128) NOT NULL, - user_id VARCHAR2(128) NOT NULL, - event_id VARCHAR2(128) NOT NULL UNIQUE, - author VARCHAR2(256){owner_id_line}, - timestamp TIMESTAMP WITH TIME ZONE NOT NULL, - {json_columns}, - content_text CLOB NOT NULL, - inserted_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL - ){table_clauses}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX idx_{self._memory_table}_session - ON {self._memory_table}(session_id)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - {fts_index} - """ + fts_index = _ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE_2.format(self._memory_table, self._memory_table) + + return _ADK_MEMORY_TABLE_DDL_FOR_TYPE_TEMPLATE_3.format( + self._memory_table, + owner_id_line, + json_columns, + table_clauses, + self._memory_table, + self._memory_table, + self._memory_table, + self._memory_table, + fts_index, + ) def _drop_memory_table_sql(self) -> "list[str]": return [ @@ -3093,10 +3031,10 @@ def _event_data_column_ddl(storage_type: JSONStorageType) -> str: def _json_column_ddl(column_name: str, storage_type: JSONStorageType) -> str: """Return an Oracle JSON column DDL fragment for the configured storage type.""" if storage_type == JSONStorageType.JSON_NATIVE: - return f"{column_name} JSON NOT NULL" + return _ADK_JSON_COLUMN_DDL_TEMPLATE.format(column_name) if storage_type == JSONStorageType.BLOB_JSON: - return f"{column_name} BLOB CHECK ({column_name} IS JSON) NOT NULL" - return f"{column_name} BLOB NOT NULL" + return _ADK_JSON_COLUMN_DDL_TEMPLATE_2.format(column_name, column_name) + return _ADK_JSON_COLUMN_DDL_TEMPLATE_3.format(column_name) def _adk_config(config: Any) -> OracleADKConfig: diff --git a/sqlspec/adapters/psycopg/adk/store.py b/sqlspec/adapters/psycopg/adk/store.py index 0173c14fb..74a6bc484 100644 --- a/sqlspec/adapters/psycopg/adk/store.py +++ b/sqlspec/adapters/psycopg/adk/store.py @@ -31,6 +31,131 @@ logger = get_logger("sqlspec.adapters.psycopg.adk.store") +_ADK_SESSIONS_TABLE_DDL_TEMPLATE = ",\n {0}" + +_ADK_SESSIONS_TABLE_DDL_TEMPLATE_2 = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL{1},\n" + " state JSONB NOT NULL DEFAULT '{{}}'::jsonb,\n" + " create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,\n" + " update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP\n" + " ) WITH (fillfactor = 80);\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{2}_app_user\n" + " ON {3}(app_name, user_id);\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{4}_update_time\n" + " ON {5}(update_time DESC);\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{6}_state\n" + " ON {7} USING GIN (state)\n" + " WHERE state != '{{}}'::jsonb;\n" + " " +) + +_ADK_EVENTS_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " session_id VARCHAR(128) NOT NULL,\n" + " invocation_id VARCHAR(256),\n" + " timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,\n" + " event_data JSONB NOT NULL{1},\n" + " FOREIGN KEY (session_id) REFERENCES {2}(id) ON DELETE CASCADE\n" + " ) WITH (fillfactor = 80);\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{3}_session\n" + " ON {4}(session_id, timestamp ASC){5};\n" + " {6}\n" + " " +) + +_ADK_APP_STATES_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " app_name VARCHAR(128) PRIMARY KEY,\n" + " state JSONB NOT NULL DEFAULT '{{}}'::jsonb,\n" + " update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP\n" + " ) WITH (fillfactor = 80);\n" + " " +) + +_ADK_USER_STATES_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL,\n" + " state JSONB NOT NULL DEFAULT '{{}}'::jsonb,\n" + " update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,\n" + " PRIMARY KEY (app_name, user_id)\n" + " ) WITH (fillfactor = 80);\n" + " " +) + +_ADK_METADATA_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " key VARCHAR(128) PRIMARY KEY,\n" + " value VARCHAR(512) NOT NULL\n" + " );\n" + " " +) + +_ADK_METADATA_SEED_SQL_TEMPLATE = ( + "\n" + " INSERT INTO {0} (key, value)\n" + " VALUES ('schema_version', '1')\n" + " ON CONFLICT (key) DO NOTHING\n" + " " +) + +_ADK_MEMORY_TABLE_DDL_TEMPLATE = ( + "\n" + " CREATE INDEX IF NOT EXISTS idx_{0}_fts\n" + " ON {1} USING GIN (to_tsvector('english', content_text));\n" + " " +) + +_ADK_MEMORY_TABLE_DDL_TEMPLATE_2 = ( + "\n" + " CREATE TABLE IF NOT EXISTS {0} (\n" + " id VARCHAR(128) PRIMARY KEY,\n" + " session_id VARCHAR(128) NOT NULL,\n" + " app_name VARCHAR(128) NOT NULL,\n" + " user_id VARCHAR(128) NOT NULL,\n" + " event_id VARCHAR(128) NOT NULL UNIQUE,\n" + " author VARCHAR(256){1},\n" + " timestamp TIMESTAMPTZ NOT NULL,\n" + " content_json JSONB NOT NULL,\n" + " content_text TEXT NOT NULL,\n" + " metadata_json JSONB,\n" + " inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP\n" + " );\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{2}_app_user_time\n" + " ON {3}(app_name, user_id, timestamp DESC);\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{4}_session\n" + " ON {5}(session_id);\n" + " {6}\n" + " " +) + +_ADK_POSTGRES_EVENT_DDL_OPTIONS_TEMPLATE = ( + "\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{0}_author_gc\n" + " ON {1}(session_id, author_gc, timestamp ASC);\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{2}_node_path_gc\n" + " ON {3}(session_id, node_path_gc, timestamp ASC);\n" + " " +) + + class PsycopgADKConfig(ADKConfig): """Psycopg-specific ADK extension settings. @@ -458,28 +583,18 @@ async def set_metadata(self, key: str, value: str) -> None: async def _sessions_table_ddl(self) -> str: 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; - """ + owner_id_line = _ADK_SESSIONS_TABLE_DDL_TEMPLATE.format(self._owner_id_column_ddl) + + return _ADK_SESSIONS_TABLE_DDL_TEMPLATE_2.format( + self._session_table, + owner_id_line, + self._session_table, + self._session_table, + self._session_table, + self._session_table, + self._session_table, + self._session_table, + ) async def _events_table_ddl(self) -> str: adk_config = _adk_config(self._config) @@ -487,55 +602,27 @@ async def _events_table_ddl(self) -> str: adk_config, self._events_table ) - return f""" - CREATE TABLE IF NOT EXISTS {self._events_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - invocation_id VARCHAR(256), - timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - event_data JSONB NOT NULL{generated_columns}, - FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE - ) WITH (fillfactor = 80); - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session - ON {self._events_table}(session_id, timestamp ASC){covering_columns}; - {generated_indexes} - """ + return _ADK_EVENTS_TABLE_DDL_TEMPLATE.format( + self._events_table, + generated_columns, + self._session_table, + self._events_table, + self._events_table, + covering_columns, + generated_indexes, + ) async def _app_states_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._app_state_table} ( - app_name VARCHAR(128) PRIMARY KEY, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = 80); - """ + return _ADK_APP_STATES_TABLE_DDL_TEMPLATE.format(self._app_state_table) async def _user_states_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._user_state_table} ( - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY (app_name, user_id) - ) WITH (fillfactor = 80); - """ + return _ADK_USER_STATES_TABLE_DDL_TEMPLATE.format(self._user_state_table) async def _metadata_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._metadata_table} ( - key VARCHAR(128) PRIMARY KEY, - value VARCHAR(512) NOT NULL - ); - """ + return _ADK_METADATA_TABLE_DDL_TEMPLATE.format(self._metadata_table) async def _metadata_seed_sql(self) -> str: - return f""" - INSERT INTO {self._metadata_table} (key, value) - VALUES ('schema_version', '1') - ON CONFLICT (key) DO NOTHING - """ + return _ADK_METADATA_SEED_SQL_TEMPLATE.format(self._metadata_table) def _drop_app_states_table_sql(self) -> str: return f"DROP TABLE IF EXISTS {self._app_state_table}" @@ -954,28 +1041,18 @@ def set_metadata(self, key: str, value: str) -> None: def _sessions_table_ddl(self) -> str: 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; - """ + owner_id_line = _ADK_SESSIONS_TABLE_DDL_TEMPLATE.format(self._owner_id_column_ddl) + + return _ADK_SESSIONS_TABLE_DDL_TEMPLATE_2.format( + self._session_table, + owner_id_line, + self._session_table, + self._session_table, + self._session_table, + self._session_table, + self._session_table, + self._session_table, + ) def _events_table_ddl(self) -> str: adk_config = _adk_config(self._config) @@ -983,55 +1060,27 @@ def _events_table_ddl(self) -> str: adk_config, self._events_table ) - return f""" - CREATE TABLE IF NOT EXISTS {self._events_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - invocation_id VARCHAR(256), - timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - event_data JSONB NOT NULL{generated_columns}, - FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE - ) WITH (fillfactor = 80); - - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session - ON {self._events_table}(session_id, timestamp ASC){covering_columns}; - {generated_indexes} - """ + return _ADK_EVENTS_TABLE_DDL_TEMPLATE.format( + self._events_table, + generated_columns, + self._session_table, + self._events_table, + self._events_table, + covering_columns, + generated_indexes, + ) def _app_states_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._app_state_table} ( - app_name VARCHAR(128) PRIMARY KEY, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = 80); - """ + return _ADK_APP_STATES_TABLE_DDL_TEMPLATE.format(self._app_state_table) def _user_states_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._user_state_table} ( - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - state JSONB NOT NULL DEFAULT '{{}}'::jsonb, - update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY (app_name, user_id) - ) WITH (fillfactor = 80); - """ + return _ADK_USER_STATES_TABLE_DDL_TEMPLATE.format(self._user_state_table) def _metadata_table_ddl(self) -> str: - return f""" - CREATE TABLE IF NOT EXISTS {self._metadata_table} ( - key VARCHAR(128) PRIMARY KEY, - value VARCHAR(512) NOT NULL - ); - """ + return _ADK_METADATA_TABLE_DDL_TEMPLATE.format(self._metadata_table) def _metadata_seed_sql(self) -> str: - return f""" - INSERT INTO {self._metadata_table} (key, value) - VALUES ('schema_version', '1') - ON CONFLICT (key) DO NOTHING - """ + return _ADK_METADATA_SEED_SQL_TEMPLATE.format(self._metadata_table) def _drop_app_states_table_sql(self) -> str: return f"DROP TABLE IF EXISTS {self._app_state_table}" @@ -1184,37 +1233,21 @@ async def _memory_table_ddl(self) -> str: """Get PostgreSQL CREATE TABLE SQL for memory entries.""" owner_id_line = "" if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" + owner_id_line = _ADK_SESSIONS_TABLE_DDL_TEMPLATE.format(self._owner_id_column_ddl) fts_index = "" if self._use_fts: - fts_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts - ON {self._memory_table} USING GIN (to_tsvector('english', content_text)); - """ - - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMPTZ NOT NULL, - content_json JSONB NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSONB, - inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session - ON {self._memory_table}(session_id); - {fts_index} - """ + fts_index = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(self._memory_table, self._memory_table) + + return _ADK_MEMORY_TABLE_DDL_TEMPLATE_2.format( + self._memory_table, + owner_id_line, + self._memory_table, + self._memory_table, + self._memory_table, + self._memory_table, + fts_index, + ) def _drop_memory_table_sql(self) -> "list[str]": """Get PostgreSQL DROP TABLE SQL statements.""" @@ -1375,37 +1408,21 @@ def _memory_table_ddl(self) -> str: """Get PostgreSQL CREATE TABLE SQL for memory entries.""" owner_id_line = "" if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" + owner_id_line = _ADK_SESSIONS_TABLE_DDL_TEMPLATE.format(self._owner_id_column_ddl) fts_index = "" if self._use_fts: - fts_index = f""" - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_fts - ON {self._memory_table} USING GIN (to_tsvector('english', content_text)); - """ - - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMPTZ NOT NULL, - content_json JSONB NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSONB, - inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session - ON {self._memory_table}(session_id); - {fts_index} - """ + fts_index = _ADK_MEMORY_TABLE_DDL_TEMPLATE.format(self._memory_table, self._memory_table) + + return _ADK_MEMORY_TABLE_DDL_TEMPLATE_2.format( + self._memory_table, + owner_id_line, + self._memory_table, + self._memory_table, + self._memory_table, + self._memory_table, + fts_index, + ) def _drop_memory_table_sql(self) -> "list[str]": """Get PostgreSQL DROP TABLE SQL statements.""" @@ -1523,14 +1540,9 @@ def _postgres_event_ddl_options(adk_config: PsycopgADKConfig, events_table: str) generated_columns = """, author_gc VARCHAR(256) GENERATED ALWAYS AS (event_data->>'author') STORED, node_path_gc TEXT GENERATED ALWAYS AS (event_data->'node_info'->>'path') STORED""" - generated_indexes = f""" - - CREATE INDEX IF NOT EXISTS idx_{events_table}_author_gc - ON {events_table}(session_id, author_gc, timestamp ASC); - - CREATE INDEX IF NOT EXISTS idx_{events_table}_node_path_gc - ON {events_table}(session_id, node_path_gc, timestamp ASC); - """ + generated_indexes = _ADK_POSTGRES_EVENT_DDL_OPTIONS_TEMPLATE.format( + events_table, events_table, events_table, events_table + ) covering_columns = "" if adk_config.get("enable_covering_indexes", False): diff --git a/tests/unit/extensions/test_adk/test_sql_template_hoist.py b/tests/unit/extensions/test_adk/test_sql_template_hoist.py new file mode 100644 index 000000000..21f988d4a --- /dev/null +++ b/tests/unit/extensions/test_adk/test_sql_template_hoist.py @@ -0,0 +1,30 @@ +"""Regression tests for adapter ADK SQL template ownership.""" + +import ast +import importlib +import inspect + +import pytest + + +@pytest.mark.parametrize( + "module_name", + [ + "sqlspec.adapters.psycopg.adk.store", + "sqlspec.adapters.cockroach_psycopg.adk.store", + "sqlspec.adapters.mysqlconnector.adk.store", + "sqlspec.adapters.oracledb.adk.store", + ], +) +def test_adk_ddl_methods_reference_module_templates(module_name: str) -> None: + module = importlib.import_module(module_name) + tree = ast.parse(inspect.getsource(module)) + + sql_owners = [ + node + for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and ("ddl" in node.name or "seed" in node.name) + ] + inline_templates = [node for owner in sql_owners for node in ast.walk(owner) if isinstance(node, ast.JoinedStr)] + + assert not inline_templates From c2dd2d5c9ede2424e3c900b338a4933ae30f8003 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 11 Jul 2026 01:02:36 +0000 Subject: [PATCH 5/6] style(adk): apply repository formatting --- sqlspec/extensions/adk/_config_utils.py | 4 +--- sqlspec/extensions/adk/artifact/service.py | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/sqlspec/extensions/adk/_config_utils.py b/sqlspec/extensions/adk/_config_utils.py index eb1110bd4..68c9badb3 100644 --- a/sqlspec/extensions/adk/_config_utils.py +++ b/sqlspec/extensions/adk/_config_utils.py @@ -103,9 +103,7 @@ def _adk_artifact_store_config(config: _ADKConfigSource) -> _ADKArtifactStoreCon return {"artifact_table": str(adk_config.get("artifact_table") or "adk_artifact")} -def _apply_owner_id( - result: "_ADKSessionStoreConfig | _ADKMemoryStoreConfig", adk_config: dict[str, Any] -) -> None: +def _apply_owner_id(result: "_ADKSessionStoreConfig | _ADKMemoryStoreConfig", adk_config: dict[str, Any]) -> None: """Copy the configured owner column into normalized store settings.""" owner_id = adk_config.get("owner_id_column") if owner_id is not None: diff --git a/sqlspec/extensions/adk/artifact/service.py b/sqlspec/extensions/adk/artifact/service.py index 99d792dd5..439485706 100644 --- a/sqlspec/extensions/adk/artifact/service.py +++ b/sqlspec/extensions/adk/artifact/service.py @@ -451,6 +451,7 @@ async def list_artifact_versions( app_name=app_name, user_id=user_id, filename=filename, session_id=session_id ) return [_record_to_artifact_version(r) for r in records] + async def get_artifact_version( self, *, From e743f101c7fd059f2b4485050037d4571a44ae93 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 11 Jul 2026 01:22:08 +0000 Subject: [PATCH 6/6] test(adk): strengthen consolidation contracts --- sqlspec/adapters/oracledb/adk/store.py | 2 +- .../test_adk/test_sql_template_hoist.py | 202 ++++++++++++++++++ .../test_adk/test_store_common_mixins.py | 92 ++++++++ 3 files changed, 295 insertions(+), 1 deletion(-) diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index 128967057..e7868a50a 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -2147,7 +2147,7 @@ def _sessions_table_ddl_for_type(self, storage_type: JSONStorageType) -> str: self._config, "session", in_memory=self._in_memory, - hash_partition_key="session_id", + hash_partition_key="id", range_partition_key="create_time", ) diff --git a/tests/unit/extensions/test_adk/test_sql_template_hoist.py b/tests/unit/extensions/test_adk/test_sql_template_hoist.py index 21f988d4a..588434774 100644 --- a/tests/unit/extensions/test_adk/test_sql_template_hoist.py +++ b/tests/unit/extensions/test_adk/test_sql_template_hoist.py @@ -1,11 +1,67 @@ +# pyright: reportPrivateUsage=false """Regression tests for adapter ADK SQL template ownership.""" import ast import importlib import inspect +from unittest.mock import MagicMock import pytest +from sqlspec.adapters.cockroach_psycopg.adk import ( + CockroachPsycopgAsyncADKMemoryStore, + CockroachPsycopgAsyncADKStore, + CockroachPsycopgSyncADKMemoryStore, + CockroachPsycopgSyncADKStore, +) +from sqlspec.adapters.mysqlconnector.adk import ( + MysqlConnectorAsyncADKMemoryStore, + MysqlConnectorAsyncADKStore, + MysqlConnectorSyncADKMemoryStore, + MysqlConnectorSyncADKStore, +) +from sqlspec.adapters.oracledb.adk import ( + JSONStorageType, + OracleAsyncADKMemoryStore, + OracleAsyncADKStore, + OracleSyncADKMemoryStore, + OracleSyncADKStore, +) +from sqlspec.adapters.psycopg.adk import ( + PsycopgAsyncADKMemoryStore, + PsycopgAsyncADKStore, + PsycopgSyncADKMemoryStore, + PsycopgSyncADKStore, +) + + +def _mock_config(adk_config: dict[str, object]) -> MagicMock: + config = MagicMock() + config.extension_config = {"adk": adk_config} + return config + + +async def _session_ddl(store: object) -> dict[str, str]: + return { + "sessions": await store._sessions_table_ddl(), # type: ignore[attr-defined] + "events": await store._events_table_ddl(), # type: ignore[attr-defined] + "app_state": await store._app_states_table_ddl(), # type: ignore[attr-defined] + "user_state": await store._user_states_table_ddl(), # type: ignore[attr-defined] + "metadata": await store._metadata_table_ddl(), # type: ignore[attr-defined] + "seed": await store._metadata_seed_sql(), # type: ignore[attr-defined] + } + + +def _sync_session_ddl(store: object) -> dict[str, str]: + return { + "sessions": store._sessions_table_ddl(), # type: ignore[attr-defined] + "events": store._events_table_ddl(), # type: ignore[attr-defined] + "app_state": store._app_states_table_ddl(), # type: ignore[attr-defined] + "user_state": store._user_states_table_ddl(), # type: ignore[attr-defined] + "metadata": store._metadata_table_ddl(), # type: ignore[attr-defined] + "seed": store._metadata_seed_sql(), # type: ignore[attr-defined] + } + @pytest.mark.parametrize( "module_name", @@ -28,3 +84,149 @@ def test_adk_ddl_methods_reference_module_templates(module_name: str) -> None: inline_templates = [node for owner in sql_owners for node in ast.walk(owner) if isinstance(node, ast.JoinedStr)] assert not inline_templates + + +@pytest.mark.anyio +async def test_psycopg_templates_bind_identically_for_sync_and_async_stores() -> None: + config = _mock_config({ + "session_table": "agent_session", + "events_table": "agent_event", + "memory_table": "agent_memory", + "owner_id_column": "tenant_id UUID REFERENCES tenant(id)", + "enable_event_generated_columns": True, + "enable_covering_indexes": True, + }) + async_store = PsycopgAsyncADKStore(config) + sync_store = PsycopgSyncADKStore(config) + async_memory = PsycopgAsyncADKMemoryStore(config) + sync_memory = PsycopgSyncADKMemoryStore(config) + + async_ddl = await _session_ddl(async_store) + sync_ddl = _sync_session_ddl(sync_store) + async_memory_ddl = await async_memory._memory_table_ddl() + sync_memory_ddl = sync_memory._memory_table_ddl() + + assert async_ddl == sync_ddl + assert async_memory_ddl == sync_memory_ddl + assert "tenant_id UUID REFERENCES tenant(id)," in async_ddl["sessions"] + assert "tenant_id UUID REFERENCES tenant(id)," in async_memory_ddl + assert "author_gc VARCHAR(256) GENERATED ALWAYS AS (event_data->>'author') STORED" in async_ddl["events"] + assert "node_path_gc TEXT GENERATED ALWAYS AS (event_data->'node_info'->>'path') STORED" in async_ddl["events"] + assert "ON agent_event(session_id, timestamp ASC) INCLUDE (invocation_id)" in async_ddl["events"] + + +@pytest.mark.anyio +async def test_cockroach_templates_bind_identically_for_sync_and_async_stores() -> None: + config = _mock_config({ + "session_table": "agent_session", + "events_table": "agent_event", + "memory_table": "agent_memory", + "owner_id_column": "tenant_id UUID", + "table_locality": "LOCALITY GLOBAL", + "enable_hash_sharded_indexes": True, + "hash_shard_bucket_count": 8, + "enable_storing_indexes": True, + "enable_memory_trigram_index": True, + }) + async_store = CockroachPsycopgAsyncADKStore(config) + sync_store = CockroachPsycopgSyncADKStore(config) + async_memory = CockroachPsycopgAsyncADKMemoryStore(config) + sync_memory = CockroachPsycopgSyncADKMemoryStore(config) + + async_ddl = await _session_ddl(async_store) + sync_ddl = _sync_session_ddl(sync_store) + async_memory_ddl = await async_memory._memory_table_ddl() + sync_memory_ddl = sync_memory._memory_table_ddl() + + assert async_ddl == sync_ddl + assert async_memory_ddl == sync_memory_ddl + assert "tenant_id UUID," in async_ddl["sessions"] + assert "LOCALITY GLOBAL" in async_ddl["sessions"] + assert "USING HASH WITH (bucket_count = 8)" in async_ddl["events"] + assert "STORING (invocation_id, event_data)" in async_ddl["events"] + assert "ON agent_memory USING GIN (content_text gin_trgm_ops)" in async_memory_ddl + + +@pytest.mark.anyio +async def test_mysqlconnector_templates_bind_identically_for_sync_and_async_stores() -> None: + config = _mock_config({ + "session_table": "agent_session", + "events_table": "agent_event", + "memory_table": "agent_memory", + "owner_id_column": "tenant_id BIGINT", + "enable_event_generated_columns": True, + "enable_covering_indexes": True, + "events_table_options": "COMMENT='agent-events'", + "memory_table_options": "COMMENT='agent-memory'", + }) + async_store = MysqlConnectorAsyncADKStore(config) + sync_store = MysqlConnectorSyncADKStore(config) + async_memory = MysqlConnectorAsyncADKMemoryStore(config) + sync_memory = MysqlConnectorSyncADKMemoryStore(config) + + async_ddl = await _session_ddl(async_store) + sync_ddl = _sync_session_ddl(sync_store) + async_memory_ddl = await async_memory._memory_table_ddl() + sync_memory_ddl = sync_memory._memory_table_ddl() + + assert async_ddl == sync_ddl + assert async_memory_ddl == sync_memory_ddl + assert "tenant_id BIGINT," in async_ddl["sessions"] + assert "tenant_id BIGINT," in async_memory_ddl + assert "author_gc VARCHAR(256) GENERATED ALWAYS AS" in async_ddl["events"] + assert "INDEX idx_agent_event_session (session_id, timestamp ASC, invocation_id)" in async_ddl["events"] + assert "COMMENT='agent-events'" in async_ddl["events"] + assert "COMMENT='agent-memory'" in async_memory_ddl + + +@pytest.mark.anyio +async def test_oracle_templates_bind_identically_for_every_json_storage_type() -> None: + config = _mock_config({ + "session_table": "agent_session", + "events_table": "agent_event", + "memory_table": "agent_memory", + "owner_id_column": "tenant_id NUMBER", + "partitioning": {"strategy": "hash", "partition_count": 8}, + }) + async_store = OracleAsyncADKStore(config) + sync_store = OracleSyncADKStore(config) + async_memory = OracleAsyncADKMemoryStore(config) + sync_memory = OracleSyncADKMemoryStore(config) + expected_columns = { + JSONStorageType.JSON_NATIVE: ("state JSON NOT NULL", "event_data JSON NOT NULL", "content_json JSON"), + JSONStorageType.BLOB_JSON: ( + "state BLOB CHECK (state IS JSON) NOT NULL", + "event_data BLOB CHECK (event_data IS JSON) NOT NULL", + "content_json BLOB CHECK (content_json IS JSON)", + ), + JSONStorageType.BLOB_PLAIN: ("state BLOB NOT NULL", "event_data BLOB NOT NULL", "content_json BLOB"), + } + + for storage_type, (state_column, event_column, memory_column) in expected_columns.items(): + async_session_ddl = async_store._sessions_table_ddl_for_type(storage_type) + sync_session_ddl = sync_store._sessions_table_ddl_for_type(storage_type) + async_events_ddl = async_store._events_table_ddl_for_type(storage_type) + sync_events_ddl = sync_store._events_table_ddl_for_type(storage_type) + async_memory_ddl = async_memory._memory_table_ddl_for_type(storage_type) + sync_memory_ddl = sync_memory._memory_table_ddl_for_type(storage_type) + + assert async_session_ddl == sync_session_ddl + assert async_events_ddl == sync_events_ddl + assert async_store._app_states_table_ddl_for_type(storage_type) == sync_store._app_states_table_ddl_for_type( + storage_type + ) + assert async_store._user_states_table_ddl_for_type(storage_type) == sync_store._user_states_table_ddl_for_type( + storage_type + ) + assert async_memory_ddl == sync_memory_ddl + assert state_column in async_session_ddl + assert event_column in async_events_ddl + assert memory_column in async_memory_ddl + assert ( + "update_time TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, tenant_id NUMBER" in async_session_ddl + ) + assert "tenant_id NUMBER," in async_memory_ddl + assert "PARTITION BY HASH (id) PARTITIONS 8" in async_session_ddl + + assert await async_store._metadata_table_ddl() == sync_store._metadata_table_ddl() + assert await async_store._metadata_seed_sql() == sync_store._metadata_seed_sql() diff --git a/tests/unit/extensions/test_adk/test_store_common_mixins.py b/tests/unit/extensions/test_adk/test_store_common_mixins.py index 6870c76f9..16d9d4c7f 100644 --- a/tests/unit/extensions/test_adk/test_store_common_mixins.py +++ b/tests/unit/extensions/test_adk/test_store_common_mixins.py @@ -1,5 +1,10 @@ """Regression tests for common ADK store state mixins.""" +import pytest + +from sqlspec.adapters.sqlite import SqliteConfig +from sqlspec.adapters.sqlite.adk import SqliteADKMemoryStore, SqliteADKStore +from sqlspec.extensions.adk.artifact._types import ArtifactRecord from sqlspec.extensions.adk.artifact.store import ( BaseAsyncADKArtifactStore, BaseSyncADKArtifactStore, @@ -9,6 +14,55 @@ from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore, _ADKStoreCommon +class _ConcreteArtifactStore(BaseSyncADKArtifactStore[SqliteConfig]): + __slots__ = () + + def insert_artifact(self, record: ArtifactRecord) -> None: + return None + + def get_artifact( + self, app_name: str, user_id: str, filename: str, session_id: str | None = None, version: int | None = None + ) -> ArtifactRecord | None: + return None + + def list_artifact_keys(self, app_name: str, user_id: str, session_id: str | None = None) -> list[str]: + return [] + + def list_artifact_versions( + self, app_name: str, user_id: str, filename: str, session_id: str | None = None + ) -> list[ArtifactRecord]: + return [] + + def delete_artifact( + self, app_name: str, user_id: str, filename: str, session_id: str | None = None + ) -> list[ArtifactRecord]: + return [] + + def get_next_version(self, app_name: str, user_id: str, filename: str, session_id: str | None = None) -> int: + return 0 + + def create_table(self) -> None: + return None + + +def _sqlite_config() -> SqliteConfig: + return SqliteConfig( + connection_config={"database": ":memory:"}, + extension_config={ + "adk": { + "session_table": "agent_sessions", + "events_table": "agent_events", + "app_state_table": "agent_app_state", + "user_state_table": "agent_user_state", + "metadata_table": "agent_metadata", + "memory_table": "agent_memory", + "artifact_table": "agent_artifact", + "owner_id_column": "tenant_id INTEGER", + } + }, + ) + + def test_session_store_bases_share_slotted_common_state() -> None: assert issubclass(BaseAsyncADKStore, _ADKStoreCommon) assert issubclass(BaseSyncADKStore, _ADKStoreCommon) @@ -28,3 +82,41 @@ def test_artifact_store_bases_share_slotted_common_state() -> None: assert issubclass(BaseSyncADKArtifactStore, _ADKArtifactStoreCommon) assert BaseAsyncADKArtifactStore.__slots__ == () assert BaseSyncADKArtifactStore.__slots__ == () + + +@pytest.mark.parametrize( + "store", + [ + pytest.param(SqliteADKStore(_sqlite_config()), id="session"), + pytest.param(SqliteADKMemoryStore(_sqlite_config()), id="memory"), + pytest.param(_ConcreteArtifactStore(_sqlite_config()), id="artifact"), + ], +) +def test_concrete_adk_stores_reject_undeclared_attributes(store: object) -> None: + with pytest.raises(AttributeError): + setattr(store, "undeclared", True) + + +def test_concrete_session_store_exposes_all_shared_properties() -> None: + config = _sqlite_config() + store = SqliteADKStore(config) + + assert { + "config": store.config, + "session_table": store.session_table, + "events_table": store.events_table, + "app_state_table": store.app_state_table, + "user_state_table": store.user_state_table, + "metadata_table": store.metadata_table, + "owner_id_column_ddl": store.owner_id_column_ddl, + "owner_id_column_name": store.owner_id_column_name, + } == { + "config": config, + "session_table": "agent_sessions", + "events_table": "agent_events", + "app_state_table": "agent_app_state", + "user_state_table": "agent_user_state", + "metadata_table": "agent_metadata", + "owner_id_column_ddl": "tenant_id INTEGER", + "owner_id_column_name": "tenant_id", + }