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
@@ -1,12 +1,14 @@
1
1
  """PostgreSQL-specific data dictionary for metadata queries via psqlpy."""
2
2
 
3
3
  import re
4
- from typing import TYPE_CHECKING, Callable, Optional, cast
4
+ from typing import TYPE_CHECKING, Any, cast
5
5
 
6
6
  from sqlspec.driver import AsyncDataDictionaryBase, AsyncDriverAdapterBase, VersionInfo
7
7
  from sqlspec.utils.logging import get_logger
8
8
 
9
9
  if TYPE_CHECKING:
10
+ from collections.abc import Callable
11
+
10
12
  from sqlspec.adapters.psqlpy.driver import PsqlpyDriver
11
13
 
12
14
  logger = get_logger("adapters.psqlpy.data_dictionary")
@@ -20,7 +22,7 @@ __all__ = ("PsqlpyAsyncDataDictionary",)
20
22
  class PsqlpyAsyncDataDictionary(AsyncDataDictionaryBase):
21
23
  """PostgreSQL-specific async data dictionary via psqlpy."""
22
24
 
23
- async def get_version(self, driver: AsyncDriverAdapterBase) -> "Optional[VersionInfo]":
25
+ async def get_version(self, driver: AsyncDriverAdapterBase) -> "VersionInfo | None":
24
26
  """Get PostgreSQL database version information.
25
27
 
26
28
  Args:
@@ -111,6 +113,50 @@ class PsqlpyAsyncDataDictionary(AsyncDataDictionaryBase):
111
113
  }
112
114
  return type_map.get(type_category, "TEXT")
113
115
 
116
+ async def get_columns(
117
+ self, driver: AsyncDriverAdapterBase, table: str, schema: "str | None" = None
118
+ ) -> "list[dict[str, Any]]":
119
+ """Get column information for a table using pg_catalog.
120
+
121
+ Args:
122
+ driver: Psqlpy async driver instance
123
+ table: Table name to query columns for
124
+ schema: Schema name (None for default 'public')
125
+
126
+ Returns:
127
+ List of column metadata dictionaries with keys:
128
+ - column_name: Name of the column
129
+ - data_type: PostgreSQL data type
130
+ - is_nullable: Whether column allows NULL (YES/NO)
131
+ - column_default: Default value if any
132
+
133
+ Notes:
134
+ Uses pg_catalog instead of information_schema to avoid psqlpy's
135
+ inability to handle the PostgreSQL 'name' type returned by information_schema.
136
+ """
137
+ psqlpy_driver = cast("PsqlpyDriver", driver)
138
+
139
+ schema_name = schema or "public"
140
+ sql = """
141
+ SELECT
142
+ a.attname::text AS column_name,
143
+ pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
144
+ CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable,
145
+ pg_catalog.pg_get_expr(d.adbin, d.adrelid)::text AS column_default
146
+ FROM pg_catalog.pg_attribute a
147
+ JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
148
+ JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
149
+ LEFT JOIN pg_catalog.pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
150
+ WHERE c.relname = $1
151
+ AND n.nspname = $2
152
+ AND a.attnum > 0
153
+ AND NOT a.attisdropped
154
+ ORDER BY a.attnum
155
+ """
156
+
157
+ result = await psqlpy_driver.execute(sql, (table, schema_name))
158
+ return result.data or []
159
+
114
160
  def list_available_features(self) -> "list[str]":
115
161
  """List available PostgreSQL feature flags.
116
162
 
@@ -6,17 +6,29 @@ and transaction management.
6
6
 
7
7
  import decimal
8
8
  import re
9
- from typing import TYPE_CHECKING, Any, Final, Optional
9
+ from typing import TYPE_CHECKING, Any, Final
10
10
 
11
- import psqlpy
12
11
  import psqlpy.exceptions
13
12
 
13
+ from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyAsyncDataDictionary
14
14
  from sqlspec.adapters.psqlpy.type_converter import PostgreSQLTypeConverter
15
15
  from sqlspec.core.cache import get_cache_config
16
16
  from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
