Test: Support coprocessor test (#5957)

ref #4609

dbms/src/Debug/MockComputeServerManager.cpp

@@ -25,7 +25,7 @@ extern const int IP_ADDRESS_NOT_ALLOWED;
 } // namespace ErrorCodes
 namespace tests
 {
-void MockComputeServerManager::addServer(String addr)
+void MockComputeServerManager::addServer(const String & addr)
 {
     MockServerConfig config;
     for (const auto & server : server_config_map)
@@ -135,4 +135,4 @@ String MockComputeServerManager::queryInfo()
     return buf.toString();
 }
 } // namespace tests
-} // namespace DB
+} // namespace DB

dbms/src/Debug/MockComputeServerManager.h

@@ -30,7 +30,7 @@ class MockComputeServerManager : public ext::Singleton<MockComputeServerManager>
 {
 public:
     /// register an server to run.
-    void addServer(String addr);
+    void addServer(const String & addr);
 
     /// call startServers to run all servers in current test.
     void startServers(const LoggerPtr & log_ptr, Context & global_context);
@@ -61,4 +61,4 @@ class MockComputeServerManager : public ext::Singleton<MockComputeServerManager>
     std::unordered_map<size_t, std::unique_ptr<FlashGrpcServerHolder>> server_map;
     std::unordered_map<size_t, MockServerConfig> server_config_map;
 };
-} // namespace DB::tests
+} // namespace DB::tests

dbms/src/Debug/dbgFuncCoprocessor.cpp

@@ -136,7 +136,6 @@ class UniqRawResReformatBlockOutputStream : public IProfilingBlockInputStream
 tipb::SelectResponse executeDAGRequest(Context & context, const tipb::DAGRequest & dag_request, RegionID region_id, UInt64 region_version, UInt64 region_conf_version, Timestamp start_ts, std::vector<std::pair<DecodedTiKVKeyPtr, DecodedTiKVKeyPtr>> & key_ranges);
 bool runAndCompareDagReq(const coprocessor::Request & req, const coprocessor::Response & res, Context & context, String & unequal_msg);
 BlockInputStreamPtr outputDAGResponse(Context & context, const DAGSchema & schema, const tipb::SelectResponse & dag_response);
-DAGSchema getSelectSchema(Context & context);
 bool dagRspEqual(Context & context, const tipb::SelectResponse & expected, const tipb::SelectResponse & actual, String & unequal_msg);
 
 DAGProperties getDAGProperties(const String & prop_string)

dbms/src/Debug/dbgFuncCoprocessor.h

@@ -87,6 +87,8 @@ QueryTasks queryPlanToQueryTasks(
 BlockInputStreamPtr executeQuery(Context & context, RegionID region_id, const DAGProperties & properties, QueryTasks & query_tasks, MakeResOutputStream & func_wrap_output_stream);
 BlockInputStreamPtr executeMPPQuery(Context & context, const DAGProperties & properties, QueryTasks & query_tasks);
 std::vector<BlockInputStreamPtr> executeMPPQueryWithMultipleContext(const DAGProperties & properties, QueryTasks & query_tasks, std::unordered_map<size_t, MockServerConfig> & server_config_map);
+DAGSchema getSelectSchema(Context & context);
+std::unique_ptr<ChunkCodec> getCodec(tipb::EncodeType encode_type);
 
 namespace Debug
 {

dbms/src/Flash/FlashService.cpp

@@ -105,7 +105,8 @@ grpc::Status FlashService::Coprocessor(
     CPUAffinityManager::getInstance().bindSelfGrpcThread();
     LOG_FMT_DEBUG(log, "Handling coprocessor request: {}", request->DebugString());
 
-    if (!security_config->checkGrpcContext(grpc_context))
+    // For coprocessor test, we don't care about security config.
+    if (unlikely(!context->isCopTest() && !security_config->checkGrpcContext(grpc_context)))
     {
         return grpc::Status(grpc::PERMISSION_DENIED, tls_err_msg);
     }
@@ -119,6 +120,8 @@ grpc::Status FlashService::Coprocessor(
         GET_METRIC(tiflash_coprocessor_response_bytes).Increment(response->ByteSizeLong());
     });
 
+    context->setMockStorage(mock_storage);
+
     grpc::Status ret = executeInThreadPool(*cop_pool, [&] {
         auto [db_context, status] = createDBContext(grpc_context);
         if (!status.ok())
@@ -430,8 +433,8 @@ std::tuple<ContextPtr, grpc::Status> FlashService::createDBContext(const grpc::S
         std::string client_ip = peer.substr(pos + 1);
         Poco::Net::SocketAddress client_address(client_ip);
 
-        // For MPP test, we don't care about security config.
-        if (!context->isMPPTest())
+        // For MPP or Cop test, we don't care about security config.
+        if (likely(!context->isTest()))
             tmp_context->setUser(user, password, client_address, quota_key);
 
         String query_id = getClientMetaVarWithDefault(grpc_context, "query_id", "");

dbms/src/Flash/tests/gtest_compute_server.cpp

@@ -381,5 +381,23 @@ try
     }
 }
 CATCH
+
+TEST_F(ComputeServerRunner, runCoprocessor)
+try
+{
+    // In coprocessor test, we only need to start 1 server.
+    startServers(1);
+    {
+        auto request = context
+                           .scan("test_db", "l_table")
+                           .build(context);
+
+        auto expected_cols = {
+            toNullableVec<String>({{"banana", {}, "banana"}}),
+            toNullableVec<String>({{"apple", {}, "banana"}})};
+        ASSERT_COLUMNS_EQ_UR(expected_cols, executeCoprocessorTask(request));
+    }
+}
+CATCH
 } // namespace tests
-} // namespace DB
+} // namespace DB

