sqlspec 0.11.1__py3-none-any.whl → 0.12.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 (155) hide show
  1. sqlspec/__init__.py +16 -3
  2. sqlspec/_serialization.py +3 -10
  3. sqlspec/_sql.py +1147 -0
  4. sqlspec/_typing.py +343 -41
  5. sqlspec/adapters/adbc/__init__.py +2 -6
  6. sqlspec/adapters/adbc/config.py +474 -149
  7. sqlspec/adapters/adbc/driver.py +330 -621
  8. sqlspec/adapters/aiosqlite/__init__.py +2 -6
  9. sqlspec/adapters/aiosqlite/config.py +143 -57
  10. sqlspec/adapters/aiosqlite/driver.py +269 -431
  11. sqlspec/adapters/asyncmy/__init__.py +3 -8
  12. sqlspec/adapters/asyncmy/config.py +247 -202
  13. sqlspec/adapters/asyncmy/driver.py +218 -436
  14. sqlspec/adapters/asyncpg/__init__.py +4 -7
  15. sqlspec/adapters/asyncpg/config.py +329 -176
  16. sqlspec/adapters/asyncpg/driver.py +417 -487
  17. sqlspec/adapters/bigquery/__init__.py +2 -2
  18. sqlspec/adapters/bigquery/config.py +407 -0
  19. sqlspec/adapters/bigquery/driver.py +600 -553
  20. sqlspec/adapters/duckdb/__init__.py +4 -1
  21. sqlspec/adapters/duckdb/config.py +432 -321
  22. sqlspec/adapters/duckdb/driver.py +392 -406
  23. sqlspec/adapters/oracledb/__init__.py +3 -8
  24. sqlspec/adapters/oracledb/config.py +625 -0
  25. sqlspec/adapters/oracledb/driver.py +548 -921
  26. sqlspec/adapters/psqlpy/__init__.py +4 -7
  27. sqlspec/adapters/psqlpy/config.py +372 -203
  28. sqlspec/adapters/psqlpy/driver.py +197 -533
  29. sqlspec/adapters/psycopg/__init__.py +3 -8
  30. sqlspec/adapters/psycopg/config.py +741 -0
  31. sqlspec/adapters/psycopg/driver.py +734 -694
  32. sqlspec/adapters/sqlite/__init__.py +2 -6
  33. sqlspec/adapters/sqlite/config.py +146 -81
  34. sqlspec/adapters/sqlite/driver.py +242 -405
  35. sqlspec/base.py +220 -784
  36. sqlspec/config.py +354 -0
  37. sqlspec/driver/__init__.py +22 -0
  38. sqlspec/driver/_async.py +252 -0
  39. sqlspec/driver/_common.py +338 -0
  40. sqlspec/driver/_sync.py +261 -0
  41. sqlspec/driver/mixins/__init__.py +17 -0
  42. sqlspec/driver/mixins/_pipeline.py +523 -0
  43. sqlspec/driver/mixins/_result_utils.py +122 -0
  44. sqlspec/driver/mixins/_sql_translator.py +35 -0
  45. sqlspec/driver/mixins/_storage.py +993 -0
  46. sqlspec/driver/mixins/_type_coercion.py +131 -0
  47. sqlspec/exceptions.py +299 -7
  48. sqlspec/extensions/aiosql/__init__.py +10 -0
  49. sqlspec/extensions/aiosql/adapter.py +474 -0
  50. sqlspec/extensions/litestar/__init__.py +1 -6
  51. sqlspec/extensions/litestar/_utils.py +1 -5
  52. sqlspec/extensions/litestar/config.py +5 -6
  53. sqlspec/extensions/litestar/handlers.py +13 -12
  54. sqlspec/extensions/litestar/plugin.py +22 -24
  55. sqlspec/extensions/litestar/providers.py +37 -55
  56. sqlspec/loader.py +528 -0
  57. sqlspec/service/__init__.py +3 -0
  58. sqlspec/service/base.py +24 -0
  59. sqlspec/service/pagination.py +26 -0
  60. sqlspec/statement/__init__.py +21 -0
  61. sqlspec/statement/builder/__init__.py +54 -0
  62. sqlspec/statement/builder/_ddl_utils.py +119 -0
  63. sqlspec/statement/builder/_parsing_utils.py +135 -0
  64. sqlspec/statement/builder/base.py +328 -0
  65. sqlspec/statement/builder/ddl.py +1379 -0
  66. sqlspec/statement/builder/delete.py +80 -0
  67. sqlspec/statement/builder/insert.py +274 -0
  68. sqlspec/statement/builder/merge.py +95 -0
  69. sqlspec/statement/builder/mixins/__init__.py +65 -0
  70. sqlspec/statement/builder/mixins/_aggregate_functions.py +151 -0
  71. sqlspec/statement/builder/mixins/_case_builder.py +91 -0
  72. sqlspec/statement/builder/mixins/_common_table_expr.py +91 -0
  73. sqlspec/statement/builder/mixins/_delete_from.py +34 -0
  74. sqlspec/statement/builder/mixins/_from.py +61 -0
  75. sqlspec/statement/builder/mixins/_group_by.py +119 -0
  76. sqlspec/statement/builder/mixins/_having.py +35 -0
  77. sqlspec/statement/builder/mixins/_insert_from_select.py +48 -0
  78. sqlspec/statement/builder/mixins/_insert_into.py +36 -0
  79. sqlspec/statement/builder/mixins/_insert_values.py +69 -0
  80. sqlspec/statement/builder/mixins/_join.py +110 -0
  81. sqlspec/statement/builder/mixins/_limit_offset.py +53 -0
  82. sqlspec/statement/builder/mixins/_merge_clauses.py +405 -0
  83. sqlspec/statement/builder/mixins/_order_by.py +46 -0
  84. sqlspec/statement/builder/mixins/_pivot.py +82 -0
  85. sqlspec/statement/builder/mixins/_returning.py +37 -0
  86. sqlspec/statement/builder/mixins/_select_columns.py +60 -0
  87. sqlspec/statement/builder/mixins/_set_ops.py +122 -0
  88. sqlspec/statement/builder/mixins/_unpivot.py +80 -0
  89. sqlspec/statement/builder/mixins/_update_from.py +54 -0
  90. sqlspec/statement/builder/mixins/_update_set.py +91 -0
  91. sqlspec/statement/builder/mixins/_update_table.py +29 -0
  92. sqlspec/statement/builder/mixins/_where.py +374 -0
  93. sqlspec/statement/builder/mixins/_window_functions.py +86 -0
  94. sqlspec/statement/builder/protocols.py +20 -0
  95. sqlspec/statement/builder/select.py +206 -0
  96. sqlspec/statement/builder/update.py +178 -0
  97. sqlspec/statement/filters.py +571 -0
  98. sqlspec/statement/parameters.py +736 -0
  99. sqlspec/statement/pipelines/__init__.py +67 -0
  100. sqlspec/statement/pipelines/analyzers/__init__.py +9 -0
  101. sqlspec/statement/pipelines/analyzers/_analyzer.py +649 -0
  102. sqlspec/statement/pipelines/base.py +315 -0
  103. sqlspec/statement/pipelines/context.py +119 -0
  104. sqlspec/statement/pipelines/result_types.py +41 -0
  105. sqlspec/statement/pipelines/transformers/__init__.py +8 -0
  106. sqlspec/statement/pipelines/transformers/_expression_simplifier.py +256 -0
  107. sqlspec/statement/pipelines/transformers/_literal_parameterizer.py +623 -0
  108. sqlspec/statement/pipelines/transformers/_remove_comments.py +66 -0
  109. sqlspec/statement/pipelines/transformers/_remove_hints.py +81 -0
  110. sqlspec/statement/pipelines/validators/__init__.py +23 -0
  111. sqlspec/statement/pipelines/validators/_dml_safety.py +275 -0
  112. sqlspec/statement/pipelines/validators/_parameter_style.py +297 -0
  113. sqlspec/statement/pipelines/validators/_performance.py +703 -0
  114. sqlspec/statement/pipelines/validators/_security.py +990 -0
  115. sqlspec/statement/pipelines/validators/base.py +67 -0
  116. sqlspec/statement/result.py +527 -0
  117. sqlspec/statement/splitter.py +701 -0
  118. sqlspec/statement/sql.py +1198 -0
  119. sqlspec/storage/__init__.py +15 -0
  120. sqlspec/storage/backends/__init__.py +0 -0
  121. sqlspec/storage/backends/base.py +166 -0
  122. sqlspec/storage/backends/fsspec.py +315 -0
  123. sqlspec/storage/backends/obstore.py +464 -0
  124. sqlspec/storage/protocol.py +170 -0
  125. sqlspec/storage/registry.py +315 -0
  126. sqlspec/typing.py +157 -36
  127. sqlspec/utils/correlation.py +155 -0
  128. sqlspec/utils/deprecation.py +3 -6
  129. sqlspec/utils/fixtures.py +6 -11
  130. sqlspec/utils/logging.py +135 -0
  131. sqlspec/utils/module_loader.py +45 -43
  132. sqlspec/utils/serializers.py +4 -0
  133. sqlspec/utils/singleton.py +6 -8
  134. sqlspec/utils/sync_tools.py +15 -27
  135. sqlspec/utils/text.py +58 -26
  136. {sqlspec-0.11.1.dist-info → sqlspec-0.12.0.dist-info}/METADATA +97 -26
  137. sqlspec-0.12.0.dist-info/RECORD +145 -0
  138. sqlspec/adapters/bigquery/config/__init__.py +0 -3
  139. sqlspec/adapters/bigquery/config/_common.py +0 -40
  140. sqlspec/adapters/bigquery/config/_sync.py +0 -87
  141. sqlspec/adapters/oracledb/config/__init__.py +0 -9
  142. sqlspec/adapters/oracledb/config/_asyncio.py +0 -186
  143. sqlspec/adapters/oracledb/config/_common.py +0 -131
  144. sqlspec/adapters/oracledb/config/_sync.py +0 -186
  145. sqlspec/adapters/psycopg/config/__init__.py +0 -19
  146. sqlspec/adapters/psycopg/config/_async.py +0 -169
  147. sqlspec/adapters/psycopg/config/_common.py +0 -56
  148. sqlspec/adapters/psycopg/config/_sync.py +0 -168
  149. sqlspec/filters.py +0 -331
  150. sqlspec/mixins.py +0 -305
  151. sqlspec/statement.py +0 -378
  152. sqlspec-0.11.1.dist-info/RECORD +0 -69
  153. {sqlspec-0.11.1.dist-info → sqlspec-0.12.0.dist-info}/WHEEL +0 -0
  154. {sqlspec-0.11.1.dist-info → sqlspec-0.12.0.dist-info}/licenses/LICENSE +0 -0
  155. {sqlspec-0.11.1.dist-info → sqlspec-0.12.0.dist-info}/licenses/NOTICE +0 -0
