o11y/pkg/modules/inframonitoring/implinframonitoring/hosts.go

package implinframonitoring

import (
	"context"
	"fmt"
	"slices"
	"strings"

	"github.com/hanzoai/o11y/pkg/datastoresql"

	"github.com/hanzo-ds/sqlbuilder"
	"github.com/hanzoai/o11y/pkg/flagger"
	"github.com/hanzoai/o11y/pkg/telemetrymetrics"
	"github.com/hanzoai/o11y/pkg/types/featuretypes"
	"github.com/hanzoai/o11y/pkg/types/inframonitoringtypes"
	qbtypes "github.com/hanzoai/o11y/pkg/types/querybuildertypes/querybuildertypesv5"
	"github.com/hanzoai/o11y/pkg/valuer"
)

// getPerGroupHostStatusCounts computes the number of active and inactive hosts per group
// for the current page. It queries the timeseries table using the provided filter
// expression (which includes user filter + status filter + page groups IN clause).
// Uses GLOBAL IN with the active-hosts subquery inside uniqExactIf for active count,
// and a simple uniqExactIf for total count. Inactive = total - active (computed in Go).
func (m *module) getPerGroupHostStatusCounts(
	ctx context.Context,
	orgID valuer.UUID,
	req *inframonitoringtypes.PostableHosts,
	metricNames []string,
	pageGroups []map[string]string,
	sinceUnixMilli int64,
) (map[string]groupHostStatusCounts, error) {

	reductionEnabled := m.fl.BooleanOrEmpty(ctx, flagger.FeatureEnableMetricsReduction, featuretypes.NewFlaggerEvaluationContext(orgID))

	// Build the full filter expression from req (user filter + status filter) and page groups.
	reqFilterExpr := ""
	if req.Filter != nil {
		reqFilterExpr = req.Filter.Expression
	}
	pageGroupsFilterExpr := buildPageGroupsFilterExpr(pageGroups)
	filterExpr := mergeFilterExpressions(reqFilterExpr, pageGroupsFilterExpr)

	// Step-floor bounds + resolve tables in one shot to match QB v5 querier.
	samplesStartMs, flooredEndMs, tsAdjustedStartMs, distributedTimeSeriesTableName, _, _, localSamplesTable := alignedMetricWindow(req.Start, req.End)

	hostNameExpr := fmt.Sprintf("JSONExtractString(labels, '%s')", inframonitoringtypes.HostNameAttrKey)

	sb := sqlbuilder.NewSelectBuilder()
	selectCols := make([]string, 0, len(req.GroupBy)+2)
	for _, key := range req.GroupBy {
		selectCols = append(selectCols,
			fmt.Sprintf("JSONExtractString(labels, %s) AS %s", sb.Var(key.Name), quoteIdentifier(key.Name)),
		)
	}

	activeHostsSQ := m.getActiveHostsQuery(metricNames, inframonitoringtypes.HostNameAttrKey, sinceUnixMilli)
	selectCols = append(selectCols,
		fmt.Sprintf("uniqExactIf(%s, %s GLOBAL IN (%s)) AS active_host_count", hostNameExpr, hostNameExpr, sb.Var(activeHostsSQ)),
		fmt.Sprintf("uniqExactIf(%s, %s != '') AS total_host_count", hostNameExpr, hostNameExpr),
	)

	sb.Select(selectCols...)

	if reductionEnabled {
		var filterClause *sqlbuilder.WhereClause
		if filterExpr != "" {
			var err error
			filterClause, err = m.buildFilterClause(ctx, &qbtypes.Filter{Expression: filterExpr}, req.Start, req.End)
			if err != nil {
				return nil, err
			}
		}

		rawSrc := sqlbuilder.NewSelectBuilder()
		rawSrc.Select("labels")
		rawSrc.From(fmt.Sprintf("%s.%s", telemetrymetrics.DBName, distributedTimeSeriesTableName))
		rawSrc.Where(
			rawSrc.In("metric_name", sqlbuilder.List(metricNames)),
			rawSrc.GE("unix_milli", tsAdjustedStartMs),
			rawSrc.LE("unix_milli", flooredEndMs),
			fmt.Sprintf("fingerprint IN (%s)", rawSrc.Var(m.buildSamplesTblFingerprintSubQuery(metricNames, localSamplesTable, samplesStartMs, flooredEndMs))),
		)
		if filterClause != nil {
			rawSrc.AddWhereClause(sqlbuilder.CopyWhereClause(filterClause))
		}

		reducedSrc := sqlbuilder.NewSelectBuilder()
		reducedSrc.Select("labels")
		reducedSrc.From(fmt.Sprintf("%s.%s", telemetrymetrics.DBName, telemetrymetrics.SeriesReducedTableName))
		reducedSrc.Where(
			reducedSrc.In("metric_name", sqlbuilder.List(metricNames)),
			reducedSrc.GE("unix_milli", tsAdjustedStartMs),
			reducedSrc.LE("unix_milli", flooredEndMs),
			fmt.Sprintf("fingerprint IN (%s)", reducedSrc.Var(m.buildReducedSamplesTblFingerprintSubQuery(metricNames, samplesStartMs, flooredEndMs))),
		)
		if filterClause != nil {
			reducedSrc.AddWhereClause(sqlbuilder.CopyWhereClause(filterClause))
		}

		sb.From(sb.BuilderAs(sqlbuilder.UnionAll(rawSrc, reducedSrc), "series"))
	} else {

		fpSB := m.buildSamplesTblFingerprintSubQuery(metricNames, localSamplesTable, samplesStartMs, flooredEndMs)

		sb.From(fmt.Sprintf("%s.%s", telemetrymetrics.DBName, distributedTimeSeriesTableName))
		sb.Where(
			sb.In("metric_name", sqlbuilder.List(metricNames)),
			sb.GE("unix_milli", tsAdjustedStartMs),
			sb.LE("unix_milli", flooredEndMs),
			fmt.Sprintf("fingerprint IN (%s)", sb.Var(fpSB)),
		)

		if filterExpr != "" {
			filterClause, err := m.buildFilterClause(ctx, &qbtypes.Filter{Expression: filterExpr}, req.Start, req.End)
			if err != nil {
				return nil, err
			}
			if filterClause != nil {
				sb.AddWhereClause(filterClause)
			}
		}
	}

	// GROUP BY
	groupByAliases := make([]string, 0, len(req.GroupBy))
	for _, key := range req.GroupBy {
		groupByAliases = append(groupByAliases, quoteIdentifier(key.Name))
	}
	sb.GroupBy(groupByAliases...)

	query, args := sb.BuildWithFlavor(datastoresql.Flavor)

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

	result := make(map[string]groupHostStatusCounts)
	for rows.Next() {
		groupVals := make([]string, len(req.GroupBy))
		scanPtrs := make([]any, 0, len(req.GroupBy)+2)
		for i := range groupVals {
			scanPtrs = append(scanPtrs, &groupVals[i])
		}

		var activeCount, totalCount uint64
		scanPtrs = append(scanPtrs, &activeCount, &totalCount)

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

		compositeKey := compositeKeyFromList(groupVals)
		result[compositeKey] = groupHostStatusCounts{
			Active:   int(activeCount),
			Inactive: int(totalCount - activeCount),
		}
	}

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

	return result, nil
}

