sqlspec 0.9.1__py3-none-any.whl → 0.10.1__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 (48) hide show
  1. sqlspec/__init__.py +2 -1
  2. sqlspec/_typing.py +24 -32
  3. sqlspec/adapters/adbc/__init__.py +2 -1
  4. sqlspec/adapters/adbc/config.py +7 -13
  5. sqlspec/adapters/adbc/driver.py +37 -30
  6. sqlspec/adapters/aiosqlite/__init__.py +2 -1
  7. sqlspec/adapters/aiosqlite/config.py +10 -12
  8. sqlspec/adapters/aiosqlite/driver.py +36 -31
  9. sqlspec/adapters/asyncmy/__init__.py +2 -1
  10. sqlspec/adapters/asyncmy/driver.py +34 -31
  11. sqlspec/adapters/asyncpg/config.py +1 -3
  12. sqlspec/adapters/asyncpg/driver.py +7 -3
  13. sqlspec/adapters/bigquery/__init__.py +4 -0
  14. sqlspec/adapters/bigquery/config/__init__.py +3 -0
  15. sqlspec/adapters/bigquery/config/_common.py +40 -0
  16. sqlspec/adapters/bigquery/config/_sync.py +87 -0
  17. sqlspec/adapters/bigquery/driver.py +701 -0
  18. sqlspec/adapters/duckdb/__init__.py +2 -1
  19. sqlspec/adapters/duckdb/config.py +17 -18
  20. sqlspec/adapters/duckdb/driver.py +38 -30
  21. sqlspec/adapters/oracledb/__init__.py +8 -1
  22. sqlspec/adapters/oracledb/config/_asyncio.py +7 -8
  23. sqlspec/adapters/oracledb/config/_sync.py +6 -7
  24. sqlspec/adapters/oracledb/driver.py +65 -62
  25. sqlspec/adapters/psqlpy/__init__.py +9 -0
  26. sqlspec/adapters/psqlpy/config.py +5 -5
  27. sqlspec/adapters/psqlpy/driver.py +34 -28
  28. sqlspec/adapters/psycopg/__init__.py +8 -1
  29. sqlspec/adapters/psycopg/config/__init__.py +10 -0
  30. sqlspec/adapters/psycopg/config/_async.py +6 -7
  31. sqlspec/adapters/psycopg/config/_sync.py +7 -8
  32. sqlspec/adapters/psycopg/driver.py +63 -53
  33. sqlspec/adapters/sqlite/__init__.py +2 -1
  34. sqlspec/adapters/sqlite/config.py +12 -11
  35. sqlspec/adapters/sqlite/driver.py +36 -29
  36. sqlspec/base.py +1 -66
  37. sqlspec/exceptions.py +9 -0
  38. sqlspec/extensions/litestar/config.py +3 -11
  39. sqlspec/extensions/litestar/handlers.py +2 -1
  40. sqlspec/extensions/litestar/plugin.py +4 -2
  41. sqlspec/mixins.py +156 -0
  42. sqlspec/typing.py +19 -1
  43. {sqlspec-0.9.1.dist-info → sqlspec-0.10.1.dist-info}/METADATA +8 -3
  44. sqlspec-0.10.1.dist-info/RECORD +67 -0
  45. sqlspec-0.9.1.dist-info/RECORD +0 -61
  46. {sqlspec-0.9.1.dist-info → sqlspec-0.10.1.dist-info}/WHEEL +0 -0
  47. {sqlspec-0.9.1.dist-info → sqlspec-0.10.1.dist-info}/licenses/LICENSE +0 -0
  48. {sqlspec-0.9.1.dist-info → sqlspec-0.10.1.dist-info}/licenses/NOTICE +0 -0
@@ -1,32 +1,39 @@
1
+ import sqlite3
1
2
  from contextlib import contextmanager
2
- from sqlite3 import Connection, Cursor
3
+ from sqlite3 import Cursor
3
4
  from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
4
5
 
5
- from sqlspec.base import SyncDriverAdapterProtocol, T
6
+ from sqlspec.base import SyncDriverAdapterProtocol
7
+ from sqlspec.mixins import SQLTranslatorMixin
6
8
 
