@renderers.dispatch_for(ops.AddColumnOp)
def _add_column(autogen_context: AutogenContext, op: ops.AddColumnOp) -> str:
- schema, tname, column = op.schema, op.table_name, op.column
+ schema, tname, column, if_not_exists = (
+ op.schema,
+ op.table_name,
+ op.column,
+ op.if_not_exists,
+ )
if autogen_context._has_batch:
template = "%(prefix)sadd_column(%(column)s)"
else:
template = "%(prefix)sadd_column(%(tname)r, %(column)s"
if schema:
template += ", schema=%(schema)r"
+ if if_not_exists is not None:
+ template += ", if_not_exists=%(if_not_exists)r"
template += ")"
text = template % {
"prefix": _alembic_autogenerate_prefix(autogen_context),
"tname": tname,
"column": _render_column(column, autogen_context),
"schema": schema,
+ "if_not_exists": if_not_exists,
}
return text
@renderers.dispatch_for(ops.DropColumnOp)
def _drop_column(autogen_context: AutogenContext, op: ops.DropColumnOp) -> str:
- schema, tname, column_name = op.schema, op.table_name, op.column_name
+ schema, tname, column_name, if_exists = (
+ op.schema,
+ op.table_name,
+ op.column_name,
+ op.if_exists,
+ )
if autogen_context._has_batch:
template = "%(prefix)sdrop_column(%(cname)r)"
template = "%(prefix)sdrop_column(%(tname)r, %(cname)r"
if schema:
template += ", schema=%(schema)r"
+ if if_exists is not None:
+ template += ", if_exists=%(if_exists)r"
template += ")"
text = template % {
"tname": _ident(tname),
"cname": _ident(column_name),
"schema": _ident(schema),
+ "if_exists": if_exists,
}
return text
name: str,
column: Column[Any],
schema: Optional[Union[quoted_name, str]] = None,
+ if_not_exists: Optional[bool] = None,
) -> None:
super().__init__(name, schema=schema)
self.column = column
+ self.if_not_exists = if_not_exists
class DropColumn(AlterTable):
def __init__(
- self, name: str, column: Column[Any], schema: Optional[str] = None
+ self,
+ name: str,
+ column: Column[Any],
+ schema: Optional[str] = None,
+ if_exists: Optional[bool] = None,
) -> None:
super().__init__(name, schema=schema)
self.column = column
+ self.if_exists = if_exists
class ColumnComment(AlterColumn):
table_name: str,
column: Column[Any],
schema: Optional[Union[str, quoted_name]] = None,
+ if_not_exists: Optional[bool] = None,
) -> None:
- self._exec(base.AddColumn(table_name, column, schema=schema))
+ self._exec(
+ base.AddColumn(
+ table_name,
+ column,
+ schema=schema,
+ if_not_exists=if_not_exists,
+ )
+ )
def drop_column(
self,
table_name: str,
column: Column[Any],
schema: Optional[str] = None,
+ if_exists: Optional[bool] = None,
**kw,
) -> None:
- self._exec(base.DropColumn(table_name, column, schema=schema))
+ self._exec(
+ base.DropColumn(
+ table_name, column, schema=schema, if_exists=if_exists
+ )
+ )
def add_constraint(self, const: Any) -> None:
if const._create_rule is None or const._create_rule(self):
table_name: str,
column: Column[Any],
schema: Optional[str] = None,
+ if_exists: Optional[bool] = None,
**kw,
) -> None:
drop_default = kw.pop("mssql_drop_default", False)
drop_fks = kw.pop("mssql_drop_foreign_key", False)
if drop_fks:
self._exec(_ExecDropFKConstraint(table_name, column, schema))
- super().drop_column(table_name, column, schema=schema, **kw)
+ super().drop_column(
+ table_name, column, schema=schema, if_exists=if_exists, **kw
+ )
def compare_server_default(
self,
from sqlalchemy import schema
from sqlalchemy import types as sqltypes
+from .base import AddColumn
from .base import alter_table
from .base import AlterColumn
from .base import ColumnDefault
from .base import ColumnName
from .base import ColumnNullable
from .base import ColumnType
+from .base import DropColumn
from .base import format_column_name
from .base import format_server_default
from .impl import DefaultImpl
)
+@compiles(AddColumn, "mysql", "mariadb")
+def _mysql_add_column(
+ element: AddColumn, compiler: MySQLDDLCompiler, **kw
+) -> str:
+
+ return "%s ADD COLUMN %s%s" % (
+ alter_table(compiler, element.table_name, element.schema),
+ "IF NOT EXISTS " if element.if_not_exists else "",
+ compiler.get_column_specification(element.column, **kw),
+ )
+
+
+@compiles(DropColumn, "mysql", "mariadb")
+def _mysql_drop_column(
+ element: DropColumn, compiler: MySQLDDLCompiler, **kw
+) -> str:
+ return "%s DROP COLUMN %s%s" % (
+ alter_table(compiler, element.table_name, element.schema),
+ "IF EXISTS " if element.if_exists else "",
+ format_column_name(compiler, element.column.name),
+ )
+
+
@compiles(MySQLModifyColumn, "mysql", "mariadb")
def _mysql_modify_column(
element: MySQLModifyColumn, compiler: MySQLDDLCompiler, **kw
from sqlalchemy.dialects.postgresql import ExcludeConstraint
from sqlalchemy.dialects.postgresql import INTEGER
from sqlalchemy.schema import CreateIndex
+from sqlalchemy.sql.compiler import DDLCompiler
from sqlalchemy.sql.elements import ColumnClause
from sqlalchemy.sql.elements import TextClause
from sqlalchemy.sql.functions import FunctionElement
from sqlalchemy.types import NULLTYPE
+from .base import AddColumn
from .base import alter_column
from .base import alter_table
from .base import AlterColumn
from .base import ColumnComment
+from .base import DropColumn
from .base import format_column_name
from .base import format_table_name
from .base import format_type
from ..util import sqla_compat
from ..util.sqla_compat import compiles
+
if TYPE_CHECKING:
from typing import Literal
self.using = using
+@compiles(AddColumn, "postgresql")
+def visit_add_column(element: AddColumn, compiler: PGDDLCompiler, **kw) -> str:
+ return "%s %s" % (
+ alter_table(compiler, element.table_name, element.schema),
+ add_column(
+ compiler,
+ element.column,
+ if_not_exists=element.if_not_exists,
+ **kw,
+ ),
+ )
+
+
+@compiles(DropColumn, "postgresql")
+def visit_drop_column(
+ element: DropColumn, compiler: PGDDLCompiler, **kw
+) -> str:
+ return "%s %s" % (
+ alter_table(compiler, element.table_name, element.schema),
+ drop_column(
+ compiler,
+ element.column.name,
+ if_exists=element.if_exists,
+ **kw,
+ ),
+ )
+
+
@compiles(RenameTable, "postgresql")
def visit_rename_table(
element: RenameTable, compiler: PGDDLCompiler, **kw
autogen_context,
wrap_in_element=isinstance(value, (TextClause, FunctionElement)),
)
+
+
+def add_column(
+ compiler: DDLCompiler,
+ column: Column[Any],
+ *,
+ if_not_exists: Optional[bool] = None,
+ **kw,
+) -> str:
+ text = "ADD COLUMN "
+ if if_not_exists:
+ text += "IF NOT EXISTS "
+
+ text += compiler.get_column_specification(column, **kw)
+
+ const = " ".join(
+ compiler.process(constraint) for constraint in column.constraints
+ )
+ if const:
+ text += " " + const
+
+ return text
+
+
+def drop_column(
+ compiler: DDLCompiler,
+ name: str,
+ *,
+ if_exists: Optional[bool] = None,
+ **kw,
+) -> str:
+ text = "DROP COLUMN "
+ if if_exists:
+ text += "IF EXISTS "
+ text += format_column_name(compiler, name)
+ return text
### end imports ###
def add_column(
- table_name: str, column: Column[Any], *, schema: Optional[str] = None
+ table_name: str,
+ column: Column[Any],
+ *,
+ schema: Optional[str] = None,
+ if_not_exists: Optional[bool] = None,
) -> None:
"""Issue an "add column" instruction using the current
migration context.
quoting of the schema outside of the default behavior, use
the SQLAlchemy construct
:class:`~sqlalchemy.sql.elements.quoted_name`.
+ :param if_not_exists: If True, adds IF NOT EXISTS operator
+ when creating the new column for compatible dialects
+
+ .. versionadded:: 1.15.3
"""
quoting of the schema outside of the default behavior, use
the SQLAlchemy construct
:class:`~sqlalchemy.sql.elements.quoted_name`.
+ :param if_exists: If True, adds IF EXISTS operator when
+ dropping the new column for compatible dialects
+
+ .. versionadded:: 1.15.3
+
:param mssql_drop_check: Optional boolean. When ``True``, on
Microsoft SQL Server only, first
drop the CHECK constraint on the column using a
then exec's a separate DROP CONSTRAINT for that default. Only
works if the column has exactly one FK constraint which refers to
it, at the moment.
-
"""
def drop_constraint(
column: Column[Any],
*,
schema: Optional[str] = None,
+ if_not_exists: Optional[bool] = None,
) -> None:
"""Issue an "add column" instruction using the current
migration context.
quoting of the schema outside of the default behavior, use
the SQLAlchemy construct
:class:`~sqlalchemy.sql.elements.quoted_name`.
+ :param if_not_exists: If True, adds IF NOT EXISTS operator
+ when creating the new column for compatible dialects
+
+ .. versionadded:: 1.15.3
""" # noqa: E501
...
quoting of the schema outside of the default behavior, use
the SQLAlchemy construct
:class:`~sqlalchemy.sql.elements.quoted_name`.
+ :param if_exists: If True, adds IF EXISTS operator when
+ dropping the new column for compatible dialects
+
+ .. versionadded:: 1.15.3
+
:param mssql_drop_check: Optional boolean. When ``True``, on
Microsoft SQL Server only, first
drop the CHECK constraint on the column using a
then exec's a separate DROP CONSTRAINT for that default. Only
works if the column has exactly one FK constraint which refers to
it, at the moment.
-
""" # noqa: E501
...
*,
insert_before: Optional[str] = None,
insert_after: Optional[str] = None,
+ if_not_exists: Optional[bool] = None,
) -> None:
"""Issue an "add column" instruction using the current
batch migration context.
column: Column[Any],
*,
schema: Optional[str] = None,
+ if_not_exists: Optional[bool] = None,
**kw: Any,
) -> None:
super().__init__(table_name, schema=schema)
self.column = column
+ self.if_not_exists = if_not_exists
self.kw = kw
def reverse(self) -> DropColumnOp:
- return DropColumnOp.from_column_and_tablename(
+ op = DropColumnOp.from_column_and_tablename(
self.schema, self.table_name, self.column
)
+ op.if_exists = self.if_not_exists
+ return op
def to_diff_tuple(
self,
column: Column[Any],
*,
schema: Optional[str] = None,
+ if_not_exists: Optional[bool] = None,
) -> None:
"""Issue an "add column" instruction using the current
migration context.
quoting of the schema outside of the default behavior, use
the SQLAlchemy construct
:class:`~sqlalchemy.sql.elements.quoted_name`.
+ :param if_not_exists: If True, adds IF NOT EXISTS operator
+ when creating the new column for compatible dialects
+
+ .. versionadded:: 1.15.3
"""
- op = cls(table_name, column, schema=schema)
+ op = cls(
+ table_name,
+ column,
+ schema=schema,
+ if_not_exists=if_not_exists,
+ )
return operations.invoke(op)
@classmethod
*,
insert_before: Optional[str] = None,
insert_after: Optional[str] = None,
+ if_not_exists: Optional[bool] = None,
) -> None:
"""Issue an "add column" instruction using the current
batch migration context.
operations.impl.table_name,
column,
schema=operations.impl.schema,
+ if_not_exists=if_not_exists,
**kw,
)
return operations.invoke(op)
column_name: str,
*,
schema: Optional[str] = None,
+ if_exists: Optional[bool] = None,
_reverse: Optional[AddColumnOp] = None,
**kw: Any,
) -> None:
super().__init__(table_name, schema=schema)
self.column_name = column_name
self.kw = kw
+ self.if_exists = if_exists
self._reverse = _reverse
def to_diff_tuple(
"original column is not present"
)
- return AddColumnOp.from_column_and_tablename(
+ op = AddColumnOp.from_column_and_tablename(
self.schema, self.table_name, self._reverse.column
)
+ op.if_not_exists = self.if_exists
+ return op
@classmethod
def from_column_and_tablename(
quoting of the schema outside of the default behavior, use
the SQLAlchemy construct
:class:`~sqlalchemy.sql.elements.quoted_name`.
+ :param if_exists: If True, adds IF EXISTS operator when
+ dropping the new column for compatible dialects
+
+ .. versionadded:: 1.15.3
+
:param mssql_drop_check: Optional boolean. When ``True``, on
Microsoft SQL Server only, first
drop the CHECK constraint on the column using a
then exec's a separate DROP CONSTRAINT for that default. Only
works if the column has exactly one FK constraint which refers to
it, at the moment.
-
"""
op = cls(table_name, column_name, schema=schema, **kw)
) -> None:
column = operation.to_column(operations.migration_context)
operations.impl.drop_column(
- operation.table_name, column, schema=operation.schema, **operation.kw
+ operation.table_name,
+ column,
+ schema=operation.schema,
+ if_exists=operation.if_exists,
+ **operation.kw,
)
column = _copy(column)
t = operations.schema_obj.table(table_name, column, schema=schema)
- operations.impl.add_column(table_name, column, schema=schema, **kw)
+ operations.impl.add_column(
+ table_name,
+ column,
+ schema=schema,
+ if_not_exists=operation.if_not_exists,
+ **kw,
+ )
for constraint in t.constraints:
if not isinstance(constraint, sa_schema.PrimaryKeyConstraint):
@overload
def tuple_rev_as_scalar(
- rev: Union[Tuple[_T, ...], List[_T]]
+ rev: Union[Tuple[_T, ...], List[_T]],
) -> Union[_T, Tuple[_T, ...], List[_T]]: ...
)
# ### end Alembic commands ###
+
+.. _cookbook_add_if_exists:
+
+
+Add IF [NOT] EXISTS to CREATE/DROP operations
+=============================================
+
+Sometimes CREATE/DROP operations take too long during a production deployment
+and it is preferable to apply them offline,
+and still keep alembic migrations aligned and/or for test environments.
+
+Using the rewriter makes it possible::
+
+ from alembic.operations import ops
+ from alembic.autogenerate import rewriter
+
+ writer = rewriter.Rewriter()
+
+ @writer.rewrites(ops.CreateTableOp)
+ @writer.rewrites(ops.CreateIndexOp)
+ def add_if_not_exists(context, revision, op):
+ op.if_not_exists = True
+ return op
+
+ @writer.rewrites(ops.DropTableOp)
+ @writer.rewrites(ops.DropIndexOp)
+ def add_if_exists(context, revision, op):
+ op.if_exists = True
+ return op
+
+Same operation is possible for ADD/DROP COLUMN on postgresql/mariadb::
+
+ @writer.rewrites(ops.AddColumnOp)
+ def add_column_if_not_exists(context, revision, op):
+ op.if_not_exists = True
+ return op
+
+ @writer.rewrites(ops.DropColumnOp)
+ def drop_column_if_not_exists(context, revision, op):
+ op.if_exists = True
+ return op
+
+
+
Don't emit CREATE TABLE statements for Views
============================================
--- /dev/null
+.. change::
+ :tags: usecase, autogenerate, postgresql
+
+ Render `if_exists`` and `if_not_exists`` parameters in
+ :class:`.AddColumnOp` and :class:`.DropColumnOp`, in an autogenerate context.
+ 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).
\ No newline at end of file
"server_default='5', nullable=True, somedialect_foobar='option'))",
)
+ def test_render_add_column_if_not_exists(self):
+ op_obj = ops.AddColumnOp(
+ "foo",
+ Column("x", Integer, server_default="5", nullable=True),
+ if_not_exists=True,
+ )
+ eq_ignore_whitespace(
+ autogenerate.render_op_text(self.autogen_context, op_obj),
+ "op.add_column('foo', sa.Column('x', sa.Integer(), "
+ "server_default='5', nullable=True), "
+ "if_not_exists=True)",
+ )
+
def test_render_add_column_system(self):
# this would never actually happen since "system" columns
# can't be added in any case. However it will render as
"op.drop_column('bar', 'x', schema='foo')",
)
+ def test_render_drop_column_if_exists(self):
+ op_obj = ops.DropColumnOp.from_column_and_tablename(
+ None, "foo", Column("x", Integer, server_default="5")
+ )
+ op_obj.if_exists = True
+ eq_ignore_whitespace(
+ autogenerate.render_op_text(self.autogen_context, op_obj),
+ "op.drop_column('foo', 'x', if_exists=True)",
+ )
+
def test_render_quoted_server_default(self):
eq_(
autogenerate.render._render_server_default(
batch.impl.operations.impl.mock_calls,
[
mock.call.drop_column(
- "tname", self.mock_schema.Column(), schema=None
+ "tname",
+ self.mock_schema.Column(),
+ schema=None,
+ if_exists=None,
)
],
)
batch.add_column(column)
assert (
- mock.call.add_column("tname", column, schema=None)
+ mock.call.add_column(
+ "tname", column, schema=None, if_not_exists=None
+ )
in batch.impl.operations.impl.mock_calls
)
"ALTER TABLE t ADD COLUMN q INTEGER COMMENT 'This is a comment'"
)
+ def test_add_column_if_not_exists(self):
+ context = op_fixture("mysql")
+ op.add_column("t", Column("c", Integer), if_not_exists=True)
+ context.assert_("ALTER TABLE t ADD COLUMN IF NOT EXISTS c INTEGER")
+
+ def test_drop_column_if_exists(self):
+ context = op_fixture("mysql")
+ op.drop_column("t", "c", if_exists=True)
+ context.assert_("ALTER TABLE t DROP COLUMN IF EXISTS c")
+
def test_rename_column(self):
context = op_fixture("mysql")
op.alter_column(
op.create_index("i", "t", ["c1", "c2"], unique=False)
context.assert_("CREATE INDEX i ON t (c1, c2)")
- def test_create_index_postgresql_if_not_exists(self):
+ def test_create_index_if_not_exists(self):
context = op_fixture("postgresql")
op.create_index("i", "t", ["c1", "c2"], if_not_exists=True)
context.assert_("CREATE INDEX IF NOT EXISTS i ON t (c1, c2)")
op.drop_index("geocoded", postgresql_concurrently=True)
context.assert_("DROP INDEX CONCURRENTLY geocoded")
- def test_drop_index_postgresql_if_exists(self):
+ def test_drop_index_if_exists(self):
context = op_fixture("postgresql")
op.drop_index("geocoded", if_exists=True)
context.assert_("DROP INDEX IF EXISTS geocoded")
"ALTER TABLE t ALTER COLUMN c TYPE INTEGER USING c::integer"
)
+ def test_add_column_if_not_exists(self):
+ context = op_fixture("postgresql")
+ op.add_column("t", Column("c", Integer), if_not_exists=True)
+ context.assert_("ALTER TABLE t ADD COLUMN IF NOT EXISTS c INTEGER")
+
+ def test_drop_column_if_exists(self):
+ context = op_fixture("postgresql")
+ op.drop_column("t", "c", if_exists=True)
+ context.assert_("ALTER TABLE t DROP COLUMN IF EXISTS c")
+
def test_col_w_pk_is_serial(self):
context = op_fixture("postgresql")
op.add_column("some_table", Column("q", Integer, primary_key=True))
projects / thirdparty / sqlalchemy / alembic.git / commitdiff
