feat: Improve EXPLAIN support for MySQL and PostgreSQL

CHANGELOG.rst

@@ -11,6 +11,11 @@ Changelog
 1.1.8
 -----
 
+Added
+^^^^^
+- ``QuerySet.union()`` — SQL UNION query support for combining results from multiple QuerySets, including support for union across different models, ``union(all=True)`` for duplicates, ``order_by()``, ``limit()``, and ``count()``.
+- Added comprehensive EXPLAIN support for MySQL and PostgreSQL.
+
 Fixed
 ^^^^^
 - ``MigrationRecorder`` now uses parameterized queries; fixes MariaDB/MySQL rejecting ISO-8601 ``applied_at`` values. (#2132)

tests/backends/test_explain.py

@@ -2,7 +2,8 @@
 
 from tests.testmodels import Tournament
 from tortoise.contrib.test import requireCapability
-from tortoise.contrib.test.condition import NotEQ
+from tortoise.contrib.test.condition import NotEQ, NotIn
+from tortoise.exceptions import UnSupportedError
 
 
 @requireCapability(dialect=NotEQ("mssql"))
@@ -18,3 +19,19 @@ async def test_explain(db):
     plan = await Tournament.all().explain()
     # This should have returned *some* information.
     assert len(str(plan)) > 20
+
+
+@requireCapability(dialect=NotIn("postgres", "mysql", "mssql"))
+@pytest.mark.asyncio
+async def test_explain_unsupported_output_fmt(db):
+    await Tournament.create(name="Test")
+    with pytest.raises(UnSupportedError, match="does not support different explain formats"):
+        await Tournament.all().explain(output_fmt="json")
+
+
+@requireCapability(dialect=NotIn("postgres", "mysql", "mssql"))
+@pytest.mark.asyncio
+async def test_explain_unsupported_options(db):
+    await Tournament.create(name="Test")
+    with pytest.raises(UnSupportedError, match="does not support explain options"):
+        await Tournament.all().explain(analyze=True)

tests/backends/test_mysql.py

@@ -3,13 +3,17 @@
 """
 
 import copy
+import json
 import os
 import ssl
 
 import pytest
 
+from tests.testmodels import Tournament
 from tortoise.backends.base.config_generator import generate_config
 from tortoise.context import TortoiseContext
+from tortoise.contrib.test import requireCapability
+from tortoise.exceptions import UnSupportedError
 
 
 def _get_db_config():
@@ -87,3 +91,65 @@ async def test_ssl_custom():
             await ctx.init(db_config, _create_db=True)
         except ConnectionError:
             pass
+
+
+@requireCapability(dialect="mysql")
+@pytest.mark.asyncio
+async def test_explain(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain()
+    data = json.loads(result[0]["EXPLAIN"])
+    assert "query_plan" in data or "query_block" in data
+
+
+@requireCapability(dialect="mysql")
+@pytest.mark.asyncio
+async def test_explain_format_traditional(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(output_fmt="traditional")
+    assert "table" in result[0]
+    assert result[0]["table"] == "tournament"
+
+
+@requireCapability(dialect="mysql")
+@pytest.mark.asyncio
+async def test_explain_format_tree(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(output_fmt="tree")
+    assert isinstance(result[0]["EXPLAIN"], str)
+    assert "->" in result[0]["EXPLAIN"]
+    assert "tournament" in result[0]["EXPLAIN"]
+
+
+@requireCapability(dialect="mysql")
+@pytest.mark.asyncio
+async def test_explain_analyze(db_simple):
+    await Tournament.create(name="Test")
+    # Older MySQL version don't support ANALYZE with JSON format, that's why we use TREE
+    result = await Tournament.all().explain(output_fmt="tree", analyze=True)
+    assert "actual" in result[0]["EXPLAIN"]
+
+
+@requireCapability(dialect="mysql")
+@pytest.mark.asyncio
+async def test_explain_analyze_false(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(analyze=False)
+    assert "query_plan" in result[0]["EXPLAIN"] or "query_block" in result[0]["EXPLAIN"]
+    assert "actual" not in result[0]["EXPLAIN"]
+
+
+@requireCapability(dialect="mysql")
+@pytest.mark.asyncio
+async def test_explain_unsupported_format(db_simple):
+    await Tournament.create(name="Test")
+    with pytest.raises(UnSupportedError, match="Unsupported explain format"):
+        await Tournament.all().explain(output_fmt="invalid")
+
+
+@requireCapability(dialect="mysql")
+@pytest.mark.asyncio
+async def test_explain_unsupported_option(db_simple):
+    await Tournament.create(name="Test")
+    with pytest.raises(UnSupportedError, match="Unsupported options"):
+        await Tournament.all().explain(unsupported_option=True)

tests/backends/test_postgres.py

@@ -2,15 +2,19 @@
 Test some PostgreSQL-specific features
 """
 
+import json
 import os
 import ssl
+import xml.etree.ElementTree as ET
 
 import pytest
+import yaml
 
 from tests.testmodels import Tournament
 from tortoise import Tortoise, connections
 from tortoise.backends.base.config_generator import generate_config
-from tortoise.exceptions import OperationalError
+from tortoise.contrib.test import requireCapability
+from tortoise.exceptions import OperationalError, UnSupportedError
 
 
 def _get_db_config():
@@ -28,11 +32,10 @@ def _get_db_config():
     return db_config, is_asyncpg, is_psycopg
 
 
+@requireCapability(dialect="postgres")
 @pytest.mark.asyncio
-async def test_schema(db_simple):
-    db_config, is_asyncpg, is_psycopg = _get_db_config()
-    if not is_asyncpg and not is_psycopg:
-        pytest.skip("PostgreSQL only")
+async def test_schema(db_isolated):
+    db_config, is_asyncpg, _ = _get_db_config()
 
     if is_asyncpg:
         from asyncpg.exceptions import InvalidSchemaNameError
@@ -75,11 +78,10 @@ async def test_schema(db_simple):
             await Tortoise._drop_databases()
 
 
+@requireCapability(dialect="postgres")
 @pytest.mark.asyncio
-async def test_ssl_true():
-    db_config, is_asyncpg, is_psycopg = _get_db_config()
-    if not is_asyncpg and not is_psycopg:
-        pytest.skip("PostgreSQL only")
+async def test_ssl_true(db_isolated):
+    db_config, _, _ = _get_db_config()
 
     db_config["connections"]["models"]["credentials"]["ssl"] = True
     ssl_failed = False
@@ -95,11 +97,10 @@ async def test_ssl_true():
             await Tortoise._drop_databases()
 
 
+@requireCapability(dialect="postgres")
 @pytest.mark.asyncio
-async def test_ssl_custom():
-    db_config, is_asyncpg, is_psycopg = _get_db_config()
-    if not is_asyncpg and not is_psycopg:
-        pytest.skip("PostgreSQL only")
+async def test_ssl_custom(db_isolated):
+    db_config, _, _ = _get_db_config()
 
     # Expect connectionerror or pass
     ssl_ctx = ssl.create_default_context()
@@ -118,11 +119,10 @@ async def test_ssl_custom():
             await Tortoise._drop_databases()
 
 
+@requireCapability(dialect="postgres")
 @pytest.mark.asyncio
-async def test_application_name():
+async def test_application_name(db_isolated):
     db_config, is_asyncpg, is_psycopg = _get_db_config()
-    if not is_asyncpg and not is_psycopg:
-        pytest.skip("PostgreSQL only")
 
     db_config["connections"]["models"]["credentials"]["application_name"] = "mytest_application"
     try:
@@ -138,3 +138,162 @@ async def test_application_name():
     finally:
         if Tortoise._inited:
             await Tortoise._drop_databases()
+
+
+def _get_query_plan(result: list):
+    query_plan = result[0]["QUERY PLAN"]
+    if isinstance(query_plan, str):
+        query_plan = json.loads(query_plan)
+    return query_plan[0]
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain()
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_format_text(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(output_fmt="text")
+    assert isinstance(result[0]["QUERY PLAN"], str)
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_format_yaml(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(output_fmt="yaml")
+    yaml.safe_dump(result[0]["QUERY PLAN"])
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_format_xml(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(output_fmt="xml")
+    ET.fromstring(result[0]["QUERY PLAN"])
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_unsupported_format(db_simple):
+    await Tournament.create(name="Test")
+    with pytest.raises(UnSupportedError) as exc_info:
+        await Tournament.all().explain(output_fmt="invalid")
+    assert "Unsupported explain format" in str(exc_info.value)
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_analyze(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(analyze=True)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Actual Loops" in query_plan["Plan"]
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_costs(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(costs=True)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Total Cost" in query_plan["Plan"]
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_buffers(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(buffers=True)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Shared Hit Blocks" in query_plan["Plan"]
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_timing(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(analyze=True, timing=True)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Actual Total Time" in query_plan["Plan"]
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_memory(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(memory=True)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Memory" in query_plan or "Memory" in str(query_plan)
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_settings(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(settings=True)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_summary(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(summary=True)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Planning Time" in query_plan
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_multiple_options(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(analyze=True, costs=True, buffers=True)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Actual Loops" in query_plan["Plan"]
+    assert "Total Cost" in query_plan["Plan"]
+    assert "Shared Hit Blocks" in query_plan["Plan"]
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_unsupported_option(db_simple):
+    await Tournament.create(name="Test")
+    with pytest.raises(UnSupportedError) as exc_info:
+        await Tournament.all().explain(unsupported_option=True)
+    assert "UNSUPPORTED_OPTION" in str(exc_info.value)
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_option_false(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain(analyze=False)
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Actual Loops" not in query_plan["Plan"]
+
+
+@requireCapability(dialect="postgres")
+@pytest.mark.asyncio
+async def test_explain_default_verbose(db_simple):
+    await Tournament.create(name="Test")
+    result = await Tournament.all().explain()
+    query_plan = _get_query_plan(result)
+    assert "Plan" in query_plan
+    assert "Output" in query_plan["Plan"]

tortoise/backends/base/executor.py

@@ -10,7 +10,7 @@
 from pypika_tortoise import JoinType, Parameter, Table
 from pypika_tortoise.queries import QueryBuilder
 
-from tortoise.exceptions import OperationalError
+from tortoise.exceptions import OperationalError, UnSupportedError
 from tortoise.expressions import Expression, ResolveContext
 from tortoise.fields.base import DatabaseDefault
 from tortoise.fields.relational import (
@@ -96,7 +96,15 @@ def __init__(
                 self.update_cache,
             ) = EXECUTOR_CACHE[key]
 
-    async def execute_explain(self, sql: str) -> Any:
+    async def execute_explain(
+        self, sql: str, output_fmt: str | None = None, **options: bool
+    ) -> Any:
+        if output_fmt:
+            raise UnSupportedError("This database does not support different explain formats")
+
+        if options:
+            raise UnSupportedError("This database does not support explain options")
+
         sql = " ".join((self.EXPLAIN_PREFIX, sql))
         return (await self.db.execute_query(sql))[1]