7
9
  if TYPE_CHECKING:
8
10
  from collections.abc import Generator, Sequence
9
11
 
10
- from sqlspec.typing import ModelDTOT, StatementParameterType
12
+ from sqlspec.typing import ModelDTOT, StatementParameterType, T
11
13
 
12
- __all__ = ("SqliteDriver",)
14
+ __all__ = ("SqliteConnection", "SqliteDriver")
13
15
 
16
+ SqliteConnection = sqlite3.Connection
14
17
 
15
- class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
18
+
19
+ class SqliteDriver(
20
+ SQLTranslatorMixin["SqliteConnection"],
21
+ SyncDriverAdapterProtocol["SqliteConnection"],
22
+ ):
16
23
  """SQLite Sync Driver Adapter."""
17
24
 
18
- connection: "Connection"
25
+ connection: "SqliteConnection"
19
26
  dialect: str = "sqlite"
20
27
 
21
- def __init__(self, connection: "Connection") -> None:
28
+ def __init__(self, connection: "SqliteConnection") -> None:
22
29
  self.connection = connection
23
30
 
24
31
  @staticmethod
25
- def _cursor(connection: "Connection", *args: Any, **kwargs: Any) -> Cursor:
32
+ def _cursor(connection: "SqliteConnection", *args: Any, **kwargs: Any) -> Cursor:
26
33
  return connection.cursor(*args, **kwargs) # type: ignore[no-any-return]
27
34
 
28
35
  @contextmanager
29
- def _with_cursor(self, connection: "Connection") -> "Generator[Cursor, None, None]":
36
+ def _with_cursor(self, connection: "SqliteConnection") -> "Generator[Cursor, None, None]":
30
37
  cursor = self._cursor(connection)
31
38
  try:
32
39
  yield cursor
@@ -41,7 +48,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
41
48
  parameters: "Optional[StatementParameterType]" = None,
42
49
  /,
43
50
  *,
44
- connection: "Optional[Connection]" = None,
51
+ connection: "Optional[SqliteConnection]" = None,
45
52
  schema_type: None = None,
46
53
  **kwargs: Any,
47
54
  ) -> "Sequence[dict[str, Any]]": ...
@@ -52,7 +59,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
52
59
  parameters: "Optional[StatementParameterType]" = None,
53
60
  /,
54
61
  *,
55
- connection: "Optional[Connection]" = None,
62
+ connection: "Optional[SqliteConnection]" = None,
56
63
  schema_type: "type[ModelDTOT]",
57
64
  **kwargs: Any,
58
65
  ) -> "Sequence[ModelDTOT]": ...
@@ -62,7 +69,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
62
69
  parameters: Optional["StatementParameterType"] = None,
63
70
  /,
64
71
  *,
65
- connection: Optional["Connection"] = None,
72
+ connection: Optional["SqliteConnection"] = None,
66
73
  schema_type: "Optional[type[ModelDTOT]]" = None,
67
74
  **kwargs: Any,
68
75
  ) -> "Sequence[Union[ModelDTOT, dict[str, Any]]]":
@@ -93,7 +100,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
93
100
  parameters: "Optional[StatementParameterType]" = None,
94
101
  /,
95
102
  *,
96
- connection: "Optional[Connection]" = None,
103
+ connection: "Optional[SqliteConnection]" = None,
97
104
  schema_type: None = None,
98
105
  **kwargs: Any,
99
106
  ) -> "dict[str, Any]": ...
@@ -104,7 +111,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
104
111
  parameters: "Optional[StatementParameterType]" = None,
105
112
  /,
106
113
  *,
107
- connection: "Optional[Connection]" = None,
114
+ connection: "Optional[SqliteConnection]" = None,
108
115
  schema_type: "type[ModelDTOT]",
109
116
  **kwargs: Any,
110
117
  ) -> "ModelDTOT": ...
@@ -114,7 +121,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
114
121
  parameters: Optional["StatementParameterType"] = None,
115
122
  /,
116
123
  *,
