Skip to content

Commit

Permalink
Render if_not_exists option for CreateTableOp, CreateIndexOp, DropT…
Browse files Browse the repository at this point in the history
…ableOp and DropIndexOp

Render ``if_exists`` and ``if_not_exists`` parameters in
:class:`.CreateTableOp`, :class:`.CreateIndexOp`, :class:`.DropTableOp` and
:class:`.DropIndexOp` in an autogenerate context.  While Alembic does not
set these parameters during an autogenerate run, they can be enabled using
a custom :class:`.Rewriter` in the ``env.py`` file, where they will now be
part of the rendered Python code in revision files.  Pull request courtesy
of Louis-Amaury Chaib (@lachaib).

Closes: #1446
Pull-request: #1446
Pull-request-sha: 90c9735

Change-Id: I6b0a5ffaf7e2d1a0a1e1f1e80ed0ee168ae2bd09
  • Loading branch information
lachaib authored and zzzeek committed Sep 23, 2024
1 parent 8fae3e1 commit 3db374f
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 0 deletions.
11 changes: 11 additions & 0 deletions alembic/autogenerate/render.py
Original file line number Diff line number Diff line change
Expand Up @@ -279,6 +279,9 @@ def _add_table(autogen_context: AutogenContext, op: ops.CreateTableOp) -> str:
prefixes = ", ".join("'%s'" % p for p in table._prefixes)
text += ",\nprefixes=[%s]" % prefixes

if op.if_not_exists is not None:
text += ",\nif_not_exists=%r" % bool(op.if_not_exists)

text += "\n)"
return text

Expand All @@ -291,6 +294,10 @@ def _drop_table(autogen_context: AutogenContext, op: ops.DropTableOp) -> str:
}
if op.schema:
text += ", schema=%r" % _ident(op.schema)

if op.if_exists is not None:
text += ", if_exists=%r" % bool(op.if_exists)

text += ")"
return text

Expand Down Expand Up @@ -324,6 +331,8 @@ def _add_index(autogen_context: AutogenContext, op: ops.CreateIndexOp) -> str:
assert index.table is not None

opts = _render_dialect_kwargs_items(autogen_context, index)
if op.if_not_exists is not None:
opts.append("if_not_exists=%r" % bool(op.if_not_exists))
text = tmpl % {
"prefix": _alembic_autogenerate_prefix(autogen_context),
"name": _render_gen_name(autogen_context, index.name),
Expand Down Expand Up @@ -356,6 +365,8 @@ def _drop_index(autogen_context: AutogenContext, op: ops.DropIndexOp) -> str:
"table_name=%(table_name)r%(schema)s%(kwargs)s)"
)
opts = _render_dialect_kwargs_items(autogen_context, index)
if op.if_exists is not None:
opts.append("if_exists=%r" % bool(op.if_exists))
text = tmpl % {
"prefix": _alembic_autogenerate_prefix(autogen_context),
"name": _render_gen_name(autogen_context, op.index_name),
Expand Down
10 changes: 10 additions & 0 deletions docs/build/unreleased/1446.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
.. change::
:tags: usecase, autogenerate

Render ``if_exists`` and ``if_not_exists`` parameters in
:class:`.CreateTableOp`, :class:`.CreateIndexOp`, :class:`.DropTableOp` and
:class:`.DropIndexOp` in an autogenerate context. While Alembic does not
set these parameters during an autogenerate run, they can be enabled using
a custom :class:`.Rewriter` in the ``env.py`` file, where they will now be
part of the rendered Python code in revision files. Pull request courtesy
of Louis-Amaury Chaib (@lachaib).
50 changes: 50 additions & 0 deletions tests/test_autogen_render.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,20 @@ def test_render_add_index(self):
"['active', 'code'], unique=False)",
)

def test_render_add_index_if_not_exists(self):
"""
autogenerate.render._add_index
"""
t = self.table()
idx = Index("test_active_code_idx", t.c.active, t.c.code)
op_obj = ops.CreateIndexOp.from_index(idx)
op_obj.if_not_exists = True
eq_ignore_whitespace(
autogenerate.render_op_text(self.autogen_context, op_obj),
"op.create_index('test_active_code_idx', 'test', "
"['active', 'code'], unique=False, if_not_exists=True)",
)

@testing.emits_warning("Can't validate argument ")
def test_render_add_index_custom_kwarg(self):
t = self.table()
Expand Down Expand Up @@ -212,6 +226,20 @@ def test_drop_index(self):
"op.drop_index('test_active_code_idx', table_name='test')",
)

def test_drop_index_if_exists(self):
"""
autogenerate.render._drop_index
"""
t = self.table()
idx = Index("test_active_code_idx", t.c.active, t.c.code)
op_obj = ops.DropIndexOp.from_index(idx)
op_obj.if_exists = True
eq_ignore_whitespace(
autogenerate.render_op_text(self.autogen_context, op_obj),
"op.drop_index('test_active_code_idx', table_name='test', "
"if_exists=True)",
)

def test_drop_index_text(self):
"""
autogenerate.render._drop_index
Expand Down Expand Up @@ -989,6 +1017,19 @@ def test_render_addtl_args(self):
"mysql_engine='InnoDB',sqlite_autoincrement=True)",
)

def test_render_if_not_exists(self):
t = self.table()
op_obj = ops.CreateTableOp.from_table(t)
op_obj.if_not_exists = True
eq_ignore_whitespace(
autogenerate.render_op_text(self.autogen_context, op_obj),
"op.create_table('test',"
"sa.Column('id', sa.Integer(), nullable=False),"
"sa.Column('active', sa.Boolean(), nullable=True),"
"sa.Column('code', sa.String(length=255), nullable=True),"
"sa.PrimaryKeyConstraint('id'),if_not_exists=True)",
)

def test_render_drop_table(self):
op_obj = ops.DropTableOp.from_table(Table("sometable", MetaData()))
eq_ignore_whitespace(
Expand All @@ -1005,6 +1046,15 @@ def test_render_drop_table_w_schema(self):
"op.drop_table('sometable', schema='foo')",
)

def test_render_drop_table_if_exists(self):
t = self.table()
op_obj = ops.DropTableOp.from_table(t)
op_obj.if_exists = True
eq_ignore_whitespace(
autogenerate.render_op_text(self.autogen_context, op_obj),
"op.drop_table('test', if_exists=True)",
)

def test_render_table_no_implicit_check(self):
m = MetaData()
t = Table("test", m, Column("x", Boolean()))
Expand Down

0 comments on commit 3db374f

Please sign in to comment.