Support if_exists and if_not_exists on create/drop table commands

Added support for :paramref:`.Operations.create_table.if_not_exists` and :paramref:`.Operations.drop_table.if_exists`, adding similar functionality to render IF [NOT] EXISTS for table operations in a similar way as with indexes. Pull request courtesy Aaron Griffin. Fixes: #1520 Closes: #1521 Pull-request: sqlalchemy/alembic#1521 Pull-request-sha: 469be01c6b5f9f42dc26017040a6fc54c4caef54 Change-Id: I5dcf44d9e906cdb84c32c4bfb6a1c63cde6324fd
zzzeek · Sep 13, 2024 · 8fae3e1 · 8fae3e1
1 parent 9d6e212
commit 8fae3e1
Show file tree

Hide file tree

Showing 7 changed files with 101 additions and 12 deletions.
diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py
@@ -362,11 +362,11 @@ def rename_table(
             base.RenameTable(old_table_name, new_table_name, schema=schema)
         )
 
-    def create_table(self, table: Table) -> None:
+    def create_table(self, table: Table, **kw: Any) -> None:
         table.dispatch.before_create(
             table, self.connection, checkfirst=False, _ddl_runner=self
         )
-        self._exec(schema.CreateTable(table))
+        self._exec(schema.CreateTable(table, **kw))
         table.dispatch.after_create(
             table, self.connection, checkfirst=False, _ddl_runner=self
         )
@@ -385,11 +385,11 @@ def create_table(self, table: Table) -> None:
             if comment and with_comment:
                 self.create_column_comment(column)
 
-    def drop_table(self, table: Table) -> None:
+    def drop_table(self, table: Table, **kw: Any) -> None:
         table.dispatch.before_drop(
             table, self.connection, checkfirst=False, _ddl_runner=self
         )
-        self._exec(schema.DropTable(table))
+        self._exec(schema.DropTable(table, **kw))
         table.dispatch.after_drop(
             table, self.connection, checkfirst=False, _ddl_runner=self
         )

