From 40563e9355f4dd2e677d65b37a2e40dfb0f8ff31 Mon Sep 17 00:00:00 2001
From: donkopotamus <derek.harland@finq.co.nz>
Date: Tue, 14 Jan 2014 12:44:05 +1300
Subject: Support mssql_clustered option in mssql dialect for both Table and
 PrimaryKeyConstraint

---
 lib/sqlalchemy/dialects/mssql/base.py | 62 ++++++++++++++++++++++++++++++-----
 test/dialect/mssql/test_compiler.py   | 22 +++++++++++++
 2 files changed, 76 insertions(+), 8 deletions(-)

diff --git a/lib/sqlalchemy/dialects/mssql/base.py b/lib/sqlalchemy/dialects/mssql/base.py
index 90fd1f383..a4997d194 100644
--- a/lib/sqlalchemy/dialects/mssql/base.py
+++ b/lib/sqlalchemy/dialects/mssql/base.py
@@ -103,22 +103,47 @@ for these types will be issued as DATETIME.
 
 .. _mssql_indexes:
 
-MSSQL-Specific Index Options
------------------------------
-
-The MSSQL dialect supports special options for :class:`.Index`.
+Clustered Index Support
+-----------------------
 
-CLUSTERED
-^^^^^^^^^^
+The MSSQL dialect supports clustered indexes (and primary keys) via the
+``mssql_clustered`` option.  This option is available to :class:`.Index`,
+:class:`.Table` and :class:`.PrimaryKeyConstraint`.
 
-The ``mssql_clustered`` option  adds the CLUSTERED keyword to the index::
+To generate a clustered index::
 
     Index("my_index", table.c.x, mssql_clustered=True)
 
-would render the index as ``CREATE CLUSTERED INDEX my_index ON table (x)``
+which renders the index as ``CREATE CLUSTERED INDEX my_index ON table (x)``.
 
 .. versionadded:: 0.8
 
+To generate a clustered primary key use either::
+
+    Table('my_table', metadata,
+          Column('x', ..., primary_key=True),
+          Column('y', ..., primary_key=True),
+          mssql_clustered=True)
+
+or::
+
+    Table('my_table', metadata,
+          Column('x', ...),
+          Column('y', ...),
+          PrimaryKeyConstraint("x", "y", mssql_clustered=True))
+
+which will render the table, for example, as::
+
+  CREATE TABLE my_table (x INTEGER NOT NULL, y INTEGER NOT NULL, PRIMARY KEY CLUSTERED (x, y))
+
+.. versionadded:: 0.9
+
+MSSQL-Specific Index Options
+-----------------------------
+
+In addition to clustering, the MSSQL dialect supports other special options
+for :class:`.Index`.
+
 INCLUDE
 ^^^^^^^
 
@@ -1023,6 +1048,27 @@ class MSDDLCompiler(compiler.DDLCompiler):
             self.preparer.format_table(drop.element.table)
             )
 
+    def visit_primary_key_constraint(self, constraint):
+        if len(constraint) == 0:
+            return ''
+        text = ""
+        if constraint.name is not None:
+            text += "CONSTRAINT %s " % \
+                    self.preparer.format_constraint(constraint)
+        text += "PRIMARY KEY "
+
+        # we allow for mssql_clustered to have been specified directly on a
+        # PrimaryKeyConstraint, or specified upon the table object.  The latter allows
+        # users to tag columns with primary_key=True and still achieve clustering.
+        if (constraint.kwargs.get("mssql_clustered") or
+            constraint.table.kwargs.get("mssql_clustered")):
+            text += "CLUSTERED "
+
+        text += "(%s)" % ', '.join(self.preparer.quote(c.name)
+                                       for c in constraint)
+        text += self.define_constraint_deferrability(constraint)
+        return text
+
 class MSIdentifierPreparer(compiler.IdentifierPreparer):
     reserved_words = RESERVED_WORDS
 
diff --git a/test/dialect/mssql/test_compiler.py b/test/dialect/mssql/test_compiler.py
index 1742f024c..40ca603a0 100644
--- a/test/dialect/mssql/test_compiler.py
+++ b/test/dialect/mssql/test_compiler.py
@@ -510,6 +510,28 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
                             "CREATE TABLE test (id INTEGER NOT NULL IDENTITY(1,1))"
                             )
 
+    def test_table_pkc_clustering_1(self):
+        metadata = MetaData()
+        tbl = Table('test', metadata,
+                    Column('x', Integer, primary_key=True, autoincrement=False),
+                    Column('y', Integer, primary_key=True, autoincrement=False),
+                    mssql_clustered=True)
+        self.assert_compile(schema.CreateTable(tbl),
+                            "CREATE TABLE test (x INTEGER NOT NULL, y INTEGER NOT NULL, "
+                            "PRIMARY KEY CLUSTERED (x, y))"
+                            )
+
+    def test_table_pkc_clustering_2(self):
+        metadata = MetaData()
+        tbl = Table('test', metadata,
+                    Column('x', Integer, autoincrement=False),
+                    Column('y', Integer, autoincrement=False),
+                    PrimaryKeyConstraint("x", "y", mssql_clustered=True))
+        self.assert_compile(schema.CreateTable(tbl),
+                            "CREATE TABLE test (x INTEGER NOT NULL, y INTEGER NOT NULL, "
+                            "PRIMARY KEY CLUSTERED (x, y))"
+                            )
+
     def test_index_clustering(self):
         metadata = MetaData()
         tbl = Table('test', metadata,
-- 
cgit v1.2.1