QueryService stats support

Author: vgvoleg

examples/query-stats/main.py

@@ -0,0 +1,44 @@
+import os
+import ydb
+
+
+SESSION_QUERY = "SELECT 1"
+TX_QUERY = "SELECT 1"
+
+
+def get_query_stats_from_session(pool: ydb.QuerySessionPool):
+    def callee(session: ydb.QuerySession):
+        with session.execute(SESSION_QUERY, stats_mode=ydb.QueryStatsMode.PROFILE):
+            pass
+
+        print(session.last_query_stats)
+
+    pool.retry_operation_sync(callee)
+
+
+def get_query_stats_from_tx(pool: ydb.QuerySessionPool):
+    def callee(tx: ydb.QueryTxContext):
+        with tx.execute(TX_QUERY, stats_mode=ydb.QueryStatsMode.PROFILE):
+            pass
+
+        print(tx.last_query_stats)
+
+    pool.retry_tx_sync(callee)
+
+
+def main():
+    driver = ydb.Driver(
+        endpoint=os.getenv("YDB_ENDPOINT", "grpc://localhost:2136"),
+        database=os.getenv("YDB_DATABASE", "/local"),
+        credentials=ydb.AnonymousCredentials(),
+    )
+
+    with driver:
+        # wait until driver become initialized
+        driver.wait(fail_fast=True, timeout=5)
+        with ydb.QuerySessionPool(driver) as pool:
+            get_query_stats_from_session(pool)
+            get_query_stats_from_tx(pool)
+
+
+main()

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/pool.py

@@ -142,7 +142,7 @@ async def retry_tx_async(
         """Special interface to execute a bunch of commands with transaction in a safe, retriable way.
 
         :param callee: A function, that works with session.
-        :param tx_mode: Transaction mode, which is a one from the following choises:
+        :param tx_mode: Transaction mode, which is a one from the following choices:
           1) QuerySerializableReadWrite() which is default mode;
           2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
           3) QuerySnapshotReadOnly();

ydb/aio/query/session.py

@@ -117,26 +117,34 @@ async def execute(
         exec_mode: base.QueryExecMode = None,
         concurrent_result_sets: bool = False,
         settings: Optional[BaseRequestSettings] = None,
+        *,
+        stats_mode: Optional[base.QueryStatsMode] = None,
     ) -> AsyncResponseContextIterator:
         """Sends a query to Query Service
 
         :param query: (YQL or SQL text) to be executed.
-        :param syntax: Syntax of the query, which is a one from the following choises:
+        :param syntax: Syntax of the query, which is a one from the following choices:
          1) QuerySyntax.YQL_V1, which is default;
          2) QuerySyntax.PG.
         :param parameters: dict with parameters and YDB types;
         :param concurrent_result_sets: A flag to allow YDB mix parts of different result sets. Default is False;
+        :param stats_mode: Mode of query statistics to gather, which is a one from the following choices:
+         1) QueryStatsMode:NONE, which is default;
+         2) QueryStatsMode.BASIC;
+         3) QueryStatsMode.FULL;
+         4) QueryStatsMode.PROFILE;
 
         :return: Iterator with result sets
         """
         self._state._check_session_ready_to_use()
 
         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 +155,7 @@ async def execute(
                 rpc_state=None,
                 response_pb=resp,
                 session_state=self._state,
+                session=self,
                 settings=self._settings,
             ),
         )

ydb/aio/query/transaction.py

@@ -29,7 +29,7 @@ def __init__(self, driver, session_state, session, tx_mode):
 
         :param driver: A driver instance
         :param session_state: A state of session
-        :param tx_mode: Transaction mode, which is a one from the following choises:
+        :param tx_mode: Transaction mode, which is a one from the following choices:
          1) QuerySerializableReadWrite() which is default mode;
          2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
          3) QuerySnapshotReadOnly();
@@ -142,32 +142,40 @@ async def execute(
         exec_mode: Optional[base.QueryExecMode] = None,
         concurrent_result_sets: Optional[bool] = False,
         settings: Optional[BaseRequestSettings] = None,
+        *,
+        stats_mode: Optional[base.QueryStatsMode] = None,
     ) -> AsyncResponseContextIterator:
         """Sends a query to Query Service
 
         :param query: (YQL or SQL text) to be executed.
         :param parameters: dict with parameters and YDB types;
         :param commit_tx: A special flag that allows transaction commit.
-        :param syntax: Syntax of the query, which is a one from the following choises:
+        :param syntax: Syntax of the query, which is a one from the following choices:
          1) QuerySyntax.YQL_V1, which is default;
          2) QuerySyntax.PG.
-        :param exec_mode: Exec mode of the query, which is a one from the following choises:
+        :param exec_mode: Exec mode of the query, which is a one from the following choices:
          1) QueryExecMode.EXECUTE, which is default;
          2) QueryExecMode.EXPLAIN;
          3) QueryExecMode.VALIDATE;
          4) QueryExecMode.PARSE.
         :param concurrent_result_sets: A flag to allow YDB mix parts of different result sets. Default is False;
+        :param stats_mode: Mode of query statistics to gather, which is a one from the following choices:
+         1) QueryStatsMode:NONE, which is default;
+         2) QueryStatsMode.BASIC;
+         3) QueryStatsMode.FULL;
+         4) QueryStatsMode.PROFILE;
 
         :return: Iterator with result sets
         """
         await self._ensure_prev_stream_finished()
 
         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/__init__.py

@@ -7,13 +7,15 @@
     "QuerySessionPool",
     "QueryClientSettings",
     "QuerySession",
+    "QueryStatsMode",
     "QueryTxContext",
 ]
 
 import logging
 
 from .base import (
     QueryClientSettings,
+    QueryStatsMode,
 )
 
 from .session import QuerySession

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/pool.py

@@ -151,7 +151,7 @@ def retry_tx_sync(
         """Special interface to execute a bunch of commands with transaction in a safe, retriable way.
 
         :param callee: A function, that works with session.
-        :param tx_mode: Transaction mode, which is a one from the following choises:
+        :param tx_mode: Transaction mode, which is a one from the following choices:
           1) QuerySerializableReadWrite() which is default mode;
           2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
           3) QuerySnapshotReadOnly();