Cache policy counts (#15244)

changes/13574-cache-policy-results

@@ -0,0 +1,2 @@
+- policy results are now cached in mysql for faster sort operations on policy counts.  counts are
+updated by the cleanups_then_aggregation cron job 1X per hour by default.

cmd/fleet/cron.go

@@ -780,6 +780,12 @@ func newCleanupsAndAggregationSchedule(
 				return ds.UpdateQueryAggregatedStats(ctx)
 			},
 		),
+		schedule.WithJob(
+			"policy_aggregated_stats",
+			func(ctx context.Context) error {
+				return ds.UpdateHostPolicyCounts(ctx)
+			},
+		),
 		schedule.WithJob(
 			"aggregated_munki_and_mdm",
 			func(ctx context.Context) error {

server/datastore/mysql/migrations/tables/20231121054530_CreatePolicyStatsTable.go

@@ -0,0 +1,45 @@
+package tables
+
+import (
+	"database/sql"
+	"fmt"
+)
+
+func init() {
+	MigrationClient.AddMigration(Up_20231121054530, Down_20231121054530)
+}
+
+func Up_20231121054530(tx *sql.Tx) error {
+	stmt := `
+	CREATE TABLE policy_stats (
+		id int(10) unsigned NOT NULL AUTO_INCREMENT PRIMARY KEY,
+		policy_id int(10) unsigned NOT NULL,
+		-- inherited_team_id is used to indicate the row contains inherited 
+		-- global policies counts for a team, otherwise it will be 0.  This allows us 
+		-- to use the UNIQUE KEY constraint with this column to avoid duplicate rows
+		-- when policies.team_id is null. 
+		-- A foreign key constraint is not used here because team 0 is not a valid team
+		inherited_team_id int(10) unsigned NOT NULL DEFAULT 0,
+		-- cached counts for the policy / team
+		passing_host_count MEDIUMINT UNSIGNED NOT NULL DEFAULT 0,
+		failing_host_count MEDIUMINT UNSIGNED NOT NULL DEFAULT 0,
+		
+		created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
+		
+		FOREIGN KEY (policy_id) REFERENCES policies(id) ON DELETE CASCADE,
+		UNIQUE KEY policy_team_unique (policy_id, inherited_team_id)
+);
+`
+
+	_, err := tx.Exec(stmt)
+	if err != nil {
+		return fmt.Errorf("failed to create policy_stats table: %w", err)
+	}
+
+	return nil
+}
+
+func Down_20231121054530(tx *sql.Tx) error {
+	return nil
+}

server/datastore/mysql/migrations/tables/20231121054530_CreatePolicyStatsTable_test.go

@@ -0,0 +1,83 @@
+package tables
+
+import (
+	"database/sql"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestUp_20231121054530(t *testing.T) {
+	db := applyUpToPrev(t)
+
+	const (
+		insertUsersStmt = `INSERT INTO users (id, name, password, salt, email) VALUES (?, ?, ?, ?, ?)`
+
+		insertPolicyStmt = `INSERT INTO policies (
+			team_id, name, query, description, author_id, platforms, critical
+		) VALUES (?, ?, ?, ?, ?, ?, ?)`
+
+		insertTeamStmt = `INSERT INTO teams (name) VALUES (?)`
+
+		deletePolicyStmt = `DELETE FROM policies WHERE id = ?`
+
+		loadPolicyStatsStmt = `SELECT
+			id, policy_id, inherited_team_id, passing_host_count, failing_host_count
+		FROM policy_stats WHERE id = ?`
+	)
+
+	// Apply current migration.
+	applyNext(t, db)
+
+	// Create a user
+	_, err := db.Exec(insertUsersStmt, 1, "user1", "password999", "salt999", "foo")
+	require.NoError(t, err)
+
+	// Create a team
+	res, err := db.Exec(insertTeamStmt, "team1")
+	require.NoError(t, err)
+	teamID, _ := res.LastInsertId()
+
+	// Create a global policy
+	res, err = db.Exec(insertPolicyStmt, nil, "global-policy", "SELECT 1;", "Global policy description", 1, "all", false)
+	require.NoError(t, err)
+	globalPolicyStatID, _ := res.LastInsertId()
+
+	// Insert a policy_stats entry for the global policy (globally)
+	_, err = db.Exec(`INSERT INTO policy_stats (policy_id, inherited_team_id, passing_host_count, failing_host_count) VALUES (?, 0, ?, ?)`, globalPolicyStatID, 100, 10)
+	require.NoError(t, err)
+
+	// Insert a policy_stats entry for the team inheriting the global policy
+	_, err = db.Exec(`INSERT INTO policy_stats (policy_id, inherited_team_id, passing_host_count, failing_host_count) VALUES (?, ?, ?, ?)`, globalPolicyStatID, teamID, 50, 5)
+	require.NoError(t, err)
+
+	// Verify the entries in the policy_stats table
+	var id int
+	var policyID int64
+	var inheritedTeamID int64
+	var passingCount, failingCount int
+
+	// Verify global policy stats (global level)
+	err = db.QueryRow(loadPolicyStatsStmt, 1).Scan(&id, &policyID, &inheritedTeamID, &passingCount, &failingCount)
+	require.NoError(t, err)
+	require.Equal(t, globalPolicyStatID, policyID)
+	require.Equal(t, int64(0), inheritedTeamID)
+
+	// Verify global policy stats (team level)
+	err = db.QueryRow(loadPolicyStatsStmt, 2).Scan(&id, &policyID, &inheritedTeamID, &passingCount, &failingCount)
+	require.NoError(t, err)
+	require.Equal(t, globalPolicyStatID, policyID)
+	require.Equal(t, teamID, inheritedTeamID)
+
+	// Verify global policy stats still exist (global level)
+	err = db.QueryRow(loadPolicyStatsStmt, globalPolicyStatID).Scan(&id, &policyID, &inheritedTeamID, &passingCount, &failingCount)
+	require.NoError(t, err)
+
+	// Delete the global policy and check that its policy_stats entry is also deleted
+	_, err = db.Exec(deletePolicyStmt, globalPolicyStatID)
+	require.NoError(t, err)
+
+	err = db.QueryRow(loadPolicyStatsStmt, globalPolicyStatID).Scan(&id, &policyID, &inheritedTeamID, &passingCount, &failingCount)
+	require.Error(t, err)
+	require.ErrorIs(t, err, sql.ErrNoRows)
+}

server/datastore/mysql/policies.go

@@ -62,10 +62,13 @@ func (ds *Datastore) PolicyByName(ctx context.Context, name string) (*fleet.Poli
 		fmt.Sprint(`SELECT `+policyCols+`,
 		COALESCE(u.name, '<deleted>') AS author_name,
 		COALESCE(u.email, '') AS author_email,
-		(select count(*) from policy_membership where policy_id=p.id and passes=true) as passing_host_count,
-		(select count(*) from policy_membership where policy_id=p.id and passes=false) as failing_host_count
+		COALESCE(ps.passing_host_count, 0) as passing_host_count,
+		COALESCE(ps.failing_host_count, 0) as failing_host_count
 		FROM policies p
 		LEFT JOIN users u ON p.author_id = u.id
+		LEFT JOIN policy_stats ps ON p.id = ps.policy_id
+			AND ((p.team_id IS NULL AND ps.inherited_team_id = 0)
+				OR (p.team_id IS NOT NULL AND ps.inherited_team_id = p.team_id))
 		WHERE p.name=?`), name)
 	if err != nil {
 		if err == sql.ErrNoRows {
@@ -87,14 +90,17 @@ func policyDB(ctx context.Context, q sqlx.QueryerContext, id uint, teamID *uint)
 	var policy fleet.Policy
 	err := sqlx.GetContext(ctx, q, &policy,
 		fmt.Sprintf(`
-		SELECT `+policyCols+`,
+		SELECT %s,
 		    COALESCE(u.name, '<deleted>') AS author_name,
 			COALESCE(u.email, '') AS author_email,
-			(select count(*) from policy_membership where policy_id=p.id and passes=true) as passing_host_count,
-			(select count(*) from policy_membership where policy_id=p.id and passes=false) as failing_host_count
+			COALESCE(ps.passing_host_count, 0) as passing_host_count,
+			COALESCE(ps.failing_host_count, 0) as failing_host_count
 		FROM policies p
 		LEFT JOIN users u ON p.author_id = u.id
-		WHERE p.id=? AND %s`, teamWhere),
+		LEFT JOIN policy_stats ps ON p.id = ps.policy_id
+		AND ((p.team_id IS NULL AND ps.inherited_team_id = 0)
+			OR (p.team_id IS NOT NULL AND ps.inherited_team_id = p.team_id))
+		WHERE p.id=? AND %s`, policyCols, teamWhere),
 		args...)
 	if err != nil {
 		if err == sql.ErrNoRows {
@@ -292,112 +298,75 @@ func (ds *Datastore) RecordPolicyQueryExecutions(ctx context.Context, host *flee
 }
 
 func (ds *Datastore) ListGlobalPolicies(ctx context.Context, opts fleet.ListOptions) ([]*fleet.Policy, error) {
-	return listPoliciesDB(ctx, ds.reader(ctx), nil, nil, opts)
+	return listPoliciesDB(ctx, ds.reader(ctx), nil, opts)
 }
 
 // returns the list of policies associated with the provided teamID, or the
 // global policies if teamID is nil. The pass/fail host counts are the totals
 // regardless of hosts' team if countsForTeamID is nil, or the totals just for
 // hosts that belong to the provided countsForTeamID if it is not nil.
-func listPoliciesDB(ctx context.Context, q sqlx.QueryerContext, teamID, countsForTeamID *uint, opts fleet.ListOptions) ([]*fleet.Policy, error) {
+func listPoliciesDB(ctx context.Context, q sqlx.QueryerContext, teamID *uint, opts fleet.ListOptions) ([]*fleet.Policy, error) {
 	var args []interface{}
 
-	var initialQuery string
-
-	// Sorting by failing host counts requires an expensive join with the
-	// policy membership table and may result in long response times
-	if opts.OrderKey == "failing_host_count" {
-		if countsForTeamID != nil {
-			initialQuery = `
-				SELECT p.id
-				FROM policies p
-				LEFT JOIN (
-					SELECT pm.policy_id,
-						COUNT(*) AS failing_host_count
-					FROM policy_membership pm
-					INNER JOIN hosts h ON pm.host_id = h.id AND pm.passes = false AND h.team_id = ?
-					GROUP BY pm.policy_id
-				) AS subq ON p.id = subq.policy_id
-			`
-			args = append(args, *countsForTeamID)
-		} else {
-			initialQuery = `
-				SELECT p.id
-				FROM policies p
-				LEFT JOIN (
-					SELECT pm.policy_id,
-						COUNT(*) AS failing_host_count
-					FROM policy_membership pm
-					WHERE pm.passes = false
-					GROUP BY pm.policy_id
-				) AS subq ON p.id = subq.policy_id
-			`
-		}
-	} else {
-		initialQuery = "SELECT id FROM policies"
-	}
-
-	if teamID != nil {
-		initialQuery += " WHERE team_id = ?"
-		args = append(args, *teamID)
-	} else {
-		initialQuery += " WHERE team_id IS NULL"
-	}
-
-	initialQuery, args = searchLike(initialQuery, args, opts.MatchQuery, policySearchColumns...)
-	initialQuery, args = appendListOptionsWithCursorToSQL(initialQuery, args, &opts)
-
-	var ids []uint
-	err := sqlx.SelectContext(ctx, q, &ids, initialQuery, args...)
-	if err != nil {
-		return nil, ctxerr.Wrap(ctx, err, "retrieving policy ids")
-	}
-
-	if len(ids) == 0 {
-		return []*fleet.Policy{}, nil
-	}
-
-	args = []interface{}{} // reset args
-
-	counts := `
-	(select count(*) from policy_membership where policy_id=p.id and passes=true) as passing_host_count,
-	(select count(*) from policy_membership where policy_id=p.id and passes=false) as failing_host_count
-`
-	if countsForTeamID != nil {
-		counts = `
-		(select count(*) from policy_membership pm inner join hosts h on pm.host_id = h.id where pm.policy_id=p.id and pm.passes=true and h.team_id = ?) as passing_host_count,
-		(select count(*) from policy_membership pm inner join hosts h on pm.host_id = h.id where pm.policy_id=p.id and pm.passes=false and h.team_id = ?) as failing_host_count
-`
-		args = append(args, *countsForTeamID, *countsForTeamID)
-	}
-
-	query := fmt.Sprintf(`
-		SELECT `+policyCols+`,
+	query := `
+		SELECT ` + policyCols + `,
 			COALESCE(u.name, '<deleted>') AS author_name,
 			COALESCE(u.email, '') AS author_email,
-			%s
+			COALESCE(ps.passing_host_count, 0) AS passing_host_count,
+			COALESCE(ps.failing_host_count, 0) AS failing_host_count
 		FROM policies p
 		LEFT JOIN users u ON p.author_id = u.id
-		WHERE p.id IN (?)`, counts)
+		LEFT JOIN policy_stats ps ON p.id = ps.policy_id AND ps.inherited_team_id = 0
+	`
 
-	args = append(args, ids)
-
-	query, args, err = sqlx.In(query, args...)
-	if err != nil {
-		return nil, ctxerr.Wrap(ctx, err, "building query to get policies by ID")
+	if teamID != nil {
+		query += " WHERE team_id = ?"
+		args = append(args, *teamID)
+	} else {
+		query += " WHERE team_id IS NULL"
 	}
 
-	// removing pagination options to avoid double pagination
-	opts.Page = 0
-	opts.PerPage = 0
-
+	query, args = searchLike(query, args, opts.MatchQuery, policySearchColumns...)
 	query, args = appendListOptionsWithCursorToSQL(query, args, &opts)
 
 	var policies []*fleet.Policy
-	err = sqlx.SelectContext(ctx, q, &policies, query, args...)
+	err := sqlx.SelectContext(ctx, q, &policies, query, args...)
 	if err != nil {
 		return nil, ctxerr.Wrap(ctx, err, "listing policies")
 	}
+
+	return policies, nil
+}
+
+// getInheritedPoliciesForTeam returns the list of global policies with the
+// passing and failing host counts for the provided teamID
+func getInheritedPoliciesForTeam(ctx context.Context, q sqlx.QueryerContext, TeamID uint, opts fleet.ListOptions) ([]*fleet.Policy, error) {
+	var args []interface{}
+
+	query := `
+        SELECT 
+            ` + policyCols + `,
+			COALESCE(u.name, '<deleted>') AS author_name,
+			COALESCE(u.email, '') AS author_email,
+            COALESCE(ps.passing_host_count, 0) as passing_host_count,
+            COALESCE(ps.failing_host_count, 0) as failing_host_count
+        FROM policies p
+        LEFT JOIN users u ON p.author_id = u.id
+        LEFT JOIN policy_stats ps ON p.id = ps.policy_id AND ps.inherited_team_id = ?
+        WHERE p.team_id IS NULL
+    `
+
+	args = append(args, TeamID)
+
+	query, args = searchLike(query, args, opts.MatchQuery, policySearchColumns...)
+	query, _ = appendListOptionsToSQL(query, &opts)
+
+	var policies []*fleet.Policy
+	err := sqlx.SelectContext(ctx, q, &policies, query, args...)
+	if err != nil {
+		return nil, ctxerr.Wrap(ctx, err, "listing inherited policies")
+	}
+
 	return policies, nil
 }
 
@@ -431,10 +400,13 @@ func (ds *Datastore) PoliciesByID(ctx context.Context, ids []uint) (map[uint]*fl
 	sql := `SELECT ` + policyCols + `,
 	  COALESCE(u.name, '<deleted>') AS author_name,
 	  COALESCE(u.email, '') AS author_email,
-	  (select count(*) from policy_membership where policy_id=p.id and passes=true) as passing_host_count,
-	  (select count(*) from policy_membership where policy_id=p.id and passes=false) as failing_host_count
+	  COALESCE(ps.passing_host_count, 0) as passing_host_count,
+	  COALESCE(ps.failing_host_count, 0) as failing_host_count
 	  FROM policies p
 	  LEFT JOIN users u ON p.author_id = u.id
+	  LEFT JOIN policy_stats ps ON p.id = ps.policy_id
+	  	AND ((p.team_id IS NULL AND ps.inherited_team_id = 0)
+				OR (p.team_id IS NOT NULL AND ps.inherited_team_id = p.team_id))
 	  WHERE p.id IN (?)`
 	query, args, err := sqlx.In(sql, ids)
 	if err != nil {
@@ -561,12 +533,12 @@ func (ds *Datastore) NewTeamPolicy(ctx context.Context, teamID uint, authorID *u
 }
 
 func (ds *Datastore) ListTeamPolicies(ctx context.Context, teamID uint, opts fleet.ListOptions, iopts fleet.ListOptions) (teamPolicies, inheritedPolicies []*fleet.Policy, err error) {
-	teamPolicies, err = listPoliciesDB(ctx, ds.reader(ctx), &teamID, nil, opts)
+	teamPolicies, err = listPoliciesDB(ctx, ds.reader(ctx), &teamID, opts)
 	if err != nil {
 		return nil, nil, err
 	}
 	// get inherited (global) policies with counts of hosts for that team
-	inheritedPolicies, err = listPoliciesDB(ctx, ds.reader(ctx), nil, &teamID, iopts)
+	inheritedPolicies, err = getInheritedPoliciesForTeam(ctx, ds.reader(ctx), teamID, iopts)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -1133,3 +1105,56 @@ func amountPolicyViolationDaysDB(ctx context.Context, tx sqlx.QueryerContext) (i
 
 	return int(counts.FailingHostCount), int(counts.TotalHostCount), nil
 }
+
+func (ds *Datastore) UpdateHostPolicyCounts(ctx context.Context) error {
+	// Update Counts for Inherited Global Policies for each Team
+	_, err := ds.writer(ctx).ExecContext(ctx, `
+		INSERT INTO policy_stats (policy_id, inherited_team_id, passing_host_count, failing_host_count)
+		SELECT
+			p.id,
+			t.id AS inherited_team_id,
+			(
+				SELECT COUNT(*) 
+				FROM policy_membership pm 
+				INNER JOIN hosts h ON pm.host_id = h.id 
+				WHERE pm.policy_id = p.id AND pm.passes = true AND h.team_id = t.id
+			) AS passing_host_count,
+			(
+				SELECT COUNT(*) 
+				FROM policy_membership pm 
+				INNER JOIN hosts h ON pm.host_id = h.id 
+				WHERE pm.policy_id = p.id AND pm.passes = false AND h.team_id = t.id
+			) AS failing_host_count
+		FROM policies p
+		CROSS JOIN teams t
+		WHERE p.team_id IS NULL
+		GROUP BY p.id, t.id
+		ON DUPLICATE KEY UPDATE 
+			passing_host_count = VALUES(passing_host_count),
+			failing_host_count = VALUES(failing_host_count);
+    `)
+	if err != nil {
+		return ctxerr.Wrap(ctx, err, "update host policy counts for inherited global policies")
+	}
+
+	// Update Counts for Global and Team Policies
+	_, err = ds.writer(ctx).ExecContext(ctx, `
+		INSERT INTO policy_stats (policy_id, inherited_team_id, passing_host_count, failing_host_count)
+		SELECT
+			p.id,
+			0 AS inherited_team_id, -- using 0 to represent global scope
+			COALESCE(SUM(IF(pm.passes IS NULL, 0, pm.passes = 1)), 0), 
+			COALESCE(SUM(IF(pm.passes IS NULL, 0, pm.passes = 0)), 0)
+		FROM policies p
+		LEFT JOIN policy_membership pm ON p.id = pm.policy_id
+		GROUP BY p.id
+		ON DUPLICATE KEY UPDATE 
+			passing_host_count = VALUES(passing_host_count),
+			failing_host_count = VALUES(failing_host_count);
+    `)
+	if err != nil {
+		return ctxerr.Wrap(ctx, err, "update host policy counts for global and team policies")
+	}
+
+	return nil
+}

server/datastore/mysql/policies_test.go

@@ -53,6 +53,7 @@ func TestPolicies(t *testing.T) {
 		{"TestListGlobalPoliciesCanPaginate", testListGlobalPoliciesCanPaginate},
 		{"TestListTeamPoliciesCanPaginate", testListTeamPoliciesCanPaginate},
 		{"TestCountPolicies", testCountPolicies},
+		{"TestUpdatePolicyHostCounts", testUpdatePolicyHostCounts},
 	}
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
@@ -277,6 +278,8 @@ func testPoliciesMembershipView(deferred bool, t *testing.T, ds *Datastore) {
 
 	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, host2, map[uint]*bool{p2.ID: nil}, time.Now(), deferred))
 
+	require.NoError(t, ds.UpdateHostPolicyCounts(ctx))
+
 	policies, err := ds.ListGlobalPolicies(ctx, fleet.ListOptions{})
 	require.NoError(t, err)
 	require.Len(t, policies, 2)
@@ -292,6 +295,8 @@ func testPoliciesMembershipView(deferred bool, t *testing.T, ds *Datastore) {
 	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, host1, map[uint]*bool{p.ID: ptr.Bool(false)}, time.Now(), deferred))
 	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, host2, map[uint]*bool{p2.ID: ptr.Bool(false)}, time.Now(), deferred))
 
+	require.NoError(t, ds.UpdateHostPolicyCounts(ctx))
+
 	policies, err = ds.ListGlobalPolicies(ctx, fleet.ListOptions{})
 	require.NoError(t, err)
 	require.Len(t, policies, 2)
@@ -349,6 +354,8 @@ func testPoliciesMembershipView(deferred bool, t *testing.T, ds *Datastore) {
 	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, host4, map[uint]*bool{t2pol.ID: ptr.Bool(false), t2pol2.ID: ptr.Bool(true), p.ID: ptr.Bool(false)}, time.Now(), deferred))
 	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, host5, map[uint]*bool{t2pol.ID: ptr.Bool(true), t2pol2.ID: ptr.Bool(true), p2.ID: ptr.Bool(true)}, time.Now(), deferred))
 
+	require.NoError(t, ds.UpdateHostPolicyCounts(ctx))
+
 	t1Pols, t1Inherited, err := ds.ListTeamPolicies(ctx, team1.ID, fleet.ListOptions{}, fleet.ListOptions{})
 	require.NoError(t, err)
 	require.Len(t, t1Pols, 1)
@@ -1028,6 +1035,11 @@ func testPoliciesByID(t *testing.T, ds *Datastore) {
 	user1 := test.NewUser(t, ds, "Alice", "alice@example.com", true)
 	policy1 := newTestPolicy(t, ds, user1, "policy1", "darwin", nil)
 	_ = newTestPolicy(t, ds, user1, "policy2", "darwin", nil)
+	host1 := newTestHostWithPlatform(t, ds, "host1", "darwin", nil)
+
+	require.NoError(t, ds.RecordPolicyQueryExecutions(context.Background(), host1, map[uint]*bool{policy1.ID: ptr.Bool(true)}, time.Now(), false))
+	require.NoError(t, ds.UpdateHostPolicyCounts(context.Background()))
+
 	policiesByID, err := ds.PoliciesByID(context.Background(), []uint{1, 2})
 	require.NoError(t, err)
 	assert.Equal(t, len(policiesByID), 2)
@@ -1035,6 +1047,7 @@ func testPoliciesByID(t *testing.T, ds *Datastore) {
 	assert.Equal(t, policiesByID[1].Name, policy1.Name)
 	assert.Equal(t, policiesByID[2].ID, uint(2))
 	assert.Equal(t, policiesByID[2].Name, "policy2")
+	assert.Equal(t, uint(1), policiesByID[1].PassingHostCount)
 
 	_, err = ds.PoliciesByID(context.Background(), []uint{1, 2, 3})
 	require.Error(t, err)
@@ -1100,7 +1113,11 @@ func testTeamPolicyTransfer(t *testing.T, ds *Datastore) {
 	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, host2, map[uint]*bool{team1Policy.ID: ptr.Bool(false), globalPolicy.ID: ptr.Bool(true)}, time.Now(), false))
 	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, host2, map[uint]*bool{team1Policy.ID: ptr.Bool(true), globalPolicy.ID: ptr.Bool(true)}, time.Now(), false))
 
+	require.NoError(t, ds.UpdateHostPolicyCounts(ctx))
+
 	checkPassingCount := func(tm1, tm1Inherited, tm2Inherited, global uint) {
+		t.Helper()
+		require.NoError(t, ds.UpdateHostPolicyCounts(ctx))
 		policies, inherited, err := ds.ListTeamPolicies(ctx, team1.ID, fleet.ListOptions{}, fleet.ListOptions{})
 		require.NoError(t, err)
 		require.Len(t, policies, 1)
@@ -2429,3 +2446,61 @@ func testCountPolicies(t *testing.T, ds *Datastore) {
 	require.NoError(t, err)
 	assert.Equal(t, 10, globalCount)
 }
+
+func testUpdatePolicyHostCounts(t *testing.T, ds *Datastore) {
+	// new policy
+	policy, err := ds.NewGlobalPolicy(context.Background(), nil, fleet.PolicyPayload{Name: "global policy 1"})
+	require.NoError(t, err)
+
+	team, err := ds.NewTeam(context.Background(), &fleet.Team{Name: "team1"})
+	require.NoError(t, err)
+
+	// create 4 team hosts
+	var teamHosts []*fleet.Host
+	for i := 0; i < 4; i++ {
+		h, err := ds.NewHost(context.Background(), &fleet.Host{OsqueryHostID: ptr.String(fmt.Sprintf("host%d", i)), NodeKey: ptr.String(fmt.Sprintf("host%d", i)), TeamID: &team.ID})
+		require.NoError(t, err)
+		teamHosts = append(teamHosts, h)
+	}
+
+	// create 4 global hosts
+	var globalHosts []*fleet.Host
+	for i := 4; i < 8; i++ {
+		h, err := ds.NewHost(context.Background(), &fleet.Host{OsqueryHostID: ptr.String(fmt.Sprintf("host%d", i)), NodeKey: ptr.String(fmt.Sprintf("host%d", i)), TeamID: nil})
+		require.NoError(t, err)
+		globalHosts = append(globalHosts, h)
+	}
+
+	// add policy responses
+	for _, h := range teamHosts {
+		res := map[uint]*bool{
+			policy.ID: ptr.Bool(true),
+		}
+		err = ds.RecordPolicyQueryExecutions(context.Background(), h, res, time.Now(), false)
+		require.NoError(t, err)
+	}
+
+	for _, h := range globalHosts {
+		res := map[uint]*bool{
+			policy.ID: ptr.Bool(true),
+		}
+		err = ds.RecordPolicyQueryExecutions(context.Background(), h, res, time.Now(), false)
+		require.NoError(t, err)
+	}
+
+	// check policy host counts before update
+	policy, err = ds.Policy(context.Background(), policy.ID)
+	require.NoError(t, err)
+	require.Equal(t, uint(0), policy.FailingHostCount)
+	require.Equal(t, uint(0), policy.PassingHostCount)
+
+	// update policy host counts
+	err = ds.UpdateHostPolicyCounts(context.Background())
+	require.NoError(t, err)
+
+	// check policy host counts
+	policy, err = ds.Policy(context.Background(), policy.ID)
+	require.NoError(t, err)
+	require.Equal(t, uint(0), policy.FailingHostCount)
+	require.Equal(t, uint(8), policy.PassingHostCount)
+}

server/fleet/datastore.go

@@ -529,6 +529,7 @@ type Datastore interface {
 	PoliciesByID(ctx context.Context, ids []uint) (map[uint]*Policy, error)
 	DeleteGlobalPolicies(ctx context.Context, ids []uint) ([]uint, error)
 	CountPolicies(ctx context.Context, teamID *uint, matchQuery string) (int, error)
+	UpdateHostPolicyCounts(ctx context.Context) error
 
 	PolicyQueriesForHost(ctx context.Context, host *Host) (map[string]string, error)
 

server/mock/datastore_mock.go

@@ -392,6 +392,8 @@ type DeleteGlobalPoliciesFunc func(ctx context.Context, ids []uint) ([]uint, err
 
 type CountPoliciesFunc func(ctx context.Context, teamID *uint, matchQuery string) (int, error)
 
+type UpdateHostPolicyCountsFunc func(ctx context.Context) error
+
 type PolicyQueriesForHostFunc func(ctx context.Context, host *fleet.Host) (map[string]string, error)
 
 type AsyncBatchInsertPolicyMembershipFunc func(ctx context.Context, batch []fleet.PolicyMembershipResult) error
@@ -1314,6 +1316,9 @@ type DataStore struct {
 	CountPoliciesFunc        CountPoliciesFunc
 	CountPoliciesFuncInvoked bool
 
+	UpdateHostPolicyCountsFunc        UpdateHostPolicyCountsFunc
+	UpdateHostPolicyCountsFuncInvoked bool
+
 	PolicyQueriesForHostFunc        PolicyQueriesForHostFunc
 	PolicyQueriesForHostFuncInvoked bool
 
@@ -3166,6 +3171,13 @@ func (s *DataStore) CountPolicies(ctx context.Context, teamID *uint, matchQuery
 	return s.CountPoliciesFunc(ctx, teamID, matchQuery)
 }
 
+func (s *DataStore) UpdateHostPolicyCounts(ctx context.Context) error {
+	s.mu.Lock()
+	s.UpdateHostPolicyCountsFuncInvoked = true
+	s.mu.Unlock()
+	return s.UpdateHostPolicyCountsFunc(ctx)
+}
+
 func (s *DataStore) PolicyQueriesForHost(ctx context.Context, host *fleet.Host) (map[string]string, error) {
 	s.mu.Lock()
 	s.PolicyQueriesForHostFuncInvoked = true