sqlspec 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/adapters/asyncmy/adk/store.py

@@ -0,0 +1,503 @@
+"""AsyncMy ADK store for Google Agent Development Kit session/event storage."""
+
+import json
+from typing import TYPE_CHECKING, Any, Final
+
+import asyncmy
+
+from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from datetime import datetime
+
+    from sqlspec.adapters.asyncmy.config import AsyncmyConfig
+
+logger = get_logger("adapters.asyncmy.adk.store")
+
+__all__ = ("AsyncmyADKStore",)
+
+MYSQL_TABLE_NOT_FOUND_ERROR: Final = 1146
+
+
+class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
+    """MySQL/MariaDB ADK store using AsyncMy driver.
+
+    Implements session and event storage for Google Agent Development Kit
+    using MySQL/MariaDB via the AsyncMy driver. Provides:
+    - Session state management with JSON storage
+    - Event history tracking with BLOB-serialized actions
+    - Microsecond-precision timestamps
+    - Foreign key constraints with cascade delete
+    - Efficient upserts using ON DUPLICATE KEY UPDATE
+
+    Args:
+        config: AsyncmyConfig with extension_config["adk"] settings.
+
+    Example:
+        from sqlspec.adapters.asyncmy import AsyncmyConfig
+        from sqlspec.adapters.asyncmy.adk import AsyncmyADKStore
+
+        config = AsyncmyConfig(
+            pool_config={"host": "localhost", ...},
+            extension_config={
+                "adk": {
+                    "session_table": "my_sessions",
+                    "events_table": "my_events",
+                    "owner_id_column": "tenant_id BIGINT NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
+                }
+            }
+        )
+        store = AsyncmyADKStore(config)
+        await store.create_tables()
+
+    Notes:
+        - MySQL JSON type used (not JSONB) - requires MySQL 5.7.8+
+        - TIMESTAMP(6) provides microsecond precision
+        - InnoDB engine required for foreign key support
+        - State merging handled at application level
+        - Configuration is read from config.extension_config["adk"]
+    """
+
+    __slots__ = ()
+
+    def __init__(self, config: "AsyncmyConfig") -> None:
+        """Initialize AsyncMy ADK store.
+
+        Args:
+            config: AsyncmyConfig instance.
+
+        Notes:
+            Configuration is read from config.extension_config["adk"]:
+            - session_table: Sessions table name (default: "adk_sessions")
+            - events_table: Events table name (default: "adk_events")
+            - owner_id_column: Optional owner FK column DDL (default: None)
+        """
+        super().__init__(config)
+
+    def _parse_owner_id_column_for_mysql(self, column_ddl: str) -> "tuple[str, str]":
+        """Parse owner ID column DDL for MySQL FOREIGN KEY syntax.
+
+        MySQL ignores inline REFERENCES syntax in column definitions.
+        This method extracts the column definition and creates a separate
+        FOREIGN KEY constraint.
+
+        Args:
+            column_ddl: Column DDL like "tenant_id BIGINT NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
+
+        Returns:
+            Tuple of (column_definition, foreign_key_constraint)
+
+        Example:
+            Input: "tenant_id BIGINT NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
+            Output: ("tenant_id BIGINT NOT NULL", "FOREIGN KEY (tenant_id) REFERENCES tenants(id) ON DELETE CASCADE")
+        """
+        import re
+
+        references_match = re.search(r"\s+REFERENCES\s+(.+)", column_ddl, re.IGNORECASE)
+
+        if not references_match:
+            return (column_ddl.strip(), "")
+
+        col_def = column_ddl[: references_match.start()].strip()
+        fk_clause = references_match.group(1).strip()
+        col_name = col_def.split()[0]
+        fk_constraint = f"FOREIGN KEY ({col_name}) REFERENCES {fk_clause}"
+
+        return (col_def, fk_constraint)
+
+    async def _get_create_sessions_table_sql(self) -> str:
+        """Get MySQL CREATE TABLE SQL for sessions.
+
+        Returns:
+            SQL statement to create adk_sessions table with indexes.
+
+        Notes:
+            - VARCHAR(128) for IDs and names (sufficient for UUIDs and app names)
+            - JSON type for state storage (MySQL 5.7.8+)
+            - TIMESTAMP(6) with microsecond precision
+            - AUTO-UPDATE on update_time
+            - Composite index on (app_name, user_id) for listing
+            - Index on update_time DESC for recent session queries
+            - Optional owner ID column for multi-tenancy
+            - MySQL requires explicit FOREIGN KEY syntax (inline REFERENCES is ignored)
+        """
+        owner_id_col = ""
+        fk_constraint = ""
+
+        if self._owner_id_column_ddl:
+            col_def, fk_def = self._parse_owner_id_column_for_mysql(self._owner_id_column_ddl)
+            owner_id_col = f"{col_def},"
+            if fk_def:
+                fk_constraint = f",\n            {fk_def}"
+
+        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_col}
+            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_{self._session_table}_app_user (app_name, user_id),
+            INDEX idx_{self._session_table}_update_time (update_time DESC){fk_constraint}
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+        """
+
+    async def _get_create_events_table_sql(self) -> str:
+        """Get MySQL CREATE TABLE SQL for events.
+
+        Returns:
+            SQL statement to create adk_events table with indexes.
+
+        Notes:
+            - VARCHAR sizes: id(128), session_id(128), invocation_id(256), author(256),
+              branch(256), error_code(256), error_message(1024)
+            - BLOB for pickled actions (up to 64KB)
+            - JSON for content, grounding_metadata, custom_metadata, long_running_tool_ids_json
+            - BOOLEAN for partial, turn_complete, interrupted
+            - Foreign key to sessions with CASCADE delete
+            - Index on (session_id, timestamp ASC) for ordered event retrieval
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._events_table} (
+            id VARCHAR(128) PRIMARY KEY,
+            session_id VARCHAR(128) NOT NULL,
+            app_name VARCHAR(128) NOT NULL,
+            user_id VARCHAR(128) NOT NULL,
+            invocation_id VARCHAR(256) NOT NULL,
+            author VARCHAR(256) NOT NULL,
+            actions BLOB NOT NULL,
+            long_running_tool_ids_json JSON,
+            branch VARCHAR(256),
+            timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
+            content JSON,
+            grounding_metadata JSON,
+            custom_metadata JSON,
+            partial BOOLEAN,
+            turn_complete BOOLEAN,
+            interrupted BOOLEAN,
+            error_code VARCHAR(256),
+            error_message VARCHAR(1024),
+            FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE,
+            INDEX idx_{self._events_table}_session (session_id, timestamp ASC)
+        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+        """
+
+    def _get_drop_tables_sql(self) -> "list[str]":
+        """Get MySQL DROP TABLE SQL statements.
+
+        Returns:
+            List of SQL statements to drop tables and indexes.
+
+        Notes:
+            Order matters: drop events table (child) before sessions (parent).
+            MySQL automatically drops indexes when dropping tables.
+        """
+        return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
+
+    async def create_tables(self) -> None:
+        """Create both sessions and events tables if they don't exist."""
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
+        logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
+
+    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.
+
+        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 INSERT with UTC_TIMESTAMP(6) for create_time and update_time.
+            State is JSON-serialized before insertion.
+            If owner_id_column is configured, owner_id must be provided.
+        """
+        state_json = json.dumps(state)
+
+        params: tuple[Any, ...]
+        if self._owner_id_column_name:
+            sql = f"""
+            INSERT INTO {self._session_table} (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
+            VALUES (%s, %s, %s, %s, %s, UTC_TIMESTAMP(6), UTC_TIMESTAMP(6))
+            """
+            params = (session_id, app_name, user_id, owner_id, state_json)
+        else:
+            sql = f"""
+            INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time)
+            VALUES (%s, %s, %s, %s, UTC_TIMESTAMP(6), UTC_TIMESTAMP(6))
+            """
+            params = (session_id, app_name, user_id, state_json)
+
+        async with self._config.provide_connection() as conn, conn.cursor() as cursor:
+            await cursor.execute(sql, params)
+            await conn.commit()
+
+        return await self.get_session(session_id)  # type: ignore[return-value]
+
+    async def get_session(self, session_id: str) -> "SessionRecord | None":
+        """Get session by ID.
+
+        Args:
+            session_id: Session identifier.
+
+        Returns:
+            Session record or None if not found.
+
+        Notes:
+            MySQL returns datetime objects for TIMESTAMP columns.
+            JSON is parsed from database storage.
+        """
+        sql = f"""
+        SELECT id, app_name, user_id, state, create_time, update_time
+        FROM {self._session_table}
+        WHERE id = %s
+        """
+
+        try:
+            async with self._config.provide_connection() as conn, conn.cursor() as cursor:
+                await cursor.execute(sql, (session_id,))
+                row = await cursor.fetchone()
+
+                if row is None:
+                    return None
+
+                session_id_val, app_name, user_id, state_json, create_time, update_time = row
+
+                return SessionRecord(
+                    id=session_id_val,
+                    app_name=app_name,
+                    user_id=user_id,
+                    state=json.loads(state_json) if isinstance(state_json, str) else state_json,
+                    create_time=create_time,
+                    update_time=update_time,
+                )
+        except asyncmy.errors.ProgrammingError as e:  # pyright: ignore[reportAttributeAccessIssue][reportAttributeAccessIssue]
+            if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return None
+            raise
+
+    async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
+        """Update session state.
+
+        Args:
+            session_id: Session identifier.
+            state: New state dictionary (replaces existing state).
+
+        Notes:
+            This replaces the entire state dictionary.
+            Uses update_time auto-update trigger.
+        """
+        state_json = json.dumps(state)
+
+        sql = f"""
+        UPDATE {self._session_table}
+        SET state = %s
+        WHERE id = %s
+        """
+
+        async with self._config.provide_connection() as conn, conn.cursor() as cursor:
+            await cursor.execute(sql, (state_json, session_id))
+            await conn.commit()
+
+    async def delete_session(self, session_id: str) -> None:
+        """Delete session and all associated events (cascade).
+
+        Args:
+            session_id: Session identifier.
+
+        Notes:
+            Foreign key constraint ensures events are cascade-deleted.
+        """
+        sql = f"DELETE FROM {self._session_table} WHERE id = %s"
+
+        async with self._config.provide_connection() as conn, conn.cursor() as cursor:
+            await cursor.execute(sql, (session_id,))
+            await conn.commit()
+
+    async 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.
+        """
+        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 = %s
+            ORDER BY update_time DESC
+            """
+            params: tuple[str, ...] = (app_name,)
+        else:
+            sql = f"""
+            SELECT id, app_name, user_id, state, create_time, update_time
+            FROM {self._session_table}
+            WHERE app_name = %s AND user_id = %s
+            ORDER BY update_time DESC
+            """
+            params = (app_name, user_id)
+
+        try:
+            async with self._config.provide_connection() as conn, conn.cursor() as cursor:
+                await cursor.execute(sql, params)
+                rows = await cursor.fetchall()
+
+                return [
+                    SessionRecord(
+                        id=row[0],
+                        app_name=row[1],
+                        user_id=row[2],
+                        state=json.loads(row[3]) if isinstance(row[3], str) else row[3],
+                        create_time=row[4],
+                        update_time=row[5],
+                    )
+                    for row in rows
+                ]
+        except asyncmy.errors.ProgrammingError as e:  # pyright: ignore[reportAttributeAccessIssue]
+            if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return []
+            raise
+
+    async def append_event(self, event_record: EventRecord) -> None:
+        """Append an event to a session.
+
+        Args:
+            event_record: Event record to store.
+
+        Notes:
+            Uses UTC_TIMESTAMP(6) for timestamp if not provided.
+            JSON fields are serialized before insertion.
+        """
+        content_json = json.dumps(event_record.get("content")) if event_record.get("content") else None
+        grounding_metadata_json = (
+            json.dumps(event_record.get("grounding_metadata")) if event_record.get("grounding_metadata") else None
+        )
+        custom_metadata_json = (
+            json.dumps(event_record.get("custom_metadata")) if event_record.get("custom_metadata") else None
+        )
+
+        sql = f"""
+        INSERT INTO {self._events_table} (
+            id, session_id, app_name, user_id, invocation_id, author, actions,
+            long_running_tool_ids_json, branch, timestamp, content,
+            grounding_metadata, custom_metadata, partial, turn_complete,
+            interrupted, error_code, error_message
+        ) VALUES (
+            %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
+        )
+        """
+
+        async with self._config.provide_connection() as conn, conn.cursor() as cursor:
+            await cursor.execute(
+                sql,
+                (
+                    event_record["id"],
+                    event_record["session_id"],
+                    event_record["app_name"],
+                    event_record["user_id"],
+                    event_record["invocation_id"],
+                    event_record["author"],
+                    event_record["actions"],
+                    event_record.get("long_running_tool_ids_json"),
+                    event_record.get("branch"),
+                    event_record["timestamp"],
+                    content_json,
+                    grounding_metadata_json,
+                    custom_metadata_json,
+                    event_record.get("partial"),
+                    event_record.get("turn_complete"),
+                    event_record.get("interrupted"),
+                    event_record.get("error_code"),
+                    event_record.get("error_message"),
+                ),
+            )
+            await conn.commit()
+
+    async def get_events(
+        self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
+    ) -> "list[EventRecord]":
+        """Get events for a session.
+
+        Args:
+            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.
+
+        Notes:
+            Uses index on (session_id, timestamp ASC).
+            Parses JSON fields and converts BLOB actions to bytes.
+        """
+        where_clauses = ["session_id = %s"]
+        params: list[Any] = [session_id]
+
+        if after_timestamp is not None:
+            where_clauses.append("timestamp > %s")
+            params.append(after_timestamp)
+
+        where_clause = " AND ".join(where_clauses)
+        limit_clause = f" LIMIT {limit}" if limit else ""
+
+        sql = f"""
+        SELECT id, session_id, app_name, user_id, invocation_id, author, actions,
+               long_running_tool_ids_json, branch, timestamp, content,
+               grounding_metadata, custom_metadata, partial, turn_complete,
+               interrupted, error_code, error_message
+        FROM {self._events_table}
+        WHERE {where_clause}
+        ORDER BY timestamp ASC{limit_clause}
+        """
+
+        try:
+            async with self._config.provide_connection() as conn, conn.cursor() as cursor:
+                await cursor.execute(sql, params)
+                rows = await cursor.fetchall()
+
+                return [
+                    EventRecord(
+                        id=row[0],
+                        session_id=row[1],
+                        app_name=row[2],
+                        user_id=row[3],
+                        invocation_id=row[4],
+                        author=row[5],
+                        actions=bytes(row[6]),
+                        long_running_tool_ids_json=row[7],
+                        branch=row[8],
+                        timestamp=row[9],
+                        content=json.loads(row[10]) if row[10] and isinstance(row[10], str) else row[10],
+                        grounding_metadata=json.loads(row[11]) if row[11] and isinstance(row[11], str) else row[11],
+                        custom_metadata=json.loads(row[12]) if row[12] and isinstance(row[12], str) else row[12],
+                        partial=row[13],
+                        turn_complete=row[14],
+                        interrupted=row[15],
+                        error_code=row[16],
+                        error_message=row[17],
+                    )
+                    for row in rows
+                ]
+        except asyncmy.errors.ProgrammingError as e:  # pyright: ignore[reportAttributeAccessIssue]
+            if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
+                return []
+            raise

sqlspec/adapters/asyncmy/config.py

@@ -3,7 +3,7 @@
 import logging
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypedDict, Union
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
 
 import asyncmy
 from asyncmy.cursors import Cursor, DictCursor  # pyright: ignore
@@ -13,20 +13,23 @@ from typing_extensions import NotRequired
 from sqlspec.adapters.asyncmy._types import AsyncmyConnection
 from sqlspec.adapters.asyncmy.driver import AsyncmyCursor, AsyncmyDriver, asyncmy_statement_config
 from sqlspec.config import AsyncDatabaseConfig
+from sqlspec.utils.serializers import from_json, to_json
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from asyncmy.cursors import Cursor, DictCursor  # pyright: ignore
     from asyncmy.pool import Pool  # pyright: ignore
 
     from sqlspec.core.statement import StatementConfig
 
 
-__all__ = ("AsyncmyConfig", "AsyncmyConnectionParams", "AsyncmyPoolParams")
+__all__ = ("AsyncmyConfig", "AsyncmyConnectionParams", "AsyncmyDriverFeatures", "AsyncmyPoolParams")
 
 logger = logging.getLogger(__name__)
 
 
-class AsyncmyConnectionParams(TypedDict, total=False):
+class AsyncmyConnectionParams(TypedDict):
     """Asyncmy connection parameters."""
 
     host: NotRequired[str]
@@ -44,11 +47,11 @@ class AsyncmyConnectionParams(TypedDict, total=False):
     ssl: NotRequired[Any]
     sql_mode: NotRequired[str]
     init_command: NotRequired[str]
-    cursor_class: NotRequired[Union[type["Cursor"], type["DictCursor"]]]
+    cursor_class: NotRequired[type["Cursor"] | type["DictCursor"]]
     extra: NotRequired[dict[str, Any]]
 
 
-class AsyncmyPoolParams(AsyncmyConnectionParams, total=False):
+class AsyncmyPoolParams(AsyncmyConnectionParams):
     """Asyncmy pool parameters."""
 
     minsize: NotRequired[int]
@@ -57,21 +60,41 @@ class AsyncmyPoolParams(AsyncmyConnectionParams, total=False):
     pool_recycle: NotRequired[int]
 
 
+class AsyncmyDriverFeatures(TypedDict):
+    """Asyncmy driver feature flags.
+
+    MySQL/MariaDB handle JSON natively, but custom serializers can be provided
+    for specialized use cases (e.g., orjson for performance, msgspec for type safety).
+
+    json_serializer: Custom JSON serializer function.
+        Defaults to sqlspec.utils.serializers.to_json.
+        Use for performance (orjson) or custom encoding.
+    json_deserializer: Custom JSON deserializer function.
+        Defaults to sqlspec.utils.serializers.from_json.
+        Use for performance (orjson) or custom decoding.
+    """
+
+    json_serializer: NotRequired["Callable[[Any], str]"]
+    json_deserializer: NotRequired["Callable[[str], Any]"]
+
+
 class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", AsyncmyDriver]):  # pyright: ignore
     """Configuration for Asyncmy database connections."""
 
     driver_type: ClassVar[type[AsyncmyDriver]] = AsyncmyDriver
     connection_type: "ClassVar[type[AsyncmyConnection]]" = AsyncmyConnection  # pyright: ignore
+    supports_transactional_ddl: ClassVar[bool] = False
 
     def __init__(
         self,
         *,
-        pool_config: "Optional[Union[AsyncmyPoolParams, dict[str, Any]]]" = None,
-        pool_instance: "Optional[AsyncmyPool]" = None,
-        migration_config: Optional[dict[str, Any]] = None,
-        statement_config: "Optional[StatementConfig]" = None,
-        driver_features: "Optional[dict[str, Any]]" = None,
-        bind_key: "Optional[str]" = None,
+        pool_config: "AsyncmyPoolParams | dict[str, Any] | None" = None,
+        pool_instance: "AsyncmyPool | None" = None,
+        migration_config: dict[str, Any] | None = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "AsyncmyDriverFeatures | dict[str, Any] | None" = None,
+        bind_key: "str | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | None" = None,
     ) -> None:
         """Initialize Asyncmy configuration.
 
@@ -80,8 +103,9 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
             pool_instance: Existing pool instance to use
             migration_config: Migration configuration
             statement_config: Statement configuration override
-            driver_features: Driver feature configuration
+            driver_features: Driver feature configuration (TypedDict or dict)
             bind_key: Optional unique identifier for this configuration
+            extension_config: Extension-specific configuration (e.g., Litestar plugin settings)
         """
         processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
         if "extra" in processed_pool_config:
@@ -96,17 +120,33 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
         if statement_config is None:
             statement_config = asyncmy_statement_config
 
+        processed_driver_features: dict[str, Any] = dict(driver_features) if driver_features else {}
+
+        if "json_serializer" not in processed_driver_features:
+            processed_driver_features["json_serializer"] = to_json
+        if "json_deserializer" not in processed_driver_features:
+            processed_driver_features["json_deserializer"] = from_json
+
         super().__init__(
             pool_config=processed_pool_config,
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config,
-            driver_features=driver_features or {},
+            driver_features=processed_driver_features,
             bind_key=bind_key,
+            extension_config=extension_config,
         )
 
     async def _create_pool(self) -> "AsyncmyPool":  # pyright: ignore
-        """Create the actual async connection pool."""
+        """Create the actual async connection pool.
+
+        MySQL/MariaDB handle JSON types natively without requiring connection-level
+        type handlers. JSON serialization is handled via type_coercion_map in the
+        driver's statement_config (see driver.py).
+
+        Future driver_features can be added here if needed (e.g., custom connection
+        initialization, specialized type handling).
+        """
         return await asyncmy.create_pool(**dict(self.pool_config))  # pyright: ignore
 
     async def _close_pool(self) -> None:
@@ -146,7 +186,7 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
 
     @asynccontextmanager
     async def provide_session(
-        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+        self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
     ) -> AsyncGenerator[AsyncmyDriver, None]:
         """Provide an async driver session context manager.
 
@@ -159,8 +199,10 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
             An AsyncmyDriver instance.
         """
         async with self.provide_connection(*args, **kwargs) as connection:
-            final_statement_config = statement_config or asyncmy_statement_config
-            yield self.driver_type(connection=connection, statement_config=final_statement_config)
+            final_statement_config = statement_config or self.statement_config or asyncmy_statement_config
+            yield self.driver_type(
+                connection=connection, statement_config=final_statement_config, driver_features=self.driver_features
+            )
 
     async def provide_pool(self, *args: Any, **kwargs: Any) -> "Pool":  # pyright: ignore
         """Provide async pool instance.