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,5 @@
1
+ """Psqlpy ADK store module."""
2
+
3
+ from sqlspec.adapters.psqlpy.adk.store import PsqlpyADKStore
4
+
5
+ __all__ = ("PsqlpyADKStore",)
@@ -0,0 +1,483 @@
1
+ """Psqlpy ADK store for Google Agent Development Kit session/event storage."""
2
+
3
+ from typing import TYPE_CHECKING, Any, Final
4
+
5
+ import psqlpy.exceptions
6
+
7
+ from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
8
+ from sqlspec.utils.logging import get_logger
9
+
10
+ if TYPE_CHECKING:
11
+ from datetime import datetime
12
+
13
+ from sqlspec.adapters.psqlpy.config import PsqlpyConfig
14
+
15
+ logger = get_logger("adapters.psqlpy.adk.store")
16
+
17
+ __all__ = ("PsqlpyADKStore",)
18
+
19
+ POSTGRES_TABLE_NOT_FOUND_SQLSTATE: Final = "42P01"
20
+
21
+
22
+ class PsqlpyADKStore(BaseAsyncADKStore["PsqlpyConfig"]):
23
+ """PostgreSQL ADK store using Psqlpy driver.
24
+
25
+ Implements session and event storage for Google Agent Development Kit
26
+ using PostgreSQL via the high-performance Rust-based psqlpy driver.
27
+
28
+ Provides:
29
+ - Session state management with JSONB storage
30
+ - Event history tracking with BYTEA-serialized actions
31
+ - Microsecond-precision timestamps with TIMESTAMPTZ
32
+ - Foreign key constraints with cascade delete
33
+ - Efficient upserts using ON CONFLICT
34
+ - GIN indexes for JSONB queries
35
+ - HOT updates with FILLFACTOR 80
36
+
37
+ Args:
38
+ config: PsqlpyConfig with extension_config["adk"] settings.
39
+
40
+ Example:
41
+ from sqlspec.adapters.psqlpy import PsqlpyConfig
42
+ from sqlspec.adapters.psqlpy.adk import PsqlpyADKStore
43
+
44
+ config = PsqlpyConfig(
45
+ pool_config={"dsn": "postgresql://..."},
46
+ extension_config={
47
+ "adk": {
48
+ "session_table": "my_sessions",
49
+ "events_table": "my_events",
50
+ "owner_id_column": "tenant_id INTEGER NOT NULL REFERENCES tenants(id) ON DELETE CASCADE"
51
+ }
52
+ }
53
+ )
54
+ store = PsqlpyADKStore(config)
55
+ await store.create_tables()
56
+
57
+ Notes:
58
+ - PostgreSQL JSONB type used for state (more efficient than JSON)
59
+ - Psqlpy automatically converts Python dicts to/from JSONB
60
+ - TIMESTAMPTZ provides timezone-aware microsecond precision
61
+ - BYTEA for pre-serialized actions from Google ADK
62
+ - GIN index on state for JSONB queries (partial index)
63
+ - FILLFACTOR 80 leaves space for HOT updates
64
+ - Uses PostgreSQL numeric parameter style ($1, $2, $3)
65
+ - Configuration is read from config.extension_config["adk"]
66
+ """
67
+
68
+ __slots__ = ()
69
+
70
+ def __init__(self, config: "PsqlpyConfig") -> None:
71
+ """Initialize Psqlpy ADK store.
72
+
73
+ Args:
74
+ config: PsqlpyConfig instance.
75
+
76
+ Notes:
77
+ Configuration is read from config.extension_config["adk"]:
78
+ - session_table: Sessions table name (default: "adk_sessions")
79
+ - events_table: Events table name (default: "adk_events")
80
+ - owner_id_column: Optional owner FK column DDL (default: None)
81
+ """
82
+ super().__init__(config)
83
+
84
+ async def _get_create_sessions_table_sql(self) -> str:
85
+ """Get PostgreSQL CREATE TABLE SQL for sessions.
86
+
87
+ Returns:
88
+ SQL statement to create adk_sessions table with indexes.
89
+
90
+ Notes:
91
+ - VARCHAR(128) for IDs and names (sufficient for UUIDs and app names)
92
+ - JSONB type for state storage with default empty object
93
+ - TIMESTAMPTZ with microsecond precision
94
+ - FILLFACTOR 80 for HOT updates (reduces table bloat)
95
+ - Composite index on (app_name, user_id) for listing
96
+ - Index on update_time DESC for recent session queries
97
+ - Partial GIN index on state for JSONB queries (only non-empty)
98
+ - Optional owner ID column for multi-tenancy or user references
99
+ """
100
+ owner_id_line = ""
101
+ if self._owner_id_column_ddl:
102
+ owner_id_line = f",\n {self._owner_id_column_ddl}"
103
+
104
+ return f"""
105
+ CREATE TABLE IF NOT EXISTS {self._session_table} (
106
+ id VARCHAR(128) PRIMARY KEY,
107
+ app_name VARCHAR(128) NOT NULL,
108
+ user_id VARCHAR(128) NOT NULL{owner_id_line},
109
+ state JSONB NOT NULL DEFAULT '{{}}'::jsonb,
110
+ create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
111
+ update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
112
+ ) WITH (fillfactor = 80);
113
+
114
+ CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user
115
+ ON {self._session_table}(app_name, user_id);
116
+
117
+ CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time
118
+ ON {self._session_table}(update_time DESC);
119
+
120
+ CREATE INDEX IF NOT EXISTS idx_{self._session_table}_state
121
+ ON {self._session_table} USING GIN (state)
122
+ WHERE state != '{{}}'::jsonb;
123
+ """
124
+
125
+ async def _get_create_events_table_sql(self) -> str:
126
+ """Get PostgreSQL CREATE TABLE SQL for events.
127
+
128
+ Returns:
129
+ SQL statement to create adk_events table with indexes.
130
+
131
+ Notes:
132
+ - VARCHAR sizes: id(128), session_id(128), invocation_id(256), author(256),
133
+ branch(256), error_code(256), error_message(1024)
134
+ - BYTEA for pre-serialized actions (no size limit)
135
+ - JSONB for content, grounding_metadata, custom_metadata, long_running_tool_ids_json
136
+ - BOOLEAN for partial, turn_complete, interrupted
137
+ - Foreign key to sessions with CASCADE delete
138
+ - Index on (session_id, timestamp ASC) for ordered event retrieval
139
+ """
140
+ return f"""
141
+ CREATE TABLE IF NOT EXISTS {self._events_table} (
142
+ id VARCHAR(128) PRIMARY KEY,
143
+ session_id VARCHAR(128) NOT NULL,
144
+ app_name VARCHAR(128) NOT NULL,
145
+ user_id VARCHAR(128) NOT NULL,
146
+ invocation_id VARCHAR(256),
147
+ author VARCHAR(256),
148
+ actions BYTEA,
149
+ long_running_tool_ids_json JSONB,
150
+ branch VARCHAR(256),
151
+ timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
152
+ content JSONB,
153
+ grounding_metadata JSONB,
154
+ custom_metadata JSONB,
155
+ partial BOOLEAN,
156
+ turn_complete BOOLEAN,
157
+ interrupted BOOLEAN,
158
+ error_code VARCHAR(256),
159
+ error_message VARCHAR(1024),
160
+ FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ON DELETE CASCADE
161
+ );
162
+
163
+ CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session
164
+ ON {self._events_table}(session_id, timestamp ASC);
165
+ """
166
+
167
+ def _get_drop_tables_sql(self) -> "list[str]":
168
+ """Get PostgreSQL DROP TABLE SQL statements.
169
+
170
+ Returns:
171
+ List of SQL statements to drop tables and indexes.
172
+
173
+ Notes:
174
+ Order matters: drop events table (child) before sessions (parent).
175
+ PostgreSQL automatically drops indexes when dropping tables.
176
+ """
177
+ return [f"DROP TABLE IF EXISTS {self._events_table}", f"DROP TABLE IF EXISTS {self._session_table}"]
178
+
179
+ async def create_tables(self) -> None:
180
+ """Create both sessions and events tables if they don't exist.
181
+
182
+ Notes:
183
+ Uses driver.execute_script() which handles multiple statements.
184
+ Creates sessions table first, then events table (FK dependency).
185
+ """
186
+ async with self._config.provide_session() as driver:
187
+ await driver.execute_script(await self._get_create_sessions_table_sql())
188
+ await driver.execute_script(await self._get_create_events_table_sql())
189
+
190
+ logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
191
+
192
+ async def create_session(
193
+ self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None
194
+ ) -> SessionRecord:
195
+ """Create a new session.
196
+
197
+ Args:
198
+ session_id: Unique session identifier.
199
+ app_name: Application name.
200
+ user_id: User identifier.
201
+ state: Initial session state.
202
+ owner_id: Optional owner ID value for owner_id_column (if configured).
203
+
204
+ Returns:
205
+ Created session record.
206
+
207
+ Notes:
208
+ Uses CURRENT_TIMESTAMP for create_time and update_time.
209
+ State is passed as dict and psqlpy converts to JSONB automatically.
210
+ If owner_id_column is configured, owner_id value must be provided.
211
+ """
212
+ async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue]
213
+ if self._owner_id_column_name:
214
+ sql = f"""
215
+ INSERT INTO {self._session_table}
216
+ (id, app_name, user_id, {self._owner_id_column_name}, state, create_time, update_time)
217
+ VALUES ($1, $2, $3, $4, $5, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
218
+ """
219
+ await conn.execute(sql, [session_id, app_name, user_id, owner_id, state])
220
+ else:
221
+ sql = f"""
222
+ INSERT INTO {self._session_table} (id, app_name, user_id, state, create_time, update_time)
223
+ VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
224
+ """
225
+ await conn.execute(sql, [session_id, app_name, user_id, state])
226
+
227
+ return await self.get_session(session_id) # type: ignore[return-value]
228
+
229
+ async def get_session(self, session_id: str) -> "SessionRecord | None":
230
+ """Get session by ID.
231
+
232
+ Args:
233
+ session_id: Session identifier.
234
+
235
+ Returns:
236
+ Session record or None if not found.
237
+
238
+ Notes:
239
+ PostgreSQL returns datetime objects for TIMESTAMPTZ columns.
240
+ JSONB is automatically parsed by psqlpy to Python dicts.
241
+ Returns None if table doesn't exist (catches database errors).
242
+ """
243
+ sql = f"""
244
+ SELECT id, app_name, user_id, state, create_time, update_time
245
+ FROM {self._session_table}
246
+ WHERE id = $1
247
+ """
248
+
249
+ try:
250
+ async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue]
251
+ result = await conn.fetch(sql, [session_id])
252
+ rows: list[dict[str, Any]] = result.result() if result else []
253
+
254
+ if not rows:
255
+ return None
256
+
257
+ row = rows[0]
258
+ return SessionRecord(
259
+ id=row["id"],
260
+ app_name=row["app_name"],
261
+ user_id=row["user_id"],
262
+ state=row["state"],
263
+ create_time=row["create_time"],
264
+ update_time=row["update_time"],
265
+ )
266
+ except psqlpy.exceptions.DatabaseError as e:
267
+ error_msg = str(e).lower()
268
+ if "does not exist" in error_msg or "relation" in error_msg:
269
+ return None
270
+ raise
271
+
272
+ async def update_session_state(self, session_id: str, state: "dict[str, Any]") -> None:
273
+ """Update session state.
274
+
275
+ Args:
276
+ session_id: Session identifier.
277
+ state: New state dictionary (replaces existing state).
278
+
279
+ Notes:
280
+ This replaces the entire state dictionary.
281
+ Uses CURRENT_TIMESTAMP for update_time.
282
+ Psqlpy automatically converts dict to JSONB.
283
+ """
284
+ sql = f"""
285
+ UPDATE {self._session_table}
286
+ SET state = $1, update_time = CURRENT_TIMESTAMP
287
+ WHERE id = $2
288
+ """
289
+
290
+ async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue]
291
+ await conn.execute(sql, [state, session_id])
292
+
293
+ async def delete_session(self, session_id: str) -> None:
294
+ """Delete session and all associated events (cascade).
295
+
296
+ Args:
297
+ session_id: Session identifier.
298
+
299
+ Notes:
300
+ Foreign key constraint ensures events are cascade-deleted.
301
+ """
302
+ sql = f"DELETE FROM {self._session_table} WHERE id = $1"
303
+
304
+ async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue]
305
+ await conn.execute(sql, [session_id])
306
+
307
+ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]":
308
+ """List sessions for an app, optionally filtered by user.
309
+
310
+ Args:
311
+ app_name: Application name.
312
+ user_id: User identifier. If None, lists all sessions for the app.
313
+
314
+ Returns:
315
+ List of session records ordered by update_time DESC.
316
+
317
+ Notes:
318
+ Uses composite index on (app_name, user_id) when user_id is provided.
319
+ Returns empty list if table doesn't exist.
320
+ """
321
+ if user_id is None:
322
+ sql = f"""
323
+ SELECT id, app_name, user_id, state, create_time, update_time
324
+ FROM {self._session_table}
325
+ WHERE app_name = $1
326
+ ORDER BY update_time DESC
327
+ """
328
+ params = [app_name]
329
+ else:
330
+ sql = f"""
331
+ SELECT id, app_name, user_id, state, create_time, update_time
332
+ FROM {self._session_table}
333
+ WHERE app_name = $1 AND user_id = $2
334
+ ORDER BY update_time DESC
335
+ """
336
+ params = [app_name, user_id]
337
+
338
+ try:
339
+ async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue]
340
+ result = await conn.fetch(sql, params)
341
+ rows: list[dict[str, Any]] = result.result() if result else []
342
+
343
+ return [
344
+ SessionRecord(
345
+ id=row["id"],
346
+ app_name=row["app_name"],
347
+ user_id=row["user_id"],
348
+ state=row["state"],
349
+ create_time=row["create_time"],
350
+ update_time=row["update_time"],
351
+ )
352
+ for row in rows
353
+ ]
354
+ except psqlpy.exceptions.DatabaseError as e:
355
+ error_msg = str(e).lower()
356
+ if "does not exist" in error_msg or "relation" in error_msg:
357
+ return []
358
+ raise
359
+
360
+ async def append_event(self, event_record: EventRecord) -> None:
361
+ """Append an event to a session.
362
+
363
+ Args:
364
+ event_record: Event record to store.
365
+
366
+ Notes:
367
+ Uses CURRENT_TIMESTAMP for timestamp if not provided.
368
+ JSONB fields are passed as dicts and psqlpy converts automatically.
369
+ BYTEA actions field stores pre-serialized data from Google ADK.
370
+ """
371
+ content_json = event_record.get("content")
372
+ grounding_metadata_json = event_record.get("grounding_metadata")
373
+ custom_metadata_json = event_record.get("custom_metadata")
374
+
375
+ sql = f"""
376
+ INSERT INTO {self._events_table} (
377
+ id, session_id, app_name, user_id, invocation_id, author, actions,
378
+ long_running_tool_ids_json, branch, timestamp, content,
379
+ grounding_metadata, custom_metadata, partial, turn_complete,
380
+ interrupted, error_code, error_message
381
+ ) VALUES (
382
+ $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18
383
+ )
384
+ """
385
+
386
+ async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue]
387
+ await conn.execute(
388
+ sql,
389
+ [
390
+ event_record["id"],
391
+ event_record["session_id"],
392
+ event_record["app_name"],
393
+ event_record["user_id"],
394
+ event_record.get("invocation_id"),
395
+ event_record.get("author"),
396
+ event_record.get("actions"),
397
+ event_record.get("long_running_tool_ids_json"),
398
+ event_record.get("branch"),
399
+ event_record["timestamp"],
400
+ content_json,
401
+ grounding_metadata_json,
402
+ custom_metadata_json,
403
+ event_record.get("partial"),
404
+ event_record.get("turn_complete"),
405
+ event_record.get("interrupted"),
406
+ event_record.get("error_code"),
407
+ event_record.get("error_message"),
408
+ ],
409
+ )
410
+
411
+ async def get_events(
412
+ self, session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None
413
+ ) -> "list[EventRecord]":
414
+ """Get events for a session.
415
+
416
+ Args:
417
+ session_id: Session identifier.
418
+ after_timestamp: Only return events after this time.
419
+ limit: Maximum number of events to return.
420
+
421
+ Returns:
422
+ List of event records ordered by timestamp ASC.
423
+
424
+ Notes:
425
+ Uses index on (session_id, timestamp ASC).
426
+ Parses JSONB fields and converts BYTEA actions to bytes.
427
+ Returns empty list if table doesn't exist.
428
+ """
429
+ where_clauses = ["session_id = $1"]
430
+ params: list[Any] = [session_id]
431
+
432
+ if after_timestamp is not None:
433
+ where_clauses.append(f"timestamp > ${len(params) + 1}")
434
+ params.append(after_timestamp)
435
+
436
+ where_clause = " AND ".join(where_clauses)
437
+ limit_clause = f" LIMIT ${len(params) + 1}" if limit else ""
438
+ if limit:
439
+ params.append(limit)
440
+
441
+ sql = f"""
442
+ SELECT id, session_id, app_name, user_id, invocation_id, author, actions,
443
+ long_running_tool_ids_json, branch, timestamp, content,
444
+ grounding_metadata, custom_metadata, partial, turn_complete,
445
+ interrupted, error_code, error_message
446
+ FROM {self._events_table}
447
+ WHERE {where_clause}
448
+ ORDER BY timestamp ASC{limit_clause}
449
+ """
450
+
451
+ try:
452
+ async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue]
453
+ result = await conn.fetch(sql, params)
454
+ rows: list[dict[str, Any]] = result.result() if result else []
455
+
456
+ return [
457
+ EventRecord(
458
+ id=row["id"],
459
+ session_id=row["session_id"],
460
+ app_name=row["app_name"],
461
+ user_id=row["user_id"],
462
+ invocation_id=row["invocation_id"],
463
+ author=row["author"],
464
+ actions=bytes(row["actions"]) if row["actions"] else b"",
465
+ long_running_tool_ids_json=row["long_running_tool_ids_json"],
466
+ branch=row["branch"],
467
+ timestamp=row["timestamp"],
468
+ content=row["content"],
469
+ grounding_metadata=row["grounding_metadata"],
470
+ custom_metadata=row["custom_metadata"],
471
+ partial=row["partial"],
472
+ turn_complete=row["turn_complete"],
473
+ interrupted=row["interrupted"],
474
+ error_code=row["error_code"],
475
+ error_message=row["error_message"],
476
+ )
477
+ for row in rows
478
+ ]
479
+ except psqlpy.exceptions.DatabaseError as e:
480
+ error_msg = str(e).lower()
481
+ if "does not exist" in error_msg or "relation" in error_msg:
482
+ return []
483
+ 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, Optional, TypedDict, Union
6
+ from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
7
7
 
