Merge "support named CHECK constraints in batch"

5 files changed, 166 insertions, 18 deletions

diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py
index ee1fe05..f5ff0bb 100644
--- a/alembic/operations/batch.py
+++ b/alembic/operations/batch.py
@@ -240,8 +240,13 @@ class ApplyBatchImpl:
         for const in self.table.constraints:
             if _is_type_bound(const):
                 continue
-            elif self.reflected and isinstance(const, CheckConstraint):
-                # TODO: we are skipping reflected CheckConstraint because
+            elif (
+                self.reflected
+                and isinstance(const, CheckConstraint)
+                and not const.name
+            ):
+                # TODO: we are skipping unnamed reflected CheckConstraint
+                # because
                 # we have no way to determine _is_type_bound() for these.
                 pass
             elif const.name:
@@ -457,6 +462,14 @@ class ApplyBatchImpl:
             existing.name = name
             existing_transfer["name"] = name
 
+            # pop named constraints for Boolean/Enum for rename
+            if (
+                "existing_type" in kw
+                and isinstance(kw["existing_type"], SchemaEventTarget)
+                and kw["existing_type"].name
+            ):
+                self.named_constraints.pop(kw["existing_type"].name, None)
+
         if type_ is not None:
             type_ = sqltypes.to_instance(type_)
             # old type is being discarded so turn off eventing
@@ -464,6 +477,7 @@ class ApplyBatchImpl:
             # erase the events set up by this type, but this is simpler.
             # we also ignore the drop_constraint that will come here from
             # Operations.implementation_for(alter_column)
