]> git.ipfire.org Git - thirdparty/sqlalchemy/alembic.git/commitdiff
Added parameters if_exists and if_not_exists for index operations.
authorАдриан Максим Александрович <adrian@tochka.com>
Tue, 11 Jul 2023 19:31:29 +0000 (15:31 -0400)
committerFederico Caselli <cfederico87@gmail.com>
Tue, 11 Jul 2023 20:27:15 +0000 (22:27 +0200)
Fixes: #151
<!-- Provide a general summary of your proposed changes in the Title field above -->

### Description
<!-- Describe your changes in detail -->

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
<!-- go over following points. check them with an `x` if they do apply, (they turn into clickable checkboxes once the PR is submitted, so no need to do everything at once)

-->

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: #<issue number>` 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: #<issue number>` 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
alembic/ddl/mssql.py
alembic/ddl/postgresql.py
alembic/op.pyi
alembic/operations/base.py
alembic/operations/batch.py
alembic/operations/ops.py
alembic/operations/toimpl.py
docs/build/unreleased/151.rst [new file with mode: 0644]
tests/test_op.py
tests/test_postgresql.py

index 31667ef8c694b68a91ddfbc7a782bad110f83142..a6102d1e3b2787c315cedf152b34ba6bf21d52b8 100644 (file)
@@ -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,
index 56dd12c35ea3f640faeba284a99bb120c7464791..e77fd1b5bb70c6309d29c3e52a7380b1883035d5 100644 (file)
@@ -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
index afabd6c010bd11edca40d36b6165e0d6bbad3548..f8ae97046690981510f5dc8b5ac1495d3bf73339 100644 (file)
@@ -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:
index c62ffc3d60e38b15fc243b08c306c3d6aa84675c..dd4703100b20b6c23458fc63f4276d9f6589b3d8 100644 (file)
@@ -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
-        ``<dialectname>_<argname>``.
-        See the documentation regarding an individual dialect at
-        :ref:`dialect_toplevel` for detail on documented arguments.
+     dialect specific, and passed in the form
+     ``<dialectname>_<argname>``.
+     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
-        ``<dialectname>_<argname>``.
-        See the documentation regarding an individual dialect at
-        :ref:`dialect_toplevel` for detail on documented arguments.
+     dialect specific, and passed in the form
+     ``<dialectname>_<argname>``.
+     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
index e2c1fd23061556fb71e51a2c52a0a505b58103e1..4f33bd0464dc1748befcf13e5f03786ed1c5310f 100644 (file)
@@ -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
-                ``<dialectname>_<argname>``.
-                See the documentation regarding an individual dialect at
-                :ref:`dialect_toplevel` for detail on documented arguments.
+             dialect specific, and passed in the form
+             ``<dialectname>_<argname>``.
+             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
-                ``<dialectname>_<argname>``.
-                See the documentation regarding an individual dialect at
-                :ref:`dialect_toplevel` for detail on documented arguments.
+             dialect specific, and passed in the form
+             ``<dialectname>_<argname>``.
+             See the documentation regarding an individual dialect at
+             :ref:`dialect_toplevel` for detail on documented arguments.
 
             """  # noqa: E501
             ...
index e4413dd3372932f1fdecd8f6885ad0a972664388..8c88e885acfdb54bfb506848c5a91b2327ae3703 100644 (file)
@@ -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,), {}))
index 4d9001212ccc56150b6cc948efd6bdcade9450ab..dcf7b6b00cfbbd4966de1521fce5ff3c6f1498bf 100644 (file)
@@ -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
-            ``<dialectname>_<argname>``.
-            See the documentation regarding an individual dialect at
-            :ref:`dialect_toplevel` for detail on documented arguments.
+         dialect specific, and passed in the form
+         ``<dialectname>_<argname>``.
+         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
-            ``<dialectname>_<argname>``.
-            See the documentation regarding an individual dialect at
-            :ref:`dialect_toplevel` for detail on documented arguments.
+         dialect specific, and passed in the form
+         ``<dialectname>_<argname>``.
+         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
index 72229c6c7c70f4d8004ff7775331b16277a943cd..ba974b6228c35a66fd4b1c103418f21b55f2e409 100644 (file)
@@ -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 (file)
index 0000000..9f2c42d
--- /dev/null
@@ -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
index 54637fd375fa1b6c280f8cd7b95c7a9d7f81adba..64632443d1348cd0ff4874709e751e201e914139 100644 (file)
@@ -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")
index 8984437b713613ce10ee37f5a6c6b5226f27adee..3a5e98380e9db77e3095d20422626b3aeb1a83eb 100644 (file)
@@ -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")