Alembic: How to migrate custom type in a model?

问题

My User model is

class User(UserMixin, db.Model):
    __tablename__ = 'users'
    # noinspection PyShadowingBuiltins
    uuid = Column('uuid', GUID(), default=uuid.uuid4, primary_key=True,
                  unique=True)
    email = Column('email', String, nullable=False, unique=True)
    _password = Column('password', String, nullable=False)
    created_on = Column('created_on', sa.types.DateTime(timezone=True),
                        default=datetime.utcnow())
    last_login = Column('last_login', sa.types.DateTime(timezone=True),
                        onupdate=datetime.utcnow())

where GUID is a custom type as described in sqlalchemy docs (Exactly same)

Now when I run

alembic revision --autogenerate -m "Added initial table"

I get my upgrade() as

def upgrade():
    ### commands auto generated by Alembic - please adjust! ###
    op.create_table('users',
    sa.Column('uuid', sa.GUID(), nullable=False),
    sa.Column('email', sa.String(), nullable=False),
    sa.Column('password', sa.String(), nullable=False),
    sa.Column('created_on', sa.DateTime(timezone=True), nullable=True),
    sa.Column('last_login', sa.DateTime(timezone=True), nullable=True),
    sa.PrimaryKeyConstraint('uuid'),
    sa.UniqueConstraint('email'),
    sa.UniqueConstraint('uuid')
    )
    ### end Alembic commands ###

but during applying upgrade -> alembic upgrade head, I see

File "alembic/versions/49cc74d0da9d_added_initial_table.py", line 20, in upgrade
    sa.Column('uuid', sa.GUID(), nullable=False),
AttributeError: 'module' object has no attribute 'GUID'

How can I make it work with GUID/custom type here?

回答1:

You can replace sa.GUID() with either sa.CHAR(32) or UUID() (after adding the import line from sqlalchemy.dialects.postgresql import UUID) depending on the dialect.

Replacing it with GUID() (after adding the import line from your.models.custom_types import GUID) will work also, but then the upgrade script is tied to your model code, which may not be a good thing.

回答2:

I had a similar problem and solved it like follows:

Let's assume you have the following module my_guid, containing (from the page you already cited, with minor naming modifications):

import uuid as uuid_package
from sqlalchemy.dialects.postgresql import UUID as PG_UUID
from sqlalchemy import TypeDecorator, CHAR

class GUID(TypeDecorator):
    impl = CHAR

    def load_dialect_impl(self, dialect):
        if dialect.name == 'postgresql':
            return dialect.type_descriptor(PG_UUID())
        else:
            return dialect.type_descriptor(CHAR(32))

    def process_bind_param(self, value, dialect):
        if value is None:
            return value
        elif dialect.name == 'postgresql':
            return str(value)
        else:
            if not isinstance(value, uuid_package.UUID):
                return "%.32x" % uuid_package.UUID(value)
            else:
                # hexstring
                return "%.32x" % value

    def process_result_value(self, value, dialect):
        if value is None:
            return value
        else:
            return uuid_package.UUID(value)

If you use this GUID in your models, you just need to add three lines at alembic/env.py:

from my_guid import GUID
import sqlalchemy as sa
sa.GUID = GUID

That worked for me. Hope that helps!

回答3:

Using the __repr__ function of the impl attribute class worked for me for most custom types. I find it cleaner to have the migration definition contained inside class instead worrying about putting imports in your env.py or scripts.py.mako. Plus, it makes it easy to move your code between modules.

Class GUID(types.TypeDecorator)
    impl = CHAR

    def __repr__(self):
        return self.impl.__repr__()

    # You type logic here.

The automigration will produce CHAR(length=XXX).

回答4:

My solution uses sqlalchemy_utils.types.uuid.UUIDType, which uses CHAR(32) or BINARY(16) to represent the UUID if you are on a database without a UUID type. You need to account for this in your migration, which must create a CHAR(32)/BINARY(16) on a DB without a UUID type and a UUIDType on databases with it.

My SQLAlchemy class looks like this:

from sqlalchemy_utils.types.uuid import UUIDType
from sqlalchemy import CHAR, Column, Integer

Base = declarative_base()

def get_uuid():
    return str(uuid.uuid4())

class Dashboard(Base):
    __tablename__ = 'dashboards'
    id = Column(Integer, primary_key=True)
    uuid = Column(UUIDType(binary=False), default=get_uuid)

and the actual batch operation looks like this (which supports SQLite, MySQL and Postgres):

from superset import db # Sets up a SQLAlchemy connection

def upgrade():
    bind = op.get_bind()
    session = db.Session(bind=bind)
    db_type = session.bind.dialect.name

    def add_uuid_column(col_name, _type):
        """Add a uuid column to a given table"""
        with op.batch_alter_table(col_name) as batch_op:
            batch_op.add_column(Column('uuid', UUIDType(binary=False), default=get_uuid))
        for s in session.query(_type):
            s.uuid = get_uuid()
            session.merge(s)

        if db_type != 'postgresql':
            with op.batch_alter_table(col_name) as batch_op:
                batch_op.alter_column('uuid', existing_type=CHAR(32),
                                      new_column_name='uuid', nullable=False)
                batch_op.create_unique_constraint('uq_uuid', ['uuid'])

        session.commit()

add_uuid_column('dashboards', Dashboard)
session.close()

Hope this helps!

来源：https://stackoverflow.com/questions/15668115/alembic-how-to-migrate-custom-type-in-a-model

标签

python

sqlalchemy

flask-sqlalchemy

alembic