17
17
  from sqlspec.core.statement import SQL, StatementConfig
18
18
  from sqlspec.driver import AsyncDriverAdapterBase
19
- from sqlspec.exceptions import SQLParsingError, SQLSpecError
19
+ from sqlspec.exceptions import (
20
+ CheckViolationError,
21
+ DatabaseConnectionError,
22
+ DataError,
23
+ ForeignKeyViolationError,
24
+ IntegrityError,
25
+ NotNullViolationError,
26
+ OperationalError,
27
+ SQLParsingError,
28
+ SQLSpecError,
29
+ TransactionError,
30
+ UniqueViolationError,
31
+ )
20
32
  from sqlspec.utils.logging import get_logger
21
33
 
22
34
  if TYPE_CHECKING:
@@ -93,7 +105,11 @@ class PsqlpyCursor:
93
105
 
94
106
 
95
107
  class PsqlpyExceptionHandler:
96
- """Async context manager for handling psqlpy database exceptions."""
108
+ """Async context manager for handling psqlpy database exceptions.
109
+
110
+ Maps PostgreSQL SQLSTATE error codes to specific SQLSpec exceptions
111
+ for better error handling in application code.
112
+ """
97
113
 
98
114
  __slots__ = ()
99
115
 
@@ -103,31 +119,85 @@ class PsqlpyExceptionHandler:
103
119
  async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
104
120
  if exc_type is None:
105
121
  return
122
+ if issubclass(exc_type, (psqlpy.exceptions.DatabaseError, psqlpy.exceptions.Error)):
123
+ self._map_postgres_exception(exc_val)
106
124
 
