sqlspec 0.26.0__py3-none-any.whl → 0.27.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of sqlspec might be problematic. Click here for more details.
- sqlspec/__init__.py +7 -15
- sqlspec/_serialization.py +55 -25
- sqlspec/_typing.py +62 -52
- sqlspec/adapters/adbc/_types.py +1 -1
- sqlspec/adapters/adbc/adk/__init__.py +5 -0
- sqlspec/adapters/adbc/adk/store.py +870 -0
- sqlspec/adapters/adbc/config.py +62 -12
- sqlspec/adapters/adbc/data_dictionary.py +52 -2
- sqlspec/adapters/adbc/driver.py +144 -45
- sqlspec/adapters/adbc/litestar/__init__.py +5 -0
- sqlspec/adapters/adbc/litestar/store.py +504 -0
- sqlspec/adapters/adbc/type_converter.py +44 -50
- sqlspec/adapters/aiosqlite/_types.py +1 -1
- sqlspec/adapters/aiosqlite/adk/__init__.py +5 -0
- sqlspec/adapters/aiosqlite/adk/store.py +527 -0
- sqlspec/adapters/aiosqlite/config.py +86 -16
- sqlspec/adapters/aiosqlite/data_dictionary.py +34 -2
- sqlspec/adapters/aiosqlite/driver.py +127 -38
- sqlspec/adapters/aiosqlite/litestar/__init__.py +5 -0
- sqlspec/adapters/aiosqlite/litestar/store.py +281 -0
- sqlspec/adapters/aiosqlite/pool.py +7 -7
- sqlspec/adapters/asyncmy/__init__.py +7 -1
- sqlspec/adapters/asyncmy/_types.py +1 -1
- sqlspec/adapters/asyncmy/adk/__init__.py +5 -0
- sqlspec/adapters/asyncmy/adk/store.py +493 -0
- sqlspec/adapters/asyncmy/config.py +59 -17
- sqlspec/adapters/asyncmy/data_dictionary.py +41 -2
- sqlspec/adapters/asyncmy/driver.py +293 -62
- sqlspec/adapters/asyncmy/litestar/__init__.py +5 -0
- sqlspec/adapters/asyncmy/litestar/store.py +296 -0
- sqlspec/adapters/asyncpg/__init__.py +2 -1
- sqlspec/adapters/asyncpg/_type_handlers.py +71 -0
- sqlspec/adapters/asyncpg/_types.py +11 -7
- sqlspec/adapters/asyncpg/adk/__init__.py +5 -0
- sqlspec/adapters/asyncpg/adk/store.py +450 -0
- sqlspec/adapters/asyncpg/config.py +57 -36
- sqlspec/adapters/asyncpg/data_dictionary.py +41 -2
- sqlspec/adapters/asyncpg/driver.py +153 -23
- sqlspec/adapters/asyncpg/litestar/__init__.py +5 -0
- sqlspec/adapters/asyncpg/litestar/store.py +253 -0
- sqlspec/adapters/bigquery/_types.py +1 -1
- sqlspec/adapters/bigquery/adk/__init__.py +5 -0
- sqlspec/adapters/bigquery/adk/store.py +576 -0
- sqlspec/adapters/bigquery/config.py +25 -11
- sqlspec/adapters/bigquery/data_dictionary.py +42 -2
- sqlspec/adapters/bigquery/driver.py +352 -144
- sqlspec/adapters/bigquery/litestar/__init__.py +5 -0
- sqlspec/adapters/bigquery/litestar/store.py +327 -0
- sqlspec/adapters/bigquery/type_converter.py +55 -23
- sqlspec/adapters/duckdb/_types.py +2 -2
- sqlspec/adapters/duckdb/adk/__init__.py +14 -0
- sqlspec/adapters/duckdb/adk/store.py +553 -0
- sqlspec/adapters/duckdb/config.py +79 -21
- sqlspec/adapters/duckdb/data_dictionary.py +41 -2
- sqlspec/adapters/duckdb/driver.py +138 -43
- sqlspec/adapters/duckdb/litestar/__init__.py +5 -0
- sqlspec/adapters/duckdb/litestar/store.py +332 -0
- sqlspec/adapters/duckdb/pool.py +5 -5
- sqlspec/adapters/duckdb/type_converter.py +51 -21
- sqlspec/adapters/oracledb/_numpy_handlers.py +133 -0
- sqlspec/adapters/oracledb/_types.py +20 -2
- sqlspec/adapters/oracledb/adk/__init__.py +5 -0
- sqlspec/adapters/oracledb/adk/store.py +1745 -0
- sqlspec/adapters/oracledb/config.py +120 -36
- sqlspec/adapters/oracledb/data_dictionary.py +87 -20
- sqlspec/adapters/oracledb/driver.py +292 -84
- sqlspec/adapters/oracledb/litestar/__init__.py +5 -0
- sqlspec/adapters/oracledb/litestar/store.py +767 -0
- sqlspec/adapters/oracledb/migrations.py +316 -25
- sqlspec/adapters/oracledb/type_converter.py +91 -16
- sqlspec/adapters/psqlpy/_type_handlers.py +44 -0
- sqlspec/adapters/psqlpy/_types.py +2 -1
- sqlspec/adapters/psqlpy/adk/__init__.py +5 -0
- sqlspec/adapters/psqlpy/adk/store.py +482 -0
- sqlspec/adapters/psqlpy/config.py +45 -19
- sqlspec/adapters/psqlpy/data_dictionary.py +41 -2
- sqlspec/adapters/psqlpy/driver.py +101 -31
- sqlspec/adapters/psqlpy/litestar/__init__.py +5 -0
- sqlspec/adapters/psqlpy/litestar/store.py +272 -0
- sqlspec/adapters/psqlpy/type_converter.py +40 -11
- sqlspec/adapters/psycopg/_type_handlers.py +80 -0
- sqlspec/adapters/psycopg/_types.py +2 -1
- sqlspec/adapters/psycopg/adk/__init__.py +5 -0
- sqlspec/adapters/psycopg/adk/store.py +944 -0
- sqlspec/adapters/psycopg/config.py +65 -37
- sqlspec/adapters/psycopg/data_dictionary.py +77 -3
- sqlspec/adapters/psycopg/driver.py +200 -78
- sqlspec/adapters/psycopg/litestar/__init__.py +5 -0
- sqlspec/adapters/psycopg/litestar/store.py +554 -0
- sqlspec/adapters/sqlite/__init__.py +2 -1
- sqlspec/adapters/sqlite/_type_handlers.py +86 -0
- sqlspec/adapters/sqlite/_types.py +1 -1
- sqlspec/adapters/sqlite/adk/__init__.py +5 -0
- sqlspec/adapters/sqlite/adk/store.py +572 -0
- sqlspec/adapters/sqlite/config.py +85 -16
- sqlspec/adapters/sqlite/data_dictionary.py +34 -2
- sqlspec/adapters/sqlite/driver.py +120 -52
- sqlspec/adapters/sqlite/litestar/__init__.py +5 -0
- sqlspec/adapters/sqlite/litestar/store.py +318 -0
- sqlspec/adapters/sqlite/pool.py +5 -5
- sqlspec/base.py +45 -26
- sqlspec/builder/__init__.py +73 -4
- sqlspec/builder/_base.py +91 -58
- sqlspec/builder/_column.py +5 -5
- sqlspec/builder/_ddl.py +98 -89
- sqlspec/builder/_delete.py +5 -4
- sqlspec/builder/_dml.py +388 -0
- sqlspec/{_sql.py → builder/_factory.py} +41 -44
- sqlspec/builder/_insert.py +5 -82
- sqlspec/builder/{mixins/_join_operations.py → _join.py} +145 -143
- sqlspec/builder/_merge.py +446 -11
- sqlspec/builder/_parsing_utils.py +9 -11
- sqlspec/builder/_select.py +1313 -25
- sqlspec/builder/_update.py +11 -42
- sqlspec/cli.py +76 -69
- sqlspec/config.py +231 -60
- sqlspec/core/__init__.py +5 -4
- sqlspec/core/cache.py +18 -18
- sqlspec/core/compiler.py +6 -8
- sqlspec/core/filters.py +37 -37
- sqlspec/core/hashing.py +9 -9
- sqlspec/core/parameters.py +76 -45
- sqlspec/core/result.py +102 -46
- sqlspec/core/splitter.py +16 -17
- sqlspec/core/statement.py +32 -31
- sqlspec/core/type_conversion.py +3 -2
- sqlspec/driver/__init__.py +1 -3
- sqlspec/driver/_async.py +95 -161
- sqlspec/driver/_common.py +133 -80
- sqlspec/driver/_sync.py +95 -162
- sqlspec/driver/mixins/_result_tools.py +20 -236
- sqlspec/driver/mixins/_sql_translator.py +4 -4
- sqlspec/exceptions.py +70 -7
- sqlspec/extensions/adk/__init__.py +53 -0
- sqlspec/extensions/adk/_types.py +51 -0
- sqlspec/extensions/adk/converters.py +172 -0
- sqlspec/extensions/adk/migrations/0001_create_adk_tables.py +144 -0
- sqlspec/extensions/adk/migrations/__init__.py +0 -0
- sqlspec/extensions/adk/service.py +181 -0
- sqlspec/extensions/adk/store.py +536 -0
- sqlspec/extensions/aiosql/adapter.py +73 -53
- sqlspec/extensions/litestar/__init__.py +21 -4
- sqlspec/extensions/litestar/cli.py +54 -10
- sqlspec/extensions/litestar/config.py +59 -266
- sqlspec/extensions/litestar/handlers.py +46 -17
- sqlspec/extensions/litestar/migrations/0001_create_session_table.py +137 -0
- sqlspec/extensions/litestar/migrations/__init__.py +3 -0
- sqlspec/extensions/litestar/plugin.py +324 -223
- sqlspec/extensions/litestar/providers.py +25 -25
- sqlspec/extensions/litestar/store.py +265 -0
- sqlspec/loader.py +30 -49
- sqlspec/migrations/base.py +200 -76
- sqlspec/migrations/commands.py +591 -62
- sqlspec/migrations/context.py +6 -9
- sqlspec/migrations/fix.py +199 -0
- sqlspec/migrations/loaders.py +47 -19
- sqlspec/migrations/runner.py +241 -75
- sqlspec/migrations/tracker.py +237 -21
- sqlspec/migrations/utils.py +51 -3
- sqlspec/migrations/validation.py +177 -0
- sqlspec/protocols.py +66 -36
- sqlspec/storage/_utils.py +98 -0
- sqlspec/storage/backends/fsspec.py +134 -106
- sqlspec/storage/backends/local.py +78 -51
- sqlspec/storage/backends/obstore.py +278 -162
- sqlspec/storage/registry.py +75 -39
- sqlspec/typing.py +14 -84
- sqlspec/utils/config_resolver.py +6 -6
- sqlspec/utils/correlation.py +4 -5
- sqlspec/utils/data_transformation.py +3 -2
- sqlspec/utils/deprecation.py +9 -8
- sqlspec/utils/fixtures.py +4 -4
- sqlspec/utils/logging.py +46 -6
- sqlspec/utils/module_loader.py +2 -2
- sqlspec/utils/schema.py +288 -0
- sqlspec/utils/serializers.py +3 -3
- sqlspec/utils/sync_tools.py +21 -17
- sqlspec/utils/text.py +1 -2
- sqlspec/utils/type_guards.py +111 -20
- sqlspec/utils/version.py +433 -0
- {sqlspec-0.26.0.dist-info → sqlspec-0.27.0.dist-info}/METADATA +40 -21
- sqlspec-0.27.0.dist-info/RECORD +207 -0
- sqlspec/builder/mixins/__init__.py +0 -55
- sqlspec/builder/mixins/_cte_and_set_ops.py +0 -253
- sqlspec/builder/mixins/_delete_operations.py +0 -50
- sqlspec/builder/mixins/_insert_operations.py +0 -282
- sqlspec/builder/mixins/_merge_operations.py +0 -698
- sqlspec/builder/mixins/_order_limit_operations.py +0 -145
- sqlspec/builder/mixins/_pivot_operations.py +0 -157
- sqlspec/builder/mixins/_select_operations.py +0 -930
- sqlspec/builder/mixins/_update_operations.py +0 -199
- sqlspec/builder/mixins/_where_clause.py +0 -1298
- sqlspec-0.26.0.dist-info/RECORD +0 -157
- sqlspec-0.26.0.dist-info/licenses/NOTICE +0 -29
- {sqlspec-0.26.0.dist-info → sqlspec-0.27.0.dist-info}/WHEEL +0 -0
- {sqlspec-0.26.0.dist-info → sqlspec-0.27.0.dist-info}/entry_points.txt +0 -0
- {sqlspec-0.26.0.dist-info → sqlspec-0.27.0.dist-info}/licenses/LICENSE +0 -0
|
@@ -0,0 +1,493 @@
|
|
|
1
|
+
"""AsyncMy ADK store for Google Agent Development Kit session/event storage."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from typing import TYPE_CHECKING, Any, Final
|
|
5
|
+
|
|
6
|
+
import asyncmy
|
|
7
|
+
|
|
8
|
+
from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
|
|
9
|
+
from sqlspec.utils.logging import get_logger
|
|
10
|
+
|
|
11
|
+
if TYPE_CHECKING:
|
|
12
|
+
from datetime import datetime
|
|
13
|
+
|
|
14
|
+
from sqlspec.adapters.asyncmy.config import AsyncmyConfig
|
|
15
|
+
|
|
16
|
+
logger = get_logger("adapters.asyncmy.adk.store")
|
|
17
|
+
|
|
18
|
+
__all__ = ("AsyncmyADKStore",)
|
|
19
|
+
|
|
20
|
+
MYSQL_TABLE_NOT_FOUND_ERROR: Final = 1146
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
|
|
24
|
+
"""MySQL/MariaDB ADK store using AsyncMy driver.
|
|
25
|
+
|
|
26
|
+
Implements session and event storage for Google Agent Development Kit
|
|
27
|
+
using MySQL/MariaDB via the AsyncMy driver. Provides:
|
|
28
|
+
- Session state management with JSON storage
|
|
29
|
+
- Event history tracking with BLOB-serialized actions
|
|
30
|
+
- Microsecond-precision timestamps
|
|
31
|
+
- Foreign key constraints with cascade delete
|
|
32
|
+
- Efficient upserts using ON DUPLICATE KEY UPDATE
|
|
33
|
+
|
|
34
|
+
Args:
|
|
35
|
+
config: AsyncmyConfig with extension_config["adk"] settings.
|
|
36
|
+
|
|
37
|
+
Example:
|
|
38
|
+
from sqlspec.adapters.asyncmy import AsyncmyConfig
|
|
39
|
+
from sqlspec.adapters.asyncmy.adk import AsyncmyADKStore
|
|
40
|
+
|
|
41
|
+
config = AsyncmyConfig(
|
|
42
|
+
pool_config={"host": "localhost", ...},
|
|
43
|
+
extension_config={
|
|
44
|
+
"adk": {
|
|
45
|
+
"session_table": "my_sessions",
|
|
46
|
+
"events_table": "my_events",
|
|
47
|
+
"owner_id_column": "tenant_id BIGINT NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
|
|
48
|
+
}
|
|
49
|
+
}
|
|
50
|
+
)
|
|
51
|
+
store = AsyncmyADKStore(config)
|
|
52
|
+
await store.create_tables()
|
|
53
|
+
|
|
54
|
+
Notes:
|
|
55
|
+
- MySQL JSON type used (not JSONB) - requires MySQL 5.7.8+
|
|
56
|
+
- TIMESTAMP(6) provides microsecond precision
|
|
57
|
+
- InnoDB engine required for foreign key support
|
|
58
|
+
- State merging handled at application level
|
|
59
|
+
- Configuration is read from config.extension_config["adk"]
|
|
60
|
+
"""
|
|
61
|
+
|
|
62
|
+
__slots__ = ()
|
|
63
|
+
|
|
64
|
+
def __init__(self, config: "AsyncmyConfig") -> None:
|
|
65
|
+
"""Initialize AsyncMy ADK store.
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
config: AsyncmyConfig instance.
|
|
69
|
+
|
|
70
|
+
Notes:
|
|
71
|
+
Configuration is read from config.extension_config["adk"]:
|
|
72
|
+
- session_table: Sessions table name (default: "adk_sessions")
|
|
73
|
+
- events_table: Events table name (default: "adk_events")
|
|
74
|
+
- owner_id_column: Optional owner FK column DDL (default: None)
|
|
75
|
+
"""
|
|
76
|
+
super().__init__(config)
|
|
77
|
+
|
|
78
|
+
def _parse_owner_id_column_for_mysql(self, column_ddl: str) -> "tuple[str, str]":
|
|
79
|
+
"""Parse owner ID column DDL for MySQL FOREIGN KEY syntax.
|
|
80
|
+
|
|
81
|
+
MySQL ignores inline REFERENCES syntax in column definitions.
|
|
82
|
+
This method extracts the column definition and creates a separate
|
|
83
|
+
FOREIGN KEY constraint.
|
|
84
|
+
|
|
85
|
+
Args:
|
|
86
|
+
column_ddl: Column DDL like "tenant_id BIGINT NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
|
|
87
|
+
|
|
88
|
+
Returns:
|
|
89
|
+
Tuple of (column_definition, foreign_key_constraint)
|
|
90
|
+
|
|
91
|
+
Example:
|
|
92
|
+
Input: "tenant_id BIGINT NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
|
|
93
|
+
Output: ("tenant_id BIGINT NOT NULL", "FOREIGN KEY (tenant_id) REFERENCES tenants(id) ON DELETE CASCADE")
|
|
94
|
+
"""
|
|
95
|
+
import re
|
|
96
|
+
|
|
97
|
+
references_match = re.search(r"\s+REFERENCES\s+(.+)", column_ddl, re.IGNORECASE)
|
|
98
|
+
|
|
99
|
+
if not references_match:
|
|
100
|
+
return (column_ddl.strip(), "")
|
|
101
|
+
|
|
102
|
+
col_def = column_ddl[: references_match.start()].strip()
|
|
103
|
+
fk_clause = references_match.group(1).strip()
|
|
104
|
+
col_name = col_def.split()[0]
|
|
105
|
+
fk_constraint = f"FOREIGN KEY ({col_name}) REFERENCES {fk_clause}"
|
|
106
|
+
|
|
107
|
+
return (col_def, fk_constraint)
|
|
108
|
+
|
|
109
|
+
def _get_create_sessions_table_sql(self) -> str:
|
|
110
|
+
"""Get MySQL CREATE TABLE SQL for sessions.
|
|
111
|
+
|
|
112
|
+
Returns:
|
|
113
|
+
SQL statement to create adk_sessions table with indexes.
|
|
114
|
+
|
|
115
|
+
Notes:
|
|
116
|
+
- VARCHAR(128) for IDs and names (sufficient for UUIDs and app names)
|
|
117
|
+
- JSON type for state storage (MySQL 5.7.8+)
|
|
118
|
+
- TIMESTAMP(6) with microsecond precision
|
|
119
|
+
- AUTO-UPDATE on update_time
|
|
120
|
+
- Composite index on (app_name, user_id) for listing
|
|
121
|
+
- Index on update_time DESC for recent session queries
|
|
122
|
+
- Optional owner ID column for multi-tenancy
|
|
123
|
+
- MySQL requires explicit FOREIGN KEY syntax (inline REFERENCES is ignored)
|
|
124
|
+
"""
|
|
125
|
+
owner_id_col = ""
|
|
126
|
+
fk_constraint = ""
|
|
127
|
+
|
|
128
|
+
if self._owner_id_column_ddl:
|
|
129
|
+
col_def, fk_def = self._parse_owner_id_column_for_mysql(self._owner_id_column_ddl)
|
|
130
|
+
owner_id_col = f"{col_def},"
|
|
131
|
+
if fk_def:
|
|
132
|
+
fk_constraint = f",\n {fk_def}"
|
|
133
|
+
|
|
134
|
+
return f"""
|
|
135
|
+
CREATE TABLE IF NOT EXISTS {self._session_table} (
|
|
136
|
+
id VARCHAR(128) PRIMARY KEY,
|
|
137
|
+
app_name VARCHAR(128) NOT NULL,
|
|
138
|
+
user_id VARCHAR(128) NOT NULL,
|
|
139
|
+
{owner_id_col}
|
|
140
|
+
state JSON NOT NULL,
|
|
141
|
+
create_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
|
|
142
|
+
update_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
|
|
143
|
+
INDEX idx_{self._session_table}_app_user (app_name, user_id),
|
|
144
|
+
INDEX idx_{self._session_table}_update_time (update_time DESC){fk_constraint}
|
|
145
|
+
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
|
|
146
|
+
"""
|
|
147
|
+
|
|
148
|
+
def _get_create_events_table_sql(self) -> str:
|
|
149
|
+
"""Get MySQL CREATE TABLE SQL for events.
|
|
150
|
+
|
|
151
|
+
Returns:
|
|
152
|
+
SQL statement to create adk_events table with indexes.
|
|
153
|
+
|
|
154
|
+
Notes:
|
|
155
|
+
- VARCHAR sizes: id(128), session_id(128), invocation_id(256), author(256),
|
|
156
|
+
branch(256), error_code(256), error_message(1024)
|
|
157
|
+
- BLOB for pickled actions (up to 64KB)
|
|
158
|
+
- JSON for content, grounding_metadata, custom_metadata, long_running_tool_ids_json
|
|
159
|
+
- BOOLEAN for partial, turn_complete, interrupted
|
|
160
|
+
- Foreign key to sessions with CASCADE delete
|
|
161
|
+
- Index on (session_id, timestamp ASC) for ordered event retrieval
|
|
162
|
+
"""
|
|
163
|
+
return f"""
|
|
164
|
+
CREATE TABLE IF NOT EXISTS {self._events_table} (
|
|
165
|
+
id VARCHAR(128) PRIMARY KEY,
|
|
166
|
+
session_id VARCHAR(128) NOT NULL,
|
|
167
|
+
app_name VARCHAR(128) NOT NULL,
|
|
168
|
+
user_id VARCHAR(128) NOT NULL,
|
|
169
|
+
invocation_id VARCHAR(256) NOT NULL,
|
|
170
|
+
author VARCHAR(256) NOT NULL,
|
|
171
|
+
actions BLOB NOT NULL,
|
|
172
|
+
long_running_tool_ids_json JSON,
|
|
173
|
+
branch VARCHAR(256),
|
|
174
|
+
timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
|
|
175
|
+
content JSON,
|
|
176
|
+
grounding_metadata JSON,
|
|
177
|
+
custom_metadata JSON,
|
|
178
|
+
partial BOOLEAN,
|
|
179
|
+
turn_complete BOOLEAN,
|
|
180
|
+
interrupted BOOLEAN,
|
|
181
|
+
error_code VARCHAR(256),
|
|
182
|
+
error_message VARCHAR(1024),
|
|
183
|
+
FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE,
|
|
184
|
+
INDEX idx_{self._events_table}_session (session_id, timestamp ASC)
|
|
185
|
+
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
|
|
186
|
+
"""
|
|
187
|
+
|
|
188
|
+
def _get_drop_tables_sql(self) -> "list[str]":
|
|
189
|
+
"""Get MySQL DROP TABLE SQL statements.
|
|
190
|
+
|
|
191
|
+
Returns:
|
|
192
|
+
List of SQL statements to drop tables and indexes.
|
|
193
|
+
|
|
194
|
+
Notes:
|
|
195
|
+
Order matters: drop events table (child) before sessions (parent).
|
|
196
|
+
MySQL automatically drops indexes when dropping tables.
|
|
197
|
+
"""
|
|
198
|
+
return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
|
|
199
|
+
|
|
200
|
+
async def create_tables(self) -> None:
|
|
201
|
+
"""Create both sessions and events tables if they don't exist."""
|
|
202
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cursor:
|
|
203
|
+
await cursor.execute(self._get_create_sessions_table_sql())
|
|
204
|
+
await cursor.execute(self._get_create_events_table_sql())
|
|
205
|
+
logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
|
|
206
|
+
|
|
207
|
+
async def create_session(
|
|
208
|
+
self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
|
|
209
|
+
) -> SessionRecord:
|
|
210
|
+
"""Create a new session.
|
|
211
|
+
|
|
212
|
+
Args:
|
|
213
|
+
session_id: Unique session identifier.
|
|
214
|
+
app_name: Application name.
|
|
215
|
+
user_id: User identifier.
|
|
216
|
+
state: Initial session state.
|
|
217
|
+
owner_id: Optional owner ID value for owner_id_column (if configured).
|
|
218
|
+
|
|
219
|
+
Returns:
|
|
220
|
+
Created session record.
|
|
221
|
+
|
|
222
|
+
Notes:
|
|
223
|
+
Uses INSERT with UTC_TIMESTAMP(6) for create_time and update_time.
|
|
224
|
+
State is JSON-serialized before insertion.
|
|
225
|
+
If owner_id_column is configured, owner_id must be provided.
|
|
226
|
+
"""
|
|
227
|
+
state_json = json.dumps(state)
|
|
228
|
+
|
|
229
|
+
params: tuple[Any, ...]
|
|
230
|
+
if self._owner_id_column_name:
|
|
231
|
+
sql = f"""
|
|
232
|
+
INSERT INTO {self._session_table} (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
|
|
233
|
+
VALUES (%s, %s, %s, %s, %s, UTC_TIMESTAMP(6), UTC_TIMESTAMP(6))
|
|
234
|
+
"""
|
|
235
|
+
params = (session_id, app_name, user_id, owner_id, state_json)
|
|
236
|
+
else:
|
|
237
|
+
sql = f"""
|
|
238
|
+
INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time)
|
|
239
|
+
VALUES (%s, %s, %s, %s, UTC_TIMESTAMP(6), UTC_TIMESTAMP(6))
|
|
240
|
+
"""
|
|
241
|
+
params = (session_id, app_name, user_id, state_json)
|
|
242
|
+
|
|
243
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cursor:
|
|
244
|
+
await cursor.execute(sql, params)
|
|
245
|
+
await conn.commit()
|
|
246
|
+
|
|
247
|
+
return await self.get_session(session_id) # type: ignore[return-value]
|
|
248
|
+
|
|
249
|
+
async def get_session(self, session_id: str) -> "SessionRecord | None":
|
|
250
|
+
"""Get session by ID.
|
|
251
|
+
|
|
252
|
+
Args:
|
|
253
|
+
session_id: Session identifier.
|
|
254
|
+
|
|
255
|
+
Returns:
|
|
256
|
+
Session record or None if not found.
|
|
257
|
+
|
|
258
|
+
Notes:
|
|
259
|
+
MySQL returns datetime objects for TIMESTAMP columns.
|
|
260
|
+
JSON is parsed from database storage.
|
|
261
|
+
"""
|
|
262
|
+
sql = f"""
|
|
263
|
+
SELECT id, app_name, user_id, state, create_time, update_time
|
|
264
|
+
FROM {self._session_table}
|
|
265
|
+
WHERE id = %s
|
|
266
|
+
"""
|
|
267
|
+
|
|
268
|
+
try:
|
|
269
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cursor:
|
|
270
|
+
await cursor.execute(sql, (session_id,))
|
|
271
|
+
row = await cursor.fetchone()
|
|
272
|
+
|
|
273
|
+
if row is None:
|
|
274
|
+
return None
|
|
275
|
+
|
|
276
|
+
session_id_val, app_name, user_id, state_json, create_time, update_time = row
|
|
277
|
+
|
|
278
|
+
return SessionRecord(
|
|
279
|
+
id=session_id_val,
|
|
280
|
+
app_name=app_name,
|
|
281
|
+
user_id=user_id,
|
|
282
|
+
state=json.loads(state_json) if isinstance(state_json, str) else state_json,
|
|
283
|
+
create_time=create_time,
|
|
284
|
+
update_time=update_time,
|
|
285
|
+
)
|
|
286
|
+
except asyncmy.errors.ProgrammingError as e: # pyright: ignore[reportAttributeAccessIssue][reportAttributeAccessIssue]
|
|
287
|
+
if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
|
|
288
|
+
return None
|
|
289
|
+
raise
|
|
290
|
+
|
|
291
|
+
async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
|
|
292
|
+
"""Update session state.
|
|
293
|
+
|
|
294
|
+
Args:
|
|
295
|
+
session_id: Session identifier.
|
|
296
|
+
state: New state dictionary (replaces existing state).
|
|
297
|
+
|
|
298
|
+
Notes:
|
|
299
|
+
This replaces the entire state dictionary.
|
|
300
|
+
Uses update_time auto-update trigger.
|
|
301
|
+
"""
|
|
302
|
+
state_json = json.dumps(state)
|
|
303
|
+
|
|
304
|
+
sql = f"""
|
|
305
|
+
UPDATE {self._session_table}
|
|
306
|
+
SET state = %s
|
|
307
|
+
WHERE id = %s
|
|
308
|
+
"""
|
|
309
|
+
|
|
310
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cursor:
|
|
311
|
+
await cursor.execute(sql, (state_json, session_id))
|
|
312
|
+
await conn.commit()
|
|
313
|
+
|
|
314
|
+
async def delete_session(self, session_id: str) -> None:
|
|
315
|
+
"""Delete session and all associated events (cascade).
|
|
316
|
+
|
|
317
|
+
Args:
|
|
318
|
+
session_id: Session identifier.
|
|
319
|
+
|
|
320
|
+
Notes:
|
|
321
|
+
Foreign key constraint ensures events are cascade-deleted.
|
|
322
|
+
"""
|
|
323
|
+
sql = f"DELETE FROM {self._session_table} WHERE id = %s"
|
|
324
|
+
|
|
325
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cursor:
|
|
326
|
+
await cursor.execute(sql, (session_id,))
|
|
327
|
+
await conn.commit()
|
|
328
|
+
|
|
329
|
+
async def list_sessions(self, app_name: str, user_id: str) -> "list[SessionRecord]":
|
|
330
|
+
"""List all sessions for a user in an app.
|
|
331
|
+
|
|
332
|
+
Args:
|
|
333
|
+
app_name: Application name.
|
|
334
|
+
user_id: User identifier.
|
|
335
|
+
|
|
336
|
+
Returns:
|
|
337
|
+
List of session records ordered by update_time DESC.
|
|
338
|
+
|
|
339
|
+
Notes:
|
|
340
|
+
Uses composite index on (app_name, user_id).
|
|
341
|
+
"""
|
|
342
|
+
sql = f"""
|
|
343
|
+
SELECT id, app_name, user_id, state, create_time, update_time
|
|
344
|
+
FROM {self._session_table}
|
|
345
|
+
WHERE app_name = %s AND user_id = %s
|
|
346
|
+
ORDER BY update_time DESC
|
|
347
|
+
"""
|
|
348
|
+
|
|
349
|
+
try:
|
|
350
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cursor:
|
|
351
|
+
await cursor.execute(sql, (app_name, user_id))
|
|
352
|
+
rows = await cursor.fetchall()
|
|
353
|
+
|
|
354
|
+
return [
|
|
355
|
+
SessionRecord(
|
|
356
|
+
id=row[0],
|
|
357
|
+
app_name=row[1],
|
|
358
|
+
user_id=row[2],
|
|
359
|
+
state=json.loads(row[3]) if isinstance(row[3], str) else row[3],
|
|
360
|
+
create_time=row[4],
|
|
361
|
+
update_time=row[5],
|
|
362
|
+
)
|
|
363
|
+
for row in rows
|
|
364
|
+
]
|
|
365
|
+
except asyncmy.errors.ProgrammingError as e: # pyright: ignore[reportAttributeAccessIssue]
|
|
366
|
+
if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
|
|
367
|
+
return []
|
|
368
|
+
raise
|
|
369
|
+
|
|
370
|
+
async def append_event(self, event_record: EventRecord) -> None:
|
|
371
|
+
"""Append an event to a session.
|
|
372
|
+
|
|
373
|
+
Args:
|
|
374
|
+
event_record: Event record to store.
|
|
375
|
+
|
|
376
|
+
Notes:
|
|
377
|
+
Uses UTC_TIMESTAMP(6) for timestamp if not provided.
|
|
378
|
+
JSON fields are serialized before insertion.
|
|
379
|
+
"""
|
|
380
|
+
content_json = json.dumps(event_record.get("content")) if event_record.get("content") else None
|
|
381
|
+
grounding_metadata_json = (
|
|
382
|
+
json.dumps(event_record.get("grounding_metadata")) if event_record.get("grounding_metadata") else None
|
|
383
|
+
)
|
|
384
|
+
custom_metadata_json = (
|
|
385
|
+
json.dumps(event_record.get("custom_metadata")) if event_record.get("custom_metadata") else None
|
|
386
|
+
)
|
|
387
|
+
|
|
388
|
+
sql = f"""
|
|
389
|
+
INSERT INTO {self._events_table} (
|
|
390
|
+
id, session_id, app_name, user_id, invocation_id, author, actions,
|
|
391
|
+
long_running_tool_ids_json, branch, timestamp, content,
|
|
392
|
+
grounding_metadata, custom_metadata, partial, turn_complete,
|
|
393
|
+
interrupted, error_code, error_message
|
|
394
|
+
) VALUES (
|
|
395
|
+
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
|
|
396
|
+
)
|
|
397
|
+
"""
|
|
398
|
+
|
|
399
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cursor:
|
|
400
|
+
await cursor.execute(
|
|
401
|
+
sql,
|
|
402
|
+
(
|
|
403
|
+
event_record["id"],
|
|
404
|
+
event_record["session_id"],
|
|
405
|
+
event_record["app_name"],
|
|
406
|
+
event_record["user_id"],
|
|
407
|
+
event_record["invocation_id"],
|
|
408
|
+
event_record["author"],
|
|
409
|
+
event_record["actions"],
|
|
410
|
+
event_record.get("long_running_tool_ids_json"),
|
|
411
|
+
event_record.get("branch"),
|
|
412
|
+
event_record["timestamp"],
|
|
413
|
+
content_json,
|
|
414
|
+
grounding_metadata_json,
|
|
415
|
+
custom_metadata_json,
|
|
416
|
+
event_record.get("partial"),
|
|
417
|
+
event_record.get("turn_complete"),
|
|
418
|
+
event_record.get("interrupted"),
|
|
419
|
+
event_record.get("error_code"),
|
|
420
|
+
event_record.get("error_message"),
|
|
421
|
+
),
|
|
422
|
+
)
|
|
423
|
+
await conn.commit()
|
|
424
|
+
|
|
425
|
+
async def get_events(
|
|
426
|
+
self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
|
|
427
|
+
) -> "list[EventRecord]":
|
|
428
|
+
"""Get events for a session.
|
|
429
|
+
|
|
430
|
+
Args:
|
|
431
|
+
session_id: Session identifier.
|
|
432
|
+
after_timestamp: Only return events after this time.
|
|
433
|
+
limit: Maximum number of events to return.
|
|
434
|
+
|
|
435
|
+
Returns:
|
|
436
|
+
List of event records ordered by timestamp ASC.
|
|
437
|
+
|
|
438
|
+
Notes:
|
|
439
|
+
Uses index on (session_id, timestamp ASC).
|
|
440
|
+
Parses JSON fields and converts BLOB actions to bytes.
|
|
441
|
+
"""
|
|
442
|
+
where_clauses = ["session_id = %s"]
|
|
443
|
+
params: list[Any] = [session_id]
|
|
444
|
+
|
|
445
|
+
if after_timestamp is not None:
|
|
446
|
+
where_clauses.append("timestamp > %s")
|
|
447
|
+
params.append(after_timestamp)
|
|
448
|
+
|
|
449
|
+
where_clause = " AND ".join(where_clauses)
|
|
450
|
+
limit_clause = f" LIMIT {limit}" if limit else ""
|
|
451
|
+
|
|
452
|
+
sql = f"""
|
|
453
|
+
SELECT id, session_id, app_name, user_id, invocation_id, author, actions,
|
|
454
|
+
long_running_tool_ids_json, branch, timestamp, content,
|
|
455
|
+
grounding_metadata, custom_metadata, partial, turn_complete,
|
|
456
|
+
interrupted, error_code, error_message
|
|
457
|
+
FROM {self._events_table}
|
|
458
|
+
WHERE {where_clause}
|
|
459
|
+
ORDER BY timestamp ASC{limit_clause}
|
|
460
|
+
"""
|
|
461
|
+
|
|
462
|
+
try:
|
|
463
|
+
async with self._config.provide_connection() as conn, conn.cursor() as cursor:
|
|
464
|
+
await cursor.execute(sql, params)
|
|
465
|
+
rows = await cursor.fetchall()
|
|
466
|
+
|
|
467
|
+
return [
|
|
468
|
+
EventRecord(
|
|
469
|
+
id=row[0],
|
|
470
|
+
session_id=row[1],
|
|
471
|
+
app_name=row[2],
|
|
472
|
+
user_id=row[3],
|
|
473
|
+
invocation_id=row[4],
|
|
474
|
+
author=row[5],
|
|
475
|
+
actions=bytes(row[6]),
|
|
476
|
+
long_running_tool_ids_json=row[7],
|
|
477
|
+
branch=row[8],
|
|
478
|
+
timestamp=row[9],
|
|
479
|
+
content=json.loads(row[10]) if row[10] and isinstance(row[10], str) else row[10],
|
|
480
|
+
grounding_metadata=json.loads(row[11]) if row[11] and isinstance(row[11], str) else row[11],
|
|
481
|
+
custom_metadata=json.loads(row[12]) if row[12] and isinstance(row[12], str) else row[12],
|
|
482
|
+
partial=row[13],
|
|
483
|
+
turn_complete=row[14],
|
|
484
|
+
interrupted=row[15],
|
|
485
|
+
error_code=row[16],
|
|
486
|
+
error_message=row[17],
|
|
487
|
+
)
|
|
488
|
+
for row in rows
|
|
489
|
+
]
|
|
490
|
+
except asyncmy.errors.ProgrammingError as e: # pyright: ignore[reportAttributeAccessIssue]
|
|
491
|
+
if "doesn't exist" in str(e) or e.args[0] == MYSQL_TABLE_NOT_FOUND_ERROR:
|
|
492
|
+
return []
|
|
493
|
+
raise
|
|
@@ -3,7 +3,7 @@
|
|
|
3
3
|
import logging
|
|
4
4
|
from collections.abc import AsyncGenerator
|
|
5
5
|
from contextlib import asynccontextmanager
|
|
6
|
-
from typing import TYPE_CHECKING, Any, ClassVar,
|
|
6
|
+
from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
|
|
7
7
|
|
|
8
8
|
import asyncmy
|
|
9
9
|
from asyncmy.cursors import Cursor, DictCursor # pyright: ignore
|
|
@@ -13,20 +13,23 @@ from typing_extensions import NotRequired
|
|
|
13
13
|
from sqlspec.adapters.asyncmy._types import AsyncmyConnection
|
|
14
14
|
from sqlspec.adapters.asyncmy.driver import AsyncmyCursor, AsyncmyDriver, asyncmy_statement_config
|
|
15
15
|
from sqlspec.config import AsyncDatabaseConfig
|
|
16
|
+
from sqlspec.utils.serializers import from_json, to_json
|
|
16
17
|
|
|
17
18
|
if TYPE_CHECKING:
|
|
19
|
+
from collections.abc import Callable
|
|
20
|
+
|
|
18
21
|
from asyncmy.cursors import Cursor, DictCursor # pyright: ignore
|
|
19
22
|
from asyncmy.pool import Pool # pyright: ignore
|
|
20
23
|
|
|
21
24
|
from sqlspec.core.statement import StatementConfig
|
|
22
25
|
|
|
23
26
|
|
|
24
|
-
__all__ = ("AsyncmyConfig", "AsyncmyConnectionParams", "AsyncmyPoolParams")
|
|
27
|
+
__all__ = ("AsyncmyConfig", "AsyncmyConnectionParams", "AsyncmyDriverFeatures", "AsyncmyPoolParams")
|
|
25
28
|
|
|
26
29
|
logger = logging.getLogger(__name__)
|
|
27
30
|
|
|
28
31
|
|
|
29
|
-
class AsyncmyConnectionParams(TypedDict
|
|
32
|
+
class AsyncmyConnectionParams(TypedDict):
|
|
30
33
|
"""Asyncmy connection parameters."""
|
|
31
34
|
|
|
32
35
|
host: NotRequired[str]
|
|
@@ -44,11 +47,11 @@ class AsyncmyConnectionParams(TypedDict, total=False):
|
|
|
44
47
|
ssl: NotRequired[Any]
|
|
45
48
|
sql_mode: NotRequired[str]
|
|
46
49
|
init_command: NotRequired[str]
|
|
47
|
-
cursor_class: NotRequired[
|
|
50
|
+
cursor_class: NotRequired[type["Cursor"] | type["DictCursor"]]
|
|
48
51
|
extra: NotRequired[dict[str, Any]]
|
|
49
52
|
|
|
50
53
|
|
|
51
|
-
class AsyncmyPoolParams(AsyncmyConnectionParams
|
|
54
|
+
class AsyncmyPoolParams(AsyncmyConnectionParams):
|
|
52
55
|
"""Asyncmy pool parameters."""
|
|
53
56
|
|
|
54
57
|
minsize: NotRequired[int]
|
|
@@ -57,21 +60,41 @@ class AsyncmyPoolParams(AsyncmyConnectionParams, total=False):
|
|
|
57
60
|
pool_recycle: NotRequired[int]
|
|
58
61
|
|
|
59
62
|
|
|
63
|
+
class AsyncmyDriverFeatures(TypedDict):
|
|
64
|
+
"""Asyncmy driver feature flags.
|
|
65
|
+
|
|
66
|
+
MySQL/MariaDB handle JSON natively, but custom serializers can be provided
|
|
67
|
+
for specialized use cases (e.g., orjson for performance, msgspec for type safety).
|
|
68
|
+
|
|
69
|
+
json_serializer: Custom JSON serializer function.
|
|
70
|
+
Defaults to sqlspec.utils.serializers.to_json.
|
|
71
|
+
Use for performance (orjson) or custom encoding.
|
|
72
|
+
json_deserializer: Custom JSON deserializer function.
|
|
73
|
+
Defaults to sqlspec.utils.serializers.from_json.
|
|
74
|
+
Use for performance (orjson) or custom decoding.
|
|
75
|
+
"""
|
|
76
|
+
|
|
77
|
+
json_serializer: NotRequired["Callable[[Any], str]"]
|
|
78
|
+
json_deserializer: NotRequired["Callable[[str], Any]"]
|
|
79
|
+
|
|
80
|
+
|
|
60
81
|
class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", AsyncmyDriver]): # pyright: ignore
|
|
61
82
|
"""Configuration for Asyncmy database connections."""
|
|
62
83
|
|
|
63
84
|
driver_type: ClassVar[type[AsyncmyDriver]] = AsyncmyDriver
|
|
64
85
|
connection_type: "ClassVar[type[AsyncmyConnection]]" = AsyncmyConnection # pyright: ignore
|
|
86
|
+
supports_transactional_ddl: ClassVar[bool] = False
|
|
65
87
|
|
|
66
88
|
def __init__(
|
|
67
89
|
self,
|
|
68
90
|
*,
|
|
69
|
-
pool_config: "
|
|
70
|
-
pool_instance: "
|
|
71
|
-
migration_config:
|
|
72
|
-
statement_config: "
|
|
73
|
-
driver_features: "
|
|
74
|
-
bind_key: "
|
|
91
|
+
pool_config: "AsyncmyPoolParams | dict[str, Any] | None" = None,
|
|
92
|
+
pool_instance: "AsyncmyPool | None" = None,
|
|
93
|
+
migration_config: dict[str, Any] | None = None,
|
|
94
|
+
statement_config: "StatementConfig | None" = None,
|
|
95
|
+
driver_features: "AsyncmyDriverFeatures | dict[str, Any] | None" = None,
|
|
96
|
+
bind_key: "str | None" = None,
|
|
97
|
+
extension_config: "dict[str, dict[str, Any]] | None" = None,
|
|
75
98
|
) -> None:
|
|
76
99
|
"""Initialize Asyncmy configuration.
|
|
77
100
|
|
|
@@ -80,8 +103,9 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
|
|
|
80
103
|
pool_instance: Existing pool instance to use
|
|
81
104
|
migration_config: Migration configuration
|
|
82
105
|
statement_config: Statement configuration override
|
|
83
|
-
driver_features: Driver feature configuration
|
|
106
|
+
driver_features: Driver feature configuration (TypedDict or dict)
|
|
84
107
|
bind_key: Optional unique identifier for this configuration
|
|
108
|
+
extension_config: Extension-specific configuration (e.g., Litestar plugin settings)
|
|
85
109
|
"""
|
|
86
110
|
processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
|
|
87
111
|
if "extra" in processed_pool_config:
|
|
@@ -96,17 +120,33 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
|
|
|
96
120
|
if statement_config is None:
|
|
97
121
|
statement_config = asyncmy_statement_config
|
|
98
122
|
|
|
123
|
+
processed_driver_features: dict[str, Any] = dict(driver_features) if driver_features else {}
|
|
124
|
+
|
|
125
|
+
if "json_serializer" not in processed_driver_features:
|
|
126
|
+
processed_driver_features["json_serializer"] = to_json
|
|
127
|
+
if "json_deserializer" not in processed_driver_features:
|
|
128
|
+
processed_driver_features["json_deserializer"] = from_json
|
|
129
|
+
|
|
99
130
|
super().__init__(
|
|
100
131
|
pool_config=processed_pool_config,
|
|
101
132
|
pool_instance=pool_instance,
|
|
102
133
|
migration_config=migration_config,
|
|
103
134
|
statement_config=statement_config,
|
|
104
|
-
driver_features=
|
|
135
|
+
driver_features=processed_driver_features,
|
|
105
136
|
bind_key=bind_key,
|
|
137
|
+
extension_config=extension_config,
|
|
106
138
|
)
|
|
107
139
|
|
|
108
140
|
async def _create_pool(self) -> "AsyncmyPool": # pyright: ignore
|
|
109
|
-
"""Create the actual async connection pool.
|
|
141
|
+
"""Create the actual async connection pool.
|
|
142
|
+
|
|
143
|
+
MySQL/MariaDB handle JSON types natively without requiring connection-level
|
|
144
|
+
type handlers. JSON serialization is handled via type_coercion_map in the
|
|
145
|
+
driver's statement_config (see driver.py).
|
|
146
|
+
|
|
147
|
+
Future driver_features can be added here if needed (e.g., custom connection
|
|
148
|
+
initialization, specialized type handling).
|
|
149
|
+
"""
|
|
110
150
|
return await asyncmy.create_pool(**dict(self.pool_config)) # pyright: ignore
|
|
111
151
|
|
|
112
152
|
async def _close_pool(self) -> None:
|
|
@@ -146,7 +186,7 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
|
|
|
146
186
|
|
|
147
187
|
@asynccontextmanager
|
|
148
188
|
async def provide_session(
|
|
149
|
-
self, *args: Any, statement_config: "
|
|
189
|
+
self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
|
|
150
190
|
) -> AsyncGenerator[AsyncmyDriver, None]:
|
|
151
191
|
"""Provide an async driver session context manager.
|
|
152
192
|
|
|
@@ -159,8 +199,10 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
|
|
|
159
199
|
An AsyncmyDriver instance.
|
|
160
200
|
"""
|
|
161
201
|
async with self.provide_connection(*args, **kwargs) as connection:
|
|
162
|
-
final_statement_config = statement_config or asyncmy_statement_config
|
|
163
|
-
yield self.driver_type(
|
|
202
|
+
final_statement_config = statement_config or self.statement_config or asyncmy_statement_config
|
|
203
|
+
yield self.driver_type(
|
|
204
|
+
connection=connection, statement_config=final_statement_config, driver_features=self.driver_features
|
|
205
|
+
)
|
|
164
206
|
|
|
165
207
|
async def provide_pool(self, *args: Any, **kwargs: Any) -> "Pool": # pyright: ignore
|
|
166
208
|
"""Provide async pool instance.
|