sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/adapters/spanner/adk/store.py

@@ -0,0 +1,758 @@
+"""Spanner ADK store."""
+
+from collections.abc import Iterable
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast
+
+from google.cloud.spanner_v1 import param_types
+
+from sqlspec.adapters.spanner.config import SpannerSyncConfig
+from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
+from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore
+from sqlspec.protocols import SpannerParamTypesProtocol
+from sqlspec.utils.serializers import from_json, to_json
+from sqlspec.utils.sync_tools import async_, run_
+
+if TYPE_CHECKING:
+    from google.cloud.spanner_v1.database import Database
+    from google.cloud.spanner_v1.transaction import Transaction
+
+    from sqlspec.config import ADKConfig
+    from sqlspec.extensions.adk import MemoryRecord
+SPANNER_PARAM_TYPES: SpannerParamTypesProtocol = cast("SpannerParamTypesProtocol", param_types)
+
+__all__ = ("SpannerSyncADKMemoryStore", "SpannerSyncADKStore")
+
+
+class SpannerSyncADKStore(BaseAsyncADKStore[SpannerSyncConfig]):
+    """Spanner ADK store backed by synchronous Spanner client."""
+
+    connector_name: ClassVar[str] = "spanner"
+
+    def __init__(self, config: SpannerSyncConfig) -> None:
+        super().__init__(config)
+        adk_config = cast("dict[str, Any]", config.extension_config.get("adk", {}))
+        self._shard_count: int = int(adk_config.get("shard_count", 0)) if adk_config.get("shard_count") else 0
+        self._session_table_options: str | None = adk_config.get("session_table_options")
+        self._events_table_options: str | None = adk_config.get("events_table_options")
+        self._expires_index_options: str | None = adk_config.get("expires_index_options")
+        self._session_row_deletion_policy = _spanner_row_deletion_policy(
+            adk_config, "session_ttl_seconds", "create_time"
+        )
+        self._events_row_deletion_policy = _spanner_row_deletion_policy(adk_config, "event_ttl_seconds", "timestamp")
+
+    def _database(self) -> "Database":
+        return self._config.get_database()
+
+    def _run_read(
+        self, sql: str, params: "dict[str, Any] | None" = None, types: "dict[str, Any] | None" = None
+    ) -> "list[Any]":
+        with self._config.provide_connection() as snapshot:
+            result_set = cast("Any", snapshot).execute_sql(sql, params=params, param_types=types)
+            return list(result_set)
+
+    def _run_write(self, statements: "list[tuple[str, dict[str, Any], dict[str, Any]]]") -> None:
+        self._database().run_in_transaction(_SpannerWriteJob(statements))  # type: ignore[no-untyped-call]
+
+    def _session_param_types(self, include_owner: bool) -> "dict[str, Any]":
+        json_type = _json_param_type()
+        types: dict[str, Any] = {
+            "id": SPANNER_PARAM_TYPES.STRING,
+            "app_name": SPANNER_PARAM_TYPES.STRING,
+            "user_id": SPANNER_PARAM_TYPES.STRING,
+            "state": json_type,
+        }
+        if include_owner and self._owner_id_column_name:
+            types["owner_id"] = SPANNER_PARAM_TYPES.STRING
+        return types
+
+    def _event_param_types(self) -> "dict[str, Any]":
+        json_type = _json_param_type()
+        return {
+            "session_id": SPANNER_PARAM_TYPES.STRING,
+            "invocation_id": SPANNER_PARAM_TYPES.STRING,
+            "author": SPANNER_PARAM_TYPES.STRING,
+            "timestamp": SPANNER_PARAM_TYPES.TIMESTAMP,
+            "event_json": json_type,
+        }
+
+    def _decode_state(self, raw: Any) -> Any:
+        if isinstance(raw, str):
+            return from_json(raw)
+        return raw
+
+    def _decode_json(self, raw: Any) -> Any:
+        if raw is None:
+            return None
+        if isinstance(raw, str):
+            return from_json(raw)
+        return raw
+
+    def _create_session(
+        self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
+    ) -> SessionRecord:
+        state_json = to_json(state)
+        params: dict[str, Any] = {"id": session_id, "app_name": app_name, "user_id": user_id, "state": state_json}
+        columns = "id, app_name, user_id, state, create_time, update_time"
+        values = "@id, @app_name, @user_id, @state, PENDING_COMMIT_TIMESTAMP(), PENDING_COMMIT_TIMESTAMP()"
+        if self._owner_id_column_name:
+            params["owner_id"] = owner_id
+            columns = f"id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time"
+            values = (
+                "@id, @app_name, @user_id, @owner_id, @state, PENDING_COMMIT_TIMESTAMP(), PENDING_COMMIT_TIMESTAMP()"
+            )
+
+        sql = f"""
+            INSERT INTO {self._session_table} ({columns})
+            VALUES ({values})
+        """
+        self._run_write([(sql, params, self._session_param_types(self._owner_id_column_name is not None))])
+
+        return {
+            "id": session_id,
+            "app_name": app_name,
+            "user_id": user_id,
+            "state": state,
+            "create_time": datetime.now(timezone.utc),
+            "update_time": datetime.now(timezone.utc),
+        }
+
+    async 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 await async_(self._create_session)(session_id, app_name, user_id, state, owner_id)
+
+    def _get_session(self, session_id: str) -> "SessionRecord | None":
+        sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time{", " + self._owner_id_column_name if self._owner_id_column_name else ""}
+            FROM {self._session_table}
+            WHERE id = @id
+        """
+        if self._shard_count > 1:
+            sql = f"{sql} AND shard_id = MOD(FARM_FINGERPRINT(@id), {self._shard_count})"
+        sql = f"{sql} LIMIT 1"
+        params = {"id": session_id}
+        rows = self._run_read(sql, params, {"id": SPANNER_PARAM_TYPES.STRING})
+        if not rows:
+            return None
+
+        row = rows[0]
+        state_value = self._decode_state(row[3])
+        record: SessionRecord = {
+            "id": row[0],
+            "app_name": row[1],
+            "user_id": row[2],
+            "state": state_value,
+            "create_time": row[4],
+            "update_time": row[5],
+        }
+        return record
+
+    async def get_session(self, session_id: str) -> "SessionRecord | None":
+        """Get session by ID."""
+        return await async_(self._get_session)(session_id)
+
+    def _update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
+        params = {"id": session_id, "state": to_json(state)}
+        json_type = _json_param_type()
+        sql = f"""
+            UPDATE {self._session_table}
+            SET state = @state, update_time = PENDING_COMMIT_TIMESTAMP()
+            WHERE id = @id
+        """
+        if self._shard_count > 1:
+            sql = f"{sql} AND shard_id = MOD(FARM_FINGERPRINT(@id), {self._shard_count})"
+        self._run_write([(sql, params, {"id": SPANNER_PARAM_TYPES.STRING, "state": json_type})])
+
+    async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
+        """Update session state."""
+        await async_(self._update_session_state)(session_id, state)
+
+    def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]":
+        sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time{", " + self._owner_id_column_name if self._owner_id_column_name else ""}
+            FROM {self._session_table}
+            WHERE app_name = @app_name
+        """
+        params: dict[str, Any] = {"app_name": app_name}
+        types: dict[str, Any] = {"app_name": SPANNER_PARAM_TYPES.STRING}
+        if user_id is not None:
+            sql = f"{sql} AND user_id = @user_id"
+            params["user_id"] = user_id
+            types["user_id"] = SPANNER_PARAM_TYPES.STRING
+        if self._shard_count > 1:
+            sql = f"{sql} AND shard_id = MOD(FARM_FINGERPRINT(id), {self._shard_count})"
+        sql = f"{sql} ORDER BY update_time DESC"
+
+        rows = self._run_read(sql, params, types)
+        records: list[SessionRecord] = []
+        for row in rows:
+            state_value = self._decode_state(row[3])
+            record: SessionRecord = {
+                "id": row[0],
+                "app_name": row[1],
+                "user_id": row[2],
+                "state": state_value,
+                "create_time": row[4],
+                "update_time": row[5],
+            }
+            records.append(record)
+        return records
+
+    async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
+        """List sessions for an app."""
+        return await async_(self._list_sessions)(app_name, user_id)
+
+    def _delete_session(self, session_id: str) -> None:
+        shard_clause = (
+            f" AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})" if self._shard_count > 1 else ""
+        )
+        delete_events_sql = f"DELETE FROM {self._events_table} WHERE session_id = @session_id{shard_clause}"
+        delete_session_sql = f"DELETE FROM {self._session_table} WHERE id = @session_id{shard_clause}"
+        params = {"session_id": session_id}
+        types = {"session_id": SPANNER_PARAM_TYPES.STRING}
+        self._run_write([(delete_events_sql, params, types), (delete_session_sql, params, types)])
+
+    async def delete_session(self, session_id: str) -> None:
+        """Delete session and associated events."""
+        await async_(self._delete_session)(session_id)
+
+    def _append_event_and_update_state(
+        self, event_record: "EventRecord", session_id: str, state: "dict[str, Any]"
+    ) -> SessionRecord:
+        """Atomically insert an event and update session state in one transaction.
+
+        Both the event INSERT and the session state UPDATE execute within a single
+        Spanner transaction so they succeed or fail together. A follow-up
+        single-use read returns the SessionRecord; we can't capture update_time
+        inside the write txn because PENDING_COMMIT_TIMESTAMP() only materialises
+        on commit.
+
+        Args:
+            event_record: Event record to store.
+            session_id: Session whose state should be updated.
+            state: Post-append durable state snapshot.
+        """
+        event_params: dict[str, Any] = {
+            "session_id": event_record["session_id"],
+            "invocation_id": event_record["invocation_id"],
+            "author": event_record["author"],
+            "timestamp": event_record["timestamp"],
+            "event_json": to_json(event_record["event_json"]),
+        }
+        insert_sql = f"""
+            INSERT INTO {self._events_table} (session_id, invocation_id, author, timestamp, event_json)
+            VALUES (@session_id, @invocation_id, @author, @timestamp, @event_json)
+        """
+
+        json_type = _json_param_type()
+        state_params: dict[str, Any] = {"id": session_id, "state": to_json(state)}
+        update_sql = f"""
+            UPDATE {self._session_table}
+            SET state = @state, update_time = PENDING_COMMIT_TIMESTAMP()
+            WHERE id = @id
+        """
+        if self._shard_count > 1:
+            update_sql = f"{update_sql} AND shard_id = MOD(FARM_FINGERPRINT(@id), {self._shard_count})"
+
+        self._run_write([
+            (insert_sql, event_params, self._event_param_types()),
+            (update_sql, state_params, {"id": SPANNER_PARAM_TYPES.STRING, "state": json_type}),
+        ])
+
+        record = self._get_session(session_id)
+        if record is None:
+            msg = f"Session {session_id} not found during append_event_and_update_state."
+            raise ValueError(msg)
+        return record
+
+    async def append_event_and_update_state(
+        self, event_record: EventRecord, session_id: str, state: "dict[str, Any]"
+    ) -> SessionRecord:
+        """Atomically append an event and update the session's durable state."""
+        return await async_(self._append_event_and_update_state)(event_record, session_id, state)
+
+    def _insert_event(self, event_record: "EventRecord") -> None:
+        event_params: dict[str, Any] = {
+            "session_id": event_record["session_id"],
+            "invocation_id": event_record["invocation_id"],
+            "author": event_record["author"],
+            "timestamp": event_record["timestamp"],
+            "event_json": to_json(event_record["event_json"]),
+        }
+        insert_sql = f"""
+            INSERT INTO {self._events_table} (session_id, invocation_id, author, timestamp, event_json)
+            VALUES (@session_id, @invocation_id, @author, @timestamp, @event_json)
+        """
+        self._run_write([(insert_sql, event_params, self._event_param_types())])
+
+    def _get_events(
+        self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
+    ) -> "list[EventRecord]":
+        sql = f"""
+            SELECT session_id, invocation_id, author, timestamp, event_json
+            FROM {self._events_table}
+            WHERE session_id = @session_id
+        """
+        if self._shard_count > 1:
+            sql = f"{sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})"
+        params: dict[str, Any] = {"session_id": session_id}
+        types: dict[str, Any] = {"session_id": SPANNER_PARAM_TYPES.STRING}
+        if after_timestamp is not None:
+            sql = f"{sql} AND timestamp > @after_timestamp"
+            params["after_timestamp"] = after_timestamp
+            types["after_timestamp"] = SPANNER_PARAM_TYPES.TIMESTAMP
+        sql = f"{sql} ORDER BY timestamp ASC"
+        if limit is not None:
+            sql = f"{sql} LIMIT @limit"
+            params["limit"] = limit
+            types["limit"] = SPANNER_PARAM_TYPES.INT64
+        rows = self._run_read(sql, params, types)
+        return [
+            {
+                "session_id": row[0],
+                "invocation_id": row[1] or "",
+                "author": row[2] or "",
+                "timestamp": row[3],
+                "event_json": row[4],
+            }
+            for row in rows
+        ]
+
+    async def get_events(
+        self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
+    ) -> "list[EventRecord]":
+        """Get events for a session."""
+        return await async_(self._get_events)(session_id, after_timestamp, limit)
+
+    def _append_event(self, event_record: EventRecord) -> None:
+        """Synchronous implementation of append_event."""
+        self._insert_event(event_record)
+
+    async def append_event(self, event_record: EventRecord) -> None:
+        """Append an event to a session."""
+        await async_(self._append_event)(event_record)
+
+    def _create_tables(self) -> None:
+        database = self._database()
+        existing_tables = {t.table_id for t in database.list_tables()}  # type: ignore[no-untyped-call]
+
+        ddl_statements: list[str] = []
+        if self._session_table not in existing_tables:
+            ddl_statements.append(run_(self._get_create_sessions_table_sql)())
+        if self._events_table not in existing_tables:
+            ddl_statements.append(run_(self._get_create_events_table_sql)())
+
+        if ddl_statements:
+            database.update_ddl(ddl_statements).result(300)  # type: ignore[no-untyped-call]
+
+    async def create_tables(self) -> None:
+        """Create tables if they don't exist."""
+        await async_(self._create_tables)()
+
+    async def _get_create_sessions_table_sql(self) -> str:
+        owner_line = ""
+        if self._owner_id_column_ddl:
+            owner_line = f",\n  {self._owner_id_column_ddl}"
+        shard_column = ""
+        pk = "PRIMARY KEY (id)"
+        if self._shard_count > 1:
+            shard_column = f",\n  shard_id INT64 AS (MOD(FARM_FINGERPRINT(id), {self._shard_count})) STORED"
+            pk = "PRIMARY KEY (shard_id, id)"
+        options = ""
+        if self._session_table_options:
+            options = f"\nOPTIONS ({self._session_table_options})"
+        return f"""
+CREATE TABLE {self._session_table} (
+  id STRING(128) NOT NULL,
+  app_name STRING(128) NOT NULL,
+  user_id STRING(128) NOT NULL{owner_line},
+  state JSON NOT NULL,
+  create_time TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true),
+  update_time TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true){shard_column}
+) {pk}{options}{self._session_row_deletion_policy}
+"""
+
+    async def _get_create_events_table_sql(self) -> str:
+        shard_column = ""
+        pk = "PRIMARY KEY (session_id, timestamp)"
+        if self._shard_count > 1:
+            shard_column = f",\n  shard_id INT64 AS (MOD(FARM_FINGERPRINT(session_id), {self._shard_count})) STORED"
+            pk = "PRIMARY KEY (shard_id, session_id, timestamp)"
+        options = ""
+        if self._events_table_options:
+            options = f"\nOPTIONS ({self._events_table_options})"
+        return f"""
+CREATE TABLE {self._events_table} (
+  session_id STRING(128) NOT NULL,
+  invocation_id STRING(256) NOT NULL,
+  author STRING(128) NOT NULL,
+  timestamp TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true),
+  event_json JSON NOT NULL{shard_column}
+) {pk}{options}{self._events_row_deletion_policy}
+"""
+
+    def _get_drop_tables_sql(self) -> "list[str]":
+        return [f"DROP TABLE {self._events_table}", f"DROP TABLE {self._session_table}"]
+
+
+class SpannerSyncADKMemoryStore(BaseAsyncADKMemoryStore[SpannerSyncConfig]):
+    """Spanner ADK memory store backed by synchronous Spanner client."""
+
+    connector_name: ClassVar[str] = "spanner"
+
+    def __init__(self, config: SpannerSyncConfig) -> None:
+        super().__init__(config)
+        adk_config = cast("ADKConfig", config.extension_config.get("adk", {}))
+        shard_count = adk_config.get("shard_count")
+        self._shard_count = int(shard_count) if isinstance(shard_count, int) else 0
+        self._memory_table_options: str | None = adk_config.get("memory_table_options")
+        self._memory_row_deletion_policy = _spanner_row_deletion_policy(
+            cast("dict[str, Any]", adk_config), "memory_ttl_seconds", "inserted_at"
+        )
+
+    def _database(self) -> "Database":
+        return self._config.get_database()
+
+    def _run_read(
+        self, sql: str, params: "dict[str, Any] | None" = None, types: "dict[str, Any] | None" = None
+    ) -> "list[Any]":
+        with self._config.provide_connection() as snapshot:
+            reader = cast("_SpannerReadProtocol", snapshot)
+            result_set = reader.execute_sql(sql, params=params, param_types=types)
+            return list(result_set)
+
+    def _run_write(self, statements: "list[tuple[str, dict[str, Any], dict[str, Any]]]") -> None:
+        self._database().run_in_transaction(_SpannerMemoryWriteJob(statements))  # type: ignore[no-untyped-call]
+
+    def _execute_update(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> int:
+        return int(self._database().run_in_transaction(_SpannerMemoryUpdateJob(sql, params, types)))  # type: ignore[no-untyped-call]
+
+    def _memory_param_types(self, include_owner: bool) -> "dict[str, Any]":
+        types: dict[str, Any] = {
+            "id": SPANNER_PARAM_TYPES.STRING,
+            "session_id": SPANNER_PARAM_TYPES.STRING,
+            "app_name": SPANNER_PARAM_TYPES.STRING,
+            "user_id": SPANNER_PARAM_TYPES.STRING,
+            "event_id": SPANNER_PARAM_TYPES.STRING,
+            "author": SPANNER_PARAM_TYPES.STRING,
+            "timestamp": SPANNER_PARAM_TYPES.TIMESTAMP,
+            "content_json": _json_param_type(),
+            "content_text": SPANNER_PARAM_TYPES.STRING,
+            "metadata_json": _json_param_type(),
+            "inserted_at": SPANNER_PARAM_TYPES.TIMESTAMP,
+        }
+        if include_owner and self._owner_id_column_name:
+            types["owner_id"] = SPANNER_PARAM_TYPES.STRING
+        return types
+
+    def _decode_json(self, raw: Any) -> Any:
+        if raw is None:
+            return None
+        if isinstance(raw, str):
+            return from_json(raw)
+        return raw
+
+    def _create_tables(self) -> None:
+        if not self._enabled:
+            return
+
+        database = self._database()
+        existing_tables = {t.table_id for t in database.list_tables()}  # type: ignore[no-untyped-call]
+
+        ddl_statements: list[str] = []
+        if self._memory_table not in existing_tables:
+            ddl_statements.extend(run_(self._get_create_memory_table_sql)())
+
+        if ddl_statements:
+            database.update_ddl(ddl_statements).result(300)  # type: ignore[no-untyped-call]
+
+    async def create_tables(self) -> None:
+        """Create tables if they don't exist."""
+        await async_(self._create_tables)()
+
+    async def _get_create_memory_table_sql(self) -> "list[str]":
+        owner_line = ""
+        if self._owner_id_column_ddl:
+            owner_line = f",\n  {self._owner_id_column_ddl}"
+
+        fts_column_line = ""
+        fts_index = ""
+        if self._use_fts:
+            fts_column_line = "\n  content_tokens TOKENLIST AS (TOKENIZE_FULLTEXT(content_text)) HIDDEN"
+            fts_index = f"CREATE SEARCH INDEX idx_{self._memory_table}_fts ON {self._memory_table}(content_tokens)"
+
+        shard_column = ""
+        pk = "PRIMARY KEY (id)"
+        if self._shard_count > 1:
+            shard_column = f",\n  shard_id INT64 AS (MOD(FARM_FINGERPRINT(id), {self._shard_count})) STORED"
+            pk = "PRIMARY KEY (shard_id, id)"
+        options = ""
+        if self._memory_table_options:
+            options = f"\nOPTIONS ({self._memory_table_options})"
+
+        table_sql = f"""
+CREATE TABLE {self._memory_table} (
+  id STRING(128) NOT NULL,
+  session_id STRING(128) NOT NULL,
+  app_name STRING(128) NOT NULL,
+  user_id STRING(128) NOT NULL,
+  event_id STRING(128) NOT NULL,
+  author STRING(256){owner_line},
+  timestamp TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true),
+  content_json JSON NOT NULL,
+  content_text STRING(MAX) NOT NULL,
+  metadata_json JSON,
+  inserted_at TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true){fts_column_line}{shard_column}
+) {pk}{options}{self._memory_row_deletion_policy}
+"""
+
+        app_user_idx = (
+            f"CREATE INDEX idx_{self._memory_table}_app_user_time "
+            f"ON {self._memory_table}(app_name, user_id, timestamp DESC)"
+        )
+        session_idx = f"CREATE INDEX idx_{self._memory_table}_session ON {self._memory_table}(session_id)"
+
+        statements = [table_sql, app_user_idx, session_idx]
+        if fts_index:
+            statements.append(fts_index)
+        return statements
+
+    def _get_drop_memory_table_sql(self) -> "list[str]":
+        """Get SQL to drop the memory table and its indexes.
+
+        Returns:
+            List of SQL statements to drop the memory table and associated indexes.
+        """
+        statements: list[str] = []
+        if self._use_fts:
+            statements.append(f"DROP SEARCH INDEX idx_{self._memory_table}_fts")
+        statements.extend([
+            f"DROP INDEX idx_{self._memory_table}_session",
+            f"DROP INDEX idx_{self._memory_table}_app_user_time",
+            f"DROP TABLE {self._memory_table}",
+        ])
+        return statements
+
+    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
+
+        inserted_count = 0
+        statements: list[tuple[str, dict[str, Any], dict[str, Any]]] = []
+
+        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 ""
+
+        insert_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
+        )
+        """
+
+        for entry in entries:
+            if self._event_exists(entry["event_id"]):
+                continue
+            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": to_json(entry["content_json"]),
+                "content_text": entry["content_text"],
+                "metadata_json": to_json(entry["metadata_json"]) if entry["metadata_json"] is not None else None,
+                "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
+            statements.append((insert_sql, params, self._memory_param_types(self._owner_id_column_name is not None)))
+            inserted_count += 1
+
+        if statements:
+            self._run_write(statements)
+        return inserted_count
+
+    async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int:
+        """Bulk insert memory entries with deduplication."""
+        return await async_(self._insert_memory_entries)(entries, owner_id)
+
+    def _event_exists(self, event_id: str) -> bool:
+        sql = f"SELECT event_id FROM {self._memory_table} WHERE event_id = @event_id LIMIT 1"
+        rows = self._run_read(sql, {"event_id": event_id}, {"event_id": SPANNER_PARAM_TYPES.STRING})
+        return bool(rows)
+
+    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
+
+        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)
+
+    async def search_entries(
+        self, query: str, app_name: str, user_id: str, limit: "int | None" = None
+    ) -> "list[MemoryRecord]":
+        """Search memory entries by text query."""
+        return await async_(self._search_entries)(query, app_name, user_id, limit)
+
+    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,
+               timestamp, content_json, content_text, metadata_json, inserted_at
+        FROM {self._memory_table}
+        WHERE app_name = @app_name
+          AND user_id = @user_id
+          AND SEARCH(content_tokens, @query)
+        ORDER BY timestamp DESC
+        LIMIT @limit
+        """
+        params = {"app_name": app_name, "user_id": user_id, "query": query, "limit": limit}
+        types = {
+            "app_name": SPANNER_PARAM_TYPES.STRING,
+            "user_id": SPANNER_PARAM_TYPES.STRING,
+            "query": SPANNER_PARAM_TYPES.STRING,
+            "limit": SPANNER_PARAM_TYPES.INT64,
+        }
+        rows = self._run_read(sql, params, types)
+        return self._rows_to_records(rows)
+
+    def _search_entries_simple(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,
+               timestamp, content_json, content_text, metadata_json, inserted_at
+        FROM {self._memory_table}
+        WHERE app_name = @app_name
+          AND user_id = @user_id
+          AND LOWER(content_text) LIKE @pattern
+        ORDER BY timestamp DESC
+        LIMIT @limit
+        """
+        pattern = f"%{query.lower()}%"
+        params = {"app_name": app_name, "user_id": user_id, "pattern": pattern, "limit": limit}
+        types = {
+            "app_name": SPANNER_PARAM_TYPES.STRING,
+            "user_id": SPANNER_PARAM_TYPES.STRING,
+            "pattern": SPANNER_PARAM_TYPES.STRING,
+            "limit": SPANNER_PARAM_TYPES.INT64,
+        }
+        rows = self._run_read(sql, params, types)
+        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"
+        params = {"session_id": session_id}
+        types = {"session_id": SPANNER_PARAM_TYPES.STRING}
+        return self._execute_update(sql, params, types)
+
+    async def delete_entries_by_session(self, session_id: str) -> int:
+        """Delete all memory entries for a specific session."""
+        return await async_(self._delete_entries_by_session)(session_id)
+
+    def _delete_entries_older_than(self, days: int) -> int:
+        cutoff = datetime.now(timezone.utc) - timedelta(days=days)
+        sql = f"DELETE FROM {self._memory_table} WHERE inserted_at < @cutoff"
+        params = {"cutoff": cutoff}
+        types = {"cutoff": SPANNER_PARAM_TYPES.TIMESTAMP}
+        return self._execute_update(sql, params, types)
+
+    async def delete_entries_older_than(self, days: int) -> int:
+        """Delete memory entries older than specified days."""
+        return await async_(self._delete_entries_older_than)(days)
+
+    def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]":
+        return [
+            {
+                "id": row[0],
+                "session_id": row[1],
+                "app_name": row[2],
+                "user_id": row[3],
+                "event_id": row[4],
+                "author": row[5],
+                "timestamp": row[6],
+                "content_json": self._decode_json(row[7]),
+                "content_text": row[8],
+                "metadata_json": self._decode_json(row[9]),
+                "inserted_at": row[10],
+            }
+            for row in rows
+        ]
+
+
+def _json_param_type() -> Any:
+    try:
+        return SPANNER_PARAM_TYPES.JSON
+    except AttributeError:
+        return SPANNER_PARAM_TYPES.STRING
+
+
+def _spanner_ttl_days(ttl_seconds: Any) -> int:
+    if not isinstance(ttl_seconds, int) or ttl_seconds <= 0:
+        return 0
+    return max(1, (ttl_seconds + 86_399) // 86_400)
+
+
+def _spanner_row_deletion_policy(adk_config: dict[str, Any], ttl_key: str, column: str) -> str:
+    retention = adk_config.get("retention")
+    if not isinstance(retention, dict):
+        return ""
+    ttl_days = _spanner_ttl_days(retention.get(ttl_key))
+    if ttl_days == 0:
+        return ""
+    return f"\nROW DELETION POLICY (OLDER_THAN({column}, INTERVAL {ttl_days} DAY))"
+
+
+class _SpannerWriteJob:
+    __slots__ = ("_statements",)
+
+    def __init__(self, statements: "list[tuple[str, dict[str, Any], dict[str, Any]]]") -> None:
+        self._statements = statements
+
+    def __call__(self, transaction: "Transaction") -> None:
+        for sql, params, types in self._statements:
+            transaction.execute_update(sql, params=params, param_types=types)  # type: ignore[no-untyped-call]
+
+
+class _SpannerMemoryWriteJob:
+    __slots__ = ("_statements",)
+
+    def __init__(self, statements: "list[tuple[str, dict[str, Any], dict[str, Any]]]") -> None:
+        self._statements = statements
+
+    def __call__(self, transaction: "Transaction") -> None:
+        for sql, params, types in self._statements:
+            transaction.execute_update(sql, params=params, param_types=types)  # type: ignore[no-untyped-call]
+
+
+class _SpannerMemoryUpdateJob:
+    __slots__ = ("_params", "_sql", "_types")
+
+    def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None:
+        self._sql = sql
+        self._params = params
+        self._types = types
+
+    def __call__(self, transaction: "Transaction") -> int:
+        return int(transaction.execute_update(self._sql, params=self._params, param_types=self._types))  # type: ignore[no-untyped-call]
+
+
+class _SpannerReadProtocol(Protocol):
+    def execute_sql(
+        self, sql: str, params: "dict[str, Any] | None" = None, param_types: "dict[str, Any] | None" = None
+    ) -> Iterable[Any]: ...