107
- if issubclass(exc_type, psqlpy.exceptions.DatabaseError):
108
- e = exc_val
109
- msg = f"Psqlpy database error: {e}"
110
- raise SQLSpecError(msg) from e
111
- if issubclass(exc_type, psqlpy.exceptions.InterfaceError):
112
- e = exc_val
113
- msg = f"Psqlpy interface error: {e}"
114
- raise SQLSpecError(msg) from e
115
- if issubclass(exc_type, psqlpy.exceptions.Error):
116
- e = exc_val
117
- error_msg = str(e).lower()
118
- if "syntax" in error_msg or "parse" in error_msg:
119
- msg = f"Psqlpy SQL syntax error: {e}"
120
- raise SQLParsingError(msg) from e
121
- msg = f"Psqlpy error: {e}"
122
- raise SQLSpecError(msg) from e
123
- if issubclass(exc_type, Exception):
124
- e = exc_val
125
- error_msg = str(e).lower()
126
- if "parse" in error_msg or "syntax" in error_msg:
127
- msg = f"SQL parsing failed: {e}"
128
- raise SQLParsingError(msg) from e
129
- msg = f"Unexpected async database operation error: {e}"
130
- raise SQLSpecError(msg) from e
125
+ def _map_postgres_exception(self, e: Any) -> None:
126
+ """Map PostgreSQL exception to SQLSpec exception.
127
+
128
+ psqlpy does not expose SQLSTATE codes directly, so we use message-based
129
+ detection to map exceptions.
130
+
131
+ Args:
132
+ e: psqlpy exception instance
133
+
134
+ Raises:
135
+ Specific SQLSpec exception based on error message patterns
136
+ """
137
+ error_msg = str(e).lower()
138
+
139
+ if "unique" in error_msg or "duplicate key" in error_msg:
140
+ self._raise_unique_violation(e, None)
141
+ elif "foreign key" in error_msg or "violates foreign key" in error_msg:
142
+ self._raise_foreign_key_violation(e, None)
143
+ elif "not null" in error_msg or ("null value" in error_msg and "violates not-null" in error_msg):
144
+ self._raise_not_null_violation(e, None)
145
+ elif "check constraint" in error_msg or "violates check constraint" in error_msg:
146
+ self._raise_check_violation(e, None)
147
+ elif "constraint" in error_msg:
148
+ self._raise_integrity_error(e, None)
149
+ elif "syntax error" in error_msg or "parse" in error_msg:
150
+ self._raise_parsing_error(e, None)
151
+ elif "connection" in error_msg or "could not connect" in error_msg:
152
+ self._raise_connection_error(e, None)
153
+ elif "deadlock" in error_msg or "serialization failure" in error_msg:
154
+ self._raise_transaction_error(e, None)
155
+ else:
156
+ self._raise_generic_error(e, None)
157
+
158
+ def _raise_unique_violation(self, e: Any, code: "str | None") -> None:
159
+ msg = f"PostgreSQL unique constraint violation: {e}"
160
+ raise UniqueViolationError(msg) from e
161
+
162
+ def _raise_foreign_key_violation(self, e: Any, code: "str | None") -> None:
163
+ msg = f"PostgreSQL foreign key constraint violation: {e}"
164
+ raise ForeignKeyViolationError(msg) from e
165
+
166
+ def _raise_not_null_violation(self, e: Any, code: "str | None") -> None:
167
+ msg = f"PostgreSQL not-null constraint violation: {e}"
168
+ raise NotNullViolationError(msg) from e
169
+
170
+ def _raise_check_violation(self, e: Any, code: "str | None") -> None:
171
+ msg = f"PostgreSQL check constraint violation: {e}"
172
+ raise CheckViolationError(msg) from e
173
+
174
+ def _raise_integrity_error(self, e: Any, code: "str | None") -> None:
175
+ msg = f"PostgreSQL integrity constraint violation: {e}"
176
+ raise IntegrityError(msg) from e
177
+
178
+ def _raise_parsing_error(self, e: Any, code: "str | None") -> None:
179
+ msg = f"PostgreSQL SQL syntax error: {e}"
180
+ raise SQLParsingError(msg) from e
181
+
182
+ def _raise_connection_error(self, e: Any, code: "str | None") -> None:
183
+ msg = f"PostgreSQL connection error: {e}"
184
+ raise DatabaseConnectionError(msg) from e
185
+
186
+ def _raise_transaction_error(self, e: Any, code: "str | None") -> None:
187
+ msg = f"PostgreSQL transaction error: {e}"
188
+ raise TransactionError(msg) from e
189
+
190
+ def _raise_data_error(self, e: Any, code: "str | None") -> None:
191
+ msg = f"PostgreSQL data error: {e}"
192
+ raise DataError(msg) from e
193
+
194
+ def _raise_operational_error(self, e: Any, code: "str | None") -> None:
195
+ msg = f"PostgreSQL operational error: {e}"
196
+ raise OperationalError(msg) from e
197
+
198
+ def _raise_generic_error(self, e: Any, code: "str | None") -> None:
199
+ msg = f"PostgreSQL database error: {e}"
200
+ raise SQLSpecError(msg) from e
131
201
 
132
202
 
133
203
  class PsqlpyDriver(AsyncDriverAdapterBase):
@@ -143,8 +213,8 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
143
213
  def __init__(
144
214
  self,
145
215
  connection: "PsqlpyConnection",
146
- statement_config: "Optional[StatementConfig]" = None,
147
- driver_features: "Optional[dict[str, Any]]" = None,
216
+ statement_config: "StatementConfig | None" = None,
217
+ driver_features: "dict[str, Any] | None" = None,
148
218
  ) -> None:
149
219
  if statement_config is None:
150
220
  cache_config = get_cache_config()
@@ -156,7 +226,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
156
226
  )
157
227
 
158
228
  super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
159
- self._data_dictionary: Optional[AsyncDataDictionaryBase] = None
229
+ self._data_dictionary: AsyncDataDictionaryBase | None = None
160
230
 
161
231
  def with_cursor(self, connection: "PsqlpyConnection") -> "PsqlpyCursor":
162
232
  """Create context manager for psqlpy cursor.
@@ -177,7 +247,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
177
247
  """
178
248
  return PsqlpyExceptionHandler()
179
249
 
