sqlspec 0.9.1__py3-none-any.whl → 0.10.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 (47) hide show
  1. sqlspec/__init__.py +2 -1
  2. sqlspec/adapters/adbc/__init__.py +2 -1
  3. sqlspec/adapters/adbc/config.py +7 -13
  4. sqlspec/adapters/adbc/driver.py +37 -30
  5. sqlspec/adapters/aiosqlite/__init__.py +2 -1
  6. sqlspec/adapters/aiosqlite/config.py +10 -12
  7. sqlspec/adapters/aiosqlite/driver.py +36 -31
  8. sqlspec/adapters/asyncmy/__init__.py +2 -1
  9. sqlspec/adapters/asyncmy/driver.py +34 -31
  10. sqlspec/adapters/asyncpg/config.py +1 -3
  11. sqlspec/adapters/asyncpg/driver.py +7 -3
  12. sqlspec/adapters/bigquery/__init__.py +4 -0
  13. sqlspec/adapters/bigquery/config/__init__.py +3 -0
  14. sqlspec/adapters/bigquery/config/_common.py +40 -0
  15. sqlspec/adapters/bigquery/config/_sync.py +87 -0
  16. sqlspec/adapters/bigquery/driver.py +701 -0
  17. sqlspec/adapters/duckdb/__init__.py +2 -1
  18. sqlspec/adapters/duckdb/config.py +17 -18
  19. sqlspec/adapters/duckdb/driver.py +38 -30
  20. sqlspec/adapters/oracledb/__init__.py +8 -1
  21. sqlspec/adapters/oracledb/config/_asyncio.py +7 -8
  22. sqlspec/adapters/oracledb/config/_sync.py +6 -7
  23. sqlspec/adapters/oracledb/driver.py +65 -62
  24. sqlspec/adapters/psqlpy/__init__.py +9 -0
  25. sqlspec/adapters/psqlpy/config.py +5 -5
  26. sqlspec/adapters/psqlpy/driver.py +34 -28
  27. sqlspec/adapters/psycopg/__init__.py +8 -1
  28. sqlspec/adapters/psycopg/config/__init__.py +10 -0
  29. sqlspec/adapters/psycopg/config/_async.py +6 -7
  30. sqlspec/adapters/psycopg/config/_sync.py +7 -8
  31. sqlspec/adapters/psycopg/driver.py +63 -53
  32. sqlspec/adapters/sqlite/__init__.py +2 -1
  33. sqlspec/adapters/sqlite/config.py +12 -11
  34. sqlspec/adapters/sqlite/driver.py +36 -29
  35. sqlspec/base.py +1 -66
  36. sqlspec/exceptions.py +9 -0
  37. sqlspec/extensions/litestar/config.py +3 -11
  38. sqlspec/extensions/litestar/handlers.py +2 -1
  39. sqlspec/extensions/litestar/plugin.py +4 -2
  40. sqlspec/mixins.py +156 -0
  41. sqlspec/typing.py +19 -1
  42. {sqlspec-0.9.1.dist-info → sqlspec-0.10.0.dist-info}/METADATA +8 -3
  43. sqlspec-0.10.0.dist-info/RECORD +67 -0
  44. sqlspec-0.9.1.dist-info/RECORD +0 -61
  45. {sqlspec-0.9.1.dist-info → sqlspec-0.10.0.dist-info}/WHEEL +0 -0
  46. {sqlspec-0.9.1.dist-info → sqlspec-0.10.0.dist-info}/licenses/LICENSE +0 -0
  47. {sqlspec-0.9.1.dist-info → sqlspec-0.10.0.dist-info}/licenses/NOTICE +0 -0
@@ -1,7 +1,8 @@
1
1
  from sqlspec.adapters.duckdb.config import DuckDBConfig
2
- from sqlspec.adapters.duckdb.driver import DuckDBDriver
2
+ from sqlspec.adapters.duckdb.driver import DuckDBConnection, DuckDBDriver
3
3
 
4
4
  __all__ = (
5
5
  "DuckDBConfig",
6
+ "DuckDBConnection",
6
7
  "DuckDBDriver",
7
8
  )
@@ -2,10 +2,9 @@ from contextlib import contextmanager
2
2
  from dataclasses import dataclass, field
3
3
  from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast
4
4
 
5
- from duckdb import DuckDBPyConnection
6
5
  from typing_extensions import Literal, NotRequired, TypedDict
