postgresql.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848
  1. # mypy: allow-untyped-defs, allow-incomplete-defs, allow-untyped-calls
  2. # mypy: no-warn-return-any, allow-any-generics
  3. from __future__ import annotations
  4. import logging
  5. import re
  6. from typing import Any
  7. from typing import cast
  8. from typing import Dict
  9. from typing import List
  10. from typing import Optional
  11. from typing import Sequence
  12. from typing import Tuple
  13. from typing import TYPE_CHECKING
  14. from typing import Union
  15. from sqlalchemy import Column
  16. from sqlalchemy import literal_column
  17. from sqlalchemy import Numeric
  18. from sqlalchemy import text
  19. from sqlalchemy import types as sqltypes
  20. from sqlalchemy.dialects.postgresql import BIGINT
  21. from sqlalchemy.dialects.postgresql import ExcludeConstraint
  22. from sqlalchemy.dialects.postgresql import INTEGER
  23. from sqlalchemy.schema import CreateIndex
  24. from sqlalchemy.sql.elements import ColumnClause
  25. from sqlalchemy.sql.elements import TextClause
  26. from sqlalchemy.sql.functions import FunctionElement
  27. from sqlalchemy.types import NULLTYPE
  28. from .base import alter_column
  29. from .base import alter_table
  30. from .base import AlterColumn
  31. from .base import ColumnComment
  32. from .base import format_column_name
  33. from .base import format_table_name
  34. from .base import format_type
  35. from .base import IdentityColumnDefault
  36. from .base import RenameTable
  37. from .impl import ComparisonResult
  38. from .impl import DefaultImpl
  39. from .. import util
  40. from ..autogenerate import render
  41. from ..operations import ops
  42. from ..operations import schemaobj
  43. from ..operations.base import BatchOperations
  44. from ..operations.base import Operations
  45. from ..util import sqla_compat
  46. from ..util.sqla_compat import compiles
  47. if TYPE_CHECKING:
  48. from typing import Literal
  49. from sqlalchemy import Index
  50. from sqlalchemy import UniqueConstraint
  51. from sqlalchemy.dialects.postgresql.array import ARRAY
  52. from sqlalchemy.dialects.postgresql.base import PGDDLCompiler
  53. from sqlalchemy.dialects.postgresql.hstore import HSTORE
  54. from sqlalchemy.dialects.postgresql.json import JSON
  55. from sqlalchemy.dialects.postgresql.json import JSONB
  56. from sqlalchemy.sql.elements import ClauseElement
  57. from sqlalchemy.sql.elements import ColumnElement
  58. from sqlalchemy.sql.elements import quoted_name
  59. from sqlalchemy.sql.schema import MetaData
  60. from sqlalchemy.sql.schema import Table
  61. from sqlalchemy.sql.type_api import TypeEngine
  62. from .base import _ServerDefault
  63. from ..autogenerate.api import AutogenContext
  64. from ..autogenerate.render import _f_name
  65. from ..runtime.migration import MigrationContext
  66. log = logging.getLogger(__name__)
  67. class PostgresqlImpl(DefaultImpl):
  68. __dialect__ = "postgresql"
  69. transactional_ddl = True
  70. type_synonyms = DefaultImpl.type_synonyms + (
  71. {"FLOAT", "DOUBLE PRECISION"},
  72. )
  73. def create_index(self, index: Index, **kw: Any) -> None:
  74. # this likely defaults to None if not present, so get()
  75. # should normally not return the default value. being
  76. # defensive in any case
  77. postgresql_include = index.kwargs.get("postgresql_include", None) or ()
  78. for col in postgresql_include:
  79. if col not in index.table.c: # type: ignore[union-attr]
  80. index.table.append_column( # type: ignore[union-attr]
  81. Column(col, sqltypes.NullType)
  82. )
  83. self._exec(CreateIndex(index, **kw))
  84. def prep_table_for_batch(self, batch_impl, table):
  85. for constraint in table.constraints:
  86. if (
  87. constraint.name is not None
  88. and constraint.name in batch_impl.named_constraints
  89. ):
  90. self.drop_constraint(constraint)
  91. def compare_server_default(
  92. self,
  93. inspector_column,
  94. metadata_column,
  95. rendered_metadata_default,
  96. rendered_inspector_default,
  97. ):
  98. # don't do defaults for SERIAL columns
  99. if (
  100. metadata_column.primary_key
  101. and metadata_column is metadata_column.table._autoincrement_column
  102. ):
  103. return False
  104. conn_col_default = rendered_inspector_default
  105. defaults_equal = conn_col_default == rendered_metadata_default
  106. if defaults_equal:
  107. return False
  108. if None in (
  109. conn_col_default,
  110. rendered_metadata_default,
  111. metadata_column.server_default,
  112. ):
  113. return not defaults_equal
  114. metadata_default = metadata_column.server_default.arg
  115. if isinstance(metadata_default, str):
  116. if not isinstance(inspector_column.type, Numeric):
  117. metadata_default = re.sub(r"^'|'$", "", metadata_default)
  118. metadata_default = f"'{metadata_default}'"
  119. metadata_default = literal_column(metadata_default)
  120. # run a real compare against the server
  121. conn = self.connection
  122. assert conn is not None
  123. return not conn.scalar(
  124. sqla_compat._select(
  125. literal_column(conn_col_default) == metadata_default
  126. )
  127. )
  128. def alter_column( # type:ignore[override]
  129. self,
  130. table_name: str,
  131. column_name: str,
  132. nullable: Optional[bool] = None,
  133. server_default: Union[_ServerDefault, Literal[False]] = False,
  134. name: Optional[str] = None,
  135. type_: Optional[TypeEngine] = None,
  136. schema: Optional[str] = None,
  137. autoincrement: Optional[bool] = None,
  138. existing_type: Optional[TypeEngine] = None,
  139. existing_server_default: Optional[_ServerDefault] = None,
  140. existing_nullable: Optional[bool] = None,
  141. existing_autoincrement: Optional[bool] = None,
  142. **kw: Any,
  143. ) -> None:
  144. using = kw.pop("postgresql_using", None)
  145. if using is not None and type_ is None:
  146. raise util.CommandError(
  147. "postgresql_using must be used with the type_ parameter"
  148. )
  149. if type_ is not None:
  150. self._exec(
  151. PostgresqlColumnType(
  152. table_name,
  153. column_name,
  154. type_,
  155. schema=schema,
  156. using=using,
  157. existing_type=existing_type,
  158. existing_server_default=existing_server_default,
  159. existing_nullable=existing_nullable,
  160. )
  161. )
  162. super().alter_column(
  163. table_name,
  164. column_name,
  165. nullable=nullable,
  166. server_default=server_default,
  167. name=name,
  168. schema=schema,
  169. autoincrement=autoincrement,
  170. existing_type=existing_type,
  171. existing_server_default=existing_server_default,
  172. existing_nullable=existing_nullable,
  173. existing_autoincrement=existing_autoincrement,
  174. **kw,
  175. )
  176. def autogen_column_reflect(self, inspector, table, column_info):
  177. if column_info.get("default") and isinstance(
  178. column_info["type"], (INTEGER, BIGINT)
  179. ):
  180. seq_match = re.match(
  181. r"nextval\('(.+?)'::regclass\)", column_info["default"]
  182. )
  183. if seq_match:
  184. info = sqla_compat._exec_on_inspector(
  185. inspector,
  186. text(
  187. "select c.relname, a.attname "
  188. "from pg_class as c join "
  189. "pg_depend d on d.objid=c.oid and "
  190. "d.classid='pg_class'::regclass and "
  191. "d.refclassid='pg_class'::regclass "
  192. "join pg_class t on t.oid=d.refobjid "
  193. "join pg_attribute a on a.attrelid=t.oid and "
  194. "a.attnum=d.refobjsubid "
  195. "where c.relkind='S' and c.relname=:seqname"
  196. ),
  197. seqname=seq_match.group(1),
  198. ).first()
  199. if info:
  200. seqname, colname = info
  201. if colname == column_info["name"]:
  202. log.info(
  203. "Detected sequence named '%s' as "
  204. "owned by integer column '%s(%s)', "
  205. "assuming SERIAL and omitting",
  206. seqname,
  207. table.name,
  208. colname,
  209. )
  210. # sequence, and the owner is this column,
  211. # its a SERIAL - whack it!
  212. del column_info["default"]
  213. def correct_for_autogen_constraints(
  214. self,
  215. conn_unique_constraints,
  216. conn_indexes,
  217. metadata_unique_constraints,
  218. metadata_indexes,
  219. ):
  220. doubled_constraints = {
  221. index
  222. for index in conn_indexes
  223. if index.info.get("duplicates_constraint")
  224. }
  225. for ix in doubled_constraints:
  226. conn_indexes.remove(ix)
  227. if not sqla_compat.sqla_2:
  228. self._skip_functional_indexes(metadata_indexes, conn_indexes)
  229. # pg behavior regarding modifiers
  230. # | # | compiled sql | returned sql | regexp. group is removed |
  231. # | - | ---------------- | -----------------| ------------------------ |
  232. # | 1 | nulls first | nulls first | - |
  233. # | 2 | nulls last | | (?<! desc)( nulls last)$ |
  234. # | 3 | asc | | ( asc)$ |
  235. # | 4 | asc nulls first | nulls first | ( asc) nulls first$ |
  236. # | 5 | asc nulls last | | ( asc nulls last)$ |
  237. # | 6 | desc | desc | - |
  238. # | 7 | desc nulls first | desc | desc( nulls first)$ |
  239. # | 8 | desc nulls last | desc nulls last | - |
  240. _default_modifiers_re = ( # order of case 2 and 5 matters
  241. re.compile("( asc nulls last)$"), # case 5
  242. re.compile("(?<! desc)( nulls last)$"), # case 2
  243. re.compile("( asc)$"), # case 3
  244. re.compile("( asc) nulls first$"), # case 4
  245. re.compile(" desc( nulls first)$"), # case 7
  246. )
  247. def _cleanup_index_expr(self, index: Index, expr: str) -> str:
  248. expr = expr.lower().replace('"', "").replace("'", "")
  249. if index.table is not None:
  250. # should not be needed, since include_table=False is in compile
  251. expr = expr.replace(f"{index.table.name.lower()}.", "")
  252. if "::" in expr:
  253. # strip :: cast. types can have spaces in them
  254. expr = re.sub(r"(::[\w ]+\w)", "", expr)
  255. while expr and expr[0] == "(" and expr[-1] == ")":
  256. expr = expr[1:-1]
  257. # NOTE: when parsing the connection expression this cleanup could
  258. # be skipped
  259. for rs in self._default_modifiers_re:
  260. if match := rs.search(expr):
  261. start, end = match.span(1)
  262. expr = expr[:start] + expr[end:]
  263. break
  264. while expr and expr[0] == "(" and expr[-1] == ")":
  265. expr = expr[1:-1]
  266. # strip casts
  267. cast_re = re.compile(r"cast\s*\(")
  268. if cast_re.match(expr):
  269. expr = cast_re.sub("", expr)
  270. # remove the as type
  271. expr = re.sub(r"as\s+[^)]+\)", "", expr)
  272. # remove spaces
  273. expr = expr.replace(" ", "")
  274. return expr
  275. def _dialect_options(
  276. self, item: Union[Index, UniqueConstraint]
  277. ) -> Tuple[Any, ...]:
  278. # only the positive case is returned by sqlalchemy reflection so
  279. # None and False are threated the same
  280. if item.dialect_kwargs.get("postgresql_nulls_not_distinct"):
  281. return ("nulls_not_distinct",)
  282. return ()
  283. def compare_indexes(
  284. self,
  285. metadata_index: Index,
  286. reflected_index: Index,
  287. ) -> ComparisonResult:
  288. msg = []
  289. unique_msg = self._compare_index_unique(
  290. metadata_index, reflected_index
  291. )
  292. if unique_msg:
  293. msg.append(unique_msg)
  294. m_exprs = metadata_index.expressions
  295. r_exprs = reflected_index.expressions
  296. if len(m_exprs) != len(r_exprs):
  297. msg.append(f"expression number {len(r_exprs)} to {len(m_exprs)}")
  298. if msg:
  299. # no point going further, return early
  300. return ComparisonResult.Different(msg)
  301. skip = []
  302. for pos, (m_e, r_e) in enumerate(zip(m_exprs, r_exprs), 1):
  303. m_compile = self._compile_element(m_e)
  304. m_text = self._cleanup_index_expr(metadata_index, m_compile)
  305. # print(f"META ORIG: {m_compile!r} CLEANUP: {m_text!r}")
  306. r_compile = self._compile_element(r_e)
  307. r_text = self._cleanup_index_expr(metadata_index, r_compile)
  308. # print(f"CONN ORIG: {r_compile!r} CLEANUP: {r_text!r}")
  309. if m_text == r_text:
  310. continue # expressions these are equal
  311. elif m_compile.strip().endswith("_ops") and (
  312. " " in m_compile or ")" in m_compile # is an expression
  313. ):
  314. skip.append(
  315. f"expression #{pos} {m_compile!r} detected "
  316. "as including operator clause."
  317. )
  318. util.warn(
  319. f"Expression #{pos} {m_compile!r} in index "
  320. f"{reflected_index.name!r} detected to include "
  321. "an operator clause. Expression compare cannot proceed. "
  322. "Please move the operator clause to the "
  323. "``postgresql_ops`` dict to enable proper compare "
  324. "of the index expressions: "
  325. "https://docs.sqlalchemy.org/en/latest/dialects/postgresql.html#operator-classes", # noqa: E501
  326. )
  327. else:
  328. msg.append(f"expression #{pos} {r_compile!r} to {m_compile!r}")
  329. m_options = self._dialect_options(metadata_index)
  330. r_options = self._dialect_options(reflected_index)
  331. if m_options != r_options:
  332. msg.extend(f"options {r_options} to {m_options}")
  333. if msg:
  334. return ComparisonResult.Different(msg)
  335. elif skip:
  336. # if there are other changes detected don't skip the index
  337. return ComparisonResult.Skip(skip)
  338. else:
  339. return ComparisonResult.Equal()
  340. def compare_unique_constraint(
  341. self,
  342. metadata_constraint: UniqueConstraint,
  343. reflected_constraint: UniqueConstraint,
  344. ) -> ComparisonResult:
  345. metadata_tup = self._create_metadata_constraint_sig(
  346. metadata_constraint
  347. )
  348. reflected_tup = self._create_reflected_constraint_sig(
  349. reflected_constraint
  350. )
  351. meta_sig = metadata_tup.unnamed
  352. conn_sig = reflected_tup.unnamed
  353. if conn_sig != meta_sig:
  354. return ComparisonResult.Different(
  355. f"expression {conn_sig} to {meta_sig}"
  356. )
  357. metadata_do = self._dialect_options(metadata_tup.const)
  358. conn_do = self._dialect_options(reflected_tup.const)
  359. if metadata_do != conn_do:
  360. return ComparisonResult.Different(
  361. f"expression {conn_do} to {metadata_do}"
  362. )
  363. return ComparisonResult.Equal()
  364. def adjust_reflected_dialect_options(
  365. self, reflected_options: Dict[str, Any], kind: str
  366. ) -> Dict[str, Any]:
  367. options: Dict[str, Any]
  368. options = reflected_options.get("dialect_options", {}).copy()
  369. if not options.get("postgresql_include"):
  370. options.pop("postgresql_include", None)
  371. return options
  372. def _compile_element(self, element: Union[ClauseElement, str]) -> str:
  373. if isinstance(element, str):
  374. return element
  375. return element.compile(
  376. dialect=self.dialect,
  377. compile_kwargs={"literal_binds": True, "include_table": False},
  378. ).string
  379. def render_ddl_sql_expr(
  380. self,
  381. expr: ClauseElement,
  382. is_server_default: bool = False,
  383. is_index: bool = False,
  384. **kw: Any,
  385. ) -> str:
  386. """Render a SQL expression that is typically a server default,
  387. index expression, etc.
  388. """
  389. # apply self_group to index expressions;
  390. # see https://github.com/sqlalchemy/sqlalchemy/blob/
  391. # 82fa95cfce070fab401d020c6e6e4a6a96cc2578/
  392. # lib/sqlalchemy/dialects/postgresql/base.py#L2261
  393. if is_index and not isinstance(expr, ColumnClause):
  394. expr = expr.self_group()
  395. return super().render_ddl_sql_expr(
  396. expr, is_server_default=is_server_default, is_index=is_index, **kw
  397. )
  398. def render_type(
  399. self, type_: TypeEngine, autogen_context: AutogenContext
  400. ) -> Union[str, Literal[False]]:
  401. mod = type(type_).__module__
  402. if not mod.startswith("sqlalchemy.dialects.postgresql"):
  403. return False
  404. if hasattr(self, "_render_%s_type" % type_.__visit_name__):
  405. meth = getattr(self, "_render_%s_type" % type_.__visit_name__)
  406. return meth(type_, autogen_context)
  407. return False
  408. def _render_HSTORE_type(
  409. self, type_: HSTORE, autogen_context: AutogenContext
  410. ) -> str:
  411. return cast(
  412. str,
  413. render._render_type_w_subtype(
  414. type_, autogen_context, "text_type", r"(.+?\(.*text_type=)"
  415. ),
  416. )
  417. def _render_ARRAY_type(
  418. self, type_: ARRAY, autogen_context: AutogenContext
  419. ) -> str:
  420. return cast(
  421. str,
  422. render._render_type_w_subtype(
  423. type_, autogen_context, "item_type", r"(.+?\()"
  424. ),
  425. )
  426. def _render_JSON_type(
  427. self, type_: JSON, autogen_context: AutogenContext
  428. ) -> str:
  429. return cast(
  430. str,
  431. render._render_type_w_subtype(
  432. type_, autogen_context, "astext_type", r"(.+?\(.*astext_type=)"
  433. ),
  434. )
  435. def _render_JSONB_type(
  436. self, type_: JSONB, autogen_context: AutogenContext
  437. ) -> str:
  438. return cast(
  439. str,
  440. render._render_type_w_subtype(
  441. type_, autogen_context, "astext_type", r"(.+?\(.*astext_type=)"
  442. ),
  443. )
  444. class PostgresqlColumnType(AlterColumn):
  445. def __init__(
  446. self, name: str, column_name: str, type_: TypeEngine, **kw
  447. ) -> None:
  448. using = kw.pop("using", None)
  449. super().__init__(name, column_name, **kw)
  450. self.type_ = sqltypes.to_instance(type_)
  451. self.using = using
  452. @compiles(RenameTable, "postgresql")
  453. def visit_rename_table(
  454. element: RenameTable, compiler: PGDDLCompiler, **kw
  455. ) -> str:
  456. return "%s RENAME TO %s" % (
  457. alter_table(compiler, element.table_name, element.schema),
  458. format_table_name(compiler, element.new_table_name, None),
  459. )
  460. @compiles(PostgresqlColumnType, "postgresql")
  461. def visit_column_type(
  462. element: PostgresqlColumnType, compiler: PGDDLCompiler, **kw
  463. ) -> str:
  464. return "%s %s %s %s" % (
  465. alter_table(compiler, element.table_name, element.schema),
  466. alter_column(compiler, element.column_name),
  467. "TYPE %s" % format_type(compiler, element.type_),
  468. "USING %s" % element.using if element.using else "",
  469. )
  470. @compiles(ColumnComment, "postgresql")
  471. def visit_column_comment(
  472. element: ColumnComment, compiler: PGDDLCompiler, **kw
  473. ) -> str:
  474. ddl = "COMMENT ON COLUMN {table_name}.{column_name} IS {comment}"
  475. comment = (
  476. compiler.sql_compiler.render_literal_value(
  477. element.comment, sqltypes.String()
  478. )
  479. if element.comment is not None
  480. else "NULL"
  481. )
  482. return ddl.format(
  483. table_name=format_table_name(
  484. compiler, element.table_name, element.schema
  485. ),
  486. column_name=format_column_name(compiler, element.column_name),
  487. comment=comment,
  488. )
  489. @compiles(IdentityColumnDefault, "postgresql")
  490. def visit_identity_column(
  491. element: IdentityColumnDefault, compiler: PGDDLCompiler, **kw
  492. ):
  493. text = "%s %s " % (
  494. alter_table(compiler, element.table_name, element.schema),
  495. alter_column(compiler, element.column_name),
  496. )
  497. if element.default is None:
  498. # drop identity
  499. text += "DROP IDENTITY"
  500. return text
  501. elif element.existing_server_default is None:
  502. # add identity options
  503. text += "ADD "
  504. text += compiler.visit_identity_column(element.default)
  505. return text
  506. else:
  507. # alter identity
  508. diff, _, _ = element.impl._compare_identity_default(
  509. element.default, element.existing_server_default
  510. )
  511. identity = element.default
  512. for attr in sorted(diff):
  513. if attr == "always":
  514. text += "SET GENERATED %s " % (
  515. "ALWAYS" if identity.always else "BY DEFAULT"
  516. )
  517. else:
  518. text += "SET %s " % compiler.get_identity_options(
  519. sqla_compat.Identity(**{attr: getattr(identity, attr)})
  520. )
  521. return text
  522. @Operations.register_operation("create_exclude_constraint")
  523. @BatchOperations.register_operation(
  524. "create_exclude_constraint", "batch_create_exclude_constraint"
  525. )
  526. @ops.AddConstraintOp.register_add_constraint("exclude_constraint")
  527. class CreateExcludeConstraintOp(ops.AddConstraintOp):
  528. """Represent a create exclude constraint operation."""
  529. constraint_type = "exclude"
  530. def __init__(
  531. self,
  532. constraint_name: sqla_compat._ConstraintName,
  533. table_name: Union[str, quoted_name],
  534. elements: Union[
  535. Sequence[Tuple[str, str]],
  536. Sequence[Tuple[ColumnClause[Any], str]],
  537. ],
  538. where: Optional[Union[ColumnElement[bool], str]] = None,
  539. schema: Optional[str] = None,
  540. _orig_constraint: Optional[ExcludeConstraint] = None,
  541. **kw,
  542. ) -> None:
  543. self.constraint_name = constraint_name
  544. self.table_name = table_name
  545. self.elements = elements
  546. self.where = where
  547. self.schema = schema
  548. self._orig_constraint = _orig_constraint
  549. self.kw = kw
  550. @classmethod
  551. def from_constraint( # type:ignore[override]
  552. cls, constraint: ExcludeConstraint
  553. ) -> CreateExcludeConstraintOp:
  554. constraint_table = sqla_compat._table_for_constraint(constraint)
  555. return cls(
  556. constraint.name,
  557. constraint_table.name,
  558. [ # type: ignore
  559. (expr, op) for expr, name, op in constraint._render_exprs
  560. ],
  561. where=cast("ColumnElement[bool] | None", constraint.where),
  562. schema=constraint_table.schema,
  563. _orig_constraint=constraint,
  564. deferrable=constraint.deferrable,
  565. initially=constraint.initially,
  566. using=constraint.using,
  567. )
  568. def to_constraint(
  569. self, migration_context: Optional[MigrationContext] = None
  570. ) -> ExcludeConstraint:
  571. if self._orig_constraint is not None:
  572. return self._orig_constraint
  573. schema_obj = schemaobj.SchemaObjects(migration_context)
  574. t = schema_obj.table(self.table_name, schema=self.schema)
  575. excl = ExcludeConstraint(
  576. *self.elements,
  577. name=self.constraint_name,
  578. where=self.where,
  579. **self.kw,
  580. )
  581. for (
  582. expr,
  583. name,
  584. oper,
  585. ) in excl._render_exprs:
  586. t.append_column(Column(name, NULLTYPE))
  587. t.append_constraint(excl)
  588. return excl
  589. @classmethod
  590. def create_exclude_constraint(
  591. cls,
  592. operations: Operations,
  593. constraint_name: str,
  594. table_name: str,
  595. *elements: Any,
  596. **kw: Any,
  597. ) -> Optional[Table]:
  598. """Issue an alter to create an EXCLUDE constraint using the
  599. current migration context.
  600. .. note:: This method is Postgresql specific, and additionally
  601. requires at least SQLAlchemy 1.0.
  602. e.g.::
  603. from alembic import op
  604. op.create_exclude_constraint(
  605. "user_excl",
  606. "user",
  607. ("period", "&&"),
  608. ("group", "="),
  609. where=("group != 'some group'"),
  610. )
  611. Note that the expressions work the same way as that of
  612. the ``ExcludeConstraint`` object itself; if plain strings are
  613. passed, quoting rules must be applied manually.
  614. :param name: Name of the constraint.
  615. :param table_name: String name of the source table.
  616. :param elements: exclude conditions.
  617. :param where: SQL expression or SQL string with optional WHERE
  618. clause.
  619. :param deferrable: optional bool. If set, emit DEFERRABLE or
  620. NOT DEFERRABLE when issuing DDL for this constraint.
  621. :param initially: optional string. If set, emit INITIALLY <value>
  622. when issuing DDL for this constraint.
  623. :param schema: Optional schema name to operate within.
  624. """
  625. op = cls(constraint_name, table_name, elements, **kw)
  626. return operations.invoke(op)
  627. @classmethod
  628. def batch_create_exclude_constraint(
  629. cls,
  630. operations: BatchOperations,
  631. constraint_name: str,
  632. *elements: Any,
  633. **kw: Any,
  634. ) -> Optional[Table]:
  635. """Issue a "create exclude constraint" instruction using the
  636. current batch migration context.
  637. .. note:: This method is Postgresql specific, and additionally
  638. requires at least SQLAlchemy 1.0.
  639. .. seealso::
  640. :meth:`.Operations.create_exclude_constraint`
  641. """
  642. kw["schema"] = operations.impl.schema
  643. op = cls(constraint_name, operations.impl.table_name, elements, **kw)
  644. return operations.invoke(op)
  645. @render.renderers.dispatch_for(CreateExcludeConstraintOp)
  646. def _add_exclude_constraint(
  647. autogen_context: AutogenContext, op: CreateExcludeConstraintOp
  648. ) -> str:
  649. return _exclude_constraint(op.to_constraint(), autogen_context, alter=True)
  650. @render._constraint_renderers.dispatch_for(ExcludeConstraint)
  651. def _render_inline_exclude_constraint(
  652. constraint: ExcludeConstraint,
  653. autogen_context: AutogenContext,
  654. namespace_metadata: MetaData,
  655. ) -> str:
  656. rendered = render._user_defined_render(
  657. "exclude", constraint, autogen_context
  658. )
  659. if rendered is not False:
  660. return rendered
  661. return _exclude_constraint(constraint, autogen_context, False)
  662. def _postgresql_autogenerate_prefix(autogen_context: AutogenContext) -> str:
  663. imports = autogen_context.imports
  664. if imports is not None:
  665. imports.add("from sqlalchemy.dialects import postgresql")
  666. return "postgresql."
  667. def _exclude_constraint(
  668. constraint: ExcludeConstraint,
  669. autogen_context: AutogenContext,
  670. alter: bool,
  671. ) -> str:
  672. opts: List[Tuple[str, Union[quoted_name, str, _f_name, None]]] = []
  673. has_batch = autogen_context._has_batch
  674. if constraint.deferrable:
  675. opts.append(("deferrable", str(constraint.deferrable)))
  676. if constraint.initially:
  677. opts.append(("initially", str(constraint.initially)))
  678. if constraint.using:
  679. opts.append(("using", str(constraint.using)))
  680. if not has_batch and alter and constraint.table.schema:
  681. opts.append(("schema", render._ident(constraint.table.schema)))
  682. if not alter and constraint.name:
  683. opts.append(
  684. ("name", render._render_gen_name(autogen_context, constraint.name))
  685. )
  686. def do_expr_where_opts():
  687. args = [
  688. "(%s, %r)"
  689. % (
  690. _render_potential_column(
  691. sqltext, # type:ignore[arg-type]
  692. autogen_context,
  693. ),
  694. opstring,
  695. )
  696. for sqltext, name, opstring in constraint._render_exprs
  697. ]
  698. if constraint.where is not None:
  699. args.append(
  700. "where=%s"
  701. % render._render_potential_expr(
  702. constraint.where, autogen_context
  703. )
  704. )
  705. args.extend(["%s=%r" % (k, v) for k, v in opts])
  706. return args
  707. if alter:
  708. args = [
  709. repr(render._render_gen_name(autogen_context, constraint.name))
  710. ]
  711. if not has_batch:
  712. args += [repr(render._ident(constraint.table.name))]
  713. args.extend(do_expr_where_opts())
  714. return "%(prefix)screate_exclude_constraint(%(args)s)" % {
  715. "prefix": render._alembic_autogenerate_prefix(autogen_context),
  716. "args": ", ".join(args),
  717. }
  718. else:
  719. args = do_expr_where_opts()
  720. return "%(prefix)sExcludeConstraint(%(args)s)" % {
  721. "prefix": _postgresql_autogenerate_prefix(autogen_context),
  722. "args": ", ".join(args),
  723. }
  724. def _render_potential_column(
  725. value: Union[
  726. ColumnClause[Any], Column[Any], TextClause, FunctionElement[Any]
  727. ],
  728. autogen_context: AutogenContext,
  729. ) -> str:
  730. if isinstance(value, ColumnClause):
  731. if value.is_literal:
  732. # like literal_column("int8range(from, to)") in ExcludeConstraint
  733. template = "%(prefix)sliteral_column(%(name)r)"
  734. else:
  735. template = "%(prefix)scolumn(%(name)r)"
  736. return template % {
  737. "prefix": render._sqlalchemy_autogenerate_prefix(autogen_context),
  738. "name": value.name,
  739. }
  740. else:
  741. return render._render_potential_expr(
  742. value,
  743. autogen_context,
  744. wrap_in_text=isinstance(value, (TextClause, FunctionElement)),
  745. )