diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py index 31667ef8..a6102d1e 100644 --- a/alembic/ddl/impl.py +++ b/alembic/ddl/impl.py @@ -379,8 +379,8 @@ def drop_table(self, table: Table) -> None: table, self.connection, checkfirst=False, _ddl_runner=self ) - def create_index(self, index: Index) -> None: - self._exec(schema.CreateIndex(index)) + def create_index(self, index: Index, **kw: Any) -> None: + self._exec(schema.CreateIndex(index, **kw)) def create_table_comment(self, table: Table) -> None: self._exec(schema.SetTableComment(table)) @@ -391,8 +391,8 @@ def drop_table_comment(self, table: Table) -> None: def create_column_comment(self, column: ColumnElement[Any]) -> None: self._exec(schema.SetColumnComment(column)) - def drop_index(self, index: Index) -> None: - self._exec(schema.DropIndex(index)) + def drop_index(self, index: Index, **kw: Any) -> None: + self._exec(schema.DropIndex(index, **kw)) def bulk_insert( self, diff --git a/alembic/ddl/mssql.py b/alembic/ddl/mssql.py index 56dd12c3..e77fd1b5 100644 --- a/alembic/ddl/mssql.py +++ b/alembic/ddl/mssql.py @@ -170,7 +170,7 @@ def alter_column( # type:ignore[override] table_name, column_name, schema=schema, name=name ) - def create_index(self, index: Index) -> None: + def create_index(self, index: Index, **kw: Any) -> None: # this likely defaults to None if not present, so get() # should normally not return the default value. being # defensive in any case @@ -179,7 +179,7 @@ def create_index(self, index: Index) -> None: for col in mssql_include: if col not in index.table.c: index.table.append_column(Column(col, sqltypes.NullType)) - self._exec(CreateIndex(index)) + self._exec(CreateIndex(index, **kw)) def bulk_insert( # type:ignore[override] self, table: Union[TableClause, Table], rows: List[dict], **kw: Any diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py index afabd6c0..f8ae9704 100644 --- a/alembic/ddl/postgresql.py +++ b/alembic/ddl/postgresql.py @@ -80,15 +80,17 @@ class PostgresqlImpl(DefaultImpl): ) identity_attrs_ignore = ("on_null", "order") - def create_index(self, index): + def create_index(self, index: Index, **kw: Any) -> None: # this likely defaults to None if not present, so get() # should normally not return the default value. being # defensive in any case postgresql_include = index.kwargs.get("postgresql_include", None) or () for col in postgresql_include: - if col not in index.table.c: - index.table.append_column(Column(col, sqltypes.NullType)) - self._exec(CreateIndex(index)) + if col not in index.table.c: # type: ignore[union-attr] + index.table.append_column( # type: ignore[union-attr] + Column(col, sqltypes.NullType) + ) + self._exec(CreateIndex(index, **kw)) def prep_table_for_batch(self, batch_impl, table): for constraint in table.constraints: diff --git a/alembic/op.pyi b/alembic/op.pyi index c62ffc3d..dd470310 100644 --- a/alembic/op.pyi +++ b/alembic/op.pyi @@ -646,6 +646,7 @@ def create_index( *, schema: Optional[str] = None, unique: bool = False, + if_not_exists: Optional[bool] = None, **kw: Any, ) -> None: r"""Issue a "create index" instruction using the current @@ -675,20 +676,24 @@ def create_index( :class:`~sqlalchemy.sql.elements.quoted_name`. :param unique: If True, create a unique index. - :param quote: - Force quoting of this column's name on or off, corresponding - to ``True`` or ``False``. When left at its default - of ``None``, the column identifier will be quoted according to - whether the name is case sensitive (identifiers with at least one - upper case character are treated as case sensitive), or if it's a - reserved word. This flag is only needed to force quoting of a - reserved word which is not known by the SQLAlchemy dialect. + :param quote: Force quoting of this column's name on or off, + corresponding to ``True`` or ``False``. When left at its default + of ``None``, the column identifier will be quoted according to + whether the name is case sensitive (identifiers with at least one + upper case character are treated as case sensitive), or if it's a + reserved word. This flag is only needed to force quoting of a + reserved word which is not known by the SQLAlchemy dialect. + + :param if_not_exists: If True, adds IF NOT EXISTS operator when + creating the new index. + + .. versionadded:: 1.12.0 :param \**kw: Additional keyword arguments not mentioned above are - dialect specific, and passed in the form - ``_``. - See the documentation regarding an individual dialect at - :ref:`dialect_toplevel` for detail on documented arguments. + dialect specific, and passed in the form + ``_``. + See the documentation regarding an individual dialect at + :ref:`dialect_toplevel` for detail on documented arguments. """ @@ -955,6 +960,7 @@ def drop_index( table_name: Optional[str] = None, *, schema: Optional[str] = None, + if_exists: Optional[bool] = None, **kw: Any, ) -> None: r"""Issue a "drop index" instruction using the current @@ -971,11 +977,17 @@ def drop_index( 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 index. + + .. versionadded:: 1.12.0 + :param \**kw: Additional keyword arguments not mentioned above are - dialect specific, and passed in the form - ``_``. - See the documentation regarding an individual dialect at - :ref:`dialect_toplevel` for detail on documented arguments. + dialect specific, and passed in the form + ``_``. + See the documentation regarding an individual dialect at + :ref:`dialect_toplevel` for detail on documented arguments. """ @@ -1228,7 +1240,7 @@ def invoke(operation: MigrateOperation) -> Any: def register_operation( name: str, sourcename: Optional[str] = None -) -> Callable[..., Any]: +) -> Callable[[_T], _T]: """Register a new operation for this class. This method is normally used to add new operations diff --git a/alembic/operations/base.py b/alembic/operations/base.py index e2c1fd23..4f33bd04 100644 --- a/alembic/operations/base.py +++ b/alembic/operations/base.py @@ -1035,6 +1035,7 @@ def create_index( *, schema: Optional[str] = None, unique: bool = False, + if_not_exists: Optional[bool] = None, **kw: Any, ) -> None: r"""Issue a "create index" instruction using the current @@ -1064,20 +1065,24 @@ def create_index( :class:`~sqlalchemy.sql.elements.quoted_name`. :param unique: If True, create a unique index. - :param quote: - Force quoting of this column's name on or off, corresponding - to ``True`` or ``False``. When left at its default - of ``None``, the column identifier will be quoted according to - whether the name is case sensitive (identifiers with at least one - upper case character are treated as case sensitive), or if it's a - reserved word. This flag is only needed to force quoting of a - reserved word which is not known by the SQLAlchemy dialect. + :param quote: Force quoting of this column's name on or off, + corresponding to ``True`` or ``False``. When left at its default + of ``None``, the column identifier will be quoted according to + whether the name is case sensitive (identifiers with at least one + upper case character are treated as case sensitive), or if it's a + reserved word. This flag is only needed to force quoting of a + reserved word which is not known by the SQLAlchemy dialect. + + :param if_not_exists: If True, adds IF NOT EXISTS operator when + creating the new index. + + .. versionadded:: 1.12.0 :param \**kw: Additional keyword arguments not mentioned above are - dialect specific, and passed in the form - ``_``. - See the documentation regarding an individual dialect at - :ref:`dialect_toplevel` for detail on documented arguments. + dialect specific, and passed in the form + ``_``. + See the documentation regarding an individual dialect at + :ref:`dialect_toplevel` for detail on documented arguments. """ # noqa: E501 ... @@ -1359,6 +1364,7 @@ def drop_index( table_name: Optional[str] = None, *, schema: Optional[str] = None, + if_exists: Optional[bool] = None, **kw: Any, ) -> None: r"""Issue a "drop index" instruction using the current @@ -1375,11 +1381,17 @@ def drop_index( 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 index. + + .. versionadded:: 1.12.0 + :param \**kw: Additional keyword arguments not mentioned above are - dialect specific, and passed in the form - ``_``. - See the documentation regarding an individual dialect at - :ref:`dialect_toplevel` for detail on documented arguments. + dialect specific, and passed in the form + ``_``. + See the documentation regarding an individual dialect at + :ref:`dialect_toplevel` for detail on documented arguments. """ # noqa: E501 ... diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py index e4413dd3..8c88e885 100644 --- a/alembic/operations/batch.py +++ b/alembic/operations/batch.py @@ -185,11 +185,11 @@ def drop_constraint(self, const: Constraint) -> None: def rename_table(self, *arg, **kw): self.batch.append(("rename_table", arg, kw)) - def create_index(self, idx: Index) -> None: - self.batch.append(("create_index", (idx,), {})) + def create_index(self, idx: Index, **kw: Any) -> None: + self.batch.append(("create_index", (idx,), kw)) - def drop_index(self, idx: Index) -> None: - self.batch.append(("drop_index", (idx,), {})) + def drop_index(self, idx: Index, **kw: Any) -> None: + self.batch.append(("drop_index", (idx,), kw)) def create_table_comment(self, table): self.batch.append(("create_table_comment", (table,), {})) diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py index 4d900121..dcf7b6b0 100644 --- a/alembic/operations/ops.py +++ b/alembic/operations/ops.py @@ -876,6 +876,7 @@ def __init__( *, schema: Optional[str] = None, unique: bool = False, + if_not_exists: Optional[bool] = None, **kw: Any, ) -> None: self.index_name = index_name @@ -883,6 +884,7 @@ def __init__( self.columns = columns self.schema = schema self.unique = unique + self.if_not_exists = if_not_exists self.kw = kw def reverse(self) -> DropIndexOp: @@ -928,6 +930,7 @@ def create_index( *, schema: Optional[str] = None, unique: bool = False, + if_not_exists: Optional[bool] = None, **kw: Any, ) -> None: r"""Issue a "create index" instruction using the current @@ -957,24 +960,34 @@ def create_index( :class:`~sqlalchemy.sql.elements.quoted_name`. :param unique: If True, create a unique index. - :param quote: - Force quoting of this column's name on or off, corresponding - to ``True`` or ``False``. When left at its default - of ``None``, the column identifier will be quoted according to - whether the name is case sensitive (identifiers with at least one - upper case character are treated as case sensitive), or if it's a - reserved word. This flag is only needed to force quoting of a - reserved word which is not known by the SQLAlchemy dialect. + :param quote: Force quoting of this column's name on or off, + corresponding to ``True`` or ``False``. When left at its default + of ``None``, the column identifier will be quoted according to + whether the name is case sensitive (identifiers with at least one + upper case character are treated as case sensitive), or if it's a + reserved word. This flag is only needed to force quoting of a + reserved word which is not known by the SQLAlchemy dialect. + + :param if_not_exists: If True, adds IF NOT EXISTS operator when + creating the new index. + + .. versionadded:: 1.12.0 :param \**kw: Additional keyword arguments not mentioned above are - dialect specific, and passed in the form - ``_``. - See the documentation regarding an individual dialect at - :ref:`dialect_toplevel` for detail on documented arguments. + dialect specific, and passed in the form + ``_``. + See the documentation regarding an individual dialect at + :ref:`dialect_toplevel` for detail on documented arguments. """ op = cls( - index_name, table_name, columns, schema=schema, unique=unique, **kw + index_name, + table_name, + columns, + schema=schema, + unique=unique, + if_not_exists=if_not_exists, + **kw, ) return operations.invoke(op) @@ -1016,12 +1029,14 @@ def __init__( table_name: Optional[str] = None, *, schema: Optional[str] = None, + if_exists: Optional[bool] = None, _reverse: Optional[CreateIndexOp] = None, **kw: Any, ) -> None: self.index_name = index_name self.table_name = table_name self.schema = schema + self.if_exists = if_exists self._reverse = _reverse self.kw = kw @@ -1065,6 +1080,7 @@ def drop_index( table_name: Optional[str] = None, *, schema: Optional[str] = None, + if_exists: Optional[bool] = None, **kw: Any, ) -> None: r"""Issue a "drop index" instruction using the current @@ -1081,14 +1097,26 @@ def drop_index( 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 index. + + .. versionadded:: 1.12.0 + :param \**kw: Additional keyword arguments not mentioned above are - dialect specific, and passed in the form - ``_``. - See the documentation regarding an individual dialect at - :ref:`dialect_toplevel` for detail on documented arguments. + dialect specific, and passed in the form + ``_``. + See the documentation regarding an individual dialect at + :ref:`dialect_toplevel` for detail on documented arguments. """ - op = cls(index_name, table_name=table_name, schema=schema, **kw) + op = cls( + index_name, + table_name=table_name, + schema=schema, + if_exists=if_exists, + **kw, + ) return operations.invoke(op) @classmethod diff --git a/alembic/operations/toimpl.py b/alembic/operations/toimpl.py index 72229c6c..ba974b62 100644 --- a/alembic/operations/toimpl.py +++ b/alembic/operations/toimpl.py @@ -5,6 +5,7 @@ from . import ops from .base import Operations from ..util.sqla_compat import _copy +from ..util.sqla_compat import sqla_2 if TYPE_CHECKING: from sqlalchemy.sql.schema import Table @@ -59,7 +60,7 @@ def _count_constraint(constraint): existing_nullable=existing_nullable, comment=comment, existing_comment=existing_comment, - **operation.kw + **operation.kw, ) if type_: @@ -95,13 +96,27 @@ def create_index( operations: "Operations", operation: "ops.CreateIndexOp" ) -> None: idx = operation.to_index(operations.migration_context) - operations.impl.create_index(idx) + kw = {} + if operation.if_not_exists is not None: + if not sqla_2: + raise NotImplementedError("SQLAlchemy 2.0+ required") + + kw["if_not_exists"] = operation.if_not_exists + operations.impl.create_index(idx, **kw) @Operations.implementation_for(ops.DropIndexOp) def drop_index(operations: "Operations", operation: "ops.DropIndexOp") -> None: + kw = {} + if operation.if_exists is not None: + if not sqla_2: + raise NotImplementedError("SQLAlchemy 2.0+ required") + + kw["if_exists"] = operation.if_exists + operations.impl.drop_index( - operation.to_index(operations.migration_context) + operation.to_index(operations.migration_context), + **kw, ) diff --git a/docs/build/unreleased/151.rst b/docs/build/unreleased/151.rst new file mode 100644 index 00000000..9f2c42dd --- /dev/null +++ b/docs/build/unreleased/151.rst @@ -0,0 +1,6 @@ +.. change:: + :tags: feature, operations + :tickets: 151 + + Added parameters if_exists and if_not_exists for index operations. + Pull request courtesy of Max Adrian. \ No newline at end of file diff --git a/tests/test_op.py b/tests/test_op.py index 54637fd3..64632443 100644 --- a/tests/test_op.py +++ b/tests/test_op.py @@ -824,6 +824,12 @@ def test_create_index(self): op.create_index("ik_test", "t1", ["foo", "bar"]) context.assert_("CREATE INDEX ik_test ON t1 (foo, bar)") + @config.requirements.sqlalchemy_2 + def test_create_index_if_not_exists(self): + context = op_fixture() + op.create_index("ik_test", "t1", ["foo", "bar"], if_not_exists=True) + context.assert_("CREATE INDEX IF NOT EXISTS ik_test ON t1 (foo, bar)") + def test_create_unique_index(self): context = op_fixture() op.create_index("ik_test", "t1", ["foo", "bar"], unique=True) @@ -880,6 +886,12 @@ def test_drop_index_schema(self): op.drop_index("ik_test", schema="foo") context.assert_("DROP INDEX foo.ik_test") + @config.requirements.sqlalchemy_2 + def test_drop_index_if_exists(self): + context = op_fixture() + op.drop_index("ik_test", if_exists=True) + context.assert_("DROP INDEX IF EXISTS ik_test") + def test_drop_table(self): context = op_fixture() op.drop_table("tb_test") diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 8984437b..3a5e9838 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -122,6 +122,12 @@ def test_create_index_postgresql_include_is_none(self): op.create_index("i", "t", ["c1", "c2"], unique=False) context.assert_("CREATE INDEX i ON t (c1, c2)") + @config.requirements.sqlalchemy_2 + def test_create_index_postgresql_if_not_exists(self): + context = op_fixture("postgresql") + op.create_index("i", "t", ["c1", "c2"], if_not_exists=True) + context.assert_("CREATE INDEX IF NOT EXISTS i ON t (c1, c2)") + @config.combinations("include_table", "no_table", argnames="include_table") def test_drop_index_postgresql_concurrently(self, include_table): context = op_fixture("postgresql") @@ -135,6 +141,12 @@ def test_drop_index_postgresql_concurrently(self, include_table): op.drop_index("geocoded", postgresql_concurrently=True) context.assert_("DROP INDEX CONCURRENTLY geocoded") + @config.requirements.sqlalchemy_2 + def test_drop_index_postgresql_if_exists(self): + context = op_fixture("postgresql") + op.drop_index("geocoded", if_exists=True) + context.assert_("DROP INDEX IF EXISTS geocoded") + def test_alter_column_type_using(self): context = op_fixture("postgresql") op.alter_column("t", "c", type_=Integer, postgresql_using="c::integer")