Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/sqlalchemy/dialects/postgresql/ext.py: 48%

1# dialects/postgresql/ext.py 

2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors 

3# <see AUTHORS file> 

4# 

5# This module is part of SQLAlchemy and is released under 

6# the MIT License: https://www.opensource.org/licenses/mit-license.php 

7# mypy: ignore-errors 

8from __future__ import annotations 

9 

10from typing import Any 

11from typing import Iterable 

12from typing import List 

13from typing import Optional 

14from typing import overload 

15from typing import Sequence 

16from typing import Tuple 

17from typing import TYPE_CHECKING 

18from typing import TypeVar 

19 

20from . import types 

21from .array import ARRAY 

22from ... import exc 

23from ...sql import coercions 

24from ...sql import elements 

25from ...sql import expression 

26from ...sql import functions 

27from ...sql import roles 

28from ...sql import schema 

29from ...sql.base import SyntaxExtension 

30from ...sql.schema import ColumnCollectionConstraint 

31from ...sql.sqltypes import TEXT 

32from ...sql.visitors import InternalTraversal 

33 

34if TYPE_CHECKING: 

35 from ...sql._typing import _ColumnExpressionArgument 

36 from ...sql._typing import _DDLColumnArgument 

37 from ...sql.elements import ClauseElement 

38 from ...sql.elements import ColumnElement 

39 from ...sql.operators import OperatorType 

40 from ...sql.selectable import FromClause 

41 from ...sql.visitors import _CloneCallableType 

42 from ...sql.visitors import _TraverseInternalsType 

43 

44_T = TypeVar("_T", bound=Any) 

45 

46 

47class aggregate_order_by(expression.ColumnElement[_T]): 

48 """Represent a PostgreSQL aggregate order by expression. 

49 

50 E.g.:: 

51 

52 from sqlalchemy.dialects.postgresql import aggregate_order_by 

53 

54 expr = func.array_agg(aggregate_order_by(table.c.a, table.c.b.desc())) 

55 stmt = select(expr) 

56 

57 would represent the expression: 

58 

59 .. sourcecode:: sql 

60 

61 SELECT array_agg(a ORDER BY b DESC) FROM table; 

62 

63 .. legacy:: An improved dialect-agnostic form of this function is now 

64 available in Core by calling the 

65 :meth:`_functions.Function.aggregate_order_by` method on any function 

66 defined by the backend as an aggregate function. 

67 

68 .. seealso:: 

69 

70 :func:`_sql.aggregate_order_by` - Core level function 

71 

72 :class:`_functions.array_agg` 

73 

74 """ 

75 

76 __visit_name__ = "aggregate_order_by" 

77 

78 stringify_dialect = "postgresql" 

79 _traverse_internals: _TraverseInternalsType = [ 

80 ("target", InternalTraversal.dp_clauseelement), 

81 ("type", InternalTraversal.dp_type), 

82 ("order_by", InternalTraversal.dp_clauseelement), 

83 ] 

84 

85 @overload 

86 def __init__( 

87 self, 

88 target: ColumnElement[_T], 

89 *order_by: _ColumnExpressionArgument[Any], 

90 ): ... 

91 

92 @overload 

93 def __init__( 

94 self, 

95 target: _ColumnExpressionArgument[_T], 

96 *order_by: _ColumnExpressionArgument[Any], 

97 ): ... 

98 

99 def __init__( 

100 self, 

101 target: _ColumnExpressionArgument[_T], 

102 *order_by: _ColumnExpressionArgument[Any], 

103 ): 

104 self.target: ClauseElement = coercions.expect( 

105 roles.ExpressionElementRole, target 

106 ) 

107 self.type = self.target.type 

108 

109 _lob = len(order_by) 

110 self.order_by: ClauseElement 

111 if _lob == 0: 

112 raise TypeError("at least one ORDER BY element is required") 

113 elif _lob == 1: 

114 self.order_by = coercions.expect( 

115 roles.ExpressionElementRole, order_by[0] 

116 ) 

117 else: 

118 self.order_by = elements.ClauseList( 

119 *order_by, _literal_as_text_role=roles.ExpressionElementRole 

120 ) 

121 

122 def self_group( 

123 self, against: Optional[OperatorType] = None 

124 ) -> ClauseElement: 

125 return self 

