go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices.go

// Copyright 2026 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package browserdevicedb

import (
	"context"

	"entgo.io/ent/dialect/sql/sqljson"

	"go.chromium.org/luci/common/errors"

	"go.chromium.org/infra/fleetconsole/api/fleetconsolerpc"
	"go.chromium.org/infra/fleetconsole/internal/database"
	"go.chromium.org/infra/fleetconsole/internal/database/queryutils"
	"go.chromium.org/infra/fleetconsole/internal/ent/entx"
	ent "go.chromium.org/infra/fleetconsole/internal/ent/generated"
	"go.chromium.org/infra/fleetconsole/internal/ent/generated/browserdevice"
	"go.chromium.org/infra/fleetconsole/internal/ent/generated/migrate"
	"go.chromium.org/infra/fleetconsole/internal/ent/generated/predicate"
	"go.chromium.org/infra/fleetconsole/internal/swarmingclient"
)

type DeviceCounts struct {
	Total       int32 `sql:"total"`
	Alive       int32 `sql:"alive"`
	Dead        int32 `sql:"dead"`
	Quarantined int32 `sql:"quarantined"`
	Maintenance int32 `sql:"maintenance"`
}

// CountDevices calculates browser devices related metrics.
func CountDevices(ctx context.Context, filter string, realms []string) (*fleetconsolerpc.CountBrowserDevicesResponse, error) {
	var rows []DeviceCounts
	client := database.GetEntClient(ctx)

	p, err := queryutils.ToEntPredicate(filter, migrate.BrowserDevicesTable)
	if err != nil {
		return nil, errors.Fmt("failed to parse filter: %w", err)
	}

	q := client.BrowserDevice.Query().Where(p)

	if realms != nil {
		realmChecks := []predicate.BrowserDevice{browserdevice.RealmIsNil()}
		if len(realms) > 0 {
			realmChecks = append(realmChecks, browserdevice.RealmIn(realms...))
		}
		q.Where(browserdevice.Or(realmChecks...))
	}

	jsonPath := sqljson.Path("state")
	err = q.
		Aggregate(
			ent.As(ent.Count(), "total"),
			ent.As(entx.CountIf(sqljson.ValueContains(browserdevice.FieldSwarmingLabels, string(swarmingclient.Alive), jsonPath)), string(swarmingclient.Alive)),
			ent.As(entx.CountIf(sqljson.ValueContains(browserdevice.FieldSwarmingLabels, string(swarmingclient.Dead), jsonPath)), string(swarmingclient.Dead)),
			ent.As(entx.CountIf(sqljson.ValueContains(browserdevice.FieldSwarmingLabels, string(swarmingclient.Quarantined), jsonPath)), string(swarmingclient.Quarantined)),
			ent.As(entx.CountIf(sqljson.ValueContains(browserdevice.FieldSwarmingLabels, string(swarmingclient.Maintenance), jsonPath)), string(swarmingclient.Maintenance))).
		Scan(ctx, &rows)

	if err != nil {
		return nil, err
	}

	if len(rows) == 0 {
		return &fleetconsolerpc.CountBrowserDevicesResponse{}, nil
	}

	result := rows[0]
	return &fleetconsolerpc.CountBrowserDevicesResponse{
		Total: result.Total,
		SwarmingState: &fleetconsolerpc.SwarmingStateCounts{
			Alive:       result.Alive,
			Dead:        result.Dead,
			Quarantined: result.Quarantined,
			Maintenance: result.Maintenance,
		},
	}, nil
}