@@ -1,167 +1,412 @@
1
+ """ADBC database configuration using TypedDict for better maintainability."""
2
+
3
+ import logging
1
4
  from contextlib import contextmanager
2
- from dataclasses import dataclass, field
3
- from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast
5
+ from dataclasses import replace
6
+ from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional
4
7
 
5
8
  from sqlspec.adapters.adbc.driver import AdbcConnection, AdbcDriver
6
- from sqlspec.base import NoPoolSyncConfig
9
+ from sqlspec.config import NoPoolSyncConfig
7
10
  from sqlspec.exceptions import ImproperConfigurationError
8
- from sqlspec.typing import Empty, EmptyType
11
+ from sqlspec.statement.sql import SQLConfig
12
+ from sqlspec.typing import DictRow, Empty
9
13
  from sqlspec.utils.module_loader import import_string
10
14
 
11
15
  if TYPE_CHECKING:
12
16
  from collections.abc import Generator
17
+ from contextlib import AbstractContextManager
18
+
19
+ from sqlglot.dialects.dialect import DialectType
20
+
21
+ logger = logging.getLogger("sqlspec.adapters.adbc")
13
22
 
23
+ CONNECTION_FIELDS = frozenset(
24
+ {
25
+ "uri",
26
+ "driver_name",
27
+ "db_kwargs",
28
+ "conn_kwargs",
29
+ "adbc_driver_manager_entrypoint",
30
+ "autocommit",
31
+ "isolation_level",
32
+ "batch_size",
33
+ "query_timeout",
34
+ "connection_timeout",
35
+ "ssl_mode",
36
+ "ssl_cert",
37
+ "ssl_key",
38
+ "ssl_ca",
39
+ "username",
40
+ "password",
41
+ "token",
42
+ "project_id",
43
+ "dataset_id",
44
+ "account",
45
+ "warehouse",
46
+ "database",
47
+ "schema",
48
+ "role",
49
+ "authorization_header",
50
+ "grpc_options",
51
+ }
52
+ )
14
53
 
