- Implemented support for :meth:`.BatchOperations.create_primary_key`

and :meth:`.BatchOperations.create_check_constraint`.
fixes #305
- table keyword arguments are copied from the original reflected table,
such as the "mysql_engine" keyword argument.

3 files changed, 25 insertions, 8 deletions

diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py
index 726df78..7135e37 100644
--- a/alembic/operations/batch.py
+++ b/alembic/operations/batch.py
@@ -23,7 +23,7 @@ class BatchOperationsImpl(object):
         self.recreate = recreate
         self.copy_from = copy_from
         self.table_args = table_args
-        self.table_kwargs = table_kwargs
+        self.table_kwargs = dict(table_kwargs)
         self.reflect_args = reflect_args
         self.reflect_kwargs = reflect_kwargs
         self.naming_convention = naming_convention
@@ -139,11 +139,15 @@ class ApplyBatchImpl(object):
         for idx in self.table.indexes:
             self.indexes[idx.name] = idx
 
+        for k in self.table.kwargs:
+            self.table_kwargs.setdefault(k, self.table.kwargs[k])
+
     def _transfer_elements_to_new_table(self):
         assert self.new_table is None, "Can only create new table once"
 
         m = MetaData()
         schema = self.table.schema
+
         self.new_table = new_table = Table(
             '_alembic_batch_temp', m,
             *(list(self.columns.values()) + list(self.table_args)),
@@ -264,6 +268,10 @@ class ApplyBatchImpl(object):
     def add_constraint(self, const):
         if not const.name:
             raise ValueError("Constraint must have a name")
+        if isinstance(const, sql_schema.PrimaryKeyConstraint):
+            if self.table.primary_key in self.unnamed_constraints:
+                self.unnamed_constraints.remove(self.table.primary_key)
+
         self.named_constraints[const.name] = const
 
     def drop_constraint(self, const):
diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
index 82fdd90..16cccb6 100644
--- a/alembic/operations/ops.py
+++ b/alembic/operations/ops.py
@@ -209,7 +209,11 @@ class CreatePrimaryKeyOp(AddConstraintOp):
             :meth:`.Operations.create_primary_key`
 
         """
-        raise NotImplementedError("not yet implemented")
+        op = cls(
+            constraint_name, operations.impl.table_name, columns,
+            schema=operations.impl.schema
+        )
+        return operations.invoke(op)
 
 
 @Operations.register_operation("create_unique_constraint")
@@ -590,7 +594,10 @@ class CreateCheckConstraintOp(AddConstraintOp):
             :meth:`.Operations.create_check_constraint`
 
         """
-        raise NotImplementedError("not yet implemented")
+        op = cls(
+            constraint_name, operations.impl.table_name,
+            condition, schema=operations.impl.schema, **kw)
+        return operations.invoke(op)
 
 
 @Operations.register_operation("create_index")
diff --git a/alembic/operations/schemaobj.py b/alembic/operations/schemaobj.py
index b590aca..f0f8105 100644
--- a/alembic/operations/schemaobj.py
+++ b/alembic/operations/schemaobj.py
@@ -12,11 +12,13 @@ class SchemaObjects(object):
     def primary_key_constraint(self, name, table_name, cols, schema=None):
         m = self.metadata()
         columns = [sa_schema.Column(n, NULLTYPE) for n in cols]
-        t1 = sa_schema.Table(table_name, m,
-                             *columns,
-                             schema=schema)
-        p = sa_schema.PrimaryKeyConstraint(*columns, name=name)
-        t1.append_constraint(p)
+        t = sa_schema.Table(
+            table_name, m,
+            *columns,
+            schema=schema)
+        p = sa_schema.PrimaryKeyConstraint(
+            *[t.c[n] for n in cols], name=name)
+        t.append_constraint(p)
         return p
 
     def foreign_key_constraint(