fix: return InfluxQL JSON marshaling errors correctly (#27318)

JSON cannot marshal NaN and Inf, so results which
contain them should return marshaling errors

Closes #20258

Author: davidby-influx

influxql/query/response_writer.go

@@ -46,21 +46,49 @@ func (f *jsonFormatter) WriteResponse(ctx context.Context, w io.Writer, resp Res
 	span, _ := tracing.StartSpanFromContext(ctx)
 	defer span.Finish()
 
+	var wErr error
 	var b []byte
-	if f.Pretty {
-		b, err = json.MarshalIndent(resp, "", "    ")
-	} else {
-		b, err = json.Marshal(resp)
-	}
+	b, err = f.marshal(resp)
 
 	if err != nil {
-		_, err = io.WriteString(w, err.Error())
+		nakedErr := unnestError(err)
+		errResult := &Result{Err: nakedErr}
+		if len(resp.Results) >= 1 && resp.Results[0] != nil {
+			errResult.StatementID = resp.Results[0].StatementID
+		}
+		// We're so deep in errors we are going to drop this one.
+		b, _ = f.marshal(Response{Results: []*Result{errResult}})
+	}
+	if _, wErr = w.Write(b); wErr == nil {
+		_, wErr = w.Write([]byte("\n"))
+	}
+	return wErr
+}
+
+func (f *jsonFormatter) marshal(resp Response) ([]byte, error) {
+	if f.Pretty {
+		return json.MarshalIndent(resp, "", "    ")
 	} else {
-		_, err = w.Write(b)
+		return json.Marshal(resp)
 	}
+}
 
-	w.Write([]byte("\n"))
-	return err
+func unnestError(err error) error {
+	for err != nil {
+		switch e := err.(type) {
+		case interface{ Unwrap() error }:
+			err = e.Unwrap()
+		case interface{ Unwrap() []error }:
+			errs := e.Unwrap()
+			if len(errs) == 0 {
+				return err
+			}
+			err = errs[0]
+		default:
+			return err
+		}
+	}
+	return nil
 }
 
 type csvFormatter struct {

influxql/query/response_writer_test.go

@@ -0,0 +1,114 @@
+package query
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"math"
+	"testing"
+
+	"github.com/influxdata/influxdb/v2/models"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestJsonFormatter_WriteResponse_Success(t *testing.T) {
+	resp := Response{
+		Results: []*Result{{
+			StatementID: 0,
+			Series: models.Rows{
+				&models.Row{
+					Name:    "cpu",
+					Columns: []string{"time", "value"},
+					Values:  [][]interface{}{{"2021-01-01T00:00:00Z", 42.0}},
+				},
+			},
+		}},
+	}
+
+	for _, pretty := range []bool{false, true} {
+		name := "compact"
+		if pretty {
+			name = "pretty"
+		}
+		t.Run(name, func(t *testing.T) {
+			var buf bytes.Buffer
+			f := &jsonFormatter{Pretty: pretty}
+			err := f.WriteResponse(context.Background(), &buf, resp)
+			require.NoError(t, err)
+
+			output := buf.Bytes()
+			var parsed map[string]interface{}
+			require.NoError(t, json.Unmarshal(output, &parsed))
+
+			results, ok := parsed["results"].([]interface{})
+			require.True(t, ok, "expected results array, got: %s", output)
+			require.NotEmpty(t, results)
+
+			result0, ok := results[0].(map[string]interface{})
+			require.True(t, ok, "expected result object, got: %s", output)
+			assert.Nil(t, result0["error"], "unexpected error in result")
+
+			series, ok := result0["series"].([]interface{})
+			require.True(t, ok, "expected series array, got: %s", output)
+			require.NotEmpty(t, series)
+
+			if pretty {
+				assert.Contains(t, string(output), "    ", "expected indented output for pretty mode")
+			}
+		})
+	}
+}
+
+func TestJsonFormatter_WriteResponse_MarshalError(t *testing.T) {
+	const statementID = 12
+	tests := []struct {
+		name    string
+		value   interface{}
+		wantErr string
+	}{
+		{"Inf", math.Inf(1), "unsupported value"},
+		{"NaN", math.NaN(), "unsupported value"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			resp := Response{
+				Results: []*Result{{
+					StatementID: statementID,
+					Series: models.Rows{
+						&models.Row{
+							Name:    "cpu",
+							Columns: []string{"time", "value"},
+							Values:  [][]interface{}{{"2021-01-01T00:00:00Z", tt.value}},
+						},
+					},
+				}},
+			}
+
+			var buf bytes.Buffer
+			f := &jsonFormatter{Pretty: false}
+			err := f.WriteResponse(context.Background(), &buf, resp)
+			require.NoError(t, err)
+
+			output := buf.Bytes()
+			var parsed map[string]interface{}
+			require.NoError(t, json.Unmarshal(output, &parsed), "output is not valid JSON: %s", output)
+
+			results, ok := parsed["results"].([]interface{})
+			require.True(t, ok, "expected results array, got: %s", output)
+			require.NotEmpty(t, results)
+
+			result0, ok := results[0].(map[string]interface{})
+			require.True(t, ok, "expected result object, got: %s", output)
+
+			errMsg, ok := result0["error"].(string)
+			require.True(t, ok, "expected error string in result, got: %s", output)
+			require.Contains(t, errMsg, tt.wantErr)
+
+			sid, ok := result0["statement_id"].(float64)
+			require.True(t, ok, "expected statement_id in error result, got: %s", output)
+			assert.Equal(t, float64(statementID), sid, "statement_id should match original")
+		})
+	}
+}