117
- connection: Optional["Connection"] = None,
124
+ connection: Optional["SqliteConnection"] = None,
118
125
  schema_type: "Optional[type[ModelDTOT]]" = None,
119
126
  **kwargs: Any,
120
127
  ) -> "Union[ModelDTOT, dict[str, Any]]":
@@ -144,7 +151,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
144
151
  parameters: "Optional[StatementParameterType]" = None,
145
152
  /,
146
153
  *,
147
- connection: "Optional[Connection]" = None,
154
+ connection: "Optional[SqliteConnection]" = None,
148
155
  schema_type: None = None,
149
156
  **kwargs: Any,
150
157
  ) -> "Optional[dict[str, Any]]": ...
@@ -155,7 +162,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
155
162
  parameters: "Optional[StatementParameterType]" = None,
156
163
  /,
157
164
  *,
158
- connection: "Optional[Connection]" = None,
165
+ connection: "Optional[SqliteConnection]" = None,
159
166
  schema_type: "type[ModelDTOT]",
160
167
  **kwargs: Any,
161
168
  ) -> "Optional[ModelDTOT]": ...
@@ -165,7 +172,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
165
172
  parameters: Optional["StatementParameterType"] = None,
166
173
  /,
167
174
  *,
168
- connection: Optional["Connection"] = None,
175
+ connection: Optional["SqliteConnection"] = None,
169
176
  schema_type: "Optional[type[ModelDTOT]]" = None,
170
177
  **kwargs: Any,
171
178
  ) -> "Optional[Union[ModelDTOT, dict[str, Any]]]":
@@ -196,7 +203,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
196
203
  parameters: "Optional[StatementParameterType]" = None,
197
204
  /,
198
205
  *,
199
- connection: "Optional[Connection]" = None,
206
+ connection: "Optional[SqliteConnection]" = None,
200
207
  schema_type: None = None,
201
208
  **kwargs: Any,
202
209
  ) -> "Any": ...
@@ -207,7 +214,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
207
214
  parameters: "Optional[StatementParameterType]" = None,
208
215
  /,
209
216
  *,
210
- connection: "Optional[Connection]" = None,
217
+ connection: "Optional[SqliteConnection]" = None,
211
218
  schema_type: "type[T]",
212
219
  **kwargs: Any,
213
220
  ) -> "T": ...
@@ -217,7 +224,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
217
224
  parameters: "Optional[StatementParameterType]" = None,
218
225
  /,
219
226
  *,
220
- connection: "Optional[Connection]" = None,
227
+ connection: "Optional[SqliteConnection]" = None,
221
228
  schema_type: "Optional[type[T]]" = None,
222
229
  **kwargs: Any,
223
230
  ) -> "Union[T, Any]":
@@ -246,7 +253,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
246
253
  parameters: "Optional[StatementParameterType]" = None,
247
254
  /,
248
255
  *,
249
- connection: "Optional[Connection]" = None,
256
+ connection: "Optional[SqliteConnection]" = None,
250
257
  schema_type: None = None,
251
258
  **kwargs: Any,
252
259
  ) -> "Optional[Any]": ...
@@ -257,7 +264,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
257
264
  parameters: "Optional[StatementParameterType]" = None,
258
265
  /,
259
266
  *,
260
- connection: "Optional[Connection]" = None,
267
+ connection: "Optional[SqliteConnection]" = None,
261
268
  schema_type: "type[T]",
262
269
  **kwargs: Any,
263
270
  ) -> "Optional[T]": ...
@@ -267,7 +274,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
267
274
  parameters: "Optional[StatementParameterType]" = None,
268
275
  /,
269
276
  *,
270
- connection: "Optional[Connection]" = None,
277
+ connection: "Optional[SqliteConnection]" = None,
271
278
  schema_type: "Optional[type[T]]" = None,
272
279
  **kwargs: Any,
273
280
  ) -> "Optional[Union[T, Any]]":
@@ -296,7 +303,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
296
303
  parameters: Optional["StatementParameterType"] = None,
297
304
  /,
298
305
  *,
