Alembic Migration: ChoiceType failing with Enum

[sreeram-dev]

from sqlalchemy_utils import PasswordType, PhoneNumberType, PhoneNumber, ChoiceType
from sqlalchemy_utils import EmailType

from sse.app import db, login

import sse.user.constants as UserConstants

class StatesEnum(enum.Enum):
    NEW_SOUTH_WALES = 1
    VICTORIA = 2
    SOUTH_AUSTRALIA = 3
    WESTERN_AUSTRALIA = 4
    TASMANIA = 5
    QUEENSLAND = 6

class User(UserMixin, db.Model):
    __abstract__ = True
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(128), unique=True)
    password = db.Column(PasswordType(
        schemes=[
            'pbkdf2_sha512',
            'md5_crypt'
        ],
        deprecated=['md5_crypt']))
    email  = db.Column(EmailType)
    # for 2-fa
    phone = db.Column(PhoneNumberType(region='AU'))
    state = db.Column(ChoiceType(StatesEnum, impl=db.Integer()))
    first_name = db.Column(db.String(15))
    last_name = db.Column(db.String(15))

    def get_hashed_password(self):
        return self.password

sa.Column('state', sqlalchemy_utils.types.choice.ChoiceType(), nullable=True), TypeError: <flask_script.commands.Command object at 0x7f00d761bdd8>: init() missing 1 required positional argument: 'choices'

[esnowkropla]

In case anyone else runs into this, I solved this by passing the enum into ChoiceType in the migration:

from this:

op.add_column('table',
              sa.Column('my_enum',
              sqlalchemy_utils.types.choice.ChoiceType(),
              nullable=False))

to this:

from app.models.enums import MyType

op.add_column('table',
              sa.Column('my_enum',
              sqlalchemy_utils.types.choice.ChoiceType(MyType),
              nullable=False))

[CodePint]

further improvement/work around. Also see: https://stackoverflow.com/questions/30132370/trouble-when-using-alembic-with-sqlalchemy-utils

my fix does not work with enum, but im sure it could be tweaked. with a bit of work, this could fix the issue for most custom column types

def render_item(type_, obj, autogen_context):
    # custom render for sqalchemy_utils ChoiceType column and params
    # May be better to use: https://stackoverflow.com/questions/30132370/trouble-when-using-alembic-with-sqlalchemy-utils
    if type_ == 'type' and isinstance(obj, sqlalchemy_utils.types.choice.ChoiceType):
        autogen_context.imports.add("import sqlalchemy_utils")
        col_type = "sqlalchemy_utils.types.choice.ChoiceType({})"
        return col_type.format(obj.choices)

    # default rendering for other objects
    return False

def run_migrations_online():
    def process_revision_directives(context, revision, directives):
        if getattr(config.cmd_opts, 'autogenerate', False):
            script = directives[0]
            if script.upgrade_ops.is_empty():
                directives[:] = []
                logger.info('No changes in schema detected.')

    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix='sqlalchemy.',
        poolclass=pool.NullPool,
    )

    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            process_revision_directives=process_revision_directives,
            # custom object rendering
            render_item=render_item,
            **current_app.extensions['migrate'].configure_args
        )

        with context.begin_transaction():
            context.run_migrations()

[CodePint]

Revisited this a year late and solved this for enum types dynamically.

# migrations/imports.py

from models.enums import *

# migrations/env.py

# https://alembic.sqlalchemy.org/en/latest/autogenerate.html#affecting-the-rendering-of-types-themselves
def render_item(type_, obj, autogen_context):
    """Apply rendering for custom sqlalchemy types"""
    if type_ == "type":
        module_name =  obj.__class__.__module__
        if module_name.startswith("sqlalchemy_utils."):
            return render_sqlalchemy_utils_type(obj, autogen_context)

    # render default
    return False

def render_sqlalchemy_utils_type(obj, autogen_context):
    class_name = obj.__class__.__name__
    import_statement = f"from sqlalchemy_utils.types import {class_name}"
    autogen_context.imports.add(import_statement)
    if class_name == 'ChoiceType':
        return render_choice_type(obj, autogen_context)
    return f"{class_name}()"

def render_choice_type(obj, autogen_context):
    choices = obj.choices
    if obj.type_impl.__class__.__name__ == 'EnumTypeImpl':
        choices = obj.type_impl.enum_class.__name__
        import_statement = f"from migrations.imports import {choices}"
        autogen_context.imports.add(import_statement)
    impl_stmt = f"sa.{obj.impl.__class__.__name__}()"
    return f"{obj.__class__.__name__}(choices={choices}, impl={impl_stmt})"

def run_migrations_online():
    """Run migrations in 'online' mode.

    In this scenario we need to create an Engine
    and associate a connection with the context.

    """

    # this callback is used to prevent an auto-migration from being generated
    # when there are no changes to the schema
    # reference: http://alembic.zzzcomputing.com/en/latest/cookbook.html
    def process_revision_directives(context, revision, directives):
        if getattr(config.cmd_opts, 'autogenerate', False):
            script = directives[0]
            if script.upgrade_ops.is_empty():
                directives[:] = []
                logger.info('No changes in schema detected.')

    connectable = current_app.extensions['migrate'].db.engine

    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            render_item=render_item,
            target_metadata=target_metadata,
            process_revision_directives=process_revision_directives,
            **current_app.extensions['migrate'].configure_args
        )

        with context.begin_transaction():
            context.run_migrations()

[Tamunonengiyeofori]

Please I am having issues migrating some DB tables with sqlalchemy and alembic because of a foreign key error. I need help to resolve it.

[kurtmckee]

@Tamunonengiyeofori, this doesn't appear to be related to this issue, nor to sqlalchemy-utils in general.

I recommend searching the internet using keywords like "sqlalchemy NoReferencedTableError".