From 4256d41645fbbf4a5d4b5245572048b53549fec3 Mon Sep 17 00:00:00 2001
From: Kaxil Naik <kaxilnaik@gmail.com>
Date: Mon, 30 Aug 2021 19:05:08 +0100
Subject: [PATCH] Fix ``DagRunState`` enum query for ``MySQLdb`` driver
 (#17886)

same as https://github.com/apache/airflow/pull/13278 but for `DagRunState` introduced in https://github.com/apache/airflow/pull/16854

closes https://github.com/apache/airflow/issues/17879

GitOrigin-RevId: a3f9c690aa80d12ff1d5c42eaaff4fced07b9429
---
 airflow/utils/state.py    |  3 ++
 tests/utils/test_state.py | 58 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 tests/utils/test_state.py

diff --git a/airflow/utils/state.py b/airflow/utils/state.py
index f84774844f7..39e89b56a96 100644
--- a/airflow/utils/state.py
+++ b/airflow/utils/state.py
@@ -67,6 +67,9 @@ class DagRunState(str, Enum):
     SUCCESS = "success"
     FAILED = "failed"
 
+    def __str__(self) -> str:
+        return self.value
+
 
 class State:
     """
diff --git a/tests/utils/test_state.py b/tests/utils/test_state.py
new file mode 100644
index 00000000000..88e4a0db20d
--- /dev/null
+++ b/tests/utils/test_state.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from airflow.models import DAG
+from airflow.models.dagrun import DagRun
+from airflow.utils.session import create_session
+from airflow.utils.state import DagRunState
+from airflow.utils.types import DagRunType
+from tests.models import DEFAULT_DATE
+
+
+def test_dagrun_state_enum_escape():
+    """
+    Make sure DagRunState.QUEUED is converted to string 'queued' when
+    referenced in DB query
+    """
+    with create_session() as session:
+        dag = DAG(dag_id='test_dagrun_state_enum_escape', start_date=DEFAULT_DATE)
+        dag.create_dagrun(
+            run_type=DagRunType.SCHEDULED,
+            state=DagRunState.QUEUED,
+            execution_date=DEFAULT_DATE,
+            start_date=DEFAULT_DATE,
+            session=session,
+        )
+
+        query = session.query(DagRun.dag_id, DagRun.state, DagRun.run_type,).filter(
+            DagRun.dag_id == dag.dag_id,
+            # make sure enum value can be used in filter queries
+            DagRun.state == DagRunState.QUEUED,
+        )
+        assert str(query.statement.compile(compile_kwargs={"literal_binds": True})) == (
+            'SELECT dag_run.dag_id, dag_run.state, dag_run.run_type \n'
+            'FROM dag_run \n'
+            "WHERE dag_run.dag_id = 'test_dagrun_state_enum_escape' AND dag_run.state = 'queued'"
+        )
+
+        rows = query.all()
+        assert len(rows) == 1
+        assert rows[0].dag_id == dag.dag_id
+        # make sure value in db is stored as `queued`, not `DagRunType.QUEUED`
+        assert rows[0].state == 'queued'
+
+        session.rollback()