dbms/src/Interpreters/Context.cpp

@@ -1866,6 +1866,16 @@ void Context::setExecutorTest()
     test_mode = executor_test;
 }
 
+bool Context::isCopTest() const
+{
+    return test_mode == cop_test;
+}
+
+void Context::setCopTest()
+{
+    test_mode = cop_test;
+}
+
 bool Context::isTest() const
 {
     return test_mode != non_test;

dbms/src/Interpreters/Context.h

@@ -161,6 +161,7 @@ class Context
     {
         non_test,
         mpp_test,
+        cop_test,
         executor_test,
         cancel_test
     };
@@ -478,6 +479,8 @@ class Context
     void setCancelTest();
     bool isExecutorTest() const;
     void setExecutorTest();
+    void setCopTest();
+    bool isCopTest() const;
     bool isTest() const;
 
     void setMockStorage(MockStorage & mock_storage_);

dbms/src/Server/MockComputeClient.h

@@ -43,6 +43,19 @@ class MockComputeClient
         }
     }
 
+    coprocessor::Response runCoprocessor(std::shared_ptr<coprocessor::Request> request)
+    {
+        coprocessor::Response response;
+        grpc::ClientContext context;
+        Status status = stub->Coprocessor(&context, *request, &response);
+        if (!status.ok())
+        {
+            throw Exception(fmt::format("Meet error while run coprocessor task, error code = {}, message = {}", status.error_code(), status.error_message()));
+        }
+
+        return response;
+    }
+
 private:
     std::unique_ptr<tikvpb::Tikv::Stub> stub{};
 };

dbms/src/TestUtils/ExecutorTestUtils.cpp

@@ -164,8 +164,6 @@ void ExecutorTest::executeAndAssertRowsEqual(const std::shared_ptr<tipb::DAGRequ
     });
 }
 
