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
sqlspec/adapters/adbc/config.py
CHANGED
|
@@ -1,8 +1,9 @@
|
|
|
1
1
|
"""ADBC database configuration."""
|
|
2
2
|
|
|
3
3
|
import logging
|
|
4
|
+
from collections.abc import Callable
|
|
4
5
|
from contextlib import contextmanager
|
|
5
|
-
from typing import TYPE_CHECKING, Any,
|
|
6
|
+
from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
|
|
6
7
|
|
|
7
8
|
from typing_extensions import NotRequired
|
|
8
9
|
|
|
@@ -22,7 +23,7 @@ if TYPE_CHECKING:
|
|
|
22
23
|
logger = logging.getLogger("sqlspec.adapters.adbc")
|
|
23
24
|
|
|
24
25
|
|
|
25
|
-
class AdbcConnectionParams(TypedDict
|
|
26
|
+
class AdbcConnectionParams(TypedDict):
|
|
26
27
|
"""ADBC connection parameters."""
|
|
27
28
|
|
|
28
29
|
uri: NotRequired[str]
|
|
@@ -54,7 +55,37 @@ class AdbcConnectionParams(TypedDict, total=False):
|
|
|
54
55
|
extra: NotRequired[dict[str, Any]]
|
|
55
56
|
|
|
56
57
|
|
|
57
|
-
|
|
58
|
+
class AdbcDriverFeatures(TypedDict):
|
|
59
|
+
"""ADBC driver feature configuration.
|
|
60
|
+
|
|
61
|
+
Controls optional type handling and serialization behavior for the ADBC adapter.
|
|
62
|
+
These features configure how data is converted between Python and Arrow types.
|
|
63
|
+
|
|
64
|
+
Attributes:
|
|
65
|
+
json_serializer: JSON serialization function to use.
|
|
66
|
+
Callable that takes Any and returns str (JSON string).
|
|
67
|
+
Default: sqlspec.utils.serializers.to_json
|
|
68
|
+
enable_cast_detection: Enable cast-aware parameter processing.
|
|
69
|
+
When True, detects SQL casts (e.g., ::JSONB) and applies appropriate
|
|
70
|
+
serialization. Currently used for PostgreSQL JSONB handling.
|
|
71
|
+
Default: True
|
|
72
|
+
strict_type_coercion: Enforce strict type coercion rules.
|
|
73
|
+
When True, raises errors for unsupported type conversions.
|
|
74
|
+
When False, attempts best-effort conversion.
|
|
75
|
+
Default: False
|
|
76
|
+
arrow_extension_types: Enable PyArrow extension type support.
|
|
77
|
+
When True, preserves Arrow extension type metadata when reading data.
|
|
78
|
+
When False, falls back to storage types.
|
|
79
|
+
Default: True
|
|
80
|
+
"""
|
|
81
|
+
|
|
82
|
+
json_serializer: "NotRequired[Callable[[Any], str]]"
|
|
83
|
+
enable_cast_detection: NotRequired[bool]
|
|
84
|
+
strict_type_coercion: NotRequired[bool]
|
|
85
|
+
arrow_extension_types: NotRequired[bool]
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
__all__ = ("AdbcConfig", "AdbcConnectionParams", "AdbcDriverFeatures")
|
|
58
89
|
|
|
59
90
|
|
|
60
91
|
class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
|
|
@@ -69,15 +100,17 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
|
|
|
69
100
|
|
|
70
101
|
driver_type: ClassVar[type[AdbcDriver]] = AdbcDriver
|
|
71
102
|
connection_type: "ClassVar[type[AdbcConnection]]" = AdbcConnection
|
|
103
|
+
supports_transactional_ddl: ClassVar[bool] = False
|
|
72
104
|
|
|
73
105
|
def __init__(
|
|
74
106
|
self,
|
|
75
107
|
*,
|
|
76
|
-
connection_config:
|
|
77
|
-
migration_config:
|
|
78
|
-
statement_config:
|
|
79
|
-
driver_features:
|
|
80
|
-
bind_key:
|
|
108
|
+
connection_config: AdbcConnectionParams | dict[str, Any] | None = None,
|
|
109
|
+
migration_config: dict[str, Any] | None = None,
|
|
110
|
+
statement_config: StatementConfig | None = None,
|
|
111
|
+
driver_features: "AdbcDriverFeatures | dict[str, Any] | None" = None,
|
|
112
|
+
bind_key: str | None = None,
|
|
113
|
+
extension_config: "dict[str, dict[str, Any]] | None" = None,
|
|
81
114
|
) -> None:
|
|
82
115
|
"""Initialize configuration.
|
|
83
116
|
|
|
@@ -85,8 +118,9 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
|
|
|
85
118
|
connection_config: Connection configuration parameters
|
|
86
119
|
migration_config: Migration configuration
|
|
87
120
|
statement_config: Default SQL statement configuration
|
|
88
|
-
driver_features: Driver feature configuration
|
|
121
|
+
driver_features: Driver feature configuration (AdbcDriverFeatures)
|
|
89
122
|
bind_key: Optional unique identifier for this configuration
|
|
123
|
+
extension_config: Extension-specific configuration (e.g., Litestar plugin settings)
|
|
90
124
|
"""
|
|
91
125
|
if connection_config is None:
|
|
92
126
|
connection_config = {}
|
|
@@ -101,12 +135,26 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
|
|
|
101
135
|
detected_dialect = str(self._get_dialect() or "sqlite")
|
|
102
136
|
statement_config = get_adbc_statement_config(detected_dialect)
|
|
103
137
|
|
|
138
|
+
from sqlspec.utils.serializers import to_json
|
|
139
|
+
|
|
140
|
+
if driver_features is None:
|
|
141
|
+
driver_features = {}
|
|
142
|
+
if "json_serializer" not in driver_features:
|
|
143
|
+
driver_features["json_serializer"] = to_json
|
|
144
|
+
if "enable_cast_detection" not in driver_features:
|
|
145
|
+
driver_features["enable_cast_detection"] = True
|
|
146
|
+
if "strict_type_coercion" not in driver_features:
|
|
147
|
+
driver_features["strict_type_coercion"] = False
|
|
148
|
+
if "arrow_extension_types" not in driver_features:
|
|
149
|
+
driver_features["arrow_extension_types"] = True
|
|
150
|
+
|
|
104
151
|
super().__init__(
|
|
105
152
|
connection_config=self.connection_config,
|
|
106
153
|
migration_config=migration_config,
|
|
107
154
|
statement_config=statement_config,
|
|
108
|
-
driver_features=driver_features
|
|
155
|
+
driver_features=dict(driver_features),
|
|
109
156
|
bind_key=bind_key,
|
|
157
|
+
extension_config=extension_config,
|
|
110
158
|
)
|
|
111
159
|
|
|
112
160
|
def _resolve_driver_name(self) -> str:
|
|
@@ -284,7 +332,7 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
|
|
|
284
332
|
connection.close()
|
|
285
333
|
|
|
286
334
|
def provide_session(
|
|
287
|
-
self, *args: Any, statement_config: "
|
|
335
|
+
self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
|
|
288
336
|
) -> "AbstractContextManager[AdbcDriver]":
|
|
289
337
|
"""Provide a driver session context manager.
|
|
290
338
|
|
|
@@ -305,7 +353,9 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
|
|
|
305
353
|
or self.statement_config
|
|
306
354
|
or get_adbc_statement_config(str(self._get_dialect() or "sqlite"))
|
|
307
355
|
)
|
|
308
|
-
yield self.driver_type(
|
|
356
|
+
yield self.driver_type(
|
|
357
|
+
connection=connection, statement_config=final_statement_config, driver_features=self.driver_features
|
|
358
|
+
)
|
|
309
359
|
|
|
310
360
|
return session_manager()
|
|
311
361
|
|
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
"""ADBC multi-dialect 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
|
|
@@ -38,7 +38,7 @@ class AdbcDataDictionary(SyncDataDictionaryBase):
|
|
|
38
38
|
"""
|
|
39
39
|
return str(cast("AdbcDriver", driver).dialect)
|
|
40
40
|
|
|
41
|
-
def get_version(self, driver: SyncDriverAdapterBase) -> "
|
|
41
|
+
def get_version(self, driver: SyncDriverAdapterBase) -> "VersionInfo | None":
|
|
42
42
|
"""Get database version information based on detected dialect.
|
|
43
43
|
|
|
44
44
|
Args:
|
|
@@ -268,6 +268,78 @@ class AdbcDataDictionary(SyncDataDictionaryBase):
|
|
|
268
268
|
|
|
269
269
|
return type_map.get(type_category, "TEXT")
|
|
270
270
|
|
|
271
|
+
def get_columns(
|
|
272
|
+
self, driver: SyncDriverAdapterBase, table: str, schema: "str | None" = None
|
|
273
|
+
) -> "list[dict[str, Any]]":
|
|
274
|
+
"""Get column information for a table based on detected dialect.
|
|
275
|
+
|
|
276
|
+
Args:
|
|
277
|
+
driver: ADBC driver instance
|
|
278
|
+
table: Table name to query columns for
|
|
279
|
+
schema: Schema name (None for default)
|
|
280
|
+
|
|
281
|
+
Returns:
|
|
282
|
+
List of column metadata dictionaries with keys:
|
|
283
|
+
- column_name: Name of the column
|
|
284
|
+
- data_type: Database data type
|
|
285
|
+
- is_nullable or nullable: Whether column allows NULL
|
|
286
|
+
- column_default or default_value: Default value if any
|
|
287
|
+
"""
|
|
288
|
+
dialect = self._get_dialect(driver)
|
|
289
|
+
adbc_driver = cast("AdbcDriver", driver)
|
|
290
|
+
|
|
291
|
+
if dialect == "sqlite":
|
|
292
|
+
result = adbc_driver.execute(f"PRAGMA table_info({table})")
|
|
293
|
+
return [
|
|
294
|
+
{
|
|
295
|
+
"column_name": row["name"] if isinstance(row, dict) else row[1],
|
|
296
|
+
"data_type": row["type"] if isinstance(row, dict) else row[2],
|
|
297
|
+
"nullable": not (row["notnull"] if isinstance(row, dict) else row[3]),
|
|
298
|
+
"default_value": row["dflt_value"] if isinstance(row, dict) else row[4],
|
|
299
|
+
}
|
|
300
|
+
for row in result.data or []
|
|
301
|
+
]
|
|
302
|
+
|
|
303
|
+
if dialect == "postgres":
|
|
304
|
+
schema_name = schema or "public"
|
|
305
|
+
sql = """
|
|
306
|
+
SELECT
|
|
307
|
+
a.attname::text AS column_name,
|
|
308
|
+
pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
|
|
309
|
+
CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable,
|
|
310
|
+
pg_catalog.pg_get_expr(d.adbin, d.adrelid)::text AS column_default
|
|
311
|
+
FROM pg_catalog.pg_attribute a
|
|
312
|
+
JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
|
|
313
|
+
JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
|
|
314
|
+
LEFT JOIN pg_catalog.pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
|
|
315
|
+
WHERE c.relname = ?
|
|
316
|
+
AND n.nspname = ?
|
|
317
|
+
AND a.attnum > 0
|
|
318
|
+
AND NOT a.attisdropped
|
|
319
|
+
ORDER BY a.attnum
|
|
320
|
+
"""
|
|
321
|
+
result = adbc_driver.execute(sql, (table, schema_name))
|
|
322
|
+
return result.data or []
|
|
323
|
+
|
|
324
|
+
if schema:
|
|
325
|
+
sql = """
|
|
326
|
+
SELECT column_name, data_type, is_nullable, column_default
|
|
327
|
+
FROM information_schema.columns
|
|
328
|
+
WHERE table_name = ? AND table_schema = ?
|
|
329
|
+
ORDER BY ordinal_position
|
|
330
|
+
"""
|
|
331
|
+
result = adbc_driver.execute(sql, (table, schema))
|
|
332
|
+
else:
|
|
333
|
+
sql = """
|
|
334
|
+
SELECT column_name, data_type, is_nullable, column_default
|
|
335
|
+
FROM information_schema.columns
|
|
336
|
+
WHERE table_name = ?
|
|
337
|
+
ORDER BY ordinal_position
|
|
338
|
+
"""
|
|
339
|
+
result = adbc_driver.execute(sql, (table,))
|
|
340
|
+
|
|
341
|
+
return result.data or []
|
|
342
|
+
|
|
271
343
|
def list_available_features(self) -> "list[str]":
|
|
272
344
|
"""List available feature flags across all supported dialects.
|
|
273
345
|
|
sqlspec/adapters/adbc/driver.py
CHANGED
|
@@ -7,20 +7,32 @@ database dialects, parameter style conversion, and transaction management.
|
|
|
7
7
|
import contextlib
|
|
8
8
|
import datetime
|
|
9
9
|
import decimal
|
|
10
|
-
from typing import TYPE_CHECKING, Any,
|
|
10
|
+
from typing import TYPE_CHECKING, Any, cast
|
|
11
11
|
|
|
12
|
-
from adbc_driver_manager.dbapi import DatabaseError, IntegrityError, OperationalError, ProgrammingError
|
|
13
12
|
from sqlglot import exp
|
|
14
13
|
|
|
15
14
|
from sqlspec.adapters.adbc.data_dictionary import AdbcDataDictionary
|
|
16
15
|
from sqlspec.adapters.adbc.type_converter import ADBCTypeConverter
|
|
17
16
|
from sqlspec.core.cache import get_cache_config
|
|
18
17
|
from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
|
|
18
|
+
from sqlspec.core.result import create_arrow_result
|
|
19
19
|
from sqlspec.core.statement import SQL, StatementConfig
|
|
20
20
|
from sqlspec.driver import SyncDriverAdapterBase
|
|
21
|
-
from sqlspec.exceptions import
|
|
21
|
+
from sqlspec.exceptions import (
|
|
22
|
+
CheckViolationError,
|
|
23
|
+
DatabaseConnectionError,
|
|
24
|
+
DataError,
|
|
25
|
+
ForeignKeyViolationError,
|
|
26
|
+
IntegrityError,
|
|
27
|
+
NotNullViolationError,
|
|
28
|
+
SQLParsingError,
|
|
29
|
+
SQLSpecError,
|
|
30
|
+
TransactionError,
|
|
31
|
+
UniqueViolationError,
|
|
32
|
+
)
|
|
22
33
|
from sqlspec.typing import Empty
|
|
23
34
|
from sqlspec.utils.logging import get_logger
|
|
35
|
+
from sqlspec.utils.module_loader import ensure_pyarrow
|
|
24
36
|
|
|
25
37
|
if TYPE_CHECKING:
|
|
26
38
|
from contextlib import AbstractContextManager
|
|
@@ -28,9 +40,12 @@ if TYPE_CHECKING:
|
|
|
28
40
|
from adbc_driver_manager.dbapi import Cursor
|
|
29
41
|
|
|
30
42
|
from sqlspec.adapters.adbc._types import AdbcConnection
|
|
31
|
-
from sqlspec.
|
|
43
|
+
from sqlspec.builder import QueryBuilder
|
|
44
|
+
from sqlspec.core import Statement, StatementFilter
|
|
45
|
+
from sqlspec.core.result import ArrowResult, SQLResult
|
|
32
46
|
from sqlspec.driver import ExecutionResult
|
|
33
47
|
from sqlspec.driver._sync import SyncDataDictionaryBase
|
|
48
|
+
from sqlspec.typing import StatementParameters
|
|
34
49
|
|
|
35
50
|
__all__ = ("AdbcCursor", "AdbcDriver", "AdbcExceptionHandler", "get_adbc_statement_config")
|
|
36
51
|
|
|
@@ -329,7 +344,7 @@ class AdbcCursor:
|
|
|
329
344
|
|
|
330
345
|
def __init__(self, connection: "AdbcConnection") -> None:
|
|
331
346
|
self.connection = connection
|
|
332
|
-
self.cursor:
|
|
347
|
+
self.cursor: Cursor | None = None
|
|
333
348
|
|
|
334
349
|
def __enter__(self) -> "Cursor":
|
|
335
350
|
self.cursor = self.connection.cursor()
|
|
@@ -342,7 +357,11 @@ class AdbcCursor:
|
|
|
342
357
|
|
|
343
358
|
|
|
344
359
|
class AdbcExceptionHandler:
|
|
345
|
-
"""Context manager for handling database exceptions.
|
|
360
|
+
"""Context manager for handling ADBC database exceptions.
|
|
361
|
+
|
|
362
|
+
ADBC propagates underlying database errors. Exception mapping
|
|
363
|
+
depends on the specific ADBC driver being used.
|
|
364
|
+
"""
|
|
346
365
|
|
|
347
366
|
__slots__ = ()
|
|
348
367
|
|
|
@@ -350,40 +369,118 @@ class AdbcExceptionHandler:
|
|
|
350
369
|
return None
|
|
351
370
|
|
|
352
371
|
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
372
|
+
_ = exc_tb
|
|
353
373
|
if exc_type is None:
|
|
354
374
|
return
|
|
375
|
+
self._map_adbc_exception(exc_val)
|
|
355
376
|
|
|
356
|
-
|
|
357
|
-
|
|
358
|
-
|
|
359
|
-
|
|
360
|
-
|
|
361
|
-
|
|
362
|
-
|
|
363
|
-
|
|
364
|
-
|
|
365
|
-
|
|
366
|
-
|
|
367
|
-
|
|
368
|
-
|
|
369
|
-
|
|
370
|
-
|
|
371
|
-
|
|
372
|
-
|
|
373
|
-
|
|
374
|
-
|
|
375
|
-
|
|
376
|
-
|
|
377
|
-
|
|
378
|
-
|
|
379
|
-
|
|
380
|
-
|
|
381
|
-
|
|
382
|
-
|
|
383
|
-
|
|
384
|
-
|
|
385
|
-
|
|
386
|
-
|
|
377
|
+
def _map_adbc_exception(self, e: Any) -> None:
|
|
378
|
+
"""Map ADBC exception to SQLSpec exception.
|
|
379
|
+
|
|
380
|
+
ADBC drivers may expose SQLSTATE codes or driver-specific codes.
|
|
381
|
+
|
|
382
|
+
Args:
|
|
383
|
+
e: ADBC exception instance
|
|
384
|
+
"""
|
|
385
|
+
sqlstate = getattr(e, "sqlstate", None)
|
|
386
|
+
|
|
387
|
+
if sqlstate:
|
|
388
|
+
self._map_sqlstate_exception(e, sqlstate)
|
|
389
|
+
else:
|
|
390
|
+
self._map_message_based_exception(e)
|
|
391
|
+
|
|
392
|
+
def _map_sqlstate_exception(self, e: Any, sqlstate: str) -> None:
|
|
393
|
+
"""Map SQLSTATE code to exception.
|
|
394
|
+
|
|
395
|
+
Args:
|
|
396
|
+
e: Exception instance
|
|
397
|
+
sqlstate: SQLSTATE error code
|
|
398
|
+
"""
|
|
399
|
+
if sqlstate == "23505":
|
|
400
|
+
self._raise_unique_violation(e)
|
|
401
|
+
elif sqlstate == "23503":
|
|
402
|
+
self._raise_foreign_key_violation(e)
|
|
403
|
+
elif sqlstate == "23502":
|
|
404
|
+
self._raise_not_null_violation(e)
|
|
405
|
+
elif sqlstate == "23514":
|
|
406
|
+
self._raise_check_violation(e)
|
|
407
|
+
elif sqlstate.startswith("23"):
|
|
408
|
+
self._raise_integrity_error(e)
|
|
409
|
+
elif sqlstate.startswith("42"):
|
|
410
|
+
self._raise_parsing_error(e)
|
|
411
|
+
elif sqlstate.startswith("08"):
|
|
412
|
+
self._raise_connection_error(e)
|
|
413
|
+
elif sqlstate.startswith("40"):
|
|
414
|
+
self._raise_transaction_error(e)
|
|
415
|
+
elif sqlstate.startswith("22"):
|
|
416
|
+
self._raise_data_error(e)
|
|
417
|
+
else:
|
|
418
|
+
self._raise_generic_error(e)
|
|
419
|
+
|
|
420
|
+
def _map_message_based_exception(self, e: Any) -> None:
|
|
421
|
+
"""Map exception using message-based detection.
|
|
422
|
+
|
|
423
|
+
Args:
|
|
424
|
+
e: Exception instance
|
|
425
|
+
"""
|
|
426
|
+
error_msg = str(e).lower()
|
|
427
|
+
|
|
428
|
+
if "unique" in error_msg or "duplicate" in error_msg:
|
|
429
|
+
self._raise_unique_violation(e)
|
|
430
|
+
elif "foreign key" in error_msg:
|
|
431
|
+
self._raise_foreign_key_violation(e)
|
|
432
|
+
elif "not null" in error_msg or "null value" in error_msg:
|
|
433
|
+
self._raise_not_null_violation(e)
|
|
434
|
+
elif "check constraint" in error_msg:
|
|
435
|
+
self._raise_check_violation(e)
|
|
436
|
+
elif "constraint" in error_msg:
|
|
437
|
+
self._raise_integrity_error(e)
|
|
438
|
+
elif "syntax" in error_msg:
|
|
439
|
+
self._raise_parsing_error(e)
|
|
440
|
+
elif "connection" in error_msg or "connect" in error_msg:
|
|
441
|
+
self._raise_connection_error(e)
|
|
442
|
+
else:
|
|
443
|
+
self._raise_generic_error(e)
|
|
444
|
+
|
|
445
|
+
def _raise_unique_violation(self, e: Any) -> None:
|
|
446
|
+
msg = f"ADBC unique constraint violation: {e}"
|
|
447
|
+
raise UniqueViolationError(msg) from e
|
|
448
|
+
|
|
449
|
+
def _raise_foreign_key_violation(self, e: Any) -> None:
|
|
450
|
+
msg = f"ADBC foreign key constraint violation: {e}"
|
|
451
|
+
raise ForeignKeyViolationError(msg) from e
|
|
452
|
+
|
|
453
|
+
def _raise_not_null_violation(self, e: Any) -> None:
|
|
454
|
+
msg = f"ADBC not-null constraint violation: {e}"
|
|
455
|
+
raise NotNullViolationError(msg) from e
|
|
456
|
+
|
|
457
|
+
def _raise_check_violation(self, e: Any) -> None:
|
|
458
|
+
msg = f"ADBC check constraint violation: {e}"
|
|
459
|
+
raise CheckViolationError(msg) from e
|
|
460
|
+
|
|
461
|
+
def _raise_integrity_error(self, e: Any) -> None:
|
|
462
|
+
msg = f"ADBC integrity constraint violation: {e}"
|
|
463
|
+
raise IntegrityError(msg) from e
|
|
464
|
+
|
|
465
|
+
def _raise_parsing_error(self, e: Any) -> None:
|
|
466
|
+
msg = f"ADBC SQL parsing error: {e}"
|
|
467
|
+
raise SQLParsingError(msg) from e
|
|
468
|
+
|
|
469
|
+
def _raise_connection_error(self, e: Any) -> None:
|
|
470
|
+
msg = f"ADBC connection error: {e}"
|
|
471
|
+
raise DatabaseConnectionError(msg) from e
|
|
472
|
+
|
|
473
|
+
def _raise_transaction_error(self, e: Any) -> None:
|
|
474
|
+
msg = f"ADBC transaction error: {e}"
|
|
475
|
+
raise TransactionError(msg) from e
|
|
476
|
+
|
|
477
|
+
def _raise_data_error(self, e: Any) -> None:
|
|
478
|
+
msg = f"ADBC data error: {e}"
|
|
479
|
+
raise DataError(msg) from e
|
|
480
|
+
|
|
481
|
+
def _raise_generic_error(self, e: Any) -> None:
|
|
482
|
+
msg = f"ADBC database error: {e}"
|
|
483
|
+
raise SQLSpecError(msg) from e
|
|
387
484
|
|
|
388
485
|
|
|
389
486
|
class AdbcDriver(SyncDriverAdapterBase):
|
|
@@ -398,8 +495,8 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
398
495
|
def __init__(
|
|
399
496
|
self,
|
|
400
497
|
connection: "AdbcConnection",
|
|
401
|
-
statement_config: "
|
|
402
|
-
driver_features: "
|
|
498
|
+
statement_config: "StatementConfig | None" = None,
|
|
499
|
+
driver_features: "dict[str, Any] | None" = None,
|
|
403
500
|
) -> None:
|
|
404
501
|
self._detected_dialect = self._get_dialect(connection)
|
|
405
502
|
|
|
@@ -412,19 +509,7 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
412
509
|
|
|
413
510
|
super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
|
|
414
511
|
self.dialect = statement_config.dialect
|
|
415
|
-
self._data_dictionary:
|
|
416
|
-
|
|
417
|
-
@staticmethod
|
|
418
|
-
def _ensure_pyarrow_installed() -> None:
|
|
419
|
-
"""Ensure PyArrow is installed.
|
|
420
|
-
|
|
421
|
-
Raises:
|
|
422
|
-
MissingDependencyError: If PyArrow is not installed
|
|
423
|
-
"""
|
|
424
|
-
from sqlspec.typing import PYARROW_INSTALLED
|
|
425
|
-
|
|
426
|
-
if not PYARROW_INSTALLED:
|
|
427
|
-
raise MissingDependencyError(package="pyarrow", install_package="arrow")
|
|
512
|
+
self._data_dictionary: SyncDataDictionaryBase | None = None
|
|
428
513
|
|
|
429
514
|
@staticmethod
|
|
430
515
|
def _get_dialect(connection: "AdbcConnection") -> str:
|
|
@@ -480,12 +565,13 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
480
565
|
parameters: Any,
|
|
481
566
|
statement_config: "StatementConfig",
|
|
482
567
|
is_many: bool = False,
|
|
483
|
-
prepared_statement:
|
|
568
|
+
prepared_statement: Any | None = None,
|
|
484
569
|
) -> Any:
|
|
485
570
|
"""Prepare parameters with cast-aware type coercion for ADBC.
|
|
486
571
|
|
|
487
572
|
For PostgreSQL, applies cast-aware parameter processing using metadata from the compiled statement.
|
|
488
573
|
This allows proper handling of JSONB casts and other type conversions.
|
|
574
|
+
Respects driver_features['enable_cast_detection'] configuration.
|
|
489
575
|
|
|
490
576
|
Args:
|
|
491
577
|
parameters: Parameters in any format
|
|
@@ -496,7 +582,9 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
496
582
|
Returns:
|
|
497
583
|
Parameters with cast-aware type coercion applied
|
|
498
584
|
"""
|
|
499
|
-
|
|
585
|
+
enable_cast_detection = self.driver_features.get("enable_cast_detection", True)
|
|
586
|
+
|
|
587
|
+
if enable_cast_detection and prepared_statement and self.dialect in {"postgres", "postgresql"} and not is_many:
|
|
500
588
|
parameter_casts = self._get_parameter_casts(prepared_statement)
|
|
501
589
|
postgres_compatible = self._handle_postgres_empty_parameters(parameters)
|
|
502
590
|
return self._prepare_parameters_with_casts(postgres_compatible, parameter_casts, statement_config)
|
|
@@ -524,6 +612,7 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
524
612
|
"""Prepare parameters with cast-aware type coercion.
|
|
525
613
|
|
|
526
614
|
Uses type coercion map for non-dict types and dialect-aware dict handling.
|
|
615
|
+
Respects driver_features configuration for JSON serialization backend.
|
|
527
616
|
|
|
528
617
|
Args:
|
|
529
618
|
parameters: Parameter values (list, tuple, or scalar)
|
|
@@ -533,7 +622,9 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
533
622
|
Returns:
|
|
534
623
|
Parameters with cast-aware type coercion applied
|
|
535
624
|
"""
|
|
536
|
-
from sqlspec.
|
|
625
|
+
from sqlspec.utils.serializers import to_json
|
|
626
|
+
|
|
627
|
+
json_encoder = self.driver_features.get("json_serializer", to_json)
|
|
537
628
|
|
|
538
629
|
if isinstance(parameters, (list, tuple)):
|
|
539
630
|
result: list[Any] = []
|
|
@@ -541,7 +632,7 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
541
632
|
cast_type = parameter_casts.get(idx, "").upper()
|
|
542
633
|
if cast_type in {"JSON", "JSONB", "TYPE.JSON", "TYPE.JSONB"}:
|
|
543
634
|
if isinstance(param, dict):
|
|
544
|
-
result.append(
|
|
635
|
+
result.append(json_encoder(param))
|
|
545
636
|
else:
|
|
546
637
|
result.append(param)
|
|
547
638
|
elif isinstance(param, dict):
|
|
@@ -575,7 +666,7 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
575
666
|
"""
|
|
576
667
|
return AdbcExceptionHandler()
|
|
577
668
|
|
|
578
|
-
def _try_special_handling(self, cursor: "Cursor", statement: SQL) -> "
|
|
669
|
+
def _try_special_handling(self, cursor: "Cursor", statement: SQL) -> "SQLResult | None":
|
|
579
670
|
"""Handle special operations.
|
|
580
671
|
|
|
581
672
|
Args:
|
|
@@ -668,7 +759,7 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
668
759
|
column_names = [col[0] for col in cursor.description or []]
|
|
669
760
|
|
|
670
761
|
if fetched_data and isinstance(fetched_data[0], tuple):
|
|
671
|
-
dict_data: list[dict[Any, Any]] = [dict(zip(column_names, row)) for row in fetched_data]
|
|
762
|
+
dict_data: list[dict[Any, Any]] = [dict(zip(column_names, row, strict=False)) for row in fetched_data]
|
|
672
763
|
else:
|
|
673
764
|
dict_data = fetched_data # type: ignore[assignment]
|
|
674
765
|
|
|
@@ -764,3 +855,80 @@ class AdbcDriver(SyncDriverAdapterBase):
|
|
|
764
855
|
if self._data_dictionary is None:
|
|
765
856
|
self._data_dictionary = AdbcDataDictionary()
|
|
766
857
|
return self._data_dictionary
|
|
858
|
+
|
|
859
|
+
def select_to_arrow(
|
|
860
|
+
self,
|
|
861
|
+
statement: "Statement | QueryBuilder",
|
|
862
|
+
/,
|
|
863
|
+
*parameters: "StatementParameters | StatementFilter",
|
|
864
|
+
statement_config: "StatementConfig | None" = None,
|
|
865
|
+
return_format: str = "table",
|
|
866
|
+
native_only: bool = False,
|
|
867
|
+
batch_size: int | None = None,
|
|
868
|
+
arrow_schema: Any = None,
|
|
869
|
+
**kwargs: Any,
|
|
870
|
+
) -> "ArrowResult":
|
|
871
|
+
"""Execute query and return results as Apache Arrow (ADBC native path).
|
|
872
|
+
|
|
873
|
+
ADBC provides zero-copy Arrow support via cursor.fetch_arrow_table().
|
|
874
|
+
This is 5-10x faster than the conversion path for large datasets.
|
|
875
|
+
|
|
876
|
+
Args:
|
|
877
|
+
statement: SQL statement, string, or QueryBuilder
|
|
878
|
+
*parameters: Query parameters or filters
|
|
879
|
+
statement_config: Optional statement configuration override
|
|
880
|
+
return_format: "table" for pyarrow.Table (default), "batch" for RecordBatch
|
|
881
|
+
native_only: Ignored for ADBC (always uses native path)
|
|
882
|
+
batch_size: Batch size hint (for future streaming implementation)
|
|
883
|
+
arrow_schema: Optional pyarrow.Schema for type casting
|
|
884
|
+
**kwargs: Additional keyword arguments
|
|
885
|
+
|
|
886
|
+
Returns:
|
|
887
|
+
ArrowResult with native Arrow data
|
|
888
|
+
|
|
889
|
+
Raises:
|
|
890
|
+
MissingDependencyError: If pyarrow not installed
|
|
891
|
+
SQLExecutionError: If query execution fails
|
|
892
|
+
|
|
893
|
+
Example:
|
|
894
|
+
>>> result = driver.select_to_arrow(
|
|
895
|
+
... "SELECT * FROM users WHERE age > $1", 18
|
|
896
|
+
... )
|
|
897
|
+
>>> df = result.to_pandas() # Fast zero-copy conversion
|
|
898
|
+
"""
|
|
899
|
+
ensure_pyarrow()
|
|
900
|
+
|
|
901
|
+
import pyarrow as pa
|
|
902
|
+
|
|
903
|
+
# Prepare statement
|
|
904
|
+
config = statement_config or self.statement_config
|
|
905
|
+
prepared_statement = self.prepare_statement(statement, parameters, statement_config=config, kwargs=kwargs)
|
|
906
|
+
|
|
907
|
+
# Use ADBC cursor for native Arrow
|
|
908
|
+
with self.with_cursor(self.connection) as cursor, self.handle_database_exceptions():
|
|
909
|
+
if cursor is None:
|
|
910
|
+
msg = "Failed to create cursor"
|
|
911
|
+
raise DatabaseConnectionError(msg)
|
|
912
|
+
|
|
913
|
+
# Get compiled SQL and parameters
|
|
914
|
+
sql, driver_params = self._get_compiled_sql(prepared_statement, config)
|
|
915
|
+
|
|
916
|
+
# Execute query
|
|
917
|
+
cursor.execute(sql, driver_params or ())
|
|
918
|
+
|
|
919
|
+
# Fetch as Arrow table (zero-copy!)
|
|
920
|
+
arrow_table = cursor.fetch_arrow_table()
|
|
921
|
+
|
|
922
|
+
# Apply schema casting if requested
|
|
923
|
+
if arrow_schema is not None:
|
|
924
|
+
arrow_table = arrow_table.cast(arrow_schema)
|
|
925
|
+
|
|
926
|
+
# Convert to batch if requested
|
|
927
|
+
if return_format == "batch":
|
|
928
|
+
batches = arrow_table.to_batches()
|
|
929
|
+
arrow_data: Any = batches[0] if batches else pa.RecordBatch.from_pydict({})
|
|
930
|
+
else:
|
|
931
|
+
arrow_data = arrow_table
|
|
932
|
+
|
|
933
|
+
# Create ArrowResult
|
|
934
|
+
return create_arrow_result(statement=prepared_statement, data=arrow_data, rows_affected=arrow_data.num_rows)
|