// buildHostRecords constructs the final list of HostRecords for a page.
// Groups that had no metric data get default values of -1.
//
// hostCounts is nil when host.name is in the groupBy — in that case, counts are
// derived directly from activeHostsMap (1/0 per host). When non-nil (custom groupBy
// without host.name), counts are looked up from the map.
func buildHostRecords(
	isHostNameInGroupBy bool,
	resp *qbtypes.QueryRangeResponse,
	pageGroups []map[string]string,
	groupBy []qbtypes.GroupByKey,
	metadataMap map[string]map[string]string,
	activeHostsMap map[string]bool,
	hostCounts map[string]groupHostStatusCounts,
) []inframonitoringtypes.HostRecord {
	metricsMap := parseFullQueryResponse(resp, groupBy)

	records := make([]inframonitoringtypes.HostRecord, 0, len(pageGroups))
	for _, labels := range pageGroups {
		compositeKey := compositeKeyFromLabels(labels, groupBy)
		hostName := labels[inframonitoringtypes.HostNameAttrKey]

		activeStatus := inframonitoringtypes.HostStatusNone
		activeHostCount := 0
		inactiveHostCount := 0
		if isHostNameInGroupBy { // derive from activeHostsMap since each row = one host
			if hostName != "" {
				if activeHostsMap[hostName] {
					activeStatus = inframonitoringtypes.HostStatusActive
					activeHostCount = 1
				} else {
					activeStatus = inframonitoringtypes.HostStatusInactive
					inactiveHostCount = 1
				}
			}
		} else { // derive from hostCounts since custom groupBy without host.name
			if counts, ok := hostCounts[compositeKey]; ok {
				activeHostCount = counts.Active
				inactiveHostCount = counts.Inactive
			}
		}

		record := inframonitoringtypes.HostRecord{
			HostName:          hostName,
			Status:            activeStatus,
			ActiveHostCount:   activeHostCount,
			InactiveHostCount: inactiveHostCount,
			CPU:               -1,
			Memory:            -1,
			Wait:              -1,
			Load15:            -1,
			DiskUsage:         -1,
			Meta:              map[string]string{},
		}

		if metrics, ok := metricsMap[compositeKey]; ok {
			if v, exists := metrics["F1"]; exists {
				record.CPU = v
			}
			if v, exists := metrics["F2"]; exists {
				record.Memory = v
			}
			if v, exists := metrics["F3"]; exists {
				record.Wait = v
			}
			if v, exists := metrics["F4"]; exists {
				record.DiskUsage = v
			}
			if v, exists := metrics["G"]; exists {
				record.Load15 = v
			}
		}

		if attrs, ok := metadataMap[compositeKey]; ok {
			for k, v := range attrs {
				record.Meta[k] = v
			}
		}

		records = append(records, record)
	}
	return records
}