-namespace
-{
 Block mergeBlocks(Blocks blocks)
 {
     if (blocks.empty())
@@ -193,7 +191,6 @@ Block mergeBlocks(Blocks blocks)
         actual_columns.push_back({std::move(actual_cols[i]), sample_block.getColumnsWithTypeAndName()[i].type, sample_block.getColumnsWithTypeAndName()[i].name, sample_block.getColumnsWithTypeAndName()[i].column_id});
     return Block(actual_columns);
 }
-} // namespace
 
 void readStream(Blocks & blocks, BlockInputStreamPtr stream)
 {

dbms/src/TestUtils/ExecutorTestUtils.h

@@ -31,6 +31,8 @@ TiDB::TP dataTypeToTP(const DataTypePtr & type);
 
 DB::ColumnsWithTypeAndName readBlock(BlockInputStreamPtr stream);
 DB::ColumnsWithTypeAndName readBlocks(std::vector<BlockInputStreamPtr> streams);
+Block mergeBlocks(Blocks blocks);
+
 
 #define WRAP_FOR_DIS_ENABLE_PLANNER_BEGIN \
     std::vector<bool> bools{false, true}; \

dbms/src/TestUtils/MPPTaskTestUtils.cpp

@@ -11,6 +11,8 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+#include <Debug/dbgFuncCoprocessor.h>
+#include <Server/MockComputeClient.h>
 #include <TestUtils/MPPTaskTestUtils.h>
 
 namespace DB::tests
@@ -84,6 +86,46 @@ ColumnsWithTypeAndName MPPTaskTestUtils::exeucteMPPTasks(QueryTasks & tasks, con
     return readBlocks(res);
 }
 
+ColumnsWithTypeAndName extractColumns(Context & context, const std::shared_ptr<tipb::SelectResponse> & dag_response)
+{
+    auto codec = getCodec(dag_response->encode_type());
+    Blocks blocks;
+    auto schema = getSelectSchema(context);
+    for (const auto & chunk : dag_response->chunks())
+        blocks.emplace_back(codec->decode(chunk.rows_data(), schema));
+    return mergeBlocks(blocks).getColumnsWithTypeAndName();
+}
+
+ColumnsWithTypeAndName MPPTaskTestUtils::executeCoprocessorTask(std::shared_ptr<tipb::DAGRequest> & dag_request)
+{
+    assert(server_num == 1);
+    auto req = std::make_shared<coprocessor::Request>();
+    req->set_tp(103); // 103 is COP_REQ_TYPE_DAG
+    auto * data = req->mutable_data();
+    dag_request->AppendToString(data);
+
+    DAGContext dag_context(*dag_request);
+
+    TiFlashTestEnv::getGlobalContext(test_meta.context_idx).setDAGContext(&dag_context);
+    TiFlashTestEnv::getGlobalContext(test_meta.context_idx).setCopTest();
+
+    MockComputeServerManager::instance().setMockStorage(context.mockStorage());
+
+    auto addr = MockComputeServerManager::instance().getServerConfigMap()[0].addr; // Since we only have started 1 server currently.
+    MockComputeClient client(
+        grpc::CreateChannel(addr, grpc::InsecureChannelCredentials()));
+    auto resp = client.runCoprocessor(req);
+    auto resp_ptr = std::make_shared<tipb::SelectResponse>();
+    if (unlikely(!resp_ptr->ParseFromString(resp.data())))
+    {
+        throw Exception("Incorrect json response data from Coprocessor.", ErrorCodes::BAD_ARGUMENTS);
+    }
+    else
+    {
+        return extractColumns(TiFlashTestEnv::getGlobalContext(test_meta.context_idx), resp_ptr);
+    }
+}
+
 String MPPTaskTestUtils::queryInfo(size_t server_id)
 {
     FmtBuffer buf;
@@ -124,4 +166,4 @@ ::testing::AssertionResult MPPTaskTestUtils::assertQueryActive(size_t start_ts)
     }
     return ::testing::AssertionSuccess();
 }
-} // namespace DB::tests
+} // namespace DB::tests

dbms/src/TestUtils/MPPTaskTestUtils.h

@@ -83,6 +83,9 @@ class MPPTaskTestUtils : public ExecutorTest
     std::tuple<size_t, std::vector<BlockInputStreamPtr>> prepareMPPStreams(DAGRequestBuilder builder);
 
     ColumnsWithTypeAndName exeucteMPPTasks(QueryTasks & tasks, const DAGProperties & properties, std::unordered_map<size_t, MockServerConfig> & server_config_map);
+
+    ColumnsWithTypeAndName executeCoprocessorTask(std::shared_ptr<tipb::DAGRequest> & dag_request);
+
     static ::testing::AssertionResult assertQueryCancelled(size_t start_ts);
     static ::testing::AssertionResult assertQueryActive(size_t start_ts);
     static String queryInfo(size_t server_id);