fleet-console: Extend ent so our code can look prettier Change-Id: I43e4f4c77845a7125414b3a5bd45f8d5bc26926a Reviewed-on: https://chromium-review.googlesource.com/c/infra/infra/+/7494081 Auto-Submit: Vaghinak Vardanyan <vaghinak@google.com> Reviewed-by: Pietro Scuttari <pietroscutta@google.com> Commit-Queue: Vaghinak Vardanyan <vaghinak@google.com> Reviewed-by: Konrad Polarczyk <polarczyk@google.com> Cr-Commit-Position: refs/heads/main@{#78483}

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
index f0c8cbe..b8f5b11 100644
--- a/go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices.go
+++ b/go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices.go

@@ -18,7 +18,6 @@
 	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"
 )
 
@@ -33,25 +32,17 @@
 // 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.
+	database.GetEntClient(ctx).BrowserDevice.
+		Query().
+		Where(p).
+		Where(entx.RealmFilter(browserdevice.FieldRealm, realms)).
 		Aggregate(
 			ent.As(ent.Count(), "total"),
 			ent.As(entx.CountIf(sqljson.ValueContains(browserdevice.FieldSwarmingLabels, string(swarmingclient.Alive), jsonPath)), string(swarmingclient.Alive)),

diff --git a/go/src/infra/fleetconsole/internal/database/browser_devicedb/list.go b/go/src/infra/fleetconsole/internal/database/browser_devicedb/list.go
index 38229b3..8211f4b 100644
--- a/go/src/infra/fleetconsole/internal/database/browser_devicedb/list.go
+++ b/go/src/infra/fleetconsole/internal/database/browser_devicedb/list.go

@@ -12,40 +12,33 @@
 	"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"
 	"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/utils"
 )
 
 // List lists the devices from the database.
 func List(ctx context.Context, filter, orderby string, offset, pageSize int, realms []string) ([]*fleetconsolerpc.BrowserDevice, bool, error) {
-	q := database.GetEntClient(ctx).BrowserDevice.Query()
-
 	p, err := queryutils.ToEntPredicate(filter, migrate.BrowserDevicesTable)
 	if err != nil {
 		return nil, false, errors.Fmt("failed to parse filter: %w", err)
 	}
-	q.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...))
-	}
 
 	o, err := queryutils.ToEntOrderBy(orderby, migrate.BrowserDevicesTable, browserdevice.FieldID)
 	if err != nil {
 		return nil, false, utils.InvalidOrderByError(err)
 	}
-	q.Order(o)
 
 	// Fetch one extra row to check whether there is more data available
-	q.Offset(offset).Limit(pageSize + 1)
+	entities, err := database.GetEntClient(ctx).BrowserDevice.
+		Query().
+		Where(p).
+		Where(entx.RealmFilter(browserdevice.FieldRealm, realms)).
+		Order(o).
+		Offset(offset).Limit(pageSize + 1).
+		All(ctx)
 
-	entities, err := q.All(ctx)
 	if err != nil {
 		return nil, false, errors.Fmt("failed to list devices: %w", err)
 	}

diff --git a/go/src/infra/fleetconsole/internal/database/browser_devicedb/total_count.go b/go/src/infra/fleetconsole/internal/database/browser_devicedb/total_count.go
index f420ca7..0bca0de 100644
--- a/go/src/infra/fleetconsole/internal/database/browser_devicedb/total_count.go
+++ b/go/src/infra/fleetconsole/internal/database/browser_devicedb/total_count.go

@@ -11,28 +11,22 @@
 
 	"go.chromium.org/infra/fleetconsole/internal/database"
 	"go.chromium.org/infra/fleetconsole/internal/database/queryutils"
+	"go.chromium.org/infra/fleetconsole/internal/ent/entx"
 	"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"
 )
 
 // TotalCount counts the devices in the database.
 func TotalCount(ctx context.Context, filter string, realms []string) (int, error) {
-	q := database.GetEntClient(ctx).BrowserDevice.Query()
 
 	p, err := queryutils.ToEntPredicate(filter, migrate.BrowserDevicesTable)
 	if err != nil {
 		return 0, errors.Fmt("failed to parse filter: %w", err)
 	}
-	q.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...))
-	}
-
-	return q.Count(ctx)
+	return database.GetEntClient(ctx).BrowserDevice.
+		Query().
+		Where(p).
+		Where(entx.RealmFilter(browserdevice.FieldRealm, realms)).
+		Count(ctx)
 }

diff --git a/go/src/infra/fleetconsole/internal/ent/entx/filter.go b/go/src/infra/fleetconsole/internal/ent/entx/filter.go
new file mode 100644
index 0000000..7e0078d
--- /dev/null
+++ b/go/src/infra/fleetconsole/internal/ent/entx/filter.go

@@ -0,0 +1,40 @@
+// 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 entx
+
+import (
+	"entgo.io/ent/dialect/sql"
+
+	"go.chromium.org/infra/fleetconsole/internal/utils"
+)
+
+// RealmFilter returns a generic predicate that works on ANY Ent query.
+// It assumes the underlying table has a column named "realm".
+func RealmFilter(col string, realms []string) func(*sql.Selector) {
+	return func(s *sql.Selector) {
+		// If nil, apply no filter
+		if realms == nil {
+			return
+		}
+
+		// s.C resolves to "table_name"."col_name" to avoid ambiguity
+		col := s.C(col)
+
+		nilCheck := sql.IsNull(col)
+		if len(realms) == 0 {
+			s.Where(nilCheck)
+			return
+		}
+
+		// Convert []string to []any for the sql builder
+		args := utils.Map(realms, func(realm string) any { return realm })
+
+		// (realm IS NULL) OR (realm IN (...))
+		s.Where(sql.Or(
+			nilCheck,
+			sql.In(col, args...),
+		))
+	}
+}