13

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?

Florian Pilz
  • 8,002
  • 4
  • 22
  • 30
daydreamer
  • 87,243
  • 191
  • 450
  • 722

6 Answers6

9

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.

sayap
  • 6,169
  • 2
  • 36
  • 40
  • but that would break in tests if I am using `sqlite` plus it wont be consistent, isn't it? – daydreamer Mar 27 '13 at 20:50
  • 1
    I really don't care much for using `sqlite` in tests, since the test environment should closely follow the production. Other than that, are you deploying the same code against more than one dialects? – sayap Mar 27 '13 at 21:05
  • Anyway, I would want the migration scripts to be consistent with the actual schema at the point of time, consistency with the code be damned since the code is forever changing. – sayap Mar 27 '13 at 21:09
  • I would think I would care for sqlite for my tests since they can be run in-memory, but what you said also makes a valid point, let me try what you said and reply you – daydreamer Mar 27 '13 at 21:47
  • I agree with you completely and have PostgreSQL for both `dev` and `test` environments and your answer helped me a lot @sayap. Thank you very much – daydreamer Mar 30 '13 at 05:49
  • Glad it helps. We should all just be grateful that we got the amazing alembic library without having to sponsor some kickstarter campaign, like what the Django guys did :P – sayap Mar 30 '13 at 07:41
4

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!

Philipp der Rautenberg
  • 2,212
  • 3
  • 25
  • 39
2

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!

rjurney
  • 4,824
  • 5
  • 41
  • 62
2

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.

Thierry BM
  • 21
  • 1
  • This was a key solution for me. In my case the problem was that the rendered `Column` migration was my `TypeDecorator` rather than the underlying PostgreSQL type with the appropriate arguments. Adding the `__repr__` override fixed that. In addition, I need to make sure my root package was imported in my `script.py.mako` – jsnow Aug 21 '23 at 18:31
1

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.

David
  • 356
  • 4
  • 6
0

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

Red-Tune-84
  • 381
  • 6
  • 17
  • 2
    I don't think this works alone, at least with Alembic 1.4.3 and SQLAlchemy ORM. Assuming your custom types are defined or imported in `app.models`. The resulting migration will just have `app.models.CHAR` instead of `app.models.GUID`. Maybe the default rendering changed since this answer? – Z. Kimble Nov 04 '20 at 16:15