sqlspec 0.25.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of sqlspec might be problematic. Click here for more details.

Files changed (199) hide show
  1. sqlspec/__init__.py +7 -15
  2. sqlspec/_serialization.py +256 -24
  3. sqlspec/_typing.py +71 -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 +870 -0
  7. sqlspec/adapters/adbc/config.py +69 -12
  8. sqlspec/adapters/adbc/data_dictionary.py +340 -0
  9. sqlspec/adapters/adbc/driver.py +266 -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 +153 -0
  13. sqlspec/adapters/aiosqlite/_types.py +1 -1
  14. sqlspec/adapters/aiosqlite/adk/__init__.py +5 -0
  15. sqlspec/adapters/aiosqlite/adk/store.py +527 -0
  16. sqlspec/adapters/aiosqlite/config.py +88 -15
  17. sqlspec/adapters/aiosqlite/data_dictionary.py +149 -0
  18. sqlspec/adapters/aiosqlite/driver.py +143 -40
  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 +2 -2
  24. sqlspec/adapters/asyncmy/adk/__init__.py +5 -0
  25. sqlspec/adapters/asyncmy/adk/store.py +493 -0
  26. sqlspec/adapters/asyncmy/config.py +68 -23
  27. sqlspec/adapters/asyncmy/data_dictionary.py +161 -0
  28. sqlspec/adapters/asyncmy/driver.py +313 -58
  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 +450 -0
  36. sqlspec/adapters/asyncpg/config.py +59 -35
  37. sqlspec/adapters/asyncpg/data_dictionary.py +173 -0
  38. sqlspec/adapters/asyncpg/driver.py +170 -25
  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 +576 -0
  44. sqlspec/adapters/bigquery/config.py +27 -10
  45. sqlspec/adapters/bigquery/data_dictionary.py +149 -0
  46. sqlspec/adapters/bigquery/driver.py +368 -142
  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 +125 -0
  50. sqlspec/adapters/duckdb/_types.py +1 -1
  51. sqlspec/adapters/duckdb/adk/__init__.py +14 -0
  52. sqlspec/adapters/duckdb/adk/store.py +553 -0
  53. sqlspec/adapters/duckdb/config.py +80 -20
  54. sqlspec/adapters/duckdb/data_dictionary.py +163 -0
  55. sqlspec/adapters/duckdb/driver.py +167 -45
  56. sqlspec/adapters/duckdb/litestar/__init__.py +5 -0
  57. sqlspec/adapters/duckdb/litestar/store.py +332 -0
  58. sqlspec/adapters/duckdb/pool.py +4 -4
  59. sqlspec/adapters/duckdb/type_converter.py +133 -0
  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 +1745 -0
  64. sqlspec/adapters/oracledb/config.py +122 -32
  65. sqlspec/adapters/oracledb/data_dictionary.py +509 -0
  66. sqlspec/adapters/oracledb/driver.py +353 -91
  67. sqlspec/adapters/oracledb/litestar/__init__.py +5 -0
  68. sqlspec/adapters/oracledb/litestar/store.py +767 -0
  69. sqlspec/adapters/oracledb/migrations.py +348 -73
  70. sqlspec/adapters/oracledb/type_converter.py +207 -0
  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 +482 -0
  75. sqlspec/adapters/psqlpy/config.py +46 -17
  76. sqlspec/adapters/psqlpy/data_dictionary.py +172 -0
  77. sqlspec/adapters/psqlpy/driver.py +123 -209
  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 +102 -0
  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 +944 -0
  85. sqlspec/adapters/psycopg/config.py +69 -35
  86. sqlspec/adapters/psycopg/data_dictionary.py +331 -0
  87. sqlspec/adapters/psycopg/driver.py +238 -81
  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 +572 -0
  95. sqlspec/adapters/sqlite/config.py +87 -15
  96. sqlspec/adapters/sqlite/data_dictionary.py +149 -0
  97. sqlspec/adapters/sqlite/driver.py +137 -54
  98. sqlspec/adapters/sqlite/litestar/__init__.py +5 -0
  99. sqlspec/adapters/sqlite/litestar/store.py +318 -0
  100. sqlspec/adapters/sqlite/pool.py +18 -9
  101. sqlspec/base.py +45 -26
  102. sqlspec/builder/__init__.py +73 -4
  103. sqlspec/builder/_base.py +162 -89
  104. sqlspec/builder/_column.py +62 -29
  105. sqlspec/builder/_ddl.py +180 -121
  106. sqlspec/builder/_delete.py +5 -4
  107. sqlspec/builder/_dml.py +388 -0
  108. sqlspec/{_sql.py → builder/_factory.py} +53 -94
  109. sqlspec/builder/_insert.py +32 -131
  110. sqlspec/builder/_join.py +375 -0
  111. sqlspec/builder/_merge.py +446 -11
  112. sqlspec/builder/_parsing_utils.py +111 -17
  113. sqlspec/builder/_select.py +1457 -24
  114. sqlspec/builder/_update.py +11 -42
  115. sqlspec/cli.py +307 -194
  116. sqlspec/config.py +252 -67
  117. sqlspec/core/__init__.py +5 -4
  118. sqlspec/core/cache.py +17 -17
  119. sqlspec/core/compiler.py +62 -9
  120. sqlspec/core/filters.py +37 -37
  121. sqlspec/core/hashing.py +9 -9
  122. sqlspec/core/parameters.py +83 -48
  123. sqlspec/core/result.py +102 -46
  124. sqlspec/core/splitter.py +16 -17
  125. sqlspec/core/statement.py +36 -30
  126. sqlspec/core/type_conversion.py +235 -0
  127. sqlspec/driver/__init__.py +7 -6
  128. sqlspec/driver/_async.py +188 -151
  129. sqlspec/driver/_common.py +285 -80
  130. sqlspec/driver/_sync.py +188 -152
  131. sqlspec/driver/mixins/_result_tools.py +20 -236
  132. sqlspec/driver/mixins/_sql_translator.py +4 -4
  133. sqlspec/exceptions.py +75 -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 +73 -53
  142. sqlspec/extensions/litestar/__init__.py +21 -4
  143. sqlspec/extensions/litestar/cli.py +54 -10
  144. sqlspec/extensions/litestar/config.py +59 -266
  145. sqlspec/extensions/litestar/handlers.py +46 -17
  146. sqlspec/extensions/litestar/migrations/0001_create_session_table.py +137 -0
  147. sqlspec/extensions/litestar/migrations/__init__.py +3 -0
  148. sqlspec/extensions/litestar/plugin.py +324 -223
  149. sqlspec/extensions/litestar/providers.py +25 -25
  150. sqlspec/extensions/litestar/store.py +265 -0
  151. sqlspec/loader.py +30 -49
  152. sqlspec/migrations/__init__.py +4 -3
  153. sqlspec/migrations/base.py +302 -39
  154. sqlspec/migrations/commands.py +611 -144
  155. sqlspec/migrations/context.py +142 -0
  156. sqlspec/migrations/fix.py +199 -0
  157. sqlspec/migrations/loaders.py +68 -23
  158. sqlspec/migrations/runner.py +543 -107
  159. sqlspec/migrations/tracker.py +237 -21
  160. sqlspec/migrations/utils.py +51 -3
  161. sqlspec/migrations/validation.py +177 -0
  162. sqlspec/protocols.py +66 -36
  163. sqlspec/storage/_utils.py +98 -0
  164. sqlspec/storage/backends/fsspec.py +134 -106
  165. sqlspec/storage/backends/local.py +78 -51
  166. sqlspec/storage/backends/obstore.py +278 -162
  167. sqlspec/storage/registry.py +75 -39
  168. sqlspec/typing.py +16 -84
  169. sqlspec/utils/config_resolver.py +153 -0
  170. sqlspec/utils/correlation.py +4 -5
  171. sqlspec/utils/data_transformation.py +3 -2
  172. sqlspec/utils/deprecation.py +9 -8
  173. sqlspec/utils/fixtures.py +4 -4
  174. sqlspec/utils/logging.py +46 -6
  175. sqlspec/utils/module_loader.py +2 -2
  176. sqlspec/utils/schema.py +288 -0
  177. sqlspec/utils/serializers.py +50 -2
  178. sqlspec/utils/sync_tools.py +21 -17
  179. sqlspec/utils/text.py +1 -2
  180. sqlspec/utils/type_guards.py +111 -20
  181. sqlspec/utils/version.py +433 -0
  182. {sqlspec-0.25.0.dist-info → sqlspec-0.27.0.dist-info}/METADATA +40 -21
  183. sqlspec-0.27.0.dist-info/RECORD +207 -0
  184. sqlspec/builder/mixins/__init__.py +0 -55
  185. sqlspec/builder/mixins/_cte_and_set_ops.py +0 -254
  186. sqlspec/builder/mixins/_delete_operations.py +0 -50
  187. sqlspec/builder/mixins/_insert_operations.py +0 -282
  188. sqlspec/builder/mixins/_join_operations.py +0 -389
  189. sqlspec/builder/mixins/_merge_operations.py +0 -592
  190. sqlspec/builder/mixins/_order_limit_operations.py +0 -152
  191. sqlspec/builder/mixins/_pivot_operations.py +0 -157
  192. sqlspec/builder/mixins/_select_operations.py +0 -936
  193. sqlspec/builder/mixins/_update_operations.py +0 -218
  194. sqlspec/builder/mixins/_where_clause.py +0 -1304
  195. sqlspec-0.25.0.dist-info/RECORD +0 -139
  196. sqlspec-0.25.0.dist-info/licenses/NOTICE +0 -29
  197. {sqlspec-0.25.0.dist-info → sqlspec-0.27.0.dist-info}/WHEEL +0 -0
  198. {sqlspec-0.25.0.dist-info → sqlspec-0.27.0.dist-info}/entry_points.txt +0 -0
  199. {sqlspec-0.25.0.dist-info → sqlspec-0.27.0.dist-info}/licenses/LICENSE +0 -0
