diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py index d2983923..25746889 100644 --- 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 index 83deac1e..92044469 100644 --- 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 index 27dd3b9e..9b52fa6f 100644 --- 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 index e6f1fb64..60b856a8 100644 --- a/alembic/operations/ops.py +++ b/alembic/operations/ops.py @@ -1159,6 +1159,7 @@ 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, @@ -1166,6 +1167,7 @@ def __init__( 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,6 +1303,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. @@ -1307,7 +1314,7 @@ def create_table( 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 index 4759f7fd..4b960049 100644 --- 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 new file mode 100644 index 00000000..1e780b7d --- /dev/null +++ b/docs/build/unreleased/1520.rst @@ -0,0 +1,5 @@ +.. change:: + :tags: usecase, operations + :tickets: 1520 + + Support if_exists and if_not_exists on create/drop table commands diff --git a/tests/test_op.py b/tests/test_op.py index 688799c9..cbb30ea5 100644 --- 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()