From a3050ba9383ad9544a541593995ed65d6b38015e Mon Sep 17 00:00:00 2001
From: Yang Xiufeng <yangxiufeng.c@gmail.com>
Date: Fri, 29 Nov 2024 13:47:41 +0800
Subject: [PATCH] fix: node_id may be empty in query response.

---
 core/src/client.rs     | 18 +++++++++++-------
 core/src/response.rs   |  3 ++-
 driver/src/rest_api.rs |  6 ++++--
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/core/src/client.rs b/core/src/client.rs
index 2b085444..4ac6f393 100644
--- a/core/src/client.rs
+++ b/core/src/client.rs
@@ -405,18 +405,20 @@ impl APIClient {
         &self,
         query_id: &str,
         next_uri: &str,
-        node_id: &str,
+        node_id: &Option<String>,
     ) -> Result<QueryResponse> {
         info!("query page: {}", next_uri);
         let endpoint = self.endpoint.join(next_uri)?;
         let headers = self.make_headers(Some(query_id))?;
         let mut builder = self.cli.get(endpoint.clone());
-        builder = self.wrap_auth_or_session_token(builder)?;
-        let request = builder
+        builder = self
+            .wrap_auth_or_session_token(builder)?
             .headers(headers.clone())
-            .header(HEADER_STICKY_NODE, node_id)
-            .timeout(self.page_request_timeout)
-            .build()?;
+            .timeout(self.page_request_timeout);
+        if let Some(node_id) = node_id {
+            builder = builder.header(HEADER_STICKY_NODE, node_id)
+        }
+        let request = builder.build()?;
 
         let response = self.query_request_helper(request, false, true).await?;
         let body = response.bytes().await?;
@@ -452,7 +454,9 @@ impl APIClient {
     pub async fn wait_for_query(&self, resp: QueryResponse) -> Result<QueryResponse> {
         info!("wait for query: {}", resp.id);
         let node_id = resp.node_id.clone();
-        self.set_last_node_id(node_id.clone());
+        if let Some(node_id) = self.last_node_id() {
+            self.set_last_node_id(node_id.clone());
+        }
         if let Some(next_uri) = &resp.next_uri {
             let schema = resp.schema;
             let mut data = resp.data;
diff --git a/core/src/response.rs b/core/src/response.rs
index 1f8b3c42..3df5d81f 100644
--- a/core/src/response.rs
+++ b/core/src/response.rs
@@ -48,7 +48,8 @@ pub struct SchemaField {
 #[derive(Deserialize, Debug)]
 pub struct QueryResponse {
     pub id: String,
-    pub node_id: String,
+    // response from gateway do not contain node_id
+    pub node_id: Option<String>,
     pub session_id: Option<String>,
     pub session: Option<SessionState>,
     pub schema: Vec<SchemaField>,
diff --git a/driver/src/rest_api.rs b/driver/src/rest_api.rs
index eb6158a9..da45c7fa 100644
--- a/driver/src/rest_api.rs
+++ b/driver/src/rest_api.rs
@@ -208,7 +208,9 @@ impl<'o> RestAPIConnection {
             return Ok(resp);
         }
         let node_id = resp.node_id.clone();
-        self.client.set_last_node_id(node_id.clone());
+        if let Some(node_id) = &node_id {
+            self.client.set_last_node_id(node_id.clone());
+        }
         let mut result = resp;
         // preserve schema since it is not included in the final response
         while let Some(next_uri) = result.next_uri {
@@ -247,7 +249,7 @@ pub struct RestAPIRows {
     data: VecDeque<Vec<Option<String>>>,
     stats: Option<ServerStats>,
     query_id: String,
-    node_id: String,
+    node_id: Option<String>,
     next_uri: Option<String>,
     next_page: Option<PageFut>,
 }