From 9285f7edc1905676ad85971ea91f84d6770cc0c4 Mon Sep 17 00:00:00 2001
From: Vitaly Baranov <vitbar@yandex-team.ru>
Date: Sat, 7 Nov 2020 21:00:55 +0300
Subject: [PATCH] Send profile info to client.

---
 src/Server/GRPCServer.cpp                    | 13 +++++++++++++
 src/Server/grpc_protos/clickhouse_grpc.proto | 11 ++++++++++-
 tests/integration/test_grpc_protocol/test.py |  9 ++++++++-
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/Server/GRPCServer.cpp b/src/Server/GRPCServer.cpp
index e9c100f7c1f..c43cef3cb8e 100644
--- a/src/Server/GRPCServer.cpp
+++ b/src/Server/GRPCServer.cpp
@@ -208,6 +208,7 @@ namespace
         void addProgressToResult();
         void addTotalsToResult(const Block & totals);
         void addExtremesToResult(const Block & extremes);
+        void addProfileInfoToResult(const BlockStreamProfileInfo & info);
         void addLogsToResult();
         void sendResult();
         void throwIfFailedToSendResult();
@@ -579,6 +580,7 @@ namespace
 
         addTotalsToResult(io.in->getTotals());
         addExtremesToResult(io.in->getExtremes());
+        addProfileInfoToResult(io.in->getProfileInfo());
     }
 
     void Call::generateOutputWithProcessors()
@@ -619,6 +621,7 @@ namespace
 
         addTotalsToResult(executor->getTotalsBlock());
         addExtremesToResult(executor->getExtremesBlock());
+        addProfileInfoToResult(executor->getProfileInfo());
     }
 
     void Call::finishQuery()
@@ -766,6 +769,16 @@ namespace
         stream->writeSuffix();
     }
 
+    void Call::addProfileInfoToResult(const BlockStreamProfileInfo & info)
+    {
+        auto & stats = *result.mutable_stats();
+        stats.set_rows(info.rows);
+        stats.set_blocks(info.blocks);
+        stats.set_allocated_bytes(info.bytes);
+        stats.set_applied_limit(info.hasAppliedLimit());
+        stats.set_rows_before_limit(info.getRowsBeforeLimit());
+    }
+
     void Call::addLogsToResult()
     {
         if (!logs_queue)
diff --git a/src/Server/grpc_protos/clickhouse_grpc.proto b/src/Server/grpc_protos/clickhouse_grpc.proto
index 0fa01645825..87b61295a58 100644
--- a/src/Server/grpc_protos/clickhouse_grpc.proto
+++ b/src/Server/grpc_protos/clickhouse_grpc.proto
@@ -46,6 +46,14 @@ message Progress {
    uint64 written_bytes = 5;
 }
 
+message Stats {
+   uint64 rows = 1;
+   uint64 blocks = 2;
+   uint64 allocated_bytes = 3;
+   bool applied_limit = 4;
+   uint64 rows_before_limit = 5;
+}
+
 message Exception {
    int32 code = 1;
    string name = 2;
@@ -59,7 +67,8 @@ message Result {
    string extremes = 3;
    repeated LogEntry logs = 4;
    Progress progress = 5;
-   Exception exception = 6;
+   Stats stats = 6;
+   Exception exception = 7;
 }
 
 service ClickHouse {
diff --git a/tests/integration/test_grpc_protocol/test.py b/tests/integration/test_grpc_protocol/test.py
index 10162a692dc..b6acd2451c2 100644
--- a/tests/integration/test_grpc_protocol/test.py
+++ b/tests/integration/test_grpc_protocol/test.py
@@ -204,4 +204,11 @@ def test_progress():
   read_bytes: 16
 }
 , output: "6\\t0\\n7\\t0\\n"
-, ]"""
+, stats {
+  rows: 8
+  blocks: 4
+  allocated_bytes: 324
+  applied_limit: true
+  rows_before_limit: 8
+}
+]"""