mssql_include in Index creation is triggering KeyError

[sqlalchemy-bot]

Migrated issue, originally created by Florian Vichot (@fvichot)

Hi,

This is on SQL Server 2014, using pyodbc and Microsoft's ODBC driver.

When running a migration containing only the following:

def upgrade():
    op.create_index(op.f('ix_mytable_a_b'), 'mytable', ['col_a', 'col_b'],
                    unique=False, mssql_include=['col_c'])

I get the following stack trace:

Traceback (most recent call last):
  File "/usr/local/bin/flask", line 11, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.6/site-packages/flask/cli.py", line 894, in main
    cli.main(args=args, prog_name=name)
  File "/usr/local/lib/python3.6/site-packages/flask/cli.py", line 557, in main
    return super(FlaskGroup, self).main(*args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/click/core.py", line 697, in main
    rv = self.invoke(ctx)
  File "/usr/local/lib/python3.6/site-packages/click/core.py", line 1066, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/usr/local/lib/python3.6/site-packages/click/core.py", line 1066, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/usr/local/lib/python3.6/site-packages/click/core.py", line 895, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/usr/local/lib/python3.6/site-packages/click/core.py", line 535, in invoke
    return callback(*args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/click/decorators.py", line 17, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/flask/cli.py", line 412, in decorator
    return __ctx.invoke(f, *args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/click/core.py", line 535, in invoke
    return callback(*args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/flask_migrate/cli.py", line 134, in upgrade
    _upgrade(directory, revision, sql, tag, x_arg)
  File "/usr/local/lib/python3.6/site-packages/flask_migrate/__init__.py", line 95, in wrapped
    f(*args, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/flask_migrate/__init__.py", line 280, in upgrade
    command.upgrade(config, revision, sql=sql, tag=tag)
  File "/usr/local/lib/python3.6/site-packages/alembic/command.py", line 254, in upgrade
    script.run_env()
  File "/usr/local/lib/python3.6/site-packages/alembic/script/base.py", line 427, in run_env
    util.load_python_file(self.dir, 'env.py')
  File "/usr/local/lib/python3.6/site-packages/alembic/util/pyfiles.py", line 81, in load_python_file
    module = load_module_py(module_id, path)
  File "/usr/local/lib/python3.6/site-packages/alembic/util/compat.py", line 82, in load_module_py
    spec.loader.exec_module(module)
  File "<frozen importlib._bootstrap_external>", line 678, in exec_module
  File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed
  File "migrations/env.py", line 89, in <module>
    run_migrations_online()
  File "migrations/env.py", line 82, in run_migrations_online
    context.run_migrations()
  File "<string>", line 8, in run_migrations
  File "/usr/local/lib/python3.6/site-packages/alembic/runtime/environment.py", line 836, in run_migrations
    self.get_context().run_migrations(**kw)
  File "/usr/local/lib/python3.6/site-packages/alembic/runtime/migration.py", line 330, in run_migrations
    step.migration_fn(**kw)
  File "/app/migrations/versions/d4b81f83b5ab_XXXXXXXX.py", line 23, in upgrade
    unique=False, mssql_include=['col_c'])
  File "<string>", line 8, in create_index
  File "<string>", line 3, in create_index
  File "/usr/local/lib/python3.6/site-packages/alembic/operations/ops.py", line 857, in create_index
    return operations.invoke(op)
  File "/usr/local/lib/python3.6/site-packages/alembic/operations/base.py", line 319, in invoke
    return fn(self, operation)
  File "/usr/local/lib/python3.6/site-packages/alembic/operations/toimpl.py", line 88, in create_index
    operations.impl.create_index(idx)
  File "/usr/local/lib/python3.6/site-packages/alembic/ddl/impl.py", line 201, in create_index
    self._exec(schema.CreateIndex(index))
  File "/usr/local/lib/python3.6/site-packages/alembic/ddl/mssql.py", line 23, in _exec
    result = super(MSSQLImpl, self)._exec(construct, *args, **kw)
  File "/usr/local/lib/python3.6/site-packages/alembic/ddl/impl.py", line 115, in _exec
    return conn.execute(construct, *multiparams, **params)
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/engine/base.py", line 948, in execute
    return meth(self, multiparams, params)
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/sql/ddl.py", line 68, in _execute_on_connection
    return connection._execute_ddl(self, multiparams, params)
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/engine/base.py", line 1003, in _execute_ddl
    if not self.schema_for_object.is_default else None)
  File "<string>", line 1, in <lambda>
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/sql/elements.py", line 442, in compile
    return self._compiler(dialect, bind=bind, **kw)
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/sql/ddl.py", line 26, in _compiler
    return dialect.ddl_compiler(dialect, self, **kw)
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/sql/compiler.py", line 219, in __init__
    self.string = self.process(self.statement, **compile_kwargs)
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/sql/compiler.py", line 245, in process
    return obj._compiler_dispatch(self, **kwargs)
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/sql/visitors.py", line 81, in _compiler_dispatch
    return meth(self, **kw)
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/dialects/mssql/base.py", line 1654, in visit_create_index
    index.dialect_options['mssql']['include']
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/dialects/mssql/base.py", line 1653, in <listcomp>
    for col in
  File "/usr/local/lib/python3.6/site-packages/sqlalchemy/util/_collections.py", line 194, in __getitem__
    return self._data[key]
KeyError: 'col_c'

I've inserted some print statements in sqlalchemy/dialects/mssql/base.py just before the crash, and indeed the column 'col_c' is missing from the table model (ie. not in the index.table.c list). But if I run the following in my create_app() before the Migrate(app, db) (from Flask-Migrate which wraps Alembic for Flask):

from .models import mytable
print(mytable.__table_args__[0].table.c)

I get the full list of columns:
["col_a", "col_b", "col_c"]

My table is defined as follows:

class mytable(db.Model):
    device_name = db.Column(db.String(50), nullable=False)

    __table_args__ = (db.Index("ix_mytable_a_b", "col_a", "col_b", mssql_include=["col_c"]),

Versions of components:

alembic==1.0.0
SQLAlchemy==1.2.11

And probably not relevant, but just in case:

Flask==1.0.2
Flask-Migrate==2.2.1
Flask-SQLAlchemy==2.3.2

[sqlalchemy-bot]

Michael Bayer (@zzzeek) wrote:

here's a workaround for the moment

        from sqlalchemy import Table, MetaData, Index
        from sqlalchemy.types import NullType
        from sqlalchemy.schema import CreateIndex

        t = Table(
            'mytable', MetaData(),
            Column('col_a', NullType), Column('col_b', NullType),
            Column('col_c', NullType)
        )

        idx = Index('ix_mytable_a_b', 'col_a', 'col_b', unique=False, mssql_include=['col_c'])
        t.append_constraint(idx)

        op.execute(CreateIndex(idx))

[sqlalchemy-bot]

Michael Bayer (@zzzeek) wrote:

https://gerrit.sqlalchemy.org/#/c/zzzeek/alembic/+/899/

[sqlalchemy-bot]

Changes by Michael Bayer (@zzzeek):

• added labels: op directives

[sqlalchemy-bot]

Changes by Michael Bayer (@zzzeek):

• set milestone to "fasttrack"

[sqlalchemy-bot]

Michael Bayer (@zzzeek) wrote:

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

→ e01041b

[sqlalchemy-bot]

Changes by Michael Bayer (@zzzeek):

• changed status to closed

[sqlalchemy-bot]

Florian Vichot (@fvichot) wrote:

Awesome, thank you for the quick fix!