How to query a table that has ENUM column and keep the ENUM type?

Viewed 713

I'm using SQLAlchemy ORM.

I have a table in SQL DB with an id column, and a column called b, which is type enum, and can take values ('example_1', 'example_2').

In Python, I have an Enum class like this:

class BTypes(enum.Enum):
    EXAMPLE_1 = 'example_1'
    EXAMPLE_2 = 'example_2'

For querying the table, I have an ORM like this:

class Example(Base):
    __tablename__ = "example"
    id = Column(Integer, primary_key=True)
    b = Column(Enum(BTypes).values_callable)

When I do session.query(Example).all(), the objects that I get back have str type for the b attribute. In other words:

data = session.query(Example).all()
print(data[0].b)
# Outputs
# example_1

I want that the Example object for the attribute b has an enum type, not str. What is the best way to achieve this?

2 Answers

Base.metadata.create_all(create_engine("sqlite://")) with:

b = Column(Enum(BTypes).values_callable)

gives me:

sqlalchemy.exc.CompileError: (in table 'example', column 'b'): Can't generate DDL for NullType(); did you forget to specify a type on this Column?

About NullType

Since Enum(BTypes).values_callable is None, SQLAlchemy defaults to NullType.

From https://docs.sqlalchemy.org/en/14/core/type_api.html#sqlalchemy.types.NullType:

The NullType can be used within SQL expression invocation without issue, it just has no behavior either at the expression construction level or at the bind-parameter/result processing level.

In other words, when we query, its value is simply assigned as-is from the database.

How to use the Enum.values_callable parameter

From https://docs.sqlalchemy.org/en/14/core/type_basics.html#sqlalchemy.types.Enum:

In order to persist the values and not the names, the Enum.values_callable parameter may be used. The value of this parameter is a user-supplied callable, which is intended to be used with a PEP-435-compliant enumerated class and returns a list of string values to be persisted. For a simple enumeration that uses string values, a callable such as lambda x: [e.value for e in x] is sufficient.

That would be:

b = Column(Enum(BTypes, values_callable=lambda x: [e.value for e in x]))

Modify your code to query table as below to get as Enum:

   class Example(Base):
        __tablename__ = "example"
         id = Column(Integer, primary_key=True)
         b = Column(Enum(BTypes))

Note that values_callable typically return list of string values. Do have a look into the documentation for more information

values_callable – A callable which will be passed the PEP-435 compliant enumerated type, which should then return a list of string values to be persisted. This allows for alternate usages such as using the string value of an enum to be persisted to the database instead of its name.

Related