Merge pull request #27 from ydb-platform/add-query-stats

add support for query stats in scan query

Author: gridnevvvit

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 2.4.0 ##
+
+* support query stats in scan query
+
 ## 2.2.0 ##
 
 * allow to refer endpoints by node id

setup.py

@@ -4,7 +4,7 @@
 
 setuptools.setup(
     name="ydb",
-    version="2.3.1",
+    version="2.4.0",
     description="YDB Python library",
     author="Yandex LLC",
     author_email="ydb@yandex-team.ru",

tests/aio/test_async_iter_stream.py

@@ -32,6 +32,52 @@ async def test_read_table(driver, database):
         i += 1
 
 
+@pytest.mark.asyncio
+async def test_scan_query_with_query_stats(driver, database):
+    description = (
+        ydb.TableDescription()
+        .with_primary_keys("key1")
+        .with_columns(
+            ydb.Column("key1", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
+            ydb.Column("value", ydb.OptionalType(ydb.PrimitiveType.Utf8)),
+        )
+    )
+
+    session = await driver.table_client.session().create()
+    await session.create_table(
+        database + "/test_scan_query_with_query_stats", description
+    )
+    await session.transaction(ydb.SerializableReadWrite()).execute(
+        """INSERT INTO `test_scan_query_with_query_stats` (`key1`, `value`) VALUES (1, "hello_world"), (2, "2")""",
+        commit_tx=True,
+    )
+
+    expected_res = [{"key1": 1, "value": "hello_world"}, {"key1": 2, "value": "2"}]
+    i = 0
+    async for resp in await driver.table_client.scan_query(
+        "select * from `test_scan_query_with_query_stats` order by key1"
+    ):
+        for row in resp.result_set.rows:
+            assert row == expected_res[i]
+            i += 1
+
+    i = 0
+    settings = ydb.ScanQuerySettings().with_collect_stats(
+        ydb.QueryStatsCollectionMode.FULL
+    )
+    last_response = None
+    async for resp in await driver.table_client.scan_query(
+        "select * from `test_scan_query_with_query_stats` order by key1",
+        settings=settings,
+    ):
+        last_response = resp
+        for row in resp.result_set.rows:
+            assert row == expected_res[i]
+            i += 1
+
+    assert len(last_response.query_stats.query_phases) > 0
+
+
 @pytest.mark.asyncio
 async def test_read_shard_table(driver, database):
     session: ydb.ISession = await driver.table_client.session().create()

ydb/table.py

@@ -1027,11 +1027,29 @@ def with_lazy_result_sets(self, enabled):
 class ScanQueryResult(object):
     def __init__(self, result, table_client_settings):
         self._result = result
+        self.query_stats = result.query_stats
         self.result_set = convert.ResultSet.from_message(
             self._result.result_set, table_client_settings
         )
 
 
+@enum.unique
+class QueryStatsCollectionMode(enum.IntEnum):
+    NONE = _apis.ydb_table.QueryStatsCollection.Mode.STATS_COLLECTION_NONE
+    BASIC = _apis.ydb_table.QueryStatsCollection.Mode.STATS_COLLECTION_BASIC
+    FULL = _apis.ydb_table.QueryStatsCollection.Mode.STATS_COLLECTION_FULL
+
+
+class ScanQuerySettings(settings_impl.BaseRequestSettings):
+    def __init__(self):
+        super(ScanQuerySettings, self).__init__()
+        self.collect_stats = None
+
+    def with_collect_stats(self, collect_stats_mode):
+        self.collect_stats = collect_stats_mode
+        return self
+
+
 class ScanQuery(object):
     def __init__(self, yql_text, parameters_types):
         self.yql_text = yql_text
@@ -1047,10 +1065,16 @@ def _scan_query_request_factory(query, parameters=None, settings=None):
     if not isinstance(query, ScanQuery):
         query = ScanQuery(query, {})
     parameters = {} if parameters is None else parameters
+    collect_stats = getattr(
+        settings,
+        "collect_stats",
+        _apis.ydb_table.QueryStatsCollection.Mode.STATS_COLLECTION_NONE,
+    )
     return _apis.ydb_table.ExecuteScanQueryRequest(
         mode=_apis.ydb_table.ExecuteScanQueryRequest.Mode.MODE_EXEC,
         query=_apis.ydb_table.Query(yql_text=query.yql_text),
         parameters=convert.parameters_to_pb(query.parameters_types, parameters),
+        collect_stats=collect_stats,
     )
 
 

ydb/ydb_version.py

@@ -1 +1 @@
-VERSION = "2.3.1"
+VERSION = "2.4.0"