From 1e18ba19c83d4dde3271c433104b5cc00888b5b6 Mon Sep 17 00:00:00 2001
From: Haojun Liao <hjliao@taosdata.com>
Date: Wed, 6 Jul 2022 17:25:39 +0800
Subject: [PATCH] fix(query): set results for local query in async APIs.

---
 include/libs/command/command.h        |  5 ++++-
 source/client/src/clientMain.c        | 28 ++++++++++++++++-----------
 source/libs/command/src/command.c     |  2 --
 source/libs/scheduler/src/scheduler.c |  3 ---
 4 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/include/libs/command/command.h b/include/libs/command/command.h
index aee6b83783..8a4ecad37d 100644
--- a/include/libs/command/command.h
+++ b/include/libs/command/command.h
@@ -13,6 +13,9 @@
  * along with this program. If not, see <http://www.gnu.org/licenses/>.
  */
 
+#ifndef TDENGINE_COMMAND_H
+#define TDENGINE_COMMAND_H
+
 #include "cmdnodes.h"
 #include "tmsg.h"
 #include "plannodes.h"
@@ -27,4 +30,4 @@ int32_t qExecExplainEnd(SExplainCtx *pCtx, SRetrieveTableRsp **pRsp);
 int32_t qExplainUpdateExecInfo(SExplainCtx *pCtx, SExplainRsp *pRspMsg, int32_t groupId, SRetrieveTableRsp **pRsp);
 void    qExplainFreeCtx(SExplainCtx *pCtx);
 
-
+#endif
diff --git a/source/client/src/clientMain.c b/source/client/src/clientMain.c
index a30d60a589..a97937aeec 100644
--- a/source/client/src/clientMain.c
+++ b/source/client/src/clientMain.c
@@ -843,19 +843,25 @@ void taos_fetch_rows_a(TAOS_RES *res, __taos_async_fn_t fp, void *param) {
   pRequest->body.param = param;
 
   SReqResultInfo *pResultInfo = &pRequest->body.resInfo;
-  if (taos_num_fields(pRequest) == 0) {
+
+  // this query has no results or error exists, return directly
+  if (taos_num_fields(pRequest) == 0 || pRequest->code != TSDB_CODE_SUCCESS) {
     pResultInfo->numOfRows = 0;
     pRequest->body.fetchFp(param, pRequest, pResultInfo->numOfRows);
     return;
   }
 
-  if (pResultInfo->pData == NULL || pResultInfo->current >= pResultInfo->numOfRows) {
-    // All data has returned to App already, no need to try again
-    if (pResultInfo->completed) {
-      pResultInfo->numOfRows = 0;
-      pRequest->body.fetchFp(param, pRequest, pResultInfo->numOfRows);
-      return;
-    }
+  // all data has returned to App already, no need to try again
+  if ((pResultInfo->pData == NULL || pResultInfo->current >= pResultInfo->numOfRows) && pResultInfo->completed) {
+    pResultInfo->numOfRows = 0;
+    pRequest->body.fetchFp(param, pRequest, pResultInfo->numOfRows);
+    return;
+  }
+
+  // it is a local executed query, no need to do async fetch
+  if (pResultInfo->current < pResultInfo->numOfRows && pRequest->body.queryJob == 0) {
+    pRequest->body.fetchFp(param, pRequest, pResultInfo->numOfRows);
+    return;
   }
 
   schedulerAsyncFetchRows(pRequest->body.queryJob, fetchCallback, pRequest);
@@ -864,14 +870,14 @@ void taos_fetch_rows_a(TAOS_RES *res, __taos_async_fn_t fp, void *param) {
 void taos_fetch_raw_block_a(TAOS_RES *res, __taos_async_fn_t fp, void *param) {
   ASSERT(res != NULL && fp != NULL);
   ASSERT(TD_RES_QUERY(res));
+
   SRequestObj *pRequest = res;
-
-  pRequest->body.resInfo.convertUcs4 = false;
-
   SReqResultInfo *pResultInfo = &pRequest->body.resInfo;
 
   // set the current block is all consumed
   pResultInfo->current = pResultInfo->numOfRows;
+  pResultInfo->convertUcs4 = false;
+
   taos_fetch_rows_a(res, fp, param);
 }
 
diff --git a/source/libs/command/src/command.c b/source/libs/command/src/command.c
index fc3e3cbc8a..7f70a78b12 100644
--- a/source/libs/command/src/command.c
+++ b/source/libs/command/src/command.c
@@ -565,7 +565,6 @@ static int32_t createSelectResultDataBlock(SNodeList* pProjects, SSDataBlock** p
 }
 
 int32_t buildSelectResultDataBlock(SNodeList* pProjects, SSDataBlock* pBlock) {
-  int32_t numOfCols = LIST_LENGTH(pProjects);
   blockDataEnsureCapacity(pBlock, 1);
 
   int32_t index = 0;
@@ -579,7 +578,6 @@ int32_t buildSelectResultDataBlock(SNodeList* pProjects, SSDataBlock* pBlock) {
   }
 
   pBlock->info.rows = 1;
-
   return TSDB_CODE_SUCCESS;
 }
 
diff --git a/source/libs/scheduler/src/scheduler.c b/source/libs/scheduler/src/scheduler.c
index e2389c2a75..7f88df14e3 100644
--- a/source/libs/scheduler/src/scheduler.c
+++ b/source/libs/scheduler/src/scheduler.c
@@ -13,13 +13,10 @@
  * along with this program. If not, see <http://www.gnu.org/licenses/>.
  */
 
-#include "catalog.h"
-#include "command.h"
 #include "query.h"
 #include "schedulerInt.h"
 #include "tmsg.h"
 #include "tref.h"
-#include "trpc.h"
 
 SSchedulerMgmt schMgmt = {
     .jobRef = -1,