fleet-console: Update CountBrowserDevices to return actual data Bug: b/466971745 Change-Id: I75d012bb2b1b41681ab8bd7d531053ee2d682fbb Reviewed-on: https://chromium-review.googlesource.com/c/infra/infra/+/7485108 Reviewed-by: Pietro Scuttari <pietroscutta@google.com> Commit-Queue: Pietro Scuttari <pietroscutta@google.com> Auto-Submit: Vaghinak Vardanyan <vaghinak@google.com> Cr-Commit-Position: refs/heads/main@{#78480}
diff --git a/go/src/infra/fleetconsole/internal/consoleserver/count_browser_devices.go b/go/src/infra/fleetconsole/internal/consoleserver/count_browser_devices.go index c3b8ce6..6e582ce 100644 --- a/go/src/infra/fleetconsole/internal/consoleserver/count_browser_devices.go +++ b/go/src/infra/fleetconsole/internal/consoleserver/count_browser_devices.go
@@ -8,16 +8,15 @@ "context" "go.chromium.org/infra/fleetconsole/api/fleetconsolerpc" + browserdevicedb "go.chromium.org/infra/fleetconsole/internal/database/browser_devicedb" + "go.chromium.org/infra/fleetconsole/internal/ufsclient" ) func (frontend *FleetConsoleFrontend) CountBrowserDevices(ctx context.Context, req *fleetconsolerpc.CountBrowserDevicesRequest) (_ *fleetconsolerpc.CountBrowserDevicesResponse, err error) { - return &fleetconsolerpc.CountBrowserDevicesResponse{ - Total: 21345, - SwarmingState: &fleetconsolerpc.SwarmingStateCounts{ - Alive: 16930, - Dead: 417, - Quarantined: 372, - Maintenance: 18, - }, - }, nil + realms, err := ufsclient.GetUserRealms(ctx, frontend.cloudProject) + if err != nil { + return nil, err + } + + return browserdevicedb.CountDevices(ctx, req.GetFilter(), realms) }
diff --git a/go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices.go b/go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices.go new file mode 100644 index 0000000..f0c8cbe --- /dev/null +++ b/go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices.go
@@ -0,0 +1,81 @@ +// 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 +}
diff --git a/go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices_test.go b/go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices_test.go new file mode 100644 index 0000000..b03cbd3 --- /dev/null +++ b/go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices_test.go
@@ -0,0 +1,91 @@ +// 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" + "database/sql/driver" + "regexp" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + + "go.chromium.org/luci/common/testing/truth/assert" + "go.chromium.org/luci/common/testing/truth/should" + "go.chromium.org/luci/server/sqldb" + + "go.chromium.org/infra/fleetconsole/api/fleetconsolerpc" +) + +func TestCountDevicesSQL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + filter string + realms []string + expectedSQL string + expectedParameters []driver.Value + }{ + { + name: "no filter, no realms", + filter: "", + realms: nil, + expectedSQL: `SELECT COUNT(*) AS "total", ` + + `COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $1) AS "alive", ` + + `COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $2) AS "dead", ` + + `COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $3) AS "quarantined", ` + + `COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $4) AS "maintenance" ` + + `FROM "browser_devices"`, + expectedParameters: []driver.Value{`"alive"`, `"dead"`, `"quarantined"`, `"maintenance"`}, + }, + { + name: "with filter, with realms", + filter: "id = 'device1'", + realms: []string{"realm1"}, + expectedSQL: `SELECT COUNT(*) AS "total", ` + + `COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $1) AS "alive", ` + + `COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $2) AS "dead", ` + + `COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $3) AS "quarantined", ` + + `COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $4) AS "maintenance" ` + + `FROM "browser_devices" ` + + `WHERE "id" = $5 AND ("browser_devices"."realm" IS NULL OR "browser_devices"."realm" IN ($6))`, + expectedParameters: []driver.Value{`"alive"`, `"dead"`, `"quarantined"`, `"maintenance"`, "'device1'", "realm1"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("an error '%q' was not expected when opening a stub database connection", err) + } + defer db.Close() + ctx = sqldb.UseDB(ctx, db) + + mock.ExpectQuery(regexp.QuoteMeta(tt.expectedSQL)). + WithArgs(tt.expectedParameters...). + WillReturnRows(sqlmock.NewRows([]string{"total", "alive", "dead", "quarantined", "maintenance"}). + AddRow(10, 5, 2, 1, 2)) + + res, err := CountDevices(ctx, tt.filter, tt.realms) + assert.NoErr(t, err) + assert.Loosely(t, res, should.Match(&fleetconsolerpc.CountBrowserDevicesResponse{ + Total: 10, + SwarmingState: &fleetconsolerpc.SwarmingStateCounts{ + Alive: 5, + Dead: 2, + Quarantined: 1, + Maintenance: 2, + }, + })) + + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("there were unfulfilled expectations: %s", err) + } + }) + } +}
diff --git a/go/src/infra/fleetconsole/internal/ent/entx/builder.go b/go/src/infra/fleetconsole/internal/ent/entx/builder.go index 9b619005..b677632 100644 --- a/go/src/infra/fleetconsole/internal/ent/entx/builder.go +++ b/go/src/infra/fleetconsole/internal/ent/entx/builder.go
@@ -7,7 +7,10 @@ import ( "fmt" + "entgo.io/ent/dialect" "entgo.io/ent/dialect/sql" + + ent "go.chromium.org/infra/fleetconsole/internal/ent/generated" ) // True returns a standard SQL TRUE predicate. @@ -32,3 +35,31 @@ func Lateral(name string) *sql.SelectTable { return sql.Table(fmt.Sprintf("LATERAL %s", name)) } + +// CountIf generates COUNT(*) FILTER (WHERE <predicate>) or +// COUNT(CASE WHEN <predicate> THEN 1 END) depending on the dialect. +func CountIf(p *sql.Predicate) ent.AggregateFunc { + return func(s *sql.Selector) string { + // Sync the predicate's internal builder with the main query. + // We set Total() so the predicate generates placeholders ($5, $6) + // that continue from where the main query left off. + p.SetDialect(s.Dialect()) + p.SetTotal(s.Total()) + + // Compile the predicate to SQL and arguments. + condition, args := p.Query() + + // Inject the arguments into the main query. + // We use sql.Expr("", args...) to trick the builder into appending + // the arguments to its internal slice WITHOUT writing any new SQL + // text to the buffer (because the expr string is empty). + s.Arg(sql.Expr("", args...)) + + if s.Dialect() == dialect.Postgres { + // Return the standard Postgres aggregation SQL. + return fmt.Sprintf("COUNT(*) FILTER (WHERE %s)", condition) + } + + return fmt.Sprintf("COUNT(CASE WHEN %s THEN 1 END)", condition) + } +}