From 0fdfe373f0afb5ba4c0bf1c632a6a46ccd165a1a Mon Sep 17 00:00:00 2001
From: xiaoyu meng <xiaoyu.meng@airbnb.com>
Date: Mon, 9 Sep 2024 12:01:55 -0700
Subject: [PATCH] fix: handle DESCRIBE SELECT (#61)

* fix: handle DESCRIBE SELECT

* format

* remove special handling in conftest.query

* format
---
 mysql_mimic/session.py |  3 +++
 tests/conftest.py      |  2 +-
 tests/test_query.py    | 17 +++++++++++++++--
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/mysql_mimic/session.py b/mysql_mimic/session.py
index 767cfc0..0fdb119 100644
--- a/mysql_mimic/session.py
+++ b/mysql_mimic/session.py
@@ -393,6 +393,9 @@ async def _show(self, expression: exp.Show) -> AllowedResult:
     async def _describe_middleware(self, q: Query) -> AllowedResult:
         """Intercept DESCRIBE statements"""
         if isinstance(q.expression, exp.Describe):
+            if isinstance(q.expression.this, exp.Select):
+                # Mysql parse treats EXPLAIN SELECT as a DESCRIBE SELECT statement
+                return await q.next()
             name = q.expression.this.name
             show = self.dialect().parse(f"SHOW COLUMNS FROM {name}")[0]
             return await self._show(show) if isinstance(show, exp.Show) else None
diff --git a/tests/conftest.py b/tests/conftest.py
index 297d1f3..27c2f63 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -77,7 +77,7 @@ async def query(
             self.waiting.set()
             await self.pause.wait()
             self.waiting.clear()
-        assert isinstance(expression, exp.Select)
+
         self.last_query_attrs = attrs
         if self.echo:
             return [(sql,)], ["sql"]
diff --git a/tests/test_query.py b/tests/test_query.py
index 5df6281..3eb664b 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -73,10 +73,10 @@ async def q4(sql: str) -> Sequence[Dict[str, Any]]:
     raise RuntimeError("Unexpected fixture param")
 
 
-# # Uncomment to make tests only use mysql-connector, which can help during debugging
+# Uncomment to make tests only use mysql-connector, which can help during debugging
 # @pytest_asyncio.fixture
 # async def query_fixture(
-#     mysql_connector_conn: MySQLConnection,
+#     mysql_connector_conn: MySQLConnectionAbstract,
 # ) -> QueryFixture:
 #     async def q1(sql: str) -> Sequence[Dict[str, Any]]:
 #         return await query(mysql_connector_conn, sql)
@@ -293,6 +293,19 @@ async def test_query_attributes(
     assert session.last_query_attrs == query_attrs
 
 
+@pytest.mark.asyncio
+async def test_describe_select(
+    session: MockSession,
+    server: MysqlServer,
+    query_fixture: QueryFixture,
+) -> None:
+    session.echo = True
+    sql = "DESCRIBE SELECT b from a"
+    with freeze_time("2023-01-01"):
+        result = await query_fixture(sql)
+        assert [{"sql": "DESCRIBE SELECT b from a"}] == list(result)
+
+
 # pylint: disable=trailing-whitespace
 @pytest.mark.asyncio
 @pytest.mark.parametrize(