From 88db76eb66bb4393dffb5384ee649afe02020a46 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D0=B4=D1=80=D0=B8=D0=B0=D0=BD=20=D0=9C=D0=B0=D0=BA?= =?UTF-8?q?=D1=81=D0=B8=D0=BC=20=D0=90=D0=BB=D0=B5=D0=BA=D1=81=D0=B0=D0=BD?= =?UTF-8?q?=D0=B4=D1=80=D0=BE=D0=B2=D0=B8=D1=87?= Date: Tue, 11 Jul 2023 15:31:29 -0400 Subject: [PATCH] Added parameters if_exists and if_not_exists for index operations. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- alembic/ddl/impl.py | 8 ++--- alembic/ddl/mssql.py | 4 +-- alembic/ddl/postgresql.py | 10 +++--- alembic/op.pyi | 46 +++++++++++++++---------- alembic/operations/base.py | 44 +++++++++++++++--------- alembic/operations/batch.py | 8 ++--- alembic/operations/ops.py | 64 +++++++++++++++++++++++++---------- alembic/operations/toimpl.py | 21 ++++++++++-- docs/build/unreleased/151.rst | 6 ++++ tests/test_op.py | 12 +++++++ tests/test_postgresql.py | 12 +++++++ 11 files changed, 167 insertions(+), 68 deletions(-) create mode 100644 docs/build/unreleased/151.rst 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")