180
- async def _try_special_handling(self, cursor: "PsqlpyConnection", statement: SQL) -> "Optional[SQLResult]":
250
+ async def _try_special_handling(self, cursor: "PsqlpyConnection", statement: SQL) -> "SQLResult | None":
181
251
  """Hook for psqlpy-specific special operations.
182
252
 
183
253
  Args:
@@ -199,17 +269,16 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
199
269
 
200
270
  Returns:
201
271
  ExecutionResult with script execution metadata
272
+
273
+ Notes:
274
+ Uses execute() with empty parameters for each statement instead of execute_batch().
275
+ execute_batch() uses simple query protocol which can break subsequent queries
276
+ that rely on extended protocol (e.g., information_schema queries with name type).
202
277
  """
203
278
  sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
204
279
  statement_config = statement.statement_config
205
-
206
- if not prepared_parameters:
207
- await cursor.execute_batch(sql)
208
- statements = self.split_script_statements(sql, statement_config, strip_trailing_semicolon=True)
209
- return self.create_execution_result(
210
- cursor, statement_count=len(statements), successful_statements=len(statements), is_script_result=True
211
- )
212
280
  statements = self.split_script_statements(sql, statement_config, strip_trailing_semicolon=True)
281
+
213
282
  successful_count = 0
214
283
  last_result = None
215
284
 
@@ -351,7 +420,5 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
351
420
  Data dictionary instance for metadata queries
