sqlspec 0.24.1__py3-none-any.whl → 0.26.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 (95) hide show
  1. sqlspec/_serialization.py +223 -21
  2. sqlspec/_sql.py +20 -62
  3. sqlspec/_typing.py +11 -0
  4. sqlspec/adapters/adbc/config.py +8 -1
  5. sqlspec/adapters/adbc/data_dictionary.py +290 -0
  6. sqlspec/adapters/adbc/driver.py +129 -20
  7. sqlspec/adapters/adbc/type_converter.py +159 -0
  8. sqlspec/adapters/aiosqlite/config.py +3 -0
  9. sqlspec/adapters/aiosqlite/data_dictionary.py +117 -0
  10. sqlspec/adapters/aiosqlite/driver.py +17 -3
  11. sqlspec/adapters/asyncmy/_types.py +1 -1
  12. sqlspec/adapters/asyncmy/config.py +11 -8
  13. sqlspec/adapters/asyncmy/data_dictionary.py +122 -0
  14. sqlspec/adapters/asyncmy/driver.py +31 -7
  15. sqlspec/adapters/asyncpg/config.py +3 -0
  16. sqlspec/adapters/asyncpg/data_dictionary.py +134 -0
  17. sqlspec/adapters/asyncpg/driver.py +19 -4
  18. sqlspec/adapters/bigquery/config.py +3 -0
  19. sqlspec/adapters/bigquery/data_dictionary.py +109 -0
  20. sqlspec/adapters/bigquery/driver.py +21 -3
  21. sqlspec/adapters/bigquery/type_converter.py +93 -0
  22. sqlspec/adapters/duckdb/_types.py +1 -1
  23. sqlspec/adapters/duckdb/config.py +2 -0
  24. sqlspec/adapters/duckdb/data_dictionary.py +124 -0
  25. sqlspec/adapters/duckdb/driver.py +32 -5
  26. sqlspec/adapters/duckdb/pool.py +1 -1
  27. sqlspec/adapters/duckdb/type_converter.py +103 -0
  28. sqlspec/adapters/oracledb/config.py +6 -0
  29. sqlspec/adapters/oracledb/data_dictionary.py +442 -0
  30. sqlspec/adapters/oracledb/driver.py +68 -9
  31. sqlspec/adapters/oracledb/migrations.py +51 -67
  32. sqlspec/adapters/oracledb/type_converter.py +132 -0
  33. sqlspec/adapters/psqlpy/config.py +3 -0
  34. sqlspec/adapters/psqlpy/data_dictionary.py +133 -0
  35. sqlspec/adapters/psqlpy/driver.py +23 -179
  36. sqlspec/adapters/psqlpy/type_converter.py +73 -0
  37. sqlspec/adapters/psycopg/config.py +8 -4
  38. sqlspec/adapters/psycopg/data_dictionary.py +257 -0
  39. sqlspec/adapters/psycopg/driver.py +40 -5
  40. sqlspec/adapters/sqlite/config.py +3 -0
  41. sqlspec/adapters/sqlite/data_dictionary.py +117 -0
  42. sqlspec/adapters/sqlite/driver.py +18 -3
  43. sqlspec/adapters/sqlite/pool.py +13 -4
  44. sqlspec/base.py +3 -4
  45. sqlspec/builder/_base.py +130 -48
  46. sqlspec/builder/_column.py +66 -24
  47. sqlspec/builder/_ddl.py +91 -41
  48. sqlspec/builder/_insert.py +40 -58
  49. sqlspec/builder/_parsing_utils.py +127 -12
  50. sqlspec/builder/_select.py +147 -2
  51. sqlspec/builder/_update.py +1 -1
  52. sqlspec/builder/mixins/_cte_and_set_ops.py +31 -23
  53. sqlspec/builder/mixins/_delete_operations.py +12 -7
  54. sqlspec/builder/mixins/_insert_operations.py +50 -36
  55. sqlspec/builder/mixins/_join_operations.py +15 -30
  56. sqlspec/builder/mixins/_merge_operations.py +210 -78
  57. sqlspec/builder/mixins/_order_limit_operations.py +4 -10
  58. sqlspec/builder/mixins/_pivot_operations.py +1 -0
  59. sqlspec/builder/mixins/_select_operations.py +44 -22
  60. sqlspec/builder/mixins/_update_operations.py +30 -37
  61. sqlspec/builder/mixins/_where_clause.py +52 -70
  62. sqlspec/cli.py +246 -140
  63. sqlspec/config.py +33 -19
  64. sqlspec/core/__init__.py +3 -2
  65. sqlspec/core/cache.py +298 -352
  66. sqlspec/core/compiler.py +61 -4
  67. sqlspec/core/filters.py +246 -213
  68. sqlspec/core/hashing.py +9 -11
  69. sqlspec/core/parameters.py +27 -10
  70. sqlspec/core/statement.py +72 -12
  71. sqlspec/core/type_conversion.py +234 -0
  72. sqlspec/driver/__init__.py +6 -3
  73. sqlspec/driver/_async.py +108 -5
  74. sqlspec/driver/_common.py +186 -17
  75. sqlspec/driver/_sync.py +108 -5
  76. sqlspec/driver/mixins/_result_tools.py +60 -7
  77. sqlspec/exceptions.py +5 -0
  78. sqlspec/loader.py +8 -9
  79. sqlspec/migrations/__init__.py +4 -3
  80. sqlspec/migrations/base.py +153 -14
  81. sqlspec/migrations/commands.py +34 -96
  82. sqlspec/migrations/context.py +145 -0
  83. sqlspec/migrations/loaders.py +25 -8
  84. sqlspec/migrations/runner.py +352 -82
  85. sqlspec/storage/backends/fsspec.py +1 -0
  86. sqlspec/typing.py +4 -0
  87. sqlspec/utils/config_resolver.py +153 -0
  88. sqlspec/utils/serializers.py +50 -2
  89. {sqlspec-0.24.1.dist-info → sqlspec-0.26.0.dist-info}/METADATA +1 -1
  90. sqlspec-0.26.0.dist-info/RECORD +157 -0
  91. sqlspec-0.24.1.dist-info/RECORD +0 -139
  92. {sqlspec-0.24.1.dist-info → sqlspec-0.26.0.dist-info}/WHEEL +0 -0
  93. {sqlspec-0.24.1.dist-info → sqlspec-0.26.0.dist-info}/entry_points.txt +0 -0
  94. {sqlspec-0.24.1.dist-info → sqlspec-0.26.0.dist-info}/licenses/LICENSE +0 -0
  95. {sqlspec-0.24.1.dist-info → sqlspec-0.26.0.dist-info}/licenses/NOTICE +0 -0
