diff --git a/migrate/changeset/ansisql.py b/migrate/changeset/ansisql.py index d711d5a..e188100 100644 --- a/migrate/changeset/ansisql.py +++ b/migrate/changeset/ansisql.py @@ -250,6 +250,8 @@ class ANSIConstraintGenerator(ANSIConstraintCommon): ) ret = "CONSTRAINT %(name)s FOREIGN KEY (%(columns)s) "\ "REFERENCES %(reftable)s (%(referenced)s)"%params + elif isinstance(cons,constraint.CheckConstraint): + ret = "CHECK (%s)"%cons.sqltext else: raise exceptions.InvalidConstraintError(cons) return ret @@ -266,6 +268,9 @@ class ANSIConstraintGenerator(ANSIConstraintCommon): def visit_migrate_foreign_key_constraint(self,*p,**k): return self._visit_constraint(*p,**k) + def visit_migrate_check_constraint(self,*p,**k): + return self._visit_constraint(*p,**k) + class ANSIConstraintDropper(ANSIConstraintCommon): def _visit_constraint(self,constraint): self.start_alter_table(constraint) @@ -279,6 +284,9 @@ class ANSIConstraintDropper(ANSIConstraintCommon): def visit_migrate_foreign_key_constraint(self,*p,**k): return self._visit_constraint(*p,**k) + def visit_migrate_check_constraint(self,*p,**k): + return self._visit_constraint(*p,**k) + class ANSIDialect(object): columngenerator = ANSIColumnGenerator columndropper = ANSIColumnDropper diff --git a/migrate/changeset/constraint.py b/migrate/changeset/constraint.py index 4b03e9f..96d12b2 100644 --- a/migrate/changeset/constraint.py +++ b/migrate/changeset/constraint.py @@ -124,3 +124,40 @@ class ForeignKeyConstraint(ConstraintChangeset,schema.ForeignKeyConstraint): def accept_schema_visitor(self,visitor,*p,**k): func = 'visit_migrate_foreign_key_constraint' return self._accept_schema_visitor(visitor,func,*p,**k) + +class CheckConstraint(ConstraintChangeset, schema.CheckConstraint): + def __init__(self, sqltext, *args, **kwargs): + cols = kwargs.pop('columns') + colnames, table = self._normalize_columns(cols) + table = kwargs.pop('table', table) + ConstraintChangeset.__init__(self, *args, **kwargs) + schema.CheckConstraint.__init__(self, sqltext, *args, **kwargs) + if table is not None: + self._set_parent(table) + self.colnames = colnames + + def _set_parent(self, table): + self.table = table + return super(ConstraintChangeset, self)._set_parent(table) + + def create(self): + from migrate.changeset.databases.visitor import get_engine_visitor + visitorcallable = get_engine_visitor(self.table.bind, + 'constraintgenerator') + _engine_run_visitor(self.table.bind, visitorcallable, self) + + def drop(self): + from migrate.changeset.databases.visitor import get_engine_visitor + visitorcallable = get_engine_visitor(self.table.bind, + 'constraintdropper') + _engine_run_visitor(self.table.bind, visitorcallable, self) + self.columns.clear() + return self + + def autoname(self): + return "%(table)s_%(cols)s_check" % \ + {"table": self.table.name, "cols": "_".join(self.colnames)} + + def accept_schema_visitor(self, visitor, *args, **kwargs): + func = 'visit_migrate_check_constraint' + return self._accept_schema_visitor(visitor, func, *args, **kwargs)