Fix using of 'mysql_length' for composite indexes

Currently, one can specify the prefix length for an index
column using 'mysql_length' keyword argument when creating
an Index instance. But in case of composite indexes the
prefix length value is applied only to the last column.

Extend the existing API in way so that 'mysql_length' argument
value can be either:
    - an integer specifying the same prefix length value
      for each column of an index
    - a (column_name --> integer value) mapping specifying
      the prefix length value for each column of an index
      separately

Fixes issue #2704.

2 files changed, 54 insertions, 7 deletions

diff --git a/lib/sqlalchemy/dialects/mysql/base.py b/lib/sqlalchemy/dialects/mysql/base.py
index 03827edb7..b7cff43eb 100644
--- a/lib/sqlalchemy/dialects/mysql/base.py
+++ b/lib/sqlalchemy/dialects/mysql/base.py
@@ -226,12 +226,15 @@ become part of the index. SQLAlchemy provides this feature via the
 ``mysql_length`` parameter::
 
     Index('my_index', my_table.c.data, mysql_length=10)
+    Index('a_b_idx', my_table.c.a, my_table.c.b, mysql_length={'a': 4, 'b': 9})
 
 Prefix lengths are given in characters for nonbinary string types and in bytes
-for binary string types. The value passed to the keyword argument will be
-simply passed through to the underlying CREATE INDEX command, so it *must* be
-an integer. MySQL only allows a length for an index if it is for a CHAR,
-VARCHAR, TEXT, BINARY, VARBINARY and BLOB.
+for binary string types. The value passed to the keyword argument *must* be
+either an integer (and, thus, specify the same prefix length value for all
+columns of the index) or a dict in which keys are column names and values are
+prefix length values for corresponding columns. MySQL only allows a length for
+a column of an index if it is for a CHAR, VARCHAR, TEXT, BINARY, VARBINARY and
+BLOB.
 
 Index Types
 ~~~~~~~~~~~~~
@@ -1519,12 +1522,27 @@ class MySQLDDLCompiler(compiler.DDLCompiler):
             text += "UNIQUE "
         text += "INDEX %s ON %s " % (name, table)
 
-        columns = ', '.join(columns)
         if 'mysql_length' in index.kwargs:
             length = index.kwargs['mysql_length']
-            text += "(%s(%d))" % (columns, length)
+
+            # length value can be an integer value specifying the same
+            # prefix length for all columns of the index
+            try:
+                columns = ', '.join(
+                    '%s(%d)' % (col, length)
+                    for col in columns
+                )
+            # otherwise it's a (column_name --> integer value) mapping
+            # specifying the prefix length for each column of the index
+            except TypeError:
+                columns = ', '.join(
+                    ('%s(%d)' % (col, length[col])
+                     if col in length else '%s' % col)
+                    for col in columns
+                )
         else:
-            text += "(%s)" % (columns)
+            columns = ', '.join(columns)
+        text += '(%s)' % columns
 
         if 'mysql_using' in index.kwargs:
             using = index.kwargs['mysql_using']
diff --git a/test/dialect/test_mysql.py b/test/dialect/test_mysql.py
index 2c459dead..18a37076c 100644
--- a/test/dialect/test_mysql.py
+++ b/test/dialect/test_mysql.py
@@ -50,6 +50,35 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
             'CREATE INDEX test_idx2 ON testtbl (data(5))',
             dialect=mysql.dialect())
 
+    def test_create_composite_index_with_length(self):
+        m = MetaData()
+        tbl = Table('testtbl', m,
+                    Column('a', String(255)),
+                    Column('b', String(255)))
+
+        idx1 = Index('test_idx1', tbl.c.a, tbl.c.b,
+                     mysql_length={'a': 10, 'b': 20})
+        idx2 = Index('test_idx2', tbl.c.a, tbl.c.b,
+                     mysql_length={'a': 15})
+        idx3 = Index('test_idx3', tbl.c.a, tbl.c.b,
+                     mysql_length=30)
+
+        self.assert_compile(
+            schema.CreateIndex(idx1),
+            'CREATE INDEX test_idx1 ON testtbl (a(10), b(20))',
+            dialect=mysql.dialect()
+        )
+        self.assert_compile(
+            schema.CreateIndex(idx2),
+            'CREATE INDEX test_idx2 ON testtbl (a(15), b)',
+            dialect=mysql.dialect()
+        )
+        self.assert_compile(
+            schema.CreateIndex(idx3),
+            'CREATE INDEX test_idx3 ON testtbl (a(30), b(30))',
+            dialect=mysql.dialect()
+        )
+
     def test_create_index_with_using(self):
         m = MetaData()
         tbl = Table('testtbl', m, Column('data', String(255)))