352
421
  """
353
422
  if self._data_dictionary is None:
354
- from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyAsyncDataDictionary
355
-
356
423
  self._data_dictionary = PsqlpyAsyncDataDictionary()
357
424
  return self._data_dictionary
@@ -0,0 +1,5 @@
1
+ """Litestar integration for psqlpy adapter."""
2
+
3
+ from sqlspec.adapters.psqlpy.litestar.store import PsqlpyStore
4
+
5
+ __all__ = ("PsqlpyStore",)
@@ -0,0 +1,272 @@
1
+ """Psqlpy session store for Litestar integration."""
2
+
3
+ from datetime import datetime, timedelta, timezone
4
+ from typing import TYPE_CHECKING
5
+
6
+ from sqlspec.extensions.litestar.store import BaseSQLSpecStore
7
+ from sqlspec.utils.logging import get_logger
8
+
9
+ if TYPE_CHECKING:
10
+ from sqlspec.adapters.psqlpy.config import PsqlpyConfig
11
+
12
+ logger = get_logger("adapters.psqlpy.litestar.store")
13
+
14
+ __all__ = ("PsqlpyStore",)
15
+
16
+
17
+ class PsqlpyStore(BaseSQLSpecStore["PsqlpyConfig"]):
18
+ """PostgreSQL session store using Psqlpy driver.
19
+
20
+ Implements server-side session storage for Litestar using PostgreSQL
21
+ via the Psqlpy driver (Rust-based async driver). Provides efficient
22
+ session management with:
23
+ - Native async PostgreSQL operations via Rust
24
+ - UPSERT support using ON CONFLICT
25
+ - Automatic expiration handling
26
+ - Efficient cleanup of expired sessions
27
+
28
+ Args:
29
+ config: PsqlpyConfig instance.
30
+
31
+ Example:
32
+ from sqlspec.adapters.psqlpy import PsqlpyConfig
33
+ from sqlspec.adapters.psqlpy.litestar.store import PsqlpyStore
34
+
35
+ config = PsqlpyConfig(pool_config={"dsn": "postgresql://..."})
36
+ store = PsqlpyStore(config)
37
+ await store.create_table()
38
+ """
39
+
40
+ __slots__ = ()
41
+
42
+ def __init__(self, config: "PsqlpyConfig") -> None:
43
+ """Initialize Psqlpy session store.
44
+
45
+ Args:
46
+ config: PsqlpyConfig instance.
47
+
48
+ Notes:
49
+ Table name is read from config.extension_config["litestar"]["session_table"].
50
+ """
51
+ super().__init__(config)
52
+
53
+ def _get_create_table_sql(self) -> str:
54
+ """Get PostgreSQL CREATE TABLE SQL with optimized schema.
55
+
56
+ Returns:
57
+ SQL statement to create the sessions table with proper indexes.
58
+
59
+ Notes:
60
+ - Uses TIMESTAMPTZ for timezone-aware expiration timestamps
61
+ - Partial index WHERE expires_at IS NOT NULL reduces index size/maintenance
62
+ - FILLFACTOR 80 leaves space for HOT updates, reducing table bloat
63
+ - Audit columns (created_at, updated_at) help with debugging
64
+ - Table name is internally controlled, not user input (S608 suppressed)
65
+ """
66
+ return f"""
67
+ CREATE TABLE IF NOT EXISTS {self._table_name} (
68
+ session_id TEXT PRIMARY KEY,
69
+ data BYTEA NOT NULL,
70
+ expires_at TIMESTAMPTZ,
71
+ created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
72
+ updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
73
+ ) WITH (fillfactor = 80);
74
+
75
+ CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at
76
+ ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL;
77
+
78
+ ALTER TABLE {self._table_name} SET (
79
+ autovacuum_vacuum_scale_factor = 0.05,
80
+ autovacuum_analyze_scale_factor = 0.02
81
+ );
82
+ """
83
+
84
+ def _get_drop_table_sql(self) -> "list[str]":
85
+ """Get PostgreSQL DROP TABLE SQL statements.
86
+
87
+ Returns:
88
+ List of SQL statements to drop indexes and table.
89
+ """
90
+ return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"]
91
+
92
+ async def create_table(self) -> None:
93
+ """Create the session table if it doesn't exist."""
94
+ sql = self._get_create_table_sql()
95
+ async with self._config.provide_session() as driver:
96
+ await driver.execute_script(sql)
97
+ logger.debug("Created session table: %s", self._table_name)
98
+
99
+ async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
100
+ """Get a session value by key.
101
+
102
+ Args:
103
+ key: Session ID to retrieve.
104
+ renew_for: If given, renew the expiry time for this duration.
105
+
106
+ Returns:
107
+ Session data as bytes if found and not expired, None otherwise.
108
+
109
+ Notes:
110
+ Uses CURRENT_TIMESTAMP instead of NOW() for SQL standard compliance.
111
+ The query planner can use the partial index for expires_at > CURRENT_TIMESTAMP.
112
+ """
113
+ sql = f"""
114
+ SELECT data, expires_at FROM {self._table_name}
115
+ WHERE session_id = $1
116
+ AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
117
+ """
118
+
119
+ async with self._config.provide_connection() as conn:
120
+ query_result = await conn.fetch(sql, [key])
121
+ rows = query_result.result()
122
+
123
+ if not rows:
124
+ return None
125
+
126
+ row = rows[0]
127
+
128
+ if renew_for is not None and row["expires_at"] is not None:
129
+ new_expires_at = self._calculate_expires_at(renew_for)
130
+ if new_expires_at is not None:
131
+ update_sql = f"""
132
+ UPDATE {self._table_name}
133
+ SET expires_at = $1, updated_at = CURRENT_TIMESTAMP
134
+ WHERE session_id = $2
135
+ """
136
+ await conn.execute(update_sql, [new_expires_at, key])
137
+
138
+ return bytes(row["data"])
139
+
140
+ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
141
+ """Store a session value.
142
+
143
+ Args:
144
+ key: Session ID.
145
+ value: Session data.
146
+ expires_in: Time until expiration.
147
+
148
+ Notes:
149
+ Uses EXCLUDED to reference the proposed insert values in ON CONFLICT.
150
+ Updates updated_at timestamp on every write for audit trail.
151
+ """
152
+ data = self._value_to_bytes(value)
153
+ expires_at = self._calculate_expires_at(expires_in)
154
+
155
+ sql = f"""
156
+ INSERT INTO {self._table_name} (session_id, data, expires_at)
157
+ VALUES ($1, $2, $3)
158
+ ON CONFLICT (session_id)
159
+ DO UPDATE SET
160
+ data = EXCLUDED.data,
161
+ expires_at = EXCLUDED.expires_at,
162
+ updated_at = CURRENT_TIMESTAMP
163
+ """
164
+
165
+ async with self._config.provide_connection() as conn:
166
+ await conn.execute(sql, [key, data, expires_at])
167
+
168
+ async def delete(self, key: str) -> None:
169
+ """Delete a session by key.
170
+
171
+ Args:
172
+ key: Session ID to delete.
173
+ """
174
+ sql = f"DELETE FROM {self._table_name} WHERE session_id = $1"
175
+
176
+ async with self._config.provide_connection() as conn:
177
+ await conn.execute(sql, [key])
178
+
179
+ async def delete_all(self) -> None:
180
+ """Delete all sessions from the store."""
181
+ sql = f"DELETE FROM {self._table_name}"
182
+
183
+ async with self._config.provide_connection() as conn:
184
+ await conn.execute(sql)
185
+ logger.debug("Deleted all sessions from table: %s", self._table_name)
186
+
187
+ async def exists(self, key: str) -> bool:
188
+ """Check if a session key exists and is not expired.
189
+
190
+ Args:
191
+ key: Session ID to check.
192
+
193
+ Returns:
194
+ True if the session exists and is not expired.
195
+
196
+ Notes:
197
+ Uses CURRENT_TIMESTAMP for consistency with get() method.
198
+ Uses fetch() instead of fetch_val() to handle zero-row case.
199
+ """
200
+ sql = f"""
201
+ SELECT 1 FROM {self._table_name}
202
+ WHERE session_id = $1
203
+ AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
204
+ """
205
+
206
+ async with self._config.provide_connection() as conn:
207
+ query_result = await conn.fetch(sql, [key])
208
+ rows = query_result.result()
209
+ return len(rows) > 0
210
+
211
+ async def expires_in(self, key: str) -> "int | None":
212
+ """Get the time in seconds until the session expires.
213
+
214
+ Args:
215
+ key: Session ID to check.
216
+
217
+ Returns:
218
+ Seconds until expiration, or None if no expiry or key doesn't exist.
219
+
220
+ Notes:
221
+ Uses fetch() to handle the case where the key doesn't exist.
222
+ """
223
+ sql = f"""
224
+ SELECT expires_at FROM {self._table_name}
225
+ WHERE session_id = $1
226
+ """
227
+
228
+ async with self._config.provide_connection() as conn:
229
+ query_result = await conn.fetch(sql, [key])
230
+ rows = query_result.result()
231
+
232
+ if not rows:
233
+ return None
234
+
235
+ expires_at = rows[0]["expires_at"]
236
+
237
+ if expires_at is None:
238
+ return None
239
+
240
+ now = datetime.now(timezone.utc)
241
+ if expires_at <= now:
242
+ return 0
243
+
244
+ delta = expires_at - now
245
+ return int(delta.total_seconds())
246
+
247
+ async def delete_expired(self) -> int:
248
+ """Delete all expired sessions.
249
+
250
+ Returns:
251
+ Number of sessions deleted.
252
+
253
+ Notes:
254
+ Uses CURRENT_TIMESTAMP for consistency.
255
+ Uses RETURNING to get deleted row count since psqlpy QueryResult
256
+ doesn't expose command tags.
257
+ For very large tables (10M+ rows), consider batching deletes
258
+ to avoid holding locks too long.
259
+ """
260
+ sql = f"""
261
+ DELETE FROM {self._table_name}
262
+ WHERE expires_at <= CURRENT_TIMESTAMP
263
+ RETURNING session_id
264
+ """
265
+
266
+ async with self._config.provide_connection() as conn:
267
+ query_result = await conn.fetch(sql, [])
268
+ rows = query_result.result()
269
+ count = len(rows)
270
+ if count > 0:
271
+ logger.debug("Cleaned up %d expired sessions", count)
272
+ return count
@@ -6,11 +6,11 @@ backward compatibility.
6
6
  """
7
7
 
8
8
  import re
9
- from typing import Any, Final, Optional
9
+ from functools import lru_cache
10
+ from typing import Any, Final
10
11
 
11
12
  from sqlspec.core.type_conversion import BaseTypeConverter
12
13
 
13
- # PostgreSQL-specific regex patterns for types not covered by base BaseTypeConverter
14
14
  PG_SPECIFIC_REGEX: Final[re.Pattern[str]] = re.compile(
15
15
  r"^(?:"
16
16
  r"(?P<interval>(?:(?:\d+\s+(?:year|month|day|hour|minute|second)s?\s*)+)|(?:P(?:\d+Y)?(?:\d+M)?(?:\d+D)?(?:T(?:\d+H)?(?:\d+M)?(?:\d+(?:\.\d+)?S)?)?))|"
@@ -19,17 +19,50 @@ PG_SPECIFIC_REGEX: Final[re.Pattern[str]] = re.compile(
19
19
  re.IGNORECASE,
20
20
  )
21
21
 
22
+ PG_SPECIAL_CHARS: Final[frozenset[str]] = frozenset({"{", "-", ":", "T", ".", "P", "[", "Y", "M", "D", "H", "S"})
23
+
22
24
 
23
25
  class PostgreSQLTypeConverter(BaseTypeConverter):
24
26
  """PostgreSQL-specific type converter with interval and array support.
