Add support of IF [NOT] EXISTS for ADD/DROP COLUMN in Postgresql

Added :paramref:`.Operations.add_column.if_not_exists` and
:paramref:`.Operations.drop_column.if_exists` to render ``IF [NOT] EXISTS``
for ``ADD COLUMN`` and ``DROP COLUMN`` operations, a feature available on
some database backends such as PostgreSQL, MariaDB, as well as third party
backends.  The parameters also support autogenerate rendering allowing them
to be turned on via a custom :class:`.Rewriter`. Pull request courtesy of
Louis-Amaury Chaib (@lachaib).

Fixes: #1626
Closes: #1627
Pull-request: #1627
Pull-request-sha: c503b04

Change-Id: I4d55f65f072b9f03698b2f45f066872b5c3e8c58

Author: lachaib

alembic/autogenerate/render.py

@@ -457,40 +457,56 @@ def _drop_constraint(
 
 @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)"
     else:
         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 % {
         "prefix": _alembic_autogenerate_prefix(autogen_context),
         "tname": _ident(tname),
         "cname": _ident(column_name),
         "schema": _ident(schema),
+        "if_exists": if_exists,
     }
     return text
 

alembic/ddl/base.py

@@ -154,17 +154,24 @@ def __init__(
         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):
@@ -189,15 +196,19 @@ def visit_rename_table(
 def visit_add_column(element: AddColumn, compiler: DDLCompiler, **kw) -> str:
     return "%s %s" % (
         alter_table(compiler, element.table_name, element.schema),
-        add_column(compiler, element.column, **kw),
+        add_column(
+            compiler, element.column, if_not_exists=element.if_not_exists, **kw
+        ),
     )
 
 
 @compiles(DropColumn)
 def visit_drop_column(element: DropColumn, compiler: DDLCompiler, **kw) -> str:
     return "%s %s" % (
         alter_table(compiler, element.table_name, element.schema),
-        drop_column(compiler, element.column.name, **kw),
+        drop_column(
+            compiler, element.column.name, if_exists=element.if_exists, **kw
+        ),
     )
 
 
@@ -320,16 +331,29 @@ def alter_table(
     return "ALTER TABLE %s" % format_table_name(compiler, name, schema)
 
 
-def drop_column(compiler: DDLCompiler, name: str, **kw) -> str:
-    return "DROP COLUMN %s" % format_column_name(compiler, name)
+def drop_column(
+    compiler: DDLCompiler, name: str, if_exists: Optional[bool] = None, **kw
+) -> str:
+    return "DROP COLUMN %s%s" % (
+        "IF EXISTS " if if_exists else "",
+        format_column_name(compiler, name),
+    )
 
 
 def alter_column(compiler: DDLCompiler, name: str) -> str:
     return "ALTER COLUMN %s" % format_column_name(compiler, name)
 
 
-def add_column(compiler: DDLCompiler, column: Column[Any], **kw) -> str:
-    text = "ADD COLUMN %s" % compiler.get_column_specification(column, **kw)
+def add_column(
+    compiler: DDLCompiler,
+    column: Column[Any],
+    if_not_exists: Optional[bool] = None,
+    **kw,
+) -> str:
+    text = "ADD COLUMN %s%s" % (
+        "IF NOT EXISTS " if if_not_exists else "",
+        compiler.get_column_specification(column, **kw),
+    )
 
     const = " ".join(
         compiler.process(constraint) for constraint in column.constraints

alembic/ddl/impl.py

@@ -256,6 +256,7 @@ def alter_column(
         self,
         table_name: str,
         column_name: str,
+        *,
         nullable: Optional[bool] = None,
         server_default: Optional[
             Union[_ServerDefault, Literal[False]]
@@ -370,18 +371,33 @@ def add_column(
         self,
         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):

alembic/ddl/mssql.py

@@ -83,10 +83,11 @@ def emit_commit(self) -> None:
         if self.as_sql and self.batch_separator:
             self.static_output(self.batch_separator)
 
-    def alter_column(  # type:ignore[override]
+    def alter_column(
         self,
         table_name: str,
         column_name: str,
+        *,
         nullable: Optional[bool] = None,
         server_default: Optional[
             Union[_ServerDefault, Literal[False]]
@@ -202,6 +203,7 @@ def drop_column(
         self,
         table_name: str,
         column: Column[Any],
+        *,
         schema: Optional[str] = None,
         **kw,
     ) -> None:

alembic/ddl/mysql.py

@@ -47,10 +47,11 @@ class MySQLImpl(DefaultImpl):
     )
     type_arg_extract = [r"character set ([\w\-_]+)", r"collate ([\w\-_]+)"]
 
-    def alter_column(  # type:ignore[override]
+    def alter_column(
         self,
         table_name: str,
         column_name: str,
+        *,
         nullable: Optional[bool] = None,
         server_default: Optional[
             Union[_ServerDefault, Literal[False]]

alembic/ddl/postgresql.py

@@ -52,6 +52,7 @@
 from ..util import sqla_compat
 from ..util.sqla_compat import compiles
 
+
 if TYPE_CHECKING:
     from typing import Literal
 
@@ -148,10 +149,11 @@ def compare_server_default(
             select(literal_column(conn_col_default) == metadata_default)
         )
 
-    def alter_column(  # type:ignore[override]
+    def alter_column(
         self,
         table_name: str,
         column_name: str,
+        *,
         nullable: Optional[bool] = None,
         server_default: Optional[
             Union[_ServerDefault, Literal[False]]

alembic/op.pyi

@@ -60,7 +60,11 @@ _C = TypeVar("_C", bound=Callable[..., Any])
 ### 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.
@@ -137,6 +141,10 @@ def add_column(
      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.16.0
 
     """
 
@@ -927,6 +935,11 @@ def drop_column(
      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.16.0
+
     :param mssql_drop_check: Optional boolean.  When ``True``, on
      Microsoft SQL Server only, first
      drop the CHECK constraint on the column using a
@@ -948,7 +961,6 @@ def drop_column(
      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(
@@ -972,7 +984,7 @@ def drop_constraint(
     :param if_exists: If True, adds IF EXISTS operator when
      dropping the constraint
 
-     .. versionadded:: 1.15.3
+     .. versionadded:: 1.16.0
 
     """
 

alembic/operations/base.py

@@ -618,6 +618,7 @@ def add_column(
             column: Column[Any],
             *,
             schema: Optional[str] = None,
+            if_not_exists: Optional[bool] = None,
         ) -> None:
             """Issue an "add column" instruction using the current
             migration context.
@@ -694,6 +695,10 @@ def add_column(
              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.16.0
 
             """  # noqa: E501
             ...
@@ -1361,6 +1366,11 @@ def drop_column(
              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.16.0
+
             :param mssql_drop_check: Optional boolean.  When ``True``, on
              Microsoft SQL Server only, first
              drop the CHECK constraint on the column using a
@@ -1382,7 +1392,6 @@ def drop_column(
              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
             ...
 
@@ -1408,7 +1417,7 @@ def drop_constraint(
             :param if_exists: If True, adds IF EXISTS operator when
              dropping the constraint
 
-             .. versionadded:: 1.15.3
+             .. versionadded:: 1.16.0
 
             """  # noqa: E501
             ...
@@ -1651,6 +1660,7 @@ def add_column(
             *,
             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.