Disable check constraints via feature flag

Author: danielfrankcom

aurora_dsql_django/features.py

@@ -31,9 +31,10 @@ class DatabaseFeatures(features.DatabaseFeatures):
     # at the end of each save operation?
     supports_forward_references = True
 
-    # Does it support foreign keys?
     supports_foreign_keys = False
 
+    supports_table_check_constraints = False
+
     # Can it create foreign key constraints inline when adding columns?
     can_create_inline_fk = False
 

aurora_dsql_django/schema.py

@@ -49,8 +49,6 @@ class DatabaseSchemaEditor(schema.DatabaseSchemaEditor):
 
     # These "ALTER TABLE" operations are not supported.
     sql_create_pk = ""
-    sql_create_check = ""
-    sql_delete_check = ""
     sql_delete_constraint = ""
     sql_delete_column = ""
 
@@ -67,12 +65,19 @@ def __enter__(self):
     def add_index(self, model, index, concurrently=False):
         if index.contains_expressions and not self.connection.features.supports_expression_indexes:
             return None
-        super().add_index(model, index, concurrently)
+        return super().add_index(model, index, concurrently)
 
     def remove_index(self, model, index, concurrently=False):
         if index.contains_expressions and not self.connection.features.supports_expression_indexes:
             return None
-        super().remove_index(model, index, concurrently)
+        return super().remove_index(model, index, concurrently)
+
+    def _check_sql(self, name, check):
+        # There is no feature check in the upstream implementation when creating
+        # a model, so we add our own check.
+        if not self.connection.features.supports_table_check_constraints:
+            return None
+        return super()._check_sql(name, check)
 
     def _index_columns(self, table, columns, col_suffixes, opclasses):
         # Aurora DSQL doesn't support PostgreSQL opclasses.

aurora_dsql_django/tests/unit/test_features.py

@@ -18,6 +18,9 @@ def test_supports_forward_references(self):
     def test_supports_foreign_keys(self):
         self.assertFalse(self.features.supports_foreign_keys)
 
+    def test_supports_check_constraints(self):
+        self.assertFalse(self.features.supports_table_check_constraints)
+
     def test_can_create_inline_fk(self):
         self.assertFalse(self.features.can_create_inline_fk)
 

aurora_dsql_django/tests/unit/test_schema.py

@@ -23,8 +23,6 @@ def test_sql_attributes(self):
             self.schema_editor.sql_update_with_default,
             "UPDATE %(table)s SET %(column)s = %(default)s WHERE %(column)s IS NULL"
         )
-        self.assertEqual(self.schema_editor.sql_create_check, "")
-        self.assertEqual(self.schema_editor.sql_delete_check, "")
         self.assertEqual(self.schema_editor.sql_delete_constraint, "")
         self.assertEqual(self.schema_editor.sql_delete_column, "")
 
@@ -93,6 +91,25 @@ def test_create_like_index_sql(self):
 
         self.assertIsNone(result)
 
+    @patch('aurora_dsql_django.schema.schema.DatabaseSchemaEditor._check_sql')
+    def test_check_sql_feature_disabled(self, mock_super_check_sql):
+        self.connection.features.supports_table_check_constraints = False
+        
+        result = self.schema_editor._check_sql("test_check", "age >= 0")
+        
+        self.assertIsNone(result)
+        mock_super_check_sql.assert_not_called()
+
+    @patch('aurora_dsql_django.schema.schema.DatabaseSchemaEditor._check_sql')
+    def test_check_sql_feature_enabled(self, mock_super_check_sql):
+        self.connection.features.supports_table_check_constraints = True
+        mock_super_check_sql.return_value = "CHECK (age >= 0)"
+        
+        result = self.schema_editor._check_sql("test_check", "age >= 0")
+        
+        mock_super_check_sql.assert_called_once_with("test_check", "age >= 0")
+        self.assertEqual(result, "CHECK (age >= 0)")
+
 
 if __name__ == '__main__':
     unittest.main()

aurora_dsql_django/tests/unit/test_wrapper.py

@@ -3,6 +3,7 @@
 import django
 from django.conf import settings
 from django.db import models
+from django.db.models import CheckConstraint, Q
 from aurora_dsql_django.base import DatabaseWrapper
 from aurora_dsql_django.features import DatabaseFeatures
 from aurora_dsql_django.schema import DatabaseSchemaEditor
@@ -16,20 +17,46 @@
     django.setup()
 
 