@@ -4,15 +4,14 @@ Provides parameter style conversion, type coercion, error handling,
4
4
  and transaction management.
5
5
  """
6
6
 
7
- import datetime
8
7
  import decimal
9
8
  import re
10
- import uuid
11
9
  from typing import TYPE_CHECKING, Any, Final, Optional
12
10
 
13
11
  import psqlpy
14
12
  import psqlpy.exceptions
15
13
 
14
+ from sqlspec.adapters.psqlpy.type_converter import PostgreSQLTypeConverter
16
15
  from sqlspec.core.cache import get_cache_config
17
16
  from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
18
17
  from sqlspec.core.statement import SQL, StatementConfig
@@ -26,11 +25,14 @@ if TYPE_CHECKING:
26
25
  from sqlspec.adapters.psqlpy._types import PsqlpyConnection
27
26
  from sqlspec.core.result import SQLResult
28
27
  from sqlspec.driver import ExecutionResult
28
+ from sqlspec.driver._async import AsyncDataDictionaryBase
29
29
 
30
30
  __all__ = ("PsqlpyCursor", "PsqlpyDriver", "PsqlpyExceptionHandler", "psqlpy_statement_config")
31
31
 
32
32
  logger = get_logger("adapters.psqlpy")
33
33
 
34
+ _type_converter = PostgreSQLTypeConverter()
35
+
34
36
  psqlpy_statement_config = StatementConfig(
35
37
  dialect="postgres",
36
38
  parameter_config=ParameterStyleConfig(
@@ -38,7 +40,7 @@ psqlpy_statement_config = StatementConfig(
38
40
  supported_parameter_styles={ParameterStyle.NUMERIC, ParameterStyle.NAMED_DOLLAR, ParameterStyle.QMARK},
39
41
  default_execution_parameter_style=ParameterStyle.NUMERIC,
40
42
  supported_execution_parameter_styles={ParameterStyle.NUMERIC},
41
- type_coercion_map={tuple: list, decimal.Decimal: float},
43
+ type_coercion_map={tuple: list, decimal.Decimal: float, str: _type_converter.convert_if_detected},
42
44
  has_native_list_expansion=False,
43
45
  needs_static_script_compilation=False,
44
46
  allow_mixed_parameter_styles=False,
@@ -52,173 +54,6 @@ psqlpy_statement_config = StatementConfig(
52
54
 
53
55
  PSQLPY_STATUS_REGEX: Final[re.Pattern[str]] = re.compile(r"^([A-Z]+)(?:\s+(\d+))?\s+(\d+)$", re.IGNORECASE)
54
56
 
55
- SPECIAL_TYPE_REGEX: Final[re.Pattern[str]] = re.compile(
56
- r"^(?:"
57
- r"(?P<uuid>[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})|"
58
- r"(?P<ipv4>(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)(?:/(?:3[0-2]|[12]?[0-9]))?)|"
59
- r"(?P<ipv6>(?:(?:[0-9a-f]{1,4}:){7}[0-9a-f]{1,4}|(?:[0-9a-f]{1,4}:){1,7}:|:(?::[0-9a-f]{1,4}){1,7}|(?:[0-9a-f]{1,4}:){1,6}:[0-9a-f]{1,4}|::(?:ffff:)?(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?))(?:/(?:12[0-8]|1[01][0-9]|[1-9]?[0-9]))?)|"
60
- r"(?P<mac>(?:[0-9a-f]{2}[:-]){5}[0-9a-f]{2}|[0-9a-f]{12})|"
61
- r"(?P<iso_datetime>\d{4}-\d{2}-\d{2}[T\s]\d{2}:\d{2}:\d{2}(?:\.\d{1,6})?(?:Z|[+-]\d{2}:?\d{2})?)|"
62
- r"(?P<iso_date>\d{4}-\d{2}-\d{2})|"
63
- r"(?P<iso_time>\d{2}:\d{2}:\d{2}(?:\.\d{1,6})?(?:Z|[+-]\d{2}:?\d{2})?)|"
64
- r"(?P<interval>(?:(?:\d+\s+(?:year|month|day|hour|minute|second)s?\s*)+)|(?:P(?:\d+Y)?(?:\d+M)?(?:\d+D)?(?:T(?:\d+H)?(?:\d+M)?(?:\d+(?:\.\d+)?S)?)?))|"
65
- r"(?P<json>\{[\s\S]*\}|\[[\s\S]*\])|"
66
- r"(?P<pg_array>\{(?:[^{}]+|\{[^{}]*\})*\})"
67
- r")$",
68
- re.IGNORECASE,
69
- )
70
-
71
-
72
- def _detect_postgresql_type(value: str) -> Optional[str]:
73
- """Detect PostgreSQL data type from string value.
74
-
75
- Args:
76
- value: String value to analyze
77
-
78
- Returns:
79
- Type name if detected, None otherwise.
80
- """
81
- match = SPECIAL_TYPE_REGEX.match(value)
82
- if not match:
83
- return None
84
-
85
- for group_name in [
86
- "uuid",
87
- "ipv4",
88
- "ipv6",
89
- "mac",
90
- "iso_datetime",
91
- "iso_date",
92
- "iso_time",
93
- "interval",
94
- "json",
95
- "pg_array",
96
- ]:
97
- if match.group(group_name):
98
- return group_name
99
-
100
- return None
101
-
102
-
103
- def _convert_uuid(value: str) -> Any:
104
- """Convert UUID string to UUID object.
105
-
106
- Args:
107
- value: UUID string to convert
108
-
109
- Returns:
110
- UUID object or original value if conversion fails
111
- """
112
- try:
113
- clean_uuid = value.replace("-", "").lower()
114
- uuid_length = 32
115
- if len(clean_uuid) == uuid_length:
116
- formatted = f"{clean_uuid[:8]}-{clean_uuid[8:12]}-{clean_uuid[12:16]}-{clean_uuid[16:20]}-{clean_uuid[20:]}"
117
- return uuid.UUID(formatted)
118
- return uuid.UUID(value)
119
- except (ValueError, AttributeError):
120
- return value
121
-
122
-
123
- def _convert_iso_datetime(value: str) -> Any:
124
- """Convert ISO datetime string to datetime object.
125
-
126
- Args:
127
- value: ISO datetime string to convert
128
-
129
- Returns:
130
- datetime object or original value if conversion fails
131
- """
132
- try:
133
- normalized = value.replace("Z", "+00:00")
134
- return datetime.datetime.fromisoformat(normalized)
135
- except ValueError:
136
- return value
137
-
138
-
139
- def _convert_iso_date(value: str) -> Any:
140
- """Convert ISO date string to date object.
141
-
142
- Args:
143
- value: ISO date string to convert
144
-
145
- Returns:
146
- date object or original value if conversion fails
147
- """
148
- try:
149
- return datetime.date.fromisoformat(value)
150
- except ValueError:
151
- return value
152
-
153
-
154
- def _validate_json(value: str) -> str:
155
- """Validate JSON string format.
156
-
157
- Args:
158
- value: JSON string to validate
159
-
160
- Returns:
161
- Original string value
162
- """
163
- from sqlspec.utils.serializers import from_json
164
-
165
- try:
166
- from_json(value)
167
- except (ValueError, TypeError):
168
- return value
169
- return value
170
-
171
-
172
- def _passthrough(value: str) -> str:
173
- """Pass value through unchanged.
174
-
175
- Args:
176
- value: String value to pass through
177
-
178
- Returns:
179
- Original value unchanged
180
- """
181
- return value
182
-
183
-
184
- _PSQLPY_TYPE_CONVERTERS: dict[str, Any] = {
185
- "uuid": _convert_uuid,
186
- "iso_datetime": _convert_iso_datetime,
187
- "iso_date": _convert_iso_date,
188
- "iso_time": _passthrough,
189
- "json": _validate_json,
190
- "pg_array": _passthrough,
191
- "ipv4": _passthrough,
192
- "ipv6": _passthrough,
193
- "mac": _passthrough,
194
- "interval": _passthrough,
195
- }
196
-
197
-
198
- def _convert_psqlpy_parameters(value: Any) -> Any:
199
- """Convert parameters for psqlpy compatibility.
200
-
201
- Args:
202
- value: Parameter value to convert
203
-
204
- Returns:
205
- Converted value suitable for psqlpy execution
206
- """
207
- if isinstance(value, str):
208
- detected_type = _detect_postgresql_type(value)
209
-
210
- if detected_type:
211
- converter = _PSQLPY_TYPE_CONVERTERS.get(detected_type)
212
- if converter:
213
- return converter(value)
214
-
215
- return value
216
-
217
- if isinstance(value, (dict, list, tuple, uuid.UUID, datetime.datetime, datetime.date)):
218
- return value
219
-
220
- return value
221
-
222
57
 
223
58
  class PsqlpyCursor:
224
59
  """Context manager for psqlpy cursor management."""
@@ -238,7 +73,7 @@ class PsqlpyCursor:
238
73
  self._in_use = True
239
74
  return self.connection
240
75
 
241
- async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
76
+ async def __aexit__(self, *_: Any) -> None:
242
77
  """Exit cursor context.
