Alembic: How to migrate custom type in a model?

Question 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.

Question 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!

Question 3

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!

Question 4

Following-up on the solution from @Red-Tune-84

class GUID(types.TypeDecorator)
  impl = CHAR

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

  # You type logic here.

does work but you might need in parallel to set in env.py the configuration user_module_prefix.

For instance, with context.configure(..., user_module_prefix="sa."), the type above will appear as sa.CHAR(...) in the alembic migration.

Question 5

Short answer (using sqlalchemy version 1.4.25):

From the documentation:

For user-defined types, that is, any custom type that is not within the sqlalchemy. module namespace, by default Alembic will use the value of __module__ for the custom type:
Column("my_column", myapp.models.utils.types.MyCustomType())
The imports for the above type again must be made present within the migration, either manually, or by adding it to script.py.mako.

So, import your my_module in script.py.mako (and you probably need to define your custom type in a file other than models.py):

...
from alembic import op
import sqlalchemy as sa
import my_module
${imports if imports else ""}
...

Long answer:

I have my custom type BinaryUUID defined in uuid_type_mysql.py and I'm importing it in models.py and using it there:

models.py:

...
from .uuid_type_mysql import BinaryUUID
...

After generating the migrations with flask db migrate I would get this on the migration file:

...
sa.Column('public_id', my_module.uuid_type_mysql.BinaryUUID(length=16), nullable=False),
...

And the problem is that the migration does not know my_module because it is not imported.

After adding import my_module to script.py.mako as suggested in the documentation, now the module is imported in the migration file:

...
from alembic import op
import sqlalchemy as sa
import my_module
...

After this, everything work fine for me with the bonus that it is not needed to manually edit the generated migration.

Question 6

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).