o11y/pkg/prometheus/datastoreprometheus/client.go

package datastoreprometheus

import (
	"context"
	"fmt"
	"math"
	"strconv"
	"strings"
	"sync"

	"github.com/hanzoai/o11y/pkg/errors"
	"github.com/hanzoai/o11y/pkg/factory"
	"github.com/hanzoai/o11y/pkg/query-service/constants"
	"github.com/hanzoai/o11y/pkg/telemetrystore"
	"github.com/hanzoai/o11y/pkg/types/ctxtypes"
	"github.com/hanzoai/o11y/pkg/types/instrumentationtypes"
	"github.com/hanzoai/o11y/pkg/types/telemetrytypes"
	promValue "github.com/prometheus/prometheus/model/value"
	"github.com/prometheus/prometheus/prompb"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/storage/remote"
)

type client struct {
	settings       factory.ScopedProviderSettings
	telemetryStore telemetrystore.TelemetryStore
}

func NewReadClient(settings factory.ScopedProviderSettings, telemetryStore telemetrystore.TelemetryStore) remote.ReadClient {
	return &client{
		settings:       settings,
		telemetryStore: telemetryStore,
	}
}

func (client *client) Read(ctx context.Context, query *prompb.Query, sortSeries bool) (storage.SeriesSet, error) {
	if len(query.Matchers) == 2 {
		var hasJob bool
		var queryString string
		for _, m := range query.Matchers {
			if m.Type == prompb.LabelMatcher_EQ && m.Name == "job" && m.Value == "rawsql" {
				hasJob = true
			}
			if m.Type == prompb.LabelMatcher_EQ && m.Name == "query" {
				queryString = m.Value
			}
		}

		if hasJob && queryString != "" {
			res, err := client.queryRaw(ctx, queryString, int64(query.EndTimestampMs))
			if err != nil {
				return nil, err
			}

			return remote.FromQueryResult(sortSeries, res), nil
		}
	}

	var metricName string
	for _, matcher := range query.Matchers {
		if matcher.Name == "__name__" {
			metricName = matcher.Value
		}
	}

	datastoreQuery, args, err := client.queryToDatastoreQuery(ctx, query, metricName, false)
	if err != nil {
		return nil, err
	}

	fingerprints, err := client.getFingerprintsFromDatastoreQuery(ctx, datastoreQuery, args)
	if err != nil {
		return nil, err
	}
	if len(fingerprints) == 0 {
		return remote.FromQueryResult(sortSeries, new(prompb.QueryResult)), nil
	}

	datastoreSubQuery, args, err := client.queryToDatastoreQuery(ctx, query, metricName, true)
	if err != nil {
		return nil, err
	}

	res := new(prompb.QueryResult)
	timeseries, err := client.querySamples(ctx, int64(query.StartTimestampMs), int64(query.EndTimestampMs), fingerprints, metricName, datastoreSubQuery, args)
	if err != nil {
		return nil, err
	}
	res.Timeseries = timeseries

	return remote.FromQueryResult(sortSeries, res), nil
}

func (c *client) ReadMultiple(ctx context.Context, queries []*prompb.Query, sortSeries bool) (storage.SeriesSet, error) {
	if len(queries) == 0 {
		return storage.EmptySeriesSet(), nil
	}
	if len(queries) == 1 {
		return c.Read(ctx, queries[0], sortSeries)
	}

	type result struct {
		ss  storage.SeriesSet
		err error
	}

	results := make([]result, len(queries))
	var wg sync.WaitGroup
	wg.Add(len(queries))
	for i, q := range queries {
		go func(i int, q *prompb.Query) {
			defer wg.Done()
			ss, err := c.Read(ctx, q, sortSeries)
			results[i] = result{ss, err}
		}(i, q)
	}
	wg.Wait()

	sets := make([]storage.SeriesSet, 0, len(queries))
	for _, r := range results {
		if r.err != nil {
			return nil, r.err
		}
		sets = append(sets, r.ss)
	}
	return storage.NewMergeSeriesSet(sets, 0, storage.ChainedSeriesMerge), nil
}

func (client *client) queryToDatastoreQuery(_ context.Context, query *prompb.Query, metricName string, subQuery bool) (string, []any, error) {
	var datastoreQuery string
	var conditions []string
	var argCount = 0
	var selectString = "fingerprint, any(labels)"
	if subQuery {
		argCount = 1
		selectString = "fingerprint"
	}

	start, end, tableName := getStartAndEndAndTableName(query.StartTimestampMs, query.EndTimestampMs)

	var args []any
	conditions = append(conditions, fmt.Sprintf("metric_name = $%d", argCount+1))
	conditions = append(conditions, "temporality IN ['Cumulative', 'Unspecified']")
	conditions = append(conditions, fmt.Sprintf("unix_milli >= %d AND unix_milli < %d", start, end))

	normalized := !constants.IsDotMetricsEnabled

	conditions = append(conditions, fmt.Sprintf("__normalized = %v", normalized))

	args = append(args, metricName)
	for _, m := range query.Matchers {
		switch m.Type {
		case prompb.LabelMatcher_EQ:
			conditions = append(conditions, fmt.Sprintf("JSONExtractString(labels, $%d) = $%d", argCount+2, argCount+3))
		case prompb.LabelMatcher_NEQ:
			conditions = append(conditions, fmt.Sprintf("JSONExtractString(labels, $%d) != $%d", argCount+2, argCount+3))
		case prompb.LabelMatcher_RE:
			conditions = append(conditions, fmt.Sprintf("match(JSONExtractString(labels, $%d), $%d)", argCount+2, argCount+3))
		case prompb.LabelMatcher_NRE:
			conditions = append(conditions, fmt.Sprintf("not match(JSONExtractString(labels, $%d), $%d)", argCount+2, argCount+3))
		default:
			return "", nil, errors.NewInvalidInputf(errors.CodeInvalidInput, "unsupported or invalid matcher type: %s", m.Type.String())
		}
		args = append(args, m.Name, m.Value)
		argCount += 2
	}

	whereClause := strings.Join(conditions, " AND ")

	datastoreQuery = fmt.Sprintf(`SELECT %s FROM %s.%s WHERE %s GROUP BY fingerprint`, selectString, databaseName, tableName, whereClause)

	return datastoreQuery, args, nil
}

