Merge remote-tracking branch 'origin/main' into feat/upgrade-066

Author: hantmac

const.go

@@ -4,6 +4,7 @@ const (
 	DatabendTenantHeader    = "X-DATABEND-TENANT"
 	DatabendWarehouseHeader = "X-DATABEND-WAREHOUSE"
 	DatabendQueryIDHeader   = "X-DATABEND-QUERY-ID"
+	DatabendQueryIDNode     = "X-DATABEND-NODE-ID"
 	Authorization           = "Authorization"
 	WarehouseRoute          = "X-DATABEND-ROUTE"
 	UserAgent               = "User-Agent"

query.go

@@ -30,6 +30,7 @@ type DataField struct {
 
 type QueryResponse struct {
 	ID      string           `json:"id"`
+	NodeID  string           `json:"node_id"`
 	Session *json.RawMessage `json:"session"`
 	Schema  *[]DataField     `json:"schema"`
 	Data    [][]string       `json:"data"`

restful.go

@@ -80,6 +80,7 @@ func (c *APIClient) NewDefaultCopyOptions() map[string]string {
 type APIClient struct {
 	SessionID string
 	QuerySeq  int64
+	NodeID    string
 	cli       *http.Client
 	rows      *nextRows
 
@@ -219,6 +220,9 @@ func (c *APIClient) doRequest(ctx context.Context, method, path string, req inte
 		if err != nil {
 			return errors.Wrap(err, "failed to make request headers")
 		}
+		if method == "GET" && len(c.NodeID) != 0 {
+			headers.Set(DatabendQueryIDNode, c.NodeID)
+		}
 		headers.Set(contentType, jsonContentType)
 		headers.Set(accept, jsonContentType)
 		httpReq.Header = headers
@@ -466,6 +470,7 @@ func (c *APIClient) startQueryRequest(ctx context.Context, request *QueryRequest
 	// try update session as long as resp is not nil, even if query failed (resp.Error != nil)
 	// e.g. transaction state need to be updated if commit fail
 	c.applySessionState(&resp)
+	c.NodeID = resp.NodeID
 	c.trackStats(&resp)
 	return &resp, nil
 }