126 

127 def get_children(self, **kwargs: Any) -> Iterable[ClauseElement]: 

128 return self.target, self.order_by 

129 

130 def _copy_internals( 

131 self, clone: _CloneCallableType = elements._clone, **kw: Any 

132 ) -> None: 

133 self.target = clone(self.target, **kw) 

134 self.order_by = clone(self.order_by, **kw) 

135 

136 @property 

137 def _from_objects(self) -> List[FromClause]: 

138 return self.target._from_objects + self.order_by._from_objects 

139 

140 

141class ExcludeConstraint(ColumnCollectionConstraint): 

142 """A table-level EXCLUDE constraint. 

143 

144 Defines an EXCLUDE constraint as described in the `PostgreSQL 

145 documentation`__. 

146 

147 __ https://www.postgresql.org/docs/current/static/sql-createtable.html#SQL-CREATETABLE-EXCLUDE 

148 

149 """ # noqa 

150 

151 __visit_name__ = "exclude_constraint" 

152 

153 where = None 

154 inherit_cache = False 

155 

156 create_drop_stringify_dialect = "postgresql" 

157 

158 @elements._document_text_coercion( 

159 "where", 

160 ":class:`.ExcludeConstraint`", 

161 ":paramref:`.ExcludeConstraint.where`", 

162 ) 

163 def __init__( 

164 self, *elements: Tuple[_DDLColumnArgument, str], **kw: Any 

165 ) -> None: 

166 r""" 

167 Create an :class:`.ExcludeConstraint` object. 

168 

169 E.g.:: 

170 

171 const = ExcludeConstraint( 

172 (Column("period"), "&&"), 

173 (Column("group"), "="), 

174 where=(Column("group") != "some group"), 

175 ops={"group": "my_operator_class"}, 

176 ) 

177 

178 The constraint is normally embedded into the :class:`_schema.Table` 

179 construct 

180 directly, or added later using :meth:`.append_constraint`:: 

181 

182 some_table = Table( 

183 "some_table", 

184 metadata, 

185 Column("id", Integer, primary_key=True), 

186 Column("period", TSRANGE()), 

187 Column("group", String), 

188 ) 

189 

190 some_table.append_constraint( 

191 ExcludeConstraint( 

192 (some_table.c.period, "&&"), 

193 (some_table.c.group, "="), 

194 where=some_table.c.group != "some group", 

195 name="some_table_excl_const", 

196 ops={"group": "my_operator_class"}, 

197 ) 

198 ) 

199 

200 The exclude constraint defined in this example requires the 

201 ``btree_gist`` extension, that can be created using the 

202 command ``CREATE EXTENSION btree_gist;``. 

203 

204 :param \*elements: 

205 

206 A sequence of two tuples of the form ``(column, operator)`` where 

207 "column" is either a :class:`_schema.Column` object, or a SQL 

208 expression element (e.g. ``func.int8range(table.from, table.to)``) 

209 or the name of a column as string, and "operator" is a string 

210 containing the operator to use (e.g. `"&&"` or `"="`). 

211 

212 In order to specify a column name when a :class:`_schema.Column` 

213 object is not available, while ensuring 

214 that any necessary quoting rules take effect, an ad-hoc 

215 :class:`_schema.Column` or :func:`_expression.column` 

216 object should be used. 

217 The ``column`` may also be a string SQL expression when 

218 passed as :func:`_expression.literal_column` or 

219 :func:`_expression.text` 

220 

221 :param name: 

222 Optional, the in-database name of this constraint. 

223 

224 :param deferrable: 

225 Optional bool. If set, emit DEFERRABLE or NOT DEFERRABLE when 

226 issuing DDL for this constraint. 

227 

228 :param initially: 

229 Optional string. If set, emit INITIALLY <value> when issuing DDL 

230 for this constraint. 

231 

232 :param info: Optional data dictionary which will be populated into the 

233 :attr:`.SchemaItem.info` attribute of this object. 

234 

235 .. versionadded:: 2.0.50 

236 

237 :param using: 

238 Optional string. If set, emit USING <index_method> when issuing DDL 

239 for this constraint. Defaults to 'gist'. 

240 

241 :param where: 

242 Optional SQL expression construct or literal SQL string. 

243 If set, emit WHERE <predicate> when issuing DDL 

244 for this constraint. 

245 

246 :param ops: 

247 Optional dictionary. Used to define operator classes for the 

248 elements; works the same way as that of the 

249 :ref:`postgresql_ops <postgresql_operator_classes>` 

250 parameter specified to the :class:`_schema.Index` construct. 

251 

252 .. seealso:: 

253 

254 :ref:`postgresql_operator_classes` - general description of how 

255 PostgreSQL operator classes are specified. 

256 

257 """ 