+def simple_quote_value(value):
+    return f"'{value}'"
+
+
 class TestWrapper(unittest.TestCase):
     """Test Aurora DSQL wrapper behavior when all parts are working together"""
 
     def setUp(self):
         self.connection = DatabaseWrapper({})
         self.connection.connection = MagicMock()
-        self.connection.connection.encoding = 'utf8'
 
         # Configure mock to use real components.
         self.connection.features = DatabaseFeatures(self.connection)
         self.schema_editor = DatabaseSchemaEditor(self.connection)
 
-    def test_foreign_key_sql_generation(self):
-        """Ensure foreign key SQL is not generated when the feature is disabled"""
+    def _assert_sql_not_generated(self, operation_func, sql_patterns, message):
+        """Helper method to verify SQL patterns are not generated"""
+        executed_sql = []
+
+        def mock_execute(sql, params=None):
+            executed_sql.append((sql, params))
+
+        # Capture SQL statements without running anything against a real DB.
+        execute_patch = patch.object(self.schema_editor, 'execute', side_effect=mock_execute)
+
+        # Work around issue caused by missing encoding configuration in test environment.
+        quote_patch = patch.object(self.schema_editor, 'quote_value', side_effect=simple_quote_value)
+
+        with execute_patch, quote_patch:
+            operation_func()
+
+        all_sql = [str(sql) for sql, _ in executed_sql]
+        if hasattr(self.schema_editor, 'deferred_sql'):
+            all_sql += [str(sql) for sql in self.schema_editor.deferred_sql]
+
+        matching_statements = [sql for sql in all_sql if any(pattern in sql for pattern in sql_patterns)]
+        self.assertListEqual([], matching_statements, message)
+
+    def test_foreign_key_operations_ignored(self):
+        """Ensure foreign key constraint operations are ignored for model creation when the feature is disabled"""
 
         class ParentModel(models.Model):
             class Meta:
@@ -41,14 +68,77 @@ class ChildModel(models.Model):
             class Meta:
                 app_label = 'test_app'
 
-        # Mock execute to capture SQL without actually running it.
-        with patch.object(self.schema_editor, 'execute'):
+        def operation():
             with self.schema_editor:
                 self.schema_editor.create_model(ChildModel)
 
-                # Check that no foreign key SQL was deferred.
-                foreign_key_statements = [sql for sql in self.schema_editor.deferred_sql if 'FOREIGN KEY' in str(sql)]
-                self.assertListEqual([], foreign_key_statements, "Should not generate foreign key SQL")
+        self._assert_sql_not_generated(
+            operation,
+            ['FOREIGN KEY', 'REFERENCES'],
+            "Should not generate foreign key SQL"
+        )
+
+    def test_check_constraint_create_model_ignored(self):
+        """Ensure check constraint operations are ignored for model creation when the feature is disabled"""
+
+        class CheckConstraintModel(models.Model):
+            age = models.IntegerField()
+
+            class Meta:
+                app_label = 'test_app'
+                constraints = [
+                    CheckConstraint(condition=Q(age__gte=0), name='age_gte_0')
+                ]
+
+        def operation():
+            with self.schema_editor:
+                self.schema_editor.create_model(CheckConstraintModel)
+
+        self._assert_sql_not_generated(
+            operation,
+            ['CHECK'],
+            "Should not generate check constraint SQL"
+        )
+
+    def test_check_constraint_add_constraint_ignored(self):
+        """Ensure add_constraint operations ignore check constraints when the feature is disabled"""
+
+        class AddCheckConstraintModel(models.Model):
+            age = models.IntegerField()
+
+            class Meta:
+                app_label = 'test_app'
+
+        constraint = CheckConstraint(condition=Q(age__gte=0), name='age_gte_0')
+
+        def operation():
+            self.schema_editor.add_constraint(AddCheckConstraintModel, constraint)
+
+        self._assert_sql_not_generated(
+            operation,
+            ['CHECK'],
+            "Should not execute check constraint SQL"
+        )
+
+    def test_check_constraint_remove_constraint_ignored(self):
+        """Ensure remove_constraint operations ignore check constraints when the feature is disabled"""
+
+        class RemoveCheckConstraintModel(models.Model):
+            age = models.IntegerField()
+
+            class Meta:
+                app_label = 'test_app'
+
+        constraint = CheckConstraint(condition=Q(age__gte=0), name='age_gte_0')
+
+        def operation():
+            self.schema_editor.remove_constraint(RemoveCheckConstraintModel, constraint)
+
+        self._assert_sql_not_generated(
+            operation,
+            ['CONSTRAINT'],
+            "Should not execute check constraint removal SQL"
+        )
 
 
 if __name__ == '__main__':