QueryService stats support

Author: vgvoleg

tests/query/test_query_session.py

@@ -4,7 +4,7 @@
 from concurrent.futures import _base as b
 from unittest import mock
 
-
+from ydb.query.base import QueryStatsMode
 from ydb.query.session import QuerySession
 
 
@@ -143,3 +143,34 @@ def cancel(self):
         assert "attach stream thread" not in thread_names
 
         _check_session_state_empty(session)
+
+    @pytest.mark.parametrize(
+        "stats_mode",
+        [
+            None,
+            QueryStatsMode.UNSPECIFIED,
+            QueryStatsMode.NONE,
+            QueryStatsMode.BASIC,
+            QueryStatsMode.FULL,
+            QueryStatsMode.PROFILE,
+        ],
+    )
+    def test_stats_mode(self, session: QuerySession, stats_mode: QueryStatsMode):
+        session.create()
+
+        for _ in session.execute("SELECT 1; SELECT 2; SELECT 3;", stats_mode=stats_mode):
+            pass
+
+        stats = session.last_query_stats
+
+        if stats_mode in [None, QueryStatsMode.NONE, QueryStatsMode.UNSPECIFIED]:
+            assert stats is None
+            return
+
+        assert stats is not None
+        assert len(stats.query_phases) > 0
+
+        if stats_mode != QueryStatsMode.BASIC:
+            assert len(stats.query_plan) > 0
+        else:
+            assert stats.query_plan == ""

tests/query/test_query_transaction.py

@@ -1,5 +1,6 @@
 import pytest
 
+from ydb.query.base import QueryStatsMode
 from ydb.query.transaction import QueryTxContext
 from ydb.query.transaction import QueryTxStateEnum
 
@@ -104,3 +105,32 @@ def test_tx_identity_after_begin_works(self, tx: QueryTxContext):
 
         assert identity.tx_id == tx.tx_id
         assert identity.session_id == tx.session_id
+
+    @pytest.mark.parametrize(
+        "stats_mode",
+        [
+            None,
+            QueryStatsMode.UNSPECIFIED,
+            QueryStatsMode.NONE,
+            QueryStatsMode.BASIC,
+            QueryStatsMode.FULL,
+            QueryStatsMode.PROFILE,
+        ],
+    )
+    def test_stats_mode(self, tx: QueryTxContext, stats_mode: QueryStatsMode):
+        for _ in tx.execute("SELECT 1; SELECT 2; SELECT 3;", commit_tx=True, stats_mode=stats_mode):
+            pass
+
+        stats = tx.last_query_stats
+
+        if stats_mode in [None, QueryStatsMode.NONE, QueryStatsMode.UNSPECIFIED]:
+            assert stats is None
+            return
+
+        assert stats is not None
+        assert len(stats.query_phases) > 0
+
+        if stats_mode != QueryStatsMode.BASIC:
+            assert len(stats.query_plan) > 0
+        else:
+            assert stats.query_plan == ""

ydb/aio/query/session.py

@@ -115,6 +115,7 @@ async def execute(
         parameters: dict = None,
         syntax: base.QuerySyntax = None,
         exec_mode: base.QueryExecMode = None,
+        stats_mode: Optional[base.QueryStatsMode] = None,
         concurrent_result_sets: bool = False,
         settings: Optional[BaseRequestSettings] = None,
     ) -> AsyncResponseContextIterator:
@@ -133,10 +134,11 @@ async def execute(
 
         stream_it = await self._execute_call(
             query=query,
+            parameters=parameters,
             commit_tx=True,
             syntax=syntax,
             exec_mode=exec_mode,
-            parameters=parameters,
+            stats_mode=stats_mode,
             concurrent_result_sets=concurrent_result_sets,
             settings=settings,
         )
@@ -147,6 +149,7 @@ async def execute(
                 rpc_state=None,
                 response_pb=resp,
                 session_state=self._state,
+                session=self,
                 settings=self._settings,
             ),
         )

ydb/aio/query/transaction.py

@@ -140,6 +140,7 @@ async def execute(
         commit_tx: Optional[bool] = False,
         syntax: Optional[base.QuerySyntax] = None,
         exec_mode: Optional[base.QueryExecMode] = None,
+        stats_mode: Optional[base.QueryStatsMode] = None,
         concurrent_result_sets: Optional[bool] = False,
         settings: Optional[BaseRequestSettings] = None,
     ) -> AsyncResponseContextIterator:
@@ -156,6 +157,11 @@ async def execute(
          2) QueryExecMode.EXPLAIN;
          3) QueryExecMode.VALIDATE;
          4) QueryExecMode.PARSE.