8
8
  from psqlpy import ConnectionPool
9
9
  from typing_extensions import NotRequired
@@ -12,6 +12,7 @@ from sqlspec.adapters.psqlpy._types import PsqlpyConnection
12
12
  from sqlspec.adapters.psqlpy.driver import PsqlpyCursor, PsqlpyDriver, psqlpy_statement_config
13
13
  from sqlspec.config import AsyncDatabaseConfig
14
14
  from sqlspec.core.statement import StatementConfig
15
+ from sqlspec.typing import PGVECTOR_INSTALLED
15
16
 
16
17
  if TYPE_CHECKING:
17
18
  from collections.abc import Callable
@@ -20,7 +21,7 @@ if TYPE_CHECKING:
20
21
  logger = logging.getLogger("sqlspec.adapters.psqlpy")
21
22
 
22
23
 
23
- class PsqlpyConnectionParams(TypedDict, total=False):
24
+ class PsqlpyConnectionParams(TypedDict):
24
25
  """Psqlpy connection parameters."""
25
26
 
26
27
  dsn: NotRequired[str]
@@ -62,7 +63,7 @@ class PsqlpyConnectionParams(TypedDict, total=False):
62
63
  load_balance_hosts: NotRequired[str]
63
64
 
64
65
 
65
- class PsqlpyPoolParams(PsqlpyConnectionParams, total=False):
66
+ class PsqlpyPoolParams(PsqlpyConnectionParams):
66
67
  """Psqlpy pool parameters."""
