From 7e68c7892ed18f8fff4ec3dda48c0fe265395f11 Mon Sep 17 00:00:00 2001 From: Charles Leifer Date: Sun, 7 Feb 2021 08:37:13 -0600 Subject: [PATCH] Preserve and handle CHECK() constraints with sqlite migrator. Fixes #2343 --- playhouse/migrate.py | 2 +- tests/migrations.py | 38 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/playhouse/migrate.py b/playhouse/migrate.py index f4b1d5414..f536a45c9 100644 --- a/playhouse/migrate.py +++ b/playhouse/migrate.py @@ -717,7 +717,7 @@ def _update_column(self, table, column_to_update, fn): new_column_defs = [] new_column_names = [] original_column_names = [] - constraint_terms = ('foreign ', 'primary ', 'constraint ') + constraint_terms = ('foreign ', 'primary ', 'constraint ', 'check ') for column_def in column_defs: column_name, = self.column_name_re.match(column_def).groups() diff --git a/tests/migrations.py b/tests/migrations.py index 359725fb8..13b1a9330 100644 --- a/tests/migrations.py +++ b/tests/migrations.py @@ -877,11 +877,45 @@ class Meta: SQL('CONSTRAINT const2 UNIQUE (foreign_data)')] +class HasChecks(TestModel): + key = TextField() + value = IntegerField() + class Meta: + constraints = [ + SQL("CHECK (key != '')"), + SQL('CHECK (value > 0)')] + + class TestSqliteColumnNameRegression(ModelTestCase): database = get_in_memory_db() - requires = [BadNames] + requires = [BadNames, HasChecks] + + def test_sqlite_check_constraints(self): + HasChecks.create(key='k1', value=1) + + migrator = SchemaMigrator.from_database(self.database) + extra = TextField(default='') + migrate(migrator.add_column('has_checks', 'extra', extra)) + + columns = self.database.get_columns('has_checks') + self.assertEqual([c.name for c in columns], + ['id', 'key', 'value', 'extra']) + + HC = Table('has_checks', ('id', 'key', 'value', 'extra')) + HC = HC.bind(self.database) + + # Sanity-check: ensure we can create a new row. + data = {'key': 'k2', 'value': 2, 'extra': 'x2'} + self.assertTrue(HC.insert(data).execute()) + + # Check constraints preserved. + data = {'key': 'k0', 'value': 0, 'extra': 'x0'} + self.assertRaises(IntegrityError, HC.insert(data).execute) + + data = {'key': '', 'value': 3, 'extra': 'x3'} + self.assertRaises(IntegrityError, HC.insert(data).execute) - def test_sqlite_column_name_regression(self): + def test_sqlite_column_name_constraint_regression(self): BadNames.create(primary_data='pd', foreign_data='fd', data='d') migrator = SchemaMigrator.from_database(self.database)