// getTopHostGroups runs a ranking query for the ordering metric, sorts the
// results, paginates, and backfills from metadataMap when the page extends
// past the metric-ranked groups.
func (m *module) getTopHostGroups(
	ctx context.Context,
	orgID valuer.UUID,
	req *inframonitoringtypes.PostableHosts,
	metadataMap map[string]map[string]string,
) ([]map[string]string, error) {
	orderByKey := req.OrderBy.Key.Name
	if orderByKey == inframonitoringtypes.HostNameAttrKey {
		return inframonitoringtypes.PaginateMetadataByName(metadataMap, req.GroupBy, req.OrderBy.Direction, req.Offset, req.Limit, inframonitoringtypes.HostNameAttrKey), nil
	}
	queryNamesForOrderBy := orderByToHostsQueryNames[orderByKey]
	// The last entry is the formula/query whose value we sort by.
	rankingQueryName := queryNamesForOrderBy[len(queryNamesForOrderBy)-1]

	topReq := &qbtypes.QueryRangeRequest{
		Start:       uint64(req.Start),
		End:         uint64(req.End),
		RequestType: qbtypes.RequestTypeScalar,
		CompositeQuery: qbtypes.CompositeQuery{
			Queries: make([]qbtypes.QueryEnvelope, 0, len(queryNamesForOrderBy)),
		},
	}

	for _, envelope := range m.newListHostsQuery().CompositeQuery.Queries {
		if !slices.Contains(queryNamesForOrderBy, envelope.GetQueryName()) {
			continue
		}
		copied := envelope
		if copied.Type == qbtypes.QueryTypeBuilder {
			existingExpr := ""
			if f := copied.GetFilter(); f != nil {
				existingExpr = f.Expression
			}
			reqFilterExpr := ""
			if req.Filter != nil {
				reqFilterExpr = req.Filter.Expression
			}
			merged := mergeFilterExpressions(existingExpr, reqFilterExpr)
			copied.SetFilter(&qbtypes.Filter{Expression: merged})
			copied.SetGroupBy(req.GroupBy)
		}
		topReq.CompositeQuery.Queries = append(topReq.CompositeQuery.Queries, copied)
	}

	resp, err := m.querier.QueryRange(ctx, orgID, topReq)
	if err != nil {
		return nil, err
	}

	allMetricGroups := parseAndSortGroups(resp, rankingQueryName, req.GroupBy, req.OrderBy.Direction)
	return paginateWithBackfill(allMetricGroups, metadataMap, req.GroupBy, req.Offset, req.Limit), nil
}

