sqlspec 0.26.0__py3-none-any.whl → 0.28.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.

Files changed (212) hide show
  1. sqlspec/__init__.py +7 -15
  2. sqlspec/_serialization.py +55 -25
  3. sqlspec/_typing.py +155 -52
  4. sqlspec/adapters/adbc/_types.py +1 -1
  5. sqlspec/adapters/adbc/adk/__init__.py +5 -0
  6. sqlspec/adapters/adbc/adk/store.py +880 -0
  7. sqlspec/adapters/adbc/config.py +62 -12
  8. sqlspec/adapters/adbc/data_dictionary.py +74 -2
  9. sqlspec/adapters/adbc/driver.py +226 -58
  10. sqlspec/adapters/adbc/litestar/__init__.py +5 -0
  11. sqlspec/adapters/adbc/litestar/store.py +504 -0
  12. sqlspec/adapters/adbc/type_converter.py +44 -50
  13. sqlspec/adapters/aiosqlite/_types.py +1 -1
  14. sqlspec/adapters/aiosqlite/adk/__init__.py +5 -0
  15. sqlspec/adapters/aiosqlite/adk/store.py +536 -0
  16. sqlspec/adapters/aiosqlite/config.py +86 -16
  17. sqlspec/adapters/aiosqlite/data_dictionary.py +34 -2
  18. sqlspec/adapters/aiosqlite/driver.py +127 -38
  19. sqlspec/adapters/aiosqlite/litestar/__init__.py +5 -0
  20. sqlspec/adapters/aiosqlite/litestar/store.py +281 -0
  21. sqlspec/adapters/aiosqlite/pool.py +7 -7
  22. sqlspec/adapters/asyncmy/__init__.py +7 -1
  23. sqlspec/adapters/asyncmy/_types.py +1 -1
  24. sqlspec/adapters/asyncmy/adk/__init__.py +5 -0
  25. sqlspec/adapters/asyncmy/adk/store.py +503 -0
  26. sqlspec/adapters/asyncmy/config.py +59 -17
  27. sqlspec/adapters/asyncmy/data_dictionary.py +41 -2
  28. sqlspec/adapters/asyncmy/driver.py +293 -62
  29. sqlspec/adapters/asyncmy/litestar/__init__.py +5 -0
  30. sqlspec/adapters/asyncmy/litestar/store.py +296 -0
  31. sqlspec/adapters/asyncpg/__init__.py +2 -1
  32. sqlspec/adapters/asyncpg/_type_handlers.py +71 -0
  33. sqlspec/adapters/asyncpg/_types.py +11 -7
  34. sqlspec/adapters/asyncpg/adk/__init__.py +5 -0
  35. sqlspec/adapters/asyncpg/adk/store.py +460 -0
  36. sqlspec/adapters/asyncpg/config.py +57 -36
  37. sqlspec/adapters/asyncpg/data_dictionary.py +48 -2
  38. sqlspec/adapters/asyncpg/driver.py +153 -23
  39. sqlspec/adapters/asyncpg/litestar/__init__.py +5 -0
  40. sqlspec/adapters/asyncpg/litestar/store.py +253 -0
  41. sqlspec/adapters/bigquery/_types.py +1 -1
  42. sqlspec/adapters/bigquery/adk/__init__.py +5 -0
  43. sqlspec/adapters/bigquery/adk/store.py +585 -0
  44. sqlspec/adapters/bigquery/config.py +36 -11
  45. sqlspec/adapters/bigquery/data_dictionary.py +42 -2
  46. sqlspec/adapters/bigquery/driver.py +489 -144
  47. sqlspec/adapters/bigquery/litestar/__init__.py +5 -0
  48. sqlspec/adapters/bigquery/litestar/store.py +327 -0
  49. sqlspec/adapters/bigquery/type_converter.py +55 -23
  50. sqlspec/adapters/duckdb/_types.py +2 -2
  51. sqlspec/adapters/duckdb/adk/__init__.py +14 -0
  52. sqlspec/adapters/duckdb/adk/store.py +563 -0
  53. sqlspec/adapters/duckdb/config.py +79 -21
  54. sqlspec/adapters/duckdb/data_dictionary.py +41 -2
  55. sqlspec/adapters/duckdb/driver.py +225 -44
  56. sqlspec/adapters/duckdb/litestar/__init__.py +5 -0
  57. sqlspec/adapters/duckdb/litestar/store.py +332 -0
  58. sqlspec/adapters/duckdb/pool.py +5 -5
  59. sqlspec/adapters/duckdb/type_converter.py +51 -21
  60. sqlspec/adapters/oracledb/_numpy_handlers.py +133 -0
  61. sqlspec/adapters/oracledb/_types.py +20 -2
  62. sqlspec/adapters/oracledb/adk/__init__.py +5 -0
  63. sqlspec/adapters/oracledb/adk/store.py +1628 -0
  64. sqlspec/adapters/oracledb/config.py +120 -36
  65. sqlspec/adapters/oracledb/data_dictionary.py +87 -20
  66. sqlspec/adapters/oracledb/driver.py +475 -86
  67. sqlspec/adapters/oracledb/litestar/__init__.py +5 -0
  68. sqlspec/adapters/oracledb/litestar/store.py +765 -0
  69. sqlspec/adapters/oracledb/migrations.py +316 -25
  70. sqlspec/adapters/oracledb/type_converter.py +91 -16
  71. sqlspec/adapters/psqlpy/_type_handlers.py +44 -0
  72. sqlspec/adapters/psqlpy/_types.py +2 -1
  73. sqlspec/adapters/psqlpy/adk/__init__.py +5 -0
  74. sqlspec/adapters/psqlpy/adk/store.py +483 -0
  75. sqlspec/adapters/psqlpy/config.py +45 -19
  76. sqlspec/adapters/psqlpy/data_dictionary.py +48 -2
  77. sqlspec/adapters/psqlpy/driver.py +108 -41
  78. sqlspec/adapters/psqlpy/litestar/__init__.py +5 -0
  79. sqlspec/adapters/psqlpy/litestar/store.py +272 -0
  80. sqlspec/adapters/psqlpy/type_converter.py +40 -11
  81. sqlspec/adapters/psycopg/_type_handlers.py +80 -0
  82. sqlspec/adapters/psycopg/_types.py +2 -1
  83. sqlspec/adapters/psycopg/adk/__init__.py +5 -0
  84. sqlspec/adapters/psycopg/adk/store.py +962 -0
  85. sqlspec/adapters/psycopg/config.py +65 -37
  86. sqlspec/adapters/psycopg/data_dictionary.py +91 -3
  87. sqlspec/adapters/psycopg/driver.py +200 -78
  88. sqlspec/adapters/psycopg/litestar/__init__.py +5 -0
  89. sqlspec/adapters/psycopg/litestar/store.py +554 -0
  90. sqlspec/adapters/sqlite/__init__.py +2 -1
  91. sqlspec/adapters/sqlite/_type_handlers.py +86 -0
  92. sqlspec/adapters/sqlite/_types.py +1 -1
  93. sqlspec/adapters/sqlite/adk/__init__.py +5 -0
  94. sqlspec/adapters/sqlite/adk/store.py +582 -0
  95. sqlspec/adapters/sqlite/config.py +85 -16
  96. sqlspec/adapters/sqlite/data_dictionary.py +34 -2
  97. sqlspec/adapters/sqlite/driver.py +120 -52
  98. sqlspec/adapters/sqlite/litestar/__init__.py +5 -0
  99. sqlspec/adapters/sqlite/litestar/store.py +318 -0
  100. sqlspec/adapters/sqlite/pool.py +5 -5
  101. sqlspec/base.py +45 -26
  102. sqlspec/builder/__init__.py +73 -4
  103. sqlspec/builder/_base.py +91 -58
  104. sqlspec/builder/_column.py +5 -5
  105. sqlspec/builder/_ddl.py +98 -89
  106. sqlspec/builder/_delete.py +5 -4
  107. sqlspec/builder/_dml.py +388 -0
  108. sqlspec/{_sql.py → builder/_factory.py} +41 -44
  109. sqlspec/builder/_insert.py +5 -82
  110. sqlspec/builder/{mixins/_join_operations.py → _join.py} +145 -143
  111. sqlspec/builder/_merge.py +446 -11
  112. sqlspec/builder/_parsing_utils.py +9 -11
  113. sqlspec/builder/_select.py +1313 -25
  114. sqlspec/builder/_update.py +11 -42
  115. sqlspec/cli.py +76 -69
  116. sqlspec/config.py +331 -62
  117. sqlspec/core/__init__.py +5 -4
  118. sqlspec/core/cache.py +18 -18
  119. sqlspec/core/compiler.py +6 -8
  120. sqlspec/core/filters.py +55 -47
  121. sqlspec/core/hashing.py +9 -9
  122. sqlspec/core/parameters.py +76 -45
  123. sqlspec/core/result.py +234 -47
  124. sqlspec/core/splitter.py +16 -17
  125. sqlspec/core/statement.py +32 -31
  126. sqlspec/core/type_conversion.py +3 -2
  127. sqlspec/driver/__init__.py +1 -3
  128. sqlspec/driver/_async.py +183 -160
  129. sqlspec/driver/_common.py +197 -109
  130. sqlspec/driver/_sync.py +189 -161
  131. sqlspec/driver/mixins/_result_tools.py +20 -236
  132. sqlspec/driver/mixins/_sql_translator.py +4 -4
  133. sqlspec/exceptions.py +70 -7
  134. sqlspec/extensions/adk/__init__.py +53 -0
  135. sqlspec/extensions/adk/_types.py +51 -0
  136. sqlspec/extensions/adk/converters.py +172 -0
  137. sqlspec/extensions/adk/migrations/0001_create_adk_tables.py +144 -0
  138. sqlspec/extensions/adk/migrations/__init__.py +0 -0
  139. sqlspec/extensions/adk/service.py +181 -0
  140. sqlspec/extensions/adk/store.py +536 -0
  141. sqlspec/extensions/aiosql/adapter.py +69 -61
  142. sqlspec/extensions/fastapi/__init__.py +21 -0
  143. sqlspec/extensions/fastapi/extension.py +331 -0
  144. sqlspec/extensions/fastapi/providers.py +543 -0
  145. sqlspec/extensions/flask/__init__.py +36 -0
  146. sqlspec/extensions/flask/_state.py +71 -0
  147. sqlspec/extensions/flask/_utils.py +40 -0
  148. sqlspec/extensions/flask/extension.py +389 -0
  149. sqlspec/extensions/litestar/__init__.py +21 -4
  150. sqlspec/extensions/litestar/cli.py +54 -10
  151. sqlspec/extensions/litestar/config.py +56 -266
  152. sqlspec/extensions/litestar/handlers.py +46 -17
  153. sqlspec/extensions/litestar/migrations/0001_create_session_table.py +137 -0
  154. sqlspec/extensions/litestar/migrations/__init__.py +3 -0
  155. sqlspec/extensions/litestar/plugin.py +349 -224
  156. sqlspec/extensions/litestar/providers.py +25 -25
  157. sqlspec/extensions/litestar/store.py +265 -0
  158. sqlspec/extensions/starlette/__init__.py +10 -0
  159. sqlspec/extensions/starlette/_state.py +25 -0
  160. sqlspec/extensions/starlette/_utils.py +52 -0
  161. sqlspec/extensions/starlette/extension.py +254 -0
  162. sqlspec/extensions/starlette/middleware.py +154 -0
  163. sqlspec/loader.py +30 -49
  164. sqlspec/migrations/base.py +200 -76
  165. sqlspec/migrations/commands.py +591 -62
  166. sqlspec/migrations/context.py +6 -9
  167. sqlspec/migrations/fix.py +199 -0
  168. sqlspec/migrations/loaders.py +47 -19
  169. sqlspec/migrations/runner.py +241 -75
  170. sqlspec/migrations/tracker.py +237 -21
  171. sqlspec/migrations/utils.py +51 -3
  172. sqlspec/migrations/validation.py +177 -0
  173. sqlspec/protocols.py +106 -36
  174. sqlspec/storage/_utils.py +85 -0
  175. sqlspec/storage/backends/fsspec.py +133 -107
  176. sqlspec/storage/backends/local.py +78 -51
  177. sqlspec/storage/backends/obstore.py +276 -168
  178. sqlspec/storage/registry.py +75 -39
  179. sqlspec/typing.py +30 -84
  180. sqlspec/utils/__init__.py +25 -4
  181. sqlspec/utils/arrow_helpers.py +81 -0
  182. sqlspec/utils/config_resolver.py +6 -6
  183. sqlspec/utils/correlation.py +4 -5
  184. sqlspec/utils/data_transformation.py +3 -2
  185. sqlspec/utils/deprecation.py +9 -8
  186. sqlspec/utils/fixtures.py +4 -4
  187. sqlspec/utils/logging.py +46 -6
  188. sqlspec/utils/module_loader.py +205 -5
  189. sqlspec/utils/portal.py +311 -0
  190. sqlspec/utils/schema.py +288 -0
  191. sqlspec/utils/serializers.py +113 -4
  192. sqlspec/utils/sync_tools.py +36 -22
  193. sqlspec/utils/text.py +1 -2
  194. sqlspec/utils/type_guards.py +136 -20
  195. sqlspec/utils/version.py +433 -0
  196. {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/METADATA +41 -22
  197. sqlspec-0.28.0.dist-info/RECORD +221 -0
  198. sqlspec/builder/mixins/__init__.py +0 -55
  199. sqlspec/builder/mixins/_cte_and_set_ops.py +0 -253
  200. sqlspec/builder/mixins/_delete_operations.py +0 -50
  201. sqlspec/builder/mixins/_insert_operations.py +0 -282
  202. sqlspec/builder/mixins/_merge_operations.py +0 -698
  203. sqlspec/builder/mixins/_order_limit_operations.py +0 -145
  204. sqlspec/builder/mixins/_pivot_operations.py +0 -157
  205. sqlspec/builder/mixins/_select_operations.py +0 -930
  206. sqlspec/builder/mixins/_update_operations.py +0 -199
  207. sqlspec/builder/mixins/_where_clause.py +0 -1298
  208. sqlspec-0.26.0.dist-info/RECORD +0 -157
  209. sqlspec-0.26.0.dist-info/licenses/NOTICE +0 -29
  210. {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/WHEEL +0 -0
  211. {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/entry_points.txt +0 -0
  212. {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,880 @@
1
+ """ADBC ADK store for Google Agent Development Kit session/event storage."""
2
+
3
+ from typing import TYPE_CHECKING, Any, Final
4
+
5
+ from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord
6
+ from sqlspec.utils.logging import get_logger
7
+ from sqlspec.utils.serializers import from_json, to_json
8
+
9
+ if TYPE_CHECKING:
10
+ from sqlspec.adapters.adbc.config import AdbcConfig
11
+
12
+ logger = get_logger("adapters.adbc.adk.store")
13
+
14
+ __all__ = ("AdbcADKStore",)
15
+
16
+ DIALECT_POSTGRESQL: Final = "postgresql"
17
+ DIALECT_SQLITE: Final = "sqlite"
18
+ DIALECT_DUCKDB: Final = "duckdb"
19
+ DIALECT_SNOWFLAKE: Final = "snowflake"
20
+ DIALECT_GENERIC: Final = "generic"
21
+
22
+ ADBC_TABLE_NOT_FOUND_PATTERNS: Final = ("no such table", "table or view does not exist", "relation does not exist")
23
+
24
+
25
+ class AdbcADKStore(BaseSyncADKStore["AdbcConfig"]):
26
+ """ADBC synchronous ADK store for Arrow Database Connectivity.
27
+
28
+ Implements session and event storage for Google Agent Development Kit
29
+ using ADBC. ADBC provides a vendor-neutral API with Arrow-native data
30
+ transfer across multiple databases (PostgreSQL, SQLite, DuckDB, etc.).
31
+
32
+ Provides:
33
+ - Session state management with JSON serialization (TEXT storage)
34
+ - Event history tracking with BLOB-serialized actions
35
+ - Timezone-aware timestamps
36
+ - Foreign key constraints with cascade delete
37
+ - Database-agnostic SQL (supports multiple backends)
38
+
39
+ Args:
40
+ config: AdbcConfig with extension_config["adk"] settings.
41
+
42
+ Example:
43
+ from sqlspec.adapters.adbc import AdbcConfig
44
+ from sqlspec.adapters.adbc.adk import AdbcADKStore
45
+
46
+ config = AdbcConfig(
47
+ connection_config={"driver_name": "sqlite", "uri": ":memory:"},
48
+ extension_config={
49
+ "adk": {
50
+ "session_table": "my_sessions",
51
+ "events_table": "my_events",
52
+ "owner_id_column": "tenant_id INTEGER REFERENCES tenants(id)"
53
+ }
54
+ }
55
+ )
56
+ store = AdbcADKStore(config)
57
+ store.create_tables()
58
+
59
+ Notes:
60
+ - TEXT for JSON storage (compatible across all ADBC backends)
61
+ - BLOB for pre-serialized actions from Google ADK
62
+ - TIMESTAMP for timezone-aware timestamps (driver-dependent precision)
63
+ - INTEGER for booleans (0/1/NULL)
64
+ - Parameter style varies by backend (?, $1, :name, etc.)
65
+ - Uses dialect-agnostic SQL for maximum compatibility
66
+ - State and JSON fields use to_json/from_json for serialization
67
+ - ADBC drivers handle parameter binding automatically
68
+ - Configuration is read from config.extension_config["adk"]
69
+ """
70
+
71
+ __slots__ = ("_dialect",)
72
+
73
+ def __init__(self, config: "AdbcConfig") -> None:
74
+ """Initialize ADBC ADK store.
75
+
76
+ Args:
77
+ config: AdbcConfig instance (any ADBC driver).
78
+
79
+ Notes:
80
+ Configuration is read from config.extension_config["adk"]:
81
+ - session_table: Sessions table name (default: "adk_sessions")
82
+ - events_table: Events table name (default: "adk_events")
83
+ - owner_id_column: Optional owner FK column DDL (default: None)
84
+ """
85
+ super().__init__(config)
86
+ self._dialect = self._detect_dialect()
87
+
88
+ @property
89
+ def dialect(self) -> str:
90
+ """Return the detected database dialect."""
91
+ return self._dialect
92
+
93
+ def _detect_dialect(self) -> str:
94
+ """Detect ADBC driver dialect from connection config.
95
+
96
+ Returns:
97
+ Dialect identifier for DDL generation.
98
+
99
+ Notes:
100
+ Reads from config.connection_config driver_name.
101
+ Falls back to generic for unknown drivers.
102
+ """
103
+ driver_name = self._config.connection_config.get("driver_name", "").lower()
104
+
105
+ if "postgres" in driver_name:
106
+ return DIALECT_POSTGRESQL
107
+ if "sqlite" in driver_name:
108
+ return DIALECT_SQLITE
109
+ if "duckdb" in driver_name:
110
+ return DIALECT_DUCKDB
111
+ if "snowflake" in driver_name:
112
+ return DIALECT_SNOWFLAKE
113
+
114
+ logger.warning(
115
+ "Unknown ADBC driver: %s. Using generic SQL dialect. "
116
+ "Consider using a direct adapter for better performance.",
117
+ driver_name,
118
+ )
119
+ return DIALECT_GENERIC
120
+
121
+ def _serialize_state(self, state: "dict[str, Any]") -> str:
122
+ """Serialize state dictionary to JSON string.
123
+
124
+ Args:
125
+ state: State dictionary to serialize.
126
+
127
+ Returns:
128
+ JSON string.
129
+ """
130
+ return to_json(state)
131
+
132
+ def _deserialize_state(self, data: Any) -> "dict[str, Any]":
133
+ """Deserialize state data from JSON string.
134
+
135
+ Args:
136
+ data: JSON string from database.
137
+
138
+ Returns:
139
+ Deserialized state dictionary.
140
+ """
141
+ if data is None:
142
+ return {}
143
+ return from_json(str(data)) # type: ignore[no-any-return]
144
+
145
+ def _serialize_json_field(self, value: Any) -> "str | None":
146
+ """Serialize optional JSON field for event storage.
147
+
148
+ Args:
149
+ value: Value to serialize (dict or None).
150
+
151
+ Returns:
152
+ Serialized JSON string or None.
153
+ """
154
+ if value is None:
155
+ return None
156
+ return to_json(value)
157
+
158
+ def _deserialize_json_field(self, data: Any) -> "dict[str, Any] | None":
159
+ """Deserialize optional JSON field from database.
160
+
161
+ Args:
162
+ data: JSON string from database or None.
163
+
164
+ Returns:
165
+ Deserialized dictionary or None.
166
+ """
167
+ if data is None:
168
+ return None
169
+ return from_json(str(data)) # type: ignore[no-any-return]
170
+
171
+ def _get_create_sessions_table_sql(self) -> str:
172
+ """Get CREATE TABLE SQL for sessions with dialect dispatch.
173
+
174
+ Returns:
175
+ SQL statement to create adk_sessions table.
176
+ """
177
+ if self._dialect == DIALECT_POSTGRESQL:
178
+ return self._get_sessions_ddl_postgresql()
179
+ if self._dialect == DIALECT_SQLITE:
180
+ return self._get_sessions_ddl_sqlite()
181
+ if self._dialect == DIALECT_DUCKDB:
182
+ return self._get_sessions_ddl_duckdb()
183
+ if self._dialect == DIALECT_SNOWFLAKE:
184
+ return self._get_sessions_ddl_snowflake()
185
+ return self._get_sessions_ddl_generic()
186
+
187
+ def _get_sessions_ddl_postgresql(self) -> str:
188
+ """PostgreSQL DDL with JSONB and TIMESTAMPTZ.
189
+
190
+ Returns:
191
+ SQL to create sessions table optimized for PostgreSQL.
192
+ """
193
+ owner_id_ddl = f", {self._owner_id_column_ddl}" if self._owner_id_column_ddl else ""
194
+ return f"""
195
+ CREATE TABLE IF NOT EXISTS {self._session_table} (
196
+ id VARCHAR(128) PRIMARY KEY,
197
+ app_name VARCHAR(128) NOT NULL,
198
+ user_id VARCHAR(128) NOT NULL{owner_id_ddl},
199
+ state JSONB NOT NULL DEFAULT '{{}}'::jsonb,
200
+ create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
201
+ update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
202
+ )
203
+ """
204
+
205
+ def _get_sessions_ddl_sqlite(self) -> str:
206
+ """SQLite DDL with TEXT and REAL timestamps.
207
+
208
+ Returns:
209
+ SQL to create sessions table optimized for SQLite.
210
+ """
211
+ owner_id_ddl = f", {self._owner_id_column_ddl}" if self._owner_id_column_ddl else ""
212
+ return f"""
213
+ CREATE TABLE IF NOT EXISTS {self._session_table} (
214
+ id TEXT PRIMARY KEY,
215
+ app_name TEXT NOT NULL,
216
+ user_id TEXT NOT NULL{owner_id_ddl},
217
+ state TEXT NOT NULL DEFAULT '{{}}',
218
+ create_time REAL NOT NULL,
219
+ update_time REAL NOT NULL
220
+ )
221
+ """
222
+
223
+ def _get_sessions_ddl_duckdb(self) -> str:
224
+ """DuckDB DDL with native JSON type.
225
+
226
+ Returns:
227
+ SQL to create sessions table optimized for DuckDB.
228
+ """
229
+ owner_id_ddl = f", {self._owner_id_column_ddl}" if self._owner_id_column_ddl else ""
230
+ return f"""
231
+ CREATE TABLE IF NOT EXISTS {self._session_table} (
232
+ id VARCHAR(128) PRIMARY KEY,
233
+ app_name VARCHAR(128) NOT NULL,
234
+ user_id VARCHAR(128) NOT NULL{owner_id_ddl},
235
+ state JSON NOT NULL,
236
+ create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
237
+ update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
238
+ )
239
+ """
240
+
241
+ def _get_sessions_ddl_snowflake(self) -> str:
242
+ """Snowflake DDL with VARIANT type.
243
+
244
+ Returns:
245
+ SQL to create sessions table optimized for Snowflake.
246
+ """
247
+ owner_id_ddl = f", {self._owner_id_column_ddl}" if self._owner_id_column_ddl else ""
248
+ return f"""
249
+ CREATE TABLE IF NOT EXISTS {self._session_table} (
250
+ id VARCHAR PRIMARY KEY,
251
+ app_name VARCHAR NOT NULL,
252
+ user_id VARCHAR NOT NULL{owner_id_ddl},
253
+ state VARIANT NOT NULL,
254
+ create_time TIMESTAMP_TZ NOT NULL DEFAULT CURRENT_TIMESTAMP(),
255
+ update_time TIMESTAMP_TZ NOT NULL DEFAULT CURRENT_TIMESTAMP()
256
+ )
257
+ """
258
+
259
+ def _get_sessions_ddl_generic(self) -> str:
260
+ """Generic SQL-92 compatible DDL fallback.
261
+
262
+ Returns:
263
+ SQL to create sessions table using generic types.
264
+ """
265
+ owner_id_ddl = f", {self._owner_id_column_ddl}" if self._owner_id_column_ddl else ""
266
+ return f"""
267
+ CREATE TABLE IF NOT EXISTS {self._session_table} (
268
+ id VARCHAR(128) PRIMARY KEY,
269
+ app_name VARCHAR(128) NOT NULL,
270
+ user_id VARCHAR(128) NOT NULL{owner_id_ddl},
271
+ state TEXT NOT NULL DEFAULT '{{}}',
272
+ create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
273
+ update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
274
+ )
275
+ """
276
+
277
+ def _get_create_events_table_sql(self) -> str:
278
+ """Get CREATE TABLE SQL for events with dialect dispatch.
279
+
280
+ Returns:
281
+ SQL statement to create adk_events table.
282
+ """
283
+ if self._dialect == DIALECT_POSTGRESQL:
284
+ return self._get_events_ddl_postgresql()
285
+ if self._dialect == DIALECT_SQLITE:
286
+ return self._get_events_ddl_sqlite()
287
+ if self._dialect == DIALECT_DUCKDB:
288
+ return self._get_events_ddl_duckdb()
289
+ if self._dialect == DIALECT_SNOWFLAKE:
290
+ return self._get_events_ddl_snowflake()
291
+ return self._get_events_ddl_generic()
292
+
293
+ def _get_events_ddl_postgresql(self) -> str:
294
+ """PostgreSQL DDL for events table.
295
+
296
+ Returns:
297
+ SQL to create events table optimized for PostgreSQL.
298
+ """
299
+ return f"""
300
+ CREATE TABLE IF NOT EXISTS {self._events_table} (
301
+ id VARCHAR(128) PRIMARY KEY,
302
+ session_id VARCHAR(128) NOT NULL,
303
+ app_name VARCHAR(128) NOT NULL,
304
+ user_id VARCHAR(128) NOT NULL,
305
+ invocation_id VARCHAR(256),
306
+ author VARCHAR(256),
307
+ actions BYTEA,
308
+ long_running_tool_ids_json TEXT,
309
+ branch VARCHAR(256),
310
+ timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
311
+ content JSONB,
312
+ grounding_metadata JSONB,
313
+ custom_metadata JSONB,
314
+ partial BOOLEAN,
315
+ turn_complete BOOLEAN,
316
+ interrupted BOOLEAN,
317
+ error_code VARCHAR(256),
318
+ error_message VARCHAR(1024),
319
+ FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
320
+ )
321
+ """
322
+
323
+ def _get_events_ddl_sqlite(self) -> str:
324
+ """SQLite DDL for events table.
325
+
326
+ Returns:
327
+ SQL to create events table optimized for SQLite.
328
+ """
329
+ return f"""
330
+ CREATE TABLE IF NOT EXISTS {self._events_table} (
331
+ id TEXT PRIMARY KEY,
332
+ session_id TEXT NOT NULL,
333
+ app_name TEXT NOT NULL,
334
+ user_id TEXT NOT NULL,
335
+ invocation_id TEXT,
336
+ author TEXT,
337
+ actions BLOB,
338
+ long_running_tool_ids_json TEXT,
339
+ branch TEXT,
340
+ timestamp REAL NOT NULL,
341
+ content TEXT,
342
+ grounding_metadata TEXT,
343
+ custom_metadata TEXT,
344
+ partial INTEGER,
345
+ turn_complete INTEGER,
346
+ interrupted INTEGER,
347
+ error_code TEXT,
348
+ error_message TEXT,
349
+ FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
350
+ )
351
+ """
352
+
353
+ def _get_events_ddl_duckdb(self) -> str:
354
+ """DuckDB DDL for events table.
355
+
356
+ Returns:
357
+ SQL to create events table optimized for DuckDB.
358
+ """
359
+ return f"""
360
+ CREATE TABLE IF NOT EXISTS {self._events_table} (
361
+ id VARCHAR(128) PRIMARY KEY,
362
+ session_id VARCHAR(128) NOT NULL,
363
+ app_name VARCHAR(128) NOT NULL,
364
+ user_id VARCHAR(128) NOT NULL,
365
+ invocation_id VARCHAR(256),
366
+ author VARCHAR(256),
367
+ actions BLOB,
368
+ long_running_tool_ids_json VARCHAR,
369
+ branch VARCHAR(256),
370
+ timestamp TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
371
+ content JSON,
372
+ grounding_metadata JSON,
373
+ custom_metadata JSON,
374
+ partial BOOLEAN,
375
+ turn_complete BOOLEAN,
376
+ interrupted BOOLEAN,
377
+ error_code VARCHAR(256),
378
+ error_message VARCHAR(1024),
379
+ FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
380
+ )
381
+ """
382
+
383
+ def _get_events_ddl_snowflake(self) -> str:
384
+ """Snowflake DDL for events table.
385
+
386
+ Returns:
387
+ SQL to create events table optimized for Snowflake.
388
+ """
389
+ return f"""
390
+ CREATE TABLE IF NOT EXISTS {self._events_table} (
391
+ id VARCHAR PRIMARY KEY,
392
+ session_id VARCHAR NOT NULL,
393
+ app_name VARCHAR NOT NULL,
394
+ user_id VARCHAR NOT NULL,
395
+ invocation_id VARCHAR,
396
+ author VARCHAR,
397
+ actions BINARY,
398
+ long_running_tool_ids_json VARCHAR,
399
+ branch VARCHAR,
400
+ timestamp TIMESTAMP_TZ NOT NULL DEFAULT CURRENT_TIMESTAMP(),
401
+ content VARIANT,
402
+ grounding_metadata VARIANT,
403
+ custom_metadata VARIANT,
404
+ partial BOOLEAN,
405
+ turn_complete BOOLEAN,
406
+ interrupted BOOLEAN,
407
+ error_code VARCHAR,
408
+ error_message VARCHAR,
409
+ FOREIGN KEY (session_id) REFERENCES {self._session_table}(id)
410
+ )
411
+ """
412
+
413
+ def _get_events_ddl_generic(self) -> str:
414
+ """Generic SQL-92 compatible DDL for events table.
415
+
416
+ Returns:
417
+ SQL to create events table using generic types.
418
+ """
419
+ return f"""
420
+ CREATE TABLE IF NOT EXISTS {self._events_table} (
421
+ id VARCHAR(128) PRIMARY KEY,
422
+ session_id VARCHAR(128) NOT NULL,
423
+ app_name VARCHAR(128) NOT NULL,
424
+ user_id VARCHAR(128) NOT NULL,
425
+ invocation_id VARCHAR(256),
426
+ author VARCHAR(256),
427
+ actions BLOB,
428
+ long_running_tool_ids_json TEXT,
429
+ branch VARCHAR(256),
430
+ timestamp TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
431
+ content TEXT,
432
+ grounding_metadata TEXT,
433
+ custom_metadata TEXT,
434
+ partial INTEGER,
435
+ turn_complete INTEGER,
436
+ interrupted INTEGER,
437
+ error_code VARCHAR(256),
438
+ error_message VARCHAR(1024),
439
+ FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
440
+ )
441
+ """
442
+
443
+ def _get_drop_tables_sql(self) -> "list[str]":
444
+ """Get DROP TABLE SQL statements.
445
+
446
+ Returns:
447
+ List of SQL statements to drop tables and indexes.
448
+
449
+ Notes:
450
+ Order matters: drop events table (child) before sessions (parent).
451
+ Most databases automatically drop indexes when dropping tables.
452
+ """
453
+ return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
454
+
455
+ def create_tables(self) -> None:
456
+ """Create both sessions and events tables if they don't exist."""
457
+ with self._config.provide_connection() as conn:
458
+ cursor = conn.cursor()
459
+ try:
460
+ self._enable_foreign_keys(cursor, conn)
461
+
462
+ cursor.execute(self._get_create_sessions_table_sql())
463
+ conn.commit()
464
+
465
+ sessions_idx_app_user = (
466
+ f"CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user "
467
+ f"ON {self._session_table}(app_name, user_id)"
468
+ )
469
+ cursor.execute(sessions_idx_app_user)
470
+ conn.commit()
471
+
472
+ sessions_idx_update = (
473
+ f"CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time "
474
+ f"ON {self._session_table}(update_time DESC)"
475
+ )
476
+ cursor.execute(sessions_idx_update)
477
+ conn.commit()
478
+
479
+ cursor.execute(self._get_create_events_table_sql())
480
+ conn.commit()
481
+
482
+ events_idx = (
483
+ f"CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session "
484
+ f"ON {self._events_table}(session_id, timestamp ASC)"
485
+ )
486
+ cursor.execute(events_idx)
487
+ conn.commit()
488
+ finally:
489
+ cursor.close() # type: ignore[no-untyped-call]
490
+
491
+ logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
492
+
493
+ def _enable_foreign_keys(self, cursor: Any, conn: Any) -> None:
494
+ """Enable foreign key constraints for SQLite.
495
+
496
+ Args:
497
+ cursor: Database cursor.
498
+ conn: Database connection.
499
+
500
+ Notes:
501
+ SQLite requires PRAGMA foreign_keys = ON to be set per connection.
502
+ This is a no-op for other databases.
503
+ """
504
+ try:
505
+ cursor.execute("PRAGMA foreign_keys = ON")
506
+ conn.commit()
507
+ except Exception:
508
+ logger.debug("Foreign key enforcement not supported or already enabled")
509
+
510
+ def create_session(
511
+ self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
512
+ ) -> SessionRecord:
513
+ """Create a new session.
514
+
515
+ Args:
516
+ session_id: Unique session identifier.
517
+ app_name: Application name.
518
+ user_id: User identifier.
519
+ state: Initial session state.
520
+ owner_id: Optional owner ID value for owner_id_column (can be None for nullable columns).
521
+
522
+ Returns:
523
+ Created session record.
524
+ """
525
+ state_json = self._serialize_state(state)
526
+
527
+ params: tuple[Any, ...]
528
+ if self._owner_id_column_name:
529
+ sql = f"""
530
+ INSERT INTO {self._session_table}
531
+ (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
532
+ VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
533
+ """
534
+ params = (session_id, app_name, user_id, owner_id, state_json)
535
+ else:
536
+ sql = f"""
537
+ INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time)
538
+ VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
539
+ """
540
+ params = (session_id, app_name, user_id, state_json)
541
+
542
+ with self._config.provide_connection() as conn:
543
+ cursor = conn.cursor()
544
+ try:
545
+ cursor.execute(sql, params)
546
+ conn.commit()
547
+ finally:
548
+ cursor.close() # type: ignore[no-untyped-call]
549
+
550
+ return self.get_session(session_id) # type: ignore[return-value]
551
+
552
+ def get_session(self, session_id: str) -> "SessionRecord | None":
553
+ """Get session by ID.
554
+
555
+ Args:
556
+ session_id: Session identifier.
557
+
558
+ Returns:
559
+ Session record or None if not found.
560
+
561
+ Notes:
562
+ State is deserialized from JSON string.
563
+ """
564
+ sql = f"""
565
+ SELECT id, app_name, user_id, state, create_time, update_time
566
+ FROM {self._session_table}
567
+ WHERE id = ?
568
+ """
569
+
570
+ try:
571
+ with self._config.provide_connection() as conn:
572
+ cursor = conn.cursor()
573
+ try:
574
+ cursor.execute(sql, (session_id,))
575
+ row = cursor.fetchone()
576
+
577
+ if row is None:
578
+ return None
579
+
580
+ return SessionRecord(
581
+ id=row[0],
582
+ app_name=row[1],
583
+ user_id=row[2],
584
+ state=self._deserialize_state(row[3]),
585
+ create_time=row[4],
586
+ update_time=row[5],
587
+ )
588
+ finally:
589
+ cursor.close() # type: ignore[no-untyped-call]
590
+ except Exception as e:
591
+ error_msg = str(e).lower()
592
+ if any(pattern in error_msg for pattern in ADBC_TABLE_NOT_FOUND_PATTERNS):
593
+ return None
594
+ raise
595
+
596
+ def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
597
+ """Update session state.
598
+
599
+ Args:
600
+ session_id: Session identifier.
601
+ state: New state dictionary (replaces existing state).
602
+
603
+ Notes:
604
+ This replaces the entire state dictionary.
605
+ Updates update_time to current timestamp.
606
+ """
607
+ state_json = self._serialize_state(state)
608
+ sql = f"""
609
+ UPDATE {self._session_table}
610
+ SET state = ?, update_time = CURRENT_TIMESTAMP
611
+ WHERE id = ?
612
+ """
613
+
614
+ with self._config.provide_connection() as conn:
615
+ cursor = conn.cursor()
616
+ try:
617
+ cursor.execute(sql, (state_json, session_id))
618
+ conn.commit()
619
+ finally:
620
+ cursor.close() # type: ignore[no-untyped-call]
621
+
622
+ def delete_session(self, session_id: str) -> None:
623
+ """Delete session and all associated events (cascade).
624
+
625
+ Args:
626
+ session_id: Session identifier.
627
+
628
+ Notes:
629
+ Foreign key constraint ensures events are cascade-deleted.
630
+ """
631
+ sql = f"DELETE FROM {self._session_table} WHERE id = ?"
632
+
633
+ with self._config.provide_connection() as conn:
634
+ cursor = conn.cursor()
635
+ try:
636
+ self._enable_foreign_keys(cursor, conn)
637
+ cursor.execute(sql, (session_id,))
638
+ conn.commit()
639
+ finally:
640
+ cursor.close() # type: ignore[no-untyped-call]
641
+
642
+ def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
643
+ """List sessions for an app, optionally filtered by user.
644
+
645
+ Args:
646
+ app_name: Application name.
647
+ user_id: User identifier. If None, lists all sessions for the app.
648
+
649
+ Returns:
650
+ List of session records ordered by update_time DESC.
651
+
652
+ Notes:
653
+ Uses composite index on (app_name, user_id) when user_id is provided.
654
+ """
655
+ if user_id is None:
656
+ sql = f"""
657
+ SELECT id, app_name, user_id, state, create_time, update_time
658
+ FROM {self._session_table}
659
+ WHERE app_name = ?
660
+ ORDER BY update_time DESC
661
+ """
662
+ params: tuple[str, ...] = (app_name,)
663
+ else:
664
+ sql = f"""
665
+ SELECT id, app_name, user_id, state, create_time, update_time
666
+ FROM {self._session_table}
667
+ WHERE app_name = ? AND user_id = ?
668
+ ORDER BY update_time DESC
669
+ """
670
+ params = (app_name, user_id)
671
+
672
+ try:
673
+ with self._config.provide_connection() as conn:
674
+ cursor = conn.cursor()
675
+ try:
676
+ cursor.execute(sql, params)
677
+ rows = cursor.fetchall()
678
+
679
+ return [
680
+ SessionRecord(
681
+ id=row[0],
682
+ app_name=row[1],
683
+ user_id=row[2],
684
+ state=self._deserialize_state(row[3]),
685
+ create_time=row[4],
686
+ update_time=row[5],
687
+ )
688
+ for row in rows
689
+ ]
690
+ finally:
691
+ cursor.close() # type: ignore[no-untyped-call]
692
+ except Exception as e:
693
+ error_msg = str(e).lower()
694
+ if any(pattern in error_msg for pattern in ADBC_TABLE_NOT_FOUND_PATTERNS):
695
+ return []
696
+ raise
697
+
698
+ def create_event(
699
+ self,
700
+ event_id: str,
701
+ session_id: str,
702
+ app_name: str,
703
+ user_id: str,
704
+ author: "str | None" = None,
705
+ actions: "bytes | None" = None,
706
+ content: "dict[str, Any] | None" = None,
707
+ **kwargs: Any,
708
+ ) -> "EventRecord":
709
+ """Create a new event.
710
+
711
+ Args:
712
+ event_id: Unique event identifier.
713
+ session_id: Session identifier.
714
+ app_name: Application name.
715
+ user_id: User identifier.
716
+ author: Event author (user/assistant/system).
717
+ actions: Pickled actions object.
718
+ content: Event content (JSON).
719
+ **kwargs: Additional optional fields.
720
+
721
+ Returns:
722
+ Created event record.
723
+
724
+ Notes:
725
+ Uses CURRENT_TIMESTAMP for timestamp if not provided.
726
+ JSON fields are serialized to JSON strings.
727
+ Boolean fields are converted to INTEGER (0/1).
728
+ """
729
+ content_json = self._serialize_json_field(content)
730
+ grounding_metadata_json = self._serialize_json_field(kwargs.get("grounding_metadata"))
731
+ custom_metadata_json = self._serialize_json_field(kwargs.get("custom_metadata"))
732
+
733
+ partial_int = self._to_int_bool(kwargs.get("partial"))
734
+ turn_complete_int = self._to_int_bool(kwargs.get("turn_complete"))
735
+ interrupted_int = self._to_int_bool(kwargs.get("interrupted"))
736
+
737
+ sql = f"""
738
+ INSERT INTO {self._events_table} (
739
+ id, session_id, app_name, user_id, invocation_id, author, actions,
740
+ long_running_tool_ids_json, branch, timestamp, content,
741
+ grounding_metadata, custom_metadata, partial, turn_complete,
742
+ interrupted, error_code, error_message
743
+ ) VALUES (
744
+ ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
745
+ )
746
+ """
747
+
748
+ timestamp = kwargs.get("timestamp")
749
+ if timestamp is None:
750
+ from datetime import datetime, timezone
751
+
752
+ timestamp = datetime.now(timezone.utc)
753
+
754
+ with self._config.provide_connection() as conn:
755
+ cursor = conn.cursor()
756
+ try:
757
+ cursor.execute(
758
+ sql,
759
+ (
760
+ event_id,
761
+ session_id,
762
+ app_name,
763
+ user_id,
764
+ kwargs.get("invocation_id"),
765
+ author,
766
+ actions,
767
+ kwargs.get("long_running_tool_ids_json"),
768
+ kwargs.get("branch"),
769
+ timestamp,
770
+ content_json,
771
+ grounding_metadata_json,
772
+ custom_metadata_json,
773
+ partial_int,
774
+ turn_complete_int,
775
+ interrupted_int,
776
+ kwargs.get("error_code"),
777
+ kwargs.get("error_message"),
778
+ ),
779
+ )
780
+ conn.commit()
781
+ finally:
782
+ cursor.close() # type: ignore[no-untyped-call]
783
+
784
+ events = self.list_events(session_id)
785
+ for event in events:
786
+ if event["id"] == event_id:
787
+ return event
788
+
789
+ msg = f"Failed to retrieve created event {event_id}"
790
+ raise RuntimeError(msg)
791
+
792
+ def list_events(self, session_id: str) -> "list[EventRecord]":
793
+ """List events for a session ordered by timestamp.
794
+
795
+ Args:
796
+ session_id: Session identifier.
797
+
798
+ Returns:
799
+ List of event records ordered by timestamp ASC.
800
+
801
+ Notes:
802
+ Uses index on (session_id, timestamp ASC).
803
+ JSON fields deserialized from JSON strings.
804
+ Converts INTEGER booleans to Python bool.
805
+ """
806
+ sql = f"""
807
+ SELECT id, session_id, app_name, user_id, invocation_id, author, actions,
808
+ long_running_tool_ids_json, branch, timestamp, content,
809
+ grounding_metadata, custom_metadata, partial, turn_complete,
810
+ interrupted, error_code, error_message
811
+ FROM {self._events_table}
812
+ WHERE session_id = ?
813
+ ORDER BY timestamp ASC
814
+ """
815
+
816
+ try:
817
+ with self._config.provide_connection() as conn:
818
+ cursor = conn.cursor()
819
+ try:
820
+ cursor.execute(sql, (session_id,))
821
+ rows = cursor.fetchall()
822
+
823
+ return [
824
+ EventRecord(
825
+ id=row[0],
826
+ session_id=row[1],
827
+ app_name=row[2],
828
+ user_id=row[3],
829
+ invocation_id=row[4],
830
+ author=row[5],
831
+ actions=bytes(row[6]) if row[6] is not None else b"",
832
+ long_running_tool_ids_json=row[7],
833
+ branch=row[8],
834
+ timestamp=row[9],
835
+ content=self._deserialize_json_field(row[10]),
836
+ grounding_metadata=self._deserialize_json_field(row[11]),
837
+ custom_metadata=self._deserialize_json_field(row[12]),
838
+ partial=self._from_int_bool(row[13]),
839
+ turn_complete=self._from_int_bool(row[14]),
840
+ interrupted=self._from_int_bool(row[15]),
841
+ error_code=row[16],
842
+ error_message=row[17],
843
+ )
844
+ for row in rows
845
+ ]
846
+ finally:
847
+ cursor.close() # type: ignore[no-untyped-call]
848
+ except Exception as e:
849
+ error_msg = str(e).lower()
850
+ if any(pattern in error_msg for pattern in ADBC_TABLE_NOT_FOUND_PATTERNS):
851
+ return []
852
+ raise
853
+
854
+ @staticmethod
855
+ def _to_int_bool(value: "bool | None") -> "int | None":
856
+ """Convert Python boolean to INTEGER (0/1).
857
+
858
+ Args:
859
+ value: Python boolean value or None.
860
+
861
+ Returns:
862
+ 1 for True, 0 for False, None for None.
863
+ """
864
+ if value is None:
865
+ return None
866
+ return 1 if value else 0
867
+
868
+ @staticmethod
869
+ def _from_int_bool(value: "int | None") -> "bool | None":
870
+ """Convert INTEGER to Python boolean.
871
+
872
+ Args:
873
+ value: INTEGER value (0, 1, or None).
874
+
875
+ Returns:
876
+ Python boolean or None.
877
+ """
878
+ if value is None:
879
+ return None
880
+ return bool(value)