sqlspec 0.14.1__py3-none-any.whl → 0.15.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of sqlspec might be problematic. Click here for more details.

Files changed (158) hide show
  1. sqlspec/__init__.py +50 -25
  2. sqlspec/__main__.py +1 -1
  3. sqlspec/__metadata__.py +1 -3
  4. sqlspec/_serialization.py +1 -2
  5. sqlspec/_sql.py +256 -120
  6. sqlspec/_typing.py +278 -142
  7. sqlspec/adapters/adbc/__init__.py +4 -3
  8. sqlspec/adapters/adbc/_types.py +12 -0
  9. sqlspec/adapters/adbc/config.py +115 -260
  10. sqlspec/adapters/adbc/driver.py +462 -367
  11. sqlspec/adapters/aiosqlite/__init__.py +18 -3
  12. sqlspec/adapters/aiosqlite/_types.py +13 -0
  13. sqlspec/adapters/aiosqlite/config.py +199 -129
  14. sqlspec/adapters/aiosqlite/driver.py +230 -269
  15. sqlspec/adapters/asyncmy/__init__.py +18 -3
  16. sqlspec/adapters/asyncmy/_types.py +12 -0
  17. sqlspec/adapters/asyncmy/config.py +80 -168
  18. sqlspec/adapters/asyncmy/driver.py +260 -225
  19. sqlspec/adapters/asyncpg/__init__.py +19 -4
  20. sqlspec/adapters/asyncpg/_types.py +17 -0
  21. sqlspec/adapters/asyncpg/config.py +82 -181
  22. sqlspec/adapters/asyncpg/driver.py +285 -383
  23. sqlspec/adapters/bigquery/__init__.py +17 -3
  24. sqlspec/adapters/bigquery/_types.py +12 -0
  25. sqlspec/adapters/bigquery/config.py +191 -258
  26. sqlspec/adapters/bigquery/driver.py +474 -646
  27. sqlspec/adapters/duckdb/__init__.py +14 -3
  28. sqlspec/adapters/duckdb/_types.py +12 -0
  29. sqlspec/adapters/duckdb/config.py +415 -351
  30. sqlspec/adapters/duckdb/driver.py +343 -413
  31. sqlspec/adapters/oracledb/__init__.py +19 -5
  32. sqlspec/adapters/oracledb/_types.py +14 -0
  33. sqlspec/adapters/oracledb/config.py +123 -379
  34. sqlspec/adapters/oracledb/driver.py +507 -560
  35. sqlspec/adapters/psqlpy/__init__.py +13 -3
  36. sqlspec/adapters/psqlpy/_types.py +11 -0
  37. sqlspec/adapters/psqlpy/config.py +93 -254
  38. sqlspec/adapters/psqlpy/driver.py +505 -234
  39. sqlspec/adapters/psycopg/__init__.py +19 -5
  40. sqlspec/adapters/psycopg/_types.py +17 -0
  41. sqlspec/adapters/psycopg/config.py +143 -403
  42. sqlspec/adapters/psycopg/driver.py +706 -872
  43. sqlspec/adapters/sqlite/__init__.py +14 -3
  44. sqlspec/adapters/sqlite/_types.py +11 -0
  45. sqlspec/adapters/sqlite/config.py +202 -118
  46. sqlspec/adapters/sqlite/driver.py +264 -303
  47. sqlspec/base.py +105 -9
  48. sqlspec/{statement/builder → builder}/__init__.py +12 -14
  49. sqlspec/{statement/builder → builder}/_base.py +120 -55
  50. sqlspec/{statement/builder → builder}/_column.py +17 -6
  51. sqlspec/{statement/builder → builder}/_ddl.py +46 -79
  52. sqlspec/{statement/builder → builder}/_ddl_utils.py +5 -10
  53. sqlspec/{statement/builder → builder}/_delete.py +6 -25
  54. sqlspec/{statement/builder → builder}/_insert.py +6 -64
  55. sqlspec/builder/_merge.py +56 -0
  56. sqlspec/{statement/builder → builder}/_parsing_utils.py +3 -10
  57. sqlspec/{statement/builder → builder}/_select.py +11 -56
  58. sqlspec/{statement/builder → builder}/_update.py +12 -18
  59. sqlspec/{statement/builder → builder}/mixins/__init__.py +10 -14
  60. sqlspec/{statement/builder → builder}/mixins/_cte_and_set_ops.py +48 -59
  61. sqlspec/{statement/builder → builder}/mixins/_insert_operations.py +22 -16
  62. sqlspec/{statement/builder → builder}/mixins/_join_operations.py +1 -3
  63. sqlspec/{statement/builder → builder}/mixins/_merge_operations.py +3 -5
  64. sqlspec/{statement/builder → builder}/mixins/_order_limit_operations.py +3 -3
  65. sqlspec/{statement/builder → builder}/mixins/_pivot_operations.py +4 -8
  66. sqlspec/{statement/builder → builder}/mixins/_select_operations.py +21 -36
  67. sqlspec/{statement/builder → builder}/mixins/_update_operations.py +3 -14
  68. sqlspec/{statement/builder → builder}/mixins/_where_clause.py +52 -79
  69. sqlspec/cli.py +4 -5
  70. sqlspec/config.py +180 -133
  71. sqlspec/core/__init__.py +63 -0
  72. sqlspec/core/cache.py +873 -0
  73. sqlspec/core/compiler.py +396 -0
  74. sqlspec/core/filters.py +828 -0
  75. sqlspec/core/hashing.py +310 -0
  76. sqlspec/core/parameters.py +1209 -0
  77. sqlspec/core/result.py +664 -0
  78. sqlspec/{statement → core}/splitter.py +321 -191
  79. sqlspec/core/statement.py +651 -0
  80. sqlspec/driver/__init__.py +7 -10
  81. sqlspec/driver/_async.py +387 -176
  82. sqlspec/driver/_common.py +527 -289
  83. sqlspec/driver/_sync.py +390 -172
  84. sqlspec/driver/mixins/__init__.py +2 -19
  85. sqlspec/driver/mixins/_result_tools.py +168 -0
  86. sqlspec/driver/mixins/_sql_translator.py +6 -3
  87. sqlspec/exceptions.py +5 -252
  88. sqlspec/extensions/aiosql/adapter.py +93 -96
  89. sqlspec/extensions/litestar/config.py +0 -1
  90. sqlspec/extensions/litestar/handlers.py +15 -26
  91. sqlspec/extensions/litestar/plugin.py +16 -14
  92. sqlspec/extensions/litestar/providers.py +17 -52
  93. sqlspec/loader.py +424 -105
  94. sqlspec/migrations/__init__.py +12 -0
  95. sqlspec/migrations/base.py +92 -68
  96. sqlspec/migrations/commands.py +24 -106
  97. sqlspec/migrations/loaders.py +402 -0
  98. sqlspec/migrations/runner.py +49 -51
  99. sqlspec/migrations/tracker.py +31 -44
  100. sqlspec/migrations/utils.py +64 -24
  101. sqlspec/protocols.py +7 -183
  102. sqlspec/storage/__init__.py +1 -1
  103. sqlspec/storage/backends/base.py +37 -40
  104. sqlspec/storage/backends/fsspec.py +136 -112
  105. sqlspec/storage/backends/obstore.py +138 -160
  106. sqlspec/storage/capabilities.py +5 -4
  107. sqlspec/storage/registry.py +57 -106
  108. sqlspec/typing.py +136 -115
  109. sqlspec/utils/__init__.py +2 -3
  110. sqlspec/utils/correlation.py +0 -3
  111. sqlspec/utils/deprecation.py +6 -6
  112. sqlspec/utils/fixtures.py +6 -6
  113. sqlspec/utils/logging.py +0 -2
  114. sqlspec/utils/module_loader.py +7 -12
  115. sqlspec/utils/singleton.py +0 -1
  116. sqlspec/utils/sync_tools.py +16 -37
  117. sqlspec/utils/text.py +12 -51
  118. sqlspec/utils/type_guards.py +443 -232
  119. {sqlspec-0.14.1.dist-info → sqlspec-0.15.0.dist-info}/METADATA +7 -2
  120. sqlspec-0.15.0.dist-info/RECORD +134 -0
  121. sqlspec/adapters/adbc/transformers.py +0 -108
  122. sqlspec/driver/connection.py +0 -207
  123. sqlspec/driver/mixins/_cache.py +0 -114
  124. sqlspec/driver/mixins/_csv_writer.py +0 -91
  125. sqlspec/driver/mixins/_pipeline.py +0 -508
  126. sqlspec/driver/mixins/_query_tools.py +0 -796
  127. sqlspec/driver/mixins/_result_utils.py +0 -138
  128. sqlspec/driver/mixins/_storage.py +0 -912
  129. sqlspec/driver/mixins/_type_coercion.py +0 -128
  130. sqlspec/driver/parameters.py +0 -138
  131. sqlspec/statement/__init__.py +0 -21
  132. sqlspec/statement/builder/_merge.py +0 -95
  133. sqlspec/statement/cache.py +0 -50
  134. sqlspec/statement/filters.py +0 -625
  135. sqlspec/statement/parameters.py +0 -956
  136. sqlspec/statement/pipelines/__init__.py +0 -210
  137. sqlspec/statement/pipelines/analyzers/__init__.py +0 -9
  138. sqlspec/statement/pipelines/analyzers/_analyzer.py +0 -646
  139. sqlspec/statement/pipelines/context.py +0 -109
  140. sqlspec/statement/pipelines/transformers/__init__.py +0 -7
  141. sqlspec/statement/pipelines/transformers/_expression_simplifier.py +0 -88
  142. sqlspec/statement/pipelines/transformers/_literal_parameterizer.py +0 -1247
  143. sqlspec/statement/pipelines/transformers/_remove_comments_and_hints.py +0 -76
  144. sqlspec/statement/pipelines/validators/__init__.py +0 -23
  145. sqlspec/statement/pipelines/validators/_dml_safety.py +0 -290
  146. sqlspec/statement/pipelines/validators/_parameter_style.py +0 -370
  147. sqlspec/statement/pipelines/validators/_performance.py +0 -714
  148. sqlspec/statement/pipelines/validators/_security.py +0 -967
  149. sqlspec/statement/result.py +0 -435
  150. sqlspec/statement/sql.py +0 -1774
  151. sqlspec/utils/cached_property.py +0 -25
  152. sqlspec/utils/statement_hashing.py +0 -203
  153. sqlspec-0.14.1.dist-info/RECORD +0 -145
  154. /sqlspec/{statement/builder → builder}/mixins/_delete_operations.py +0 -0
  155. {sqlspec-0.14.1.dist-info → sqlspec-0.15.0.dist-info}/WHEEL +0 -0
  156. {sqlspec-0.14.1.dist-info → sqlspec-0.15.0.dist-info}/entry_points.txt +0 -0
  157. {sqlspec-0.14.1.dist-info → sqlspec-0.15.0.dist-info}/licenses/LICENSE +0 -0
  158. {sqlspec-0.14.1.dist-info → sqlspec-0.15.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
- PsycopgAsyncConnection,
14
+ PsycopgAsyncCursor,
13
15
  PsycopgAsyncDriver,
14
- PsycopgSyncConnection,
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 psycopg import Connection
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
- POOL_FIELDS = CONNECTION_FIELDS.union(
48
- {
49
- "min_size",
50
- "max_size",
51
- "name",
52
- "timeout",
53
- "max_waiting",
54
- "max_lifetime",
55
- "max_idle",
56
- "reconnect_timeout",
57
- "num_workers",
58
- "configure",
59
- "kwargs",
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
- is_async: ClassVar[bool] = False
70
- supports_connection_pooling: ClassVar[bool] = True
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
- statement_config: "Optional[SQLConfig]" = None,
85
- default_row_type: "type[DictRow]" = DictRow,
86
- # Connection parameters
87
- conninfo: Optional[str] = None,
88
- host: Optional[str] = None,
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
- default_row_type: Default row type for results
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
- # Gather non-None parameters from all fields (connection + pool)
217
- config = {
218
- field: getattr(self, field)
219
- for field in POOL_FIELDS
220
- if getattr(self, field, None) is not None and getattr(self, field) is not Empty
221
- }
222
-
223
- # Merge extras and kwargs
224
- config.update(self.extras)
225
- if self.kwargs:
226
- config.update(self.kwargs)
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.pool_config_dict.copy()
117
+ all_config = dict(self.pool_config)
238
118
 
239
- # Separate pool-specific parameters that ConnectionPool accepts directly
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
- pool_params["configure"] = all_config.pop("configure", configure_connection)
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
- pool_params = {k: v for k, v in pool_params.items() if v is not None}
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
- # If conninfo is provided, use it directly
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
- # row_factory is already popped out earlier
272
- all_config.pop("row_factory", None)
273
- all_config.pop("kwargs", None)
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(self, *args: Any, **kwargs: Any) -> "Generator[PsycopgSyncDriver, None, None]":
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
- statement_config = self.statement_config
346
- # Inject parameter style info if not already set
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
- is_async: ClassVar[bool] = True
386
- supports_connection_pooling: ClassVar[bool] = True
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
- statement_config: "Optional[SQLConfig]" = None,
402
- default_row_type: "type[DictRow]" = DictRow,
403
- # Connection parameters
404
- conninfo: Optional[str] = None,
405
- host: Optional[str] = None,
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
- default_row_type: Default row type for results
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
- # Gather non-None parameters from all fields (connection + pool)
534
- config = {
535
- field: getattr(self, field)
536
- for field in POOL_FIELDS
537
- if getattr(self, field, None) is not None and getattr(self, field) is not Empty
538
- }
539
-
540
- # Merge extras and kwargs
541
- config.update(self.extras)
542
- if self.kwargs:
543
- config.update(self.kwargs)
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.pool_config_dict.copy()
295
+ all_config = dict(self.pool_config)
553
296
 
554
- # Separate pool-specific parameters that AsyncConnectionPool accepts directly
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
- pool_params["configure"] = all_config.pop("configure", configure_connection)
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
- pool_params = {k: v for k, v in pool_params.items() if v is not None}
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
- # If conninfo is provided, use it directly
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
- # row_factory is already popped out earlier
587
- all_config.pop("row_factory", None)
588
- all_config.pop("kwargs", None)
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(self, *args: Any, **kwargs: Any) -> "AsyncGenerator[PsycopgAsyncDriver, None]":
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
- statement_config = self.statement_config
653
- # Inject parameter style info if not already set
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