added result format flag to explain

fixed container waits

Author: rekby

tests/aio/query/test_query_session.py

@@ -3,6 +3,7 @@
 import pytest
 
 import ydb
+from ydb import QueryExplainResultFormat
 from ydb.aio.query.session import QuerySession
 
 
@@ -139,10 +140,17 @@ async def test_explain(self, pool: ydb.aio.query.QuerySessionPool):
             async def callee(session: QuerySession):
                 nonlocal plan_fullscan, plan_lookup
 
-                plan_fullscan = await session.explain("SELECT * FROM test_explain")
+                plan = await session.explain("SELECT * FROM test_explain", result_format=QueryExplainResultFormat.STR)
+                isinstance(plan, str)
+                assert "FullScan" in plan
+
+                plan_fullscan = await session.explain(
+                    "SELECT * FROM test_explain", result_format=QueryExplainResultFormat.DICT
+                )
                 plan_lookup = await session.explain(
                     "SELECT * FROM test_explain WHERE id = $id",
                     {"$id": 1},
+                    result_format=QueryExplainResultFormat.DICT,
                 )
 
             await pool.retry_operation_async(callee)

tests/aio/query/test_query_session_pool.py

@@ -7,6 +7,7 @@
 from typing import Optional
 
 from tests.conftest import wait_container_ready_async
+from ydb import QueryExplainResultFormat
 from ydb.aio.query.pool import QuerySessionPool
 from ydb.aio.query.session import QuerySession, QuerySessionStateEnum
 from ydb.aio.query.transaction import QueryTxContext
@@ -189,13 +190,22 @@ 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 = await pool.explain_with_retries(
+                "SELECT * FROM test_explain", result_format=QueryExplainResultFormat.STR
+            )
+            isinstance(plan, str)
+            assert "FullScan" in plan
+
+            plan = await pool.explain_with_retries(
+                "SELECT * FROM test_explain", result_format=QueryExplainResultFormat.DICT
+            )
             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},
+                result_format=QueryExplainResultFormat.DICT,
             )
             plan_string = json.dumps(plan)
             assert "Lookup" in plan_string

tests/aio/test_connection_pool.py

@@ -109,8 +109,8 @@ async def restart_docker():
         await asyncio.gather(*coros, return_exceptions=False)
 
     docker_project.start()
-    await driver.stop()
     await wait_container_ready_async(driver)
+    await driver.stop()
 
 
 @pytest.mark.asyncio
@@ -134,8 +134,8 @@ async def test_disconnect_by_call(endpoint, database, docker_project):
     await asyncio.sleep(5)
     assert len(driver._store.connections) == 0
     docker_project.start()
-    await driver.stop()
     await wait_container_ready_async(driver)
+    await driver.stop()
 
 
 @pytest.mark.asyncio

tests/conftest.py

@@ -37,7 +37,7 @@ async def wait_container_ready_async(driver):
     async with ydb.aio.SessionPool(driver, 1) as pool:
 
         started_at = time.time()
-        while time.time() - started_at < 120:
+        while time.time() - started_at < 30:
             try:
                 async with pool.checkout() as session:
                     await session.execute_scheme("create table `.sys_health/test_table` (A int32, primary key(A));")

tests/query/test_query_session.py

@@ -7,7 +7,7 @@
 from unittest import mock
 
 from ydb import QuerySessionPool
-from ydb.query.base import QueryStatsMode
+from ydb.query.base import QueryStatsMode, QueryExplainResultFormat
 from ydb.query.session import QuerySession
 
 
@@ -187,10 +187,19 @@ def test_explain(self, pool: QuerySessionPool):
 
             def callee(session: QuerySession):
                 nonlocal plan_fullscan, plan_lookup
-                plan_fullscan = session.explain("SELECT * FROM test_explain")
+
+                plan = session.explain("SELECT * FROM test_explain", result_format=QueryExplainResultFormat.STR)
+                isinstance(plan, str)
+                assert "FullScan" in plan
+
+                plan_fullscan = session.explain(
+                    "SELECT * FROM test_explain", result_format=QueryExplainResultFormat.DICT
+                )
+
                 plan_lookup = session.explain(
                     "SELECT * FROM test_explain WHERE id = $id",
                     {"$id": 1},
+                    result_format=QueryExplainResultFormat.DICT,
                 )
 
             pool.retry_operation_sync(callee)

tests/query/test_query_session_pool.py

@@ -8,6 +8,7 @@
 from typing import Optional
 
 from tests.conftest import wait_container_ready
+from ydb import QueryExplainResultFormat
 from ydb.query.pool import QuerySessionPool
 from ydb.query.session import QuerySession, QuerySessionStateEnum
 from ydb.query.transaction import QueryTxContext