15
- __all__ = ("AdbcConfig",)
54
+ __all__ = ("CONNECTION_FIELDS", "AdbcConfig")
16
55
 
17
56
 
18
- @dataclass
19
- class AdbcConfig(NoPoolSyncConfig["AdbcConnection", "AdbcDriver"]):
20
- """Configuration for ADBC connections.
57
+ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
58
+ """Enhanced ADBC configuration with universal database connectivity.
21
59
 
22
- This class provides configuration options for ADBC database connections using the
23
- ADBC Driver Manager.([1](https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html))
60
+ ADBC (Arrow Database Connectivity) provides a unified interface for connecting
61
+ to multiple database systems with high-performance Arrow-native data transfer.
62
+
63
+ This configuration supports:
64
+ - Universal driver detection and loading
65
+ - High-performance Arrow data streaming
66
+ - Bulk ingestion operations
67
+ - Multiple database backends (PostgreSQL, SQLite, DuckDB, BigQuery, Snowflake, etc.)
68
+ - Intelligent driver path resolution
69
+ - Cloud database integrations
24
70
  """
25
71
 
26
- uri: "Union[str, EmptyType]" = Empty
27
- """Database URI"""
28
- driver_name: "Union[str, EmptyType]" = Empty
29
- """Full dotted path to the ADBC driver's connect function (e.g., 'adbc_driver_sqlite.dbapi.connect')"""
30
- db_kwargs: "Optional[dict[str, Any]]" = None
31
- """Additional database-specific connection parameters"""
32
- conn_kwargs: "Optional[dict[str, Any]]" = None
33
- """Additional database-specific connection parameters"""
34
- connection_type: "type[AdbcConnection]" = field(init=False, default_factory=lambda: AdbcConnection)
35
- """Type of the connection object"""
36
- driver_type: "type[AdbcDriver]" = field(init=False, default_factory=lambda: AdbcDriver) # type: ignore[type-abstract,unused-ignore]
37
- """Type of the driver object"""
38
- pool_instance: None = field(init=False, default=None, hash=False)
39
- """No connection pool is used for ADBC connections"""
40
-
41
- def _set_adbc(self) -> str:
42
- """Identify the driver type based on the URI (if provided) or preset driver name.
72
+ __slots__ = (
73
+ "_dialect",
74
+ "account",
75
+ "adbc_driver_manager_entrypoint",
76
+ "authorization_header",
77
+ "autocommit",
78
+ "batch_size",
79
+ "conn_kwargs",
80
+ "connection_timeout",
81
+ "database",
82
+ "dataset_id",
83
+ "db_kwargs",
84
+ "default_row_type",
85
+ "driver_name",
86
+ "extras",
87
+ "grpc_options",
88
+ "isolation_level",
89
+ "on_connection_create",
90
+ "password",
91
+ "pool_instance",
92
+ "project_id",
93
+ "query_timeout",
94
+ "role",
95
+ "schema",
96
+ "ssl_ca",
97
+ "ssl_cert",
98
+ "ssl_key",
99
+ "ssl_mode",
100
+ "statement_config",
101
+ "token",
102
+ "uri",
103
+ "username",
104
+ "warehouse",
105
+ )
43
106
 
44
- Raises:
45
- ImproperConfigurationError: If the driver name is not recognized or supported.
107
+ is_async: ClassVar[bool] = False
108
+ supports_connection_pooling: ClassVar[bool] = False
109
+ driver_type: type[AdbcDriver] = AdbcDriver
110
+ connection_type: type[AdbcConnection] = AdbcConnection
111
+
112
+ # Parameter style support information - dynamic based on driver
113
+ # These are used as defaults when driver cannot be determined
114
+ supported_parameter_styles: ClassVar[tuple[str, ...]] = ("qmark",)
115
+ """ADBC parameter styles depend on the underlying driver."""
116
+
117
+ preferred_parameter_style: ClassVar[str] = "qmark"
118
+ """ADBC default parameter style is ? (qmark)."""
119
+
120
+ def __init__(
121
+ self,
122
+ statement_config: Optional[SQLConfig] = None,
123
+ default_row_type: type[DictRow] = DictRow,
124
+ on_connection_create: Optional[Callable[[AdbcConnection], None]] = None,
125
+ # Core connection parameters
126
+ uri: Optional[str] = None,
127
+ driver_name: Optional[str] = None,
128
+ # Database-specific parameters
129
+ db_kwargs: Optional[dict[str, Any]] = None,
130
+ conn_kwargs: Optional[dict[str, Any]] = None,
131
+ # Driver-specific configurations
132
+ adbc_driver_manager_entrypoint: Optional[str] = None,
133
+ # Connection options
134
+ autocommit: Optional[bool] = None,
135
+ isolation_level: Optional[str] = None,
136
+ # Performance options
137
+ batch_size: Optional[int] = None,
138
+ query_timeout: Optional[int] = None,
139
+ connection_timeout: Optional[int] = None,
140
+ # Security options
141
+ ssl_mode: Optional[str] = None,
142
+ ssl_cert: Optional[str] = None,
143
+ ssl_key: Optional[str] = None,
144
+ ssl_ca: Optional[str] = None,
145
+ # Authentication
146
+ username: Optional[str] = None,
147
+ password: Optional[str] = None,
148
+ token: Optional[str] = None,
149
+ # Cloud-specific options
150
+ project_id: Optional[str] = None,
151
+ dataset_id: Optional[str] = None,
152
+ account: Optional[str] = None,
153
+ warehouse: Optional[str] = None,
154
+ database: Optional[str] = None,
155
+ schema: Optional[str] = None,
156
+ role: Optional[str] = None,
157
+ # Flight SQL specific
158
+ authorization_header: Optional[str] = None,
159
+ grpc_options: Optional[dict[str, Any]] = None,
160
+ **kwargs: Any,
161
+ ) -> None:
162
+ """Initialize ADBC configuration with universal connectivity features.
163
+
164
+ Args:
165
+ statement_config: Default SQL statement configuration
166
+ instrumentation: Instrumentation configuration
167
+ default_row_type: Default row type for results
168
+ on_connection_create: Callback executed when connection is created
169
+ uri: Database URI (e.g., 'postgresql://...', 'sqlite://...', 'bigquery://...')
170
+ driver_name: Full dotted path to ADBC driver connect function or driver alias
171
+ driver: Backward compatibility alias for driver_name
172
+ db_kwargs: Additional database-specific connection parameters
173
+ conn_kwargs: Additional connection-specific parameters
174
+ adbc_driver_manager_entrypoint: Override for driver manager entrypoint
175
+ autocommit: Enable autocommit mode
176
+ isolation_level: Transaction isolation level
177
+ batch_size: Batch size for bulk operations
178
+ query_timeout: Query timeout in seconds
179
+ connection_timeout: Connection timeout in seconds
180
+ ssl_mode: SSL mode for secure connections
181
+ ssl_cert: SSL certificate path
182
+ ssl_key: SSL private key path
183
+ ssl_ca: SSL certificate authority path
184
+ username: Database username
185
+ password: Database password
186
+ token: Authentication token (for cloud services)
187
+ project_id: Project ID (BigQuery)
188
+ dataset_id: Dataset ID (BigQuery)
189
+ account: Account identifier (Snowflake)
190
+ warehouse: Warehouse name (Snowflake)
191
+ database: Database name
192
+ schema: Schema name
193
+ role: Role name (Snowflake)
194
+ authorization_header: Authorization header for Flight SQL
195
+ grpc_options: gRPC specific options for Flight SQL
196
+ **kwargs: Additional parameters (stored in extras)
197
+
198
+ Example:
199
+ >>> # PostgreSQL via ADBC
200
+ >>> config = AdbcConfig(
201
+ ... uri="postgresql://user:pass@localhost/db",
202
+ ... driver_name="adbc_driver_postgresql",
203
+ ... )
204
+
205
+ >>> # DuckDB via ADBC
206
+ >>> config = AdbcConfig(
207
+ ... uri="duckdb://mydata.db",
208
+ ... driver_name="duckdb",
209
+ ... db_kwargs={"read_only": False},
210
+ ... )
211
+
212
+ >>> # BigQuery via ADBC
213
+ >>> config = AdbcConfig(
214
+ ... driver_name="bigquery",
215
+ ... project_id="my-project",
216
+ ... dataset_id="my_dataset",
217
+ ... )
218
+ """
219
+
220
+ # Store connection parameters as instance attributes
221
+ self.uri = uri
222
+ self.driver_name = driver_name
223
+ self.db_kwargs = db_kwargs
224
+ self.conn_kwargs = conn_kwargs
225
+ self.adbc_driver_manager_entrypoint = adbc_driver_manager_entrypoint
226
+ self.autocommit = autocommit
227
+ self.isolation_level = isolation_level
228
+ self.batch_size = batch_size
229
+ self.query_timeout = query_timeout
230
+ self.connection_timeout = connection_timeout
231
+ self.ssl_mode = ssl_mode
232
+ self.ssl_cert = ssl_cert
233
+ self.ssl_key = ssl_key
234
+ self.ssl_ca = ssl_ca
235
+ self.username = username
236
+ self.password = password
237
+ self.token = token
238
+ self.project_id = project_id
239
+ self.dataset_id = dataset_id
240
+ self.account = account
241
+ self.warehouse = warehouse
242
+ self.database = database
243
+ self.schema = schema
244
+ self.role = role
245
+ self.authorization_header = authorization_header
246
+ self.grpc_options = grpc_options
247
+
248
+ self.extras = kwargs or {}
249
+
250
+ # Store other config
251
+ self.statement_config = statement_config or SQLConfig()
252
+ self.default_row_type = default_row_type
253
+ self.on_connection_create = on_connection_create
254
+ self._dialect: DialectType = None
255
+ super().__init__()
256
+
257
+ def _resolve_driver_name(self) -> str:
258
+ """Resolve and normalize the ADBC driver name.
259
+
260
+ Supports both full driver paths and convenient aliases.
46
261
 