258 columns = [] 

259 render_exprs = [] 

260 self.operators = {} 

261 

262 expressions, operators = zip(*elements) 

263 

264 for (expr, column, strname, add_element), operator in zip( 

265 coercions.expect_col_expression_collection( 

266 roles.DDLConstraintColumnRole, expressions 

267 ), 

268 operators, 

269 ): 

270 if add_element is not None: 

271 columns.append(add_element) 

272 

273 name = column.name if column is not None else strname 

274 

275 if name is not None: 

276 # backwards compat 

277 self.operators[name] = operator 

278 

279 render_exprs.append((expr, name, operator)) 

280 

281 self._render_exprs = render_exprs 

282 

283 ColumnCollectionConstraint.__init__( 

284 self, 

285 *columns, 

286 name=kw.get("name"), 

287 deferrable=kw.get("deferrable"), 

288 initially=kw.get("initially"), 

289 info=kw.get("info"), 

290 ) 

291 self.using = kw.get("using", "gist") 

292 where = kw.get("where") 

293 if where is not None: 

294 self.where = coercions.expect(roles.StatementOptionRole, where) 

295 

296 self.ops = kw.get("ops", {}) 

297 

298 def _set_parent(self, table, **kw): 

299 super()._set_parent(table) 

300 

301 self._render_exprs = [ 

302 ( 

303 expr if not isinstance(expr, str) else table.c[expr], 

304 name, 

305 operator, 

306 ) 

307 for expr, name, operator in (self._render_exprs) 

308 ] 

309 

310 def _copy(self, target_table=None, **kw): 

311 elements = [ 

312 ( 

313 schema._copy_expression(expr, self.parent, target_table), 

314 operator, 

315 ) 

316 for expr, _, operator in self._render_exprs 

317 ] 

318 c = self.__class__( 

319 *elements, 

320 name=self.name, 

321 deferrable=self.deferrable, 

322 initially=self.initially, 

323 where=self.where, 

324 using=self.using, 

325 ) 

326 c.dispatch._update(self.dispatch) 

327 return c 

328 

329 

330def array_agg(*arg, **kw): 

331 """PostgreSQL-specific form of :class:`_functions.array_agg`, ensures 

332 return type is :class:`_postgresql.ARRAY` and not 

333 the plain :class:`_types.ARRAY`, unless an explicit ``type_`` 

334 is passed. 

335 

336 """ 

337 kw["_default_array_type"] = ARRAY 

338 return functions.func.array_agg(*arg, **kw) 

339 

340 

341class _regconfig_fn(functions.GenericFunction[_T]): 

342 inherit_cache = True 

343 

344 def __init__(self, *args, **kwargs): 

345 args = list(args) 

346 if len(args) > 1: 

347 initial_arg = coercions.expect( 

348 roles.ExpressionElementRole, 

349 args.pop(0), 

350 name=getattr(self, "name", None), 

351 apply_propagate_attrs=self, 

352 type_=types.REGCONFIG, 

353 ) 

354 initial_arg = [initial_arg] 

355 else: 

356 initial_arg = [] 

357 

358 addtl_args = [ 

359 coercions.expect( 

360 roles.ExpressionElementRole, 

361 c, 

362 name=getattr(self, "name", None), 

363 apply_propagate_attrs=self, 

364 ) 

365 for c in args 

366 ] 

367 super().__init__(*(initial_arg + addtl_args), **kwargs) 

368 

369 

370class to_tsvector(_regconfig_fn): 

