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,17 +370,27 @@ 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,
         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/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,28 @@ 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 +874,35 @@ 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,8 @@ 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 databases that support it.
 
     """
 
@@ -910,6 +916,7 @@ def drop_column(
     column_name: str,
     *,
     schema: Optional[str] = None,
+    if_exists: Optional[bool] = None,
     **kw: Any,
 ) -> None:
     """Issue a "drop column" instruction using the current
@@ -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 column if the database supports it.
 
     """
 

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,8 @@ 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 databases that support it.
 
             """  # noqa: E501
             ...
@@ -1346,6 +1349,7 @@ def drop_column(
             column_name: str,
             *,
             schema: Optional[str] = None,
+            if_exists: Optional[bool] = None,
             **kw: Any,
         ) -> None:
             """Issue a "drop column" instruction using the current
@@ -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 column if the database supports it.
 
             """  # 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.
@@ -1839,7 +1846,13 @@ def create_unique_constraint(
             """  # noqa: E501
             ...
 
-        def drop_column(self, column_name: str, **kw: Any) -> None:
+        def drop_column(
+            self,
+            column_name: str,
+            *,
+            if_exists: Optional[bool] = None,
+            **kw: Any,
+        ) -> None:
             """Issue a "drop column" instruction using the current
             batch migration context.
 

alembic/operations/ops.py

@@ -2034,16 +2034,20 @@ def __init__(
         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,
@@ -2074,6 +2078,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.
@@ -2153,7 +2158,9 @@ def add_column(
 
         """
 
-        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
@@ -2200,12 +2207,14 @@ def __init__(
         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(
@@ -2225,9 +2234,11 @@ def reverse(self) -> AddColumnOp:
                 "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(

alembic/operations/toimpl.py

@@ -92,7 +92,11 @@ def drop_column(
 ) -> 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,
     )
 
 
@@ -167,7 +171,13 @@ def add_column(operations: "Operations", operation: "ops.AddColumnOp") -> None:
         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):