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.
- sqlspec/__init__.py +7 -15
- sqlspec/_serialization.py +55 -25
- sqlspec/_typing.py +155 -52
- sqlspec/adapters/adbc/_types.py +1 -1
- sqlspec/adapters/adbc/adk/__init__.py +5 -0
- sqlspec/adapters/adbc/adk/store.py +880 -0
- sqlspec/adapters/adbc/config.py +62 -12
- sqlspec/adapters/adbc/data_dictionary.py +74 -2
- sqlspec/adapters/adbc/driver.py +226 -58
- sqlspec/adapters/adbc/litestar/__init__.py +5 -0
- sqlspec/adapters/adbc/litestar/store.py +504 -0
- sqlspec/adapters/adbc/type_converter.py +44 -50
- sqlspec/adapters/aiosqlite/_types.py +1 -1
- sqlspec/adapters/aiosqlite/adk/__init__.py +5 -0
- sqlspec/adapters/aiosqlite/adk/store.py +536 -0
- sqlspec/adapters/aiosqlite/config.py +86 -16
- sqlspec/adapters/aiosqlite/data_dictionary.py +34 -2
- sqlspec/adapters/aiosqlite/driver.py +127 -38
- sqlspec/adapters/aiosqlite/litestar/__init__.py +5 -0
- sqlspec/adapters/aiosqlite/litestar/store.py +281 -0
- sqlspec/adapters/aiosqlite/pool.py +7 -7
- sqlspec/adapters/asyncmy/__init__.py +7 -1
- sqlspec/adapters/asyncmy/_types.py +1 -1
- sqlspec/adapters/asyncmy/adk/__init__.py +5 -0
- sqlspec/adapters/asyncmy/adk/store.py +503 -0
- sqlspec/adapters/asyncmy/config.py +59 -17
- sqlspec/adapters/asyncmy/data_dictionary.py +41 -2
- sqlspec/adapters/asyncmy/driver.py +293 -62
- sqlspec/adapters/asyncmy/litestar/__init__.py +5 -0
- sqlspec/adapters/asyncmy/litestar/store.py +296 -0
- sqlspec/adapters/asyncpg/__init__.py +2 -1
- sqlspec/adapters/asyncpg/_type_handlers.py +71 -0
- sqlspec/adapters/asyncpg/_types.py +11 -7
- sqlspec/adapters/asyncpg/adk/__init__.py +5 -0
- sqlspec/adapters/asyncpg/adk/store.py +460 -0
- sqlspec/adapters/asyncpg/config.py +57 -36
- sqlspec/adapters/asyncpg/data_dictionary.py +48 -2
- sqlspec/adapters/asyncpg/driver.py +153 -23
- sqlspec/adapters/asyncpg/litestar/__init__.py +5 -0
- sqlspec/adapters/asyncpg/litestar/store.py +253 -0
- sqlspec/adapters/bigquery/_types.py +1 -1
- sqlspec/adapters/bigquery/adk/__init__.py +5 -0
- sqlspec/adapters/bigquery/adk/store.py +585 -0
- sqlspec/adapters/bigquery/config.py +36 -11
- sqlspec/adapters/bigquery/data_dictionary.py +42 -2
- sqlspec/adapters/bigquery/driver.py +489 -144
- sqlspec/adapters/bigquery/litestar/__init__.py +5 -0
- sqlspec/adapters/bigquery/litestar/store.py +327 -0
- sqlspec/adapters/bigquery/type_converter.py +55 -23
- sqlspec/adapters/duckdb/_types.py +2 -2
- sqlspec/adapters/duckdb/adk/__init__.py +14 -0
- sqlspec/adapters/duckdb/adk/store.py +563 -0
- sqlspec/adapters/duckdb/config.py +79 -21
- sqlspec/adapters/duckdb/data_dictionary.py +41 -2
- sqlspec/adapters/duckdb/driver.py +225 -44
- sqlspec/adapters/duckdb/litestar/__init__.py +5 -0
- sqlspec/adapters/duckdb/litestar/store.py +332 -0
- sqlspec/adapters/duckdb/pool.py +5 -5
- sqlspec/adapters/duckdb/type_converter.py +51 -21
- sqlspec/adapters/oracledb/_numpy_handlers.py +133 -0
- sqlspec/adapters/oracledb/_types.py +20 -2
- sqlspec/adapters/oracledb/adk/__init__.py +5 -0
- sqlspec/adapters/oracledb/adk/store.py +1628 -0
- sqlspec/adapters/oracledb/config.py +120 -36
- sqlspec/adapters/oracledb/data_dictionary.py +87 -20
- sqlspec/adapters/oracledb/driver.py +475 -86
- sqlspec/adapters/oracledb/litestar/__init__.py +5 -0
- sqlspec/adapters/oracledb/litestar/store.py +765 -0
- sqlspec/adapters/oracledb/migrations.py +316 -25
- sqlspec/adapters/oracledb/type_converter.py +91 -16
- sqlspec/adapters/psqlpy/_type_handlers.py +44 -0
- sqlspec/adapters/psqlpy/_types.py +2 -1
- sqlspec/adapters/psqlpy/adk/__init__.py +5 -0
- sqlspec/adapters/psqlpy/adk/store.py +483 -0
- sqlspec/adapters/psqlpy/config.py +45 -19
- sqlspec/adapters/psqlpy/data_dictionary.py +48 -2
- sqlspec/adapters/psqlpy/driver.py +108 -41
- sqlspec/adapters/psqlpy/litestar/__init__.py +5 -0
- sqlspec/adapters/psqlpy/litestar/store.py +272 -0
- sqlspec/adapters/psqlpy/type_converter.py +40 -11
- sqlspec/adapters/psycopg/_type_handlers.py +80 -0
- sqlspec/adapters/psycopg/_types.py +2 -1
- sqlspec/adapters/psycopg/adk/__init__.py +5 -0
- sqlspec/adapters/psycopg/adk/store.py +962 -0
- sqlspec/adapters/psycopg/config.py +65 -37
- sqlspec/adapters/psycopg/data_dictionary.py +91 -3
- sqlspec/adapters/psycopg/driver.py +200 -78
- sqlspec/adapters/psycopg/litestar/__init__.py +5 -0
- sqlspec/adapters/psycopg/litestar/store.py +554 -0
- sqlspec/adapters/sqlite/__init__.py +2 -1
- sqlspec/adapters/sqlite/_type_handlers.py +86 -0
- sqlspec/adapters/sqlite/_types.py +1 -1
- sqlspec/adapters/sqlite/adk/__init__.py +5 -0
- sqlspec/adapters/sqlite/adk/store.py +582 -0
- sqlspec/adapters/sqlite/config.py +85 -16
- sqlspec/adapters/sqlite/data_dictionary.py +34 -2
- sqlspec/adapters/sqlite/driver.py +120 -52
- sqlspec/adapters/sqlite/litestar/__init__.py +5 -0
- sqlspec/adapters/sqlite/litestar/store.py +318 -0
- sqlspec/adapters/sqlite/pool.py +5 -5
- sqlspec/base.py +45 -26
- sqlspec/builder/__init__.py +73 -4
- sqlspec/builder/_base.py +91 -58
- sqlspec/builder/_column.py +5 -5
- sqlspec/builder/_ddl.py +98 -89
- sqlspec/builder/_delete.py +5 -4
- sqlspec/builder/_dml.py +388 -0
- sqlspec/{_sql.py → builder/_factory.py} +41 -44
- sqlspec/builder/_insert.py +5 -82
- sqlspec/builder/{mixins/_join_operations.py → _join.py} +145 -143
- sqlspec/builder/_merge.py +446 -11
- sqlspec/builder/_parsing_utils.py +9 -11
- sqlspec/builder/_select.py +1313 -25
- sqlspec/builder/_update.py +11 -42
- sqlspec/cli.py +76 -69
- sqlspec/config.py +331 -62
- sqlspec/core/__init__.py +5 -4
- sqlspec/core/cache.py +18 -18
- sqlspec/core/compiler.py +6 -8
- sqlspec/core/filters.py +55 -47
- sqlspec/core/hashing.py +9 -9
- sqlspec/core/parameters.py +76 -45
- sqlspec/core/result.py +234 -47
- sqlspec/core/splitter.py +16 -17
- sqlspec/core/statement.py +32 -31
- sqlspec/core/type_conversion.py +3 -2
- sqlspec/driver/__init__.py +1 -3
- sqlspec/driver/_async.py +183 -160
- sqlspec/driver/_common.py +197 -109
- sqlspec/driver/_sync.py +189 -161
- sqlspec/driver/mixins/_result_tools.py +20 -236
- sqlspec/driver/mixins/_sql_translator.py +4 -4
- sqlspec/exceptions.py +70 -7
- sqlspec/extensions/adk/__init__.py +53 -0
- sqlspec/extensions/adk/_types.py +51 -0
- sqlspec/extensions/adk/converters.py +172 -0
- sqlspec/extensions/adk/migrations/0001_create_adk_tables.py +144 -0
- sqlspec/extensions/adk/migrations/__init__.py +0 -0
- sqlspec/extensions/adk/service.py +181 -0
- sqlspec/extensions/adk/store.py +536 -0
- sqlspec/extensions/aiosql/adapter.py +69 -61
- sqlspec/extensions/fastapi/__init__.py +21 -0
- sqlspec/extensions/fastapi/extension.py +331 -0
- sqlspec/extensions/fastapi/providers.py +543 -0
- sqlspec/extensions/flask/__init__.py +36 -0
- sqlspec/extensions/flask/_state.py +71 -0
- sqlspec/extensions/flask/_utils.py +40 -0
- sqlspec/extensions/flask/extension.py +389 -0
- sqlspec/extensions/litestar/__init__.py +21 -4
- sqlspec/extensions/litestar/cli.py +54 -10
- sqlspec/extensions/litestar/config.py +56 -266
- sqlspec/extensions/litestar/handlers.py +46 -17
- sqlspec/extensions/litestar/migrations/0001_create_session_table.py +137 -0
- sqlspec/extensions/litestar/migrations/__init__.py +3 -0
- sqlspec/extensions/litestar/plugin.py +349 -224
- sqlspec/extensions/litestar/providers.py +25 -25
- sqlspec/extensions/litestar/store.py +265 -0
- sqlspec/extensions/starlette/__init__.py +10 -0
- sqlspec/extensions/starlette/_state.py +25 -0
- sqlspec/extensions/starlette/_utils.py +52 -0
- sqlspec/extensions/starlette/extension.py +254 -0
- sqlspec/extensions/starlette/middleware.py +154 -0
- sqlspec/loader.py +30 -49
- sqlspec/migrations/base.py +200 -76
- sqlspec/migrations/commands.py +591 -62
- sqlspec/migrations/context.py +6 -9
- sqlspec/migrations/fix.py +199 -0
- sqlspec/migrations/loaders.py +47 -19
- sqlspec/migrations/runner.py +241 -75
- sqlspec/migrations/tracker.py +237 -21
- sqlspec/migrations/utils.py +51 -3
- sqlspec/migrations/validation.py +177 -0
- sqlspec/protocols.py +106 -36
- sqlspec/storage/_utils.py +85 -0
- sqlspec/storage/backends/fsspec.py +133 -107
- sqlspec/storage/backends/local.py +78 -51
- sqlspec/storage/backends/obstore.py +276 -168
- sqlspec/storage/registry.py +75 -39
- sqlspec/typing.py +30 -84
- sqlspec/utils/__init__.py +25 -4
- sqlspec/utils/arrow_helpers.py +81 -0
- sqlspec/utils/config_resolver.py +6 -6
- sqlspec/utils/correlation.py +4 -5
- sqlspec/utils/data_transformation.py +3 -2
- sqlspec/utils/deprecation.py +9 -8
- sqlspec/utils/fixtures.py +4 -4
- sqlspec/utils/logging.py +46 -6
- sqlspec/utils/module_loader.py +205 -5
- sqlspec/utils/portal.py +311 -0
- sqlspec/utils/schema.py +288 -0
- sqlspec/utils/serializers.py +113 -4
- sqlspec/utils/sync_tools.py +36 -22
- sqlspec/utils/text.py +1 -2
- sqlspec/utils/type_guards.py +136 -20
- sqlspec/utils/version.py +433 -0
- {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/METADATA +41 -22
- sqlspec-0.28.0.dist-info/RECORD +221 -0
- sqlspec/builder/mixins/__init__.py +0 -55
- sqlspec/builder/mixins/_cte_and_set_ops.py +0 -253
- sqlspec/builder/mixins/_delete_operations.py +0 -50
- sqlspec/builder/mixins/_insert_operations.py +0 -282
- sqlspec/builder/mixins/_merge_operations.py +0 -698
- sqlspec/builder/mixins/_order_limit_operations.py +0 -145
- sqlspec/builder/mixins/_pivot_operations.py +0 -157
- sqlspec/builder/mixins/_select_operations.py +0 -930
- sqlspec/builder/mixins/_update_operations.py +0 -199
- sqlspec/builder/mixins/_where_clause.py +0 -1298
- sqlspec-0.26.0.dist-info/RECORD +0 -157
- sqlspec-0.26.0.dist-info/licenses/NOTICE +0 -29
- {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/WHEEL +0 -0
- {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/entry_points.txt +0 -0
- {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/licenses/LICENSE +0 -0
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
|
|
3
3
|
from collections.abc import Sequence
|
|
4
4
|
from contextlib import contextmanager
|
|
5
|
-
from typing import TYPE_CHECKING, Any, ClassVar,
|
|
5
|
+
from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
|
|
6
6
|
|
|
7
7
|
from typing_extensions import NotRequired
|
|
8
8
|
|
|
@@ -26,7 +26,7 @@ __all__ = (
|
|
|
26
26
|
)
|
|
27
27
|
|
|
28
28
|
|
|
29
|
-
class DuckDBConnectionParams(TypedDict
|
|
29
|
+
class DuckDBConnectionParams(TypedDict):
|
|
30
30
|
"""DuckDB connection parameters."""
|
|
31
31
|
|
|
32
32
|
database: NotRequired[str]
|
|
@@ -65,7 +65,7 @@ class DuckDBConnectionParams(TypedDict, total=False):
|
|
|
65
65
|
extra: NotRequired[dict[str, Any]]
|
|
66
66
|
|
|
67
67
|
|
|
68
|
-
class DuckDBPoolParams(DuckDBConnectionParams
|
|
68
|
+
class DuckDBPoolParams(DuckDBConnectionParams):
|
|
69
69
|
"""Complete pool configuration for DuckDB adapter.
|
|
70
70
|
|
|
71
71
|
Combines standardized pool parameters with DuckDB-specific connection parameters.
|
|
@@ -77,7 +77,7 @@ class DuckDBPoolParams(DuckDBConnectionParams, total=False):
|
|
|
77
77
|
pool_recycle_seconds: NotRequired[int]
|
|
78
78
|
|
|
79
79
|
|
|
80
|
-
class DuckDBExtensionConfig(TypedDict
|
|
80
|
+
class DuckDBExtensionConfig(TypedDict):
|
|
81
81
|
"""DuckDB extension configuration for auto-management."""
|
|
82
82
|
|
|
83
83
|
name: str
|
|
@@ -93,7 +93,7 @@ class DuckDBExtensionConfig(TypedDict, total=False):
|
|
|
93
93
|
"""Force reinstallation of the extension."""
|
|
94
94
|
|
|
95
95
|
|
|
96
|
-
class DuckDBSecretConfig(TypedDict
|
|
96
|
+
class DuckDBSecretConfig(TypedDict):
|
|
97
97
|
"""DuckDB secret configuration for AI/API integrations."""
|
|
98
98
|
|
|
99
99
|
secret_type: str
|
|
@@ -109,15 +109,25 @@ class DuckDBSecretConfig(TypedDict, total=False):
|
|
|
109
109
|
"""Scope of the secret (LOCAL or PERSISTENT)."""
|
|
110
110
|
|
|
111
111
|
|
|
112
|
-
class DuckDBDriverFeatures(TypedDict
|
|
113
|
-
"""TypedDict for DuckDB driver features configuration.
|
|
112
|
+
class DuckDBDriverFeatures(TypedDict):
|
|
113
|
+
"""TypedDict for DuckDB driver features configuration.
|
|
114
|
+
|
|
115
|
+
Attributes:
|
|
116
|
+
extensions: List of extensions to install/load on connection creation.
|
|
117
|
+
secrets: List of secrets to create for AI/API integrations.
|
|
118
|
+
on_connection_create: Callback executed when connection is created.
|
|
119
|
+
json_serializer: Custom JSON serializer for dict/list parameter conversion.
|
|
120
|
+
Defaults to sqlspec.utils.serializers.to_json if not provided.
|
|
121
|
+
enable_uuid_conversion: Enable automatic UUID string conversion.
|
|
122
|
+
When True (default), UUID strings are automatically converted to UUID objects.
|
|
123
|
+
When False, UUID strings are treated as regular strings.
|
|
124
|
+
"""
|
|
114
125
|
|
|
115
126
|
extensions: NotRequired[Sequence[DuckDBExtensionConfig]]
|
|
116
|
-
"""List of extensions to install/load on connection creation."""
|
|
117
127
|
secrets: NotRequired[Sequence[DuckDBSecretConfig]]
|
|
118
|
-
"
|
|
119
|
-
|
|
120
|
-
|
|
128
|
+
on_connection_create: NotRequired["Callable[[DuckDBConnection], DuckDBConnection | None]"]
|
|
129
|
+
json_serializer: NotRequired["Callable[[Any], str]"]
|
|
130
|
+
enable_uuid_conversion: NotRequired[bool]
|
|
121
131
|
|
|
122
132
|
|
|
123
133
|
class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, DuckDBDriver]):
|
|
@@ -131,27 +141,66 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
|
|
|
131
141
|
- Auto configuration settings
|
|
132
142
|
- Arrow integration
|
|
133
143
|
- Direct file querying capabilities
|
|
144
|
+
- Configurable type handlers for JSON serialization and UUID conversion
|
|
134
145
|
|
|
135
146
|
DuckDB Connection Pool Configuration:
|
|
136
147
|
- Default pool size is 1-4 connections (DuckDB uses single connection by default)
|
|
137
148
|
- Connection recycling is set to 24 hours by default (set to 0 to disable)
|
|
138
149
|
- Shared memory databases use `:memory:shared_db` for proper concurrency
|
|
150
|
+
|
|
151
|
+
Type Handler Configuration via driver_features:
|
|
152
|
+
- `json_serializer`: Custom JSON serializer for dict/list parameters.
|
|
153
|
+
Defaults to `sqlspec.utils.serializers.to_json` if not provided.
|
|
154
|
+
Example: `json_serializer=msgspec.json.encode(...).decode('utf-8')`
|
|
155
|
+
|
|
156
|
+
- `enable_uuid_conversion`: Enable automatic UUID string conversion (default: True).
|
|
157
|
+
When True, UUID strings in query results are automatically converted to UUID objects.
|
|
158
|
+
When False, UUID strings are treated as regular strings.
|
|
159
|
+
|
|
160
|
+
Example:
|
|
161
|
+
>>> import msgspec
|
|
162
|
+
>>> from sqlspec.adapters.duckdb import DuckDBConfig
|
|
163
|
+
>>>
|
|
164
|
+
>>> # Custom JSON serializer
|
|
165
|
+
>>> def custom_json(obj):
|
|
166
|
+
... return msgspec.json.encode(obj).decode("utf-8")
|
|
167
|
+
>>>
|
|
168
|
+
>>> config = DuckDBConfig(
|
|
169
|
+
... pool_config={"database": ":memory:"},
|
|
170
|
+
... driver_features={
|
|
171
|
+
... "json_serializer": custom_json,
|
|
172
|
+
... "enable_uuid_conversion": False,
|
|
173
|
+
... },
|
|
174
|
+
... )
|
|
139
175
|
"""
|
|
140
176
|
|
|
141
177
|
driver_type: "ClassVar[type[DuckDBDriver]]" = DuckDBDriver
|
|
142
178
|
connection_type: "ClassVar[type[DuckDBConnection]]" = DuckDBConnection
|
|
179
|
+
supports_transactional_ddl: "ClassVar[bool]" = True
|
|
143
180
|
|
|
144
181
|
def __init__(
|
|
145
182
|
self,
|
|
146
183
|
*,
|
|
147
|
-
pool_config: "
|
|
148
|
-
pool_instance: "
|
|
149
|
-
migration_config:
|
|
150
|
-
statement_config: "
|
|
151
|
-
driver_features: "
|
|
152
|
-
bind_key: "
|
|
184
|
+
pool_config: "DuckDBPoolParams | dict[str, Any] | None" = None,
|
|
185
|
+
pool_instance: "DuckDBConnectionPool | None" = None,
|
|
186
|
+
migration_config: dict[str, Any] | None = None,
|
|
187
|
+
statement_config: "StatementConfig | None" = None,
|
|
188
|
+
driver_features: "DuckDBDriverFeatures | dict[str, Any] | None" = None,
|
|
189
|
+
bind_key: "str | None" = None,
|
|
190
|
+
extension_config: "dict[str, dict[str, Any]] | None" = None,
|
|
153
191
|
) -> None:
|
|
154
|
-
"""Initialize DuckDB configuration.
|
|
192
|
+
"""Initialize DuckDB configuration.
|
|
193
|
+
|
|
194
|
+
Args:
|
|
195
|
+
pool_config: Pool configuration parameters
|
|
196
|
+
pool_instance: Pre-created pool instance
|
|
197
|
+
migration_config: Migration configuration
|
|
198
|
+
statement_config: Statement configuration override
|
|
199
|
+
driver_features: DuckDB-specific driver features including json_serializer
|
|
200
|
+
and enable_uuid_conversion options
|
|
201
|
+
bind_key: Optional unique identifier for this configuration
|
|
202
|
+
extension_config: Extension-specific configuration (e.g., Litestar plugin settings)
|
|
203
|
+
"""
|
|
155
204
|
if pool_config is None:
|
|
156
205
|
pool_config = {}
|
|
157
206
|
if "database" not in pool_config:
|
|
@@ -160,13 +209,18 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
|
|
|
160
209
|
if pool_config.get("database") in {":memory:", ""}:
|
|
161
210
|
pool_config["database"] = ":memory:shared_db"
|
|
162
211
|
|
|
212
|
+
processed_features = dict(driver_features) if driver_features else {}
|
|
213
|
+
if "enable_uuid_conversion" not in processed_features:
|
|
214
|
+
processed_features["enable_uuid_conversion"] = True
|
|
215
|
+
|
|
163
216
|
super().__init__(
|
|
164
217
|
bind_key=bind_key,
|
|
165
218
|
pool_config=dict(pool_config),
|
|
166
219
|
pool_instance=pool_instance,
|
|
167
220
|
migration_config=migration_config,
|
|
168
221
|
statement_config=statement_config or duckdb_statement_config,
|
|
169
|
-
driver_features=
|
|
222
|
+
driver_features=processed_features,
|
|
223
|
+
extension_config=extension_config,
|
|
170
224
|
)
|
|
171
225
|
|
|
172
226
|
def _get_connection_config_dict(self) -> "dict[str, Any]":
|
|
@@ -247,7 +301,7 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
|
|
|
247
301
|
|
|
248
302
|
@contextmanager
|
|
249
303
|
def provide_session(
|
|
250
|
-
self, *args: Any, statement_config: "
|
|
304
|
+
self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
|
|
251
305
|
) -> "Generator[DuckDBDriver, None, None]":
|
|
252
306
|
"""Provide a DuckDB driver session context manager.
|
|
253
307
|
|
|
@@ -260,7 +314,11 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
|
|
|
260
314
|
A context manager that yields a DuckDBDriver instance.
|
|
261
315
|
"""
|
|
262
316
|
with self.provide_connection(*args, **kwargs) as connection:
|
|
263
|
-
driver = self.driver_type(
|
|
317
|
+
driver = self.driver_type(
|
|
318
|
+
connection=connection,
|
|
319
|
+
statement_config=statement_config or self.statement_config,
|
|
320
|
+
driver_features=self.driver_features,
|
|
321
|
+
)
|
|
264
322
|
yield driver
|
|
265
323
|
|
|
266
324
|
def get_signature_namespace(self) -> "dict[str, type[Any]]":
|
|
@@ -1,12 +1,14 @@
|
|
|
1
1
|
"""DuckDB-specific data dictionary for metadata queries."""
|
|
2
2
|
|
|
3
3
|
import re
|
|
4
|
-
from typing import TYPE_CHECKING,
|
|
4
|
+
from typing import TYPE_CHECKING, Any, cast
|
|
5
5
|
|
|
6
6
|
from sqlspec.driver import SyncDataDictionaryBase, SyncDriverAdapterBase, 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.duckdb.driver import DuckDBDriver
|
|
11
13
|
|
|
12
14
|
logger = get_logger("adapters.duckdb.data_dictionary")
|
|
@@ -20,7 +22,7 @@ __all__ = ("DuckDBSyncDataDictionary",)
|
|
|
20
22
|
class DuckDBSyncDataDictionary(SyncDataDictionaryBase):
|
|
21
23
|
"""DuckDB-specific sync data dictionary."""
|
|
22
24
|
|
|
23
|
-
def get_version(self, driver: SyncDriverAdapterBase) -> "
|
|
25
|
+
def get_version(self, driver: SyncDriverAdapterBase) -> "VersionInfo | None":
|
|
24
26
|
"""Get DuckDB database version information.
|
|
25
27
|
|
|
26
28
|
Args:
|
|
@@ -102,6 +104,43 @@ class DuckDBSyncDataDictionary(SyncDataDictionaryBase):
|
|
|
102
104
|
}
|
|
103
105
|
return type_map.get(type_category, "VARCHAR")
|
|
104
106
|
|
|
107
|
+
def get_columns(
|
|
108
|
+
self, driver: SyncDriverAdapterBase, table: str, schema: "str | None" = None
|
|
109
|
+
) -> "list[dict[str, Any]]":
|
|
110
|
+
"""Get column information for a table using information_schema.
|
|
111
|
+
|
|
112
|
+
Args:
|
|
113
|
+
driver: DuckDB driver instance
|
|
114
|
+
table: Table name to query columns for
|
|
115
|
+
schema: Schema name (None for default)
|
|
116
|
+
|
|
117
|
+
Returns:
|
|
118
|
+
List of column metadata dictionaries with keys:
|
|
119
|
+
- column_name: Name of the column
|
|
120
|
+
- data_type: DuckDB data type
|
|
121
|
+
- nullable: Whether column allows NULL (YES/NO)
|
|
122
|
+
- column_default: Default value if any
|
|
123
|
+
"""
|
|
124
|
+
duckdb_driver = cast("DuckDBDriver", driver)
|
|
125
|
+
|
|
126
|
+
if schema:
|
|
127
|
+
sql = f"""
|
|
128
|
+
SELECT column_name, data_type, is_nullable, column_default
|
|
129
|
+
FROM information_schema.columns
|
|
130
|
+
WHERE table_name = '{table}' AND table_schema = '{schema}'
|
|
131
|
+
ORDER BY ordinal_position
|
|
132
|
+
"""
|
|
133
|
+
else:
|
|
134
|
+
sql = f"""
|
|
135
|
+
SELECT column_name, data_type, is_nullable, column_default
|
|
136
|
+
FROM information_schema.columns
|
|
137
|
+
WHERE table_name = '{table}'
|
|
138
|
+
ORDER BY ordinal_position
|
|
139
|
+
"""
|
|
140
|
+
|
|
141
|
+
result = duckdb_driver.execute(sql)
|
|
142
|
+
return result.data or []
|
|
143
|
+
|
|
105
144
|
def list_available_features(self) -> "list[str]":
|
|
106
145
|
"""List available DuckDB feature flags.
|
|
107
146
|
|
|
@@ -2,9 +2,9 @@
|
|
|
2
2
|
|
|
3
3
|
import datetime
|
|
4
4
|
from decimal import Decimal
|
|
5
|
-
from typing import TYPE_CHECKING, Any, Final
|
|
5
|
+
from typing import TYPE_CHECKING, Any, Final
|
|
6
6
|
|
|
7
|
-
import duckdb
|
|
7
|
+
import duckdb
|
|
8
8
|
from sqlglot import exp
|
|
9
9
|
|
|
10
10
|
from sqlspec.adapters.duckdb.data_dictionary import DuckDBSyncDataDictionary
|
|
@@ -13,7 +13,19 @@ from sqlspec.core.cache import get_cache_config
|
|
|
13
13
|
from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
|
|
14
14
|
from sqlspec.core.statement import SQL, StatementConfig
|
|
15
15
|
from sqlspec.driver import SyncDriverAdapterBase
|
|
16
|
-
from sqlspec.exceptions import
|
|
16
|
+
from sqlspec.exceptions import (
|
|
17
|
+
CheckViolationError,
|
|
18
|
+
DatabaseConnectionError,
|
|
19
|
+
DataError,
|
|
20
|
+
ForeignKeyViolationError,
|
|
21
|
+
IntegrityError,
|
|
22
|
+
NotFoundError,
|
|
23
|
+
NotNullViolationError,
|
|
24
|
+
OperationalError,
|
|
25
|
+
SQLParsingError,
|
|
26
|
+
SQLSpecError,
|
|
27
|
+
UniqueViolationError,
|
|
28
|
+
)
|
|
17
29
|
from sqlspec.utils.logging import get_logger
|
|
18
30
|
from sqlspec.utils.serializers import to_json
|
|
19
31
|
|
|
@@ -21,9 +33,12 @@ if TYPE_CHECKING:
|
|
|
21
33
|
from contextlib import AbstractContextManager
|
|
22
34
|
|
|
23
35
|
from sqlspec.adapters.duckdb._types import DuckDBConnection
|
|
24
|
-
from sqlspec.
|
|
36
|
+
from sqlspec.builder import QueryBuilder
|
|
37
|
+
from sqlspec.core import Statement, StatementFilter
|
|
38
|
+
from sqlspec.core.result import ArrowResult, SQLResult
|
|
25
39
|
from sqlspec.driver import ExecutionResult
|
|
26
40
|
from sqlspec.driver._sync import SyncDataDictionaryBase
|
|
41
|
+
from sqlspec.typing import StatementParameters
|
|
27
42
|
|
|
28
43
|
__all__ = ("DuckDBCursor", "DuckDBDriver", "DuckDBExceptionHandler", "duckdb_statement_config")
|
|
29
44
|
|
|
@@ -70,7 +85,7 @@ class DuckDBCursor:
|
|
|
70
85
|
|
|
71
86
|
def __init__(self, connection: "DuckDBConnection") -> None:
|
|
72
87
|
self.connection = connection
|
|
73
|
-
self.cursor:
|
|
88
|
+
self.cursor: Any | None = None
|
|
74
89
|
|
|
75
90
|
def __enter__(self) -> Any:
|
|
76
91
|
self.cursor = self.connection.cursor()
|
|
@@ -84,8 +99,8 @@ class DuckDBCursor:
|
|
|
84
99
|
class DuckDBExceptionHandler:
|
|
85
100
|
"""Context manager for handling DuckDB database exceptions.
|
|
86
101
|
|
|
87
|
-
|
|
88
|
-
SQLSpec
|
|
102
|
+
Uses exception type and message-based detection to map DuckDB errors
|
|
103
|
+
to specific SQLSpec exceptions for better error handling.
|
|
89
104
|
"""
|
|
90
105
|
|
|
91
106
|
__slots__ = ()
|
|
@@ -94,41 +109,93 @@ class DuckDBExceptionHandler:
|
|
|
94
109
|
return None
|
|
95
110
|
|
|
96
111
|
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
112
|
+
_ = exc_tb
|
|
97
113
|
if exc_type is None:
|
|
98
114
|
return
|
|
115
|
+
self._map_duckdb_exception(exc_type, exc_val)
|
|
99
116
|
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
if
|
|
113
|
-
e
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
e
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
117
|
+
def _map_duckdb_exception(self, exc_type: Any, e: Any) -> None:
|
|
118
|
+
"""Map DuckDB exception to SQLSpec exception.
|
|
119
|
+
|
|
120
|
+
Uses exception type and message-based detection.
|
|
121
|
+
|
|
122
|
+
Args:
|
|
123
|
+
exc_type: Exception type
|
|
124
|
+
e: Exception instance
|
|
125
|
+
"""
|
|
126
|
+
error_msg = str(e).lower()
|
|
127
|
+
exc_name = exc_type.__name__ if hasattr(exc_type, "__name__") else str(exc_type)
|
|
128
|
+
|
|
129
|
+
if "constraintexception" in exc_name.lower():
|
|
130
|
+
self._handle_constraint_exception(e, error_msg)
|
|
131
|
+
elif "catalogexception" in exc_name.lower():
|
|
132
|
+
self._raise_not_found_error(e)
|
|
133
|
+
elif "parserexception" in exc_name.lower() or "binderexception" in exc_name.lower():
|
|
134
|
+
self._raise_parsing_error(e)
|
|
135
|
+
elif "ioexception" in exc_name.lower():
|
|
136
|
+
self._raise_operational_error(e)
|
|
137
|
+
elif "conversionexception" in exc_name.lower() or "type mismatch" in error_msg:
|
|
138
|
+
self._raise_data_error(e)
|
|
139
|
+
else:
|
|
140
|
+
self._raise_generic_error(e)
|
|
141
|
+
|
|
142
|
+
def _handle_constraint_exception(self, e: Any, error_msg: str) -> None:
|
|
143
|
+
"""Handle constraint exceptions using message-based detection.
|
|
144
|
+
|
|
145
|
+
Args:
|
|
146
|
+
e: Exception instance
|
|
147
|
+
error_msg: Lowercase error message
|
|
148
|
+
"""
|
|
149
|
+
if "unique" in error_msg or "duplicate" in error_msg:
|
|
150
|
+
self._raise_unique_violation(e)
|
|
151
|
+
elif "foreign key" in error_msg or "violates foreign key" in error_msg:
|
|
152
|
+
self._raise_foreign_key_violation(e)
|
|
153
|
+
elif "not null" in error_msg or "null value" in error_msg:
|
|
154
|
+
self._raise_not_null_violation(e)
|
|
155
|
+
elif "check constraint" in error_msg or "check condition" in error_msg:
|
|
156
|
+
self._raise_check_violation(e)
|
|
157
|
+
else:
|
|
158
|
+
self._raise_integrity_error(e)
|
|
159
|
+
|
|
160
|
+
def _raise_unique_violation(self, e: Any) -> None:
|
|
161
|
+
msg = f"DuckDB unique constraint violation: {e}"
|
|
162
|
+
raise UniqueViolationError(msg) from e
|
|
163
|
+
|
|
164
|
+
def _raise_foreign_key_violation(self, e: Any) -> None:
|
|
165
|
+
msg = f"DuckDB foreign key constraint violation: {e}"
|
|
166
|
+
raise ForeignKeyViolationError(msg) from e
|
|
167
|
+
|
|
168
|
+
def _raise_not_null_violation(self, e: Any) -> None:
|
|
169
|
+
msg = f"DuckDB not-null constraint violation: {e}"
|
|
170
|
+
raise NotNullViolationError(msg) from e
|
|
171
|
+
|
|
172
|
+
def _raise_check_violation(self, e: Any) -> None:
|
|
173
|
+
msg = f"DuckDB check constraint violation: {e}"
|
|
174
|
+
raise CheckViolationError(msg) from e
|
|
175
|
+
|
|
176
|
+
def _raise_integrity_error(self, e: Any) -> None:
|
|
177
|
+
msg = f"DuckDB integrity constraint violation: {e}"
|
|
178
|
+
raise IntegrityError(msg) from e
|
|
179
|
+
|
|
180
|
+
def _raise_not_found_error(self, e: Any) -> None:
|
|
181
|
+
msg = f"DuckDB catalog error: {e}"
|
|
182
|
+
raise NotFoundError(msg) from e
|
|
183
|
+
|
|
184
|
+
def _raise_parsing_error(self, e: Any) -> None:
|
|
185
|
+
msg = f"DuckDB SQL parsing error: {e}"
|
|
186
|
+
raise SQLParsingError(msg) from e
|
|
187
|
+
|
|
188
|
+
def _raise_operational_error(self, e: Any) -> None:
|
|
189
|
+
msg = f"DuckDB operational error: {e}"
|
|
190
|
+
raise OperationalError(msg) from e
|
|
191
|
+
|
|
192
|
+
def _raise_data_error(self, e: Any) -> None:
|
|
193
|
+
msg = f"DuckDB data error: {e}"
|
|
194
|
+
raise DataError(msg) from e
|
|
195
|
+
|
|
196
|
+
def _raise_generic_error(self, e: Any) -> None:
|
|
197
|
+
msg = f"DuckDB database error: {e}"
|
|
198
|
+
raise SQLSpecError(msg) from e
|
|
132
199
|
|
|
133
200
|
|
|
134
201
|
class DuckDBDriver(SyncDriverAdapterBase):
|
|
@@ -148,8 +215,8 @@ class DuckDBDriver(SyncDriverAdapterBase):
|
|
|
148
215
|
def __init__(
|
|
149
216
|
self,
|
|
150
217
|
connection: "DuckDBConnection",
|
|
151
|
-
statement_config: "
|
|
152
|
-
driver_features: "
|
|
218
|
+
statement_config: "StatementConfig | None" = None,
|
|
219
|
+
driver_features: "dict[str, Any] | None" = None,
|
|
153
220
|
) -> None:
|
|
154
221
|
if statement_config is None:
|
|
155
222
|
cache_config = get_cache_config()
|
|
@@ -161,8 +228,40 @@ class DuckDBDriver(SyncDriverAdapterBase):
|
|
|
161
228
|
)
|
|
162
229
|
statement_config = updated_config
|
|
163
230
|
|
|
231
|
+
if driver_features:
|
|
232
|
+
json_serializer = driver_features.get("json_serializer")
|
|
233
|
+
enable_uuid_conversion = driver_features.get("enable_uuid_conversion", True)
|
|
234
|
+
|
|
235
|
+
if json_serializer or not enable_uuid_conversion:
|
|
236
|
+
type_converter = DuckDBTypeConverter(enable_uuid_conversion=enable_uuid_conversion)
|
|
237
|
+
type_coercion_map = dict(statement_config.parameter_config.type_coercion_map)
|
|
238
|
+
|
|
239
|
+
if json_serializer:
|
|
240
|
+
type_coercion_map[dict] = json_serializer
|
|
241
|
+
type_coercion_map[list] = json_serializer
|
|
242
|
+
|
|
243
|
+
if not enable_uuid_conversion:
|
|
244
|
+
type_coercion_map[str] = type_converter.convert_if_detected
|
|
245
|
+
|
|
246
|
+
param_config = statement_config.parameter_config
|
|
247
|
+
updated_param_config = ParameterStyleConfig(
|
|
248
|
+
default_parameter_style=param_config.default_parameter_style,
|
|
249
|
+
supported_parameter_styles=param_config.supported_parameter_styles,
|
|
250
|
+
supported_execution_parameter_styles=param_config.supported_execution_parameter_styles,
|
|
251
|
+
default_execution_parameter_style=param_config.default_execution_parameter_style,
|
|
252
|
+
type_coercion_map=type_coercion_map,
|
|
253
|
+
has_native_list_expansion=param_config.has_native_list_expansion,
|
|
254
|
+
needs_static_script_compilation=param_config.needs_static_script_compilation,
|
|
255
|
+
allow_mixed_parameter_styles=param_config.allow_mixed_parameter_styles,
|
|
256
|
+
preserve_parameter_format=param_config.preserve_parameter_format,
|
|
257
|
+
preserve_original_params_for_many=param_config.preserve_original_params_for_many,
|
|
258
|
+
output_transformer=param_config.output_transformer,
|
|
259
|
+
ast_transformer=param_config.ast_transformer,
|
|
260
|
+
)
|
|
261
|
+
statement_config = statement_config.replace(parameter_config=updated_param_config)
|
|
262
|
+
|
|
164
263
|
super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
|
|
165
|
-
self._data_dictionary:
|
|
264
|
+
self._data_dictionary: SyncDataDictionaryBase | None = None
|
|
166
265
|
|
|
167
266
|
def with_cursor(self, connection: "DuckDBConnection") -> "DuckDBCursor":
|
|
168
267
|
"""Create context manager for DuckDB cursor.
|
|
@@ -183,7 +282,7 @@ class DuckDBDriver(SyncDriverAdapterBase):
|
|
|
183
282
|
"""
|
|
184
283
|
return DuckDBExceptionHandler()
|
|
185
284
|
|
|
186
|
-
def _try_special_handling(self, cursor: Any, statement: SQL) -> "
|
|
285
|
+
def _try_special_handling(self, cursor: Any, statement: SQL) -> "SQLResult | None":
|
|
187
286
|
"""Handle DuckDB-specific special operations.
|
|
188
287
|
|
|
189
288
|
DuckDB does not require special operation handling, so this method
|
|
@@ -298,7 +397,7 @@ class DuckDBDriver(SyncDriverAdapterBase):
|
|
|
298
397
|
column_names = [col[0] for col in cursor.description or []]
|
|
299
398
|
|
|
300
399
|
if fetched_data and isinstance(fetched_data[0], tuple):
|
|
301
|
-
dict_data = [dict(zip(column_names, row)) for row in fetched_data]
|
|
400
|
+
dict_data = [dict(zip(column_names, row, strict=False)) for row in fetched_data]
|
|
302
401
|
else:
|
|
303
402
|
dict_data = fetched_data
|
|
304
403
|
|
|
@@ -352,3 +451,85 @@ class DuckDBDriver(SyncDriverAdapterBase):
|
|
|
352
451
|
if self._data_dictionary is None:
|
|
353
452
|
self._data_dictionary = DuckDBSyncDataDictionary()
|
|
354
453
|
return self._data_dictionary
|
|
454
|
+
|
|
455
|
+
def select_to_arrow(
|
|
456
|
+
self,
|
|
457
|
+
statement: "Statement | QueryBuilder",
|
|
458
|
+
/,
|
|
459
|
+
*parameters: "StatementParameters | StatementFilter",
|
|
460
|
+
statement_config: "StatementConfig | None" = None,
|
|
461
|
+
return_format: str = "table",
|
|
462
|
+
native_only: bool = False,
|
|
463
|
+
batch_size: int | None = None,
|
|
464
|
+
arrow_schema: Any = None,
|
|
465
|
+
**kwargs: Any,
|
|
466
|
+
) -> "ArrowResult":
|
|
467
|
+
"""Execute query and return results as Apache Arrow (DuckDB native path).
|
|
468
|
+
|
|
469
|
+
DuckDB provides native Arrow support via cursor.arrow().
|
|
470
|
+
This is the fastest path due to DuckDB's columnar architecture.
|
|
471
|
+
|
|
472
|
+
Args:
|
|
473
|
+
statement: SQL statement, string, or QueryBuilder
|
|
474
|
+
*parameters: Query parameters or filters
|
|
475
|
+
statement_config: Optional statement configuration override
|
|
476
|
+
return_format: "table" for pyarrow.Table (default), "batch" for RecordBatch
|
|
477
|
+
native_only: Ignored for DuckDB (always uses native path)
|
|
478
|
+
batch_size: Batch size hint (for future streaming implementation)
|
|
479
|
+
arrow_schema: Optional pyarrow.Schema for type casting
|
|
480
|
+
**kwargs: Additional keyword arguments
|
|
481
|
+
|
|
482
|
+
Returns:
|
|
483
|
+
ArrowResult with native Arrow data
|
|
484
|
+
|
|
485
|
+
Raises:
|
|
486
|
+
MissingDependencyError: If pyarrow not installed
|
|
487
|
+
SQLExecutionError: If query execution fails
|
|
488
|
+
|
|
489
|
+
Example:
|
|
490
|
+
>>> result = driver.select_to_arrow(
|
|
491
|
+
... "SELECT * FROM users WHERE age > ?", 18
|
|
492
|
+
... )
|
|
493
|
+
>>> df = result.to_pandas() # Fast zero-copy conversion
|
|
494
|
+
"""
|
|
495
|
+
from sqlspec.utils.module_loader import ensure_pyarrow
|
|
496
|
+
|
|
497
|
+
ensure_pyarrow()
|
|
498
|
+
|
|
499
|
+
import pyarrow as pa
|
|
500
|
+
|
|
501
|
+
from sqlspec.core.result import create_arrow_result
|
|
502
|
+
|
|
503
|
+
# Prepare statement
|
|
504
|
+
config = statement_config or self.statement_config
|
|
505
|
+
prepared_statement = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs)
|
|
506
|
+
|
|
507
|
+
# Execute query and get native Arrow
|
|
508
|
+
with self.with_cursor(self.connection) as cursor, self.handle_database_exceptions():
|
|
509
|
+
if cursor is None:
|
|
510
|
+
msg = "Failed to create cursor"
|
|
511
|
+
raise DatabaseConnectionError(msg)
|
|
512
|
+
|
|
513
|
+
# Get compiled SQL and parameters
|
|
514
|
+
sql, driver_params = self._get_compiled_sql(prepared_statement, config)
|
|
515
|
+
|
|
516
|
+
# Execute query
|
|
517
|
+
cursor.execute(sql, driver_params or ())
|
|
518
|
+
|
|
519
|
+
# DuckDB native Arrow (zero-copy!)
|
|
520
|
+
arrow_reader = cursor.arrow()
|
|
521
|
+
arrow_table = arrow_reader.read_all()
|
|
522
|
+
|
|
523
|
+
# Apply schema casting if requested
|
|
524
|
+
if arrow_schema is not None:
|
|
525
|
+
arrow_table = arrow_table.cast(arrow_schema)
|
|
526
|
+
|
|
527
|
+
# Convert to batch if requested
|
|
528
|
+
if return_format == "batch":
|
|
529
|
+
batches = arrow_table.to_batches()
|
|
530
|
+
arrow_data: Any = batches[0] if batches else pa.RecordBatch.from_pydict({})
|
|
531
|
+
else:
|
|
532
|
+
arrow_data = arrow_table
|
|
533
|
+
|
|
534
|
+
# Create ArrowResult
|
|
535
|
+
return create_arrow_result(statement=prepared_statement, data=arrow_data, rows_affected=arrow_data.num_rows)
|