From cb2b6043e09882eff83f4a60a75e4fbb68df9b20 Mon Sep 17 00:00:00 2001 From: Daniil Forshev Date: Tue, 20 Oct 2026 18:40:20 +0500 Subject: [PATCH 1/3] refactor(stream search): schema aware grpc gateway marshaler --- pkg/seqproxyapi/v1/marshaler.go | 83 +------ pkg/seqproxyapi/v1/marshaler_test.go | 77 +------ pkg/seqproxyapi/v1/stream_render.go | 165 ++++++++++++++ pkg/seqproxyapi/v1/stream_render_test.go | 235 ++++++++++++++++++++ proxyapi/grpc_stream_search.go | 67 ++---- proxyapi/ingestor.go | 1 + tests/integration_tests/integration_test.go | 12 +- 7 files changed, 435 insertions(+), 205 deletions(-) create mode 100644 pkg/seqproxyapi/v1/stream_render.go create mode 100644 pkg/seqproxyapi/v1/stream_render_test.go diff --git a/pkg/seqproxyapi/v1/marshaler.go b/pkg/seqproxyapi/v1/marshaler.go index 3f93fcfce..dd660a200 100644 --- a/pkg/seqproxyapi/v1/marshaler.go +++ b/pkg/seqproxyapi/v1/marshaler.go @@ -10,9 +10,6 @@ import ( "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/types/known/durationpb" "google.golang.org/protobuf/types/known/timestamppb" - - "github.com/ozontech/seq-db/query/encoding" - "github.com/ozontech/seq-db/seq" ) var ( @@ -62,11 +59,7 @@ type rawValueBucket struct { // MarshalJSON overrides "value" field to encode math.NaN and math.Inf as string. func (b *Aggregation_Bucket) MarshalJSON() ([]byte, error) { - val := json.RawMessage(strconv.FormatFloat(b.Value, 'f', -1, 64)) - // Convert math.NaN and math.Inf to quoted string. - if math.IsNaN(b.Value) || math.IsInf(b.Value, 0) { - val = json.RawMessage(strconv.Quote(string(val))) - } + val := marshalFloat64(b.Value) bucket := rawValueBucket{ TestAggBucket: (*TestAggBucket)(b), @@ -267,73 +260,11 @@ func (i *AsyncSearchesListItem) MarshalJSON() ([]byte, error) { return pbMarshaller.Marshal(i) } -// streamSearchDocColumns is the number of columns carried by a document record -// produced by StreamSearch: [id, time, data]. -const streamSearchDocColumns = 3 - -// streamSearchAggColumns is the number of columns carried by an aggregation -// record: [key, value, ts, quantiles, not_exists]. -const streamSearchAggColumns = 5 - -// MarshalJSON formats a StreamSearch record into a human-readable JSON array. -// The record kind is distinguished by the column count: -// -// - documents (3 cells): [id, time(nanoseconds, little-endian uint64), data] -// - aggregation buckets (5 cells): [key, value(little-endian float64), -// ts(nanoseconds, little-endian uint64), quantiles(little-endian float64 -// array), not_exists(little-endian uint64)] -// -// For any other layout the raw cells are emitted as-is (base64 for bytes). -func (r *Record) MarshalJSON() ([]byte, error) { - cells := r.GetRawData() - switch len(cells) { - case streamSearchDocColumns: - // cells[1] is a little-endian uint64 storing the document MID (nanoseconds). - var ts time.Time - if len(cells[1]) == 8 { - ts = seq.MID(encoding.Uint64FromBytes(cells[1])).Time() - } - return json.Marshal([]any{ - string(cells[0]), // id - ts.UTC().Format(time.RFC3339Nano), - json.RawMessage(cells[2]), // data, inlined as JSON - }) - case streamSearchAggColumns: - return marshalAggRecord(cells) - default: - return json.Marshal(cells) - } -} - -// marshalAggRecord renders an aggregation record: [key, value, ts, quantiles, not_exists]. -func marshalAggRecord(cells [][]byte) ([]byte, error) { - // cells[1] is a little-endian float64 storing the aggregation value. - var value float64 - if len(cells[1]) == 8 { - value = encoding.Float64FromBytes(cells[1]) - } - val := json.RawMessage(strconv.FormatFloat(value, 'f', -1, 64)) - if math.IsNaN(value) || math.IsInf(value, 0) { - val = json.RawMessage(strconv.Quote(string(val))) - } - // cells[2] is a little-endian uint64 storing the bucket MID (nanoseconds). - // A zero MID (no timestamp) is rendered as an empty string. - formattedTime := "" - if len(cells[2]) == 8 { - if ns := encoding.Uint64FromBytes(cells[2]); ns != 0 { - formattedTime = time.Unix(0, int64(ns)).UTC().Format(time.RFC3339Nano) - } - } - // cells[4] is a little-endian uint64 storing the not_exists count. - var notExists uint64 - if len(cells[4]) == 8 { - notExists = encoding.Uint64FromBytes(cells[4]) +// marshalFloat64 renders a float64 as a JSON number, quoting math.NaN and math.Inf as strings. +func marshalFloat64(v float64) json.RawMessage { + val := json.RawMessage(strconv.FormatFloat(v, 'f', -1, 64)) + if math.IsNaN(v) || math.IsInf(v, 0) { + return json.RawMessage(strconv.Quote(string(val))) } - return json.Marshal([]any{ - string(cells[0]), // key - val, // value - formattedTime, // ts - encoding.Float64ArrayFromBytes(cells[3]), // quantiles - notExists, // not_exists - }) + return val } diff --git a/pkg/seqproxyapi/v1/marshaler_test.go b/pkg/seqproxyapi/v1/marshaler_test.go index 13e1c8153..c79a638f5 100644 --- a/pkg/seqproxyapi/v1/marshaler_test.go +++ b/pkg/seqproxyapi/v1/marshaler_test.go @@ -9,8 +9,6 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/protobuf/types/known/durationpb" "google.golang.org/protobuf/types/known/timestamppb" - - "github.com/ozontech/seq-db/query/encoding" ) func TestAggregationBucketMarshalJSON(t *testing.T) { @@ -152,74 +150,9 @@ func TestFetchAsyncSearchResultResponseMarshalJSON(t *testing.T) { func TestRecordMarshalJSON(t *testing.T) { r := require.New(t) - t.Run("documents", func(t *testing.T) { - ns := time.Date(2025, 7, 8, 10, 19, 8, 742000000, time.UTC).UnixNano() - rec := &Record{RawData: [][]byte{ - []byte("46e48be997010000-e70163d0fa7582e4"), - encoding.Uint64ToBytes(uint64(ns)), - []byte(`{"message":"some_message","level":3}`), - }} - - raw, err := json.Marshal(rec) - r.NoError(err) - r.Equal( - `["46e48be997010000-e70163d0fa7582e4","2025-07-08T10:19:08.742Z",{"message":"some_message","level":3}]`, - string(raw), - ) - }) - - t.Run("aggregation buckets", func(t *testing.T) { - tsNs := time.Date(2025, 7, 8, 10, 19, 8, 742000000, time.UTC).UnixNano() - rec := &Record{RawData: [][]byte{ - []byte("service-a"), - encoding.Float64ToBytes(42.5), - encoding.Uint64ToBytes(uint64(tsNs)), - encoding.Float64ArrayToBytes(nil), // no quantiles - encoding.Uint64ToBytes(7), - }} - - raw, err := json.Marshal(rec) - r.NoError(err) - r.Equal(`["service-a",42.5,"2025-07-08T10:19:08.742Z",[],7]`, string(raw)) - }) - - t.Run("aggregation bucket with quantiles", func(t *testing.T) { - tsNs := time.Date(2025, 7, 8, 10, 19, 8, 742000000, time.UTC).UnixNano() - rec := &Record{RawData: [][]byte{ - []byte("service-a"), - encoding.Float64ToBytes(42.5), - encoding.Uint64ToBytes(uint64(tsNs)), - encoding.Float64ArrayToBytes([]float64{5.5, 9.1}), - encoding.Uint64ToBytes(0), - }} - - raw, err := json.Marshal(rec) - r.NoError(err) - r.Equal(`["service-a",42.5,"2025-07-08T10:19:08.742Z",[5.5,9.1],0]`, string(raw)) - }) - - t.Run("aggregation bucket with NaN", func(t *testing.T) { - rec := &Record{RawData: [][]byte{ - []byte("service-a"), - encoding.Float64ToBytes(math.NaN()), - make([]byte, 8), - encoding.Float64ArrayToBytes(nil), - make([]byte, 8), - }} - - raw, err := json.Marshal(rec) - r.NoError(err) - r.Equal(`["service-a","NaN","",[],0]`, string(raw)) - }) - - t.Run("unknown layout falls back to raw bytes", func(t *testing.T) { - // 4 cells used to be the aggregation layout; now it is an unknown one. - rec := &Record{RawData: [][]byte{[]byte("only"), []byte("two"), []byte("three"), []byte("four")}} - raw, err := json.Marshal(rec) - r.NoError(err) - // encoding/json marshals [][]byte as an array of base64 strings. - var decoded [][]byte - r.NoError(json.Unmarshal(raw, &decoded)) - r.Equal(rec.RawData, decoded) - }) + // Without a schema from the stream header, a record renders as raw cells with base64-encoded bytes. + rec := &Record{RawData: [][]byte{[]byte("only"), []byte("two")}} + raw, err := json.Marshal(rec) + r.NoError(err) + r.Equal(`{"raw_data":["b25seQ==","dHdv"]}`, string(raw)) } diff --git a/pkg/seqproxyapi/v1/stream_render.go b/pkg/seqproxyapi/v1/stream_render.go new file mode 100644 index 000000000..df0daed79 --- /dev/null +++ b/pkg/seqproxyapi/v1/stream_render.go @@ -0,0 +1,165 @@ +package seqproxyapi + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "time" + + "google.golang.org/protobuf/proto" + + "github.com/ozontech/seq-db/query/encoding" +) + +func NewStreamSearchRewriter() func(ctx context.Context, resp proto.Message) (any, error) { + var ( + mu sync.RWMutex + schemas = make(map[context.Context]*rowSchema) // key per stream + ) + + return func(ctx context.Context, resp proto.Message) (any, error) { + res, ok := resp.(*StreamSearchResponse) + if !ok { + return resp, nil + } + + switch { + case res.GetHeader() != nil: + mu.Lock() + schemas[ctx] = newRowSchema(res.GetHeader().GetTyping()) + mu.Unlock() + go func() { + // clean up stream aborted mid way + <-ctx.Done() + mu.Lock() + delete(schemas, ctx) + mu.Unlock() + }() + case res.GetData() != nil: + mu.RLock() + schema := schemas[ctx] + mu.RUnlock() + if schema != nil { + return &renderedData{schema: schema, batch: res.GetData().GetBatch()}, nil + } + case res.GetSummary() != nil: + // delete the entry for completed stream + mu.Lock() + delete(schemas, ctx) + mu.Unlock() + } + + return res, nil + } +} + +type renderedData struct { + schema *rowSchema + batch *RecordsBatch +} + +func (d *renderedData) MarshalJSON() ([]byte, error) { + records := make([]json.RawMessage, len(d.batch.GetRecords())) + for i, rec := range d.batch.GetRecords() { + raw, err := d.schema.renderRecord(rec.GetRawData()) + if err != nil { + // schema mismatch, return raw base64 instead of failing the whole stream. + if raw, err = json.Marshal(rec.GetRawData()); err != nil { + return nil, err + } + } + records[i] = raw + } + return json.Marshal(renderedDataJSON{ + Data: renderedDataBodyJSON{ + Batch: renderedBatchJSON{Records: records}, + }, + }) +} + +type renderedDataJSON struct { + Data renderedDataBodyJSON `json:"data,omitempty"` +} + +type renderedDataBodyJSON struct { + Batch renderedBatchJSON `json:"batch,omitempty"` +} + +type renderedBatchJSON struct { + Records []json.RawMessage `json:"records,omitempty"` +} + +// rowSchema renders raw data of StreamSearch records as human readable values according to the Typing from header. +type rowSchema []*Typing + +func newRowSchema(typing []*Typing) *rowSchema { + s := rowSchema(typing) + return &s +} + +func (s *rowSchema) renderRecord(cells [][]byte) (json.RawMessage, error) { + row, err := s.render(cells) + if err != nil { + return nil, err + } + return json.Marshal(row) +} + +func (s *rowSchema) render(cells [][]byte) ([]any, error) { + if len(cells) != len(*s) { + return nil, fmt.Errorf("record has %d cells, schema declares %d columns", len(cells), len(*s)) + } + row := make([]any, len(cells)) + for i, cell := range cells { + row[i] = (*s)[i].renderCell(cell) + } + return row, nil +} + +// timestampTitles holds uint64 column titles that carry time as nanoseconds since epoch; +// so timestamp rendering is keyed on the column title instead of the type alone. +var timestampTitles = map[string]struct{}{ + "time": {}, + "ts": {}, +} + +func (t *Typing) renderCell(cell []byte) any { + switch t.Type { + case DataType_DATA_TYPE_SEQ_ID, DataType_DATA_TYPE_STRING: + return string(cell) + case DataType_DATA_TYPE_BYTES: + return cell + case DataType_DATA_TYPE_RAW_DOCUMENT: + return json.RawMessage(cell) + case DataType_DATA_TYPE_UINT32: + return encoding.Uint32FromBytes(cell) + case DataType_DATA_TYPE_UINT64: + if _, ok := timestampTitles[t.Title]; ok { + return renderTimestamp(cell) + } + return encoding.Uint64FromBytes(cell) + case DataType_DATA_TYPE_INT32: + return encoding.Int32FromBytes(cell) + case DataType_DATA_TYPE_INT64: + return encoding.Int64FromBytes(cell) + case DataType_DATA_TYPE_FLOAT64: + return marshalFloat64(encoding.Float64FromBytes(cell)) + case DataType_DATA_TYPE_FLOAT64_ARRAY: + return encoding.Float64ArrayFromBytes(cell) + case DataType_DATA_TYPE_STRING_ARRAY: + return encoding.StringArrayFromBytes(cell) + default: + return cell + } +} + +func renderTimestamp(cell []byte) string { + if len(cell) != 8 { + return "" + } + if ns := encoding.Uint64FromBytes(cell); ns != 0 { + return time.Unix(0, int64(ns)).UTC().Format(time.RFC3339Nano) + } + return "" +} diff --git a/pkg/seqproxyapi/v1/stream_render_test.go b/pkg/seqproxyapi/v1/stream_render_test.go new file mode 100644 index 000000000..b4b403464 --- /dev/null +++ b/pkg/seqproxyapi/v1/stream_render_test.go @@ -0,0 +1,235 @@ +package seqproxyapi + +import ( + "context" + "encoding/json" + "math" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/ozontech/seq-db/query/encoding" +) + +func TestRowSchemaRender(t *testing.T) { + r := require.New(t) + + tsNs := time.Date(2025, 7, 8, 10, 19, 8, 742000000, time.UTC).UnixNano() + + test := func(typing []*Typing, cells [][]byte, expected string) { + t.Helper() + + schema := newRowSchema(typing) + raw, err := schema.renderRecord(cells) + r.NoError(err) + r.Equal(expected, string(raw)) + } + + t.Run("documents", func(t *testing.T) { + test( + []*Typing{ + {Title: "id", Type: DataType_DATA_TYPE_SEQ_ID}, + {Title: "time", Type: DataType_DATA_TYPE_UINT64}, + {Title: "data", Type: DataType_DATA_TYPE_RAW_DOCUMENT}, + }, + [][]byte{ + []byte("46e48be997010000-e70163d0fa7582e4"), + encoding.Uint64ToBytes(uint64(tsNs)), + []byte(`{"message":"some_message","level":3}`), + }, + `["46e48be997010000-e70163d0fa7582e4","2025-07-08T10:19:08.742Z",{"message":"some_message","level":3}]`, + ) + }) + + t.Run("aggregation bucket", func(t *testing.T) { + test( + []*Typing{ + {Title: "token", Type: DataType_DATA_TYPE_STRING}, + {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, + {Title: "ts", Type: DataType_DATA_TYPE_UINT64}, + {Title: "quantiles", Type: DataType_DATA_TYPE_FLOAT64_ARRAY}, + {Title: "not_exists", Type: DataType_DATA_TYPE_UINT64}, + }, + [][]byte{ + []byte("service-a"), + encoding.Float64ToBytes(42.5), + encoding.Uint64ToBytes(uint64(tsNs)), + encoding.Float64ArrayToBytes(nil), // no quantiles + encoding.Uint64ToBytes(7), + }, + `["service-a",42.5,"2025-07-08T10:19:08.742Z",[],7]`, + ) + }) + + t.Run("aggregation bucket with quantiles", func(t *testing.T) { + test( + []*Typing{ + {Title: "token", Type: DataType_DATA_TYPE_STRING}, + {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, + {Title: "ts", Type: DataType_DATA_TYPE_UINT64}, + {Title: "quantiles", Type: DataType_DATA_TYPE_FLOAT64_ARRAY}, + {Title: "not_exists", Type: DataType_DATA_TYPE_UINT64}, + }, + [][]byte{ + []byte("service-a"), + encoding.Float64ToBytes(42.5), + encoding.Uint64ToBytes(uint64(tsNs)), + encoding.Float64ArrayToBytes([]float64{5.5, 9.1}), + encoding.Uint64ToBytes(0), + }, + `["service-a",42.5,"2025-07-08T10:19:08.742Z",[5.5,9.1],0]`, + ) + }) + + t.Run("aggregation bucket with NaN and no ts", func(t *testing.T) { + test( + []*Typing{ + {Title: "token", Type: DataType_DATA_TYPE_STRING}, + {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, + {Title: "ts", Type: DataType_DATA_TYPE_UINT64}, + {Title: "quantiles", Type: DataType_DATA_TYPE_FLOAT64_ARRAY}, + {Title: "not_exists", Type: DataType_DATA_TYPE_UINT64}, + }, + [][]byte{ + []byte("service-a"), + encoding.Float64ToBytes(math.NaN()), + make([]byte, 8), + encoding.Float64ArrayToBytes(nil), + make([]byte, 8), + }, + `["service-a","NaN","",[],0]`, + ) + }) + + t.Run("count without ts renders zero timestamp as empty string", func(t *testing.T) { + // A flexible schema may omit ts, quantiles and not_exists entirely. + test( + []*Typing{ + {Title: "token", Type: DataType_DATA_TYPE_STRING}, + {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, + }, + [][]byte{ + []byte("service-a"), + encoding.Float64ToBytes(3), + }, + `["service-a",3]`, + ) + }) + + t.Run("uint64 not a timestamp renders as number", func(t *testing.T) { + test( + []*Typing{{Title: "total", Type: DataType_DATA_TYPE_UINT64}}, + [][]byte{encoding.Uint64ToBytes(42)}, + `[42]`, + ) + }) +} + +func TestRowSchemaRenderMismatch(t *testing.T) { + r := require.New(t) + + schema := newRowSchema([]*Typing{ + {Title: "token", Type: DataType_DATA_TYPE_STRING}, + {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, + }) + + _, err := schema.renderRecord([][]byte{[]byte("service-a")}) + r.Error(err) + + // The data message renders record as raw base64 cells instead of failing. + raw, err := json.Marshal(&renderedData{ + schema: schema, + batch: &RecordsBatch{Records: []*Record{{RawData: [][]byte{[]byte("service-a")}}}}, + }) + r.NoError(err) + r.Equal( + `{"data":{"batch":{"records":[["c2VydmljZS1h"]]}}}`, + string(raw), + ) +} + +func TestStreamSearchRewriter(t *testing.T) { + r := require.New(t) + + rewrite := NewStreamSearchRewriter() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // Messages of other RPCs pass through untouched. + unary := &StatusResponse{} + out, err := rewrite(ctx, unary) + r.NoError(err) + r.Equal(unary, out) + + // The header carries the schema and passes through unchanged. + header := &StreamSearchResponse{ResponseType: &StreamSearchResponse_Header{Header: &ResponseHeader{ + Typing: []*Typing{ + {Title: "token", Type: DataType_DATA_TYPE_STRING}, + {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, + }, + }}} + out, err = rewrite(ctx, header) + r.NoError(err) + r.Equal(header, out) + + // A data message of the same stream is rewritten into a rendered batch. + data := &StreamSearchResponse{ResponseType: &StreamSearchResponse_Data{Data: &ResponseData{ + Batch: &RecordsBatch{Records: []*Record{ + {RawData: [][]byte{[]byte("service-a"), encoding.Float64ToBytes(3)}}, + {RawData: [][]byte{[]byte("service-b"), encoding.Float64ToBytes(4.5)}}, + }}, + }}} + out, err = rewrite(ctx, data) + r.NoError(err) + + raw, err := json.Marshal(map[string]any{"result": out}) + r.NoError(err) + r.Equal( + `{"result":{"data":{"batch":{"records":[["service-a",3],["service-b",4.5]]}}}}`, + string(raw), + ) + + // The summary closes the stream and drops the captured schema. + summary := &StreamSearchResponse{ResponseType: &StreamSearchResponse_Summary{Summary: &ResponseSummary{}}} + out, err = rewrite(ctx, summary) + r.NoError(err) + r.Equal(summary, out) + + // Without a captured schema a data message passes through and renders its records as raw base64. + out, err = rewrite(ctx, data) + r.NoError(err) + raw, err = json.Marshal(map[string]any{"result": out}) + r.NoError(err) + r.Contains(string(raw), "c2VydmljZS1h") // == base64("service-a") +} + +func TestStreamSearchRewriterSchemaCleanup(t *testing.T) { + r := require.New(t) + + rewrite := NewStreamSearchRewriter() + ctx, cancel := context.WithCancel(context.Background()) + + _, err := rewrite(ctx, &StreamSearchResponse{ResponseType: &StreamSearchResponse_Header{ + Header: &ResponseHeader{Typing: []*Typing{ + {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, + }}, + }}) + r.NoError(err) + + // Aborting the stream (context cancellation) drops the captured schema, + // so a late data message is no longer rewritten. + cancel() + + r.Eventually(func() bool { + out, err := rewrite(ctx, &StreamSearchResponse{ResponseType: &StreamSearchResponse_Data{ + Data: &ResponseData{Batch: &RecordsBatch{Records: []*Record{}}}}, + }) + return err == nil && !isRenderedData(out) + }, time.Second, 10*time.Millisecond) +} + +func isRenderedData(v any) bool { + _, ok := v.(*renderedData) + return ok +} diff --git a/proxyapi/grpc_stream_search.go b/proxyapi/grpc_stream_search.go index 20c80cb2f..eeb8c724e 100644 --- a/proxyapi/grpc_stream_search.go +++ b/proxyapi/grpc_stream_search.go @@ -96,9 +96,8 @@ func (g *grpcV1) StreamSearch(stream seqproxyapi.SeqProxyApi_StreamSearchServer) var typing []*seqproxyapi.Typing var toRecord func(*query.Record) *seqproxyapi.Record if searchReq.Agg != nil { - typing = aggsTyping() - mapper := newAggRecordMapper(aggOutSchema(searchReq.Agg)) - toRecord = mapper.toRecord + typing = schemaToTyping(aggOutSchema(searchReq.Agg)) + toRecord = aggToRecord } else { typing = docsTyping() toRecord = docToRecord @@ -304,20 +303,20 @@ func sendRecords(stream seqproxyapi.SeqProxyApi_StreamSearchServer, records []*s func docsTyping() []*seqproxyapi.Typing { return []*seqproxyapi.Typing{ {Title: "id", Type: seqproxyapi.DataType_DATA_TYPE_SEQ_ID}, - {Title: "time", Type: seqproxyapi.DataType_DATA_TYPE_UINT64}, + {Title: "time", Type: seqproxyapi.DataType_DATA_TYPE_UINT64}, // extract time from id {Title: "data", Type: seqproxyapi.DataType_DATA_TYPE_RAW_DOCUMENT}, } } -// hardcoded schema -func aggsTyping() []*seqproxyapi.Typing { - return []*seqproxyapi.Typing{ - {Title: "key", Type: seqproxyapi.DataType_DATA_TYPE_STRING}, - {Title: "value", Type: seqproxyapi.DataType_DATA_TYPE_FLOAT64}, - {Title: "ts", Type: seqproxyapi.DataType_DATA_TYPE_UINT64}, - {Title: "quantiles", Type: seqproxyapi.DataType_DATA_TYPE_FLOAT64_ARRAY}, - {Title: "not_exists", Type: seqproxyapi.DataType_DATA_TYPE_UINT64}, +func schemaToTyping(out *query.Schema) []*seqproxyapi.Typing { + typing := make([]*seqproxyapi.Typing, 0, out.Len()) + for _, c := range out.Cols() { + typing = append(typing, &seqproxyapi.Typing{ + Title: c.Name, + Type: seqproxyapi.MustProtoDataType(c.Type), + }) } + return typing } var ( @@ -338,50 +337,16 @@ func docToRecord(r *query.Record) *seqproxyapi.Record { } } -type aggRecordMapper struct { - key query.Column[string] - value query.Column[float64] - ts *query.Column[uint64] - quantiles *query.Column[[]float64] - notExists *query.Column[uint64] -} - func aggOutSchema(agg *search.AggQuery) *query.Schema { return query.AggResultSchemaFor(agg.Func, agg.Interval > 0) } -func newAggRecordMapper(out *query.Schema) aggRecordMapper { - return aggRecordMapper{ - key: out.MustColumn[string]("token"), - value: out.MustColumn[float64]("value"), - ts: query.OptionalColumn[uint64](out, "ts"), - quantiles: query.OptionalColumn[[]float64](out, "quantiles"), - notExists: query.OptionalColumn[uint64](out, "not_exists"), - } -} - -func (m aggRecordMapper) toRecord(r *query.Record) *seqproxyapi.Record { - var ts uint64 - if m.ts != nil { - ts = m.ts.Val(r) - } - var quantiles []float64 - if m.quantiles != nil { - quantiles = m.quantiles.Val(r) - } - var notExists uint64 - if m.notExists != nil { - notExists = m.notExists.Val(r) - } - return &seqproxyapi.Record{ - RawData: [][]byte{ - m.key.RawData(r), // key - m.value.RawData(r), // value - encoding.Uint64ToBytes(ts), // ts - encoding.Float64ArrayToBytes(quantiles), // quantiles - encoding.Uint64ToBytes(notExists), // not_exists - }, +func aggToRecord(r *query.Record) *seqproxyapi.Record { + cells := make([][]byte, len(r.Vals)) + for i, v := range r.Vals { + cells[i] = v.RawData() } + return &seqproxyapi.Record{RawData: cells} } func buildSearchReq(q *seqproxyapi.StreamSearchQuery) (*search.StreamSearchRequest, error) { diff --git a/proxyapi/ingestor.go b/proxyapi/ingestor.go index 7ca8aec16..75bd518a8 100644 --- a/proxyapi/ingestor.go +++ b/proxyapi/ingestor.go @@ -115,6 +115,7 @@ func NewIngestor(config IngestorConfig, store *storeapiclient.Store) (*Ingestor, grpcGateway := runtime.NewServeMux( runtime.WithMarshalerOption(runtime.MIMEWildcard, &humanReadableMarshaler{}), + runtime.WithForwardResponseRewriter(seqproxyapi.NewStreamSearchRewriter()), ) err := seqproxyapi.RegisterSeqProxyApiHandlerFromEndpoint(ctx, grpcGateway, config.API.GatewayAddr, []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}) diff --git a/tests/integration_tests/integration_test.go b/tests/integration_tests/integration_test.go index 51bc29323..fa7fdcf15 100644 --- a/tests/integration_tests/integration_test.go +++ b/tests/integration_tests/integration_test.go @@ -2211,8 +2211,8 @@ func (s *IntegrationTestSuite) TestStreamSearch() { case *seqproxyapi.StreamSearchResponse_Data: for _, rec := range v.Data.GetBatch().GetRecords() { raw := rec.GetRawData() - // [key:STRING, value:FLOAT64, ts:UINT64, quantiles:FLOAT64_ARRAY, not_exists:UINT64] — proxy agg schema. - r.Len(raw, 5) + // [key:STRING, value:FLOAT64, ts:UINT64] + r.Len(raw, 3) value := math.Float64frombits(binary.LittleEndian.Uint64(raw[1])) ts := binary.LittleEndian.Uint64(raw[2]) got[ts] += value @@ -2262,11 +2262,11 @@ func (s *IntegrationTestSuite) TestStreamSearch() { case *seqproxyapi.StreamSearchResponse_Data: for _, rec := range v.Data.GetBatch().GetRecords() { raw := rec.GetRawData() - // [key:STRING, value:FLOAT64, ts:UINT64, quantiles:FLOAT64_ARRAY, not_exists:UINT64]. - r.Len(raw, 5) + // [key:STRING, value:FLOAT64, quantiles:FLOAT64_ARRAY, not_exists:UINT64]. + r.Len(raw, 4) gotRecords++ value = math.Float64frombits(binary.LittleEndian.Uint64(raw[1])) - quantiles = encoding.Float64ArrayFromBytes(raw[3]) + quantiles = encoding.Float64ArrayFromBytes(raw[2]) } case *seqproxyapi.StreamSearchResponse_Summary: r.Equal(1, gotRecords, "one bucket per `service` group") @@ -2310,7 +2310,7 @@ func (s *IntegrationTestSuite) TestStreamSearch() { case *seqproxyapi.StreamSearchResponse_Data: for _, rec := range v.Data.GetBatch().GetRecords() { raw := rec.GetRawData() - r.Len(raw, 5) + r.Len(raw, 3) gotRecords++ value = math.Float64frombits(binary.LittleEndian.Uint64(raw[1])) } From e24eb8320f435eac2956b4f46372147d127d3375 Mon Sep 17 00:00:00 2001 From: Daniil Forshev Date: Mon, 26 Oct 2026 20:17:20 +0500 Subject: [PATCH 2/3] refactor(stream search): NewStreamSearchRewriter - cleanup schemas lazily --- pkg/seqproxyapi/v1/stream_render.go | 13 +++++----- pkg/seqproxyapi/v1/stream_render_test.go | 31 ++++++++++++++---------- 2 files changed, 24 insertions(+), 20 deletions(-) diff --git a/pkg/seqproxyapi/v1/stream_render.go b/pkg/seqproxyapi/v1/stream_render.go index df0daed79..c79c60f0d 100644 --- a/pkg/seqproxyapi/v1/stream_render.go +++ b/pkg/seqproxyapi/v1/stream_render.go @@ -28,14 +28,13 @@ func NewStreamSearchRewriter() func(ctx context.Context, resp proto.Message) (an case res.GetHeader() != nil: mu.Lock() schemas[ctx] = newRowSchema(res.GetHeader().GetTyping()) + for c := range schemas { + // clean up streams aborted mid way + if c.Err() != nil { + delete(schemas, c) + } + } mu.Unlock() - go func() { - // clean up stream aborted mid way - <-ctx.Done() - mu.Lock() - delete(schemas, ctx) - mu.Unlock() - }() case res.GetData() != nil: mu.RLock() schema := schemas[ctx] diff --git a/pkg/seqproxyapi/v1/stream_render_test.go b/pkg/seqproxyapi/v1/stream_render_test.go index b4b403464..2c9aea3b2 100644 --- a/pkg/seqproxyapi/v1/stream_render_test.go +++ b/pkg/seqproxyapi/v1/stream_render_test.go @@ -207,26 +207,31 @@ func TestStreamSearchRewriter(t *testing.T) { func TestStreamSearchRewriterSchemaCleanup(t *testing.T) { r := require.New(t) + newHeader := func(_ context.Context) *StreamSearchResponse { + return &StreamSearchResponse{ResponseType: &StreamSearchResponse_Header{ + Header: &ResponseHeader{Typing: []*Typing{ + {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, + }}, + }} + } + rewrite := NewStreamSearchRewriter() ctx, cancel := context.WithCancel(context.Background()) - _, err := rewrite(ctx, &StreamSearchResponse{ResponseType: &StreamSearchResponse_Header{ - Header: &ResponseHeader{Typing: []*Typing{ - {Title: "value", Type: DataType_DATA_TYPE_FLOAT64}, - }}, - }}) + _, err := rewrite(ctx, newHeader(ctx)) r.NoError(err) - // Aborting the stream (context cancellation) drops the captured schema, - // so a late data message is no longer rewritten. + // Aborting the stream (context cancellation) drops the captured schema on + // the next stream's header, so a late data message is no longer rewritten. cancel() + _, err = rewrite(context.Background(), newHeader(context.Background())) + r.NoError(err) - r.Eventually(func() bool { - out, err := rewrite(ctx, &StreamSearchResponse{ResponseType: &StreamSearchResponse_Data{ - Data: &ResponseData{Batch: &RecordsBatch{Records: []*Record{}}}}, - }) - return err == nil && !isRenderedData(out) - }, time.Second, 10*time.Millisecond) + out, err := rewrite(ctx, &StreamSearchResponse{ResponseType: &StreamSearchResponse_Data{ + Data: &ResponseData{Batch: &RecordsBatch{Records: []*Record{}}}}, + }) + r.NoError(err) + r.False(isRenderedData(out)) } func isRenderedData(v any) bool { From c1236dd6d25e01d8cfb3d17b181bcb875e00d59a Mon Sep 17 00:00:00 2001 From: Daniil Forshev Date: Tue, 3 Nov 2026 18:31:17 +0500 Subject: [PATCH 3/3] fix(stream search): fix custom marshaling --- pkg/seqproxyapi/v1/stream_render.go | 14 ++++++++++---- pkg/seqproxyapi/v1/stream_render_test.go | 4 ++-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/pkg/seqproxyapi/v1/stream_render.go b/pkg/seqproxyapi/v1/stream_render.go index c79c60f0d..447dce798 100644 --- a/pkg/seqproxyapi/v1/stream_render.go +++ b/pkg/seqproxyapi/v1/stream_render.go @@ -70,15 +70,21 @@ func (d *renderedData) MarshalJSON() ([]byte, error) { } records[i] = raw } - return json.Marshal(renderedDataJSON{ - Data: renderedDataBodyJSON{ - Batch: renderedBatchJSON{Records: records}, + return json.Marshal(renderedResponseTypeJSON{ + ResponseType: renderedDataJSON{ + Data: renderedDataBodyJSON{ + Batch: renderedBatchJSON{Records: records}, + }, }, }) } +type renderedResponseTypeJSON struct { + ResponseType renderedDataJSON `json:"ResponseType,omitempty"` +} + type renderedDataJSON struct { - Data renderedDataBodyJSON `json:"data,omitempty"` + Data renderedDataBodyJSON `json:"Data,omitempty"` } type renderedDataBodyJSON struct { diff --git a/pkg/seqproxyapi/v1/stream_render_test.go b/pkg/seqproxyapi/v1/stream_render_test.go index 2c9aea3b2..0741b4f7c 100644 --- a/pkg/seqproxyapi/v1/stream_render_test.go +++ b/pkg/seqproxyapi/v1/stream_render_test.go @@ -144,7 +144,7 @@ func TestRowSchemaRenderMismatch(t *testing.T) { }) r.NoError(err) r.Equal( - `{"data":{"batch":{"records":[["c2VydmljZS1h"]]}}}`, + `{"ResponseType":{"Data":{"batch":{"records":[["c2VydmljZS1h"]]}}}}`, string(raw), ) } @@ -186,7 +186,7 @@ func TestStreamSearchRewriter(t *testing.T) { raw, err := json.Marshal(map[string]any{"result": out}) r.NoError(err) r.Equal( - `{"result":{"data":{"batch":{"records":[["service-a",3],["service-b",4.5]]}}}}`, + `{"result":{"ResponseType":{"Data":{"batch":{"records":[["service-a",3],["service-b",4.5]]}}}}}`, string(raw), )