243
78
 
244
79
  Args:
@@ -246,7 +81,6 @@ class PsqlpyCursor:
246
81
  exc_val: Exception value
247
82
  exc_tb: Exception traceback
248
83
  """
249
- _ = (exc_type, exc_val, exc_tb)
250
84
  self._in_use = False
251
85
 
252
86
  def is_in_use(self) -> bool:
@@ -303,7 +137,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
303
137
  and transaction management.
304
138
  """
305
139
 
306
- __slots__ = ()
140
+ __slots__ = ("_data_dictionary",)
307
141
  dialect = "postgres"
308
142
 
309
143
  def __init__(
@@ -322,6 +156,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
322
156
  )
323
157
 
324
158
  super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
159
+ self._data_dictionary: Optional[AsyncDataDictionaryBase] = None
325
160
 
326
161
  def with_cursor(self, connection: "PsqlpyConnection") -> "PsqlpyCursor":
327
162
  """Create context manager for psqlpy cursor.
@@ -404,10 +239,9 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
404
239
  formatted_parameters = []
405
240
  for param_set in prepared_parameters:
406
241
  if isinstance(param_set, (list, tuple)):
407
- converted_params = [_convert_psqlpy_parameters(param) for param in param_set]
408
- formatted_parameters.append(converted_params)
242
+ formatted_parameters.append(list(param_set))
409
243
  else:
410
- formatted_parameters.append([_convert_psqlpy_parameters(param_set)])
244
+ formatted_parameters.append([param_set])
411
245
 
412
246
  await cursor.execute_many(sql, formatted_parameters)
413
247
 
@@ -427,9 +261,6 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
427
261
  """