371 """The PostgreSQL ``to_tsvector`` SQL function. 

372 

373 This function applies automatic casting of the REGCONFIG argument 

374 to use the :class:`_postgresql.REGCONFIG` datatype automatically, 

375 and applies a return type of :class:`_postgresql.TSVECTOR`. 

376 

377 Assuming the PostgreSQL dialect has been imported, either by invoking 

378 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL 

379 engine using ``create_engine("postgresql...")``, 

380 :class:`_postgresql.to_tsvector` will be used automatically when invoking 

381 ``sqlalchemy.func.to_tsvector()``, ensuring the correct argument and return 

382 type handlers are used at compile and execution time. 

383 

384 .. versionadded:: 2.0.0rc1 

385 

386 """ 

387 

388 inherit_cache = True 

389 type = types.TSVECTOR 

390 

391 

392class to_tsquery(_regconfig_fn): 

393 """The PostgreSQL ``to_tsquery`` SQL function. 

394 

395 This function applies automatic casting of the REGCONFIG argument 

396 to use the :class:`_postgresql.REGCONFIG` datatype automatically, 

397 and applies a return type of :class:`_postgresql.TSQUERY`. 

398 

399 Assuming the PostgreSQL dialect has been imported, either by invoking 

400 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL 

401 engine using ``create_engine("postgresql...")``, 

402 :class:`_postgresql.to_tsquery` will be used automatically when invoking 

403 ``sqlalchemy.func.to_tsquery()``, ensuring the correct argument and return 

404 type handlers are used at compile and execution time. 

405 

406 .. versionadded:: 2.0.0rc1 

407 

408 """ 

409 

410 inherit_cache = True 

411 type = types.TSQUERY 

412 

413 

414class plainto_tsquery(_regconfig_fn): 

415 """The PostgreSQL ``plainto_tsquery`` SQL function. 

416 

417 This function applies automatic casting of the REGCONFIG argument 

418 to use the :class:`_postgresql.REGCONFIG` datatype automatically, 

419 and applies a return type of :class:`_postgresql.TSQUERY`. 

420 

421 Assuming the PostgreSQL dialect has been imported, either by invoking 

422 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL 

423 engine using ``create_engine("postgresql...")``, 

424 :class:`_postgresql.plainto_tsquery` will be used automatically when 

425 invoking ``sqlalchemy.func.plainto_tsquery()``, ensuring the correct 

426 argument and return type handlers are used at compile and execution time. 

427 

428 .. versionadded:: 2.0.0rc1 

429 

430 """ 

431 

432 inherit_cache = True 

433 type = types.TSQUERY 

434 

435 

436class phraseto_tsquery(_regconfig_fn): 

437 """The PostgreSQL ``phraseto_tsquery`` SQL function. 

438 

439 This function applies automatic casting of the REGCONFIG argument 

440 to use the :class:`_postgresql.REGCONFIG` datatype automatically, 

441 and applies a return type of :class:`_postgresql.TSQUERY`. 

442 

443 Assuming the PostgreSQL dialect has been imported, either by invoking 

444 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL 

445 engine using ``create_engine("postgresql...")``, 

446 :class:`_postgresql.phraseto_tsquery` will be used automatically when 

447 invoking ``sqlalchemy.func.phraseto_tsquery()``, ensuring the correct 

448 argument and return type handlers are used at compile and execution time. 

449 

450 .. versionadded:: 2.0.0rc1 

451 

452 """ 

453 

454 inherit_cache = True 

455 type = types.TSQUERY 

456 

457 

458class websearch_to_tsquery(_regconfig_fn): 

459 """The PostgreSQL ``websearch_to_tsquery`` SQL function. 

460 

461 This function applies automatic casting of the REGCONFIG argument 

462 to use the :class:`_postgresql.REGCONFIG` datatype automatically, 

463 and applies a return type of :class:`_postgresql.TSQUERY`. 

464 

465 Assuming the PostgreSQL dialect has been imported, either by invoking 

466 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL 

467 engine using ``create_engine("postgresql...")``, 

468 :class:`_postgresql.websearch_to_tsquery` will be used automatically when 

469 invoking ``sqlalchemy.func.websearch_to_tsquery()``, ensuring the correct 

470 argument and return type handlers are used at compile and execution time. 

471 

472 .. versionadded:: 2.0.0rc1 

473 

474 """ 

475 

476 inherit_cache = True 

477 type = types.TSQUERY 

478 

479 

480class ts_headline(_regconfig_fn): 

