From 079ba1536d8d697eac801d9d43de93c3803932e2 Mon Sep 17 00:00:00 2001
From: William Chen <13495049+CPWstatic@users.noreply.github.com>
Date: Sun, 27 Jun 2021 00:34:11 +0800
Subject: [PATCH] Support kill multi query.

---
 src/common/clients/meta/MetaClient.cpp | 7 +++----
 src/common/clients/meta/MetaClient.h   | 3 ++-
 src/common/interface/common.thrift     | 4 ++--
 src/common/interface/meta.thrift       | 4 ++--
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/common/clients/meta/MetaClient.cpp b/src/common/clients/meta/MetaClient.cpp
index 8eff2b47b..49fccb55d 100644
--- a/src/common/clients/meta/MetaClient.cpp
+++ b/src/common/clients/meta/MetaClient.cpp
@@ -3718,11 +3718,10 @@ folly::Future<StatusOr<cpp2::ExecResp>> MetaClient::removeSession(SessionID sess
     return future;
 }
 
-folly::Future<StatusOr<cpp2::ExecResp>> MetaClient::killQuery(SessionID sessionId,
-                                                              ExecutionPlanID epId) {
+folly::Future<StatusOr<cpp2::ExecResp>> MetaClient::killQuery(
+    std::unordered_map<SessionID, std::unordered_set<ExecutionPlanID>> killQueries) {
     cpp2::KillQueryReq req;
-    req.set_session_id(sessionId);
-    req.set_ep_id(epId);
+    req.set_kill_queries(std::move(killQueries));
     folly::Promise<StatusOr<cpp2::ExecResp>> promise;
     auto future = promise.getFuture();
     getResponse(std::move(req),
diff --git a/src/common/clients/meta/MetaClient.h b/src/common/clients/meta/MetaClient.h
index 28ef46178..ade2a958b 100644
--- a/src/common/clients/meta/MetaClient.h
+++ b/src/common/clients/meta/MetaClient.h
@@ -460,7 +460,8 @@ class MetaClient {
 
     folly::Future<StatusOr<cpp2::ExecResp>> removeSession(SessionID sessionId);
 
-    folly::Future<StatusOr<cpp2::ExecResp>> killQuery(SessionID sessionId, ExecutionPlanID epId);
+    folly::Future<StatusOr<cpp2::ExecResp>> killQuery(
+        std::unordered_map<SessionID, std::unordered_set<ExecutionPlanID>> killQueries);
 
     // Opeartions for cache.
     StatusOr<GraphSpaceID> getSpaceIdByNameFromCache(const std::string& name);
diff --git a/src/common/interface/common.thrift b/src/common/interface/common.thrift
index 010cd581e..2f52a193b 100644
--- a/src/common/interface/common.thrift
+++ b/src/common/interface/common.thrift
@@ -328,13 +328,13 @@ enum ErrorCode {
     E_RESTORE_FAILURE                 = -2068,
 
     E_SESSION_NOT_FOUND               = -2069,
-    E_QUERY_NOT_FOUND                 = -2070,
 
     // ListClusterInfo Failure
     E_LIST_CLUSTER_FAILURE              = -2070,
     E_LIST_CLUSTER_GET_ABS_PATH_FAILURE = -2071,
-    E_GET_META_DIR_FAILURE              = -2072,    
+    E_GET_META_DIR_FAILURE              = -2072,
 
+    E_QUERY_NOT_FOUND                 = -2073,
 
     // 3xxx for storaged
     E_CONSENSUS_ERROR                 = -3001,
diff --git a/src/common/interface/meta.thrift b/src/common/interface/meta.thrift
index 906388335..ffaab1faf 100644
--- a/src/common/interface/meta.thrift
+++ b/src/common/interface/meta.thrift
@@ -1138,8 +1138,8 @@ struct RemoveSessionReq {
 }
 
 struct KillQueryReq {
-    1: common.SessionID       session_id,
-    2: common.ExecutionPlanID ep_id,
+    1: map<common.SessionID, set<common.ExecutionPlanID> (cpp.template = "std::unordered_set")>
+        (cpp.template = "std::unordered_map") kill_queries,
 }
 
 struct ReportTaskReq {