loki/pkg/querier/queryrange/extensions.go

package queryrange

import "github.com/grafana/loki/pkg/querier/queryrange/queryrangebase"

// To satisfy queryrange.Response interface(https://github.com/cortexproject/cortex/blob/21bad57b346c730d684d6d0205efef133422ab28/pkg/querier/queryrange/query_range.go#L88)
// we need to have following method as well on response types:
// GetHeaders() []*queryrange.PrometheusResponseHeader.
// This could have been done by adding "Headers" field with custom type in proto definition for various Response types which doesn't work because
// gogoproto doesn't generate getters for custom types so adding them here.
// See issue https://github.com/gogo/protobuf/issues/477 for more details.
// It also has issue of generating slices without pointer to the custom type for Repeated customtype fields, see https://github.com/gogo/protobuf/issues/478
// which is why we also have to do conversion from non-pointer to pointer type.

func (m *LokiLabelNamesResponse) GetHeaders() []*queryrangebase.PrometheusResponseHeader {
	if m != nil {
		return convertPrometheusResponseHeadersToPointers(m.Headers)
	}
	return nil
}

func (m *LokiLabelNamesResponse) SetHeader(name, value string) {
	m.Headers = setHeader(m.Headers, name, value)
}

func (m *LokiLabelNamesResponse) WithHeaders(h []queryrangebase.PrometheusResponseHeader) queryrangebase.Response {
	m.Headers = h
	return m
}

func (m *LokiSeriesResponse) GetHeaders() []*queryrangebase.PrometheusResponseHeader {
	if m != nil {
		return convertPrometheusResponseHeadersToPointers(m.Headers)
	}
	return nil
}

func (m *LokiSeriesResponse) SetHeader(name, value string) {
	m.Headers = setHeader(m.Headers, name, value)
}

func (m *LokiSeriesResponse) WithHeaders(h []queryrangebase.PrometheusResponseHeader) queryrangebase.Response {
	m.Headers = h
	return m
}

func (m *LokiPromResponse) GetHeaders() []*queryrangebase.PrometheusResponseHeader {
	if m != nil {
		return m.Response.GetHeaders()
	}
	return nil
}

func (m *LokiPromResponse) SetHeader(name, value string) {
	m.Response.SetHeader(name, value)
}

func (m *LokiPromResponse) WithHeaders(h []queryrangebase.PrometheusResponseHeader) queryrangebase.Response {
	m.Response.Headers = convertPrometheusResponseHeadersToPointers(h)
	return m
}

func (m *LokiResponse) GetHeaders() []*queryrangebase.PrometheusResponseHeader {
	if m != nil {
		return convertPrometheusResponseHeadersToPointers(m.Headers)
	}
	return nil
}

func (m *LokiResponse) SetHeader(name, value string) {
	m.Headers = setHeader(m.Headers, name, value)
}

func (m *LokiResponse) WithHeaders(h []queryrangebase.PrometheusResponseHeader) queryrangebase.Response {
	m.Headers = h
	return m
}

func convertPrometheusResponseHeadersToPointers(h []queryrangebase.PrometheusResponseHeader) []*queryrangebase.PrometheusResponseHeader {
	if h == nil {
		return nil
	}

	resp := make([]*queryrangebase.PrometheusResponseHeader, len(h))
	for i := range h {
		resp[i] = &h[i]
	}

	return resp
}

// setHeader returns the passed headers with the new key-valur pair. Existing
// entries with the same key are overridden. The order is *not* maintained.
func setHeader(headers []queryrangebase.PrometheusResponseHeader, key, value string) []queryrangebase.PrometheusResponseHeader {
	for i, h := range headers {
		if h.Name == key {
			headers[i].Values = []string{value}
			return headers
		}
	}

	return append(headers, queryrangebase.PrometheusResponseHeader{Name: key, Values: []string{value}})
}

// GetHeaders returns the HTTP headers in the response.
func (m *IndexStatsResponse) GetHeaders() []*queryrangebase.PrometheusResponseHeader {
	if m != nil {
		return convertPrometheusResponseHeadersToPointers(m.Headers)
	}
	return nil
}

func (m *IndexStatsResponse) SetHeader(name, value string) {
	m.Headers = setHeader(m.Headers, name, value)
}

func (m *IndexStatsResponse) WithHeaders(h []queryrangebase.PrometheusResponseHeader) queryrangebase.Response {
	m.Headers = h
	return m
}

// GetHeaders returns the HTTP headers in the response.
func (m *VolumeResponse) GetHeaders() []*queryrangebase.PrometheusResponseHeader {
	if m != nil {
		return convertPrometheusResponseHeadersToPointers(m.Headers)
	}
	return nil
}

func (m *VolumeResponse) SetHeader(name, value string) {
	m.Headers = setHeader(m.Headers, name, value)
}

func (m *VolumeResponse) WithHeaders(h []queryrangebase.PrometheusResponseHeader) queryrangebase.Response {
	m.Headers = h
	return m
}

// GetHeaders returns the HTTP headers in the response.
func (m *TopKSketchesResponse) GetHeaders() []*queryrangebase.PrometheusResponseHeader {
	if m != nil {
		return convertPrometheusResponseHeadersToPointers(m.Headers)
	}
	return nil
}

func (m *TopKSketchesResponse) SetHeader(name, value string) {
	m.Headers = setHeader(m.Headers, name, value)
}

func (m *TopKSketchesResponse) WithHeaders(h []queryrangebase.PrometheusResponseHeader) queryrangebase.Response {
	m.Headers = h
	return m
}

// GetHeaders returns the HTTP headers in the response.
func (m *QuantileSketchResponse) GetHeaders() []*queryrangebase.PrometheusResponseHeader {
	if m != nil {
		return convertPrometheusResponseHeadersToPointers(m.Headers)
	}
	return nil
}

func (m *QuantileSketchResponse) SetHeader(name, value string) {
	m.Headers = setHeader(m.Headers, name, value)
}

func (m *QuantileSketchResponse) WithHeaders(h []queryrangebase.PrometheusResponseHeader) queryrangebase.Response {
	m.Headers = h
	return m
}