428
262
  sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
429
263
 
430
- if prepared_parameters:
431
- prepared_parameters = [_convert_psqlpy_parameters(param) for param in prepared_parameters]
432
-
433
264
  if statement.returns_rows():
434
265
  query_result = await cursor.fetch(sql, prepared_parameters or [])
435
266
  dict_rows: list[dict[str, Any]] = query_result.result() if query_result else []
@@ -511,3 +342,16 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
511
342
  except psqlpy.exceptions.DatabaseError as e:
512
343
  msg = f"Failed to commit psqlpy transaction: {e}"
513
344
  raise SQLSpecError(msg) from e
345
+
346
+ @property
347
+ def data_dictionary(self) -> "AsyncDataDictionaryBase":
348
+ """Get the data dictionary for this driver.
349
+
350
+ Returns:
351
+ Data dictionary instance for metadata queries
352
+ """
353
+ if self._data_dictionary is None:
354
+ from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyAsyncDataDictionary
355
+
356
+ self._data_dictionary = PsqlpyAsyncDataDictionary()
357
+ return self._data_dictionary
@@ -0,0 +1,73 @@
1
+ """PostgreSQL-specific type conversion for psqlpy adapter.
2
+
3
+ Provides specialized type handling for PostgreSQL databases, including
4
+ PostgreSQL-specific types like intervals and arrays while preserving
5
+ backward compatibility.
6
+ """
7
+
8
+ import re
9
+ from typing import Any, Final, Optional
10
+
11
+ from sqlspec.core.type_conversion import BaseTypeConverter
12
+
13
+ # PostgreSQL-specific regex patterns for types not covered by base BaseTypeConverter
14
+ PG_SPECIFIC_REGEX: Final[re.Pattern[str]] = re.compile(
15
+ r"^(?:"
16
+ r"(?P<interval>(?:(?:\d+\s+(?:year|month|day|hour|minute|second)s?\s*)+)|(?:P(?:\d+Y)?(?:\d+M)?(?:\d+D)?(?:T(?:\d+H)?(?:\d+M)?(?:\d+(?:\.\d+)?S)?)?))|"
17
+ r"(?P<pg_array>\{(?:[^{}]+|\{[^{}]*\})*\})"
18
+ r")$",
19
+ re.IGNORECASE,
20
+ )
21
+
22
+
23
+ class PostgreSQLTypeConverter(BaseTypeConverter):
24
+ """PostgreSQL-specific type converter with interval and array support.
25
+
26
+ Extends the base BaseTypeConverter with PostgreSQL-specific functionality
27
+ while maintaining backward compatibility for interval and array types.
28
+ """
29
+
30
+ __slots__ = ()
31
+
32
+ def detect_type(self, value: str) -> Optional[str]:
33
+ """Detect types including PostgreSQL-specific types.
34
+
35
+ Args:
36
+ value: String value to analyze.
37
+
38
+ Returns:
39
+ Type name if detected, None otherwise.
40
+ """
41
+ # First try generic types (UUID, JSON, datetime, etc.)
42
+ detected_type = super().detect_type(value)
43
+ if detected_type:
44
+ return detected_type
45
+
46
+ # Then check PostgreSQL-specific types
47
+ match = PG_SPECIFIC_REGEX.match(value)
48
+ if match:
49
+ for group_name in ["interval", "pg_array"]:
50
+ if match.group(group_name):
51
+ return group_name
52
+
53
+ return None
54
+
55
+ def convert_value(self, value: str, detected_type: str) -> Any:
56
+ """Convert value with PostgreSQL-specific handling.
57
+
58
+ Args:
59
+ value: String value to convert.
60
+ detected_type: Detected type name.
61
+
62
+ Returns:
63
+ Converted value or original string for PostgreSQL-specific types.
64
+ """
65
+ # For PostgreSQL-specific types, preserve as strings for backward compatibility
66
+ if detected_type in ("interval", "pg_array"):
67
+ return value # Pass through as strings - psqlpy will handle casting
68
+
69
+ # Use base converter for standard types
70
+ return super().convert_value(value, detected_type)
71
+
72
+
73
+ __all__ = ("PG_SPECIFIC_REGEX", "PostgreSQLTypeConverter")
@@ -88,6 +88,7 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
88
88
  migration_config: Optional[dict[str, Any]] = None,
