Add explain method to query session and pool

Author: rekby

tests/aio/query/test_query_session.py

@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 import ydb
@@ -125,3 +127,30 @@ async def test_terminated_stream_raises_ydb_error(self, session: QuerySession):
             async with await session.execute("select 1") as results:
                 async for _ in results:
                     pass
+
+    @pytest.mark.asyncio
+    async def test_explain(self, pool: ydb.aio.query.QuerySessionPool):
+        await pool.execute_with_retries("DROP TABLE IF EXISTS test_explain")
+        await pool.execute_with_retries("CREATE TABLE test_explain (id Int64, PRIMARY KEY (id))")
+        try:
+            plan_fullscan = ""
+            plan_lookup = ""
+
+            async def callee(session: QuerySession):
+                nonlocal plan_fullscan, plan_lookup
+
+                plan_fullscan = await session.explain("SELECT * FROM test_explain")
+                plan_lookup = await session.explain(
+                    "SELECT * FROM test_explain WHERE id = $id",
+                    {"$id": 1},
+                )
+
+            await pool.retry_operation_async(callee)
+
+            plan_fulltext_string = json.dumps(plan_fullscan)
+            assert "FullScan" in plan_fulltext_string
+
+            plan_lookup_string = json.dumps(plan_lookup)
+            assert "Lookup" in plan_lookup_string
+        finally:
+            await pool.execute_with_retries("DROP TABLE test_explain")

tests/aio/query/test_query_session_pool.py

@@ -1,4 +1,6 @@
 import asyncio
+import json
+
 import pytest
 import ydb
 
@@ -179,3 +181,21 @@ async def acquire_session():
 
             assert len(ids) == 1
             assert pool._current_size == 1
+
+    @pytest.mark.asyncio
+    async def test_explain_with_retries(self, pool: QuerySessionPool):
+        await pool.execute_with_retries("DROP TABLE IF EXISTS test_explain")
+        await pool.execute_with_retries("CREATE TABLE test_explain (id Int64, PRIMARY KEY (id))")
+        try:
+            plan = await pool.explain_with_retries("SELECT * FROM test_explain")
+            plan_string = json.dumps(plan)
+            assert "FullScan" in plan_string
+
+            plan = await pool.explain_with_retries(
+                "SELECT * FROM test_explain WHERE id = $id",
+                {"$id": 1},
+            )
+            plan_string = json.dumps(plan)
+            assert "Lookup" in plan_string
+        finally:
+            await pool.execute_with_retries("DROP TABLE test_explain")

tests/query/test_query_session.py

@@ -1,9 +1,12 @@
+import json
+
 import pytest
 import threading
 import time
 from concurrent.futures import _base as b
 from unittest import mock
 
+from ydb import QuerySessionPool
 from ydb.query.base import QueryStatsMode
 from ydb.query.session import QuerySession
 
@@ -174,3 +177,28 @@ def test_stats_mode(self, session: QuerySession, stats_mode: QueryStatsMode):
             assert len(stats.query_plan) > 0
         else:
             assert stats.query_plan == ""
+
+    def test_explain(self, pool: QuerySessionPool):
+        pool.execute_with_retries("DROP TABLE IF EXISTS test_explain")
+        pool.execute_with_retries("CREATE TABLE test_explain (id Int64, PRIMARY KEY (id))")
+        try:
+            plan_fullscan = ""
+            plan_lookup = ""
+
+            def callee(session: QuerySession):
+                nonlocal plan_fullscan, plan_lookup
+                plan_fullscan = session.explain("SELECT * FROM test_explain")
+                plan_lookup = session.explain(
+                    "SELECT * FROM test_explain WHERE id = $id",
+                    {"$id": 1},
+                )
+
+            pool.retry_operation_sync(callee)
+
+            plan_fulltext_string = json.dumps(plan_fullscan)
+            assert "FullScan" in plan_fulltext_string
+
+            plan_lookup_string = json.dumps(plan_lookup)
+            assert "Lookup" in plan_lookup_string
+        finally:
+            pool.execute_with_retries("DROP TABLE test_explain")

tests/query/test_query_session_pool.py

@@ -1,3 +1,5 @@
+import json
+
 import pytest
 import ydb
 import time