25
27
 
26
28
  Extends the base BaseTypeConverter with PostgreSQL-specific functionality
27
29
  while maintaining backward compatibility for interval and array types.
30
+ Includes per-instance LRU cache for improved performance.
28
31
  """
29
32
 
30
- __slots__ = ()
33
+ __slots__ = ("_convert_cache",)
34
+
35
+ def __init__(self, cache_size: int = 5000) -> None:
36
+ """Initialize converter with per-instance conversion cache.
37
+
38
+ Args:
39
+ cache_size: Maximum number of string values to cache (default: 5000)
40
+ """
41
+ super().__init__()
42
+
43
+ @lru_cache(maxsize=cache_size)
44
+ def _cached_convert(value: str) -> Any:
45
+ if not value or not any(c in value for c in PG_SPECIAL_CHARS):
46
+ return value
47
+ detected_type = self.detect_type(value)
48
+ return self.convert_value(value, detected_type) if detected_type else value
49
+
50
+ self._convert_cache = _cached_convert
51
+
52
+ def convert_if_detected(self, value: Any) -> Any:
53
+ """Convert string if special type detected (cached).
54
+
55
+ Args:
56
+ value: Value to potentially convert
57
+
58
+ Returns:
59
+ Converted value or original value
60
+ """
61
+ if not isinstance(value, str):
62
+ return value
63
+ return self._convert_cache(value)
31
64
 
32
- def detect_type(self, value: str) -> Optional[str]:
65
+ def detect_type(self, value: str) -> str | None:
33
66
  """Detect types including PostgreSQL-specific types.
