fix: recreate http request on each auth retry and clean up client state (#213)

- Move http.NewRequest inside auth retry loop in doRequest/doRequestRaw
  so the request body is not consumed on the first attempt
- Use headers.Clone() per iteration to avoid header mutation across retries
- Move snapshotClientState and clone helpers to clientstate.go
- Clear resp.Data after JSON materialization to prevent dual-path state
- Fix materializeJSONQueryRows early-return guard for empty typedRows slice
- Add internal consistency check in rows.Next for dual typed/raw row state
- Expose RowCount() and CellString() as public accessors on QueryResponse

Author: youngsofun

arrow.go

@@ -92,11 +92,6 @@ func (c *APIClient) doRequestRaw(
 	}
 
 	url := c.makeURL(path)
-	httpReq, err := http.NewRequest(method, url, bytes.NewBuffer(reqBody))
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to create http request")
-	}
-	httpReq = httpReq.WithContext(ctx)
 
 	headers, err := c.makeHeaders(ctx)
 	if err != nil {
@@ -113,11 +108,6 @@ func (c *APIClient) doRequestRaw(
 		acceptType = jsonContentType
 	}
 	headers.Set(accept, acceptType)
-	httpReq.Header = headers
-
-	if len(c.host) > 0 {
-		httpReq.Host = c.host
-	}
 
 	authRetryLimit := 2
 	for i := 1; i <= authRetryLimit; i++ {
@@ -127,6 +117,16 @@ func (c *APIClient) doRequestRaw(
 		default:
 		}
 
+		httpReq, err := http.NewRequest(method, url, bytes.NewBuffer(reqBody))
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to create http request")
+		}
+		httpReq = httpReq.WithContext(ctx)
+		httpReq.Header = headers.Clone()
+		if len(c.host) > 0 {
+			httpReq.Host = c.host
+		}
+
 		httpResp, err := c.cli.Do(httpReq)
 		if err != nil {
 			return nil, errors.Wrap(ErrDoRequest, err.Error())
@@ -164,6 +164,7 @@ func (c *APIClient) doRequestRaw(
 		}, nil
 	}
 
+	// unreachable: loop always returns inside on the final iteration
 	return nil, errors.Errorf("failed to do request after %d retries", authRetryLimit)
 }
 
@@ -615,6 +616,9 @@ func formatArrowTimestampTZDecimalValue(value decimal128.Num) (string, error) {
 	return ts.In(zone).Format("2006-01-02 15:04:05.000000 -0700"), nil
 }
 
+// arrowTimestampTZDecimalToTime decodes a Timestamp_Tz value encoded as Decimal128.
+// The server encodes: LowBits = microseconds since Unix epoch (uint64 cast to int64),
+// HighBits = timezone offset in seconds from UTC (uint64 cast to int32).
 func arrowTimestampTZDecimalToTime(value decimal128.Num) time.Time {
 	ts := time.UnixMicro(clampArrowTimestampMicros(int64(value.LowBits())))
 	zone := time.FixedZone("", int(int32(value.HighBits())))
@@ -678,52 +682,6 @@ func clampArrowTimestampMicros(value int64) int64 {
 	}
 }
 
