sqlspec 0.14.1__py3-none-any.whl → 0.16.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/__init__.py +50 -25
- sqlspec/__main__.py +1 -1
- sqlspec/__metadata__.py +1 -3
- sqlspec/_serialization.py +1 -2
- sqlspec/_sql.py +480 -121
- sqlspec/_typing.py +278 -142
- sqlspec/adapters/adbc/__init__.py +4 -3
- sqlspec/adapters/adbc/_types.py +12 -0
- sqlspec/adapters/adbc/config.py +115 -260
- sqlspec/adapters/adbc/driver.py +462 -367
- sqlspec/adapters/aiosqlite/__init__.py +18 -3
- sqlspec/adapters/aiosqlite/_types.py +13 -0
- sqlspec/adapters/aiosqlite/config.py +199 -129
- sqlspec/adapters/aiosqlite/driver.py +230 -269
- sqlspec/adapters/asyncmy/__init__.py +18 -3
- sqlspec/adapters/asyncmy/_types.py +12 -0
- sqlspec/adapters/asyncmy/config.py +80 -168
- sqlspec/adapters/asyncmy/driver.py +260 -225
- sqlspec/adapters/asyncpg/__init__.py +19 -4
- sqlspec/adapters/asyncpg/_types.py +17 -0
- sqlspec/adapters/asyncpg/config.py +82 -181
- sqlspec/adapters/asyncpg/driver.py +285 -383
- sqlspec/adapters/bigquery/__init__.py +17 -3
- sqlspec/adapters/bigquery/_types.py +12 -0
- sqlspec/adapters/bigquery/config.py +191 -258
- sqlspec/adapters/bigquery/driver.py +474 -646
- sqlspec/adapters/duckdb/__init__.py +14 -3
- sqlspec/adapters/duckdb/_types.py +12 -0
- sqlspec/adapters/duckdb/config.py +415 -351
- sqlspec/adapters/duckdb/driver.py +343 -413
- sqlspec/adapters/oracledb/__init__.py +19 -5
- sqlspec/adapters/oracledb/_types.py +14 -0
- sqlspec/adapters/oracledb/config.py +123 -379
- sqlspec/adapters/oracledb/driver.py +507 -560
- sqlspec/adapters/psqlpy/__init__.py +13 -3
- sqlspec/adapters/psqlpy/_types.py +11 -0
- sqlspec/adapters/psqlpy/config.py +93 -254
- sqlspec/adapters/psqlpy/driver.py +505 -234
- sqlspec/adapters/psycopg/__init__.py +19 -5
- sqlspec/adapters/psycopg/_types.py +17 -0
- sqlspec/adapters/psycopg/config.py +143 -403
- sqlspec/adapters/psycopg/driver.py +706 -872
- sqlspec/adapters/sqlite/__init__.py +14 -3
- sqlspec/adapters/sqlite/_types.py +11 -0
- sqlspec/adapters/sqlite/config.py +202 -118
- sqlspec/adapters/sqlite/driver.py +264 -303
- sqlspec/base.py +105 -9
- sqlspec/{statement/builder → builder}/__init__.py +12 -14
- sqlspec/{statement/builder → builder}/_base.py +120 -55
- sqlspec/{statement/builder → builder}/_column.py +17 -6
- sqlspec/{statement/builder → builder}/_ddl.py +46 -79
- sqlspec/{statement/builder → builder}/_ddl_utils.py +5 -10
- sqlspec/{statement/builder → builder}/_delete.py +6 -25
- sqlspec/{statement/builder → builder}/_insert.py +18 -65
- sqlspec/builder/_merge.py +56 -0
- sqlspec/{statement/builder → builder}/_parsing_utils.py +8 -11
- sqlspec/{statement/builder → builder}/_select.py +11 -56
- sqlspec/{statement/builder → builder}/_update.py +12 -18
- sqlspec/{statement/builder → builder}/mixins/__init__.py +10 -14
- sqlspec/{statement/builder → builder}/mixins/_cte_and_set_ops.py +48 -59
- sqlspec/{statement/builder → builder}/mixins/_insert_operations.py +34 -18
- sqlspec/{statement/builder → builder}/mixins/_join_operations.py +1 -3
- sqlspec/{statement/builder → builder}/mixins/_merge_operations.py +19 -9
- sqlspec/{statement/builder → builder}/mixins/_order_limit_operations.py +3 -3
- sqlspec/{statement/builder → builder}/mixins/_pivot_operations.py +4 -8
- sqlspec/{statement/builder → builder}/mixins/_select_operations.py +25 -38
- sqlspec/{statement/builder → builder}/mixins/_update_operations.py +15 -16
- sqlspec/{statement/builder → builder}/mixins/_where_clause.py +210 -137
- sqlspec/cli.py +4 -5
- sqlspec/config.py +180 -133
- sqlspec/core/__init__.py +63 -0
- sqlspec/core/cache.py +873 -0
- sqlspec/core/compiler.py +396 -0
- sqlspec/core/filters.py +830 -0
- sqlspec/core/hashing.py +310 -0
- sqlspec/core/parameters.py +1209 -0
- sqlspec/core/result.py +664 -0
- sqlspec/{statement → core}/splitter.py +321 -191
- sqlspec/core/statement.py +666 -0
- sqlspec/driver/__init__.py +7 -10
- sqlspec/driver/_async.py +387 -176
- sqlspec/driver/_common.py +527 -289
- sqlspec/driver/_sync.py +390 -172
- sqlspec/driver/mixins/__init__.py +2 -19
- sqlspec/driver/mixins/_result_tools.py +164 -0
- sqlspec/driver/mixins/_sql_translator.py +6 -3
- sqlspec/exceptions.py +5 -252
- sqlspec/extensions/aiosql/adapter.py +93 -96
- sqlspec/extensions/litestar/cli.py +1 -1
- sqlspec/extensions/litestar/config.py +0 -1
- sqlspec/extensions/litestar/handlers.py +15 -26
- sqlspec/extensions/litestar/plugin.py +18 -16
- sqlspec/extensions/litestar/providers.py +17 -52
- sqlspec/loader.py +424 -105
- sqlspec/migrations/__init__.py +12 -0
- sqlspec/migrations/base.py +92 -68
- sqlspec/migrations/commands.py +24 -106
- sqlspec/migrations/loaders.py +402 -0
- sqlspec/migrations/runner.py +49 -51
- sqlspec/migrations/tracker.py +31 -44
- sqlspec/migrations/utils.py +64 -24
- sqlspec/protocols.py +7 -183
- sqlspec/storage/__init__.py +1 -1
- sqlspec/storage/backends/base.py +37 -40
- sqlspec/storage/backends/fsspec.py +136 -112
- sqlspec/storage/backends/obstore.py +138 -160
- sqlspec/storage/capabilities.py +5 -4
- sqlspec/storage/registry.py +57 -106
- sqlspec/typing.py +136 -115
- sqlspec/utils/__init__.py +2 -3
- sqlspec/utils/correlation.py +0 -3
- sqlspec/utils/deprecation.py +6 -6
- sqlspec/utils/fixtures.py +6 -6
- sqlspec/utils/logging.py +0 -2
- sqlspec/utils/module_loader.py +7 -12
- sqlspec/utils/singleton.py +0 -1
- sqlspec/utils/sync_tools.py +17 -38
- sqlspec/utils/text.py +12 -51
- sqlspec/utils/type_guards.py +443 -232
- {sqlspec-0.14.1.dist-info → sqlspec-0.16.0.dist-info}/METADATA +7 -2
- sqlspec-0.16.0.dist-info/RECORD +134 -0
- sqlspec/adapters/adbc/transformers.py +0 -108
- sqlspec/driver/connection.py +0 -207
- sqlspec/driver/mixins/_cache.py +0 -114
- sqlspec/driver/mixins/_csv_writer.py +0 -91
- sqlspec/driver/mixins/_pipeline.py +0 -508
- sqlspec/driver/mixins/_query_tools.py +0 -796
- sqlspec/driver/mixins/_result_utils.py +0 -138
- sqlspec/driver/mixins/_storage.py +0 -912
- sqlspec/driver/mixins/_type_coercion.py +0 -128
- sqlspec/driver/parameters.py +0 -138
- sqlspec/statement/__init__.py +0 -21
- sqlspec/statement/builder/_merge.py +0 -95
- sqlspec/statement/cache.py +0 -50
- sqlspec/statement/filters.py +0 -625
- sqlspec/statement/parameters.py +0 -956
- sqlspec/statement/pipelines/__init__.py +0 -210
- sqlspec/statement/pipelines/analyzers/__init__.py +0 -9
- sqlspec/statement/pipelines/analyzers/_analyzer.py +0 -646
- sqlspec/statement/pipelines/context.py +0 -109
- sqlspec/statement/pipelines/transformers/__init__.py +0 -7
- sqlspec/statement/pipelines/transformers/_expression_simplifier.py +0 -88
- sqlspec/statement/pipelines/transformers/_literal_parameterizer.py +0 -1247
- sqlspec/statement/pipelines/transformers/_remove_comments_and_hints.py +0 -76
- sqlspec/statement/pipelines/validators/__init__.py +0 -23
- sqlspec/statement/pipelines/validators/_dml_safety.py +0 -290
- sqlspec/statement/pipelines/validators/_parameter_style.py +0 -370
- sqlspec/statement/pipelines/validators/_performance.py +0 -714
- sqlspec/statement/pipelines/validators/_security.py +0 -967
- sqlspec/statement/result.py +0 -435
- sqlspec/statement/sql.py +0 -1774
- sqlspec/utils/cached_property.py +0 -25
- sqlspec/utils/statement_hashing.py +0 -203
- sqlspec-0.14.1.dist-info/RECORD +0 -145
- /sqlspec/{statement/builder → builder}/mixins/_delete_operations.py +0 -0
- {sqlspec-0.14.1.dist-info → sqlspec-0.16.0.dist-info}/WHEEL +0 -0
- {sqlspec-0.14.1.dist-info → sqlspec-0.16.0.dist-info}/entry_points.txt +0 -0
- {sqlspec-0.14.1.dist-info → sqlspec-0.16.0.dist-info}/licenses/LICENSE +0 -0
- {sqlspec-0.14.1.dist-info → sqlspec-0.16.0.dist-info}/licenses/NOTICE +0 -0
|
@@ -3,241 +3,120 @@
|
|
|
3
3
|
import contextlib
|
|
4
4
|
import logging
|
|
5
5
|
from contextlib import asynccontextmanager
|
|
6
|
-
from typing import TYPE_CHECKING, Any, ClassVar, Optional, cast
|
|
6
|
+
from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypedDict, Union, cast
|
|
7
7
|
|
|
8
8
|
from psycopg.rows import dict_row
|
|
9
9
|
from psycopg_pool import AsyncConnectionPool, ConnectionPool
|
|
10
|
+
from typing_extensions import NotRequired
|
|
10
11
|
|
|
12
|
+
from sqlspec.adapters.psycopg._types import PsycopgAsyncConnection, PsycopgSyncConnection
|
|
11
13
|
from sqlspec.adapters.psycopg.driver import (
|
|
12
|
-
|
|
14
|
+
PsycopgAsyncCursor,
|
|
13
15
|
PsycopgAsyncDriver,
|
|
14
|
-
|
|
16
|
+
PsycopgSyncCursor,
|
|
15
17
|
PsycopgSyncDriver,
|
|
18
|
+
psycopg_statement_config,
|
|
16
19
|
)
|
|
17
20
|
from sqlspec.config import AsyncDatabaseConfig, SyncDatabaseConfig
|
|
18
|
-
from sqlspec.statement.sql import SQLConfig
|
|
19
|
-
from sqlspec.typing import DictRow, Empty
|
|
20
21
|
|
|
21
22
|
if TYPE_CHECKING:
|
|
22
23
|
from collections.abc import AsyncGenerator, Callable, Generator
|
|
23
24
|
|
|
24
|
-
from
|
|
25
|
+
from sqlspec.core.statement import StatementConfig
|
|
26
|
+
|
|
25
27
|
|
|
26
28
|
logger = logging.getLogger("sqlspec.adapters.psycopg")
|
|
27
29
|
|
|
28
|
-
CONNECTION_FIELDS = frozenset(
|
|
29
|
-
{
|
|
30
|
-
"conninfo",
|
|
31
|
-
"host",
|
|
32
|
-
"port",
|
|
33
|
-
"user",
|
|
34
|
-
"password",
|
|
35
|
-
"dbname",
|
|
36
|
-
"connect_timeout",
|
|
37
|
-
"options",
|
|
38
|
-
"application_name",
|
|
39
|
-
"sslmode",
|
|
40
|
-
"sslcert",
|
|
41
|
-
"sslkey",
|
|
42
|
-
"sslrootcert",
|
|
43
|
-
"autocommit",
|
|
44
|
-
}
|
|
45
|
-
)
|
|
46
30
|
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
31
|
+
class PsycopgConnectionParams(TypedDict, total=False):
|
|
32
|
+
"""Psycopg connection parameters."""
|
|
33
|
+
|
|
34
|
+
conninfo: NotRequired[str]
|
|
35
|
+
host: NotRequired[str]
|
|
36
|
+
port: NotRequired[int]
|
|
37
|
+
user: NotRequired[str]
|
|
38
|
+
password: NotRequired[str]
|
|
39
|
+
dbname: NotRequired[str]
|
|
40
|
+
connect_timeout: NotRequired[int]
|
|
41
|
+
options: NotRequired[str]
|
|
42
|
+
application_name: NotRequired[str]
|
|
43
|
+
sslmode: NotRequired[str]
|
|
44
|
+
sslcert: NotRequired[str]
|
|
45
|
+
sslkey: NotRequired[str]
|
|
46
|
+
sslrootcert: NotRequired[str]
|
|
47
|
+
autocommit: NotRequired[bool]
|
|
48
|
+
extra: NotRequired[dict[str, Any]]
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
class PsycopgPoolParams(PsycopgConnectionParams, total=False):
|
|
52
|
+
"""Psycopg pool parameters."""
|
|
53
|
+
|
|
54
|
+
min_size: NotRequired[int]
|
|
55
|
+
max_size: NotRequired[int]
|
|
56
|
+
name: NotRequired[str]
|
|
57
|
+
timeout: NotRequired[float]
|
|
58
|
+
max_waiting: NotRequired[int]
|
|
59
|
+
max_lifetime: NotRequired[float]
|
|
60
|
+
max_idle: NotRequired[float]
|
|
61
|
+
reconnect_timeout: NotRequired[float]
|
|
62
|
+
num_workers: NotRequired[int]
|
|
63
|
+
configure: NotRequired["Callable[..., Any]"]
|
|
64
|
+
kwargs: NotRequired[dict[str, Any]]
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
__all__ = (
|
|
68
|
+
"PsycopgAsyncConfig",
|
|
69
|
+
"PsycopgAsyncCursor",
|
|
70
|
+
"PsycopgConnectionParams",
|
|
71
|
+
"PsycopgPoolParams",
|
|
72
|
+
"PsycopgSyncConfig",
|
|
73
|
+
"PsycopgSyncCursor",
|
|
61
74
|
)
|
|
62
75
|
|
|
63
|
-
__all__ = ("CONNECTION_FIELDS", "POOL_FIELDS", "PsycopgAsyncConfig", "PsycopgSyncConfig")
|
|
64
|
-
|
|
65
76
|
|
|
66
77
|
class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool, PsycopgSyncDriver]):
|
|
67
78
|
"""Configuration for Psycopg synchronous database connections with direct field-based configuration."""
|
|
68
79
|
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
# Driver class reference for dialect resolution
|
|
73
|
-
driver_type: type[PsycopgSyncDriver] = PsycopgSyncDriver
|
|
74
|
-
connection_type: type[PsycopgSyncConnection] = PsycopgSyncConnection
|
|
75
|
-
# Parameter style support information
|
|
76
|
-
supported_parameter_styles: ClassVar[tuple[str, ...]] = ("pyformat_positional", "pyformat_named")
|
|
77
|
-
"""Psycopg supports %s (positional) and %(name)s (named) parameter styles."""
|
|
78
|
-
|
|
79
|
-
default_parameter_style: ClassVar[str] = "pyformat_positional"
|
|
80
|
-
"""Psycopg's native parameter style is %s (pyformat positional)."""
|
|
80
|
+
driver_type: "ClassVar[type[PsycopgSyncDriver]]" = PsycopgSyncDriver
|
|
81
|
+
connection_type: "ClassVar[type[PsycopgSyncConnection]]" = PsycopgSyncConnection
|
|
81
82
|
|
|
82
83
|
def __init__(
|
|
83
84
|
self,
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
port: Optional[int] = None,
|
|
90
|
-
user: Optional[str] = None,
|
|
91
|
-
password: Optional[str] = None,
|
|
92
|
-
dbname: Optional[str] = None,
|
|
93
|
-
connect_timeout: Optional[float] = None,
|
|
94
|
-
options: Optional[str] = None,
|
|
95
|
-
application_name: Optional[str] = None,
|
|
96
|
-
sslmode: Optional[str] = None,
|
|
97
|
-
sslcert: Optional[str] = None,
|
|
98
|
-
sslkey: Optional[str] = None,
|
|
99
|
-
sslrootcert: Optional[str] = None,
|
|
100
|
-
autocommit: Optional[bool] = None,
|
|
101
|
-
# Pool parameters
|
|
102
|
-
min_size: Optional[int] = None,
|
|
103
|
-
max_size: Optional[int] = None,
|
|
104
|
-
name: Optional[str] = None,
|
|
105
|
-
timeout: Optional[float] = None,
|
|
106
|
-
max_waiting: Optional[int] = None,
|
|
107
|
-
max_lifetime: Optional[float] = None,
|
|
108
|
-
max_idle: Optional[float] = None,
|
|
109
|
-
reconnect_timeout: Optional[float] = None,
|
|
110
|
-
num_workers: Optional[int] = None,
|
|
111
|
-
configure: Optional["Callable[[Connection[Any]], None]"] = None,
|
|
112
|
-
kwargs: Optional[dict[str, Any]] = None,
|
|
113
|
-
# User-defined extras
|
|
114
|
-
extras: Optional[dict[str, Any]] = None,
|
|
115
|
-
**additional_kwargs: Any,
|
|
85
|
+
*,
|
|
86
|
+
pool_config: "Optional[Union[PsycopgPoolParams, dict[str, Any]]]" = None,
|
|
87
|
+
pool_instance: Optional["ConnectionPool"] = None,
|
|
88
|
+
statement_config: "Optional[StatementConfig]" = None,
|
|
89
|
+
migration_config: Optional[dict[str, Any]] = None,
|
|
116
90
|
) -> None:
|
|
117
91
|
"""Initialize Psycopg synchronous configuration.
|
|
118
92
|
|
|
119
93
|
Args:
|
|
94
|
+
pool_config: Pool configuration parameters (TypedDict or dict)
|
|
95
|
+
pool_instance: Existing pool instance to use
|
|
120
96
|
statement_config: Default SQL statement configuration
|
|
121
|
-
|
|
122
|
-
conninfo: Connection string in libpq format
|
|
123
|
-
host: Database server host
|
|
124
|
-
port: Database server port
|
|
125
|
-
user: Database user
|
|
126
|
-
password: Database password
|
|
127
|
-
dbname: Database name
|
|
128
|
-
connect_timeout: Connection timeout in seconds
|
|
129
|
-
options: Command-line options to send to the server
|
|
130
|
-
application_name: Application name for logging and statistics
|
|
131
|
-
sslmode: SSL mode (disable, prefer, require, etc.)
|
|
132
|
-
sslcert: SSL client certificate file
|
|
133
|
-
sslkey: SSL client private key file
|
|
134
|
-
sslrootcert: SSL root certificate file
|
|
135
|
-
autocommit: Enable autocommit mode
|
|
136
|
-
min_size: Minimum number of connections in the pool
|
|
137
|
-
max_size: Maximum number of connections in the pool
|
|
138
|
-
name: Name of the connection pool
|
|
139
|
-
timeout: Timeout for acquiring connections
|
|
140
|
-
max_waiting: Maximum number of waiting clients
|
|
141
|
-
max_lifetime: Maximum connection lifetime
|
|
142
|
-
max_idle: Maximum idle time for connections
|
|
143
|
-
reconnect_timeout: Time between reconnection attempts
|
|
144
|
-
num_workers: Number of background workers
|
|
145
|
-
configure: Callback to configure new connections
|
|
146
|
-
kwargs: Additional connection parameters
|
|
147
|
-
extras: Additional connection parameters not explicitly defined
|
|
148
|
-
**additional_kwargs: Additional parameters (stored in extras)
|
|
149
|
-
"""
|
|
150
|
-
# Store connection parameters as instance attributes
|
|
151
|
-
self.conninfo = conninfo
|
|
152
|
-
self.host = host
|
|
153
|
-
self.port = port
|
|
154
|
-
self.user = user
|
|
155
|
-
self.password = password
|
|
156
|
-
self.dbname = dbname
|
|
157
|
-
self.connect_timeout = connect_timeout
|
|
158
|
-
self.options = options
|
|
159
|
-
self.application_name = application_name
|
|
160
|
-
self.sslmode = sslmode
|
|
161
|
-
self.sslcert = sslcert
|
|
162
|
-
self.sslkey = sslkey
|
|
163
|
-
self.sslrootcert = sslrootcert
|
|
164
|
-
self.autocommit = autocommit
|
|
165
|
-
|
|
166
|
-
# Store pool parameters as instance attributes
|
|
167
|
-
self.min_size = min_size
|
|
168
|
-
self.max_size = max_size
|
|
169
|
-
self.name = name
|
|
170
|
-
self.timeout = timeout
|
|
171
|
-
self.max_waiting = max_waiting
|
|
172
|
-
self.max_lifetime = max_lifetime
|
|
173
|
-
self.max_idle = max_idle
|
|
174
|
-
self.reconnect_timeout = reconnect_timeout
|
|
175
|
-
self.num_workers = num_workers
|
|
176
|
-
self.configure = configure
|
|
177
|
-
self.kwargs = kwargs or {}
|
|
178
|
-
|
|
179
|
-
self.extras = extras or {}
|
|
180
|
-
self.extras.update(additional_kwargs)
|
|
181
|
-
|
|
182
|
-
# Store other config
|
|
183
|
-
self.statement_config = statement_config or SQLConfig()
|
|
184
|
-
self.default_row_type = default_row_type
|
|
185
|
-
|
|
186
|
-
super().__init__()
|
|
187
|
-
|
|
188
|
-
@property
|
|
189
|
-
def connection_config_dict(self) -> dict[str, Any]:
|
|
190
|
-
"""Return the connection configuration as a dict for psycopg operations.
|
|
191
|
-
|
|
192
|
-
Returns only connection-specific parameters.
|
|
193
|
-
"""
|
|
194
|
-
# Gather non-None parameters from connection fields only
|
|
195
|
-
config = {
|
|
196
|
-
field: getattr(self, field)
|
|
197
|
-
for field in CONNECTION_FIELDS
|
|
198
|
-
if getattr(self, field, None) is not None and getattr(self, field) is not Empty
|
|
199
|
-
}
|
|
200
|
-
|
|
201
|
-
# Merge extras and kwargs
|
|
202
|
-
config.update(self.extras)
|
|
203
|
-
if self.kwargs:
|
|
204
|
-
config.update(self.kwargs)
|
|
97
|
+
migration_config: Migration configuration
|
|
205
98
|
|
|
206
|
-
config["row_factory"] = dict_row
|
|
207
|
-
|
|
208
|
-
return config
|
|
209
|
-
|
|
210
|
-
@property
|
|
211
|
-
def pool_config_dict(self) -> dict[str, Any]:
|
|
212
|
-
"""Return the pool configuration as a dict for psycopg pool operations.
|
|
213
|
-
|
|
214
|
-
Returns all configuration parameters including connection and pool-specific parameters.
|
|
215
99
|
"""
|
|
216
|
-
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
|
|
221
|
-
|
|
222
|
-
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
config["row_factory"] = dict_row
|
|
229
|
-
|
|
230
|
-
return config
|
|
100
|
+
processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
|
|
101
|
+
if "extra" in processed_pool_config:
|
|
102
|
+
extras = processed_pool_config.pop("extra")
|
|
103
|
+
processed_pool_config.update(extras)
|
|
104
|
+
|
|
105
|
+
super().__init__(
|
|
106
|
+
pool_config=processed_pool_config,
|
|
107
|
+
pool_instance=pool_instance,
|
|
108
|
+
migration_config=migration_config,
|
|
109
|
+
statement_config=statement_config or psycopg_statement_config,
|
|
110
|
+
)
|
|
231
111
|
|
|
232
112
|
def _create_pool(self) -> "ConnectionPool":
|
|
233
113
|
"""Create the actual connection pool."""
|
|
234
114
|
logger.info("Creating Psycopg connection pool", extra={"adapter": "psycopg"})
|
|
235
115
|
|
|
236
116
|
try:
|
|
237
|
-
all_config = self.
|
|
117
|
+
all_config = dict(self.pool_config)
|
|
238
118
|
|
|
239
|
-
|
|
240
|
-
pool_params = {
|
|
119
|
+
pool_parameters = {
|
|
241
120
|
"min_size": all_config.pop("min_size", 4),
|
|
242
121
|
"max_size": all_config.pop("max_size", None),
|
|
243
122
|
"name": all_config.pop("name", None),
|
|
@@ -249,29 +128,34 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
|
|
|
249
128
|
"num_workers": all_config.pop("num_workers", 3),
|
|
250
129
|
}
|
|
251
130
|
|
|
252
|
-
# Capture autocommit setting before configuring the pool
|
|
253
131
|
autocommit_setting = all_config.get("autocommit")
|
|
254
132
|
|
|
255
133
|
def configure_connection(conn: "PsycopgSyncConnection") -> None:
|
|
256
134
|
conn.row_factory = dict_row
|
|
257
|
-
# Apply autocommit setting if specified
|
|
258
135
|
if autocommit_setting is not None:
|
|
259
136
|
conn.autocommit = autocommit_setting
|
|
260
137
|
|
|
261
|
-
|
|
138
|
+
try:
|
|
139
|
+
import pgvector.psycopg
|
|
140
|
+
|
|
141
|
+
pgvector.psycopg.register_vector(conn)
|
|
142
|
+
logger.debug("pgvector registered successfully for psycopg sync connection")
|
|
143
|
+
except ImportError:
|
|
144
|
+
pass
|
|
145
|
+
except Exception as e:
|
|
146
|
+
logger.debug("Failed to register pgvector for psycopg sync: %s", e)
|
|
147
|
+
|
|
148
|
+
pool_parameters["configure"] = all_config.pop("configure", configure_connection)
|
|
262
149
|
|
|
263
|
-
|
|
150
|
+
pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None}
|
|
264
151
|
|
|
265
152
|
conninfo = all_config.pop("conninfo", None)
|
|
266
153
|
if conninfo:
|
|
267
|
-
|
|
268
|
-
# Don't pass kwargs when using conninfo string
|
|
269
|
-
pool = ConnectionPool(conninfo, open=True, **pool_params)
|
|
154
|
+
pool = ConnectionPool(conninfo, open=True, **pool_parameters)
|
|
270
155
|
else:
|
|
271
|
-
|
|
272
|
-
all_config.
|
|
273
|
-
|
|
274
|
-
pool = ConnectionPool("", kwargs=all_config, open=True, **pool_params)
|
|
156
|
+
kwargs = all_config.pop("kwargs", {})
|
|
157
|
+
all_config.update(kwargs)
|
|
158
|
+
pool = ConnectionPool("", kwargs=all_config, open=True, **pool_parameters)
|
|
275
159
|
|
|
276
160
|
logger.info("Psycopg connection pool created successfully", extra={"adapter": "psycopg"})
|
|
277
161
|
except Exception as e:
|
|
@@ -287,7 +171,6 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
|
|
|
287
171
|
logger.info("Closing Psycopg connection pool", extra={"adapter": "psycopg"})
|
|
288
172
|
|
|
289
173
|
try:
|
|
290
|
-
# This avoids the "cannot join current thread" error during garbage collection
|
|
291
174
|
if hasattr(self.pool_instance, "_closed"):
|
|
292
175
|
self.pool_instance._closed = True
|
|
293
176
|
|
|
@@ -331,29 +214,22 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
|
|
|
331
214
|
conn.close()
|
|
332
215
|
|
|
333
216
|
@contextlib.contextmanager
|
|
334
|
-
def provide_session(
|
|
217
|
+
def provide_session(
|
|
218
|
+
self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
|
|
219
|
+
) -> "Generator[PsycopgSyncDriver, None, None]":
|
|
335
220
|
"""Provide a driver session context manager.
|
|
336
221
|
|
|
337
222
|
Args:
|
|
338
223
|
*args: Additional arguments.
|
|
224
|
+
statement_config: Optional statement configuration override.
|
|
339
225
|
**kwargs: Additional keyword arguments.
|
|
340
226
|
|
|
341
227
|
Yields:
|
|
342
228
|
A PsycopgSyncDriver instance.
|
|
343
229
|
"""
|
|
344
230
|
with self.provide_connection(*args, **kwargs) as conn:
|
|
345
|
-
|
|
346
|
-
|
|
347
|
-
if statement_config.allowed_parameter_styles is None:
|
|
348
|
-
from dataclasses import replace
|
|
349
|
-
|
|
350
|
-
statement_config = replace(
|
|
351
|
-
statement_config,
|
|
352
|
-
allowed_parameter_styles=self.supported_parameter_styles,
|
|
353
|
-
default_parameter_style=self.default_parameter_style,
|
|
354
|
-
)
|
|
355
|
-
driver = self.driver_type(connection=conn, config=statement_config)
|
|
356
|
-
yield driver
|
|
231
|
+
final_statement_config = statement_config or self.statement_config
|
|
232
|
+
yield self.driver_type(connection=conn, statement_config=final_statement_config)
|
|
357
233
|
|
|
358
234
|
def provide_pool(self, *args: Any, **kwargs: Any) -> "ConnectionPool":
|
|
359
235
|
"""Provide pool instance.
|
|
@@ -375,184 +251,50 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
|
|
|
375
251
|
Dictionary mapping type names to types.
|
|
376
252
|
"""
|
|
377
253
|
namespace = super().get_signature_namespace()
|
|
378
|
-
namespace.update({"PsycopgSyncConnection": PsycopgSyncConnection})
|
|
254
|
+
namespace.update({"PsycopgSyncConnection": PsycopgSyncConnection, "PsycopgSyncCursor": PsycopgSyncCursor})
|
|
379
255
|
return namespace
|
|
380
256
|
|
|
381
257
|
|
|
382
258
|
class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnectionPool, PsycopgAsyncDriver]):
|
|
383
259
|
"""Configuration for Psycopg asynchronous database connections with direct field-based configuration."""
|
|
384
260
|
|
|
385
|
-
|
|
386
|
-
|
|
387
|
-
|
|
388
|
-
# Driver class reference for dialect resolution
|
|
389
|
-
driver_type: type[PsycopgAsyncDriver] = PsycopgAsyncDriver
|
|
390
|
-
connection_type: type[PsycopgAsyncConnection] = PsycopgAsyncConnection
|
|
391
|
-
|
|
392
|
-
# Parameter style support information
|
|
393
|
-
supported_parameter_styles: ClassVar[tuple[str, ...]] = ("pyformat_positional", "pyformat_named")
|
|
394
|
-
"""Psycopg supports %s (pyformat_positional) and %(name)s (pyformat_named) parameter styles."""
|
|
395
|
-
|
|
396
|
-
default_parameter_style: ClassVar[str] = "pyformat_positional"
|
|
397
|
-
"""Psycopg's preferred parameter style is %s (pyformat_positional)."""
|
|
261
|
+
driver_type: ClassVar[type[PsycopgAsyncDriver]] = PsycopgAsyncDriver
|
|
262
|
+
connection_type: "ClassVar[type[PsycopgAsyncConnection]]" = PsycopgAsyncConnection
|
|
398
263
|
|
|
399
264
|
def __init__(
|
|
400
265
|
self,
|
|
401
|
-
|
|
402
|
-
|
|
403
|
-
|
|
404
|
-
|
|
405
|
-
|
|
406
|
-
port: Optional[int] = None,
|
|
407
|
-
user: Optional[str] = None,
|
|
408
|
-
password: Optional[str] = None,
|
|
409
|
-
dbname: Optional[str] = None,
|
|
410
|
-
connect_timeout: Optional[float] = None,
|
|
411
|
-
options: Optional[str] = None,
|
|
412
|
-
application_name: Optional[str] = None,
|
|
413
|
-
sslmode: Optional[str] = None,
|
|
414
|
-
sslcert: Optional[str] = None,
|
|
415
|
-
sslkey: Optional[str] = None,
|
|
416
|
-
sslrootcert: Optional[str] = None,
|
|
417
|
-
autocommit: Optional[bool] = None,
|
|
418
|
-
# Pool parameters
|
|
419
|
-
min_size: Optional[int] = None,
|
|
420
|
-
max_size: Optional[int] = None,
|
|
421
|
-
name: Optional[str] = None,
|
|
422
|
-
timeout: Optional[float] = None,
|
|
423
|
-
max_waiting: Optional[int] = None,
|
|
424
|
-
max_lifetime: Optional[float] = None,
|
|
425
|
-
max_idle: Optional[float] = None,
|
|
426
|
-
reconnect_timeout: Optional[float] = None,
|
|
427
|
-
num_workers: Optional[int] = None,
|
|
428
|
-
configure: Optional["Callable[[Connection[Any]], None]"] = None,
|
|
429
|
-
kwargs: Optional[dict[str, Any]] = None,
|
|
430
|
-
# User-defined extras
|
|
431
|
-
extras: Optional[dict[str, Any]] = None,
|
|
432
|
-
**additional_kwargs: Any,
|
|
266
|
+
*,
|
|
267
|
+
pool_config: "Optional[Union[PsycopgPoolParams, dict[str, Any]]]" = None,
|
|
268
|
+
pool_instance: "Optional[AsyncConnectionPool]" = None,
|
|
269
|
+
migration_config: "Optional[dict[str, Any]]" = None,
|
|
270
|
+
statement_config: "Optional[StatementConfig]" = None,
|
|
433
271
|
) -> None:
|
|
434
272
|
"""Initialize Psycopg asynchronous configuration.
|
|
435
273
|
|
|
436
274
|
Args:
|
|
275
|
+
pool_config: Pool configuration parameters (TypedDict or dict)
|
|
276
|
+
pool_instance: Existing pool instance to use
|
|
437
277
|
statement_config: Default SQL statement configuration
|
|
438
|
-
|
|
439
|
-
conninfo: Connection string in libpq format
|
|
440
|
-
host: Database server host
|
|
441
|
-
port: Database server port
|
|
442
|
-
user: Database user
|
|
443
|
-
password: Database password
|
|
444
|
-
dbname: Database name
|
|
445
|
-
connect_timeout: Connection timeout in seconds
|
|
446
|
-
options: Command-line options to send to the server
|
|
447
|
-
application_name: Application name for logging and statistics
|
|
448
|
-
sslmode: SSL mode (disable, prefer, require, etc.)
|
|
449
|
-
sslcert: SSL client certificate file
|
|
450
|
-
sslkey: SSL client private key file
|
|
451
|
-
sslrootcert: SSL root certificate file
|
|
452
|
-
autocommit: Enable autocommit mode
|
|
453
|
-
min_size: Minimum number of connections in the pool
|
|
454
|
-
max_size: Maximum number of connections in the pool
|
|
455
|
-
name: Name of the connection pool
|
|
456
|
-
timeout: Timeout for acquiring connections
|
|
457
|
-
max_waiting: Maximum number of waiting clients
|
|
458
|
-
max_lifetime: Maximum connection lifetime
|
|
459
|
-
max_idle: Maximum idle time for connections
|
|
460
|
-
reconnect_timeout: Time between reconnection attempts
|
|
461
|
-
num_workers: Number of background workers
|
|
462
|
-
configure: Callback to configure new connections
|
|
463
|
-
kwargs: Additional connection parameters
|
|
464
|
-
extras: Additional connection parameters not explicitly defined
|
|
465
|
-
**additional_kwargs: Additional parameters (stored in extras)
|
|
466
|
-
"""
|
|
467
|
-
# Store connection parameters as instance attributes
|
|
468
|
-
self.conninfo = conninfo
|
|
469
|
-
self.host = host
|
|
470
|
-
self.port = port
|
|
471
|
-
self.user = user
|
|
472
|
-
self.password = password
|
|
473
|
-
self.dbname = dbname
|
|
474
|
-
self.connect_timeout = connect_timeout
|
|
475
|
-
self.options = options
|
|
476
|
-
self.application_name = application_name
|
|
477
|
-
self.sslmode = sslmode
|
|
478
|
-
self.sslcert = sslcert
|
|
479
|
-
self.sslkey = sslkey
|
|
480
|
-
self.sslrootcert = sslrootcert
|
|
481
|
-
self.autocommit = autocommit
|
|
482
|
-
|
|
483
|
-
# Store pool parameters as instance attributes
|
|
484
|
-
self.min_size = min_size
|
|
485
|
-
self.max_size = max_size
|
|
486
|
-
self.name = name
|
|
487
|
-
self.timeout = timeout
|
|
488
|
-
self.max_waiting = max_waiting
|
|
489
|
-
self.max_lifetime = max_lifetime
|
|
490
|
-
self.max_idle = max_idle
|
|
491
|
-
self.reconnect_timeout = reconnect_timeout
|
|
492
|
-
self.num_workers = num_workers
|
|
493
|
-
self.configure = configure
|
|
494
|
-
self.kwargs = kwargs or {}
|
|
495
|
-
|
|
496
|
-
self.extras = extras or {}
|
|
497
|
-
self.extras.update(additional_kwargs)
|
|
498
|
-
|
|
499
|
-
# Store other config
|
|
500
|
-
self.statement_config = statement_config or SQLConfig()
|
|
501
|
-
self.default_row_type = default_row_type
|
|
502
|
-
|
|
503
|
-
super().__init__()
|
|
504
|
-
|
|
505
|
-
@property
|
|
506
|
-
def connection_config_dict(self) -> dict[str, Any]:
|
|
507
|
-
"""Return the connection configuration as a dict for psycopg operations.
|
|
508
|
-
|
|
509
|
-
Returns only connection-specific parameters.
|
|
510
|
-
"""
|
|
511
|
-
# Gather non-None parameters from connection fields only
|
|
512
|
-
config = {
|
|
513
|
-
field: getattr(self, field)
|
|
514
|
-
for field in CONNECTION_FIELDS
|
|
515
|
-
if getattr(self, field, None) is not None and getattr(self, field) is not Empty
|
|
516
|
-
}
|
|
517
|
-
|
|
518
|
-
# Merge extras and kwargs
|
|
519
|
-
config.update(self.extras)
|
|
520
|
-
if self.kwargs:
|
|
521
|
-
config.update(self.kwargs)
|
|
522
|
-
|
|
523
|
-
config["row_factory"] = dict_row
|
|
524
|
-
|
|
525
|
-
return config
|
|
526
|
-
|
|
527
|
-
@property
|
|
528
|
-
def pool_config_dict(self) -> dict[str, Any]:
|
|
529
|
-
"""Return the pool configuration as a dict for psycopg pool operations.
|
|
530
|
-
|
|
531
|
-
Returns all configuration parameters including connection and pool-specific parameters.
|
|
278
|
+
migration_config: Migration configuration
|
|
532
279
|
"""
|
|
533
|
-
|
|
534
|
-
|
|
535
|
-
|
|
536
|
-
|
|
537
|
-
|
|
538
|
-
|
|
539
|
-
|
|
540
|
-
|
|
541
|
-
|
|
542
|
-
|
|
543
|
-
|
|
544
|
-
|
|
545
|
-
config["row_factory"] = dict_row
|
|
546
|
-
|
|
547
|
-
return config
|
|
280
|
+
processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
|
|
281
|
+
if "extra" in processed_pool_config:
|
|
282
|
+
extras = processed_pool_config.pop("extra")
|
|
283
|
+
processed_pool_config.update(extras)
|
|
284
|
+
|
|
285
|
+
super().__init__(
|
|
286
|
+
pool_config=processed_pool_config,
|
|
287
|
+
pool_instance=pool_instance,
|
|
288
|
+
migration_config=migration_config,
|
|
289
|
+
statement_config=statement_config or psycopg_statement_config,
|
|
290
|
+
)
|
|
548
291
|
|
|
549
292
|
async def _create_pool(self) -> "AsyncConnectionPool":
|
|
550
293
|
"""Create the actual async connection pool."""
|
|
551
294
|
|
|
552
|
-
all_config = self.
|
|
295
|
+
all_config = dict(self.pool_config)
|
|
553
296
|
|
|
554
|
-
|
|
555
|
-
pool_params = {
|
|
297
|
+
pool_parameters = {
|
|
556
298
|
"min_size": all_config.pop("min_size", 4),
|
|
557
299
|
"max_size": all_config.pop("max_size", None),
|
|
558
300
|
"name": all_config.pop("name", None),
|
|
@@ -569,24 +311,30 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
|
|
|
569
311
|
|
|
570
312
|
async def configure_connection(conn: "PsycopgAsyncConnection") -> None:
|
|
571
313
|
conn.row_factory = dict_row
|
|
572
|
-
# Apply autocommit setting if specified (async version requires await)
|
|
573
314
|
if autocommit_setting is not None:
|
|
574
315
|
await conn.set_autocommit(autocommit_setting)
|
|
575
316
|
|
|
576
|
-
|
|
317
|
+
try:
|
|
318
|
+
from pgvector.psycopg import register_vector_async
|
|
319
|
+
|
|
320
|
+
await register_vector_async(conn)
|
|
321
|
+
logger.debug("pgvector registered successfully for psycopg async connection")
|
|
322
|
+
except ImportError:
|
|
323
|
+
pass
|
|
324
|
+
except Exception as e:
|
|
325
|
+
logger.debug("Failed to register pgvector for psycopg async: %s", e)
|
|
326
|
+
|
|
327
|
+
pool_parameters["configure"] = all_config.pop("configure", configure_connection)
|
|
577
328
|
|
|
578
|
-
|
|
329
|
+
pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None}
|
|
579
330
|
|
|
580
331
|
conninfo = all_config.pop("conninfo", None)
|
|
581
332
|
if conninfo:
|
|
582
|
-
|
|
583
|
-
# Don't pass kwargs when using conninfo string
|
|
584
|
-
pool = AsyncConnectionPool(conninfo, open=False, **pool_params)
|
|
333
|
+
pool = AsyncConnectionPool(conninfo, open=False, **pool_parameters)
|
|
585
334
|
else:
|
|
586
|
-
|
|
587
|
-
all_config.
|
|
588
|
-
|
|
589
|
-
pool = AsyncConnectionPool("", kwargs=all_config, open=False, **pool_params)
|
|
335
|
+
kwargs = all_config.pop("kwargs", {})
|
|
336
|
+
all_config.update(kwargs)
|
|
337
|
+
pool = AsyncConnectionPool("", kwargs=all_config, open=False, **pool_parameters)
|
|
590
338
|
|
|
591
339
|
await pool.open()
|
|
592
340
|
|
|
@@ -598,7 +346,6 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
|
|
|
598
346
|
return
|
|
599
347
|
|
|
600
348
|
try:
|
|
601
|
-
# This avoids the "cannot join current thread" error during garbage collection
|
|
602
349
|
if hasattr(self.pool_instance, "_closed"):
|
|
603
350
|
self.pool_instance._closed = True
|
|
604
351
|
|
|
@@ -638,29 +385,22 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
|
|
|
638
385
|
await conn.close()
|
|
639
386
|
|
|
640
387
|
@asynccontextmanager
|
|
641
|
-
async def provide_session(
|
|
388
|
+
async def provide_session(
|
|
389
|
+
self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
|
|
390
|
+
) -> "AsyncGenerator[PsycopgAsyncDriver, None]":
|
|
642
391
|
"""Provide an async driver session context manager.
|
|
643
392
|
|
|
644
393
|
Args:
|
|
645
394
|
*args: Additional arguments.
|
|
395
|
+
statement_config: Optional statement configuration override.
|
|
646
396
|
**kwargs: Additional keyword arguments.
|
|
647
397
|
|
|
648
398
|
Yields:
|
|
649
399
|
A PsycopgAsyncDriver instance.
|
|
650
400
|
"""
|
|
651
401
|
async with self.provide_connection(*args, **kwargs) as conn:
|
|
652
|
-
|
|
653
|
-
|
|
654
|
-
if statement_config.allowed_parameter_styles is None:
|
|
655
|
-
from dataclasses import replace
|
|
656
|
-
|
|
657
|
-
statement_config = replace(
|
|
658
|
-
statement_config,
|
|
659
|
-
allowed_parameter_styles=self.supported_parameter_styles,
|
|
660
|
-
default_parameter_style=self.default_parameter_style,
|
|
661
|
-
)
|
|
662
|
-
driver = self.driver_type(connection=conn, config=statement_config)
|
|
663
|
-
yield driver
|
|
402
|
+
final_statement_config = statement_config or psycopg_statement_config
|
|
403
|
+
yield self.driver_type(connection=conn, statement_config=final_statement_config)
|
|
664
404
|
|
|
665
405
|
async def provide_pool(self, *args: Any, **kwargs: Any) -> "AsyncConnectionPool":
|
|
666
406
|
"""Provide async pool instance.
|
|
@@ -682,5 +422,5 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
|
|
|
682
422
|
Dictionary mapping type names to types.
|
|
683
423
|
"""
|
|
684
424
|
namespace = super().get_signature_namespace()
|
|
685
|
-
namespace.update({"PsycopgAsyncConnection": PsycopgAsyncConnection})
|
|
425
|
+
namespace.update({"PsycopgAsyncConnection": PsycopgAsyncConnection, "PsycopgAsyncCursor": PsycopgAsyncCursor})
|
|
686
426
|
return namespace
|