299
- connection: Optional["Connection"] = None,
306
+ connection: Optional["SqliteConnection"] = None,
300
307
  **kwargs: Any,
301
308
  ) -> int:
302
309
  """Insert, update, or delete data from the database.
@@ -321,7 +328,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
321
328
  parameters: "Optional[StatementParameterType]" = None,
322
329
  /,
323
330
  *,
324
- connection: "Optional[Connection]" = None,
331
+ connection: "Optional[SqliteConnection]" = None,
325
332
  schema_type: None = None,
326
333
  **kwargs: Any,
327
334
  ) -> "dict[str, Any]": ...
@@ -332,7 +339,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
332
339
  parameters: "Optional[StatementParameterType]" = None,
333
340
  /,
334
341
  *,
335
- connection: "Optional[Connection]" = None,
342
+ connection: "Optional[SqliteConnection]" = None,
336
343
  schema_type: "type[ModelDTOT]",
337
344
  **kwargs: Any,
338
345
  ) -> "ModelDTOT": ...
@@ -342,7 +349,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
342
349
  parameters: Optional["StatementParameterType"] = None,
343
350
  /,
344
351
  *,
345
- connection: Optional["Connection"] = None,
352
+ connection: Optional["SqliteConnection"] = None,
346
353
  schema_type: "Optional[type[ModelDTOT]]" = None,
347
354
  **kwargs: Any,
348
355
  ) -> "Optional[Union[dict[str, Any], ModelDTOT]]":
@@ -386,7 +393,7 @@ class SqliteDriver(SyncDriverAdapterProtocol["Connection"]):
386
393
  parameters: Optional["StatementParameterType"] = None,
387
394
  /,
388
395
  *,
389
- connection: Optional["Connection"] = None,
396
+ connection: Optional["SqliteConnection"] = None,
390
397
  **kwargs: Any,
391
398
  ) -> str:
392
399
  """Execute a script.