47
262
  Returns:
48
- str: The driver name to be used for the connection.
263
+ The normalized driver connect function path.
264
+
265
+ Raises:
266
+ ImproperConfigurationError: If driver cannot be determined.
49
267
  """
268
+ driver_name = self.driver_name
269
+ uri = self.uri
50
270
 
51
- if isinstance(self.driver_name, str):
52
- if self.driver_name != "adbc_driver_sqlite.dbapi.connect" and self.driver_name in {
53
- "sqlite",
54
- "sqlite3",
55
- "adbc_driver_sqlite",
56
- }:
57
- self.driver_name = "adbc_driver_sqlite.dbapi.connect"
58
- elif self.driver_name != "adbc_driver_duckdb.dbapi.connect" and self.driver_name in {
59
- "duckdb",
60
- "adbc_driver_duckdb",
61
- }:
62
- self.driver_name = "adbc_driver_duckdb.dbapi.connect"
63
- elif self.driver_name != "adbc_driver_postgresql.dbapi.connect" and self.driver_name in {
64
- "postgres",
65
- "adbc_driver_postgresql",
66
- "postgresql",
67
- "pg",
68
- }:
69
- self.driver_name = "adbc_driver_postgresql.dbapi.connect"
70
- elif self.driver_name != "adbc_driver_snowflake.dbapi.connect" and self.driver_name in {
71
- "snowflake",
72
- "adbc_driver_snowflake",
73
- "sf",
74
- }:
75
- self.driver_name = "adbc_driver_snowflake.dbapi.connect"
76
- elif self.driver_name != "adbc_driver_bigquery.dbapi.connect" and self.driver_name in {
77
- "bigquery",
78
- "adbc_driver_bigquery",
79
- "bq",
80
- }:
81
- self.driver_name = "adbc_driver_bigquery.dbapi.connect"
82
- elif self.driver_name != "adbc_driver_flightsql.dbapi.connect" and self.driver_name in {
83
- "flightsql",
84
- "adbc_driver_flightsql",
85
- "grpc",
86
- }:
87
- self.driver_name = "adbc_driver_flightsql.dbapi.connect"
88
- return self.driver_name
89
-
90
- # If driver_name wasn't explicit, try to determine from URI
91
- if isinstance(self.uri, str) and self.uri.startswith("postgresql://"):
92
- self.driver_name = "adbc_driver_postgresql.dbapi.connect"
93
- elif isinstance(self.uri, str) and self.uri.startswith("sqlite://"):
94
- self.driver_name = "adbc_driver_sqlite.dbapi.connect"
95
- elif isinstance(self.uri, str) and self.uri.startswith("grpc://"):
96
- self.driver_name = "adbc_driver_flightsql.dbapi.connect"
97
- elif isinstance(self.uri, str) and self.uri.startswith("snowflake://"):
98
- self.driver_name = "adbc_driver_snowflake.dbapi.connect"
99
- elif isinstance(self.uri, str) and self.uri.startswith("bigquery://"):
100
- self.driver_name = "adbc_driver_bigquery.dbapi.connect"
101
- elif isinstance(self.uri, str) and self.uri.startswith("duckdb://"):
102
- self.driver_name = "adbc_driver_duckdb.dbapi.connect"
103
-
104
- # Check if we successfully determined a driver name
105
- if self.driver_name is Empty or not isinstance(self.driver_name, str):
106
- msg = (
107
- "Could not determine ADBC driver connect path. Please specify 'driver_name' "
108
- "(e.g., 'adbc_driver_sqlite.dbapi.connect') or provide a supported 'uri'. "
109
- f"URI: {self.uri}, Driver Name: {self.driver_name}"
110
- )
111
- raise ImproperConfigurationError(msg)
112
- return self.driver_name
271
+ # If explicit driver path is provided, normalize it
272
+ if isinstance(driver_name, str):
273
+ # Handle convenience aliases
274
+ driver_aliases = {
275
+ "sqlite": "adbc_driver_sqlite.dbapi.connect",
276
+ "sqlite3": "adbc_driver_sqlite.dbapi.connect",
277
+ "adbc_driver_sqlite": "adbc_driver_sqlite.dbapi.connect",
278
+ "duckdb": "adbc_driver_duckdb.dbapi.connect",
279
+ "adbc_driver_duckdb": "adbc_driver_duckdb.dbapi.connect",
280
+ "postgres": "adbc_driver_postgresql.dbapi.connect",
281
+ "postgresql": "adbc_driver_postgresql.dbapi.connect",
282
+ "pg": "adbc_driver_postgresql.dbapi.connect",
283
+ "adbc_driver_postgresql": "adbc_driver_postgresql.dbapi.connect",
284
+ "snowflake": "adbc_driver_snowflake.dbapi.connect",
285
+ "sf": "adbc_driver_snowflake.dbapi.connect",
286
+ "adbc_driver_snowflake": "adbc_driver_snowflake.dbapi.connect",
287
+ "bigquery": "adbc_driver_bigquery.dbapi.connect",
288
+ "bq": "adbc_driver_bigquery.dbapi.connect",
289
+ "adbc_driver_bigquery": "adbc_driver_bigquery.dbapi.connect",
290
+ "flightsql": "adbc_driver_flightsql.dbapi.connect",
291
+ "adbc_driver_flightsql": "adbc_driver_flightsql.dbapi.connect",
292
+ "grpc": "adbc_driver_flightsql.dbapi.connect",
293
+ }
113
294
 
114
- @property
115
- def connection_config_dict(self) -> "dict[str, Any]":
116
- """Return the connection configuration as a dict.
295
+ resolved_driver = driver_aliases.get(driver_name, driver_name)
296
+
297
+ # Ensure it ends with .dbapi.connect
298
+ if not resolved_driver.endswith(".dbapi.connect"):
299
+ resolved_driver = f"{resolved_driver}.dbapi.connect"
300
+
301
+ return resolved_driver
117
302
 