-func (c *APIClient) snapshotClientState() func() {
-	querySeq := c.QuerySeq
-	routeHint := c.routeHint
-	nodeID := c.nodeID
-	stateRestored := c.stateRestored
-	sessionStateRaw := cloneRawMessage(c.sessionStateRaw)
-	sessionState := cloneSessionState(c.sessionState)
-
-	return func() {
-		c.QuerySeq = querySeq
-		c.routeHint = routeHint
-		c.nodeID = nodeID
-		c.stateRestored = stateRestored
-		c.sessionStateRaw = sessionStateRaw
-		c.sessionState = sessionState
-	}
-}
-
-func cloneRawMessage(raw *json.RawMessage) *json.RawMessage {
-	if raw == nil {
-		return nil
-	}
-	cloned := json.RawMessage(append([]byte(nil), (*raw)...))
-	return &cloned
-}
-
-func cloneSessionState(state *SessionState) *SessionState {
-	if state == nil {
-		return nil
-	}
-
-	cloned := *state
-	if state.SecondaryRoles != nil {
-		roles := append([]string(nil), (*state.SecondaryRoles)...)
-		cloned.SecondaryRoles = &roles
-	}
-	if state.Settings != nil {
-		settings := make(map[string]string, len(state.Settings))
-		for key, value := range state.Settings {
-			settings[key] = value
-		}
-		cloned.Settings = settings
-	}
-	return &cloned
-}
-
 func formatArrowScalar(value interface{}) (string, error) {
 	switch v := value.(type) {
 	case nil:

arrow_test.go

@@ -442,9 +442,10 @@ func TestQuerySyncFallsBackToJSONWhenArrowRequested(t *testing.T) {
 
 	resp, err := client.QuerySync(context.Background(), "SELECT 1")
 	require.NoError(t, err)
-	require.Len(t, resp.Data, 1)
-	require.NotNil(t, resp.Data[0][0])
-	assert.Equal(t, "1", *resp.Data[0][0])
+	require.Equal(t, 1, resp.RowCount())
+	v, ok := resp.CellString(0, 0)
+	require.True(t, ok)
+	assert.Equal(t, "1", v)
 
 	mu.Lock()
 	defer mu.Unlock()
@@ -495,9 +496,10 @@ func TestQuerySyncUsesJSONForRestoredState(t *testing.T) {
 
 	resp, err := client.QuerySync(context.Background(), "SELECT 1")
 	require.NoError(t, err)
-	require.Len(t, resp.Data, 1)
-	require.NotNil(t, resp.Data[0][0])
-	assert.Equal(t, "1", *resp.Data[0][0])
+	require.Equal(t, 1, resp.RowCount())
+	v, ok := resp.CellString(0, 0)
+	require.True(t, ok)
+	assert.Equal(t, "1", v)
 
 	mu.Lock()
 	defer mu.Unlock()

client.go

@@ -251,30 +251,19 @@ func (c *APIClient) doRequest(ctx context.Context, method, path string, req inte
 	}
 
 	url := c.makeURL(path)
-	httpReq, err := http.NewRequest(method, url, bytes.NewBuffer(reqBody))
-	if err != nil {
-		return errors.Wrap(err, "failed to create http request")
-	}
-
-	httpReq = httpReq.WithContext(ctx)
 
 	headers, err := c.makeHeaders(ctx)
-	if needSticky && len(c.nodeID) != 0 {
-		headers.Set(DatabendQueryStickyNode, c.nodeID)
-	}
 	if err != nil {
 		return errors.Wrap(err, "failed to make request headers")
 	}
+	if needSticky && len(c.nodeID) != 0 {
+		headers.Set(DatabendQueryStickyNode, c.nodeID)
+	}
 	if method == "GET" && len(c.nodeID) != 0 {
 		headers.Set(DatabendQueryIDNode, c.nodeID)
 	}
 	headers.Set(contentType, jsonContentType)
 	headers.Set(accept, jsonContentType)
-	httpReq.Header = headers
-
-	if len(c.host) > 0 {
-		httpReq.Host = c.host
-	}
 
 	authRetryLimit := 2
 	for i := 1; i <= authRetryLimit; i++ {
@@ -284,6 +273,16 @@ func (c *APIClient) doRequest(ctx context.Context, method, path string, req inte
 		default:
 		}
 
+		httpReq, err := http.NewRequest(method, url, bytes.NewBuffer(reqBody))
+		if err != nil {
+			return errors.Wrap(err, "failed to create http request")
+		}
+		httpReq = httpReq.WithContext(ctx)
+		httpReq.Header = headers.Clone()
+		if len(c.host) > 0 {
+			httpReq.Host = c.host
+		}
+
 		httpResp, err := c.cli.Do(httpReq)
 		if err != nil {
 			return errors.Wrap(ErrDoRequest, err.Error())

clientstate.go

@@ -5,6 +5,52 @@ import (
 	"net/http"
 )
 
+func (c *APIClient) snapshotClientState() func() {
+	querySeq := c.QuerySeq
+	routeHint := c.routeHint
+	nodeID := c.nodeID
+	stateRestored := c.stateRestored
+	sessionStateRaw := cloneRawMessage(c.sessionStateRaw)
+	sessionState := cloneSessionState(c.sessionState)
+
+	return func() {
+		c.QuerySeq = querySeq
+		c.routeHint = routeHint
+		c.nodeID = nodeID
+		c.stateRestored = stateRestored
+		c.sessionStateRaw = sessionStateRaw
+		c.sessionState = sessionState
+	}
+}
+
+func cloneRawMessage(raw *json.RawMessage) *json.RawMessage {
+	if raw == nil {
+		return nil
+	}
+	cloned := json.RawMessage(append([]byte(nil), (*raw)...))
+	return &cloned
+}
+
+func cloneSessionState(state *SessionState) *SessionState {
+	if state == nil {
+		return nil
+	}
+
+	cloned := *state
+	if state.SecondaryRoles != nil {
+		roles := append([]string(nil), (*state.SecondaryRoles)...)
+		cloned.SecondaryRoles = &roles
+	}
+	if state.Settings != nil {
+		settings := make(map[string]string, len(state.Settings))
+		for key, value := range state.Settings {
+			settings[key] = value
+		}
+		cloned.Settings = settings
+	}
+	return &cloned
+}
+
 type APIClientState struct {
 	SessionID    string
 	QuerySeq     int64

query.go

@@ -62,6 +62,10 @@ func (r *QueryResponse) ReadFinished() bool {
 	return r.NextURI == "" || strings.Contains(r.NextURI, "/final")
 }
 
+func (r *QueryResponse) RowCount() int {
+	return r.bufferedRowCount()
+}
+
 func (r *QueryResponse) bufferedRowCount() int {
 	if r == nil {
 		return 0
@@ -85,6 +89,10 @@ func (r *QueryResponse) cellValue(rowIdx, colIdx int) (driver.Value, bool) {
 	return nil, false
 }
 
+func (r *QueryResponse) CellString(rowIdx, colIdx int) (string, bool) {
+	return r.cellString(rowIdx, colIdx)
+}
+
 func (r *QueryResponse) cellString(rowIdx, colIdx int) (string, bool) {
 	value, ok := r.cellValue(rowIdx, colIdx)
 	if !ok || value == nil {

result_rows.go

@@ -28,7 +28,7 @@ func queryResponseColumnTypeOptions(settings *Settings) (*ColumnTypeOptions, err
 }
 
 func materializeJSONQueryRows(resp *QueryResponse) error {
-	if resp == nil || resp.typedRows != nil || len(resp.Data) == 0 {
+	if resp == nil || (resp.typedRows != nil && len(resp.typedRows) > 0) || len(resp.Data) == 0 {
 		return nil
 	}
 	if resp.Schema == nil || len(*resp.Schema) != len(resp.Data[0]) {
@@ -62,6 +62,7 @@ func materializeJSONQueryRows(resp *QueryResponse) error {
 		rows = append(rows, row)
 	}
 	resp.typedRows = rows
+	resp.Data = nil
 	return nil
 }
 

rows.go

@@ -173,6 +173,9 @@ func (r *nextRows) Next(dest []driver.Value) error {
 		typedRow = r.respData.typedRows[0]
 		r.respData.typedRows = r.respData.typedRows[1:]
 	}
+	if len(typedRow) != 0 && len(lineData) != 0 {
+		return errors.New("query error: internal error, both typed and raw row data are set")
+	}
 	if len(typedRow) != 0 {
 		if len(typedRow) != len(r.columns) {
 			return errors.New("query error: internal error, typed data and schema not match")

tests/resume_query_test.go

@@ -43,6 +43,7 @@ func (s *DatabendTestSuite) TestResumeQueryWithSessionState() {
 	finalResp, err := secondClient.PollUntilQueryEnd(ctx, resumeResp)
 	s.Require().NoError(err)
 	s.Require().NotNil(finalResp)
+	s.Greater(finalResp.RowCount(), 0)
 	s.NoError(secondClient.CloseQuery(ctx, finalResp))
 }
 
@@ -65,10 +66,10 @@ func (s *DatabendTestSuite) TestSessionSettingLoadWithState() {
 	client2.WithState(state)
 	resp, err := client2.QuerySync(ctx, fmt.Sprintf("SELECT value FROM system.settings WHERE name = '%s'", settingKey))
 	s.Require().NoError(err)
-	s.Require().Greater(len(resp.Data), 0)
-	s.Require().Greater(len(resp.Data[0]), 0)
-	s.Require().NotNil(resp.Data[0][0])
-	s.Equal(fmt.Sprintf("%d", settingValue), *resp.Data[0][0])
+	s.Require().Greater(resp.RowCount(), 0)
+	v, ok := resp.CellString(0, 0)
+	s.Require().True(ok)
+	s.Equal(fmt.Sprintf("%d", settingValue), v)
 
 	roundedState := client2.GetState()
 	s.Require().NotNil(roundedState)