sqlspec/base.py CHANGED
@@ -20,16 +20,14 @@ from typing import (
20
20
 
21
21
  from sqlspec.exceptions import NotFoundError
22
22
  from sqlspec.statement import SQLStatement
23
- from sqlspec.typing import ModelDTOT, StatementParameterType
23
+ from sqlspec.typing import ConnectionT, ModelDTOT, PoolT, StatementParameterType, T
24
24
  from sqlspec.utils.sync_tools import maybe_async_
25
25
 
26
26
  if TYPE_CHECKING:
27
27
  from contextlib import AbstractAsyncContextManager, AbstractContextManager
28
28
 
29
- from pyarrow import Table as ArrowTable
30
29
 
31
30
  __all__ = (
32
- "AsyncArrowBulkOperationsMixin",
33
31
  "AsyncDatabaseConfig",
34
32
  "AsyncDriverAdapterProtocol",
35
33
  "CommonDriverAttributes",
@@ -39,15 +37,10 @@ __all__ = (
39
37
  "NoPoolSyncConfig",
40
38
  "SQLSpec",
41
39
  "SQLStatement",
42
- "SyncArrowBulkOperationsMixin",
43
40
  "SyncDatabaseConfig",
44
41
  "SyncDriverAdapterProtocol",
45
42
  )
46
43
 
47
- T = TypeVar("T")
48
- ConnectionT = TypeVar("ConnectionT")
49
- PoolT = TypeVar("PoolT")
50
- PoolT_co = TypeVar("PoolT_co", covariant=True)
51
44
  AsyncConfigT = TypeVar("AsyncConfigT", bound="Union[AsyncDatabaseConfig[Any, Any, Any], NoPoolAsyncConfig[Any, Any]]")
52
45
  SyncConfigT = TypeVar("SyncConfigT", bound="Union[SyncDatabaseConfig[Any, Any, Any], NoPoolSyncConfig[Any, Any]]")
53
46
  ConfigT = TypeVar(
@@ -558,35 +551,6 @@ class CommonDriverAttributes(Generic[ConnectionT]):
558
551
  return stmt.process()
559
552
 
560
553
 
561
- class SyncArrowBulkOperationsMixin(Generic[ConnectionT]):
562
- """Mixin for sync drivers supporting bulk Apache Arrow operations."""
563
-
564
- __supports_arrow__: "ClassVar[bool]" = True
565
-
566
- @abstractmethod
567
- def select_arrow( # pyright: ignore[reportUnknownParameterType]
568
- self,
569
- sql: str,
570
- parameters: "Optional[StatementParameterType]" = None,
571
- /,
572
- *,
573
- connection: "Optional[ConnectionT]" = None,
574
- **kwargs: Any,
575
- ) -> "ArrowTable": # pyright: ignore[reportUnknownReturnType]
576
- """Execute a SQL query and return results as an Apache Arrow Table.
577
-
578
- Args:
579
- sql: The SQL query string.
580
- parameters: Parameters for the query.
581
- connection: Optional connection override.
582
- **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
583
-
584
- Returns:
585
- An Apache Arrow Table containing the query results.
586
- """
587
- raise NotImplementedError
588
-
589
-
590
554
  class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generic[ConnectionT]):
591
555
  connection: "ConnectionT"
592
556
 
@@ -844,35 +808,6 @@ class SyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generi
844
808
  ) -> str: ...
845
809
 
846
810
 
847
- class AsyncArrowBulkOperationsMixin(Generic[ConnectionT]):
848
- """Mixin for async drivers supporting bulk Apache Arrow operations."""
849
-
850
- __supports_arrow__: "ClassVar[bool]" = True
851
-
852
- @abstractmethod
853
- async def select_arrow( # pyright: ignore[reportUnknownParameterType]
854
- self,
855
- sql: str,
856
- parameters: "Optional[StatementParameterType]" = None,
857
- /,
858
- *,
859
- connection: "Optional[ConnectionT]" = None,
860
- **kwargs: Any,
861
- ) -> "ArrowTable": # pyright: ignore[reportUnknownReturnType]
862
- """Execute a SQL query and return results as an Apache Arrow Table.
863
-
864
- Args:
865
- sql: The SQL query string.
866
- parameters: Parameters for the query.
867
- connection: Optional connection override.
868
- **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
869
-
870
- Returns:
871
- An Apache Arrow Table containing the query results.
872
- """
873
- raise NotImplementedError
874
-
875
-
876
811
  class AsyncDriverAdapterProtocol(CommonDriverAttributes[ConnectionT], ABC, Generic[ConnectionT]):
877
812
  connection: "ConnectionT"
878
813
 
sqlspec/exceptions.py CHANGED
@@ -78,6 +78,15 @@ class SQLParsingError(SQLSpecError):
78
78
  super().__init__(message)
79
79
 
80
80
 
81
+ class SQLConversionError(SQLSpecError):
82
+ """Issues converting SQL statements."""
83
+
84
+ def __init__(self, message: Optional[str] = None) -> None:
85
+ if message is None:
86
+ message = "Issues converting SQL statement."
87
+ super().__init__(message)
88
+
89
+
81
90
  class ParameterStyleMismatchError(SQLSpecError):
82
91
  """Error when parameter style doesn't match SQL placeholder style.
83
92
 
@@ -1,10 +1,6 @@
1
1
  from dataclasses import dataclass, field
2
2
  from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
3
3
 
4
- from sqlspec.base import (
5
- ConnectionT,
6
- PoolT,
7
- )
8
4
  from sqlspec.exceptions import ImproperConfigurationError
9
5
  from sqlspec.extensions.litestar.handlers import (
10
6
  autocommit_handler_maker,
@@ -23,13 +19,9 @@ if TYPE_CHECKING:
23
19
  from litestar.datastructures.state import State
24
20
  from litestar.types import BeforeMessageSendHookHandler, Scope
25
21
 
26
- from sqlspec.base import (
27
- AsyncConfigT,
28
- ConnectionT,
29
- DriverT,
30
- PoolT,
31
- SyncConfigT,
32
- )
22
+ from sqlspec.base import AsyncConfigT, DriverT, SyncConfigT
23
+ from sqlspec.typing import ConnectionT, PoolT
24
+
33
25
 
34
26
  CommitMode = Literal["manual", "autocommit", "autocommit_include_redirect"]
35
27
  DEFAULT_COMMIT_MODE: CommitMode = "manual"
@@ -19,7 +19,8 @@ if TYPE_CHECKING:
19
19
  from litestar.datastructures.state import State
20
20
  from litestar.types import Message, Scope
21
21
 
22
- from sqlspec.base import ConnectionT, DatabaseConfigProtocol, DriverT, PoolT
22
+ from sqlspec.base import DatabaseConfigProtocol, DriverT
23
+ from sqlspec.typing import ConnectionT, PoolT
23
24
 
24
25
 
25
26
  SESSION_TERMINUS_ASGI_EVENTS = {HTTP_RESPONSE_START, HTTP_DISCONNECT, WEBSOCKET_DISCONNECT, WEBSOCKET_CLOSE}
@@ -5,15 +5,14 @@ from litestar.plugins import InitPluginProtocol
5
5
 
6
6
  from sqlspec.base import (
7
7
  AsyncConfigT,
8
- ConnectionT,
9
8
  DatabaseConfigProtocol,
10
9
  DriverT,
11
- PoolT,
12
10
  SyncConfigT,
13
11
  )
14
12
  from sqlspec.base import SQLSpec as SQLSpecBase
15
13
  from sqlspec.exceptions import ImproperConfigurationError
16
14
  from sqlspec.extensions.litestar.config import DatabaseConfig
15
+ from sqlspec.typing import ConnectionT, PoolT
17
16
 
18
17
  if TYPE_CHECKING:
19
18
  from click import Group
@@ -85,6 +84,9 @@ class SQLSpec(InitPluginProtocol, SQLSpecBase):
85
84
  )
86
85
  for c in self._plugin_configs:
87
86
  c.annotation = self.add_config(c.config)
87
+ app_config.signature_types.append(c.annotation)
88
+ app_config.signature_types.append(c.config.connection_type) # type: ignore[union-attr]
89
+ app_config.signature_types.append(c.config.driver_type) # type: ignore[union-attr]
88
90
  app_config.before_send.append(c.before_send_handler)
89
91
  app_config.lifespan.append(c.lifespan_handler) # pyright: ignore[reportUnknownMemberType]
90
92
  app_config.dependencies.update(
sqlspec/mixins.py ADDED
@@ -0,0 +1,156 @@
1
+ from abc import abstractmethod
2
+ from typing import (
3
+ TYPE_CHECKING,
4
+ Any,
5
+ ClassVar,
6
+ Generic,
7
+ Optional,
8
+ )
9
+
10
+ from sqlglot import parse_one
11
+ from sqlglot.dialects.dialect import DialectType
12
+
13
+ from sqlspec.exceptions import SQLConversionError, SQLParsingError
14
+ from sqlspec.typing import ConnectionT, StatementParameterType
15
+
16
+ if TYPE_CHECKING:
17
+ from sqlspec.typing import ArrowTable
18
+
19
+ __all__ = (
20
+ "AsyncArrowBulkOperationsMixin",
21
+ "AsyncParquetExportMixin",
22
+ "SQLTranslatorMixin",
23
+ "SyncArrowBulkOperationsMixin",
24
+ "SyncParquetExportMixin",
25
+ )
26
+
27
+
28
+ class SyncArrowBulkOperationsMixin(Generic[ConnectionT]):
29
+ """Mixin for sync drivers supporting bulk Apache Arrow operations."""
30
+
31
+ __supports_arrow__: "ClassVar[bool]" = True
32
+
33
+ @abstractmethod
34
+ def select_arrow( # pyright: ignore[reportUnknownParameterType]
35
+ self,
36
+ sql: str,
37
+ parameters: "Optional[StatementParameterType]" = None,
38
+ /,
39
+ *,
40
+ connection: "Optional[ConnectionT]" = None,
41
+ **kwargs: Any,
42
+ ) -> "ArrowTable": # pyright: ignore[reportUnknownReturnType]
43
+ """Execute a SQL query and return results as an Apache Arrow Table.
44
+
45
+ Args:
46
+ sql: The SQL query string.
47
+ parameters: Parameters for the query.
48
+ connection: Optional connection override.
49
+ **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
50
+
51
+ Returns:
52
+ An Apache Arrow Table containing the query results.
53
+ """
54
+ raise NotImplementedError
55
+
56
+
57
+ class AsyncArrowBulkOperationsMixin(Generic[ConnectionT]):
58
+ """Mixin for async drivers supporting bulk Apache Arrow operations."""
59
+
60
+ __supports_arrow__: "ClassVar[bool]" = True
61
+
62
+ @abstractmethod
63
+ async def select_arrow( # pyright: ignore[reportUnknownParameterType]
64
+ self,
65
+ sql: str,
66
+ parameters: "Optional[StatementParameterType]" = None,
67
+ /,
68
+ *,
69
+ connection: "Optional[ConnectionT]" = None,
70
+ **kwargs: Any,
71
+ ) -> "ArrowTable": # pyright: ignore[reportUnknownReturnType]
72
+ """Execute a SQL query and return results as an Apache Arrow Table.
73
+
74
+ Args:
75
+ sql: The SQL query string.
76
+ parameters: Parameters for the query.
77
+ connection: Optional connection override.
78
+ **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
79
+
80
+ Returns:
81
+ An Apache Arrow Table containing the query results.
82
+ """
83
+ raise NotImplementedError
84
+
85
+
86
+ class SyncParquetExportMixin(Generic[ConnectionT]):
87
+ """Mixin for sync drivers supporting Parquet export."""
88
+
89
+ @abstractmethod
90
+ def select_to_parquet(
91
+ self,
92
+ sql: str,
93
+ parameters: "Optional[StatementParameterType]" = None,
94
+ /,
95
+ *,
96
+ connection: "Optional[ConnectionT]" = None,
97
+ **kwargs: Any,
98
+ ) -> None:
99
+ """Export a SQL query to a Parquet file."""
100
+ raise NotImplementedError
101
+
102
+
103
+ class AsyncParquetExportMixin(Generic[ConnectionT]):
104
+ """Mixin for async drivers supporting Parquet export."""
105
+
106
+ @abstractmethod
107
+ async def select_to_parquet(
108
+ self,
109
+ sql: str,
110
+ parameters: "Optional[StatementParameterType]" = None,
111
+ /,
112
+ *,
113
+ connection: "Optional[ConnectionT]" = None,
114
+ **kwargs: Any,
115
+ ) -> None:
116
+ """Export a SQL query to a Parquet file."""
117
+ raise NotImplementedError
118
+
119
+
120
+ class SQLTranslatorMixin(Generic[ConnectionT]):
121
+ """Mixin for drivers supporting SQL translation."""
122
+
123
+ dialect: str
124
+
125
+ def convert_to_dialect(
126
+ self,
127
+ sql: str,
128
+ to_dialect: DialectType = None,
129
+ pretty: bool = True,
130
+ ) -> str:
131
+ """Convert a SQL query to a different dialect.
132
+
133
+ Args:
134
+ sql: The SQL query string to convert.
135
+ to_dialect: The target dialect to convert to.
136
+ pretty: Whether to pretty-print the SQL query.
137
+
138
+ Returns:
139
+ The converted SQL query string.
140
+
141
+ Raises:
142
+ SQLParsingError: If the SQL query cannot be parsed.
143
+ SQLConversionError: If the SQL query cannot be converted to the target dialect.
144
+ """
145
+ try:
146
+ parsed = parse_one(sql, dialect=self.dialect)
147
+ except Exception as e:
148
+ error_msg = f"Failed to parse SQL: {e!s}"
149
+ raise SQLParsingError(error_msg) from e
150
+ if to_dialect is None:
151
+ to_dialect = self.dialect
152
+ try:
153
+ return parsed.sql(dialect=to_dialect, pretty=pretty)
154
+ except Exception as e:
155
+ error_msg = f"Failed to convert SQL to {to_dialect}: {e!s}"
156
+ raise SQLConversionError(error_msg) from e
sqlspec/typing.py CHANGED
@@ -33,8 +33,26 @@ PYDANTIC_USE_FAILFAST = False # leave permanently disabled for now
33
33
 
34
34
 
35
35
  T = TypeVar("T")
36
+ ConnectionT = TypeVar("ConnectionT")
37
+ """Type variable for connection types.
36
38
 
37
- ModelT = TypeVar("ModelT", bound="Union[Struct, BaseModel, DataclassProtocol]")
39
+ :class:`~sqlspec.typing.ConnectionT`
40
+ """
41
+ PoolT = TypeVar("PoolT")
42
+ """Type variable for pool types.
43
+
44
+ :class:`~sqlspec.typing.PoolT`
45
+ """
46
+ PoolT_co = TypeVar("PoolT_co", covariant=True)
47
+ """Type variable for covariant pool types.
48
+
49
+ :class:`~sqlspec.typing.PoolT_co`
50
+ """
51
+ ModelT = TypeVar("ModelT", bound="Union[dict[str, Any], Struct, BaseModel, DataclassProtocol]")
52
+ """Type variable for model types.
53
+
54
+ :class:`dict[str, Any]` | :class:`msgspec.Struct` | :class:`pydantic.BaseModel` | :class:`DataclassProtocol`
55
+ """
38
56
 
39
57
  FilterTypeT = TypeVar("FilterTypeT", bound="StatementFilter")
40
58
  """Type variable for filter types.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: sqlspec
3
- Version: 0.9.1
3
+ Version: 0.10.1
4
4
  Summary: SQL Experiments in Python
5
5
  Author-email: Cody Fincher <cody@litestar.dev>
6
6
  Maintainer-email: Litestar Developers <hello@litestar.dev>
@@ -41,6 +41,7 @@ Requires-Dist: oracledb; extra == 'oracledb'
41
41
  Provides-Extra: orjson
42
42
  Requires-Dist: orjson; extra == 'orjson'
43
43
  Provides-Extra: performance
44
+ Requires-Dist: msgspec; extra == 'performance'
44
45
  Requires-Dist: sqlglot[rs]; extra == 'performance'
45
46
  Provides-Extra: psqlpy
46
47
  Requires-Dist: psqlpy; extra == 'psqlpy'
@@ -131,7 +132,11 @@ etl_config = sql.add_config(
131
132
  )
132
133
  )
