houqp commented on a change in pull request #13278:
URL: https://github.com/apache/airflow/pull/13278#discussion_r548452559



##########
File path: airflow/utils/types.py
##########
@@ -16,8 +16,31 @@
 # under the License.
 import enum
 
+from sqlalchemy.types import TypeDecorator, String
 
-class DagRunType(str, enum.Enum):
+
+class EnumString(TypeDecorator):
+    """
+    Declare db column with this type to make the column compatible with string
+    and string based enum values when building the sqlalchemy ORM query. It can
+    be used just like sqlalchemy.types.String, for example:
+
+    ```
+    class Table(Base):
+        __tablename__ = "t"
+        run_type = Column(EnumString(50), nullable=False)
+    ```
+    """
+    impl = String
+
+    def process_bind_param(self, value, dialect):
+        if isinstance(value, enum.Enum):
+            return value.value
+        else:
+            return value
+
+
+class DagRunType(enum.Enum):

Review comment:
       Another way to address this problem is to fix it at db driver layer. 
i.e. changing mysqlclient to support enum types. If we can get this change 
merged upstream, then we can at least safely use str subclassed enum type for 
both postgres and mysql everywhere, but since the conversion is not done within 
sqlalchemy anymore, we need to double check other db drivers.
   
   However, adding enum support to mysqlclient will take a bit more time. If we 
don't need to push this fix out sooner, then we can probably wait for that.
   
   Or we can go with custom `EnumString` as a short term fix, then work on 
mysqlclient path as a follow up.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to