7
6
 
8
- from sqlspec.adapters.duckdb.driver import DuckDBDriver
7
+ from sqlspec.adapters.duckdb.driver import DuckDBConnection, DuckDBDriver
9
8
  from sqlspec.base import NoPoolSyncConfig
10
9
  from sqlspec.exceptions import ImproperConfigurationError
11
10
  from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
@@ -69,7 +68,7 @@ class SecretConfig(TypedDict):
69
68
 
70
69
 
71
70
  @dataclass
72
- class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
71
+ class DuckDBConfig(NoPoolSyncConfig["DuckDBConnection", "DuckDBDriver"]):
73
72
  """Configuration for DuckDB database connections.
74
73
 
75
74
  This class provides configuration options for DuckDB database connections, wrapping all parameters
@@ -96,10 +95,10 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
96
95
  """A dictionary of secrets to store in the connection for later retrieval."""
97
96
  auto_update_extensions: "bool" = False
98
97
  """Whether to automatically update on connection creation"""
99
- on_connection_create: "Optional[Callable[[DuckDBPyConnection], Optional[DuckDBPyConnection]]]" = None
98
+ on_connection_create: "Optional[Callable[[DuckDBConnection], Optional[DuckDBConnection]]]" = None
100
99
  """A callable to be called after the connection is created."""
101
- connection_type: "type[DuckDBPyConnection]" = field(init=False, default_factory=lambda: DuckDBPyConnection)
102
- """The type of connection to create. Defaults to DuckDBPyConnection."""
100
+ connection_type: "type[DuckDBConnection]" = field(init=False, default_factory=lambda: DuckDBConnection)
101
+ """The type of connection to create. Defaults to DuckDBConnection."""
103
102
  driver_type: "type[DuckDBDriver]" = field(init=False, default_factory=lambda: DuckDBDriver) # type: ignore[type-abstract,unused-ignore]
104
103
  """The type of driver to use. Defaults to DuckDBDriver."""
105
104
  pool_instance: "None" = field(init=False, default=None)
@@ -139,7 +138,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
139
138
  raise ImproperConfigurationError(msg) from e
140
139
  self.extensions.extend(config_exts)
141
140
 
142
- def _configure_connection(self, connection: "DuckDBPyConnection") -> None:
141
+ def _configure_connection(self, connection: "DuckDBConnection") -> None:
143
142
  """Configure the connection.
144
143
 
145
144
  Args:
@@ -148,7 +147,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
148
147
  for key, value in cast("dict[str,Any]", self.config).items():
149
148
  connection.execute(f"SET {key}='{value}'")
150
149
 
151
- def _configure_extensions(self, connection: "DuckDBPyConnection") -> None:
150
+ def _configure_extensions(self, connection: "DuckDBConnection") -> None:
152
151
  """Configure extensions for the connection.
153
152
 
154
153
  Args:
@@ -165,7 +164,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
165
164
  connection.execute("update extensions")
166
165
 
167
166
  @staticmethod
168
- def _secret_exists(connection: "DuckDBPyConnection", name: "str") -> bool:
167
+ def _secret_exists(connection: "DuckDBConnection", name: "str") -> bool:
169
168
  """Check if a secret exists in the connection.
170
169
 
171
170
  Args:
@@ -179,7 +178,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
179
178
  return results is not None
180
179
 
181
180
  @classmethod
182
- def _is_community_extension(cls, connection: "DuckDBPyConnection", name: "str") -> bool:
181
+ def _is_community_extension(cls, connection: "DuckDBConnection", name: "str") -> bool:
183
182
  """Check if an extension is a community extension.
184
183
 
185
184
  Args:
@@ -195,7 +194,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
195
194
  return results is None
196
195
 
197
196
  @classmethod
198
- def _extension_installed(cls, connection: "DuckDBPyConnection", name: "str") -> bool:
197
+ def _extension_installed(cls, connection: "DuckDBConnection", name: "str") -> bool:
199
198
  """Check if a extension exists in the connection.
200
199
 
201
200
  Args:
@@ -211,7 +210,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
211
210
  return results is not None
212
211
 
213
212
  @classmethod