// applyHostsActiveStatusFilter MODIFIES req.Filter.Expression to include an IN/NOT IN
// clause based on FilterByStatus and the set of active hosts.
// Returns true if the caller should short-circuit with an empty result (eg. ACTIVE
// requested but no hosts are active).
func (m *module) applyHostsActiveStatusFilter(req *inframonitoringtypes.PostableHosts, activeHostsMap map[string]bool) (shouldShortCircuit bool) {
	if req.Filter == nil || (req.Filter.FilterByStatus != inframonitoringtypes.HostStatusActive && req.Filter.FilterByStatus != inframonitoringtypes.HostStatusInactive) {
		return false
	}

	activeHosts := make([]string, 0, len(activeHostsMap))
	for host := range activeHostsMap {
		activeHosts = append(activeHosts, fmt.Sprintf("'%s'", host))
	}

	if len(activeHosts) == 0 {
		return req.Filter.FilterByStatus == inframonitoringtypes.HostStatusActive
	}

	op := "IN"
	if req.Filter.FilterByStatus == inframonitoringtypes.HostStatusInactive {
		op = "NOT IN"
	}
	statusClause := fmt.Sprintf("%s %s (%s)", inframonitoringtypes.HostNameAttrKey, op, strings.Join(activeHosts, ", "))
	req.Filter.Expression = mergeFilterExpressions(req.Filter.Expression, statusClause)
	return false
}

func (m *module) getHostsTableMetadata(ctx context.Context, orgID valuer.UUID, req *inframonitoringtypes.PostableHosts) (map[string]map[string]string, error) {
	var nonGroupByAttrs []string
	for _, key := range hostAttrKeysForMetadata {
		if !isKeyInGroupByAttrs(req.GroupBy, key) {
			nonGroupByAttrs = append(nonGroupByAttrs, key)
		}
	}
	var filter *qbtypes.Filter
	if req.Filter != nil {
		filter = &req.Filter.Filter
	}
	metadataMap, err := m.getMetadata(ctx, orgID, hostsTableMetricNamesList, req.GroupBy, nonGroupByAttrs, filter, req.Start, req.End)
	if err != nil {
		return nil, err
	}
	return metadataMap, nil
}

// getActiveHostsQuery builds a SelectBuilder that returns distinct host names
// with metrics reported in the last 10 minutes. The builder is not executed —
// callers can either execute it (getActiveHosts) or embed it as a subquery
// (getPerGroupActiveInactiveHostCounts).
func (m *module) getActiveHostsQuery(metricNames []string, hostNameAttr string, sinceUnixMilli int64) *sqlbuilder.SelectBuilder {
	sb := sqlbuilder.NewSelectBuilder()
	sb.Distinct()
	sb.Select("attr_string_value")
	sb.From(fmt.Sprintf("%s.%s", telemetrymetrics.DBName, telemetrymetrics.AttributeTableName))
	sb.Where(
		sb.In("metric_name", sqlbuilder.List(metricNames)),
		sb.E("attr_name", hostNameAttr),
		sb.NE("attr_string_value", ""),
		sb.GE("last_reported_unix_milli", sinceUnixMilli),
	)

	return sb
}

// getActiveHosts returns a set of host names that have reported metrics recently.
// It queries distributed_metadata for hosts where last_reported_unix_milli >= sinceUnixMilli.
func (m *module) getActiveHosts(ctx context.Context, metricNames []string, hostNameAttr string, sinceUnixMilli int64) (map[string]bool, error) {
	sb := m.getActiveHostsQuery(metricNames, hostNameAttr, sinceUnixMilli)
	query, args := sb.BuildWithFlavor(datastoresql.Flavor)

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

	activeHosts := make(map[string]bool)
	for rows.Next() {
		var hostName string
		if err := rows.Scan(&hostName); err != nil {
			return nil, err
		}
		if hostName != "" {
			activeHosts[hostName] = true
		}
	}

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

	return activeHosts, nil
}