+        :param stats_mode: Mode of query statistics to gather, which is a one from the following choises:
+         1) QueryStatsMode:NONE, which is default;
+         2) QueryStatsMode.BASIC;
+         3) QueryStatsMode.FULL;
+         4) QueryStatsMode.PROFILE;
         :param concurrent_result_sets: A flag to allow YDB mix parts of different result sets. Default is False;
 
         :return: Iterator with result sets
@@ -164,10 +170,11 @@ async def execute(
 
         stream_it = await self._execute_call(
             query=query,
+            parameters=parameters,
             commit_tx=commit_tx,
             syntax=syntax,
             exec_mode=exec_mode,
-            parameters=parameters,
+            stats_mode=stats_mode,
             concurrent_result_sets=concurrent_result_sets,
             settings=settings,
         )

ydb/query/base.py

@@ -25,6 +25,7 @@
 
 if typing.TYPE_CHECKING:
     from .transaction import BaseQueryTxContext
+    from .session import BaseQuerySession
 
 
 class QuerySyntax(enum.IntEnum):
@@ -41,7 +42,7 @@ class QueryExecMode(enum.IntEnum):
     EXECUTE = 50
 
 
-class StatsMode(enum.IntEnum):
+class QueryStatsMode(enum.IntEnum):
     UNSPECIFIED = 0
     NONE = 10
     BASIC = 20
@@ -132,12 +133,13 @@ def create_execute_query_request(
     tx_mode: Optional[BaseQueryTxMode],
     syntax: Optional[QuerySyntax],
     exec_mode: Optional[QueryExecMode],
+    stats_mode: Optional[QueryStatsMode],
     parameters: Optional[dict],
     concurrent_result_sets: Optional[bool],
 ) -> ydb_query.ExecuteQueryRequest:
     syntax = QuerySyntax.YQL_V1 if not syntax else syntax
     exec_mode = QueryExecMode.EXECUTE if not exec_mode else exec_mode
-    stats_mode = StatsMode.NONE  # TODO: choise is not supported yet
+    stats_mode = QueryStatsMode.NONE if stats_mode is None else stats_mode
 
     tx_control = None
     if not tx_id and not tx_mode:
@@ -189,6 +191,7 @@ def wrap_execute_query_response(
     response_pb: _apis.ydb_query.ExecuteQueryResponsePart,
     session_state: IQuerySessionState,
     tx: Optional["BaseQueryTxContext"] = None,
+    session: Optional["BaseQuerySession"] = None,
     commit_tx: Optional[bool] = False,
     settings: Optional[QueryClientSettings] = None,
 ) -> convert.ResultSet:
@@ -198,6 +201,12 @@ def wrap_execute_query_response(
     elif tx and response_pb.tx_meta and not tx.tx_id:
         tx._move_to_beginned(response_pb.tx_meta.id)
 
+    if response_pb.HasField("exec_stats"):
+        if tx is not None:
+            tx._last_query_stats = response_pb.exec_stats
+        if session is not None:
+            session._last_query_stats = response_pb.exec_stats
+
     if response_pb.HasField("result_set"):
         return convert.ResultSet.from_message(response_pb.result_set, settings)
 

ydb/query/session.py

@@ -147,6 +147,12 @@ def __init__(self, driver: common_utils.SupportedDriverType, settings: Optional[
             .with_timeout(DEFAULT_ATTACH_LONG_TIMEOUT)
         )
 
+        self._last_query_stats = None
+
+    @property
+    def last_query_stats(self):
+        return self._last_query_stats
+
     def _get_client_settings(
         self,
         driver: common_utils.SupportedDriverType,
@@ -189,22 +195,26 @@ def _attach_call(self) -> Iterable[_apis.ydb_query.SessionState]:
     def _execute_call(
         self,
         query: str,
+        parameters: dict = None,
         commit_tx: bool = False,
         syntax: base.QuerySyntax = None,
         exec_mode: base.QueryExecMode = None,
-        parameters: dict = None,
+        stats_mode: Optional[base.QueryStatsMode] = None,
         concurrent_result_sets: bool = False,
         settings: Optional[BaseRequestSettings] = None,
     ) -> Iterable[_apis.ydb_query.ExecuteQueryResponsePart]:
+        self._last_query_stats = None
+
         request = base.create_execute_query_request(
             query=query,
-            session_id=self._state.session_id,
+            parameters=parameters,
             commit_tx=commit_tx,
+            session_id=self._state.session_id,
             tx_mode=None,
             tx_id=None,
             syntax=syntax,
             exec_mode=exec_mode,
-            parameters=parameters,
+            stats_mode=stats_mode,
             concurrent_result_sets=concurrent_result_sets,
         )
 
@@ -320,6 +330,7 @@ def execute(
         syntax: base.QuerySyntax = None,
         exec_mode: base.QueryExecMode = None,
         concurrent_result_sets: bool = False,
+        stats_mode: Optional[base.QueryStatsMode] = None,
         settings: Optional[BaseRequestSettings] = None,
     ) -> base.SyncResponseContextIterator:
         """Sends a query to Query Service
@@ -337,10 +348,11 @@ def execute(
 
         stream_it = self._execute_call(
             query=query,
+            parameters=parameters,
             commit_tx=True,
             syntax=syntax,
             exec_mode=exec_mode,
-            parameters=parameters,
+            stats_mode=stats_mode,
             concurrent_result_sets=concurrent_result_sets,
             settings=settings,
         )
@@ -351,6 +363,7 @@ def execute(
                 rpc_state=None,
                 response_pb=resp,
                 session_state=self._state,
+                session=self,
                 settings=self._settings,
             ),
         )

ydb/query/transaction.py

@@ -210,6 +210,7 @@ def __init__(self, driver, session_state, session, tx_mode):
         self.session = session
         self._prev_stream = None
         self._external_error = None
+        self._last_query_stats = None
 
     @property
     def session_id(self) -> str:
@@ -229,6 +230,10 @@ def tx_id(self) -> Optional[str]:
         """
         return self._tx_state.tx_id
 
+    @property
+    def last_query_stats(self):
+        return self._last_query_stats
+
     def _tx_identity(self) -> _ydb_topic.TransactionIdentity:
         if not self.tx_id:
             raise RuntimeError("Unable to get tx identity without started tx.")
@@ -283,25 +288,29 @@ def _rollback_call(self, settings: Optional[BaseRequestSettings]) -> "BaseQueryT
     def _execute_call(
         self,
         query: str,
+        parameters: Optional[dict],
         commit_tx: Optional[bool],
         syntax: Optional[base.QuerySyntax],
         exec_mode: Optional[base.QueryExecMode],
-        parameters: Optional[dict],
+        stats_mode: Optional[base.QueryStatsMode],
         concurrent_result_sets: Optional[bool],
         settings: Optional[BaseRequestSettings],
     ) -> Iterable[_apis.ydb_query.ExecuteQueryResponsePart]:
         self._tx_state._check_tx_ready_to_use()
         self._check_external_error_set()
 
+        self._last_query_stats = None
+
         request = base.create_execute_query_request(
             query=query,
-            session_id=self._session_state.session_id,
+            parameters=parameters,
             commit_tx=commit_tx,
+            session_id=self._session_state.session_id,
             tx_id=self._tx_state.tx_id,
             tx_mode=self._tx_state.tx_mode,
             syntax=syntax,
             exec_mode=exec_mode,
-            parameters=parameters,
+            stats_mode=stats_mode,
             concurrent_result_sets=concurrent_result_sets,
         )
 
@@ -449,6 +458,7 @@ def execute(
         commit_tx: Optional[bool] = False,
         syntax: Optional[base.QuerySyntax] = None,
         exec_mode: Optional[base.QueryExecMode] = None,
+        stats_mode: Optional[base.QueryStatsMode] = None,
         concurrent_result_sets: Optional[bool] = False,
         settings: Optional[BaseRequestSettings] = None,
     ) -> base.SyncResponseContextIterator:
@@ -465,6 +475,11 @@ def execute(
          2) QueryExecMode.EXPLAIN;
          3) QueryExecMode.VALIDATE;
          4) QueryExecMode.PARSE.
+        :param stats_mode: Mode of query statistics to gather, which is a one from the following choises:
+         1) QueryStatsMode:NONE, which is default;
+         2) QueryStatsMode.BASIC;
+         3) QueryStatsMode.FULL;
+         4) QueryStatsMode.PROFILE;
         :param concurrent_result_sets: A flag to allow YDB mix parts of different result sets. Default is False;
         :param settings: An additional request settings BaseRequestSettings;
 
@@ -477,6 +492,7 @@ def execute(
             commit_tx=commit_tx,
             syntax=syntax,
             exec_mode=exec_mode,
+            stats_mode=stats_mode,
             parameters=parameters,
             concurrent_result_sets=concurrent_result_sets,
             settings=settings,