34
67
 
35
68
  Args:
@@ -38,12 +71,10 @@ class PostgreSQLTypeConverter(BaseTypeConverter):
38
71
  Returns:
39
72
  Type name if detected, None otherwise.
40
73
  """
41
- # First try generic types (UUID, JSON, datetime, etc.)
42
74
  detected_type = super().detect_type(value)
43
75
  if detected_type:
44
76
  return detected_type
45
77
 
46
- # Then check PostgreSQL-specific types
47
78
  match = PG_SPECIFIC_REGEX.match(value)
48
79
  if match:
49
80
  for group_name in ["interval", "pg_array"]:
@@ -62,12 +93,10 @@ class PostgreSQLTypeConverter(BaseTypeConverter):
62
93
  Returns:
63
94
  Converted value or original string for PostgreSQL-specific types.
64
95
  """
65
- # For PostgreSQL-specific types, preserve as strings for backward compatibility
66
- if detected_type in ("interval", "pg_array"):
67
- return value # Pass through as strings - psqlpy will handle casting
96
+ if detected_type in {"interval", "pg_array"}:
97
+ return value
68
98
 
69
- # Use base converter for standard types
70
99
  return super().convert_value(value, detected_type)
71
100
 
72
101
 
73
- __all__ = ("PG_SPECIFIC_REGEX", "PostgreSQLTypeConverter")
102
+ __all__ = ("PG_SPECIAL_CHARS", "PG_SPECIFIC_REGEX", "PostgreSQLTypeConverter")