sqlspec 0.3.0__py3-none-any.whl → 0.5.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/_serialization.py CHANGED
@@ -13,7 +13,7 @@ try:
13
13
 
14
14
  except ImportError:
15
15
  try:
16
- from orjson import dumps as _encode_json # pyright: ignore[reportMissingImports]
16
+ from orjson import dumps as _encode_json # pyright: ignore[reportMissingImports,reportUnknownVariableType]
17
17
  from orjson import loads as decode_json # type: ignore[no-redef]
18
18
 
19
19
  def encode_json(data: Any) -> str:
sqlspec/_typing.py CHANGED
@@ -5,20 +5,30 @@ This is used to ensure compatibility when one or more of the libraries are insta
5
5
 
6
6
  from __future__ import annotations
7
7
 
8
+ from enum import Enum
8
9
  from typing import (
9
10
  Any,
10
11
  ClassVar,
12
+ Final,
11
13
  Protocol,
14
+ Union,
12
15
  cast,
13
16
  runtime_checkable,
14
17
  )
15
18
 
16
- from typing_extensions import TypeVar, dataclass_transform
19
+ from typing_extensions import Literal, TypeVar, dataclass_transform
20
+
21
+
22
+ @runtime_checkable
23
+ class DataclassProtocol(Protocol):
24
+ """Protocol for instance checking dataclasses."""
25
+
26
+ __dataclass_fields__: ClassVar[dict[str, Any]]
27
+
17
28
 
18
29
  T = TypeVar("T")
19
30
  T_co = TypeVar("T_co", covariant=True)
20
31
 
21
-
22
32
  try:
23
33
  from pydantic import BaseModel, FailFast, TypeAdapter
24
34
 
@@ -99,11 +109,57 @@ except ImportError:
99
109
  UNSET = UnsetType.UNSET # pyright: ignore[reportConstantRedefinition]
100
110
  MSGSPEC_INSTALLED = False # pyright: ignore[reportConstantRedefinition]
101
111
 
