CODEWITHSUNDEEP
pythonpostgresqlsqlalchemy
glyphobet
glyphobet

Reputation: 1562

Combining PostgreSQL Enum with a TypeDecorator

Is there a way to have an Enum class that will both get created automatically when creating a new DB, like the standard Enum, but also connect it to a TypeDecorator that runs process_bind_param?

This first code block creates an Enum type that automatically lowercases input before saving it, but, unlike a normal Enum, the specific PostgreSQL enum type isn't created automatically in the database, so running a create_all() causes an error when creating the table, because the language_code type doesn't exist in the PostgreSQL schema. 

class LowercaseEnum(sqlalchemy.types.TypeDecorator):
    '''Converts input to lowercase.'''

    impl = sqlalchemy.types.Enum

    def process_bind_param(self, value, dialect):
        return value.lower()

class Foo(Model):
    id = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True)
    language_code = sqlalchemy.Column(LowercaseEnum(*['de', 'en'], name='language_code'))

If, on the other hand, I define my custom type like in this second code block, then it gets created automatically when calling create_all(), but the process_bind_param method is never called when sending values to the database, so the lowercasing doesn't work. 

class LowercaseEnum(sqlalchemy.dialects.postgresql.ENUM, sqlalchemy.types.TypeDecorator):
    '''Converts input to lowercase.'''

    impl = sqlalchemy.types.Enum

    def process_bind_param(self, value, dialect):
        return value.lower()

I have also tried several different combinations of inheriting from TypeDecorator in a separate class that is mixed-in, and also switching up sqlalchemy.types.Enum and sqlalchemy.dialects.postgresql.ENUM, but it seems no matter what I do, I either get a class that gets created automatically, or a class that runs process_bind_param, but never both.

Upvotes: 4

Views: 873

Answers (1)

glyphobet
glyphobet

Reputation: 1562

This appears to work, although it would still be great if anyone knows a better solution or sees a problem with this one. 

class LowercaseEnum(sqlalchemy.types.TypeDecorator, sqlalchemy.types.SchemaType):
    '''Converts input to lowercase.'''

    impl = sqlalchemy.dialects.postgresql.ENUM

    def _set_table(self, table, column):
        self.impl._set_table(table, column)

    def process_bind_param(self, value, dialect):
        return value.lower()

Upvotes: 2

Related Questions