118
- Omits the 'uri' key for known in-memory database types.
303
+ # Auto-detect from URI if no explicit driver
304
+ if isinstance(uri, str):
305
+ if uri.startswith("postgresql://"):
306
+ return "adbc_driver_postgresql.dbapi.connect"
307
+ if uri.startswith("sqlite://"):
308
+ return "adbc_driver_sqlite.dbapi.connect"
309
+ if uri.startswith("duckdb://"):
310
+ return "adbc_driver_duckdb.dbapi.connect"
311
+ if uri.startswith("grpc://"):
312
+ return "adbc_driver_flightsql.dbapi.connect"
313
+ if uri.startswith("snowflake://"):
314
+ return "adbc_driver_snowflake.dbapi.connect"
315
+ if uri.startswith("bigquery://"):
316
+ return "adbc_driver_bigquery.dbapi.connect"
317
+
318
+ # Could not determine driver
319
+ msg = (
320
+ "Could not determine ADBC driver connect path. Please specify 'driver_name' "
321
+ "(e.g., 'adbc_driver_postgresql' or 'postgresql') or provide a supported 'uri'. "
322
+ f"URI: {uri}, Driver Name: {driver_name}"
323
+ )
324
+ raise ImproperConfigurationError(msg)
325
+
326
+ def _get_connect_func(self) -> Callable[..., AdbcConnection]:
327
+ """Get the ADBC driver connect function.
119
328
 
120
329
  Returns:
121
- A string keyed dict of config kwargs for the adbc_driver_manager.dbapi.connect function.
330
+ The driver connect function.
331
+
332
+ Raises:
333
+ ImproperConfigurationError: If driver cannot be loaded.
122
334
  """