func (client *client) getFingerprintsFromDatastoreQuery(ctx context.Context, query string, args []any) (map[uint64][]prompb.Label, error) {
	ctx = client.withDatastorePrometheusContext(ctx, "getFingerprintsFromDatastoreQuery")
	rows, err := client.telemetryStore.Datastore().Query(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	fingerprints := make(map[uint64][]prompb.Label)

	var fingerprint uint64
	var labelString string
	for rows.Next() {
		if err = rows.Scan(&fingerprint, &labelString); err != nil {
			return nil, err
		}

		labels, _, err := unmarshalLabels(labelString, fingerprint)
		if err != nil {
			return nil, err
		}

		fingerprints[fingerprint] = labels
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	return fingerprints, nil
}

// buildSamplesQuery renders the samples SQL (and args) that fetches data
// points for the series selected by subQuery.
func buildSamplesQuery(start int64, end int64, metricName string, subQuery string, args []any) (string, []any) {
	argCount := len(args)

	query := fmt.Sprintf(`
		SELECT metric_name, fingerprint, unix_milli, value, flags
			FROM %s.%s
			WHERE metric_name = $1 AND fingerprint GLOBAL IN (%s) AND unix_milli >= $%s AND unix_milli <= $%s ORDER BY fingerprint, unix_milli;`,
		databaseName, samplesTableName, subQuery, strconv.Itoa(argCount+2), strconv.Itoa(argCount+3))
	query = strings.TrimSpace(query)

	allArgs := append([]any{metricName}, args...)
	allArgs = append(allArgs, start, end)
	return query, allArgs
}

func (client *client) querySamples(ctx context.Context, start int64, end int64, fingerprints map[uint64][]prompb.Label, metricName string, subQuery string, args []any) ([]*prompb.TimeSeries, error) {
	ctx = client.withDatastorePrometheusContext(ctx, "querySamples")

	query, allArgs := buildSamplesQuery(start, end, metricName, subQuery, args)

	rows, err := client.telemetryStore.Datastore().Query(ctx, query, allArgs...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var res []*prompb.TimeSeries
	var ts *prompb.TimeSeries
	var fingerprint, prevFingerprint uint64
	var timestampMs, prevTimestamp int64
	var value float64
	var flags uint32

	prevTimestamp = math.MinInt64

	for rows.Next() {
		if err := rows.Scan(&metricName, &fingerprint, &timestampMs, &value, &flags); err != nil {
			return nil, err
		}

		// collect samples in time series
		if fingerprint != prevFingerprint {
			// add collected time series to result
			prevFingerprint = fingerprint
			if ts != nil {
				res = append(res, ts)
			}

			labels := fingerprints[fingerprint]
			ts = &prompb.TimeSeries{
				Labels: labels,
			}
			prevTimestamp = math.MinInt64
		}

		if flags&1 == 1 {
			value = math.Float64frombits(promValue.StaleNaN)
		}

		if timestampMs == prevTimestamp {
			continue
		}
		prevTimestamp = timestampMs

		// add samples to current time series
		ts.Samples = append(ts.Samples, prompb.Sample{
			Timestamp: timestampMs,
			Value:     value,
		})
	}

	// add last time series
	if ts != nil {
		res = append(res, ts)
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	return res, nil
}

func (client *client) queryRaw(ctx context.Context, query string, ts int64) (*prompb.QueryResult, error) {
	ctx = client.withDatastorePrometheusContext(ctx, "queryRaw")

	rows, err := client.telemetryStore.Datastore().Query(ctx, query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	columns := rows.Columns()
	var res prompb.QueryResult
	targets := make([]any, len(columns))
	for i := range targets {
		targets[i] = new(scanner)
	}

	for rows.Next() {
		if err = rows.Scan(targets...); err != nil {
			return nil, err
		}

		labels := make([]prompb.Label, 0, len(columns))
		var value float64
		for i, c := range columns {
			v := targets[i].(*scanner)
			switch c {
			case "value":
				value = v.f
			default:
				labels = append(labels, prompb.Label{
					Name:  c,
					Value: v.s,
				})
			}
		}

		res.Timeseries = append(res.Timeseries, &prompb.TimeSeries{
			Labels: labels,
			Samples: []prompb.Sample{{
				Value:     value,
				Timestamp: ts,
			}},
		})
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return &res, nil
}

func (client *client) withDatastorePrometheusContext(ctx context.Context, functionName string) context.Context {
	comments := map[string]string{
		instrumentationtypes.TelemetrySignal:  telemetrytypes.SignalMetrics.StringValue(),
		instrumentationtypes.CodeNamespace:    "datastore-prometheus",
		instrumentationtypes.CodeFunctionName: functionName,
	}
	return ctxtypes.NewContextWithCommentVals(ctx, comments)
}