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

Fixes sqlalchemy#1626

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):

alembic/ddl/impl.py

@@ -370,8 +370,16 @@ def add_column(
         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,
@@ -380,7 +388,14 @@ def drop_column(
         schema: Optional[str] = None,
         **kw,
     ) -> None:
-        self._exec(base.DropColumn(table_name, column, schema=schema))
+        self._exec(
+            base.DropColumn(
+                table_name,
+                column,
+                schema=schema,
+                **kw,
+            )
+        )
 
     def add_constraint(self, const: Any) -> None:
         if const._create_rule is None or const._create_rule(self):

alembic/ddl/mysql.py

@@ -12,12 +12,14 @@
 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
@@ -391,6 +393,29 @@ def _mysql_alter_default(
     )
 
 
+@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

alembic/ddl/postgresql.py

@@ -27,15 +27,18 @@
 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
@@ -52,6 +55,7 @@
 from ..util import sqla_compat
 from ..util.sqla_compat import compiles
 
+
 if TYPE_CHECKING:
     from typing import Literal
 
@@ -512,6 +516,34 @@ def __init__(
         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
@@ -848,3 +880,39 @@ def _render_potential_column(
             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

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,9 @@ 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
+
 
     """
 
@@ -946,6 +953,8 @@ 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.
+    :param if_exists: If True, adds IF EXISTS operator when
+     dropping the new column for compatible dialects
 
     """
 

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,9 @@ 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
+
 
             """  # noqa: E501
             ...
@@ -1382,6 +1386,8 @@ 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.
+            :param if_exists: If True, adds IF EXISTS operator when
+             dropping the new column for compatible dialects
 
             """  # noqa: E501
             ...
@@ -1646,6 +1652,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.