214
- def _extension_loaded(cls, connection: "DuckDBPyConnection", name: "str") -> bool:
213
+ def _extension_loaded(cls, connection: "DuckDBConnection", name: "str") -> bool:
215
214
  """Check if a extension is loaded in the connection.
216
215
 
217
216
  Args:
@@ -229,7 +228,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
229
228
  @classmethod
230
229
  def _configure_secrets(
231
230
  cls,
232
- connection: "DuckDBPyConnection",
231
+ connection: "DuckDBConnection",
233
232
  secrets: "Sequence[SecretConfig]",
234
233
  ) -> None:
235
234
  """Configure persistent secrets for the connection.
@@ -258,7 +257,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
258
257
  raise ImproperConfigurationError(msg) from e
259
258
 
260
259
  @classmethod
261
- def _configure_extension(cls, connection: "DuckDBPyConnection", extension: "ExtensionConfig") -> None:
260
+ def _configure_extension(cls, connection: "DuckDBConnection", extension: "ExtensionConfig") -> None:
262
261
  """Configure a single extension for the connection.
263
262
 
264
263
  Args:
@@ -320,6 +319,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
320
319
  "auto_update_extensions",
321
320
  "driver_type",
322
321
  "connection_type",
322
+ "connection_instance",
323
323
  },
324
324
  convert_nested=False,
325
325
  )
@@ -327,7 +327,7 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
327
327
  config["database"] = ":memory:"
328
328
  return config
329
329
 
330
- def create_connection(self) -> "DuckDBPyConnection":
330
+ def create_connection(self) -> "DuckDBConnection":
331
331
  """Create and return a new database connection with configured extensions.
332
332
 
333
333
  Returns:
@@ -349,11 +349,10 @@ class DuckDBConfig(NoPoolSyncConfig["DuckDBPyConnection", "DuckDBDriver"]):
349
349
  except Exception as e:
350
350
  msg = f"Could not configure the DuckDB connection. Error: {e!s}"
351
351
  raise ImproperConfigurationError(msg) from e
352
- else:
353
- return connection
352
+ return connection
354
353
 
355
354
  @contextmanager