112
+ try:
113
+ from litestar.dto.data_structures import DTOData # pyright: ignore[reportUnknownVariableType]
114
+
115
+ LITESTAR_INSTALLED = True
116
+ except ImportError:
117
+
118
+ @runtime_checkable
119
+ class DTOData(Protocol[T]): # type: ignore[no-redef]
120
+ """Placeholder implementation"""
121
+
122
+ __slots__ = ("_backend", "_data_as_builtins")
123
+
124
+ def __init__(self, backend: Any, data_as_builtins: Any) -> None:
125
+ """Placeholder init"""
126
+
127
+ def create_instance(self, **kwargs: Any) -> T:
128
+ """Placeholder implementation"""
129
+ return cast("T", kwargs)
130
+
131
+ def update_instance(self, instance: T, **kwargs: Any) -> T:
132
+ """Placeholder implementation"""
133
+ return cast("T", kwargs)
134
+
135
+ def as_builtins(self) -> Any:
136
+ """Placeholder implementation"""
137
+ return {}
138
+
139
+ LITESTAR_INSTALLED = False # pyright: ignore[reportConstantRedefinition]
140
+
141
+
142
+ class EmptyEnum(Enum):
143
+ """A sentinel enum used as placeholder."""
144
+
145
+ EMPTY = 0
146
+
147
+
148
+ EmptyType = Union[Literal[EmptyEnum.EMPTY], UnsetType]
149
+ Empty: Final = EmptyEnum.EMPTY
150
+
151
+
102
152
  __all__ = (
153
+ "LITESTAR_INSTALLED",
103
154
  "MSGSPEC_INSTALLED",
104
155
  "PYDANTIC_INSTALLED",
105
156
  "UNSET",
106
157
  "BaseModel",
158
+ "DTOData",
159
+ "DataclassProtocol",
160
+ "Empty",
161
+ "EmptyEnum",
162
+ "EmptyType",
107
163
  "FailFast",
108
164
  "Struct",
109
165
  "TypeAdapter",
@@ -2,33 +2,31 @@ from __future__ import annotations
2
2
 
3
3
  from contextlib import contextmanager
4
4
  from dataclasses import dataclass
5
- from typing import TYPE_CHECKING, TypeVar
5
+ from typing import TYPE_CHECKING
6
6
 
7
- from sqlspec.config import GenericDatabaseConfig
8
- from sqlspec.utils.empty import Empty
7
+ from sqlspec.base import GenericDatabaseConfig, NoPoolConfig
8
+ from sqlspec.typing import Empty, EmptyType
9
9
 
10
10
  if TYPE_CHECKING:
11
11
  from collections.abc import Generator
12
12
  from typing import Any
13
13
 
14
- from adbc_driver_manager.dbapi import Connection, Cursor
15
-
16
- from sqlspec.utils.empty import EmptyType
14
+ from adbc_driver_manager.dbapi import Connection
17
15
 
18
16
  __all__ = ("AdbcDatabaseConfig",)
19
17
 
20
- ConnectionT = TypeVar("ConnectionT", bound="Connection")
21
- CursorT = TypeVar("CursorT", bound="Cursor")
22
-
23
18
 
24
19
  @dataclass
25
- class AdbcDatabaseConfig(GenericDatabaseConfig):
20
+ class AdbcDatabaseConfig(NoPoolConfig["Connection"], GenericDatabaseConfig):
26
21
  """Configuration for ADBC connections.
27
22
 
28
23
  This class provides configuration options for ADBC database connections using the
29
24
  ADBC Driver Manager.([1](https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html))
30
25
  """
31
26
 
27
+ __supports_connection_pooling = False
28
+ __is_async = False
29
+
32
30
  uri: str | EmptyType = Empty
33
31
  """Database URI"""
34
32
  driver_name: str | EmptyType = Empty
@@ -4,10 +4,9 @@ from contextlib import asynccontextmanager
4
4
  from dataclasses import dataclass
5
5
  from typing import TYPE_CHECKING, Any
6
6
 
7
- from sqlspec.config import GenericDatabaseConfig
7
+ from sqlspec.base import GenericDatabaseConfig, NoPoolConfig
8
8
  from sqlspec.exceptions import ImproperConfigurationError
9
- from sqlspec.utils.dataclass import simple_asdict
10
- from sqlspec.utils.empty import Empty, EmptyType
9
+ from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
11
10
 
12
11
  if TYPE_CHECKING:
13
12
  from collections.abc import AsyncGenerator
@@ -20,7 +19,7 @@ __all__ = ("AiosqliteConfig",)
20
19
 
21
20
 
22
21
  @dataclass
23
- class AiosqliteConfig(GenericDatabaseConfig):
22
+ class AiosqliteConfig(NoPoolConfig["Connection"], GenericDatabaseConfig):
24
23
  """Configuration for Aiosqlite database connections.
25
24
 
26
25
  This class provides configuration options for Aiosqlite database connections, wrapping all parameters
@@ -60,7 +59,7 @@ class AiosqliteConfig(GenericDatabaseConfig):
60
59
  Returns:
61
60
  A string keyed dict of config kwargs for the aiosqlite.connect() function.
62
61
  """
63
- return simple_asdict(self, exclude_empty=True, convert_nested=False)
62
+ return dataclass_to_dict(self, exclude_empty=True, convert_nested=False)
64
63
 
65
64
  async def create_connection(self) -> Connection:
66
65
  """Create and return a new database connection.
@@ -79,22 +78,6 @@ class AiosqliteConfig(GenericDatabaseConfig):
79
78
  msg = f"Could not configure the Aiosqlite connection. Error: {e!s}"
80
79
  raise ImproperConfigurationError(msg) from e
81
80
 
82
- @asynccontextmanager
83
- async def lifespan(self, *args: Any, **kwargs: Any) -> AsyncGenerator[None, None]:
84
- """Manage the lifecycle of a database connection.
85
-
86
- Yields:
87
- None
88
-
89
- Raises:
90
- ImproperConfigurationError: If the connection could not be established.
91
- """
92
- connection = await self.create_connection()
93
- try:
94
- yield
95
- finally:
96
- await connection.close()
97
-
98
81
  @asynccontextmanager
99
82
  async def provide_connection(self, *args: Any, **kwargs: Any) -> AsyncGenerator[Connection, None]:
100
83
  """Create and provide a database connection.
@@ -4,17 +4,18 @@ from contextlib import asynccontextmanager
4
4
  from dataclasses import dataclass
5
5
  from typing import TYPE_CHECKING, TypeVar
6
6
 
7
+ from asyncmy.connection import Connection
8
+ from asyncmy.pool import Pool
9
+
10
+ from sqlspec.base import DatabaseConfigProtocol, GenericDatabaseConfig, GenericPoolConfig
7
11
  from sqlspec.exceptions import ImproperConfigurationError
8
- from sqlspec.utils.dataclass import simple_asdict
9
- from sqlspec.utils.empty import Empty, EmptyType
12
+ from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
10
13
 
11
14
  if TYPE_CHECKING:
12
15
  from collections.abc import AsyncGenerator
13
16
  from typing import Any
14
17
 
15
- from asyncmy.connection import Connection
16
18
  from asyncmy.cursors import Cursor, DictCursor
17
- from asyncmy.pool import Pool
18
19
 
19
20
  __all__ = (
20
21
  "AsyncMyConfig",
@@ -26,7 +27,7 @@ T = TypeVar("T")
26
27
 
27
28
 
28
29
  @dataclass
29
- class AsyncmyPoolConfig:
30
+ class AsyncmyPoolConfig(GenericPoolConfig):
30
31
  """Configuration for Asyncmy's connection pool.
31
32
 
32
33
  This class provides configuration options for Asyncmy database connection pools.
@@ -101,13 +102,16 @@ class AsyncmyPoolConfig:
101
102
  Returns:
102
103
  A string keyed dict of config kwargs for the Asyncmy create_pool function.
103
104
  """
104
- return simple_asdict(self, exclude_empty=True, convert_nested=False)
105
+ return dataclass_to_dict(self, exclude_empty=True, convert_nested=False)
105
106
 
106
107
 
107
108
  @dataclass
108
- class AsyncMyConfig:
109
+ class AsyncMyConfig(DatabaseConfigProtocol[Connection, Pool], GenericDatabaseConfig):
109
110
  """Asyncmy Configuration."""
110
111
 
112
+ __is_async__ = True
113
+ __supports_connection_pooling__ = True
114
+
111
115
  pool_config: AsyncmyPoolConfig | None = None
112
116
  """Asyncmy Pool configuration"""
113
117
 
@@ -125,7 +129,7 @@ class AsyncMyConfig:
125
129
  A string keyed dict of config kwargs for the Asyncmy create_pool function.
126
130
  """
127
131
  if self.pool_config:
128
- return simple_asdict(self.pool_config, exclude_empty=True, convert_nested=False)
132
+ return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
129
133
  msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
130
134
  raise ImproperConfigurationError(msg)
131
135
 
@@ -162,23 +166,6 @@ class AsyncMyConfig:
162
166
  """
163
167
  return await self.create_pool()
164
168
 
165
- @asynccontextmanager
166
- async def lifespan(self, *args: Any, **kwargs: Any) -> AsyncGenerator[None, None]:
167
- """Manage the lifecycle of a database connection pool.
168
-
169
- Yields:
170
- None
171
-
172
- Raises:
173
- ImproperConfigurationError: If the pool could not be established.
174
- """
175
- pool = await self.create_pool()
176
- try:
177
- yield
178
- finally:
179
- pool.close()
180
- await pool.wait_closed()
181
-
182
169
  @asynccontextmanager
183
170
  async def provide_connection(self, *args: Any, **kwargs: Any) -> AsyncGenerator[Connection, None]:
184
171
  """Create and provide a database connection.
@@ -2,24 +2,24 @@ from __future__ import annotations
2
2
 
3
3
  from contextlib import asynccontextmanager
4
4
  from dataclasses import dataclass
5
- from typing import TYPE_CHECKING, TypeVar
5
+ from typing import TYPE_CHECKING, TypeVar, Union
6
6
 
7
7
  from asyncpg import Record
8
8
  from asyncpg import create_pool as asyncpg_create_pool
9
+ from asyncpg.connection import Connection
10
+ from asyncpg.pool import Pool, PoolConnectionProxy
11
+ from typing_extensions import TypeAlias
9
12
 
10
13
  from sqlspec._serialization import decode_json, encode_json
11
- from sqlspec.config import GenericDatabaseConfig, GenericPoolConfig
14
+ from sqlspec.base import DatabaseConfigProtocol, GenericDatabaseConfig, GenericPoolConfig
12
15
  from sqlspec.exceptions import ImproperConfigurationError
13
- from sqlspec.utils.dataclass import simple_asdict
14
- from sqlspec.utils.empty import Empty, EmptyType
16
+ from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
15
17
 
16
18
  if TYPE_CHECKING:
17
19
  from asyncio import AbstractEventLoop
18
20
  from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
19
21
  from typing import Any
20
22
 
21
- from asyncpg.connection import Connection
22
- from asyncpg.pool import Pool, PoolConnectionProxy
23
23
 
24
24
  __all__ = (
25
25
  "AsyncPgConfig",
@@ -29,6 +29,8 @@ __all__ = (
29
29
 
30
30
  T = TypeVar("T")
31
31
 
32
+ PgConnection: TypeAlias = Union[Connection, PoolConnectionProxy]
33
+
32
34
 
33
35
  @dataclass
34
36
  class AsyncPgPoolConfig(GenericPoolConfig):
@@ -71,9 +73,12 @@ class AsyncPgPoolConfig(GenericPoolConfig):
71
73
 
72
74
 
73
75
  @dataclass
74
- class AsyncPgConfig(GenericDatabaseConfig):
76
+ class AsyncPgConfig(DatabaseConfigProtocol[PgConnection, Pool], GenericDatabaseConfig):
75
77
  """Asyncpg Configuration."""
76
78
 
79
+ __is_async__ = True
80
+ __supports_connection_pooling__ = True
81
+
77
82
  pool_config: AsyncPgPoolConfig | None = None
78
83
  """Asyncpg Pool configuration"""
79
84
  json_deserializer: Callable[[str], Any] = decode_json
@@ -98,7 +103,7 @@ class AsyncPgConfig(GenericDatabaseConfig):
98
103
  function.
99
104
  """
100
105
  if self.pool_config:
101
- return simple_asdict(self.pool_config, exclude_empty=True, convert_nested=False)
106
+ return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
102
107
  msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
103
108
  raise ImproperConfigurationError(msg)
104
109
 
@@ -124,16 +129,7 @@ class AsyncPgConfig(GenericDatabaseConfig):
124
129
  )
125
130
  return self.pool_instance
126
131
 
127
- @asynccontextmanager
128
- async def lifespan(self, *args: Any, **kwargs) -> AsyncGenerator[None, None]:
129
- db_pool = await self.create_pool()
130
- try:
131
- yield
132
- finally:
133
- db_pool.terminate()
134
- await db_pool.close()
135
-
136
- def provide_pool(self, *args: Any, **kwargs) -> Awaitable[Pool]:
132
+ def provide_pool(self, *args: Any, **kwargs: Any) -> Awaitable[Pool]:
137
133
  """Create a pool instance.
138
134
 
139
135
  Returns:
@@ -142,9 +138,7 @@ class AsyncPgConfig(GenericDatabaseConfig):
142
138
  return self.create_pool()
143
139
 
144
140
  @asynccontextmanager
145
- async def provide_connection(
146
- self, *args: Any, **kwargs: Any
147
- ) -> AsyncGenerator[Connection | PoolConnectionProxy, None]:
141
+ async def provide_connection(self, *args: Any, **kwargs: Any) -> AsyncGenerator[PoolConnectionProxy, None]:
148
142
  """Create a connection instance.
149
143
 
150
144
  Returns:
@@ -0,0 +1,3 @@
1
+ from sqlspec.adapters.duckdb.config import DuckDBConfig
2
+
3
+ __all__ = ("DuckDBConfig",)
@@ -2,23 +2,71 @@ from __future__ import annotations
2
2
 
3
3
  from contextlib import contextmanager
4
4
  from dataclasses import dataclass
5
- from typing import TYPE_CHECKING, Any
5
+ from typing import TYPE_CHECKING, Any, cast
6
6
 
7
- from sqlspec.config import GenericDatabaseConfig
7
+ from duckdb import DuckDBPyConnection
8
+
9
+ from sqlspec.base import GenericDatabaseConfig, NoPoolConfig
8
10
  from sqlspec.exceptions import ImproperConfigurationError
9
- from sqlspec.utils.dataclass import simple_asdict
10
- from sqlspec.utils.empty import Empty, EmptyType
11
+ from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
11
12
 
12
13
  if TYPE_CHECKING:
13
- from collections.abc import Generator
14
+ from collections.abc import Generator, Sequence
14
15
 
15
- from duckdb import DuckDBPyConnection
16
16
 
17
- __all__ = ("DuckDBConfig",)
17
+ __all__ = ("DuckDBConfig", "ExtensionConfig")
18
18
 
19
19
 
20
20
  @dataclass
21
- class DuckDBConfig(GenericDatabaseConfig):
21
+ class ExtensionConfig:
22
+ """Configuration for a DuckDB extension.
23
+
24
+ This class provides configuration options for DuckDB extensions, including installation
25
+ and post-install configuration settings.
26
+
27
+ For details see: https://duckdb.org/docs/extensions/overview
28
+ """
29
+
30
+ name: str
31
+ """The name of the extension to install"""
32
+ config: dict[str, Any] | None = None
33
+ """Optional configuration settings to apply after installation"""
34
+ force_install: bool = False
35
+ """Whether to force reinstall if already present"""
36
+ repository: str | None = None
37
+ """Optional repository name to install from"""
38
+ repository_url: str | None = None
39
+ """Optional repository URL to install from"""
40
+ version: str | None = None
41
+ """Optional version of the extension to install"""
42
+
43
+ @classmethod
44
+ def from_dict(cls, name: str, config: dict[str, Any] | bool | None = None) -> ExtensionConfig:
45
+ """Create an ExtensionConfig from a configuration dictionary.
46
+
47
+ Args:
48
+ name: The name of the extension
49
+ config: Configuration dictionary that may contain settings
50
+
51
+ Returns:
52
+ A new ExtensionConfig instance
53
+ """
54
+ if config is None:
55
+ return cls(name=name)
56
+
57
+ if not isinstance(config, dict):
58
+ config = {"force_install": bool(config)}
59
+
60
+ install_args = {
61
+ key: config.pop(key)
62
+ for key in ["force_install", "repository", "repository_url", "version", "config", "name"]
63
+ if key in config
64
+ }
65
+ return cls(name=name, **install_args)
66
+
67
+
68
+ @dataclass
69
+ class DuckDBConfig(NoPoolConfig[DuckDBPyConnection], GenericDatabaseConfig):
22
70
  """Configuration for DuckDB database connections.
23
71
 
24
72
  This class provides configuration options for DuckDB database connections, wrapping all parameters
@@ -39,6 +87,73 @@ class DuckDBConfig(GenericDatabaseConfig):
39
87
  For details see: https://duckdb.org/docs/api/python/overview#connection-options
40
88
  """
41
89
 
90
+ extensions: Sequence[ExtensionConfig] | EmptyType = Empty
91
+ """A sequence of extension configurations to install and configure upon connection creation."""
92
+
93
+ def __post_init__(self) -> None:
94
+ """Post-initialization validation and processing.
95
+
96
+ This method handles merging extension configurations from both the extensions field
97
+ and the config dictionary, if present. The config['extensions'] field can be either:
98
+ - A dictionary mapping extension names to their configurations
99
+ - A list of extension names (which will be installed with force_install=True)
100
+
101
+ Raises:
102
+ ImproperConfigurationError: If there are duplicate extension configurations.
103
+ """
104
+ if self.config is Empty:
105
+ self.config = {}
106
+
107
+ if self.extensions is Empty:
108
+ self.extensions = []
109
+ # this is purely for mypy
110
+ assert isinstance(self.config, dict) # noqa: S101
111
+ assert isinstance(self.extensions, list) # noqa: S101
112
+
113
+ _e = self.config.pop("extensions", {})
114
+ if not isinstance(_e, (dict, list, tuple)):
115
+ msg = "When configuring extensions in the 'config' dictionary, the value must be a dictionary or sequence of extension names"
116
+ raise ImproperConfigurationError(msg)
117
+ if not isinstance(_e, dict):
118
+ _e = {str(ext): {"force_install": False} for ext in _e} # pyright: ignore[reportUnknownVariableType,reportUnknownArgumentType]
119
+
120
+ if len(set(_e.keys()).intersection({ext.name for ext in self.extensions})) > 0: # pyright: ignore[ reportUnknownArgumentType]
121
+ msg = "Configuring the same extension in both 'extensions' and as a key in 'config['extensions']' is not allowed"
122
+ raise ImproperConfigurationError(msg)
123
+
124
+ self.extensions.extend([ExtensionConfig.from_dict(name, ext_config) for name, ext_config in _e.items()]) # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType]
125
+
126
+ def _configure_extensions(self, connection: DuckDBPyConnection) -> None:
127
+ """Configure extensions for the connection.
128
+
129
+ Args:
130
+ connection: The DuckDB connection to configure extensions for.
131
+
132
+ Raises:
133
+ ImproperConfigurationError: If extension installation or configuration fails.
134
+ """
135
+ if self.extensions is Empty:
136
+ return
137
+
138
+ for extension in cast("list[ExtensionConfig]", self.extensions):
139
+ try:
140
+ if extension.force_install:
141
+ connection.install_extension(
142
+ extension=extension.name,
143
+ force_install=extension.force_install,
144
+ repository=extension.repository,
145
+ repository_url=extension.repository_url,
146
+ version=extension.version,
147
+ )
148
+ connection.load_extension(extension.name)
149
+
150
+ if extension.config:
151
+ for key, value in extension.config.items():
152
+ connection.execute(f"SET {key}={value}")
153
+ except Exception as e:
154
+ msg = f"Failed to configure extension {extension.name}. Error: {e!s}"
155
+ raise ImproperConfigurationError(msg) from e
156
+
42
157
  @property
43
158
  def connection_config_dict(self) -> dict[str, Any]:
44
159
  """Return the connection configuration as a dict.
@@ -46,44 +161,30 @@ class DuckDBConfig(GenericDatabaseConfig):
46
161
  Returns:
47
162
  A string keyed dict of config kwargs for the duckdb.connect() function.
48
163
  """
49
- config = simple_asdict(self, exclude_empty=True, convert_nested=False)
164
+ config = dataclass_to_dict(self, exclude_empty=True, exclude={"extensions"}, convert_nested=False)
50
165
  if not config.get("database"):
51
166
  config["database"] = ":memory:"
52
167
  return config
53
168
 
54
169
  def create_connection(self) -> DuckDBPyConnection:
55
- """Create and return a new database connection.
170
+ """Create and return a new database connection with configured extensions.
56
171
 
57
172
  Returns:
58
- A new DuckDB connection instance.
173
+ A new DuckDB connection instance with extensions installed and configured.
59
174
 
60
175
  Raises:
61
- ImproperConfigurationError: If the connection could not be established.
176
+ ImproperConfigurationError: If the connection could not be established or extensions could not be configured.
62
177
  """
63
178
  import duckdb
64
179
 
65
180
  try:
66
- return duckdb.connect(**self.connection_config_dict)
181
+ connection = duckdb.connect(**self.connection_config_dict) # pyright: ignore[reportUnknownMemberType]
182
+ self._configure_extensions(connection)
183
+ return connection
67
184
  except Exception as e:
68
185
  msg = f"Could not configure the DuckDB connection. Error: {e!s}"
69
186
  raise ImproperConfigurationError(msg) from e
70
187
 
71
- @contextmanager
72
- def lifespan(self, *args: Any, **kwargs: Any) -> Generator[None, None, None]:
73
- """Manage the lifecycle of a database connection.
74
-
75
- Yields:
76
- None
77
-
78
- Raises:
79
- ImproperConfigurationError: If the connection could not be established.
80
- """
81
- connection = self.create_connection()
82
- try:
83
- yield
84
- finally:
85
- connection.close()
86
-
87
188
  @contextmanager
88
189
  def provide_connection(self, *args: Any, **kwargs: Any) -> Generator[DuckDBPyConnection, None, None]:
89
190
  """Create and provide a database connection.
@@ -1,4 +1,4 @@
1
- from .config import (
1
+ from sqlspec.adapters.oracledb.config import (
2
2
  OracleAsyncDatabaseConfig,
3
3
  OracleAsyncPoolConfig,
4
4
  OracleSyncDatabaseConfig,
@@ -13,7 +13,7 @@ from sqlspec.adapters.oracledb.config._common import (
13
13
  OracleGenericPoolConfig,
14
14
  )
15
15
  from sqlspec.exceptions import ImproperConfigurationError
16
- from sqlspec.utils.dataclass import simple_asdict
16
+ from sqlspec.typing import dataclass_to_dict
17
17
 
18
18
  if TYPE_CHECKING:
19
19
  from collections.abc import AsyncGenerator, Awaitable
@@ -26,16 +26,24 @@ __all__ = (
26
26
 
27
27
 
28
28
  @dataclass
29
- class OracleAsyncPoolConfig(OracleGenericPoolConfig[AsyncConnectionPool, AsyncConnection]):
29
+ class OracleAsyncPoolConfig(OracleGenericPoolConfig[AsyncConnection, AsyncConnectionPool]):
30
30
  """Async Oracle Pool Config"""
31
31
 
32
32
 
33
33
  @dataclass
34
- class OracleAsyncDatabaseConfig(OracleGenericDatabaseConfig[AsyncConnectionPool, AsyncConnection]):
34
+ class OracleAsyncDatabaseConfig(OracleGenericDatabaseConfig[AsyncConnection, AsyncConnectionPool]):
35
35
  """Async Oracle database Configuration."""
36
36
 
37
+ __is_async__ = True
38
+ __supports_connection_pooling__ = True
39
+
37
40
  pool_config: OracleAsyncPoolConfig | None = None
38
41
  """Oracle Pool configuration"""
42
+ pool_instance: AsyncConnectionPool | None = None
43
+ """Optional pool to use.
44
+
45
+ If set, the plugin will use the provided pool rather than instantiate one.
46
+ """
39
47
 
40
48
  @property
41
49
  def pool_config_dict(self) -> dict[str, Any]:
@@ -46,7 +54,7 @@ class OracleAsyncDatabaseConfig(OracleGenericDatabaseConfig[AsyncConnectionPool,
46
54
  function.
47
55
  """
48
56
  if self.pool_config is not None:
49
- return simple_asdict(self.pool_config, exclude_empty=True, convert_nested=False)
57
+ return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
50
58
  msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
51
59
  raise ImproperConfigurationError(msg)
52
60
 
@@ -65,20 +73,12 @@ class OracleAsyncDatabaseConfig(OracleGenericDatabaseConfig[AsyncConnectionPool,
65
73
 
66
74
  pool_config = self.pool_config_dict
67
75
  self.pool_instance = oracledb_create_pool(**pool_config)
68
- if self.pool_instance is None:
76
+ if self.pool_instance is None: # pyright: ignore[reportUnnecessaryComparison]
69
77
  msg = "Could not configure the 'pool_instance'. Please check your configuration."
70
78
  raise ImproperConfigurationError(msg)
71
79
  return self.pool_instance
72
80
 
73
- @asynccontextmanager
74
- async def lifespan(self, *args: Any, **kwargs) -> AsyncGenerator[None, None]:
75
- db_pool = await self.create_pool()
76
- try:
77
- yield
78
- finally:
79
- await db_pool.close(force=True)
80
-
81
- def provide_pool(self, *args: Any, **kwargs) -> Awaitable[AsyncConnectionPool]:
81
+ def provide_pool(self, *args: Any, **kwargs: Any) -> Awaitable[AsyncConnectionPool]:
82
82
  """Create a pool instance.
83
83
 
84
84
  Returns: