Skip to content

Commit

Permalink
Support if_exists and if_not_exists on create/drop table commands
Browse files Browse the repository at this point in the history
  • Loading branch information
agriffin-grow committed Aug 14, 2024
1 parent e210c24 commit 88c7b3f
Show file tree
Hide file tree
Showing 5 changed files with 71 additions and 10 deletions.
8 changes: 4 additions & 4 deletions alembic/ddl/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand All @@ -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
)
Expand Down
21 changes: 19 additions & 2 deletions alembic/operations/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 table.
.. versionadded:: 1.13.4
:param \**kw: Other keyword arguments are passed to the underlying
:class:`sqlalchemy.schema.Table` object created for the command.
Expand Down Expand Up @@ -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.
Expand All @@ -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.4
:param \**kw: Other keyword arguments are passed to the underlying
:class:`sqlalchemy.schema.Table` object created for the command.
Expand Down
18 changes: 16 additions & 2 deletions alembic/operations/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.4
: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)


Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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.4
: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)


Expand Down
16 changes: 14 additions & 2 deletions alembic/operations/toimpl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)


Expand Down Expand Up @@ -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


Expand Down
18 changes: 18 additions & 0 deletions tests/test_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -907,6 +907,11 @@ def test_drop_table_schema(self):
op.drop_table("tb_test", schema="foo")
context.assert_("DROP TABLE foo.tb_test")

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(
Expand Down Expand Up @@ -1079,6 +1084,19 @@ def test_create_table_two_fk(self):
"FOREIGN KEY(foo_bar) REFERENCES foo (bar))"
)

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

Expand Down

0 comments on commit 88c7b3f

Please sign in to comment.