356
- def provide_connection(self, *args: Any, **kwargs: Any) -> "Generator[DuckDBPyConnection, None, None]":
355
+ def provide_connection(self, *args: Any, **kwargs: Any) -> "Generator[DuckDBConnection, None, None]":
357
356
  """Create and provide a database connection.
358
357
 
359
358
  Yields:
@@ -1,37 +1,45 @@
1
1
  from contextlib import contextmanager
2
2
  from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
3
3
 
4
- from sqlspec.base import SyncArrowBulkOperationsMixin, SyncDriverAdapterProtocol, T
4
+ from duckdb import DuckDBPyConnection
5
+
6
+ from sqlspec.base import SyncDriverAdapterProtocol
7
+ from sqlspec.mixins import SQLTranslatorMixin, SyncArrowBulkOperationsMixin
8
+ from sqlspec.typing import ArrowTable, StatementParameterType
5
9
 
6
10
  if TYPE_CHECKING:
7
11
  from collections.abc import Generator, Sequence
8
12
 
9
- from duckdb import DuckDBPyConnection
13
+ from sqlspec.typing import ArrowTable, ModelDTOT, StatementParameterType, T
10
14
 
11
- from sqlspec.typing import ArrowTable, ModelDTOT, StatementParameterType
15
+ __all__ = ("DuckDBConnection", "DuckDBDriver")
12
16
 
13
- __all__ = ("DuckDBDriver",)
17
+ DuckDBConnection = DuckDBPyConnection
14
18
 
15
19
 
16
- class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDriverAdapterProtocol["DuckDBPyConnection"]):
20
+ class DuckDBDriver(
21
+ SyncArrowBulkOperationsMixin["DuckDBConnection"],
22
+ SQLTranslatorMixin["DuckDBConnection"],
23
+ SyncDriverAdapterProtocol["DuckDBConnection"],
24
+ ):
17
25
  """DuckDB Sync Driver Adapter."""
18
26
 
19
- connection: "DuckDBPyConnection"
27
+ connection: "DuckDBConnection"
20
28
  use_cursor: bool = True
21
29
  dialect: str = "duckdb"
22
30
 
23
- def __init__(self, connection: "DuckDBPyConnection", use_cursor: bool = True) -> None:
31
+ def __init__(self, connection: "DuckDBConnection", use_cursor: bool = True) -> None:
24
32
  self.connection = connection
25
33
  self.use_cursor = use_cursor
26
34
 
27
35
  # --- Helper Methods --- #
28
- def _cursor(self, connection: "DuckDBPyConnection") -> "DuckDBPyConnection":
36
+ def _cursor(self, connection: "DuckDBConnection") -> "DuckDBConnection":
29
37
  if self.use_cursor:
30
38
  return connection.cursor()
31
39
  return connection
32
40
 
33
41
  @contextmanager
34
- def _with_cursor(self, connection: "DuckDBPyConnection") -> "Generator[DuckDBPyConnection, None, None]":
42
+ def _with_cursor(self, connection: "DuckDBConnection") -> "Generator[DuckDBConnection, None, None]":
35
43
  if self.use_cursor:
36
44
  cursor = self._cursor(connection)
37
45
  try:
@@ -49,7 +57,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
49
57
  parameters: "Optional[StatementParameterType]" = None,
50
58
  /,
51
59
  *,
52
- connection: "Optional[DuckDBPyConnection]" = None,
60
+ connection: "Optional[DuckDBConnection]" = None,
53
61
  schema_type: None = None,
54
62
  **kwargs: Any,
55
63
  ) -> "Sequence[dict[str, Any]]": ...
@@ -60,7 +68,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
60
68
  parameters: "Optional[StatementParameterType]" = None,
61
69
  /,
62
70
  *,
63
- connection: "Optional[DuckDBPyConnection]" = None,
71
+ connection: "Optional[DuckDBConnection]" = None,
64
72
  schema_type: "type[ModelDTOT]",
65
73
  **kwargs: Any,
66
74
  ) -> "Sequence[ModelDTOT]": ...
@@ -70,7 +78,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
70
78
  parameters: "Optional[StatementParameterType]" = None,
71
79
  /,
72
80
  *,
73
- connection: "Optional[DuckDBPyConnection]" = None,
81
+ connection: "Optional[DuckDBConnection]" = None,
74
82
  schema_type: "Optional[type[ModelDTOT]]" = None,
75
83
  **kwargs: Any,
76
84
  ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
@@ -95,7 +103,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
95
103
  parameters: "Optional[StatementParameterType]" = None,
96
104
  /,
97
105
  *,
98
- connection: "Optional[DuckDBPyConnection]" = None,
106
+ connection: "Optional[DuckDBConnection]" = None,
99
107
  schema_type: None = None,
100
108
  **kwargs: Any,
101
109
  ) -> "dict[str, Any]": ...
@@ -106,7 +114,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
106
114
  parameters: "Optional[StatementParameterType]" = None,
107
115
  /,
108
116
  *,
109
- connection: "Optional[DuckDBPyConnection]" = None,
117
+ connection: "Optional[DuckDBConnection]" = None,
110
118
  schema_type: "type[ModelDTOT]",
111
119
  **kwargs: Any,
112
120
  ) -> "ModelDTOT": ...
@@ -116,7 +124,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
116
124
  parameters: Optional["StatementParameterType"] = None,
117
125
  /,
118
126
  *,
119
- connection: Optional["DuckDBPyConnection"] = None,
127
+ connection: Optional["DuckDBConnection"] = None,
120
128
  schema_type: "Optional[type[ModelDTOT]]" = None,
121
129
  **kwargs: Any,
122
130
  ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -140,7 +148,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
140
148
  parameters: "Optional[StatementParameterType]" = None,
141
149
  /,
142
150
  *,
143
- connection: "Optional[DuckDBPyConnection]" = None,
151
+ connection: "Optional[DuckDBConnection]" = None,
144
152
  schema_type: None = None,
145
153
  **kwargs: Any,
146
154
  ) -> "Optional[dict[str, Any]]": ...
@@ -151,7 +159,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
151
159
  parameters: "Optional[StatementParameterType]" = None,
152
160
  /,
153
161
  *,
154
- connection: "Optional[DuckDBPyConnection]" = None,
162
+ connection: "Optional[DuckDBConnection]" = None,
155
163
  schema_type: "type[ModelDTOT]",
156
164
  **kwargs: Any,
157
165
  ) -> "Optional[ModelDTOT]": ...
@@ -161,7 +169,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
161
169
  parameters: Optional["StatementParameterType"] = None,
162
170
  /,
163
171
  *,
164
- connection: Optional["DuckDBPyConnection"] = None,
172
+ connection: Optional["DuckDBConnection"] = None,
165
173
  schema_type: "Optional[type[ModelDTOT]]" = None,
166
174
  **kwargs: Any,
167
175
  ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -185,7 +193,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
185
193
  parameters: "Optional[StatementParameterType]" = None,
186
194
  /,
187
195
  *,
188
- connection: "Optional[DuckDBPyConnection]" = None,
196
+ connection: "Optional[DuckDBConnection]" = None,
189
197
  schema_type: None = None,
190
198
  **kwargs: Any,
191
199
  ) -> "Any": ...
@@ -196,7 +204,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
196
204
  parameters: "Optional[StatementParameterType]" = None,
197
205
  /,
198
206
  *,
199
- connection: "Optional[DuckDBPyConnection]" = None,
207
+ connection: "Optional[DuckDBConnection]" = None,
200
208
  schema_type: "type[T]",
201
209
  **kwargs: Any,
202
210
  ) -> "T": ...
@@ -206,7 +214,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
206
214
  parameters: "Optional[StatementParameterType]" = None,
207
215
  /,
208
216
  *,
209
- connection: "Optional[DuckDBPyConnection]" = None,
217
+ connection: "Optional[DuckDBConnection]" = None,
210
218
  schema_type: "Optional[type[T]]" = None,
211
219
  **kwargs: Any,
212
220
  ) -> "Union[T, Any]":
@@ -227,7 +235,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
227
235
  parameters: "Optional[StatementParameterType]" = None,
228
236
  /,
229
237
  *,
230
- connection: "Optional[DuckDBPyConnection]" = None,
238
+ connection: "Optional[DuckDBConnection]" = None,
231
239
  schema_type: None = None,
232
240
  **kwargs: Any,
233
241
  ) -> "Optional[Any]": ...
@@ -238,7 +246,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
238
246
  parameters: "Optional[StatementParameterType]" = None,
239
247
  /,
240
248
  *,
241
- connection: "Optional[DuckDBPyConnection]" = None,
249
+ connection: "Optional[DuckDBConnection]" = None,
242
250
  schema_type: "type[T]",
243
251
  **kwargs: Any,
244
252
  ) -> "Optional[T]": ...
@@ -248,7 +256,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
248
256
  parameters: "Optional[StatementParameterType]" = None,
249
257
  /,
250
258
  *,
251
- connection: "Optional[DuckDBPyConnection]" = None,
259
+ connection: "Optional[DuckDBConnection]" = None,
252
260
  schema_type: "Optional[type[T]]" = None,
253
261
  **kwargs: Any,
254
262
  ) -> "Optional[Union[T, Any]]":
@@ -269,7 +277,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
269
277
  parameters: Optional["StatementParameterType"] = None,
270
278
  /,
271
279
  *,
272
- connection: Optional["DuckDBPyConnection"] = None,
280
+ connection: Optional["DuckDBConnection"] = None,
273
281
  **kwargs: Any,
274
282
  ) -> int:
275
283
  connection = self._connection(connection)
@@ -285,7 +293,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
285
293
  parameters: "Optional[StatementParameterType]" = None,
286
294
  /,
287
295
  *,
288
- connection: "Optional[DuckDBPyConnection]" = None,
296
+ connection: "Optional[DuckDBConnection]" = None,
289
297
  schema_type: None = None,
290
298
  **kwargs: Any,
291
299
  ) -> "dict[str, Any]": ...
@@ -296,7 +304,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
296
304
  parameters: "Optional[StatementParameterType]" = None,
297
305
  /,
298
306
  *,
299
- connection: "Optional[DuckDBPyConnection]" = None,
307
+ connection: "Optional[DuckDBConnection]" = None,
300
308
  schema_type: "type[ModelDTOT]",
301
309
  **kwargs: Any,
302
310
  ) -> "ModelDTOT": ...
@@ -306,7 +314,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
306
314
  parameters: "Optional[StatementParameterType]" = None,
307
315
  /,
308
316
  *,
309
- connection: "Optional[DuckDBPyConnection]" = None,
317
+ connection: "Optional[DuckDBConnection]" = None,
310
318
  schema_type: "Optional[type[ModelDTOT]]" = None,
311
319
  **kwargs: Any,
312
320
  ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -328,7 +336,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
328
336
  parameters: Optional["StatementParameterType"] = None,
329
337
  /,
330
338
  *,
331
- connection: Optional["DuckDBPyConnection"] = None,
339
+ connection: Optional["DuckDBConnection"] = None,
332
340
  **kwargs: Any,
333
341
  ) -> str:
334
342
  connection = self._connection(connection)
@@ -345,7 +353,7 @@ class DuckDBDriver(SyncArrowBulkOperationsMixin["DuckDBPyConnection"], SyncDrive
345
353
  parameters: "Optional[StatementParameterType]" = None,
346
354
  /,
347
355
  *,
348
- connection: "Optional[DuckDBPyConnection]" = None,
356
+ connection: "Optional[DuckDBConnection]" = None,
349
357
  **kwargs: Any,
350
358
  ) -> "ArrowTable":
351
359
  connection = self._connection(connection)
@@ -4,13 +4,20 @@ from sqlspec.adapters.oracledb.config import (
4
4
  OracleSyncConfig,
5
5
  OracleSyncPoolConfig,
6
6
  )
7
- from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver
7
+ from sqlspec.adapters.oracledb.driver import (
8
+ OracleAsyncConnection,
9
+ OracleAsyncDriver,
10
+ OracleSyncConnection,
11
+ OracleSyncDriver,
12
+ )
8
13
 
9
14
  __all__ = (
10
15
  "OracleAsyncConfig",
16
+ "OracleAsyncConnection",
11
17
  "OracleAsyncDriver",
12
18
  "OracleAsyncPoolConfig",
13
19
  "OracleSyncConfig",
20
+ "OracleSyncConnection",
14
21
  "OracleSyncDriver",
15
22
  "OracleSyncPoolConfig",
16
23
  )
@@ -3,10 +3,9 @@ from dataclasses import dataclass, field
3
3
  from typing import TYPE_CHECKING, Any, Optional, cast
4
4
 
5
5
  from oracledb import create_pool_async as oracledb_create_pool # pyright: ignore[reportUnknownVariableType]
6
- from oracledb.connection import AsyncConnection
7
6
 
8
7
  from sqlspec.adapters.oracledb.config._common import OracleGenericPoolConfig
9
- from sqlspec.adapters.oracledb.driver import OracleAsyncDriver
8
+ from sqlspec.adapters.oracledb.driver import OracleAsyncConnection, OracleAsyncDriver
10
9
  from sqlspec.base import AsyncDatabaseConfig
11
10
  from sqlspec.exceptions import ImproperConfigurationError
12
11
  from sqlspec.typing import dataclass_to_dict
@@ -24,12 +23,12 @@ __all__ = (
24
23
 
25
24
 
26
25
  @dataclass
27
- class OracleAsyncPoolConfig(OracleGenericPoolConfig["AsyncConnection", "AsyncConnectionPool"]):
26
+ class OracleAsyncPoolConfig(OracleGenericPoolConfig["OracleAsyncConnection", "AsyncConnectionPool"]):
28
27
  """Async Oracle Pool Config"""
29
28
 
30
29
 
31
30
  @dataclass
32
- class OracleAsyncConfig(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionPool", "OracleAsyncDriver"]):
31
+ class OracleAsyncConfig(AsyncDatabaseConfig["OracleAsyncConnection", "AsyncConnectionPool", "OracleAsyncDriver"]):
33
32
  """Oracle Async database Configuration.