67
68
 
68
69
  hosts: NotRequired[list[str]]
@@ -73,7 +74,19 @@ class PsqlpyPoolParams(PsqlpyConnectionParams, total=False):
73
74
  extra: NotRequired[dict[str, Any]]
74
75
 
75
76
 
76
- __all__ = ("PsqlpyConfig", "PsqlpyConnectionParams", "PsqlpyCursor", "PsqlpyPoolParams")
77
+ class PsqlpyDriverFeatures(TypedDict):
78
+ """Psqlpy driver feature flags.
79
+
80
+ enable_pgvector: Enable automatic pgvector extension support for vector similarity search.
81
+ Requires pgvector-python package installed.
82
+ Defaults to True when pgvector is installed.
83
+ Provides automatic conversion between NumPy arrays and PostgreSQL vector types.
84
+ """
85
+
86
+ enable_pgvector: NotRequired[bool]
87
+
88
+
89
+ __all__ = ("PsqlpyConfig", "PsqlpyConnectionParams", "PsqlpyCursor", "PsqlpyDriverFeatures", "PsqlpyPoolParams")
77
90
 
78
91
 
79
92
  class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyDriver]):
@@ -81,38 +94,47 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
81
94
 
82
95
  driver_type: ClassVar[type[PsqlpyDriver]] = PsqlpyDriver