123
- config = {}
124
- db_kwargs = self.db_kwargs or {}
125
- conn_kwargs = self.conn_kwargs or {}
126
- if isinstance(self.uri, str) and self.uri.startswith("sqlite://"):
127
- db_kwargs["uri"] = self.uri.replace("sqlite://", "")
128
- elif isinstance(self.uri, str) and self.uri.startswith("duckdb://"):
129
- db_kwargs["path"] = self.uri.replace("duckdb://", "")
130
- elif isinstance(self.uri, str):
131
- db_kwargs["uri"] = self.uri
132
- if isinstance(self.driver_name, str) and self.driver_name.startswith("adbc_driver_bigquery"):
133
- config["db_kwargs"] = db_kwargs
134
- else:
135
- config = db_kwargs
136
- if conn_kwargs:
137
- config["conn_kwargs"] = conn_kwargs
138
- return config
335
+ driver_path = self._resolve_driver_name()
139
336
 
140
- def _get_connect_func(self) -> "Callable[..., AdbcConnection]":
141
- self._set_adbc()
142
- driver_path = cast("str", self.driver_name)
143
337
  try:
144
338
  connect_func = import_string(driver_path)
145
339
  except ImportError as e:
146
- # Check if the error is likely due to missing suffix and try again
147
- if ".dbapi.connect" not in driver_path:
148
- try:
149
- driver_path += ".dbapi.connect"
150
- connect_func = import_string(driver_path)
151
- except ImportError as e2:
152
- msg = f"Failed to import ADBC connect function from '{self.driver_name}' or '{driver_path}'. Is the driver installed and the path correct? Original error: {e} / {e2}"
153
- raise ImproperConfigurationError(msg) from e2
154
- else:
155
- # Original import failed, and suffix was already present or added
156
- msg = f"Failed to import ADBC connect function from '{driver_path}'. Is the driver installed and the path correct? Original error: {e}"
157
- raise ImproperConfigurationError(msg) from e
340
+ driver_path_with_suffix = f"{driver_path}.dbapi.connect"
341
+ try:
342
+ connect_func = import_string(driver_path_with_suffix)
343
+ except ImportError as e2:
344
+ msg = (
345
+ f"Failed to import ADBC connect function from '{driver_path}' or "
346
+ f"'{driver_path_with_suffix}'. Is the driver installed? "
347
+ f"Original errors: {e} / {e2}"
348
+ )
349
+ raise ImproperConfigurationError(msg) from e2
350
+
158
351
  if not callable(connect_func):