34
33
 
35
34
  This class provides the base configuration for Oracle database connections, extending
@@ -49,7 +48,7 @@ class OracleAsyncConfig(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionP
49
48
 
50
49
  If set, the plugin will use the provided pool rather than instantiate one.
51
50
  """
52
- connection_type: "type[AsyncConnection]" = field(init=False, default_factory=lambda: AsyncConnection)
51
+ connection_type: "type[OracleAsyncConnection]" = field(init=False, default_factory=lambda: OracleAsyncConnection)
53
52
  """Connection class to use.
54
53
 
55
54
  Defaults to :class:`AsyncConnection`.
@@ -111,7 +110,7 @@ class OracleAsyncConfig(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionP
111
110
  msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
112
111
  raise ImproperConfigurationError(msg)
113
112
 
114
- async def create_connection(self) -> "AsyncConnection":
113
+ async def create_connection(self) -> "OracleAsyncConnection":
115
114
  """Create and return a new oracledb async connection from the pool.
116
115
 
117
116
  Returns:
@@ -122,7 +121,7 @@ class OracleAsyncConfig(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionP
122
121
  """
123
122
  try:
124
123
  pool = await self.provide_pool()
125
- return cast("AsyncConnection", await pool.acquire()) # type: ignore[no-any-return,unused-ignore]
124
+ return cast("OracleAsyncConnection", await pool.acquire()) # type: ignore[no-any-return,unused-ignore]
126
125
  except Exception as e:
127
126
  msg = f"Could not configure the Oracle async connection. Error: {e!s}"
128
127
  raise ImproperConfigurationError(msg) from e
@@ -160,7 +159,7 @@ class OracleAsyncConfig(AsyncDatabaseConfig["AsyncConnection", "AsyncConnectionP
160
159
  return self.create_pool()
161
160
 
162
161
  @asynccontextmanager
163
- async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[AsyncConnection, None]":
162
+ async def provide_connection(self, *args: "Any", **kwargs: "Any") -> "AsyncGenerator[OracleAsyncConnection, None]":
164
163
  """Create a connection instance.
165
164
 
166
165
  Yields:
@@ -3,10 +3,9 @@ from dataclasses import dataclass, field
3
3
  from typing import TYPE_CHECKING, Any, Optional
4
4
 
5
5
  from oracledb import create_pool as oracledb_create_pool # pyright: ignore[reportUnknownVariableType]
6
- from oracledb.connection import Connection
7
6
 
8
7
  from sqlspec.adapters.oracledb.config._common import OracleGenericPoolConfig
9
- from sqlspec.adapters.oracledb.driver import OracleSyncDriver
8
+ from sqlspec.adapters.oracledb.driver import OracleSyncConnection, OracleSyncDriver
10
9
  from sqlspec.base import SyncDatabaseConfig
11
10
  from sqlspec.exceptions import ImproperConfigurationError
12
11
  from sqlspec.typing import dataclass_to_dict
@@ -24,12 +23,12 @@ __all__ = (
24
23
 
25
24
 
26
25
  @dataclass
27
- class OracleSyncPoolConfig(OracleGenericPoolConfig["Connection", "ConnectionPool"]):
26
+ class OracleSyncPoolConfig(OracleGenericPoolConfig["OracleSyncConnection", "ConnectionPool"]):
28
27
  """Sync Oracle Pool Config"""
29
28
 
30
29
 
31
30
  @dataclass
32
- class OracleSyncConfig(SyncDatabaseConfig["Connection", "ConnectionPool", "OracleSyncDriver"]):
31
+ class OracleSyncConfig(SyncDatabaseConfig["OracleSyncConnection", "ConnectionPool", "OracleSyncDriver"]):
33
32
  """Oracle Sync database Configuration.
34
33
 
35
34
  This class provides the base configuration for Oracle database connections, extending
@@ -49,7 +48,7 @@ class OracleSyncConfig(SyncDatabaseConfig["Connection", "ConnectionPool", "Oracl
49
48
 
50
49
  If set, the plugin will use the provided pool rather than instantiate one.
51
50
  """
52
- connection_type: "type[Connection]" = field(init=False, default_factory=lambda: Connection) # pyright: ignore
51
+ connection_type: "type[OracleSyncConnection]" = field(init=False, default_factory=lambda: OracleSyncConnection) # pyright: ignore
53
52
  """Connection class to use.
54
53
 
55
54
  Defaults to :class:`Connection`.
@@ -111,7 +110,7 @@ class OracleSyncConfig(SyncDatabaseConfig["Connection", "ConnectionPool", "Oracl
111
110
  msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
112
111
  raise ImproperConfigurationError(msg)
113
112
 
114
- def create_connection(self) -> "Connection":
113
+ def create_connection(self) -> "OracleSyncConnection":
115
114
  """Create and return a new oracledb connection from the pool.
116
115
 
117
116
  Returns:
@@ -160,7 +159,7 @@ class OracleSyncConfig(SyncDatabaseConfig["Connection", "ConnectionPool", "Oracl
160
159
  return self.create_pool()
161
160
 
162
161
  @contextmanager
163
- def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[Connection, None, None]":
162
+ def provide_connection(self, *args: "Any", **kwargs: "Any") -> "Generator[OracleSyncConnection, None, None]":
164
163
  """Create a connection instance.
165
164
 
166
165
  Yields: