From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Thu, 4 Apr 2013 19:56:04 +0000 (-0400)
Subject: Added :func:`.operations.create_primary_key`
X-Git-Tag: rel_0_5_0~8
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=a41e4677e355806131ee16b97520f146c3c971a3;p=thirdparty%2Fsqlalchemy%2Falembic.git

Added :func:`.operations.create_primary_key`
operation, will genenerate an ADD CONSTRAINT
for a primary key.
#93
---

diff --git a/alembic/operations.py b/alembic/operations.py
index f97ec6e6..1e21bc77 100644
--- a/alembic/operations.py
+++ b/alembic/operations.py
@@ -50,6 +50,17 @@ class Operations(object):
         yield op
         alembic.op._remove_proxy()
 
+
+    def _primary_key_constraint(self, name, table_name, cols, schema=None):
+        m = sa_schema.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)
+        return p
+
     def _foreign_key_constraint(self, name, source, referent,
                                     local_cols, remote_cols,
                                     onupdate=None, ondelete=None,
@@ -388,6 +399,46 @@ class Operations(object):
         )
 
 
+    def create_primary_key(self, name, table_name, cols, schema=None):
+        """Issue a "create primary key" instruction using the current
+        migration context.
+
+        e.g.::
+
+            from alembic import op
+            op.create_primary_key(
+                        "pk_my_table", "my_table",
+                        ["id", "version"]
+                    )
+
+        This internally generates a :class:`~sqlalchemy.schema.Table` object
+        containing the necessary columns, then generates a new
+        :class:`~sqlalchemy.schema.PrimaryKeyConstraint`
+        object which it then associates with the :class:`~sqlalchemy.schema.Table`.
+        Any event listeners associated with this action will be fired
+        off normally.   The :class:`~sqlalchemy.schema.AddConstraint`
+        construct is ultimately used to generate the ALTER statement.
+
+        .. versionadded:: 0.5.0
+
+        :param name: Name of the primary key constraint.  The name is necessary
+         so that an ALTER statement can be emitted.  For setups that
+         use an automated naming scheme such as that described at
+         `NamingConventions <http://www.sqlalchemy.org/trac/wiki/UsageRecipes/NamingConventions>`_,
+         ``name`` here can be ``None``, as the event listener will
+         apply the name to the constraint object when it is associated
+         with the table.
+        :param table_name: String name of the target table.
+        :param cols: a list of string column names to be applied to the
+         primary key constraint.
+        :param schema: Optional schema name of the table.
+
+        """
+        self.impl.add_constraint(
+                    self._primary_key_constraint(name, table_name, cols,
+                                schema)
+                )
+
     def create_foreign_key(self, name, source, referent, local_cols,
                            remote_cols, onupdate=None, ondelete=None,
                            source_schema=None, referent_schema=None):
diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst
index aee64a75..f60ac595 100644
--- a/docs/build/changelog.rst
+++ b/docs/build/changelog.rst
@@ -6,6 +6,14 @@ Changelog
 .. changelog::
     :version: 0.5.0
 
+    .. change::
+        :tags: feature
+        :tickets: 93
+
+      Added :func:`.operations.create_primary_key`
+      operation, will genenerate an ADD CONSTRAINT
+      for a primary key.
+
     .. change::
         :tags: bug, mssql
         :tickets: 109
diff --git a/tests/test_op.py b/tests/test_op.py
index d5a6abe9..1118ae49 100644
--- a/tests/test_op.py
+++ b/tests/test_op.py
@@ -338,7 +338,7 @@ def test_add_foreign_key_schema():
     context = op_fixture()
     op.create_foreign_key('fk_test', 't1', 't2',
                     ['foo', 'bar'], ['bat', 'hoho'],
-                   source_schema='foo2', referent_schema='bar2')
+                    source_schema='foo2', referent_schema='bar2')
     context.assert_(
         "ALTER TABLE foo2.t1 ADD CONSTRAINT fk_test FOREIGN KEY(foo, bar) "
             "REFERENCES bar2.t2 (bat, hoho)"
@@ -372,6 +372,20 @@ def test_add_foreign_key_self_referential():
         "FOREIGN KEY(foo) REFERENCES t1 (bar)"
     )
 
+def test_add_primary_key_constraint():
+    context = op_fixture()
+    op.create_primary_key("pk_test", "t1", ["foo", "bar"])
+    context.assert_(
+        "ALTER TABLE t1 ADD CONSTRAINT pk_test PRIMARY KEY (foo, bar)"
+    )
+
+def test_add_primary_key_constraint_schema():
+    context = op_fixture()
+    op.create_primary_key("pk_test", "t1", ["foo"], schema="bar")
+    context.assert_(
+        "ALTER TABLE bar.t1 ADD CONSTRAINT pk_test PRIMARY KEY (foo)"
+    )
+
 def test_add_check_constraint():
     context = op_fixture()
     op.create_check_constraint(