@@ -150,8 +151,8 @@ def test_no_session_leak(self, driver_sync, docker_project):
         assert pool._current_size == 0
 
         docker_project.start()
+        wait_container_ready(driver_sync)
         pool.stop()
-        wait_container_ready(driver)
 
     def test_execute_with_retries_async(self, pool: QuerySessionPool):
         fut = pool.execute_with_retries_async("select 1;")
@@ -209,13 +210,21 @@ 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 = pool.explain_with_retries("SELECT * FROM test_explain", result_format=QueryExplainResultFormat.STR)
+            isinstance(plan, str)
+            assert "FullScan" in plan
+
+            plan = pool.explain_with_retries("SELECT * FROM test_explain", result_format=QueryExplainResultFormat.DICT)
+            assert isinstance(plan, dict)
+
             plan_string = json.dumps(plan)
             assert "FullScan" in plan_string
 
             plan = pool.explain_with_retries(
                 "SELECT * FROM test_explain WHERE id = $id",
                 {"$id": 1},
+                result_format=ydb.QueryExplainResultFormat.DICT,
             )
             plan_string = json.dumps(plan)
             assert "Lookup" in plan_string

ydb/aio/query/pool.py

@@ -6,6 +6,7 @@
     List,
     Dict,
     Any,
+    Union,
 )
 
 from .session import (
@@ -15,7 +16,7 @@
     RetrySettings,
     retry_operation_async,
 )
-from ...query.base import BaseQueryTxMode
+from ...query.base import BaseQueryTxMode, QueryExplainResultFormat
 from ...query.base import QueryClientSettings
 from ... import convert
 from ..._grpc.grpcwrapper import common_utils
@@ -201,19 +202,21 @@ async def explain_with_retries(
         query: str,
         parameters: Optional[dict] = None,
         *,
+        result_format: QueryExplainResultFormat = QueryExplainResultFormat.STR,
         retry_settings: Optional[RetrySettings] = None,
-    ) -> Dict[str, Any]:
+    ) -> Union[str, Dict[str, Any]]:
         """
         Explain a query in retriable way. No real query execution will happen.
 
         :param query: A query, yql or sql text.
         :param parameters: dict with parameters and YDB types;
+        :param result_format: Return format: string or dict.
         :param retry_settings: RetrySettings object.
         :return: Parsed query plan.
         """
 
         async def callee(session: QuerySession):
-            return await session.explain(query, parameters)
+            return await session.explain(query, parameters, result_format=result_format)
 
         return await self.retry_operation_async(callee, retry_settings)
 

ydb/aio/query/session.py

@@ -5,6 +5,7 @@
     Optional,
     Dict,
     Any,
+    Union,
 )
 
 from .base import AsyncResponseContextIterator
@@ -166,10 +167,16 @@ async def execute(
             error_converter=stream_error_converter,
         )
 
-    async def explain(self, query: str, parameters: Optional[dict] = None) -> Dict[str, Any]:
+    async def explain(
+        self,
+        query: str,
+        parameters: Optional[dict] = None,
+        result_format: base.QueryExplainResultFormat = base.QueryExplainResultFormat.STR,
+    ) -> Union[str, Dict[str, Any]]:
         """Explains query result
         :param query: YQL or SQL query.
         :param parameters: dict with parameters and YDB types;
+        :param result_format: Return format: string or dict.
         :return: Parsed query plan.
         """
 
@@ -179,6 +186,9 @@ async def explain(self, query: str, parameters: Optional[dict] = None) -> Dict[s
         async for _ in res:
             pass
 
-        plan_text = self.last_query_stats.query_plan
-        plan = json.loads(plan_text)
+        plan = self.last_query_stats.query_plan
+
+        if result_format == base.QueryExplainResultFormat.DICT:
+            plan = json.loads(plan)
+
         return plan

ydb/query/__init__.py

@@ -1,5 +1,6 @@
 __all__ = [
     "BaseQueryTxMode",
+    "QueryExplainResultFormat",
     "QueryOnlineReadOnly",
     "QuerySerializableReadWrite",
     "QuerySnapshotReadOnly",
@@ -15,6 +16,7 @@
 
 from .base import (
     QueryClientSettings,
+    QueryExplainResultFormat,
     QueryStatsMode,
 )
 

ydb/query/base.py

@@ -42,6 +42,11 @@ class QueryExecMode(enum.IntEnum):
     EXECUTE = 50
 
 
+class QueryExplainResultFormat(enum.Enum):
+    STR = 0
+    DICT = 10
+
+
 class QueryStatsMode(enum.IntEnum):
     UNSPECIFIED = 0
     NONE = 10