Source code for sqlspec.adapters.arrow_odbc.adk.store

"""arrow-odbc ADK stores for Google Agent Development Kit session storage."""

from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast

from typing_extensions import NotRequired

from sqlspec.config import ADKConfig
from sqlspec.exceptions import SQLSpecError
from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options
from sqlspec.extensions.adk.memory import BaseSyncADKMemoryStore, StoredMemory
from sqlspec.utils.serializers import from_json, to_json

if TYPE_CHECKING:
    from collections.abc import Sequence
    from datetime import timedelta

    from sqlspec.adapters.arrow_odbc.config import ArrowOdbcConfig
    from sqlspec.extensions.adk import SessionOrderBy
else:
    ArrowOdbcConfig = Any


__all__ = ("ArrowOdbcADKConfig", "ArrowOdbcADKMemoryStore", "ArrowOdbcADKStore")

MSSQL_SCHEMA: Final[str] = "dbo"
JSON_COLUMN_TYPE: Final[str] = "NVARCHAR(MAX)"


class ArrowOdbcADKConfig(ADKConfig):
    """arrow-odbc ADK extension settings."""

    native_json: NotRequired[bool]
    """Accepted for parity with SQL Server adapters; arrow-odbc uses NVARCHAR(MAX)."""



[docs]
class ArrowOdbcADKStore(BaseSyncADKStore["ArrowOdbcConfig"]):
    """Synchronous SQL Server ADK session/event store using arrow-odbc."""

    connector_name: ClassVar[str] = "arrow_odbc"
    __slots__ = ()


[docs]
    def create_tables(self) -> None:
        """Create the ADK tables and indexes the data dictionary reports as missing."""
        if not self.create_schema_enabled:
            self.reconcile_schema()
            return

        with self._config.provide_session() as driver:
            dd = driver.data_dictionary
            existing_tables = _casefold_names(dd.get_tables(driver, schema=MSSQL_SCHEMA), "table_name")
            existing_indexes = _casefold_names(dd.get_indexes(driver, schema=MSSQL_SCHEMA), "index_name")
            table_ddls = (
                (self._session_table, self._sessions_table_ddl()),
                (self._events_table, self._events_table_ddl()),
                (self._app_state_table, self._app_states_table_ddl()),
                (self._user_state_table, self._user_states_table_ddl()),
                (self._metadata_table, self._metadata_table_ddl()),
            )
            for table, ddl in table_ddls:
                if _bare_name(table) not in existing_tables:
                    driver.execute(ddl)
            for index_name, index_table, columns in self._index_specs():
                if _bare_name(index_name) not in existing_indexes:
                    driver.execute(_create_index_sql(index_table, index_name, columns))
            driver.commit()

[docs]

def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None ) -> StoredSession: """Create a new ADK session.""" owner_column = f", {_quote_identifier(self._owner_id_column_name)}" if self._owner_id_column_name else "" owner_param = ", ?" if self._owner_id_column_name else "" params: tuple[Any, ...] if self._owner_id_column_name: params = (session_id, app_name, user_id, owner_id, to_json(state)) else: params = (session_id, app_name, user_id, to_json(state)) with self._config.provide_session() as driver: driver.execute( f""" INSERT INTO {_table_ref(self._session_table)} ( id, app_name, user_id{owner_column}, state, create_time, update_time ) VALUES (?, ?, ?{owner_param}, ?, SYSUTCDATETIME(), SYSUTCDATETIME()) """, params, ) row = driver.select_one_or_none(_session_select_sql(self._session_table), (app_name, user_id, session_id)) driver.commit() if row is None: msg = "Failed to fetch created session" raise RuntimeError(msg) return _session_record_from_row(row)