@@ -200,3 +202,20 @@ def test_async_methods_after_stop_raise(self, pool: QuerySessionPool):
         pool.stop()
         with pytest.raises(ydb.SessionPoolClosed):
             pool.execute_with_retries_async("select 1;")
+
+    def test_explain_with_retries(self, pool: QuerySessionPool):
+        pool.execute_with_retries("DROP TABLE IF EXISTS test_explain")
+        pool.execute_with_retries("CREATE TABLE test_explain (id Int64, PRIMARY KEY (id))")
+        try:
+            plan = pool.explain_with_retries("SELECT * FROM test_explain")
+            plan_string = json.dumps(plan)
+            assert "FullScan" in plan_string
+
+            plan = pool.explain_with_retries(
+                "SELECT * FROM test_explain WHERE id = $id",
+                {"$id": 1},
+            )
+            plan_string = json.dumps(plan)
+            assert "Lookup" in plan_string
+        finally:
+            pool.execute_with_retries("DROP TABLE test_explain")

ydb/aio/query/pool.py

@@ -4,6 +4,8 @@
     Callable,
     Optional,
     List,
+    Dict,
+    Any,
 )
 
 from .session import (
@@ -194,6 +196,18 @@ async def wrapped_callee():
 
         return await retry_operation_async(wrapped_callee, retry_settings)
 
+    async def explain_with_retries(
+        self,
+        query: str,
+        parameters: Optional[dict] = None,
+        *,
+        retry_settings: Optional[RetrySettings] = None,
+    ) -> Dict[str, Any]:
+        async def callee(session: QuerySession):
+            return await session.explain(query, parameters)
+
+        return await self.retry_operation_async(callee, retry_settings)
+
     async def stop(self):
         self._should_stop.set()
 

ydb/aio/query/session.py

@@ -1,7 +1,10 @@
 import asyncio
+import json
 
 from typing import (
     Optional,
+    Dict,
+    Any,
 )
 
 from .base import AsyncResponseContextIterator
@@ -162,3 +165,16 @@ async def execute(
             ),
             error_converter=stream_error_converter,
         )
+
+    async def explain(self, query: str, parameters: Optional[dict] = None) -> Dict[str, Any]:
+        """Explains query result
+        :param query: YQL or SQL query.
+        """
+
+        res = await self.execute(query, parameters, exec_mode=base.QueryExecMode.EXPLAIN)
+        async for _ in res:
+            pass
+
+        plan_text = self.last_query_stats.query_plan
+        plan = json.loads(plan_text)
+        return plan

ydb/query/pool.py

@@ -4,6 +4,8 @@
     Callable,
     Optional,
     List,
+    Dict,
+    Any,
 )
 import time
 import threading
@@ -270,6 +272,18 @@ def execute_with_retries_async(
             **kwargs,
         )
 
+    def explain_with_retries(
+        self,
+        query: str,
+        parameters: Optional[dict] = None,
+        *,
+        retry_settings: Optional[RetrySettings] = None,
+    ) -> Dict[str, Any]:
+        def callee(session: QuerySession):
+            return session.explain(query, parameters)
+
+        return self.retry_operation_sync(callee, retry_settings)
+
     def stop(self, timeout=None):
         acquire_timeout = timeout if timeout is not None else -1
         acquired = self._lock.acquire(timeout=acquire_timeout)

ydb/query/session.py

@@ -1,10 +1,13 @@
 import abc
 import enum
+import json
 import logging
 import threading
 from typing import (
     Iterable,
     Optional,
+    Dict,
+    Any,
 )
 
 from . import base
@@ -371,3 +374,17 @@ def execute(
                 settings=self._settings,
             ),
         )
+
+    def explain(self, query: str, parameters: dict = None) -> Dict[str, Any]:
+        """Explains query result
+        :param query: YQL or SQL query.
+        :param parameters: dict with parameters and YDB types;
+        """
+
+        res = self.execute(query, parameters, exec_mode=base.QueryExecMode.EXPLAIN)
+        for _ in res:
+            pass
+
+        plan_text = self.last_query_stats.query_plan
+        plan = json.loads(plan_text)
+        return plan