sqlspec 0.1.1__py3-none-any.whl → 0.4.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/__metadata__.py +1 -1
- sqlspec/_serialization.py +1 -1
- sqlspec/_typing.py +138 -0
- sqlspec/adapters/adbc/config.py +52 -0
- sqlspec/adapters/aiosqlite/__init__.py +3 -0
- sqlspec/adapters/aiosqlite/config.py +95 -0
- sqlspec/adapters/asyncmy/__init__.py +3 -0
- sqlspec/adapters/asyncmy/config.py +176 -0
- sqlspec/adapters/asyncpg/__init__.py +0 -0
- sqlspec/adapters/asyncpg/config.py +145 -0
- sqlspec/adapters/duckdb/__init__.py +0 -0
- sqlspec/adapters/duckdb/config.py +201 -0
- sqlspec/adapters/oracledb/__init__.py +13 -0
- sqlspec/adapters/oracledb/config/__init__.py +9 -0
- sqlspec/adapters/oracledb/config/_asyncio.py +95 -0
- sqlspec/adapters/oracledb/config/_common.py +151 -0
- sqlspec/adapters/oracledb/config/_sync.py +95 -0
- sqlspec/adapters/psycopg/__init__.py +0 -0
- sqlspec/adapters/psycopg/config/__init__.py +9 -0
- sqlspec/adapters/psycopg/config/_async.py +75 -0
- sqlspec/adapters/psycopg/config/_common.py +73 -0
- sqlspec/adapters/psycopg/config/_sync.py +75 -0
- sqlspec/adapters/sqlite/__init__.py +0 -0
- sqlspec/adapters/sqlite/config.py +92 -0
- sqlspec/config.py +16 -0
- sqlspec/exceptions.py +29 -0
- sqlspec/extensions/__init__.py +0 -0
- sqlspec/extensions/litestar/__init__.py +0 -0
- sqlspec/extensions/litestar/plugin.py +34 -0
- sqlspec/filters.py +35 -28
- sqlspec/typing.py +415 -0
- sqlspec-0.4.0.dist-info/METADATA +84 -0
- sqlspec-0.4.0.dist-info/RECORD +39 -0
- {sqlspec-0.1.1.dist-info → sqlspec-0.4.0.dist-info}/WHEEL +1 -1
- sqlspec-0.4.0.dist-info/licenses/NOTICE +29 -0
- sqlspec/types/empty.py +0 -18
- sqlspec/types/protocols.py +0 -117
- sqlspec/utils/dataclass.py +0 -130
- sqlspec-0.1.1.dist-info/METADATA +0 -25
- sqlspec-0.1.1.dist-info/RECORD +0 -14
- /sqlspec/{types → adapters}/__init__.py +0 -0
- /sqlspec/{utils → adapters/adbc}/__init__.py +0 -0
sqlspec/__metadata__.py
CHANGED
sqlspec/_serialization.py
CHANGED
|
@@ -14,7 +14,7 @@ try:
|
|
|
14
14
|
except ImportError:
|
|
15
15
|
try:
|
|
16
16
|
from orjson import dumps as _encode_json # pyright: ignore[reportMissingImports]
|
|
17
|
-
from orjson import loads as decode_json # type: ignore[no-redef
|
|
17
|
+
from orjson import loads as decode_json # type: ignore[no-redef]
|
|
18
18
|
|
|
19
19
|
def encode_json(data: Any) -> str:
|
|
20
20
|
return _encode_json(data).decode("utf-8") # type: ignore[no-any-return]
|
sqlspec/_typing.py
ADDED
|
@@ -0,0 +1,138 @@
|
|
|
1
|
+
"""This is a simple wrapper around a few important classes in each library.
|
|
2
|
+
|
|
3
|
+
This is used to ensure compatibility when one or more of the libraries are installed.
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
from enum import Enum
|
|
9
|
+
from typing import (
|
|
10
|
+
Any,
|
|
11
|
+
ClassVar,
|
|
12
|
+
Final,
|
|
13
|
+
Protocol,
|
|
14
|
+
Union,
|
|
15
|
+
cast,
|
|
16
|
+
runtime_checkable,
|
|
17
|
+
)
|
|
18
|
+
|
|
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
|
+
|
|
28
|
+
|
|
29
|
+
T = TypeVar("T")
|
|
30
|
+
T_co = TypeVar("T_co", covariant=True)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
try:
|
|
34
|
+
from pydantic import BaseModel, FailFast, TypeAdapter
|
|
35
|
+
|
|
36
|
+
PYDANTIC_INSTALLED = True
|
|
37
|
+
except ImportError:
|
|
38
|
+
|
|
39
|
+
@runtime_checkable
|
|
40
|
+
class BaseModel(Protocol): # type: ignore[no-redef]
|
|
41
|
+
"""Placeholder Implementation"""
|
|
42
|
+
|
|
43
|
+
model_fields: ClassVar[dict[str, Any]]
|
|
44
|
+
|
|
45
|
+
def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
|
|
46
|
+
"""Placeholder"""
|
|
47
|
+
return {}
|
|
48
|
+
|
|
49
|
+
@runtime_checkable
|
|
50
|
+
class TypeAdapter(Protocol[T_co]): # type: ignore[no-redef]
|
|
51
|
+
"""Placeholder Implementation"""
|
|
52
|
+
|
|
53
|
+
def __init__(
|
|
54
|
+
self,
|
|
55
|
+
type: Any, # noqa: A002
|
|
56
|
+
*,
|
|
57
|
+
config: Any | None = None,
|
|
58
|
+
_parent_depth: int = 2,
|
|
59
|
+
module: str | None = None,
|
|
60
|
+
) -> None:
|
|
61
|
+
"""Init"""
|
|
62
|
+
|
|
63
|
+
def validate_python(
|
|
64
|
+
self,
|
|
65
|
+
object: Any, # noqa: A002
|
|
66
|
+
/,
|
|
67
|
+
*,
|
|
68
|
+
strict: bool | None = None,
|
|
69
|
+
from_attributes: bool | None = None,
|
|
70
|
+
context: dict[str, Any] | None = None,
|
|
71
|
+
) -> T_co:
|
|
72
|
+
"""Stub"""
|
|
73
|
+
return cast("T_co", object)
|
|
74
|
+
|
|
75
|
+
@runtime_checkable
|
|
76
|
+
class FailFast(Protocol): # type: ignore[no-redef]
|
|
77
|
+
"""Placeholder Implementation for FailFast"""
|
|
78
|
+
|
|
79
|
+
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
80
|
+
"""Init"""
|
|
81
|
+
|
|
82
|
+
PYDANTIC_INSTALLED = False # pyright: ignore[reportConstantRedefinition]
|
|
83
|
+
|
|
84
|
+
try:
|
|
85
|
+
from msgspec import (
|
|
86
|
+
UNSET,
|
|
87
|
+
Struct,
|
|
88
|
+
UnsetType, # pyright: ignore[reportAssignmentType]
|
|
89
|
+
convert,
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
MSGSPEC_INSTALLED: bool = True
|
|
93
|
+
except ImportError:
|
|
94
|
+
import enum
|
|
95
|
+
|
|
96
|
+
@dataclass_transform()
|
|
97
|
+
@runtime_checkable
|
|
98
|
+
class Struct(Protocol): # type: ignore[no-redef]
|
|
99
|
+
"""Placeholder Implementation"""
|
|
100
|
+
|
|
101
|
+
__struct_fields__: ClassVar[tuple[str, ...]]
|
|
102
|
+
|
|
103
|
+
def convert(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-redef]
|
|
104
|
+
"""Placeholder implementation"""
|
|
105
|
+
return {}
|
|
106
|
+
|
|
107
|
+
class UnsetType(enum.Enum): # type: ignore[no-redef]
|
|
108
|
+
UNSET = "UNSET"
|
|
109
|
+
|
|
110
|
+
UNSET = UnsetType.UNSET # pyright: ignore[reportConstantRedefinition]
|
|
111
|
+
MSGSPEC_INSTALLED = False # pyright: ignore[reportConstantRedefinition]
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
class EmptyEnum(Enum):
|
|
115
|
+
"""A sentinel enum used as placeholder."""
|
|
116
|
+
|
|
117
|
+
EMPTY = 0
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
EmptyType = Union[Literal[EmptyEnum.EMPTY], UnsetType]
|
|
121
|
+
Empty: Final = EmptyEnum.EMPTY
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
__all__ = (
|
|
125
|
+
"MSGSPEC_INSTALLED",
|
|
126
|
+
"PYDANTIC_INSTALLED",
|
|
127
|
+
"UNSET",
|
|
128
|
+
"BaseModel",
|
|
129
|
+
"DataclassProtocol",
|
|
130
|
+
"Empty",
|
|
131
|
+
"EmptyEnum",
|
|
132
|
+
"EmptyType",
|
|
133
|
+
"FailFast",
|
|
134
|
+
"Struct",
|
|
135
|
+
"TypeAdapter",
|
|
136
|
+
"UnsetType",
|
|
137
|
+
"convert",
|
|
138
|
+
)
|
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from contextlib import contextmanager
|
|
4
|
+
from dataclasses import dataclass
|
|
5
|
+
from typing import TYPE_CHECKING, TypeVar
|
|
6
|
+
|
|
7
|
+
from sqlspec.config import GenericDatabaseConfig
|
|
8
|
+
from sqlspec.typing import Empty, EmptyType
|
|
9
|
+
|
|
10
|
+
if TYPE_CHECKING:
|
|
11
|
+
from collections.abc import Generator
|
|
12
|
+
from typing import Any
|
|
13
|
+
|
|
14
|
+
from adbc_driver_manager.dbapi import Connection, Cursor
|
|
15
|
+
|
|
16
|
+
__all__ = ("AdbcDatabaseConfig",)
|
|
17
|
+
|
|
18
|
+
ConnectionT = TypeVar("ConnectionT", bound="Connection")
|
|
19
|
+
CursorT = TypeVar("CursorT", bound="Cursor")
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@dataclass
|
|
23
|
+
class AdbcDatabaseConfig(GenericDatabaseConfig):
|
|
24
|
+
"""Configuration for ADBC connections.
|
|
25
|
+
|
|
26
|
+
This class provides configuration options for ADBC database connections using the
|
|
27
|
+
ADBC Driver Manager.([1](https://arrow.apache.org/adbc/current/python/api/adbc_driver_manager.html))
|
|
28
|
+
"""
|
|
29
|
+
|
|
30
|
+
uri: str | EmptyType = Empty
|
|
31
|
+
"""Database URI"""
|
|
32
|
+
driver_name: str | EmptyType = Empty
|
|
33
|
+
"""Name of the ADBC driver to use"""
|
|
34
|
+
db_kwargs: dict[str, Any] | None = None
|
|
35
|
+
"""Additional database-specific connection parameters"""
|
|
36
|
+
|
|
37
|
+
@property
|
|
38
|
+
def connection_params(self) -> dict[str, Any]:
|
|
39
|
+
"""Return the connection parameters as a dict."""
|
|
40
|
+
return {
|
|
41
|
+
k: v
|
|
42
|
+
for k, v in {"uri": self.uri, "driver": self.driver_name, **(self.db_kwargs or {})}.items()
|
|
43
|
+
if v is not Empty
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
@contextmanager
|
|
47
|
+
def provide_connection(self, *args: Any, **kwargs: Any) -> Generator[Connection, None, None]:
|
|
48
|
+
"""Create and provide a database connection."""
|
|
49
|
+
from adbc_driver_manager.dbapi import connect
|
|
50
|
+
|
|
51
|
+
with connect(**self.connection_params) as connection:
|
|
52
|
+
yield connection
|
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from contextlib import asynccontextmanager
|
|
4
|
+
from dataclasses import dataclass
|
|
5
|
+
from typing import TYPE_CHECKING, Any
|
|
6
|
+
|
|
7
|
+
from sqlspec.config import GenericDatabaseConfig
|
|
8
|
+
from sqlspec.exceptions import ImproperConfigurationError
|
|
9
|
+
from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
|
|
10
|
+
|
|
11
|
+
if TYPE_CHECKING:
|
|
12
|
+
from collections.abc import AsyncGenerator
|
|
13
|
+
from sqlite3 import Connection as SQLite3Connection
|
|
14
|
+
from typing import Literal
|
|
15
|
+
|
|
16
|
+
from aiosqlite import Connection
|
|
17
|
+
|
|
18
|
+
__all__ = ("AiosqliteConfig",)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@dataclass
|
|
22
|
+
class AiosqliteConfig(GenericDatabaseConfig):
|
|
23
|
+
"""Configuration for Aiosqlite database connections.
|
|
24
|
+
|
|
25
|
+
This class provides configuration options for Aiosqlite database connections, wrapping all parameters
|
|
26
|
+
available to aiosqlite.connect().
|
|
27
|
+
|
|
28
|
+
For details see: https://github.com/omnilib/aiosqlite/blob/main/aiosqlite/__init__.pyi
|
|
29
|
+
"""
|
|
30
|
+
|
|
31
|
+
database: str
|
|
32
|
+
"""The path to the database file to be opened. Pass ":memory:" to open a connection to a database that resides in RAM instead of on disk."""
|
|
33
|
+
|
|
34
|
+
timeout: float | EmptyType = Empty
|
|
35
|
+
"""How many seconds the connection should wait before raising an OperationalError when a table is locked. If another thread or process has acquired a shared lock, a wait for the specified timeout occurs."""
|
|
36
|
+
|
|
37
|
+
detect_types: int | EmptyType = Empty
|
|
38
|
+
"""Control whether and how data types are detected. It can be 0 (default) or a combination of PARSE_DECLTYPES and PARSE_COLNAMES."""
|
|
39
|
+
|
|
40
|
+
isolation_level: Literal["DEFERRED", "IMMEDIATE", "EXCLUSIVE"] | None | EmptyType = Empty
|
|
41
|
+
"""The isolation_level of the connection. This can be None for autocommit mode or one of "DEFERRED", "IMMEDIATE" or "EXCLUSIVE"."""
|
|
42
|
+
|
|
43
|
+
check_same_thread: bool | EmptyType = Empty
|
|
44
|
+
"""If True (default), ProgrammingError is raised if the database connection is used by a thread other than the one that created it. If False, the connection may be shared across multiple threads."""
|
|
45
|
+
|
|
46
|
+
factory: type[SQLite3Connection] | EmptyType = Empty
|
|
47
|
+
"""A custom Connection class factory. If given, must be a callable that returns a Connection instance."""
|
|
48
|
+
|
|
49
|
+
cached_statements: int | EmptyType = Empty
|
|
50
|
+
"""The number of statements that SQLite will cache for this connection. The default is 128."""
|
|
51
|
+
|
|
52
|
+
uri: bool | EmptyType = Empty
|
|
53
|
+
"""If set to True, database is interpreted as a URI with supported options."""
|
|
54
|
+
|
|
55
|
+
@property
|
|
56
|
+
def connection_config_dict(self) -> dict[str, Any]:
|
|
57
|
+
"""Return the connection configuration as a dict.
|
|
58
|
+
|
|
59
|
+
Returns:
|
|
60
|
+
A string keyed dict of config kwargs for the aiosqlite.connect() function.
|
|
61
|
+
"""
|
|
62
|
+
return dataclass_to_dict(self, exclude_empty=True, convert_nested=False)
|
|
63
|
+
|
|
64
|
+
async def create_connection(self) -> Connection:
|
|
65
|
+
"""Create and return a new database connection.
|
|
66
|
+
|
|
67
|
+
Returns:
|
|
68
|
+
A new Aiosqlite connection instance.
|
|
69
|
+
|
|
70
|
+
Raises:
|
|
71
|
+
ImproperConfigurationError: If the connection could not be established.
|
|
72
|
+
"""
|
|
73
|
+
import aiosqlite
|
|
74
|
+
|
|
75
|
+
try:
|
|
76
|
+
return await aiosqlite.connect(**self.connection_config_dict)
|
|
77
|
+
except Exception as e:
|
|
78
|
+
msg = f"Could not configure the Aiosqlite connection. Error: {e!s}"
|
|
79
|
+
raise ImproperConfigurationError(msg) from e
|
|
80
|
+
|
|
81
|
+
@asynccontextmanager
|
|
82
|
+
async def provide_connection(self, *args: Any, **kwargs: Any) -> AsyncGenerator[Connection, None]:
|
|
83
|
+
"""Create and provide a database connection.
|
|
84
|
+
|
|
85
|
+
Yields:
|
|
86
|
+
An Aiosqlite connection instance.
|
|
87
|
+
|
|
88
|
+
Raises:
|
|
89
|
+
ImproperConfigurationError: If the connection could not be established.
|
|
90
|
+
"""
|
|
91
|
+
connection = await self.create_connection()
|
|
92
|
+
try:
|
|
93
|
+
yield connection
|
|
94
|
+
finally:
|
|
95
|
+
await connection.close()
|
|
@@ -0,0 +1,176 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from contextlib import asynccontextmanager
|
|
4
|
+
from dataclasses import dataclass
|
|
5
|
+
from typing import TYPE_CHECKING, TypeVar
|
|
6
|
+
|
|
7
|
+
from sqlspec.exceptions import ImproperConfigurationError
|
|
8
|
+
from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
|
|
9
|
+
|
|
10
|
+
if TYPE_CHECKING:
|
|
11
|
+
from collections.abc import AsyncGenerator
|
|
12
|
+
from typing import Any
|
|
13
|
+
|
|
14
|
+
from asyncmy.connection import Connection
|
|
15
|
+
from asyncmy.cursors import Cursor, DictCursor
|
|
16
|
+
from asyncmy.pool import Pool
|
|
17
|
+
|
|
18
|
+
__all__ = (
|
|
19
|
+
"AsyncMyConfig",
|
|
20
|
+
"AsyncmyPoolConfig",
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
T = TypeVar("T")
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
@dataclass
|
|
28
|
+
class AsyncmyPoolConfig:
|
|
29
|
+
"""Configuration for Asyncmy's connection pool.
|
|
30
|
+
|
|
31
|
+
This class provides configuration options for Asyncmy database connection pools.
|
|
32
|
+
|
|
33
|
+
For details see: https://github.com/long2ice/asyncmy
|
|
34
|
+
"""
|
|
35
|
+
|
|
36
|
+
host: str | EmptyType = Empty
|
|
37
|
+
"""Host where the database server is located."""
|
|
38
|
+
|
|
39
|
+
user: str | EmptyType = Empty
|
|
40
|
+
"""The username used to authenticate with the database."""
|
|
41
|
+
|
|
42
|
+
password: str | EmptyType = Empty
|
|
43
|
+
"""The password used to authenticate with the database."""
|
|
44
|
+
|
|
45
|
+
database: str | EmptyType = Empty
|
|
46
|
+
"""The database name to use."""
|
|
47
|
+
|
|
48
|
+
port: int | EmptyType = Empty
|
|
49
|
+
"""The TCP/IP port of the MySQL server. Must be an integer."""
|
|
50
|
+
|
|
51
|
+
unix_socket: str | EmptyType = Empty
|
|
52
|
+
"""The location of the Unix socket file."""
|
|
53
|
+
|
|
54
|
+
charset: str | EmptyType = Empty
|
|
55
|
+
"""The character set to use for the connection."""
|
|
56
|
+
|
|
57
|
+
connect_timeout: float | EmptyType = Empty
|
|
58
|
+
"""Timeout before throwing an error when connecting."""
|
|
59
|
+
|
|
60
|
+
read_default_file: str | EmptyType = Empty
|
|
61
|
+
"""MySQL configuration file to read."""
|
|
62
|
+
|
|
63
|
+
read_default_group: str | EmptyType = Empty
|
|
64
|
+
"""Group to read from the configuration file."""
|
|
65
|
+
|
|
66
|
+
autocommit: bool | EmptyType = Empty
|
|
67
|
+
"""If True, autocommit mode will be enabled."""
|
|
68
|
+
|
|
69
|
+
local_infile: bool | EmptyType = Empty
|
|
70
|
+
"""If True, enables LOAD LOCAL INFILE."""
|
|
71
|
+
|
|
72
|
+
ssl: dict[str, Any] | bool | EmptyType = Empty
|
|
73
|
+
"""If present, a dictionary of SSL connection parameters, or just True."""
|
|
74
|
+
|
|
75
|
+
sql_mode: str | EmptyType = Empty
|
|
76
|
+
"""Default SQL_MODE to use."""
|
|
77
|
+
|
|
78
|
+
init_command: str | EmptyType = Empty
|
|
79
|
+
"""Initial SQL statement to execute once connected."""
|
|
80
|
+
|
|
81
|
+
cursor_class: type[Cursor] | type[DictCursor] | EmptyType = Empty
|
|
82
|
+
"""Custom cursor class to use."""
|
|
83
|
+
|
|
84
|
+
minsize: int | EmptyType = Empty
|
|
85
|
+
"""Minimum number of connections to keep in the pool."""
|
|
86
|
+
|
|
87
|
+
maxsize: int | EmptyType = Empty
|
|
88
|
+
"""Maximum number of connections allowed in the pool."""
|
|
89
|
+
|
|
90
|
+
echo: bool | EmptyType = Empty
|
|
91
|
+
"""If True, logging will be enabled for all SQL statements."""
|
|
92
|
+
|
|
93
|
+
pool_recycle: int | EmptyType = Empty
|
|
94
|
+
"""Number of seconds after which a connection is recycled."""
|
|
95
|
+
|
|
96
|
+
@property
|
|
97
|
+
def pool_config_dict(self) -> dict[str, Any]:
|
|
98
|
+
"""Return the pool configuration as a dict.
|
|
99
|
+
|
|
100
|
+
Returns:
|
|
101
|
+
A string keyed dict of config kwargs for the Asyncmy create_pool function.
|
|
102
|
+
"""
|
|
103
|
+
return dataclass_to_dict(self, exclude_empty=True, convert_nested=False)
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
@dataclass
|
|
107
|
+
class AsyncMyConfig:
|
|
108
|
+
"""Asyncmy Configuration."""
|
|
109
|
+
|
|
110
|
+
pool_config: AsyncmyPoolConfig | None = None
|
|
111
|
+
"""Asyncmy Pool configuration"""
|
|
112
|
+
|
|
113
|
+
pool_instance: Pool | None = None
|
|
114
|
+
"""Optional pool to use.
|
|
115
|
+
|
|
116
|
+
If set, the plugin will use the provided pool rather than instantiate one.
|
|
117
|
+
"""
|
|
118
|
+
|
|
119
|
+
@property
|
|
120
|
+
def pool_config_dict(self) -> dict[str, Any]:
|
|
121
|
+
"""Return the pool configuration as a dict.
|
|
122
|
+
|
|
123
|
+
Returns:
|
|
124
|
+
A string keyed dict of config kwargs for the Asyncmy create_pool function.
|
|
125
|
+
"""
|
|
126
|
+
if self.pool_config:
|
|
127
|
+
return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
|
|
128
|
+
msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
|
|
129
|
+
raise ImproperConfigurationError(msg)
|
|
130
|
+
|
|
131
|
+
async def create_pool(self) -> Pool:
|
|
132
|
+
"""Return a pool. If none exists yet, create one.
|
|
133
|
+
|
|
134
|
+
Returns:
|
|
135
|
+
Getter that returns the pool instance used by the plugin.
|
|
136
|
+
|
|
137
|
+
Raises:
|
|
138
|
+
ImproperConfigurationError: If the pool could not be created.
|
|
139
|
+
"""
|
|
140
|
+
if self.pool_instance is not None:
|
|
141
|
+
return self.pool_instance
|
|
142
|
+
|
|
143
|
+
if self.pool_config is None:
|
|
144
|
+
msg = "One of 'pool_config' or 'pool_instance' must be provided."
|
|
145
|
+
raise ImproperConfigurationError(msg)
|
|
146
|
+
|
|
147
|
+
try:
|
|
148
|
+
import asyncmy
|
|
149
|
+
|
|
150
|
+
self.pool_instance = await asyncmy.create_pool(**self.pool_config_dict)
|
|
151
|
+
return self.pool_instance
|
|
152
|
+
except Exception as e:
|
|
153
|
+
msg = f"Could not configure the Asyncmy pool. Error: {e!s}"
|
|
154
|
+
raise ImproperConfigurationError(msg) from e
|
|
155
|
+
|
|
156
|
+
async def provide_pool(self, *args: Any, **kwargs: Any) -> Pool:
|
|
157
|
+
"""Create a pool instance.
|
|
158
|
+
|
|
159
|
+
Returns:
|
|
160
|
+
A Pool instance.
|
|
161
|
+
"""
|
|
162
|
+
return await self.create_pool()
|
|
163
|
+
|
|
164
|
+
@asynccontextmanager
|
|
165
|
+
async def provide_connection(self, *args: Any, **kwargs: Any) -> AsyncGenerator[Connection, None]:
|
|
166
|
+
"""Create and provide a database connection.
|
|
167
|
+
|
|
168
|
+
Yields:
|
|
169
|
+
An Asyncmy connection instance.
|
|
170
|
+
|
|
171
|
+
Raises:
|
|
172
|
+
ImproperConfigurationError: If the connection could not be established.
|
|
173
|
+
"""
|
|
174
|
+
pool = await self.provide_pool(*args, **kwargs)
|
|
175
|
+
async with pool.acquire() as connection:
|
|
176
|
+
yield connection
|
|
File without changes
|
|
@@ -0,0 +1,145 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from contextlib import asynccontextmanager
|
|
4
|
+
from dataclasses import dataclass
|
|
5
|
+
from typing import TYPE_CHECKING, TypeVar
|
|
6
|
+
|
|
7
|
+
from asyncpg import Record
|
|
8
|
+
from asyncpg import create_pool as asyncpg_create_pool
|
|
9
|
+
|
|
10
|
+
from sqlspec._serialization import decode_json, encode_json
|
|
11
|
+
from sqlspec.config import GenericDatabaseConfig, GenericPoolConfig
|
|
12
|
+
from sqlspec.exceptions import ImproperConfigurationError
|
|
13
|
+
from sqlspec.typing import Empty, EmptyType, dataclass_to_dict
|
|
14
|
+
|
|
15
|
+
if TYPE_CHECKING:
|
|
16
|
+
from asyncio import AbstractEventLoop
|
|
17
|
+
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
|
|
18
|
+
from typing import Any
|
|
19
|
+
|
|
20
|
+
from asyncpg.connection import Connection
|
|
21
|
+
from asyncpg.pool import Pool, PoolConnectionProxy
|
|
22
|
+
|
|
23
|
+
__all__ = (
|
|
24
|
+
"AsyncPgConfig",
|
|
25
|
+
"AsyncPgPoolConfig",
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
T = TypeVar("T")
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@dataclass
|
|
33
|
+
class AsyncPgPoolConfig(GenericPoolConfig):
|
|
34
|
+
"""Configuration for Asyncpg's :class:`Pool <asyncpg.pool.Pool>`.
|
|
35
|
+
|
|
36
|
+
For details see: https://magicstack.github.io/asyncpg/current/api/index.html#connection-pools
|
|
37
|
+
"""
|
|
38
|
+
|
|
39
|
+
dsn: str
|
|
40
|
+
"""Connection arguments specified using as a single string in the following format: ``postgres://user:pass@host:port/database?option=value``
|
|
41
|
+
"""
|
|
42
|
+
connect_kwargs: dict[Any, Any] | None | EmptyType = Empty
|
|
43
|
+
"""A dictionary of arguments which will be passed directly to the ``connect()`` method as keyword arguments.
|
|
44
|
+
"""
|
|
45
|
+
connection_class: type[Connection] | None | EmptyType = Empty
|
|
46
|
+
"""The class to use for connections. Must be a subclass of Connection
|
|
47
|
+
"""
|
|
48
|
+
record_class: type[Record] | EmptyType = Empty
|
|
49
|
+
"""If specified, the class to use for records returned by queries on the connections in this pool. Must be a subclass of Record."""
|
|
50
|
+
|
|
51
|
+
min_size: int | EmptyType = Empty
|
|
52
|
+
"""The number of connections to keep open inside the connection pool."""
|
|
53
|
+
max_size: int | EmptyType = Empty
|
|
54
|
+
"""The number of connections to allow in connection pool “overflow”, that is connections that can be opened above
|
|
55
|
+
and beyond the pool_size setting, which defaults to 10."""
|
|
56
|
+
|
|
57
|
+
max_queries: int | EmptyType = Empty
|
|
58
|
+
"""Number of queries after a connection is closed and replaced with a new connection.
|
|
59
|
+
"""
|
|
60
|
+
max_inactive_connection_lifetime: float | EmptyType = Empty
|
|
61
|
+
"""Number of seconds after which inactive connections in the pool will be closed. Pass 0 to disable this mechanism."""
|
|
62
|
+
|
|
63
|
+
setup: Coroutine[None, type[Connection], Any] | EmptyType = Empty
|
|
64
|
+
"""A coroutine to prepare a connection right before it is returned from Pool.acquire(). An example use case would be to automatically set up notifications listeners for all connections of a pool."""
|
|
65
|
+
init: Coroutine[None, type[Connection], Any] | EmptyType = Empty
|
|
66
|
+
"""A coroutine to prepare a connection right before it is returned from Pool.acquire(). An example use case would be to automatically set up notifications listeners for all connections of a pool."""
|
|
67
|
+
|
|
68
|
+
loop: AbstractEventLoop | EmptyType = Empty
|
|
69
|
+
"""An asyncio event loop instance. If None, the default event loop will be used."""
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
@dataclass
|
|
73
|
+
class AsyncPgConfig(GenericDatabaseConfig):
|
|
74
|
+
"""Asyncpg Configuration."""
|
|
75
|
+
|
|
76
|
+
pool_config: AsyncPgPoolConfig | None = None
|
|
77
|
+
"""Asyncpg Pool configuration"""
|
|
78
|
+
json_deserializer: Callable[[str], Any] = decode_json
|
|
79
|
+
"""For dialects that support the :class:`JSON <sqlalchemy.types.JSON>` datatype, this is a Python callable that will
|
|
80
|
+
convert a JSON string to a Python object. By default, this is set to SQLSpec's
|
|
81
|
+
:attr:`decode_json() <sqlspec._serialization.decode_json>` function."""
|
|
82
|
+
json_serializer: Callable[[Any], str] = encode_json
|
|
83
|
+
"""For dialects that support the JSON datatype, this is a Python callable that will render a given object as JSON.
|
|
84
|
+
By default, SQLSpec's :attr:`encode_json() <sqlspec._serialization.encode_json>` is used."""
|
|
85
|
+
pool_instance: Pool | None = None
|
|
86
|
+
"""Optional pool to use.
|
|
87
|
+
|
|
88
|
+
If set, the plugin will use the provided pool rather than instantiate one.
|
|
89
|
+
"""
|
|
90
|
+
|
|
91
|
+
@property
|
|
92
|
+
def pool_config_dict(self) -> dict[str, Any]:
|
|
93
|
+
"""Return the pool configuration as a dict.
|
|
94
|
+
|
|
95
|
+
Returns:
|
|
96
|
+
A string keyed dict of config kwargs for the Asyncpg :func:`create_pool <asyncpg.pool.create_pool>`
|
|
97
|
+
function.
|
|
98
|
+
"""
|
|
99
|
+
if self.pool_config:
|
|
100
|
+
return dataclass_to_dict(self.pool_config, exclude_empty=True, convert_nested=False)
|
|
101
|
+
msg = "'pool_config' methods can not be used when a 'pool_instance' is provided."
|
|
102
|
+
raise ImproperConfigurationError(msg)
|
|
103
|
+
|
|
104
|
+
async def create_pool(self) -> Pool:
|
|
105
|
+
"""Return a pool. If none exists yet, create one.
|
|
106
|
+
|
|
107
|
+
Returns:
|
|
108
|
+
Getter that returns the pool instance used by the plugin.
|
|
109
|
+
"""
|
|
110
|
+
if self.pool_instance is not None:
|
|
111
|
+
return self.pool_instance
|
|
112
|
+
|
|
113
|
+
if self.pool_config is None:
|
|
114
|
+
msg = "One of 'pool_config' or 'pool_instance' must be provided."
|
|
115
|
+
raise ImproperConfigurationError(msg)
|
|
116
|
+
|
|
117
|
+
pool_config = self.pool_config_dict
|
|
118
|
+
self.pool_instance = await asyncpg_create_pool(**pool_config)
|
|
119
|
+
if self.pool_instance is None:
|
|
120
|
+
msg = "Could not configure the 'pool_instance'. Please check your configuration."
|
|
121
|
+
raise ImproperConfigurationError(
|
|
122
|
+
msg,
|
|
123
|
+
)
|
|
124
|
+
return self.pool_instance
|
|
125
|
+
|
|
126
|
+
def provide_pool(self, *args: Any, **kwargs: Any) -> Awaitable[Pool]:
|
|
127
|
+
"""Create a pool instance.
|
|
128
|
+
|
|
129
|
+
Returns:
|
|
130
|
+
A Pool instance.
|
|
131
|
+
"""
|
|
132
|
+
return self.create_pool()
|
|
133
|
+
|
|
134
|
+
@asynccontextmanager
|
|
135
|
+
async def provide_connection(
|
|
136
|
+
self, *args: Any, **kwargs: Any
|
|
137
|
+
) -> AsyncGenerator[Connection | PoolConnectionProxy, None]:
|
|
138
|
+
"""Create a connection instance.
|
|
139
|
+
|
|
140
|
+
Returns:
|
|
141
|
+
A connection instance.
|
|
142
|
+
"""
|
|
143
|
+
db_pool = await self.provide_pool(*args, **kwargs)
|
|
144
|
+
async with db_pool.acquire() as connection:
|
|
145
|
+
yield connection
|
|
File without changes
|