83
96
  connection_type: "ClassVar[type[PsqlpyConnection]]" = PsqlpyConnection
97
+ supports_transactional_ddl: "ClassVar[bool]" = True
84
98
 
85
99
  def __init__(
86
100
  self,
87
101
  *,
88
- pool_config: Optional[Union[PsqlpyPoolParams, dict[str, Any]]] = None,
89
- pool_instance: Optional[ConnectionPool] = None,
90
- migration_config: Optional[dict[str, Any]] = None,
91
- statement_config: Optional[StatementConfig] = None,
92
- driver_features: Optional[dict[str, Any]] = None,
93
- bind_key: Optional[str] = None,
102
+ pool_config: PsqlpyPoolParams | dict[str, Any] | None = None,
103
+ pool_instance: ConnectionPool | None = None,
104
+ migration_config: dict[str, Any] | None = None,
105
+ statement_config: StatementConfig | None = None,
106
+ driver_features: "PsqlpyDriverFeatures | dict[str, Any] | None" = None,
107
+ bind_key: str | None = None,
108
+ extension_config: "dict[str, dict[str, Any]] | None" = None,
94
109
  ) -> None:
95
110
  """Initialize Psqlpy configuration.
96
111
 
97
112
  Args:
98
- pool_config: Pool configuration parameters
99
- pool_instance: Existing connection pool instance to use
100
- migration_config: Migration configuration
101
- statement_config: SQL statement configuration
102
- driver_features: Driver feature configuration
103
- bind_key: Optional unique identifier for this configuration
113
+ pool_config: Pool configuration parameters.
114
+ pool_instance: Existing connection pool instance to use.
115
+ migration_config: Migration configuration.
116
+ statement_config: SQL statement configuration.
117
+ driver_features: Driver feature configuration (TypedDict or dict).
118
+ bind_key: Optional unique identifier for this configuration.
119
+ extension_config: Extension-specific configuration (e.g., Litestar plugin settings).
104
120
  """
105
121
  processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
106
122
  if "extra" in processed_pool_config:
107
123
  extras = processed_pool_config.pop("extra")
108
124
  processed_pool_config.update(extras)
125
+
126
+ processed_driver_features: dict[str, Any] = dict(driver_features) if driver_features else {}
127
+ if "enable_pgvector" not in processed_driver_features:
128
+ processed_driver_features["enable_pgvector"] = PGVECTOR_INSTALLED
129
+
109
130
  super().__init__(
110
131
  pool_config=processed_pool_config,
111
132
  pool_instance=pool_instance,
112
133
  migration_config=migration_config,
113
134
  statement_config=statement_config or psqlpy_statement_config,
114
- driver_features=driver_features or {},
135
+ driver_features=processed_driver_features,
115
136
  bind_key=bind_key,
137
+ extension_config=extension_config,
116
138
  )
117
139
 
118
140
  def _get_pool_config_dict(self) -> dict[str, Any]:
@@ -185,7 +207,7 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
185
207
 
186
208
  @asynccontextmanager
187
209
  async def provide_session(
188
- self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
210
+ self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
189
211
  ) -> AsyncGenerator[PsqlpyDriver, None]:
190
212
  """Provide an async driver session context manager.
191
213
 
@@ -198,7 +220,11 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
198
220
  A PsqlpyDriver instance.
199
221
  """
200
222
  async with self.provide_connection(*args, **kwargs) as conn:
201
- yield self.driver_type(connection=conn, statement_config=statement_config or self.statement_config)
223
+ yield self.driver_type(
224
+ connection=conn,
225
+ statement_config=statement_config or self.statement_config,
226
+ driver_features=self.driver_features,
227
+ )
202
228
 
203
229
  async def provide_pool(self, *args: Any, **kwargs: Any) -> ConnectionPool:
204
230
  """Provide async pool instance.