Add special handling for SQL Server create_index mssql_includes

Fixed issue where usage of the SQL Server ``mssql_include`` option within a
:meth:`.Operations.create_index` would raise a KeyError, as the additional
column(s) need to be added to the table object used by the construct
internally.

Change-Id: If58fa35b9db8af473a9654e5a2c8861741810511
Fixes: #513

3 files changed, 27 insertions, 0 deletions

diff --git a/alembic/ddl/mssql.py b/alembic/ddl/mssql.py
index f10c5e6..bf6f5e9 100644
--- a/alembic/ddl/mssql.py
+++ b/alembic/ddl/mssql.py
@@ -6,6 +6,8 @@ from .base import alter_table, AddColumn, ColumnName, RenameTable,\
     format_table_name, format_column_name, ColumnNullable, alter_column,\
     format_server_default, ColumnDefault, format_type, ColumnType
 from sqlalchemy.sql.expression import ClauseElement, Executable
+from sqlalchemy.schema import CreateIndex, Column
+from sqlalchemy import types as sqltypes
 
 
 class MSSQLImpl(DefaultImpl):
@@ -87,6 +89,13 @@ class MSSQLImpl(DefaultImpl):
                 schema=schema,
                 name=name)
 
+    def create_index(self, index):
+        mssql_include = index.kwargs.get('mssql_include', ())
+        for col in mssql_include:
+            if col not in index.table.c:
+                index.table.append_column(Column(col, sqltypes.NullType))
+        self._exec(CreateIndex(index))
+
     def bulk_insert(self, table, rows, **kw):
         if self.as_sql:
             self._exec(
@@ -231,3 +240,4 @@ def visit_rename_table(element, compiler, **kw):
         format_table_name(compiler, element.table_name, element.schema),
         format_table_name(compiler, element.new_table_name, None)
     )
+
diff --git a/docs/build/unreleased/513.rst b/docs/build/unreleased/513.rst
new file mode 100644
index 0000000..87fe332
--- /dev/null
+++ b/docs/build/unreleased/513.rst
@@ -0,0 +1,8 @@
+.. change::
+   :tags: bug, mssql
+   :tickets: 513
+
+   Fixed issue where usage of the SQL Server ``mssql_include`` option within a
+   :meth:`.Operations.create_index` would raise a KeyError, as the additional
+   column(s) need to be added to the table object used by the construct
+   internally.
diff --git a/tests/test_mssql.py b/tests/test_mssql.py
index 805c23b..a5e013e 100644
--- a/tests/test_mssql.py
+++ b/tests/test_mssql.py
@@ -275,3 +275,12 @@ class OpTest(TestBase):
         context.assert_(
             "EXEC sp_rename 'y.t.c', x, 'COLUMN'"
         )
+
+    def test_create_index_mssql_include(self):
+        context = op_fixture('mssql')
+        op.create_index(
+            op.f('ix_mytable_a_b'), 'mytable', ['col_a', 'col_b'],
+            unique=False, mssql_include=['col_c'])
+        context.assert_contains(
+            "CREATE INDEX ix_mytable_a_b ON mytable "
+            "(col_a, col_b) INCLUDE (col_c)")