render.py 34 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097
  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. from io import StringIO
  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 Tuple
  12. from typing import TYPE_CHECKING
  13. from typing import Union
  14. from mako.pygen import PythonPrinter
  15. from sqlalchemy import schema as sa_schema
  16. from sqlalchemy import sql
  17. from sqlalchemy import types as sqltypes
  18. from sqlalchemy.sql.elements import conv
  19. from sqlalchemy.sql.elements import quoted_name
  20. from .. import util
  21. from ..operations import ops
  22. from ..util import sqla_compat
  23. if TYPE_CHECKING:
  24. from typing import Literal
  25. from sqlalchemy.sql.base import DialectKWArgs
  26. from sqlalchemy.sql.elements import ColumnElement
  27. from sqlalchemy.sql.elements import TextClause
  28. from sqlalchemy.sql.schema import CheckConstraint
  29. from sqlalchemy.sql.schema import Column
  30. from sqlalchemy.sql.schema import Constraint
  31. from sqlalchemy.sql.schema import FetchedValue
  32. from sqlalchemy.sql.schema import ForeignKey
  33. from sqlalchemy.sql.schema import ForeignKeyConstraint
  34. from sqlalchemy.sql.schema import Index
  35. from sqlalchemy.sql.schema import MetaData
  36. from sqlalchemy.sql.schema import PrimaryKeyConstraint
  37. from sqlalchemy.sql.schema import UniqueConstraint
  38. from sqlalchemy.sql.sqltypes import ARRAY
  39. from sqlalchemy.sql.type_api import TypeEngine
  40. from alembic.autogenerate.api import AutogenContext
  41. from alembic.config import Config
  42. from alembic.operations.ops import MigrationScript
  43. from alembic.operations.ops import ModifyTableOps
  44. from alembic.util.sqla_compat import Computed
  45. from alembic.util.sqla_compat import Identity
  46. MAX_PYTHON_ARGS = 255
  47. def _render_gen_name(
  48. autogen_context: AutogenContext,
  49. name: sqla_compat._ConstraintName,
  50. ) -> Optional[Union[quoted_name, str, _f_name]]:
  51. if isinstance(name, conv):
  52. return _f_name(_alembic_autogenerate_prefix(autogen_context), name)
  53. else:
  54. return sqla_compat.constraint_name_or_none(name)
  55. def _indent(text: str) -> str:
  56. text = re.compile(r"^", re.M).sub(" ", text).strip()
  57. text = re.compile(r" +$", re.M).sub("", text)
  58. return text
  59. def _render_python_into_templatevars(
  60. autogen_context: AutogenContext,
  61. migration_script: MigrationScript,
  62. template_args: Dict[str, Union[str, Config]],
  63. ) -> None:
  64. imports = autogen_context.imports
  65. for upgrade_ops, downgrade_ops in zip(
  66. migration_script.upgrade_ops_list, migration_script.downgrade_ops_list
  67. ):
  68. template_args[upgrade_ops.upgrade_token] = _indent(
  69. _render_cmd_body(upgrade_ops, autogen_context)
  70. )
  71. template_args[downgrade_ops.downgrade_token] = _indent(
  72. _render_cmd_body(downgrade_ops, autogen_context)
  73. )
  74. template_args["imports"] = "\n".join(sorted(imports))
  75. default_renderers = renderers = util.Dispatcher()
  76. def _render_cmd_body(
  77. op_container: ops.OpContainer,
  78. autogen_context: AutogenContext,
  79. ) -> str:
  80. buf = StringIO()
  81. printer = PythonPrinter(buf)
  82. printer.writeline(
  83. "# ### commands auto generated by Alembic - please adjust! ###"
  84. )
  85. has_lines = False
  86. for op in op_container.ops:
  87. lines = render_op(autogen_context, op)
  88. has_lines = has_lines or bool(lines)
  89. for line in lines:
  90. printer.writeline(line)
  91. if not has_lines:
  92. printer.writeline("pass")
  93. printer.writeline("# ### end Alembic commands ###")
  94. return buf.getvalue()
  95. def render_op(
  96. autogen_context: AutogenContext, op: ops.MigrateOperation
  97. ) -> List[str]:
  98. renderer = renderers.dispatch(op)
  99. lines = util.to_list(renderer(autogen_context, op))
  100. return lines
  101. def render_op_text(
  102. autogen_context: AutogenContext, op: ops.MigrateOperation
  103. ) -> str:
  104. return "\n".join(render_op(autogen_context, op))
  105. @renderers.dispatch_for(ops.ModifyTableOps)
  106. def _render_modify_table(
  107. autogen_context: AutogenContext, op: ModifyTableOps
  108. ) -> List[str]:
  109. opts = autogen_context.opts
  110. render_as_batch = opts.get("render_as_batch", False)
  111. if op.ops:
  112. lines = []
  113. if render_as_batch:
  114. with autogen_context._within_batch():
  115. lines.append(
  116. "with op.batch_alter_table(%r, schema=%r) as batch_op:"
  117. % (op.table_name, op.schema)
  118. )
  119. for t_op in op.ops:
  120. t_lines = render_op(autogen_context, t_op)
  121. lines.extend(t_lines)
  122. lines.append("")
  123. else:
  124. for t_op in op.ops:
  125. t_lines = render_op(autogen_context, t_op)
  126. lines.extend(t_lines)
  127. return lines
  128. else:
  129. return []
  130. @renderers.dispatch_for(ops.CreateTableCommentOp)
  131. def _render_create_table_comment(
  132. autogen_context: AutogenContext, op: ops.CreateTableCommentOp
  133. ) -> str:
  134. if autogen_context._has_batch:
  135. templ = (
  136. "{prefix}create_table_comment(\n"
  137. "{indent}{comment},\n"
  138. "{indent}existing_comment={existing}\n"
  139. ")"
  140. )
  141. else:
  142. templ = (
  143. "{prefix}create_table_comment(\n"
  144. "{indent}'{tname}',\n"
  145. "{indent}{comment},\n"
  146. "{indent}existing_comment={existing},\n"
  147. "{indent}schema={schema}\n"
  148. ")"
  149. )
  150. return templ.format(
  151. prefix=_alembic_autogenerate_prefix(autogen_context),
  152. tname=op.table_name,
  153. comment="%r" % op.comment if op.comment is not None else None,
  154. existing="%r" % op.existing_comment
  155. if op.existing_comment is not None
  156. else None,
  157. schema="'%s'" % op.schema if op.schema is not None else None,
  158. indent=" ",
  159. )
  160. @renderers.dispatch_for(ops.DropTableCommentOp)
  161. def _render_drop_table_comment(
  162. autogen_context: AutogenContext, op: ops.DropTableCommentOp
  163. ) -> str:
  164. if autogen_context._has_batch:
  165. templ = (
  166. "{prefix}drop_table_comment(\n"
  167. "{indent}existing_comment={existing}\n"
  168. ")"
  169. )
  170. else:
  171. templ = (
  172. "{prefix}drop_table_comment(\n"
  173. "{indent}'{tname}',\n"
  174. "{indent}existing_comment={existing},\n"
  175. "{indent}schema={schema}\n"
  176. ")"
  177. )
  178. return templ.format(
  179. prefix=_alembic_autogenerate_prefix(autogen_context),
  180. tname=op.table_name,
  181. existing="%r" % op.existing_comment
  182. if op.existing_comment is not None
  183. else None,
  184. schema="'%s'" % op.schema if op.schema is not None else None,
  185. indent=" ",
  186. )
  187. @renderers.dispatch_for(ops.CreateTableOp)
  188. def _add_table(autogen_context: AutogenContext, op: ops.CreateTableOp) -> str:
  189. table = op.to_table()
  190. args = [
  191. col
  192. for col in [
  193. _render_column(col, autogen_context) for col in table.columns
  194. ]
  195. if col
  196. ] + sorted(
  197. [
  198. rcons
  199. for rcons in [
  200. _render_constraint(
  201. cons, autogen_context, op._namespace_metadata
  202. )
  203. for cons in table.constraints
  204. ]
  205. if rcons is not None
  206. ]
  207. )
  208. if len(args) > MAX_PYTHON_ARGS:
  209. args_str = "*[" + ",\n".join(args) + "]"
  210. else:
  211. args_str = ",\n".join(args)
  212. text = "%(prefix)screate_table(%(tablename)r,\n%(args)s" % {
  213. "tablename": _ident(op.table_name),
  214. "prefix": _alembic_autogenerate_prefix(autogen_context),
  215. "args": args_str,
  216. }
  217. if op.schema:
  218. text += ",\nschema=%r" % _ident(op.schema)
  219. comment = table.comment
  220. if comment:
  221. text += ",\ncomment=%r" % _ident(comment)
  222. info = table.info
  223. if info:
  224. text += f",\ninfo={info!r}"
  225. for k in sorted(op.kw):
  226. text += ",\n%s=%r" % (k.replace(" ", "_"), op.kw[k])
  227. if table._prefixes:
  228. prefixes = ", ".join("'%s'" % p for p in table._prefixes)
  229. text += ",\nprefixes=[%s]" % prefixes
  230. text += "\n)"
  231. return text
  232. @renderers.dispatch_for(ops.DropTableOp)
  233. def _drop_table(autogen_context: AutogenContext, op: ops.DropTableOp) -> str:
  234. text = "%(prefix)sdrop_table(%(tname)r" % {
  235. "prefix": _alembic_autogenerate_prefix(autogen_context),
  236. "tname": _ident(op.table_name),
  237. }
  238. if op.schema:
  239. text += ", schema=%r" % _ident(op.schema)
  240. text += ")"
  241. return text
  242. def _render_dialect_kwargs_items(
  243. autogen_context: AutogenContext, item: DialectKWArgs
  244. ) -> list[str]:
  245. return [
  246. f"{key}={_render_potential_expr(val, autogen_context)}"
  247. for key, val in item.dialect_kwargs.items()
  248. ]
  249. @renderers.dispatch_for(ops.CreateIndexOp)
  250. def _add_index(autogen_context: AutogenContext, op: ops.CreateIndexOp) -> str:
  251. index = op.to_index()
  252. has_batch = autogen_context._has_batch
  253. if has_batch:
  254. tmpl = (
  255. "%(prefix)screate_index(%(name)r, [%(columns)s], "
  256. "unique=%(unique)r%(kwargs)s)"
  257. )
  258. else:
  259. tmpl = (
  260. "%(prefix)screate_index(%(name)r, %(table)r, [%(columns)s], "
  261. "unique=%(unique)r%(schema)s%(kwargs)s)"
  262. )
  263. assert index.table is not None
  264. opts = _render_dialect_kwargs_items(autogen_context, index)
  265. text = tmpl % {
  266. "prefix": _alembic_autogenerate_prefix(autogen_context),
  267. "name": _render_gen_name(autogen_context, index.name),
  268. "table": _ident(index.table.name),
  269. "columns": ", ".join(
  270. _get_index_rendered_expressions(index, autogen_context)
  271. ),
  272. "unique": index.unique or False,
  273. "schema": (", schema=%r" % _ident(index.table.schema))
  274. if index.table.schema
  275. else "",
  276. "kwargs": ", " + ", ".join(opts) if opts else "",
  277. }
  278. return text
  279. @renderers.dispatch_for(ops.DropIndexOp)
  280. def _drop_index(autogen_context: AutogenContext, op: ops.DropIndexOp) -> str:
  281. index = op.to_index()
  282. has_batch = autogen_context._has_batch
  283. if has_batch:
  284. tmpl = "%(prefix)sdrop_index(%(name)r%(kwargs)s)"
  285. else:
  286. tmpl = (
  287. "%(prefix)sdrop_index(%(name)r, "
  288. "table_name=%(table_name)r%(schema)s%(kwargs)s)"
  289. )
  290. opts = _render_dialect_kwargs_items(autogen_context, index)
  291. text = tmpl % {
  292. "prefix": _alembic_autogenerate_prefix(autogen_context),
  293. "name": _render_gen_name(autogen_context, op.index_name),
  294. "table_name": _ident(op.table_name),
  295. "schema": ((", schema=%r" % _ident(op.schema)) if op.schema else ""),
  296. "kwargs": ", " + ", ".join(opts) if opts else "",
  297. }
  298. return text
  299. @renderers.dispatch_for(ops.CreateUniqueConstraintOp)
  300. def _add_unique_constraint(
  301. autogen_context: AutogenContext, op: ops.CreateUniqueConstraintOp
  302. ) -> List[str]:
  303. return [_uq_constraint(op.to_constraint(), autogen_context, True)]
  304. @renderers.dispatch_for(ops.CreateForeignKeyOp)
  305. def _add_fk_constraint(
  306. autogen_context: AutogenContext, op: ops.CreateForeignKeyOp
  307. ) -> str:
  308. args = [repr(_render_gen_name(autogen_context, op.constraint_name))]
  309. if not autogen_context._has_batch:
  310. args.append(repr(_ident(op.source_table)))
  311. args.extend(
  312. [
  313. repr(_ident(op.referent_table)),
  314. repr([_ident(col) for col in op.local_cols]),
  315. repr([_ident(col) for col in op.remote_cols]),
  316. ]
  317. )
  318. kwargs = [
  319. "referent_schema",
  320. "onupdate",
  321. "ondelete",
  322. "initially",
  323. "deferrable",
  324. "use_alter",
  325. "match",
  326. ]
  327. if not autogen_context._has_batch:
  328. kwargs.insert(0, "source_schema")
  329. for k in kwargs:
  330. if k in op.kw:
  331. value = op.kw[k]
  332. if value is not None:
  333. args.append("%s=%r" % (k, value))
  334. return "%(prefix)screate_foreign_key(%(args)s)" % {
  335. "prefix": _alembic_autogenerate_prefix(autogen_context),
  336. "args": ", ".join(args),
  337. }
  338. @renderers.dispatch_for(ops.CreatePrimaryKeyOp)
  339. def _add_pk_constraint(constraint, autogen_context):
  340. raise NotImplementedError()
  341. @renderers.dispatch_for(ops.CreateCheckConstraintOp)
  342. def _add_check_constraint(constraint, autogen_context):
  343. raise NotImplementedError()
  344. @renderers.dispatch_for(ops.DropConstraintOp)
  345. def _drop_constraint(
  346. autogen_context: AutogenContext, op: ops.DropConstraintOp
  347. ) -> str:
  348. prefix = _alembic_autogenerate_prefix(autogen_context)
  349. name = _render_gen_name(autogen_context, op.constraint_name)
  350. schema = _ident(op.schema) if op.schema else None
  351. type_ = _ident(op.constraint_type) if op.constraint_type else None
  352. params_strs = []
  353. params_strs.append(repr(name))
  354. if not autogen_context._has_batch:
  355. params_strs.append(repr(_ident(op.table_name)))
  356. if schema is not None:
  357. params_strs.append(f"schema={schema!r}")
  358. if type_ is not None:
  359. params_strs.append(f"type_={type_!r}")
  360. return f"{prefix}drop_constraint({', '.join(params_strs)})"
  361. @renderers.dispatch_for(ops.AddColumnOp)
  362. def _add_column(autogen_context: AutogenContext, op: ops.AddColumnOp) -> str:
  363. schema, tname, column = op.schema, op.table_name, op.column
  364. if autogen_context._has_batch:
  365. template = "%(prefix)sadd_column(%(column)s)"
  366. else:
  367. template = "%(prefix)sadd_column(%(tname)r, %(column)s"
  368. if schema:
  369. template += ", schema=%(schema)r"
  370. template += ")"
  371. text = template % {
  372. "prefix": _alembic_autogenerate_prefix(autogen_context),
  373. "tname": tname,
  374. "column": _render_column(column, autogen_context),
  375. "schema": schema,
  376. }
  377. return text
  378. @renderers.dispatch_for(ops.DropColumnOp)
  379. def _drop_column(autogen_context: AutogenContext, op: ops.DropColumnOp) -> str:
  380. schema, tname, column_name = op.schema, op.table_name, op.column_name
  381. if autogen_context._has_batch:
  382. template = "%(prefix)sdrop_column(%(cname)r)"
  383. else:
  384. template = "%(prefix)sdrop_column(%(tname)r, %(cname)r"
  385. if schema:
  386. template += ", schema=%(schema)r"
  387. template += ")"
  388. text = template % {
  389. "prefix": _alembic_autogenerate_prefix(autogen_context),
  390. "tname": _ident(tname),
  391. "cname": _ident(column_name),
  392. "schema": _ident(schema),
  393. }
  394. return text
  395. @renderers.dispatch_for(ops.AlterColumnOp)
  396. def _alter_column(
  397. autogen_context: AutogenContext, op: ops.AlterColumnOp
  398. ) -> str:
  399. tname = op.table_name
  400. cname = op.column_name
  401. server_default = op.modify_server_default
  402. type_ = op.modify_type
  403. nullable = op.modify_nullable
  404. comment = op.modify_comment
  405. autoincrement = op.kw.get("autoincrement", None)
  406. existing_type = op.existing_type
  407. existing_nullable = op.existing_nullable
  408. existing_comment = op.existing_comment
  409. existing_server_default = op.existing_server_default
  410. schema = op.schema
  411. indent = " " * 11
  412. if autogen_context._has_batch:
  413. template = "%(prefix)salter_column(%(cname)r"
  414. else:
  415. template = "%(prefix)salter_column(%(tname)r, %(cname)r"
  416. text = template % {
  417. "prefix": _alembic_autogenerate_prefix(autogen_context),
  418. "tname": tname,
  419. "cname": cname,
  420. }
  421. if existing_type is not None:
  422. text += ",\n%sexisting_type=%s" % (
  423. indent,
  424. _repr_type(existing_type, autogen_context),
  425. )
  426. if server_default is not False:
  427. rendered = _render_server_default(server_default, autogen_context)
  428. text += ",\n%sserver_default=%s" % (indent, rendered)
  429. if type_ is not None:
  430. text += ",\n%stype_=%s" % (indent, _repr_type(type_, autogen_context))
  431. if nullable is not None:
  432. text += ",\n%snullable=%r" % (indent, nullable)
  433. if comment is not False:
  434. text += ",\n%scomment=%r" % (indent, comment)
  435. if existing_comment is not None:
  436. text += ",\n%sexisting_comment=%r" % (indent, existing_comment)
  437. if nullable is None and existing_nullable is not None:
  438. text += ",\n%sexisting_nullable=%r" % (indent, existing_nullable)
  439. if autoincrement is not None:
  440. text += ",\n%sautoincrement=%r" % (indent, autoincrement)
  441. if server_default is False and existing_server_default:
  442. rendered = _render_server_default(
  443. existing_server_default, autogen_context
  444. )
  445. text += ",\n%sexisting_server_default=%s" % (indent, rendered)
  446. if schema and not autogen_context._has_batch:
  447. text += ",\n%sschema=%r" % (indent, schema)
  448. text += ")"
  449. return text
  450. class _f_name:
  451. def __init__(self, prefix: str, name: conv) -> None:
  452. self.prefix = prefix
  453. self.name = name
  454. def __repr__(self) -> str:
  455. return "%sf(%r)" % (self.prefix, _ident(self.name))
  456. def _ident(name: Optional[Union[quoted_name, str]]) -> Optional[str]:
  457. """produce a __repr__() object for a string identifier that may
  458. use quoted_name() in SQLAlchemy 0.9 and greater.
  459. The issue worked around here is that quoted_name() doesn't have
  460. very good repr() behavior by itself when unicode is involved.
  461. """
  462. if name is None:
  463. return name
  464. elif isinstance(name, quoted_name):
  465. return str(name)
  466. elif isinstance(name, str):
  467. return name
  468. def _render_potential_expr(
  469. value: Any,
  470. autogen_context: AutogenContext,
  471. *,
  472. wrap_in_text: bool = True,
  473. is_server_default: bool = False,
  474. is_index: bool = False,
  475. ) -> str:
  476. if isinstance(value, sql.ClauseElement):
  477. if wrap_in_text:
  478. template = "%(prefix)stext(%(sql)r)"
  479. else:
  480. template = "%(sql)r"
  481. return template % {
  482. "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
  483. "sql": autogen_context.migration_context.impl.render_ddl_sql_expr(
  484. value, is_server_default=is_server_default, is_index=is_index
  485. ),
  486. }
  487. else:
  488. return repr(value)
  489. def _get_index_rendered_expressions(
  490. idx: Index, autogen_context: AutogenContext
  491. ) -> List[str]:
  492. return [
  493. repr(_ident(getattr(exp, "name", None)))
  494. if isinstance(exp, sa_schema.Column)
  495. else _render_potential_expr(exp, autogen_context, is_index=True)
  496. for exp in idx.expressions
  497. ]
  498. def _uq_constraint(
  499. constraint: UniqueConstraint,
  500. autogen_context: AutogenContext,
  501. alter: bool,
  502. ) -> str:
  503. opts: List[Tuple[str, Any]] = []
  504. has_batch = autogen_context._has_batch
  505. if constraint.deferrable:
  506. opts.append(("deferrable", str(constraint.deferrable)))
  507. if constraint.initially:
  508. opts.append(("initially", str(constraint.initially)))
  509. if not has_batch and alter and constraint.table.schema:
  510. opts.append(("schema", _ident(constraint.table.schema)))
  511. if not alter and constraint.name:
  512. opts.append(
  513. ("name", _render_gen_name(autogen_context, constraint.name))
  514. )
  515. dialect_options = _render_dialect_kwargs_items(autogen_context, constraint)
  516. if alter:
  517. args = [repr(_render_gen_name(autogen_context, constraint.name))]
  518. if not has_batch:
  519. args += [repr(_ident(constraint.table.name))]
  520. args.append(repr([_ident(col.name) for col in constraint.columns]))
  521. args.extend(["%s=%r" % (k, v) for k, v in opts])
  522. args.extend(dialect_options)
  523. return "%(prefix)screate_unique_constraint(%(args)s)" % {
  524. "prefix": _alembic_autogenerate_prefix(autogen_context),
  525. "args": ", ".join(args),
  526. }
  527. else:
  528. args = [repr(_ident(col.name)) for col in constraint.columns]
  529. args.extend(["%s=%r" % (k, v) for k, v in opts])
  530. args.extend(dialect_options)
  531. return "%(prefix)sUniqueConstraint(%(args)s)" % {
  532. "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
  533. "args": ", ".join(args),
  534. }
  535. def _user_autogenerate_prefix(autogen_context, target):
  536. prefix = autogen_context.opts["user_module_prefix"]
  537. if prefix is None:
  538. return "%s." % target.__module__
  539. else:
  540. return prefix
  541. def _sqlalchemy_autogenerate_prefix(autogen_context: AutogenContext) -> str:
  542. return autogen_context.opts["sqlalchemy_module_prefix"] or ""
  543. def _alembic_autogenerate_prefix(autogen_context: AutogenContext) -> str:
  544. if autogen_context._has_batch:
  545. return "batch_op."
  546. else:
  547. return autogen_context.opts["alembic_module_prefix"] or ""
  548. def _user_defined_render(
  549. type_: str, object_: Any, autogen_context: AutogenContext
  550. ) -> Union[str, Literal[False]]:
  551. if "render_item" in autogen_context.opts:
  552. render = autogen_context.opts["render_item"]
  553. if render:
  554. rendered = render(type_, object_, autogen_context)
  555. if rendered is not False:
  556. return rendered
  557. return False
  558. def _render_column(
  559. column: Column[Any], autogen_context: AutogenContext
  560. ) -> str:
  561. rendered = _user_defined_render("column", column, autogen_context)
  562. if rendered is not False:
  563. return rendered
  564. args: List[str] = []
  565. opts: List[Tuple[str, Any]] = []
  566. if column.server_default:
  567. rendered = _render_server_default( # type:ignore[assignment]
  568. column.server_default, autogen_context
  569. )
  570. if rendered:
  571. if _should_render_server_default_positionally(
  572. column.server_default
  573. ):
  574. args.append(rendered)
  575. else:
  576. opts.append(("server_default", rendered))
  577. if (
  578. column.autoincrement is not None
  579. and column.autoincrement != sqla_compat.AUTOINCREMENT_DEFAULT
  580. ):
  581. opts.append(("autoincrement", column.autoincrement))
  582. if column.nullable is not None:
  583. opts.append(("nullable", column.nullable))
  584. if column.system:
  585. opts.append(("system", column.system))
  586. comment = column.comment
  587. if comment:
  588. opts.append(("comment", "%r" % comment))
  589. # TODO: for non-ascii colname, assign a "key"
  590. return "%(prefix)sColumn(%(name)r, %(type)s, %(args)s%(kwargs)s)" % {
  591. "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
  592. "name": _ident(column.name),
  593. "type": _repr_type(column.type, autogen_context),
  594. "args": ", ".join([str(arg) for arg in args]) + ", " if args else "",
  595. "kwargs": (
  596. ", ".join(
  597. ["%s=%s" % (kwname, val) for kwname, val in opts]
  598. + [
  599. "%s=%s"
  600. % (key, _render_potential_expr(val, autogen_context))
  601. for key, val in sqla_compat._column_kwargs(column).items()
  602. ]
  603. )
  604. ),
  605. }
  606. def _should_render_server_default_positionally(server_default: Any) -> bool:
  607. return sqla_compat._server_default_is_computed(
  608. server_default
  609. ) or sqla_compat._server_default_is_identity(server_default)
  610. def _render_server_default(
  611. default: Optional[
  612. Union[FetchedValue, str, TextClause, ColumnElement[Any]]
  613. ],
  614. autogen_context: AutogenContext,
  615. repr_: bool = True,
  616. ) -> Optional[str]:
  617. rendered = _user_defined_render("server_default", default, autogen_context)
  618. if rendered is not False:
  619. return rendered
  620. if sqla_compat._server_default_is_computed(default):
  621. return _render_computed(cast("Computed", default), autogen_context)
  622. elif sqla_compat._server_default_is_identity(default):
  623. return _render_identity(cast("Identity", default), autogen_context)
  624. elif isinstance(default, sa_schema.DefaultClause):
  625. if isinstance(default.arg, str):
  626. default = default.arg
  627. else:
  628. return _render_potential_expr(
  629. default.arg, autogen_context, is_server_default=True
  630. )
  631. if isinstance(default, str) and repr_:
  632. default = repr(re.sub(r"^'|'$", "", default))
  633. return cast(str, default)
  634. def _render_computed(
  635. computed: Computed, autogen_context: AutogenContext
  636. ) -> str:
  637. text = _render_potential_expr(
  638. computed.sqltext, autogen_context, wrap_in_text=False
  639. )
  640. kwargs = {}
  641. if computed.persisted is not None:
  642. kwargs["persisted"] = computed.persisted
  643. return "%(prefix)sComputed(%(text)s, %(kwargs)s)" % {
  644. "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
  645. "text": text,
  646. "kwargs": (", ".join("%s=%s" % pair for pair in kwargs.items())),
  647. }
  648. def _render_identity(
  649. identity: Identity, autogen_context: AutogenContext
  650. ) -> str:
  651. kwargs = sqla_compat._get_identity_options_dict(
  652. identity, dialect_kwargs=True
  653. )
  654. return "%(prefix)sIdentity(%(kwargs)s)" % {
  655. "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
  656. "kwargs": (", ".join("%s=%s" % pair for pair in kwargs.items())),
  657. }
  658. def _repr_type(
  659. type_: TypeEngine,
  660. autogen_context: AutogenContext,
  661. _skip_variants: bool = False,
  662. ) -> str:
  663. rendered = _user_defined_render("type", type_, autogen_context)
  664. if rendered is not False:
  665. return rendered
  666. if hasattr(autogen_context.migration_context, "impl"):
  667. impl_rt = autogen_context.migration_context.impl.render_type(
  668. type_, autogen_context
  669. )
  670. else:
  671. impl_rt = None
  672. mod = type(type_).__module__
  673. imports = autogen_context.imports
  674. if mod.startswith("sqlalchemy.dialects"):
  675. match = re.match(r"sqlalchemy\.dialects\.(\w+)", mod)
  676. assert match is not None
  677. dname = match.group(1)
  678. if imports is not None:
  679. imports.add("from sqlalchemy.dialects import %s" % dname)
  680. if impl_rt:
  681. return impl_rt
  682. else:
  683. return "%s.%r" % (dname, type_)
  684. elif impl_rt:
  685. return impl_rt
  686. elif not _skip_variants and sqla_compat._type_has_variants(type_):
  687. return _render_Variant_type(type_, autogen_context)
  688. elif mod.startswith("sqlalchemy."):
  689. if "_render_%s_type" % type_.__visit_name__ in globals():
  690. fn = globals()["_render_%s_type" % type_.__visit_name__]
  691. return fn(type_, autogen_context)
  692. else:
  693. prefix = _sqlalchemy_autogenerate_prefix(autogen_context)
  694. return "%s%r" % (prefix, type_)
  695. else:
  696. prefix = _user_autogenerate_prefix(autogen_context, type_)
  697. return "%s%r" % (prefix, type_)
  698. def _render_ARRAY_type(type_: ARRAY, autogen_context: AutogenContext) -> str:
  699. return cast(
  700. str,
  701. _render_type_w_subtype(
  702. type_, autogen_context, "item_type", r"(.+?\()"
  703. ),
  704. )
  705. def _render_Variant_type(
  706. type_: TypeEngine, autogen_context: AutogenContext
  707. ) -> str:
  708. base_type, variant_mapping = sqla_compat._get_variant_mapping(type_)
  709. base = _repr_type(base_type, autogen_context, _skip_variants=True)
  710. assert base is not None and base is not False # type: ignore[comparison-overlap] # noqa:E501
  711. for dialect in sorted(variant_mapping):
  712. typ = variant_mapping[dialect]
  713. base += ".with_variant(%s, %r)" % (
  714. _repr_type(typ, autogen_context, _skip_variants=True),
  715. dialect,
  716. )
  717. return base
  718. def _render_type_w_subtype(
  719. type_: TypeEngine,
  720. autogen_context: AutogenContext,
  721. attrname: str,
  722. regexp: str,
  723. prefix: Optional[str] = None,
  724. ) -> Union[Optional[str], Literal[False]]:
  725. outer_repr = repr(type_)
  726. inner_type = getattr(type_, attrname, None)
  727. if inner_type is None:
  728. return False
  729. inner_repr = repr(inner_type)
  730. inner_repr = re.sub(r"([\(\)])", r"\\\1", inner_repr)
  731. sub_type = _repr_type(getattr(type_, attrname), autogen_context)
  732. outer_type = re.sub(regexp + inner_repr, r"\1%s" % sub_type, outer_repr)
  733. if prefix:
  734. return "%s%s" % (prefix, outer_type)
  735. mod = type(type_).__module__
  736. if mod.startswith("sqlalchemy.dialects"):
  737. match = re.match(r"sqlalchemy\.dialects\.(\w+)", mod)
  738. assert match is not None
  739. dname = match.group(1)
  740. return "%s.%s" % (dname, outer_type)
  741. elif mod.startswith("sqlalchemy"):
  742. prefix = _sqlalchemy_autogenerate_prefix(autogen_context)
  743. return "%s%s" % (prefix, outer_type)
  744. else:
  745. return None
  746. _constraint_renderers = util.Dispatcher()
  747. def _render_constraint(
  748. constraint: Constraint,
  749. autogen_context: AutogenContext,
  750. namespace_metadata: Optional[MetaData],
  751. ) -> Optional[str]:
  752. try:
  753. renderer = _constraint_renderers.dispatch(constraint)
  754. except ValueError:
  755. util.warn("No renderer is established for object %r" % constraint)
  756. return "[Unknown Python object %r]" % constraint
  757. else:
  758. return renderer(constraint, autogen_context, namespace_metadata)
  759. @_constraint_renderers.dispatch_for(sa_schema.PrimaryKeyConstraint)
  760. def _render_primary_key(
  761. constraint: PrimaryKeyConstraint,
  762. autogen_context: AutogenContext,
  763. namespace_metadata: Optional[MetaData],
  764. ) -> Optional[str]:
  765. rendered = _user_defined_render("primary_key", constraint, autogen_context)
  766. if rendered is not False:
  767. return rendered
  768. if not constraint.columns:
  769. return None
  770. opts = []
  771. if constraint.name:
  772. opts.append(
  773. ("name", repr(_render_gen_name(autogen_context, constraint.name)))
  774. )
  775. return "%(prefix)sPrimaryKeyConstraint(%(args)s)" % {
  776. "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
  777. "args": ", ".join(
  778. [repr(c.name) for c in constraint.columns]
  779. + ["%s=%s" % (kwname, val) for kwname, val in opts]
  780. ),
  781. }
  782. def _fk_colspec(
  783. fk: ForeignKey,
  784. metadata_schema: Optional[str],
  785. namespace_metadata: MetaData,
  786. ) -> str:
  787. """Implement a 'safe' version of ForeignKey._get_colspec() that
  788. won't fail if the remote table can't be resolved.
  789. """
  790. colspec = fk._get_colspec()
  791. tokens = colspec.split(".")
  792. tname, colname = tokens[-2:]
  793. if metadata_schema is not None and len(tokens) == 2:
  794. table_fullname = "%s.%s" % (metadata_schema, tname)
  795. else:
  796. table_fullname = ".".join(tokens[0:-1])
  797. if (
  798. not fk.link_to_name
  799. and fk.parent is not None
  800. and fk.parent.table is not None
  801. ):
  802. # try to resolve the remote table in order to adjust for column.key.
  803. # the FK constraint needs to be rendered in terms of the column
  804. # name.
  805. if table_fullname in namespace_metadata.tables:
  806. col = namespace_metadata.tables[table_fullname].c.get(colname)
  807. if col is not None:
  808. colname = _ident(col.name) # type: ignore[assignment]
  809. colspec = "%s.%s" % (table_fullname, colname)
  810. return colspec
  811. def _populate_render_fk_opts(
  812. constraint: ForeignKeyConstraint, opts: List[Tuple[str, str]]
  813. ) -> None:
  814. if constraint.onupdate:
  815. opts.append(("onupdate", repr(constraint.onupdate)))
  816. if constraint.ondelete:
  817. opts.append(("ondelete", repr(constraint.ondelete)))
  818. if constraint.initially:
  819. opts.append(("initially", repr(constraint.initially)))
  820. if constraint.deferrable:
  821. opts.append(("deferrable", repr(constraint.deferrable)))
  822. if constraint.use_alter:
  823. opts.append(("use_alter", repr(constraint.use_alter)))
  824. if constraint.match:
  825. opts.append(("match", repr(constraint.match)))
  826. @_constraint_renderers.dispatch_for(sa_schema.ForeignKeyConstraint)
  827. def _render_foreign_key(
  828. constraint: ForeignKeyConstraint,
  829. autogen_context: AutogenContext,
  830. namespace_metadata: MetaData,
  831. ) -> Optional[str]:
  832. rendered = _user_defined_render("foreign_key", constraint, autogen_context)
  833. if rendered is not False:
  834. return rendered
  835. opts = []
  836. if constraint.name:
  837. opts.append(
  838. ("name", repr(_render_gen_name(autogen_context, constraint.name)))
  839. )
  840. _populate_render_fk_opts(constraint, opts)
  841. apply_metadata_schema = namespace_metadata.schema
  842. return (
  843. "%(prefix)sForeignKeyConstraint([%(cols)s], "
  844. "[%(refcols)s], %(args)s)"
  845. % {
  846. "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
  847. "cols": ", ".join(
  848. repr(_ident(f.parent.name)) for f in constraint.elements
  849. ),
  850. "refcols": ", ".join(
  851. repr(_fk_colspec(f, apply_metadata_schema, namespace_metadata))
  852. for f in constraint.elements
  853. ),
  854. "args": ", ".join(
  855. ["%s=%s" % (kwname, val) for kwname, val in opts]
  856. ),
  857. }
  858. )
  859. @_constraint_renderers.dispatch_for(sa_schema.UniqueConstraint)
  860. def _render_unique_constraint(
  861. constraint: UniqueConstraint,
  862. autogen_context: AutogenContext,
  863. namespace_metadata: Optional[MetaData],
  864. ) -> str:
  865. rendered = _user_defined_render("unique", constraint, autogen_context)
  866. if rendered is not False:
  867. return rendered
  868. return _uq_constraint(constraint, autogen_context, False)
  869. @_constraint_renderers.dispatch_for(sa_schema.CheckConstraint)
  870. def _render_check_constraint(
  871. constraint: CheckConstraint,
  872. autogen_context: AutogenContext,
  873. namespace_metadata: Optional[MetaData],
  874. ) -> Optional[str]:
  875. rendered = _user_defined_render("check", constraint, autogen_context)
  876. if rendered is not False:
  877. return rendered
  878. # detect the constraint being part of
  879. # a parent type which is probably in the Table already.
  880. # ideally SQLAlchemy would give us more of a first class
  881. # way to detect this.
  882. if (
  883. constraint._create_rule
  884. and hasattr(constraint._create_rule, "target")
  885. and isinstance(
  886. constraint._create_rule.target,
  887. sqltypes.TypeEngine,
  888. )
  889. ):
  890. return None
  891. opts = []
  892. if constraint.name:
  893. opts.append(
  894. ("name", repr(_render_gen_name(autogen_context, constraint.name)))
  895. )
  896. return "%(prefix)sCheckConstraint(%(sqltext)s%(opts)s)" % {
  897. "prefix": _sqlalchemy_autogenerate_prefix(autogen_context),
  898. "opts": ", " + (", ".join("%s=%s" % (k, v) for k, v in opts))
  899. if opts
  900. else "",
  901. "sqltext": _render_potential_expr(
  902. constraint.sqltext, autogen_context, wrap_in_text=False
  903. ),
  904. }
  905. @renderers.dispatch_for(ops.ExecuteSQLOp)
  906. def _execute_sql(autogen_context: AutogenContext, op: ops.ExecuteSQLOp) -> str:
  907. if not isinstance(op.sqltext, str):
  908. raise NotImplementedError(
  909. "Autogenerate rendering of SQL Expression language constructs "
  910. "not supported here; please use a plain SQL string"
  911. )
  912. return "op.execute(%r)" % op.sqltext
  913. renderers = default_renderers.branch()