89
89
  statement_config: "Optional[StatementConfig]" = None,
90
90
  driver_features: "Optional[dict[str, Any]]" = None,
91
+ bind_key: "Optional[str]" = None,
91
92
  ) -> None:
92
93
  """Initialize Psycopg synchronous configuration.
93
94
 
@@ -97,6 +98,7 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
97
98
  migration_config: Migration configuration
98
99
  statement_config: Default SQL statement configuration
99
100
  driver_features: Optional driver feature configuration
101
+ bind_key: Optional unique identifier for this configuration
100
102
  """
101
103
  processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
102
104
  if "extra" in processed_pool_config:
@@ -109,6 +111,7 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
109
111
  migration_config=migration_config,
110
112
  statement_config=statement_config or psycopg_statement_config,
111
113
  driver_features=driver_features or {},
114
+ bind_key=bind_key,
112
115
  )
113
116
 
114
117
  def _create_pool(self) -> "ConnectionPool":
@@ -173,8 +176,7 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
173
176
  logger.info("Closing Psycopg connection pool", extra={"adapter": "psycopg"})
174
177
 
175
178
  try:
176
- if hasattr(self.pool_instance, "_closed"):
177
- self.pool_instance._closed = True
179
+ self.pool_instance._closed = True # pyright: ignore[reportPrivateUsage]
178
180
 