159
352
  msg = f"The path '{driver_path}' did not resolve to a callable function."
160
353
  raise ImproperConfigurationError(msg)
354
+
161
355
  return connect_func # type: ignore[no-any-return]
162
356
 
163
- def create_connection(self) -> "AdbcConnection":
164
- """Create and return a new database connection using the specific driver.
357
+ def _get_dialect(self) -> "DialectType":
358
+ """Get the SQL dialect type based on the ADBC driver.
359
+
360
+ Returns:
361
+ The SQL dialect type for the ADBC driver.
362
+ """
363
+ try:
364
+ driver_path = self._resolve_driver_name()
365
+ except ImproperConfigurationError:
366
+ return None
367
+
368
+ dialect_map = {
369
+ "postgres": "postgres",
370
+ "sqlite": "sqlite",
371
+ "duckdb": "duckdb",
372
+ "bigquery": "bigquery",
373
+ "snowflake": "snowflake",
374
+ "flightsql": "sqlite",
375
+ "grpc": "sqlite",
376
+ }
377
+ for keyword, dialect in dialect_map.items():
378
+ if keyword in driver_path:
379
+ return dialect
380
+ return None
381
+
382
+ def _get_parameter_styles(self) -> tuple[tuple[str, ...], str]:
383
+ """Get parameter styles based on the underlying driver.
384
+
385
+ Returns:
386
+ Tuple of (supported_parameter_styles, preferred_parameter_style)
387
+ """
388
+ try:
389
+ driver_path = self._resolve_driver_name()
390
+
391
+ # Map driver paths to parameter styles
392
+ if "postgresql" in driver_path:
393
+ return (("numeric",), "numeric") # $1, $2, ...
394
+ if "sqlite" in driver_path:
395
+ return (("qmark", "named_colon"), "qmark") # ? or :name
396
+ if "duckdb" in driver_path:
397
+ return (("qmark", "numeric"), "qmark") # ? or $1
398
+ if "bigquery" in driver_path:
399
+ return (("named_at",), "named_at") # @name
400
+ if "snowflake" in driver_path:
401
+ return (("qmark", "numeric"), "qmark") # ? or :1
402
+
403
+ except Exception:
404
+ # If we can't determine driver, use defaults
405
+ return (self.supported_parameter_styles, self.preferred_parameter_style)
406
+ return (("qmark",), "qmark")
407
+
408
+ def create_connection(self) -> AdbcConnection:
409
+ """Create and return a new ADBC connection using the specified driver.
165
410
 
166
411
  Returns:
167
412
  A new ADBC connection instance.
@@ -169,39 +414,119 @@ class AdbcConfig(NoPoolSyncConfig["AdbcConnection", "AdbcDriver"]):
169
414
  Raises:
170
415
  ImproperConfigurationError: If the connection could not be established.