sqlspec/__init__.py CHANGED
@@ -2,7 +2,6 @@
2
2
 
3
3
  from sqlspec import adapters, base, builder, core, driver, exceptions, extensions, loader, migrations, typing, utils
4
4
  from sqlspec.__metadata__ import __version__
5
- from sqlspec._sql import SQLFactory, sql
6
5
  from sqlspec.base import SQLSpec
7
6
  from sqlspec.builder import (
8
7
  Column,
@@ -15,7 +14,9 @@ from sqlspec.builder import (
15
14
  Merge,
16
15
  QueryBuilder,
17
16
  Select,
17
+ SQLFactory,
18
18
  Update,
19
+ sql,
19
20
  )
20
21
  from sqlspec.config import AsyncDatabaseConfig, SyncDatabaseConfig
21
22
  from sqlspec.core import (
@@ -34,16 +35,10 @@ from sqlspec.core import (
34
35
  from sqlspec.core import filters as filters
35
36
  from sqlspec.driver import AsyncDriverAdapterBase, ExecutionResult, SyncDriverAdapterBase
36
37
  from sqlspec.loader import SQLFile, SQLFileLoader
37
- from sqlspec.typing import (
38
- ConnectionT,
39
- DictRow,
40
- ModelDTOT,
41
- ModelT,
42
- PoolT,
43
- RowT,
44
- StatementParameters,
45
- SupportedSchemaModel,
46
- )
38
+ from sqlspec.typing import ConnectionT, PoolT, SchemaT, StatementParameters, SupportedSchemaModel
39
+ from sqlspec.utils.logging import suppress_erroneous_sqlglot_log_messages
40
+
41
+ suppress_erroneous_sqlglot_log_messages()
47
42
 
48
43
  __all__ = (
49
44
  "SQL",
@@ -57,26 +52,23 @@ __all__ = (
57
52
  "ConnectionT",
58
53
  "CreateTable",
59
54
  "Delete",
60
- "DictRow",
61
55
  "DropTable",
62
56
  "ExecutionResult",
63
57
  "FunctionColumn",
64
58
  "Insert",
65
59
  "Merge",
66
- "ModelDTOT",
67
- "ModelT",
68
60
  "ParameterConverter",
69
61
  "ParameterProcessor",
70
62
  "ParameterStyle",
71
63
  "ParameterStyleConfig",
72
64
  "PoolT",
73
65
  "QueryBuilder",
74
- "RowT",
75
66
  "SQLFactory",
76
67
  "SQLFile",
77
68
  "SQLFileLoader",
78
69
  "SQLResult",
79
70
  "SQLSpec",
71
+ "SchemaT",
80
72
  "Select",
81
73
  "Statement",
82
74
  "StatementConfig",
sqlspec/_serialization.py CHANGED
@@ -1,11 +1,37 @@
1
+ """Enhanced serialization module with byte-aware encoding and class-based architecture.
2
+
3
+ Provides a Protocol-based serialization system that users can extend.
4
+ Supports msgspec, orjson, and standard library JSON with automatic fallback.
5
+
6
+ Features optional numpy array serialization when numpy is installed.
7
+ Arrays are automatically converted to lists during JSON encoding.
8
+ """
9
+
10
+ import contextlib
1
11
  import datetime
2
12
  import enum
3
- from typing import Any
13
+ import json
14
+ from abc import ABC, abstractmethod
15
+ from typing import Any, Final, Literal, Protocol, overload
16
+
17
+ from sqlspec._typing import NUMPY_INSTALLED
18
+ from sqlspec.typing import MSGSPEC_INSTALLED, ORJSON_INSTALLED, PYDANTIC_INSTALLED, BaseModel
19
+
4
20
 
5
- from sqlspec.typing import PYDANTIC_INSTALLED, BaseModel
21
+ def _type_to_string(value: Any) -> Any: # pragma: no cover
22
+ """Convert special types to strings for JSON serialization.
6
23
 
24
+ Handles datetime, date, enums, Pydantic models, and numpy arrays.
7
25
 
8
- def _type_to_string(value: Any) -> str: # pragma: no cover
26
+ Args:
27
+ value: Value to convert.
28
+
29
+ Returns:
30
+ Serializable representation of the value (string, list, dict, etc.).
31
+
32
+ Raises:
33
+ TypeError: If value cannot be serialized.
34
+ """
9
35
  if isinstance(value, datetime.datetime):
10
36
  return convert_datetime_to_gmt_iso(value)
11
37
  if isinstance(value, datetime.date):
@@ -14,41 +40,233 @@ def _type_to_string(value: Any) -> str: # pragma: no cover
14
40
  return str(value.value)
15
41
  if PYDANTIC_INSTALLED and isinstance(value, BaseModel):
16
42
  return value.model_dump_json()
43
+ if NUMPY_INSTALLED:
44
+ import numpy as np
45
+
46
+ if isinstance(value, np.ndarray):
47
+ return value.tolist()
17
48
  try:
18
49
  return str(value)
19
50
  except Exception as exc:
20
- raise TypeError from exc
51
+ msg = f"Cannot serialize {type(value).__name__}"
52
+ raise TypeError(msg) from exc
21
53
 
22
54
 
23
- try:
24
- from msgspec.json import Decoder, Encoder
55
+ class JSONSerializer(Protocol):
56
+ """Protocol for JSON serialization implementations.
25
57
 
26
- encoder, decoder = Encoder(enc_hook=_type_to_string), Decoder()
27
- decode_json = decoder.decode
58
+ Users can implement this protocol to create custom serializers.
59
+ """
28
60
 
29
- def encode_json(data: Any) -> str: # pragma: no cover
30
- return encoder.encode(data).decode("utf-8")
61
+ def encode(self, data: Any, *, as_bytes: bool = False) -> str | bytes:
62
+ """Encode data to JSON.
31
63
 
32
- except ImportError:
33
- try:
34
- from orjson import ( # pyright: ignore[reportMissingImports]
64
+ Args:
65
+ data: Data to encode.
66
+ as_bytes: Whether to return bytes instead of string.
67
+
68
+ Returns:
69
+ JSON string or bytes depending on as_bytes parameter.
70
+ """
71
+ ...
72
+
73
+ def decode(self, data: str | bytes, *, decode_bytes: bool = True) -> Any:
74
+ """Decode from JSON.
75
+
76
+ Args:
77
+ data: JSON string or bytes to decode.
78
+ decode_bytes: Whether to decode bytes input.
79
+
80
+ Returns:
81
+ Decoded Python object.
82
+ """
83
+ ...
84
+
85
+
86
+ class BaseJSONSerializer(ABC):
87
+ """Base class for JSON serializers with common functionality."""
88
+
89
+ __slots__ = ()
90
+
91
+ @abstractmethod
92
+ def encode(self, data: Any, *, as_bytes: bool = False) -> str | bytes:
93
+ """Encode data to JSON."""
94
+ ...
95
+
96
+ @abstractmethod
97
+ def decode(self, data: str | bytes, *, decode_bytes: bool = True) -> Any:
98
+ """Decode from JSON."""
99
+ ...
100
+
101
+
102
+ class MsgspecSerializer(BaseJSONSerializer):
103
+ """Msgspec-based JSON serializer."""
104
+
105
+ __slots__ = ("_decoder", "_encoder")
106
+
107
+ def __init__(self) -> None:
108
+ """Initialize msgspec encoder and decoder."""
109
+ from msgspec.json import Decoder, Encoder
110
+
111
+ self._encoder: Final[Encoder] = Encoder(enc_hook=_type_to_string)
112
+ self._decoder: Final[Decoder] = Decoder()
113
+
114
+ def encode(self, data: Any, *, as_bytes: bool = False) -> str | bytes:
115
+ """Encode data using msgspec."""
116
+ try:
117
+ if as_bytes:
118
+ return self._encoder.encode(data)
119
+ return self._encoder.encode(data).decode("utf-8")
120
+ except (TypeError, ValueError):
121
+ if ORJSON_INSTALLED:
122
+ return OrjsonSerializer().encode(data, as_bytes=as_bytes)
123
+ return StandardLibSerializer().encode(data, as_bytes=as_bytes)
124
+
125
+ def decode(self, data: str | bytes, *, decode_bytes: bool = True) -> Any:
126
+ """Decode data using msgspec."""
127
+ if isinstance(data, bytes):
128
+ if decode_bytes:
129
+ try:
130
+ return self._decoder.decode(data)
131
+ except (TypeError, ValueError):
132
+ if ORJSON_INSTALLED:
133
+ return OrjsonSerializer().decode(data, decode_bytes=decode_bytes)
134
+ return StandardLibSerializer().decode(data, decode_bytes=decode_bytes)
135
+ return data
136
+
137
+ try:
138
+ return self._decoder.decode(data.encode("utf-8"))
139
+ except (TypeError, ValueError):
140
+ if ORJSON_INSTALLED:
141
+ return OrjsonSerializer().decode(data, decode_bytes=decode_bytes)
142
+ return StandardLibSerializer().decode(data, decode_bytes=decode_bytes)
143
+
144
+
145
+ class OrjsonSerializer(BaseJSONSerializer):
146
+ """Orjson-based JSON serializer with native datetime/UUID support.
147
+
148
+ Automatically enables numpy serialization if numpy is installed.
149
+ """
150
+
151
+ __slots__ = ()
152
+
153
+ def encode(self, data: Any, *, as_bytes: bool = False) -> str | bytes:
154
+ """Encode data using orjson.
155
+
156
+ Args:
157
+ data: Data to encode.
158
+ as_bytes: Whether to return bytes instead of string.
159
+
160
+ Returns:
161
+ JSON string or bytes depending on as_bytes parameter.
162
+ """
163
+ from orjson import (
35
164
  OPT_NAIVE_UTC, # pyright: ignore[reportUnknownVariableType]
36
- OPT_SERIALIZE_NUMPY, # pyright: ignore[reportUnknownVariableType]
37
165
  OPT_SERIALIZE_UUID, # pyright: ignore[reportUnknownVariableType]
38
166
  )
39
- from orjson import dumps as _encode_json # pyright: ignore[reportUnknownVariableType,reportMissingImports]
40
- from orjson import loads as decode_json # type: ignore[no-redef,assignment,unused-ignore]
167
+ from orjson import dumps as _orjson_dumps # pyright: ignore[reportMissingImports]
168
+
169
+ options = OPT_NAIVE_UTC | OPT_SERIALIZE_UUID
41
170
 
42
- def encode_json(data: Any) -> str: # pragma: no cover
43
- return _encode_json(
44
- data, default=_type_to_string, option=OPT_SERIALIZE_NUMPY | OPT_NAIVE_UTC | OPT_SERIALIZE_UUID
45
- ).decode("utf-8")
171
+ if NUMPY_INSTALLED:
172
+ from orjson import OPT_SERIALIZE_NUMPY # pyright: ignore[reportUnknownVariableType]
46
173
 
47
- except ImportError:
48
- from json import dumps as encode_json # type: ignore[assignment]
49
- from json import loads as decode_json # type: ignore[assignment]
174
+ options |= OPT_SERIALIZE_NUMPY
50
175
 
51
- __all__ = ("convert_date_to_iso", "convert_datetime_to_gmt_iso", "decode_json", "encode_json")
176
+ result = _orjson_dumps(data, default=_type_to_string, option=options)
177
+ return result if as_bytes else result.decode("utf-8")
178
+
179
+ def decode(self, data: str | bytes, *, decode_bytes: bool = True) -> Any:
180
+ """Decode data using orjson."""
181
+ from orjson import loads as _orjson_loads # pyright: ignore[reportMissingImports]
182
+
183
+ if isinstance(data, bytes):
184
+ if decode_bytes:
185
+ return _orjson_loads(data)
186
+ return data
187
+ return _orjson_loads(data)
188
+
189
+
190
+ class StandardLibSerializer(BaseJSONSerializer):
191
+ """Standard library JSON serializer as fallback."""
192
+
193
+ __slots__ = ()
194
+
195
+ def encode(self, data: Any, *, as_bytes: bool = False) -> str | bytes:
196
+ """Encode data using standard library json."""
197
+ json_str = json.dumps(data, default=_type_to_string)
198
+ return json_str.encode("utf-8") if as_bytes else json_str
199
+
200
+ def decode(self, data: str | bytes, *, decode_bytes: bool = True) -> Any:
201
+ """Decode data using standard library json."""
202
+ if isinstance(data, bytes):
203
+ if decode_bytes:
204
+ return json.loads(data.decode("utf-8"))
205
+ return data
206
+ return json.loads(data)
207
+
208
+
209
+ _default_serializer: JSONSerializer | None = None
210
+
211
+
212
+ def get_default_serializer() -> JSONSerializer:
213
+ """Get the default serializer based on available libraries.
214
+
215
+ Priority: msgspec > orjson > stdlib
216
+
217
+ Returns:
218
+ The best available JSON serializer.
219
+ """
220
+ global _default_serializer
221
+
222
+ if _default_serializer is None:
223
+ if MSGSPEC_INSTALLED:
224
+ with contextlib.suppress(ImportError):
225
+ _default_serializer = MsgspecSerializer()
226
+
227
+ if _default_serializer is None and ORJSON_INSTALLED:
228
+ with contextlib.suppress(ImportError):
229
+ _default_serializer = OrjsonSerializer()
230
+
231
+ if _default_serializer is None:
232
+ _default_serializer = StandardLibSerializer()
233
+
234
+ assert _default_serializer is not None
235
+ return _default_serializer
236
+
237
+
238
+ @overload
239
+ def encode_json(data: Any, *, as_bytes: Literal[False] = ...) -> str: ... # pragma: no cover
240
+
241
+
242
+ @overload
243
+ def encode_json(data: Any, *, as_bytes: Literal[True]) -> bytes: ... # pragma: no cover
244
+
245
+
246
+ def encode_json(data: Any, *, as_bytes: bool = False) -> str | bytes:
247
+ """Encode to JSON, optionally returning bytes.
248
+
249
+ Args:
250
+ data: The data to encode.
251
+ as_bytes: Whether to return bytes instead of string.
252
+
253
+ Returns:
254
+ JSON string or bytes depending on as_bytes parameter.
255
+ """
256
+ return get_default_serializer().encode(data, as_bytes=as_bytes)
257
+
258
+
259
+ def decode_json(data: str | bytes, *, decode_bytes: bool = True) -> Any:
260
+ """Decode from JSON string or bytes efficiently.
261
+
262
+ Args:
263
+ data: JSON string or bytes to decode.
264
+ decode_bytes: Whether to decode bytes input.
265
+
266
+ Returns:
267
+ Decoded Python object.
268
+ """
269
+ return get_default_serializer().decode(data, decode_bytes=decode_bytes)
52
270
 
53
271
 
54
272
  def convert_datetime_to_gmt_iso(dt: datetime.datetime) -> str: # pragma: no cover
@@ -75,3 +293,17 @@ def convert_date_to_iso(dt: datetime.date) -> str: # pragma: no cover
75
293
  The ISO formatted date string.
76
294
  """
77
295
  return dt.isoformat()
296
+
297
+
298
+ __all__ = (
299
+ "BaseJSONSerializer",
300
+ "JSONSerializer",
301
+ "MsgspecSerializer",
302
+ "OrjsonSerializer",
303
+ "StandardLibSerializer",
304
+ "convert_date_to_iso",
305
+ "convert_datetime_to_gmt_iso",
306
+ "decode_json",
307
+ "encode_json",
308
+ "get_default_serializer",
309
+ )
sqlspec/_typing.py CHANGED
@@ -6,9 +6,9 @@ from collections.abc import Iterable, Mapping
6
6
  from dataclasses import dataclass
7
7
  from enum import Enum
8
8
  from importlib.util import find_spec
9
- from typing import Any, ClassVar, Final, Optional, Protocol, Union, cast, runtime_checkable
9
+ from typing import Any, ClassVar, Final, Literal, Protocol, cast, runtime_checkable
10
10
 
11
- from typing_extensions import Literal, TypeVar, dataclass_transform
11
+ from typing_extensions import TypeVar, dataclass_transform
12
12
 
13
13
 
14
14
  @runtime_checkable
@@ -38,15 +38,15 @@ class BaseModelStub:
38
38
  self,
39
39
  /,
40
40
  *,
41
- include: "Optional[Any]" = None, # noqa: ARG002
42
- exclude: "Optional[Any]" = None, # noqa: ARG002
43
- context: "Optional[Any]" = None, # noqa: ARG002
41
+ include: "Any | None" = None, # noqa: ARG002
42
+ exclude: "Any | None" = None, # noqa: ARG002
43
+ context: "Any | None" = None, # noqa: ARG002
44
44
  by_alias: bool = False, # noqa: ARG002
45
45
  exclude_unset: bool = False, # noqa: ARG002
46
46
  exclude_defaults: bool = False, # noqa: ARG002
47
47
  exclude_none: bool = False, # noqa: ARG002
48
48
  round_trip: bool = False, # noqa: ARG002
49
- warnings: "Union[bool, Literal['none', 'warn', 'error']]" = True, # noqa: ARG002
49
+ warnings: "bool | Literal['none', 'warn', 'error']" = True, # noqa: ARG002
50
50
  serialize_as_any: bool = False, # noqa: ARG002
51
51
  ) -> "dict[str, Any]":
52
52
  """Placeholder implementation."""
@@ -56,15 +56,15 @@ class BaseModelStub:
56
56
  self,
57
57
  /,
58
58
  *,
59
- include: "Optional[Any]" = None, # noqa: ARG002
60
- exclude: "Optional[Any]" = None, # noqa: ARG002
61
- context: "Optional[Any]" = None, # noqa: ARG002
59
+ include: "Any | None" = None, # noqa: ARG002
60
+ exclude: "Any | None" = None, # noqa: ARG002
61
+ context: "Any | None" = None, # noqa: ARG002
62
62
  by_alias: bool = False, # noqa: ARG002
63
63
  exclude_unset: bool = False, # noqa: ARG002
64
64
  exclude_defaults: bool = False, # noqa: ARG002
65
65
  exclude_none: bool = False, # noqa: ARG002
66
66
  round_trip: bool = False, # noqa: ARG002
67
- warnings: "Union[bool, Literal['none', 'warn', 'error']]" = True, # noqa: ARG002
67
+ warnings: "bool | Literal['none', 'warn', 'error']" = True, # noqa: ARG002
68
68
  serialize_as_any: bool = False, # noqa: ARG002
69
69
  ) -> str:
70
70
  """Placeholder implementation."""
@@ -78,9 +78,9 @@ class TypeAdapterStub:
78
78
  self,
79
79
  type: Any, # noqa: A002
80
80
  *,
81
- config: "Optional[Any]" = None, # noqa: ARG002
81
+ config: "Any | None" = None, # noqa: ARG002
82
82
  _parent_depth: int = 2, # noqa: ARG002
83
- module: "Optional[str]" = None, # noqa: ARG002
83
+ module: "str | None" = None, # noqa: ARG002
84
84
  ) -> None:
85
85
  """Initialize."""
86
86
  self._type = type
@@ -90,10 +90,10 @@ class TypeAdapterStub:
90
90
  object: Any,
91
91
  /,
92
92
  *,
93
- strict: "Optional[bool]" = None, # noqa: ARG002
94
- from_attributes: "Optional[bool]" = None, # noqa: ARG002
95
- context: "Optional[dict[str, Any]]" = None, # noqa: ARG002
96
- experimental_allow_partial: "Union[bool, Literal['off', 'on', 'trailing-strings']]" = False, # noqa: ARG002
93
+ strict: "bool | None" = None, # noqa: ARG002
94
+ from_attributes: "bool | None" = None, # noqa: ARG002
95
+ context: "dict[str, Any] | None" = None, # noqa: ARG002
96
+ experimental_allow_partial: "bool | Literal['off', 'on', 'trailing-strings']" = False, # noqa: ARG002
97
97
  ) -> Any:
98
98
  """Validate Python object."""
99
99
  return object
@@ -143,8 +143,8 @@ def convert_stub( # noqa: PLR0913
143
143
  *,
144
144
  strict: bool = True, # noqa: ARG001
145
145
  from_attributes: bool = False, # noqa: ARG001
146
- dec_hook: "Optional[Any]" = None, # noqa: ARG001
147
- builtin_types: "Optional[Any]" = None, # noqa: ARG001
146
+ dec_hook: "Any | None" = None, # noqa: ARG001
147
+ builtin_types: "Any | None" = None, # noqa: ARG001
148
148
  str_keys: bool = False, # noqa: ARG001
149
149
  ) -> Any:
150
150
  """Placeholder implementation."""
@@ -177,6 +177,14 @@ except ImportError:
177
177
  MSGSPEC_INSTALLED = False # pyright: ignore[reportConstantRedefinition]
178
178
 
179
179
 
180
+ try:
181
+ import orjson # noqa: F401
182
+
183
+ ORJSON_INSTALLED = True # pyright: ignore[reportConstantRedefinition]
184
+ except ImportError:
185
+ ORJSON_INSTALLED = False # pyright: ignore[reportConstantRedefinition]
186
+
187
+
180
188
  # Always define stub type for DTOData
181
189
  @runtime_checkable
182
190
  class DTODataStub(Protocol[T]):
@@ -300,7 +308,7 @@ class EmptyEnum(Enum):
300
308
  EMPTY = 0
301
309
 
302
310
 
303
- EmptyType = Union[Literal[EmptyEnum.EMPTY], UnsetType]
311
+ EmptyType = Literal[EmptyEnum.EMPTY] | UnsetType
304
312
  Empty: Final = EmptyEnum.EMPTY
305
313
 
306
314
 
@@ -328,18 +336,18 @@ class ArrowTableResult(Protocol):
328
336
  def from_arrays(
329
337
  self,
330
338
  arrays: list[Any],
331
- names: "Optional[list[str]]" = None,
332
- schema: "Optional[Any]" = None,
333
- metadata: "Optional[Mapping[str, Any]]" = None,
339
+ names: "list[str] | None" = None,
340
+ schema: "Any | None" = None,
341
+ metadata: "Mapping[str, Any] | None" = None,
334
342
  ) -> Any:
335
343
  return None
336
344
 
337
345
  def from_pydict(
338
- self, mapping: dict[str, Any], schema: "Optional[Any]" = None, metadata: "Optional[Mapping[str, Any]]" = None
346
+ self, mapping: dict[str, Any], schema: "Any | None" = None, metadata: "Mapping[str, Any] | None" = None
339
347
  ) -> Any:
340
348
  return None
341
349
 
342
- def from_batches(self, batches: Iterable[Any], schema: Optional[Any] = None) -> Any:
350
+ def from_batches(self, batches: Iterable[Any], schema: Any | None = None) -> Any:
343
351
  return None
344
352
 
345
353
 
@@ -365,7 +373,7 @@ class ArrowRecordBatchResult(Protocol):
365
373
  def column(self, i: int) -> Any:
366
374
  return None
367
375
 
368
- def slice(self, offset: int = 0, length: "Optional[int]" = None) -> Any:
376
+ def slice(self, offset: int = 0, length: "int | None" = None) -> Any:
369
377
  return None
370
378
 
371
379
 
@@ -381,6 +389,24 @@ except ImportError:
381
389
  PYARROW_INSTALLED = False # pyright: ignore[reportConstantRedefinition]
382
390
 
383
391
 
392
+ @runtime_checkable
393
+ class NumpyArrayStub(Protocol):
394
+ """Protocol stub for numpy.ndarray when numpy is not installed.
395
+
396
+ Provides minimal interface for type checking and serialization support.
397
+ """
398
+
399
+ def tolist(self) -> "list[Any]":
400
+ """Convert array to Python list."""
401
+ ...
402
+
403
+
404
+ try:
405
+ from numpy import ndarray as NumpyArray # noqa: N812
406
+ except ImportError:
407
+ NumpyArray = NumpyArrayStub # type: ignore[assignment,misc]
408
+
409
+
384
410
  try:
385
411
  from opentelemetry import trace # pyright: ignore[reportMissingImports, reportAssignmentType]
386
412
  from opentelemetry.trace import ( # pyright: ignore[reportMissingImports, reportAssignmentType]
@@ -401,16 +427,16 @@ except ImportError:
401
427
  def record_exception(
402
428
  self,
403
429
  exception: "Exception",
404
- attributes: "Optional[Mapping[str, Any]]" = None,
405
- timestamp: "Optional[int]" = None,
430
+ attributes: "Mapping[str, Any] | None" = None,
431
+ timestamp: "int | None" = None,
406
432
  escaped: bool = False,
407
433
  ) -> None:
408
434
  return None
409
435
 
410
- def set_status(self, status: Any, description: "Optional[str]" = None) -> None:
436
+ def set_status(self, status: Any, description: "str | None" = None) -> None:
411
437
  return None
412
438
 
413
- def end(self, end_time: "Optional[int]" = None) -> None:
439
+ def end(self, end_time: "int | None" = None) -> None:
414
440
  return None
415
441
 
416
442
  def __enter__(self) -> "Span":
@@ -437,8 +463,8 @@ except ImportError:
437
463
  def get_tracer(
438
464
  self,
439
465
  instrumenting_module_name: str,
440
- instrumenting_library_version: "Optional[str]" = None,
441
- schema_url: "Optional[str]" = None,
466
+ instrumenting_library_version: "str | None" = None,
467
+ schema_url: "str | None" = None,
442
468
  tracer_provider: Any = None,
443
469
  ) -> Tracer:
444
470
  return Tracer() # type: ignore[abstract] # pragma: no cover
@@ -514,9 +540,6 @@ try:
514
540
  from aiosql.types import ( # pyright: ignore[reportMissingImports, reportAssignmentType]
515
541
  AsyncDriverAdapterProtocol as AiosqlAsyncProtocol, # pyright: ignore[reportMissingImports, reportAssignmentType]
516
542
  )
517
- from aiosql.types import ( # pyright: ignore[reportMissingImports, reportAssignmentType]
518
- DriverAdapterProtocol as AiosqlProtocol, # pyright: ignore[reportMissingImports, reportAssignmentType]
519
- )
520
543
  from aiosql.types import ParamType as AiosqlParamType # pyright: ignore[reportMissingImports, reportAssignmentType]
521
544
  from aiosql.types import (
522
545
  SQLOperationType as AiosqlSQLOperationType, # pyright: ignore[reportMissingImports, reportAssignmentType]
@@ -545,7 +568,7 @@ except ImportError:
545
568
  aiosql = _AiosqlShim() # type: ignore[assignment]
546
569
 
547
570
  # Placeholder types for aiosql protocols
548
- AiosqlParamType = Union[dict[str, Any], list[Any], tuple[Any, ...], None] # type: ignore[misc]
571
+ AiosqlParamType = Any # type: ignore[misc]
549
572
 
550
573
  class AiosqlSQLOperationType(Enum): # type: ignore[no-redef]
551
574
  """Enumeration of aiosql operation types."""
@@ -558,12 +581,6 @@ except ImportError:
558
581
  SELECT_ONE = 5
559
582
  SELECT_VALUE = 6
560
583
 
561
- @runtime_checkable
562
- class AiosqlProtocol(Protocol): # type: ignore[no-redef]
563
- """Placeholder for aiosql DriverAdapterProtocol"""
564
-
565
- def process_sql(self, query_name: str, op_type: Any, sql: str) -> str: ...
566
-
567
584
  @runtime_checkable
568
585
  class AiosqlSyncProtocol(Protocol): # type: ignore[no-redef]
569
586
  """Placeholder for aiosql SyncDriverAdapterProtocol"""
@@ -572,16 +589,16 @@ except ImportError:
572
589
 
573
590
  def process_sql(self, query_name: str, op_type: Any, sql: str) -> str: ...
574
591
  def select(
575
- self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Optional[Any]" = None
592
+ self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Any | None" = None
576
593
  ) -> Any: ...
577
594
  def select_one(
578
- self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Optional[Any]" = None
579
- ) -> "Optional[Any]": ...
580
- def select_value(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Optional[Any]": ...
595
+ self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Any | None" = None
596
+ ) -> "Any | None": ...
597
+ def select_value(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Any | None": ...
581
598
  def select_cursor(self, conn: Any, query_name: str, sql: str, parameters: Any) -> Any: ...
582
599
  def insert_update_delete(self, conn: Any, query_name: str, sql: str, parameters: Any) -> int: ...
583
600
  def insert_update_delete_many(self, conn: Any, query_name: str, sql: str, parameters: Any) -> int: ...
584
- def insert_returning(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Optional[Any]": ...
601
+ def insert_returning(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Any | None": ...
585
602
 
586
603
  @runtime_checkable
587
604
  class AiosqlAsyncProtocol(Protocol): # type: ignore[no-redef]
@@ -591,16 +608,16 @@ except ImportError:
591
608
 
592
609
  def process_sql(self, query_name: str, op_type: Any, sql: str) -> str: ...
593
610
  async def select(
594
- self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Optional[Any]" = None
611
+ self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Any | None" = None
595
612
  ) -> Any: ...
596
613
  async def select_one(
597
- self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Optional[Any]" = None
598
- ) -> "Optional[Any]": ...
599
- async def select_value(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Optional[Any]": ...
614
+ self, conn: Any, query_name: str, sql: str, parameters: Any, record_class: "Any | None" = None
615
+ ) -> "Any | None": ...
616
+ async def select_value(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Any | None": ...
600
617
  async def select_cursor(self, conn: Any, query_name: str, sql: str, parameters: Any) -> Any: ...
601
618
  async def insert_update_delete(self, conn: Any, query_name: str, sql: str, parameters: Any) -> None: ...
602
619
  async def insert_update_delete_many(self, conn: Any, query_name: str, sql: str, parameters: Any) -> None: ...
603
- async def insert_returning(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Optional[Any]": ...
620
+ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameters: Any) -> "Any | None": ...
604
621
 
605
622
  AIOSQL_INSTALLED = False # pyright: ignore[reportConstantRedefinition] # pyright: ignore[reportConstantRedefinition]
606
623
 
@@ -621,6 +638,7 @@ __all__ = (
621
638
  "NUMPY_INSTALLED",
622
639
  "OBSTORE_INSTALLED",
623
640
  "OPENTELEMETRY_INSTALLED",
641
+ "ORJSON_INSTALLED",
624
642
  "PGVECTOR_INSTALLED",
625
643
  "PROMETHEUS_INSTALLED",
626
644
  "PYARROW_INSTALLED",
@@ -629,7 +647,6 @@ __all__ = (
629
647
  "UNSET_STUB",
630
648
  "AiosqlAsyncProtocol",
631
649
  "AiosqlParamType",
632
- "AiosqlProtocol",
633
650
  "AiosqlSQLOperationType",
634
651
  "AiosqlSyncProtocol",
635
652
  "ArrowRecordBatch",
@@ -651,6 +668,8 @@ __all__ = (
651
668
  "FailFastStub",
652
669
  "Gauge",
653
670
  "Histogram",
671
+ "NumpyArray",
672
+ "NumpyArrayStub",
654
673
  "Span",
655
674
  "Status",
656
675
  "StatusCode",