133
134
  with sql.provide_session(etl_config) as session:
134
- result = session.select_one("SELECT open_prompt(?)", data.message, schema_type=ChatMessage)
135
+ result = session.select_one(
136
+ "SELECT open_prompt(?)",
137
+ "Can you write a haiku about DuckDB?",
138
+ schema_type=ChatMessage
139
+ )
135
140
  print(result) # result is a ChatMessage pydantic model
136
141
  ```
137
142
 
@@ -249,7 +254,7 @@ This list is not final. If you have a driver you'd like to see added, please ope
249
254
  | [`oracledb`](https://oracle.github.io/python-oracledb/) | Oracle | Async | ✅ |
250
255
  | [`oracledb`](https://oracle.github.io/python-oracledb/) | Oracle | Sync | ✅ |
251
256
  | [`duckdb`](https://duckdb.org/) | DuckDB | Sync | ✅ |
252
- | [`bigquery`](https://googleapis.dev/python/bigquery/latest/index.html) | BigQuery | Sync | 🗓️ |
257
+ | [`bigquery`](https://googleapis.dev/python/bigquery/latest/index.html) | BigQuery | Sync | |
253
258
  | [`spanner`](https://googleapis.dev/python/spanner/latest/index.html) | Spanner | Sync | 🗓️ |
254
259
  | [`sqlserver`](https://docs.microsoft.com/en-us/sql/connect/python/pyodbc/python-sql-driver-for-pyodbc?view=sql-server-ver16) | SQL Server | Sync | 🗓️ |
255
260
  | [`mysql`](https://dev.mysql.com/doc/connector-python/en/connector-python-api-mysql-connector-python.html) | MySQL | Sync | 🗓️ |