179
181
  self.pool_instance.close()
180
182
  logger.info("Psycopg connection pool closed successfully", extra={"adapter": "psycopg"})
@@ -271,6 +273,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
271
273
  migration_config: "Optional[dict[str, Any]]" = None,
272
274
  statement_config: "Optional[StatementConfig]" = None,
273
275
  driver_features: "Optional[dict[str, Any]]" = None,
276
+ bind_key: "Optional[str]" = None,
274
277
  ) -> None:
275
278
  """Initialize Psycopg asynchronous configuration.
276
279
 
@@ -280,6 +283,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
280
283
  migration_config: Migration configuration
281
284
  statement_config: Default SQL statement configuration
282
285
  driver_features: Optional driver feature configuration
286
+ bind_key: Optional unique identifier for this configuration
283
287
  """
284
288
  processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
285
289
  if "extra" in processed_pool_config:
@@ -292,6 +296,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
292
296
  migration_config=migration_config,
293
297
  statement_config=statement_config or psycopg_statement_config,
294
298
  driver_features=driver_features or {},
299
+ bind_key=bind_key,
295
300
  )
296
301
 
297
302
  async def _create_pool(self) -> "AsyncConnectionPool":
@@ -350,8 +355,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
350
355
  return
351
356
 
352
357
  try:
353
- if hasattr(self.pool_instance, "_closed"):
354
- self.pool_instance._closed = True
358
+ self.pool_instance._closed = True # pyright: ignore[reportPrivateUsage]
355
359
 
356
360
  await self.pool_instance.close()
357
361
  finally:
