diff --git a/platypush/entities/procedures.py b/platypush/entities/procedures.py index eef3427bbc..14c6ad14d3 100644 --- a/platypush/entities/procedures.py +++ b/platypush/entities/procedures.py @@ -1,8 +1,9 @@ import logging +from enum import Enum from sqlalchemy import ( Column, - Enum, + Enum as DbEnum, ForeignKey, Integer, JSON, @@ -16,6 +17,12 @@ from . import Entity logger = logging.getLogger(__name__) +class ProcedureType(Enum): + PYTHON = 'python' + CONFIG = 'config' + DB = 'db' + + if not is_defined('procedure'): class Procedure(Entity): @@ -30,7 +37,13 @@ if not is_defined('procedure'): ) args = Column(JSON, nullable=False, default=[]) procedure_type = Column( - Enum('python', 'config', name='procedure_type'), nullable=False + DbEnum( + *[m.value for m in ProcedureType.__members__.values()], + name='procedure_type', + create_constraint=True, + validate_strings=True, + ), + nullable=False, ) module = Column(String) source = Column(String)