+
             if isinstance(existing.type, SchemaEventTarget):
                 existing.type._create_events = (  # type:ignore[attr-defined]
                     existing.type.create_constraint  # type:ignore[attr-defined] # noqa
@@ -572,6 +586,14 @@ class ApplyBatchImpl:
         del self.column_transfers[column.name]
         self.existing_ordering.remove(column.name)
 
+        # pop named constraints for Boolean/Enum for rename
+        if (
+            "existing_type" in kw
+            and isinstance(kw["existing_type"], SchemaEventTarget)
+            and kw["existing_type"].name
+        ):
+            self.named_constraints.pop(kw["existing_type"].name, None)
+
     def create_column_comment(self, column):
         """the batch table creation function will issue create_column_comment
         on the real "impl" as part of the create table process.
diff --git a/docs/build/batch.rst b/docs/build/batch.rst
index e638eb1..00ccc0f 100644
--- a/docs/build/batch.rst
+++ b/docs/build/batch.rst
@@ -178,6 +178,8 @@ them directly, which can be via the
         ):
         batch_op.add_column(Column('foo', Integer))
 
+.. _batch_schematype_constraints:
+
 Changing the Type of Boolean, Enum and other implicit CHECK datatypes
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -190,16 +192,62 @@ Alembic handles dropping and creating the CHECK constraints here automatically,
 including in the case of batch mode.  When changing the type of an existing
 column, what's necessary is that the existing type be specified fully::
 
-  with self.op.batch_alter_table("some_table"):
+  with self.op.batch_alter_table("some_table") as batch_op:
       batch_op.alter_column(
           'q', type_=Integer,
           existing_type=Boolean(create_constraint=True, constraint_name="ck1"))
 
+When dropping a column that includes a named CHECK constraint, as of Alembic
+1.7 this named constraint must also be provided using a similar form, as there
+is no ability for Alembic to otherwise link this reflected CHECK constraint as
+belonging to a particular column::
+
+    with self.op.batch_alter_table("some_table") as batch_op:
+        batch_op.drop_column(
+            'q',
+            existing_type=Boolean(create_constraint=True, constraint_name="ck1"))
+        )
+
+.. versionchanged:: 1.7  The :meth:`.BatchOperations.drop_column` operation can
+   accept an ``existing_type`` directive where a "schema type" such as
+   :class:`~sqlalchemy.types.Boolean` and :class:`~sqlalchemy.types.Enum` may
+   be specified such that an associated named constraint can be removed.
+
+.. _batch_check_constraints:
+
 Including CHECK constraints
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-SQLAlchemy currently doesn't reflect CHECK constraints on any backend.
-So again these must be stated explicitly if they are to be included in the
+As of Alembic 1.7, **named** CHECK constraints are automatically included
+in batch mode, as modern SQLAlchemy versions are capable of reflecting these
+constraints like any other constraint.
+
+Note that when dropping or renaming a column that is mentioned in a named
+CHECK constraint, this CHECK constraint must be explicitly dropped first,
+as Alembic has no means of linking a reflected CHECK constraint to that
+column.  Supposing column ``q`` of ``some_table`` were mentioned in a CHECK
+constraint named ``ck1``.  In order to drop this column, we have to drop
+the check constraint also::
+
+    with self.op.batch_alter_table("some_table") as batch_op:
+        batch_op.drop_constraint("ck1", "check")
+        batch_op.drop_column('q')
+
+.. versionchanged:: 1.7  Named CHECK constraints participate in batch mode
+   in the same way as any other kind of constraint. This requires that column
+   drops or renames now include explicit directives to drop an existing named
+   constraint which refers to this column, as it will otherwise not be
+   automatically detected as being associated with that particular column.
+
+   Unnamed CHECK constraints continue to be silently omitted from the table
+   recreate operation.
+
+For **unnamed** CHECK constraints, these are still not automatically included
+as part of the batch process as they are often present due to the use of the
+:class:`~sqlalchemy.types.Boolean` or :class:`~sqlalchemy.types.Enum`
+datatypes, which up through SQLAlchemy 1.3 would generate CHECK constraints
+automatically and cannot be tracked to the reflected table. Therefore unnamed
+constraints can be stated explicitly if they are to be included in the
 recreated table::
 
     with op.batch_alter_table("some_table", table_args=[
diff --git a/docs/build/unreleased/884.rst b/docs/build/unreleased/884.rst
new file mode 100644
index 0000000..d357d99
--- /dev/null
+++ b/docs/build/unreleased/884.rst
@@ -0,0 +1,24 @@
+.. change::
+    :tags: usecase, batch
+    :tickets: 884
+
+    Named CHECK constraints are now supported by batch mode, and will
+    automatically be part of the recreated table assuming they are named. They
+    also can be explicitly dropped using ``op.drop_constraint()``. For
+    "unnamed" CHECK constraints, these are still skipped as they cannot be
+    distinguished from the CHECK constraints that are generated by the
+    ``Boolean`` and ``Enum`` datatypes.
+
+    Note that this change may require adjustments to migrations that drop or
+    rename columns which feature an associated named check constraint, such
+    that an additional ``op.drop_constraint()`` directive should be added for
+    that named constraint as there will no longer be an associated column
+    for it; for the ``Boolean`` and ``Enum`` datatypes, an ``existing_type``
+    keyword may be passed to ``BatchOperations.drop_constraint`` as well.
+
+    .. seealso::
+
+      :ref:`batch_schematype_constraints`
+
+      :ref:`batch_check_constraints`
+
diff --git a/tests/requirements.py b/tests/requirements.py
index 011c620..61c3562 100644
--- a/tests/requirements.py
+++ b/tests/requirements.py
@@ -39,6 +39,17 @@ class DefaultRequirements(SuiteRequirements):
         )
 
     @property
+    def non_native_boolean_check_constraint(self):
+        """backend creates a check constraint for booleans if enabled"""
+
+        return exclusions.only_on(
+            exclusions.LambdaPredicate(
+                lambda config: not config.db.dialect.supports_native_boolean
+                and config.db.dialect.non_native_boolean_check_constraint
+            )
+        )
+
+    @property
     def check_constraints_w_enforcement(self):
         return exclusions.fails_on(["mysql", "mariadb"])
 
diff --git a/tests/test_batch.py b/tests/test_batch.py
index 1e0a86f..ab53283 100644
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -6,7 +6,6 @@ from sqlalchemy import CheckConstraint
 from sqlalchemy import Column
 from sqlalchemy import DateTime
 from sqlalchemy import Enum
-from sqlalchemy import exc
 from sqlalchemy import ForeignKey
 from sqlalchemy import ForeignKeyConstraint
 from sqlalchemy import func
@@ -1344,6 +1343,17 @@ class BatchRoundTripTest(TestBase):
             t.create(self.conn)
             return t
 
+    def _ck_constraint_fixture(self):
+        with self.conn.begin():
+            t = Table(
+                "ck_table",
+                self.metadata,
+                Column("id", Integer, nullable=False),
+                CheckConstraint("id is not NULL", name="ck"),
+            )
+            t.create(self.conn)
+            return t
+
     def _datetime_server_default_fixture(self):
         return func.datetime("now", "localtime")
 
@@ -1427,7 +1437,9 @@ class BatchRoundTripTest(TestBase):
     def test_drop_col_schematype(self):
         self._boolean_fixture()
         with self.op.batch_alter_table("hasbool") as batch_op:
-            batch_op.drop_column("x")
+            batch_op.drop_column(
+                "x", existing_type=Boolean(create_constraint=True, name="ck1")
+            )
         insp = inspect(config.db)
 
         assert "x" not in (c["name"] for c in insp.get_columns("hasbool"))
@@ -1659,22 +1671,28 @@ class BatchRoundTripTest(TestBase):
             eq_(pk_const["name"], "newpk")
         eq_(pk_const["constrained_columns"], ["a", "b"])
 
-    @config.requirements.check_constraints_w_enforcement
+    @config.requirements.check_constraint_reflection
     def test_add_ck_constraint(self):
         with self.op.batch_alter_table("foo", recreate="always") as batch_op:
             batch_op.create_check_constraint("newck", text("x > 0"))
 
-        # we dont support reflection of CHECK constraints
-        # so test this by just running invalid data in
-        foo = self.metadata.tables["foo"]
-
-        assert_raises_message(
-            exc.IntegrityError,
-            "newck",
-            self.conn.execute,
-            foo.insert(),
-            {"id": 6, "data": 5, "x": -2},
+        ck_consts = inspect(self.conn).get_check_constraints("foo")
+        ck_consts[0]["sqltext"] = re.sub(
+            r"[\'\"`\(\)]", "", ck_consts[0]["sqltext"]
         )
+        eq_(ck_consts, [{"sqltext": "x > 0", "name": "newck"}])
+
+    @config.requirements.check_constraint_reflection
+    def test_drop_ck_constraint(self):
+        self._ck_constraint_fixture()
+
+        with self.op.batch_alter_table(
+            "ck_table", recreate="always"
+        ) as batch_op:
+            batch_op.drop_constraint("ck", "check")
+
+        ck_consts = inspect(self.conn).get_check_constraints("ck_table")
+        eq_(ck_consts, [])
 
     @config.requirements.unnamed_constraints
     def test_drop_foreign_key(self):
@@ -1854,6 +1872,31 @@ class BatchRoundTripTest(TestBase):
             [{"id": 1, "bflag": True}, {"id": 2, "bflag": False}], "bar"
         )
 
+    #    @config.requirements.check_constraint_reflection
+    def test_rename_column_boolean_named_ck(self):
+        bar = Table(
+            "bar",
+            self.metadata,
+            Column("id", Integer, primary_key=True),
+            Column("flag", Boolean(create_constraint=True, name="ck1")),
+            mysql_engine="InnoDB",
+        )
+        with self.conn.begin():
+            bar.create(self.conn)
+            self.conn.execute(bar.insert(), {"id": 1, "flag": True})
+            self.conn.execute(bar.insert(), {"id": 2, "flag": False})
+
+        with self.op.batch_alter_table("bar", recreate="always") as batch_op:
+            batch_op.alter_column(
+                "flag",
+                new_column_name="bflag",
+                existing_type=Boolean(create_constraint=True, name="ck1"),
+            )
+
+        self._assert_data(
+            [{"id": 1, "bflag": True}, {"id": 2, "bflag": False}], "bar"
+        )
+
     @config.requirements.non_native_boolean
     def test_rename_column_non_native_boolean_no_ck(self):
         bar = Table(