481 """The PostgreSQL ``ts_headline`` SQL function. 

482 

483 This function applies automatic casting of the REGCONFIG argument 

484 to use the :class:`_postgresql.REGCONFIG` datatype automatically, 

485 and applies a return type of :class:`_types.TEXT`. 

486 

487 Assuming the PostgreSQL dialect has been imported, either by invoking 

488 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL 

489 engine using ``create_engine("postgresql...")``, 

490 :class:`_postgresql.ts_headline` will be used automatically when invoking 

491 ``sqlalchemy.func.ts_headline()``, ensuring the correct argument and return 

492 type handlers are used at compile and execution time. 

493 

494 .. versionadded:: 2.0.0rc1 

495 

496 """ 

497 

498 inherit_cache = True 

499 type = TEXT 

500 

501 def __init__(self, *args, **kwargs): 

502 args = list(args) 

503 

504 # parse types according to 

505 # https://www.postgresql.org/docs/current/textsearch-controls.html#TEXTSEARCH-HEADLINE 

506 if len(args) < 2: 

507 # invalid args; don't do anything 

508 has_regconfig = False 

509 elif ( 

510 isinstance(args[1], elements.ColumnElement) 

511 and args[1].type._type_affinity is types.TSQUERY 

512 ): 

513 # tsquery is second argument, no regconfig argument 

514 has_regconfig = False 

515 else: 

516 has_regconfig = True 

517 

518 if has_regconfig: 

519 initial_arg = coercions.expect( 

520 roles.ExpressionElementRole, 

521 args.pop(0), 

522 apply_propagate_attrs=self, 

523 name=getattr(self, "name", None), 

524 type_=types.REGCONFIG, 

525 ) 

526 initial_arg = [initial_arg] 

527 else: 

528 initial_arg = [] 

529 

530 addtl_args = [ 

531 coercions.expect( 

532 roles.ExpressionElementRole, 

533 c, 

534 name=getattr(self, "name", None), 

535 apply_propagate_attrs=self, 

536 ) 

537 for c in args 

538 ] 

539 super().__init__(*(initial_arg + addtl_args), **kwargs) 

540 

541 

542def distinct_on(*expr: _ColumnExpressionArgument[Any]) -> DistinctOnClause: 

543 """apply a DISTINCT_ON to a SELECT statement 

544 

545 e.g.:: 

546 

547 stmt = select(tbl).ext(distinct_on(t.c.some_col)) 

548 

549 this supersedes the previous approach of using 

550 ``select(tbl).distinct(t.c.some_col))`` to apply a similar construct. 

551 

552 .. versionadded:: 2.1 

553 

554 """ 

555 return DistinctOnClause(expr) 

556 

557 

558class DistinctOnClause(SyntaxExtension, expression.ClauseElement): 

559 stringify_dialect = "postgresql" 

560 __visit_name__ = "postgresql_distinct_on" 

561 

562 _traverse_internals: _TraverseInternalsType = [ 

563 ("_distinct_on", InternalTraversal.dp_clauseelement_tuple), 

564 ] 

565 

566 def __init__(self, distinct_on: Sequence[_ColumnExpressionArgument[Any]]): 

567 self._distinct_on = tuple( 

568 coercions.expect(roles.ByOfRole, e, apply_propagate_attrs=self) 

569 for e in distinct_on 

570 ) 

571 

572 def apply_to_select(self, select_stmt: expression.Select[Any]) -> None: 

573 if select_stmt._distinct_on: 

574 raise exc.InvalidRequestError( 

575 "Cannot mix ``select.ext(distinct_on(...))`` and " 

576 "``select.distinct(...)``" 

577 ) 

578 # mark this select as a distinct 

579 select_stmt.distinct.non_generative(select_stmt) 

580 

581 select_stmt.apply_syntax_extension_point( 

582 self._merge_other_distinct, "pre_columns" 

583 ) 

584 

585 def _merge_other_distinct( 

586 self, existing: Sequence[elements.ClauseElement] 

587 ) -> Sequence[elements.ClauseElement]: 

588 res = [] 

589 to_merge = () 

590 for e in existing: 

591 if isinstance(e, DistinctOnClause): 

592 to_merge += e._distinct_on 

593 else: 

594 res.append(e) 

595 if to_merge: 

596 res.append(DistinctOnClause(to_merge + self._distinct_on)) 

597 else: 

598 res.append(self) 

599 return res