From: Адриан Максим Александрович Date: Tue, 11 Jul 2023 19:31:29 +0000 (-0400) Subject: Added parameters if_exists and if_not_exists for index operations. X-Git-Tag: rel_1_11_2~1 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=88db76eb66bb4393dffb5384ee649afe02020a46;p=thirdparty%2Fsqlalchemy%2Falembic.git Added parameters if_exists and if_not_exists for index operations. Fixes: #151 ### Description As mentioned in #151, IF EXISTS/IF NOT EXISTS syntax was implemented in SQLAlchemy 2.0. This request adds an ability to use them for index operations. If the issue implies to implement all the possible cases with these directives, I could continue working on it 🙂 ### Checklist This pull request is: - [ ] A documentation / typographical error fix - Good to go, no issue or tests are needed - [ ] A short code fix - please include the issue number, and create an issue if none exists, which must include a complete example of the issue. one line code fixes without an issue and demonstration will not be accepted. - Please include: `Fixes: #` in the commit message - please include tests. one line code fixes without tests will not be accepted. - [x] A new feature implementation - please include the issue number, and create an issue if none exists, which must include a complete example of how the feature would look. - Please include: `Fixes: #` in the commit message - please include tests. **Have a nice day!** Closes: #1260 Pull-request: https://github.com/sqlalchemy/alembic/pull/1260 Pull-request-sha: 5ed62d17c349dc175521993c5a62b842f6ac624d Change-Id: Ic0fec21e20d4a868e9e29275e540f3e09918d1ff --- 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 @@ class DefaultImpl(metaclass=ImplMeta): 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 @@ class DefaultImpl(metaclass=ImplMeta): 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 @@ class MSSQLImpl(DefaultImpl): 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 @@ class MSSQLImpl(DefaultImpl): 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 @@ class Operations(AbstractOperations): *, 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 @@ class Operations(AbstractOperations): :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 @@ class Operations(AbstractOperations): 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 @@ class Operations(AbstractOperations): 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 @@ class BatchOperationsImpl: 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 @@ class CreateIndexOp(MigrateOperation): *, 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 @@ class CreateIndexOp(MigrateOperation): 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 @@ class CreateIndexOp(MigrateOperation): *, 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 @@ class CreateIndexOp(MigrateOperation): :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 @@ class DropIndexOp(MigrateOperation): 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 @@ class DropIndexOp(MigrateOperation): 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 @@ class DropIndexOp(MigrateOperation): 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 sqlalchemy import schema as sa_schema 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 alter_column( 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 @@ class OpTest(TestBase): 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 @@ class OpTest(TestBase): 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 @@ class PostgresqlOpTest(TestBase): 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 @@ class PostgresqlOpTest(TestBase): 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")