171
416
  """
417
+
172
418
  try:
173
419
  connect_func = self._get_connect_func()
174
- return connect_func(**self.connection_config_dict)
420
+ connection = connect_func(**self.connection_config_dict)
421
+
422
+ if self.on_connection_create:
423
+ self.on_connection_create(connection)
175
424
  except Exception as e:
176
- # Include driver name in error message for better context
177
- driver_name = self.driver_name if isinstance(self.driver_name, str) else "Unknown/Missing"
178
- # Use the potentially modified driver_path from _get_connect_func if available,
179
- # otherwise fallback to self.driver_name for the error message.
180
- # This requires _get_connect_func to potentially return the used path or store it.
181
- # For simplicity now, we stick to self.driver_name in the message.
182
- msg = f"Could not configure the ADBC connection using driver path '{driver_name}'. Error: {e!s}"
425
+ driver_name = self.driver_name or "Unknown"
426
+ msg = f"Could not configure ADBC connection using driver '{driver_name}'. Error: {e}"
183
427
  raise ImproperConfigurationError(msg) from e
428
+ return connection
184
429
 
185
430
  @contextmanager
186
- def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[AdbcConnection, None, None]":
187
- """Create and provide a database connection using the specific driver.
431
+ def provide_connection(self, *args: Any, **kwargs: Any) -> "Generator[AdbcConnection, None, None]":
432
+ """Provide an ADBC connection context manager.
433
+
434
+ Args:
435
+ *args: Additional arguments.
436
+ **kwargs: Additional keyword arguments.
188
437
 
189
438
  Yields:
190
- Connection: A database connection instance.
439
+ An ADBC connection instance.
191
440
  """
192
-
193
441
  connection = self.create_connection()
194
442
  try:
195
443
  yield connection
196
444
  finally:
197
445
  connection.close()
198
446
 
199
- @contextmanager
200
- def provide_session(self, *args: Any, **kwargs: Any) -> "Generator[AdbcDriver, None, None]":
201
- """Create and provide a database session.
447
+ def provide_session(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[AdbcDriver]":
448
+ """Provide an ADBC driver session context manager.
202
449
 
203
- Yields:
204
- An ADBC driver instance with an active connection.
450
+ Args:
451
+ *args: Additional arguments.
452
+ **kwargs: Additional keyword arguments.
453
+
454
+ Returns:
455
+ A context manager that yields an AdbcDriver instance.
456
+ """
457
+
458
+ @contextmanager
459
+ def session_manager() -> "Generator[AdbcDriver, None, None]":
460
+ with self.provide_connection(*args, **kwargs) as connection:
461
+ # Get parameter styles based on the actual driver
462
+ supported_styles, preferred_style = self._get_parameter_styles()
463
+
464
+ # Create statement config with parameter style info if not already set
465
+ statement_config = self.statement_config
466
+ if statement_config.allowed_parameter_styles is None:
467
+ statement_config = replace(
468
+ statement_config,
469
+ allowed_parameter_styles=supported_styles,
470
+ target_parameter_style=preferred_style,
471
+ )
472
+
473
+ driver = self.driver_type(connection=connection, config=statement_config)
474
+ yield driver
475
+
476
+ return session_manager()
477
+
478
+ @property
479
+ def connection_config_dict(self) -> dict[str, Any]:
480
+ """Get the connection configuration dictionary.
481
+
482
+ Returns:
483
+ The connection configuration dictionary.
205
484
  """
206
- with self.provide_connection(*args, **kwargs) as connection:
207
- yield self.driver_type(connection)
485
+ # Gather non-None connection parameters
486
+ config = {
487
+ field: getattr(self, field)
488
+ for field in CONNECTION_FIELDS
489
+ if getattr(self, field, None) is not None and getattr(self, field) is not Empty
490
+ }
491
+
492
+ # Merge extras parameters
493
+ config.update(self.extras)
494
+
495
+ # Process URI based on driver type
496
+ if "driver_name" in config:
497
+ driver_name = config["driver_name"]
498
+
499
+ if "uri" in config:
500
+ uri = config["uri"]
501
+
502
+ # SQLite: strip sqlite:// prefix
503
+ if driver_name in {"sqlite", "sqlite3", "adbc_driver_sqlite"} and uri.startswith("sqlite://"): # pyright: ignore
504
+ config["uri"] = uri[9:] # Remove "sqlite://" # pyright: ignore
505
+
506
+ # DuckDB: convert uri to path
507
+ elif driver_name in {"duckdb", "adbc_driver_duckdb"} and uri.startswith("duckdb://"): # pyright: ignore
508
+ config["path"] = uri[9:] # Remove "duckdb://" # pyright: ignore
509
+ config.pop("uri", None)
510
+
511
+ # BigQuery: wrap certain parameters in db_kwargs
512
+ if driver_name in {"bigquery", "bq", "adbc_driver_bigquery"}:
513
+ bigquery_params = ["project_id", "dataset_id", "token"]
514
+ db_kwargs = config.get("db_kwargs", {})
515
+
516
+ for param in bigquery_params:
517
+ if param in config and param != "db_kwargs":
518
+ db_kwargs[param] = config.pop(param) # pyright: ignore
519
+
520
+ if db_kwargs:
521
+ config["db_kwargs"] = db_kwargs
522
+
523
+ # For other drivers (like PostgreSQL), merge db_kwargs into top level
524
+ elif "db_kwargs" in config and driver_name not in {"bigquery", "bq", "adbc_driver_bigquery"}:
525
+ db_kwargs = config.pop("db_kwargs")
526
+ if isinstance(db_kwargs, dict):
527
+ config.update(db_kwargs)
528
+
529
+ # Remove driver_name from config as it's not a connection parameter
530
+ config.pop("driver_name", None)
531
+
532
+ return config