diff --git a/alembic/op.pyi b/alembic/op.pyi
@@ -747,7 +747,12 @@ def create_primary_key(
 
     """
 
-def create_table(table_name: str, *columns: SchemaItem, **kw: Any) -> Table:
+def create_table(
+    table_name: str,
+    *columns: SchemaItem,
+    if_not_exists: Optional[bool] = None,
+    **kw: Any,
+) -> Table:
     r"""Issue a "create table" instruction using the current migration
     context.
 
@@ -818,6 +823,10 @@ def create_table(table_name: str, *columns: SchemaItem, **kw: Any) -> Table:
      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 table.
+
+     .. versionadded:: 1.13.3
     :param \**kw: Other keyword arguments are passed to the underlying
      :class:`sqlalchemy.schema.Table` object created for the command.
 
@@ -998,7 +1007,11 @@ def drop_index(
     """
 
 def drop_table(
-    table_name: str, *, schema: Optional[str] = None, **kw: Any
+    table_name: str,
+    *,
+    schema: Optional[str] = None,
+    if_exists: Optional[bool] = None,
+    **kw: Any,
 ) -> None:
     r"""Issue a "drop table" instruction using the current
     migration context.
@@ -1013,6 +1026,10 @@ def drop_table(
      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 table.
+
+     .. versionadded:: 1.13.3
     :param \**kw: Other keyword arguments are passed to the underlying
      :class:`sqlalchemy.schema.Table` object created for the command.
 

diff --git a/alembic/operations/base.py b/alembic/operations/base.py
@@ -1175,7 +1175,11 @@ def create_primary_key(
             ...
 
         def create_table(
-            self, table_name: str, *columns: SchemaItem, **kw: Any
+            self,
+            table_name: str,
+            *columns: SchemaItem,
+            if_not_exists: Optional[bool] = None,
+            **kw: Any,
         ) -> Table:
             r"""Issue a "create table" instruction using the current migration
             context.
@@ -1247,6 +1251,10 @@ def create_table(
              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 table.
+
+             .. versionadded:: 1.13.3
             :param \**kw: Other keyword arguments are passed to the underlying
              :class:`sqlalchemy.schema.Table` object created for the command.
 
@@ -1438,7 +1446,12 @@ def drop_index(
             ...
 
         def drop_table(
-            self, table_name: str, *, schema: Optional[str] = None, **kw: Any
+            self,
+            table_name: str,
+            *,
+            schema: Optional[str] = None,
+            if_exists: Optional[bool] = None,
+            **kw: Any,
         ) -> None:
             r"""Issue a "drop table" instruction using the current
             migration context.
@@ -1453,6 +1466,10 @@ def drop_table(
              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 table.
+
+             .. versionadded:: 1.13.3
             :param \**kw: Other keyword arguments are passed to the underlying
              :class:`sqlalchemy.schema.Table` object created for the command.
 

diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
@@ -1159,13 +1159,15 @@ def __init__(
         columns: Sequence[SchemaItem],
         *,
         schema: Optional[str] = None,
+        if_not_exists: Optional[bool] = None,
         _namespace_metadata: Optional[MetaData] = None,
         _constraints_included: bool = False,
         **kw: Any,
     ) -> None:
         self.table_name = table_name
         self.columns = columns
         self.schema = schema
+        self.if_not_exists = if_not_exists
         self.info = kw.pop("info", {})
         self.comment = kw.pop("comment", None)
         self.prefixes = kw.pop("prefixes", None)
@@ -1228,6 +1230,7 @@ def create_table(
         operations: Operations,
         table_name: str,
         *columns: SchemaItem,
+        if_not_exists: Optional[bool] = None,
         **kw: Any,
     ) -> Table:
         r"""Issue a "create table" instruction using the current migration
@@ -1300,14 +1303,18 @@ def create_table(
          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 table.
+
+         .. versionadded:: 1.13.3
         :param \**kw: Other keyword arguments are passed to the underlying
          :class:`sqlalchemy.schema.Table` object created for the command.
 
         :return: the :class:`~sqlalchemy.schema.Table` object corresponding
          to the parameters given.
 
         """
-        op = cls(table_name, columns, **kw)
+        op = cls(table_name, columns, if_not_exists=if_not_exists, **kw)
         return operations.invoke(op)
 
 
@@ -1320,11 +1327,13 @@ def __init__(
         table_name: str,
         *,
         schema: Optional[str] = None,
+        if_exists: Optional[bool] = None,
         table_kw: Optional[MutableMapping[Any, Any]] = None,
         _reverse: Optional[CreateTableOp] = None,
     ) -> None:
         self.table_name = table_name
         self.schema = schema
+        self.if_exists = if_exists
         self.table_kw = table_kw or {}
         self.comment = self.table_kw.pop("comment", None)
         self.info = self.table_kw.pop("info", None)
@@ -1385,6 +1394,7 @@ def drop_table(
         table_name: str,
         *,
         schema: Optional[str] = None,
+        if_exists: Optional[bool] = None,
         **kw: Any,
     ) -> None:
         r"""Issue a "drop table" instruction using the current
@@ -1400,11 +1410,15 @@ def drop_table(
          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 table.
+
+         .. versionadded:: 1.13.3
         :param \**kw: Other keyword arguments are passed to the underlying
          :class:`sqlalchemy.schema.Table` object created for the command.
 
         """
-        op = cls(table_name, schema=schema, table_kw=kw)
+        op = cls(table_name, schema=schema, if_exists=if_exists, table_kw=kw)
         operations.invoke(op)
 
 

diff --git a/alembic/operations/toimpl.py b/alembic/operations/toimpl.py
@@ -79,8 +79,14 @@ def _count_constraint(constraint):
 
 @Operations.implementation_for(ops.DropTableOp)
 def drop_table(operations: "Operations", operation: "ops.DropTableOp") -> None:
+    kw = {}
+    if operation.if_exists is not None:
+        if not sqla_14:
+            raise NotImplementedError("SQLAlchemy 1.4+ required")
+
+        kw["if_exists"] = operation.if_exists
     operations.impl.drop_table(
-        operation.to_table(operations.migration_context)
+        operation.to_table(operations.migration_context), **kw
     )
 
 
@@ -127,8 +133,14 @@ def drop_index(operations: "Operations", operation: "ops.DropIndexOp") -> None:
 def create_table(
     operations: "Operations", operation: "ops.CreateTableOp"
 ) -> "Table":
+    kw = {}
+    if operation.if_not_exists is not None:
+        if not sqla_14:
+            raise NotImplementedError("SQLAlchemy 1.4+ required")
+
+        kw["if_not_exists"] = operation.if_not_exists
     table = operation.to_table(operations.migration_context)
-    operations.impl.create_table(table)
+    operations.impl.create_table(table, **kw)
     return table
 
 

diff --git a/docs/build/unreleased/1520.rst b/docs/build/unreleased/1520.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: usecase, operations
+    :tickets: 1520
+
+    Added support for :paramref:`.Operations.create_table.if_not_exists` and
+    :paramref:`.Operations.drop_table.if_exists`, adding similar functionality
+    to render IF [NOT] EXISTS for table operations in a similar way as with
+    indexes. Pull request courtesy Aaron Griffin.
+
diff --git a/tests/test_op.py b/tests/test_op.py
@@ -907,6 +907,12 @@ def test_drop_table_schema(self):
         op.drop_table("tb_test", schema="foo")
         context.assert_("DROP TABLE foo.tb_test")
 
+    @config.requirements.sqlalchemy_14
+    def test_drop_table_if_exists(self):
+        context = op_fixture()
+        op.drop_table("tb_test", if_exists=True)
+        context.assert_("DROP TABLE IF EXISTS tb_test")
+
     def test_create_table_selfref(self):
         context = op_fixture()
         op.create_table(
@@ -1079,6 +1085,20 @@ def test_create_table_two_fk(self):
             "FOREIGN KEY(foo_bar) REFERENCES foo (bar))"
         )
 
+    @config.requirements.sqlalchemy_14
+    def test_create_table_if_not_exists(self):
+        context = op_fixture()
+        op.create_table(
+            "some_table",
+            Column("id", Integer, primary_key=True),
+            if_not_exists=True,
+        )
+        context.assert_(
+            "CREATE TABLE IF NOT EXISTS some_table ("
+            "id INTEGER NOT NULL, "
+            "PRIMARY KEY (id))"
+        )
+
     def test_execute_delete(self):
         context = op_fixture()