@@ -0,0 +1,257 @@
1
+ """PostgreSQL-specific data dictionary for metadata queries via psycopg."""
2
+
3
+ import re
4
+ from typing import TYPE_CHECKING, Optional, cast
5
+
6
+ from sqlspec.driver import (
7
+ AsyncDataDictionaryBase,
8
+ AsyncDriverAdapterBase,
9
+ SyncDataDictionaryBase,
10
+ SyncDriverAdapterBase,
11
+ VersionInfo,
12
+ )
13
+ from sqlspec.utils.logging import get_logger
14
+
15
+ if TYPE_CHECKING:
16
+ from collections.abc import Callable
17
+
18
+ from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver
19
+
20
+ logger = get_logger("adapters.psycopg.data_dictionary")
21
+
22
+ # Compiled regex patterns
23
+ POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?")
24
+
25
+ __all__ = ("PostgresAsyncDataDictionary", "PostgresSyncDataDictionary")
26
+
27
+
28
+ class PostgresSyncDataDictionary(SyncDataDictionaryBase):
29
+ """PostgreSQL-specific sync data dictionary."""
30
+
31
+ def get_version(self, driver: SyncDriverAdapterBase) -> "Optional[VersionInfo]":
32
+ """Get PostgreSQL database version information.
33
+
34
+ Args:
35
+ driver: Sync database driver instance
36
+
37
+ Returns:
38
+ PostgreSQL version information or None if detection fails
39
+ """
40
+ version_str = cast("PsycopgSyncDriver", driver).select_value("SELECT version()")
41
+ if not version_str:
42
+ logger.warning("No PostgreSQL version information found")
43
+ return None
44
+
45
+ # Parse version like "PostgreSQL 15.3 on x86_64-pc-linux-gnu..."
46
+ version_match = POSTGRES_VERSION_PATTERN.search(str(version_str))
47
+ if not version_match:
48
+ logger.warning("Could not parse PostgreSQL version: %s", version_str)
49
+ return None
50
+
51
+ major = int(version_match.group(1))
52
+ minor = int(version_match.group(2))
53
+ patch = int(version_match.group(3)) if version_match.group(3) else 0
54
+
55
+ version_info = VersionInfo(major, minor, patch)
56
+ logger.debug("Detected PostgreSQL version: %s", version_info)
57
+ return version_info
58
+
59
+ def get_feature_flag(self, driver: SyncDriverAdapterBase, feature: str) -> bool:
60
+ """Check if PostgreSQL database supports a specific feature.
61
+
62
+ Args:
63
+ driver: Sync database driver instance
64
+ feature: Feature name to check
65
+
66
+ Returns:
67
+ True if feature is supported, False otherwise
68
+ """
69
+ version_info = self.get_version(driver)
70
+ if not version_info:
71
+ return False
72
+
73
+ feature_checks: dict[str, Callable[[VersionInfo], bool]] = {
74
+ "supports_json": lambda v: v >= VersionInfo(9, 2, 0),
75
+ "supports_jsonb": lambda v: v >= VersionInfo(9, 4, 0),
76
+ "supports_uuid": lambda _: True, # UUID extension widely available
77
+ "supports_arrays": lambda _: True, # PostgreSQL has excellent array support
78
+ "supports_returning": lambda v: v >= VersionInfo(8, 2, 0),
79
+ "supports_upsert": lambda v: v >= VersionInfo(9, 5, 0), # ON CONFLICT
80
+ "supports_window_functions": lambda v: v >= VersionInfo(8, 4, 0),
81
+ "supports_cte": lambda v: v >= VersionInfo(8, 4, 0),
82
+ "supports_transactions": lambda _: True,
83
+ "supports_prepared_statements": lambda _: True,
84
+ "supports_schemas": lambda _: True,
85
+ "supports_partitioning": lambda v: v >= VersionInfo(10, 0, 0),
86
+ }
87
+
88
+ if feature in feature_checks:
89
+ return bool(feature_checks[feature](version_info))
90
+
91
+ return False
92
+
93
+ def get_optimal_type(self, driver: SyncDriverAdapterBase, type_category: str) -> str:
94
+ """Get optimal PostgreSQL type for a category.
95
+
96
+ Args:
97
+ driver: Sync database driver instance
98
+ type_category: Type category
99
+
100
+ Returns:
101
+ PostgreSQL-specific type name
102
+ """
103
+ version_info = self.get_version(driver)
104
+
105
+ if type_category == "json":
106
+ if version_info and version_info >= VersionInfo(9, 4, 0):
107
+ return "JSONB" # Prefer JSONB over JSON
108
+ if version_info and version_info >= VersionInfo(9, 2, 0):
109
+ return "JSON"
110
+ return "TEXT"
111
+
112
+ type_map = {
113
+ "uuid": "UUID",
114
+ "boolean": "BOOLEAN",
115
+ "timestamp": "TIMESTAMP WITH TIME ZONE",
116
+ "text": "TEXT",
117
+ "blob": "BYTEA",
118
+ "array": "ARRAY",
119
+ }
120
+ return type_map.get(type_category, "TEXT")
121
+
122
+ def list_available_features(self) -> "list[str]":
123
+ """List available PostgreSQL feature flags.
124
+
125
+ Returns:
126
+ List of supported feature names
127
+ """
128
+ return [
129
+ "supports_json",
130
+ "supports_jsonb",
131
+ "supports_uuid",
132
+ "supports_arrays",
133
+ "supports_returning",
134
+ "supports_upsert",
135
+ "supports_window_functions",
136
+ "supports_cte",
137
+ "supports_transactions",
138
+ "supports_prepared_statements",
139
+ "supports_schemas",
140
+ "supports_partitioning",
141
+ ]
142
+
143
+
144
+ class PostgresAsyncDataDictionary(AsyncDataDictionaryBase):
145
+ """PostgreSQL-specific async data dictionary."""
146
+
147
+ async def get_version(self, driver: AsyncDriverAdapterBase) -> "Optional[VersionInfo]":
148
+ """Get PostgreSQL database version information.
149
+
150
+ Args:
151
+ driver: Async database driver instance
152
+
153
+ Returns:
154
+ PostgreSQL version information or None if detection fails
155
+ """
156
+ version_str = await cast("PsycopgAsyncDriver", driver).select_value("SELECT version()")
157
+ if not version_str:
158
+ logger.warning("No PostgreSQL version information found")
159
+ return None
160
+
161
+ # Parse version like "PostgreSQL 15.3 on x86_64-pc-linux-gnu..."
162
+ version_match = POSTGRES_VERSION_PATTERN.search(str(version_str))
163
+ if not version_match:
164
+ logger.warning("Could not parse PostgreSQL version: %s", version_str)
165
+ return None
166
+
167
+ major = int(version_match.group(1))
168
+ minor = int(version_match.group(2))
169
+ patch = int(version_match.group(3)) if version_match.group(3) else 0
170
+
171
+ version_info = VersionInfo(major, minor, patch)
172
+ logger.debug("Detected PostgreSQL version: %s", version_info)
173
+ return version_info
174
+
175
+ async def get_feature_flag(self, driver: AsyncDriverAdapterBase, feature: str) -> bool:
176
+ """Check if PostgreSQL database supports a specific feature.
177
+
178
+ Args:
179
+ driver: Async database driver instance
180
+ feature: Feature name to check
181
+
182
+ Returns:
183
+ True if feature is supported, False otherwise
184
+ """
185
+ version_info = await self.get_version(driver)
186
+ if not version_info:
187
+ return False
188
+
189
+ feature_checks: dict[str, Callable[[VersionInfo], bool]] = {
190
+ "supports_json": lambda v: v >= VersionInfo(9, 2, 0),
191
+ "supports_jsonb": lambda v: v >= VersionInfo(9, 4, 0),
192
+ "supports_uuid": lambda _: True, # UUID extension widely available
193
+ "supports_arrays": lambda _: True, # PostgreSQL has excellent array support
194
+ "supports_returning": lambda v: v >= VersionInfo(8, 2, 0),
195
+ "supports_upsert": lambda v: v >= VersionInfo(9, 5, 0), # ON CONFLICT
196
+ "supports_window_functions": lambda v: v >= VersionInfo(8, 4, 0),
197
+ "supports_cte": lambda v: v >= VersionInfo(8, 4, 0),
198
+ "supports_transactions": lambda _: True,
199
+ "supports_prepared_statements": lambda _: True,
200
+ "supports_schemas": lambda _: True,
201
+ "supports_partitioning": lambda v: v >= VersionInfo(10, 0, 0),
202
+ }
203
+
204
+ if feature in feature_checks:
205
+ return bool(feature_checks[feature](version_info))
206
+
207
+ return False
208
+
209
+ async def get_optimal_type(self, driver: AsyncDriverAdapterBase, type_category: str) -> str:
210
+ """Get optimal PostgreSQL type for a category.
211
+
212
+ Args:
213
+ driver: Async database driver instance
214
+ type_category: Type category
215
+
216
+ Returns:
217
+ PostgreSQL-specific type name
218
+ """
219
+ version_info = await self.get_version(driver)
220
+
221
+ if type_category == "json":
222
+ if version_info and version_info >= VersionInfo(9, 4, 0):
223
+ return "JSONB" # Prefer JSONB over JSON
224
+ if version_info and version_info >= VersionInfo(9, 2, 0):
225
+ return "JSON"
226
+ return "TEXT"
227
+
228
+ type_map = {
229
+ "uuid": "UUID",
230
+ "boolean": "BOOLEAN",
231
+ "timestamp": "TIMESTAMP WITH TIME ZONE",
232
+ "text": "TEXT",
233
+ "blob": "BYTEA",
234
+ "array": "ARRAY",
235
+ }
236
+ return type_map.get(type_category, "TEXT")
237
+
238
+ def list_available_features(self) -> "list[str]":
239
+ """List available PostgreSQL feature flags.
240
+
241
+ Returns:
242
+ List of supported feature names
243
+ """
244
+ return [
245
+ "supports_json",
246
+ "supports_jsonb",
247
+ "supports_uuid",
248
+ "supports_arrays",
249
+ "supports_returning",
250
+ "supports_upsert",
251
+ "supports_window_functions",
252
+ "supports_cte",
253
+ "supports_transactions",
254
+ "supports_prepared_statements",
255
+ "supports_schemas",
256
+ "supports_partitioning",
257
+ ]