Cleanup query results after host is transferred to another team (#18712)

#18079

- [X] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://fleetdm.com/docs/contributing/committing-changes#changes-files)
for more information.
- [X] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements)
- ~[ ] Added support on fleet's osquery simulator `cmd/osquery-perf` for
new osquery data ingestion features.~
- [X] Added/updated tests
- [X] If database migrations are included, checked table schema to
confirm autoupdate
- For database migrations:
- ~[ ] Checked schema for all modified table for columns that will
auto-update timestamps during migration.~
- ~[ ] Confirmed that updating the timestamps is acceptable, and will
not cause unwanted side effects.~
- ~[ ] Ensured the correct collation is explicitly set for character
columns (`COLLATE utf8mb4_unicode_ci`).~
- [X] Manual QA for all new/changed functionality
  - ~For Orbit and Fleet Desktop changes:~
- ~[ ] Manual QA must be performed in the three main OSs, macOS, Windows
and Linux.~
- ~[ ] Auto-update manual QA, from released version of component to new
version (see [tools/tuf/test](../tools/tuf/test/README.md)).~

changes/18079-query-results-bug-transfer-hosts

@@ -0,0 +1 @@
+* Fixed bug where hosts query results were not cleared after transferring the host to other teams.

server/datastore/mysql/hosts.go

@@ -2717,6 +2717,9 @@ func (ds *Datastore) AddHostsToTeam(ctx context.Context, teamID *uint, hostIDs [
 		if err := cleanupPolicyMembershipOnTeamChange(ctx, tx, hostIDs); err != nil {
 			return ctxerr.Wrap(ctx, err, "AddHostsToTeam delete policy membership")
 		}
+		if err := cleanupQueryResultsOnTeamChange(ctx, tx, hostIDs); err != nil {
+			return ctxerr.Wrap(ctx, err, "AddHostsToTeam delete query results")
+		}
 
 		query, args, err := sqlx.In(`UPDATE hosts SET team_id = ? WHERE id IN (?)`, teamID, hostIDs)
 		if err != nil {

server/datastore/mysql/hosts_test.go

@@ -165,6 +165,7 @@ func TestHosts(t *testing.T) {
 		{"HostHealth", testHostHealth},
 		{"GetHostOrbitInfo", testGetHostOrbitInfo},
 		{"HostnamesByIdentifiers", testHostnamesByIdentifiers},
+		{"HostsAddToTeamCleansUpTeamQueryResults", testHostsAddToTeamCleansUpTeamQueryResults},
 	}
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
@@ -8860,3 +8861,200 @@ func testHostnamesByIdentifiers(t *testing.T, ds *Datastore) {
 		})
 	}
 }
+
+func testHostsAddToTeamCleansUpTeamQueryResults(t *testing.T, ds *Datastore) {
+	ctx := context.Background()
+
+	team1, err := ds.NewTeam(ctx, &fleet.Team{Name: "team1"})
+	require.NoError(t, err)
+	team2, err := ds.NewTeam(ctx, &fleet.Team{Name: "team2"})
+	require.NoError(t, err)
+
+	hostCount := 1
+	newHost := func(teamID *uint) *fleet.Host {
+		h, err := ds.NewHost(ctx, &fleet.Host{
+			OsqueryHostID: ptr.String(fmt.Sprintf("foobar%d", hostCount)),
+			NodeKey:       ptr.String(fmt.Sprintf("nodekey%d", hostCount)),
+			TeamID:        teamID,
+		})
+		require.NoError(t, err)
+		hostCount++
+		return h
+	}
+	newQuery := func(name string, teamID *uint) *fleet.Query {
+		q, err := ds.NewQuery(ctx, &fleet.Query{
+			Name:    name,
+			Query:   "SELECT 1:",
+			TeamID:  teamID,
+			Logging: fleet.LoggingSnapshot,
+		})
+		require.NoError(t, err)
+		return q
+	}
+
+	h0 := newHost(nil)
+	h1 := newHost(&team1.ID)
+	h2 := newHost(&team2.ID)
+	h3 := newHost(&team2.ID)
+
+	hostStaticOnTeam1 := newHost(&team1.ID) // host that we won't move
+
+	query0Global := newQuery("query0Global", nil)
+	query1Team1 := newQuery("query1Team1", &team1.ID)
+	query2Team2 := newQuery("query2Team2", &team2.ID)
+
+	// Transfer h2 from team2 to team1 and back without any query results yet.
+	err = ds.AddHostsToTeam(ctx, &team1.ID, []uint{h2.ID})
+	require.NoError(t, err)
+	err = ds.AddHostsToTeam(ctx, &team2.ID, []uint{h2.ID})
+	require.NoError(t, err)
+
+	data := ptr.RawMessage(json.RawMessage(`{"foo": "bar"}`))
+	h0Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  h0.ID,
+			QueryID: query0Global.ID,
+			Data:    data,
+		},
+	}
+	h1Global0Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  h1.ID,
+			QueryID: query0Global.ID,
+			Data:    data,
+		},
+	}
+	h1Query1Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  h1.ID,
+			QueryID: query1Team1.ID,
+			Data:    data,
+		},
+	}
+	h2Global0Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  h2.ID,
+			QueryID: query0Global.ID,
+			Data:    data,
+		},
+	}
+	h2Query2Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  h2.ID,
+			QueryID: query2Team2.ID,
+			Data:    data,
+		},
+	}
+	h3Global0Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  h3.ID,
+			QueryID: query0Global.ID,
+			Data:    data,
+		},
+	}
+	h3Query2Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  h3.ID,
+			QueryID: query2Team2.ID,
+			Data:    data,
+		},
+	}
+	h4Global0Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  hostStaticOnTeam1.ID,
+			QueryID: query0Global.ID,
+			Data:    data,
+		},
+	}
+	h4Query1Results := []*fleet.ScheduledQueryResultRow{
+		{
+			HostID:  hostStaticOnTeam1.ID,
+			QueryID: query1Team1.ID,
+			Data:    data,
+		},
+	}
+	for _, results := range [][]*fleet.ScheduledQueryResultRow{
+		h0Results,
+		h1Global0Results,
+		h1Query1Results,
+		h2Global0Results,
+		h2Query2Results,
+		h3Global0Results,
+		h3Query2Results,
+		h4Global0Results,
+		h4Query1Results,
+	} {
+		err = ds.OverwriteQueryResultRows(ctx, results)
+		require.NoError(t, err)
+	}
+
+	tf := fleet.TeamFilter{
+		User: &fleet.User{
+			GlobalRole: ptr.String(fleet.RoleAdmin),
+		},
+	}
+
+	rows, err := ds.QueryResultRows(ctx, query0Global.ID, tf)
+	require.NoError(t, err)
+	require.Len(t, rows, 5)
+	rows, err = ds.QueryResultRows(ctx, query1Team1.ID, tf)
+	require.NoError(t, err)
+	require.Len(t, rows, 2)
+	rows, err = ds.QueryResultRows(ctx, query2Team2.ID, tf)
+	require.NoError(t, err)
+	require.Len(t, rows, 2)
+
+	// Transfer h2 from team2 to team1.
+	err = ds.AddHostsToTeam(ctx, &team1.ID, []uint{h2.ID})
+	require.NoError(t, err)
+	// Transfer h1 from team1 to team2.
+	err = ds.AddHostsToTeam(ctx, &team2.ID, []uint{h1.ID})
+	require.NoError(t, err)
+	// Transfer h3 from team2 to global.
+	err = ds.AddHostsToTeam(ctx, nil, []uint{h3.ID})
+	require.NoError(t, err)
+
+	// No global query results should be deleted
+	rows, err = ds.QueryResultRows(ctx, query0Global.ID, tf)
+	require.NoError(t, err)
+	require.Len(t, rows, 5)
+	// Results for h1 should be gone, and results for hostStaticOnTeam1 should be here.
+	rows, err = ds.QueryResultRows(ctx, query1Team1.ID, tf)
+	require.NoError(t, err)
+	require.Len(t, rows, 1)
+	require.Equal(t, hostStaticOnTeam1.ID, rows[0].HostID)
+	// Results for h2 and h3 should be gone.
+	rows, err = ds.QueryResultRows(ctx, query2Team2.ID, tf)
+	require.NoError(t, err)
+	require.Empty(t, rows)
+
+	// h1 should have only the global result.
+	h1, err = ds.Host(ctx, h1.ID)
+	require.NoError(t, err)
+	require.Len(t, h1.PackStats, 1)
+	require.Len(t, h1.PackStats[0].QueryStats, 1)
+	require.Equal(t, query0Global.ID, h1.PackStats[0].QueryStats[0].ScheduledQueryID)
+
+	// h2 should have only the global result.
+	h2, err = ds.Host(ctx, h2.ID)
+	require.NoError(t, err)
+	require.Len(t, h2.PackStats, 1)
+	require.Len(t, h2.PackStats[0].QueryStats, 1)
+	require.Equal(t, query0Global.ID, h2.PackStats[0].QueryStats[0].ScheduledQueryID)
+
+	// h3 should have only the global result.
+	h3, err = ds.Host(ctx, h3.ID)
+	require.NoError(t, err)
+	require.Len(t, h3.PackStats, 1)
+	require.Len(t, h3.PackStats[0].QueryStats, 1)
+	require.Equal(t, query0Global.ID, h3.PackStats[0].QueryStats[0].ScheduledQueryID)
+
+	// hostStaticOnTeam1 should have the global result and the team1 result.
+	hostStaticOnTeam1, err = ds.Host(ctx, hostStaticOnTeam1.ID)
+	require.NoError(t, err)
+	require.Len(t, hostStaticOnTeam1.PackStats, 2)
+	require.Len(t, hostStaticOnTeam1.PackStats[0].QueryStats, 1)
+	require.Equal(t, query0Global.ID, hostStaticOnTeam1.PackStats[0].QueryStats[0].ScheduledQueryID)
+	require.Len(t, hostStaticOnTeam1.PackStats[1].QueryStats, 1)
+	require.Equal(t, query1Team1.ID, hostStaticOnTeam1.PackStats[1].QueryStats[0].ScheduledQueryID)
+}

server/datastore/mysql/migrations/tables/20240430111727_CleanupQueryResults.go

@@ -0,0 +1,30 @@
+package tables
+
+import (
+	"database/sql"
+	"fmt"
+)
+
+func init() {
+	MigrationClient.AddMigration(Up_20240430111727, Down_20240430111727)
+}
+
+func Up_20240430111727(tx *sql.Tx) error {
+	// This cleanup correspond to the following bug: https://github.com/fleetdm/fleet/issues/18079.
+	// The following deletes "team query results" that do not match the host's team.
+	_, err := tx.Exec(`
+		DELETE qr
+		FROM query_results qr
+		JOIN queries q ON (q.id=qr.query_id)
+		JOIN hosts h ON (h.id=qr.host_id)
+		WHERE q.team_id IS NOT NULL AND q.team_id != COALESCE(h.team_id, 0);
+	`)
+	if err != nil {
+		return fmt.Errorf("failed to delete query_results %w", err)
+	}
+	return nil
+}
+
+func Down_20240430111727(tx *sql.Tx) error {
+	return nil
+}

server/datastore/mysql/migrations/tables/20240430111727_CleanupQueryResults_test.go

@@ -0,0 +1,98 @@
+package tables
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestUp_20240430111727(t *testing.T) {
+	db := applyUpToPrev(t)
+
+	hostID := 1
+	newTeam := func(name string) uint {
+		return uint(execNoErrLastID(t, db,
+			`INSERT INTO teams (name) VALUES (?);`,
+			name,
+		))
+	}
+	newHost := func(teamID *uint) uint {
+		id := fmt.Sprintf("%d", hostID)
+		hostID++
+		return uint(execNoErrLastID(t, db,
+			`INSERT INTO hosts (osquery_host_id, node_key, team_id) VALUES (?, ?, ?);`,
+			id, id, teamID,
+		))
+	}
+	newQuery := func(name string, teamID *uint) uint {
+		return uint(execNoErrLastID(t, db,
+			`INSERT INTO queries (name, description, logging_type, team_id, query, saved) VALUES (?, '', 'snapshot', ?, 'SELECT 1;', 1);`,
+			name, teamID,
+		))
+	}
+	newQueryResults := func(queryID, hostID uint, resultCount int) {
+		var args []interface{}
+		for i := 0; i < resultCount; i++ {
+			args = append(args, queryID, hostID, fmt.Sprintf(`{"foo": "bar%d"}`, i))
+		}
+		values := strings.TrimSuffix(strings.Repeat("(?, ?, ?, NOW()),", resultCount), ",")
+		_, err := db.Exec(fmt.Sprintf(`INSERT INTO query_results (query_id, host_id, data, last_fetched) VALUES %s;`, values),
+			args...,
+		)
+		require.NoError(t, err)
+	}
+
+	team1ID := newTeam("team1")
+	team2ID := newTeam("team2")
+	host1GlobalID := newHost(nil)
+	host2Team1ID := newHost(&team1ID)
+	host3Team2ID := newHost(&team2ID)
+	query1GlobalID := newQuery("query1Global", nil)
+	query2Team1ID := newQuery("query2Team1", &team1ID)
+	query3Team2ID := newQuery("query3Team2", &team2ID)
+
+	newQueryResults(query1GlobalID, host1GlobalID, 1)
+	newQueryResults(query1GlobalID, host2Team1ID, 2)
+	newQueryResults(query1GlobalID, host3Team2ID, 3)
+
+	newQueryResults(query2Team1ID, host1GlobalID, 4)
+	newQueryResults(query2Team1ID, host2Team1ID, 5)
+	newQueryResults(query2Team1ID, host3Team2ID, 6)
+
+	newQueryResults(query3Team2ID, host1GlobalID, 7)
+	newQueryResults(query3Team2ID, host2Team1ID, 8)
+	newQueryResults(query3Team2ID, host3Team2ID, 9)
+
+	// Apply current migration.
+	applyNext(t, db)
+
+	getQueryResultsCount := func(queryID, hostID uint) int {
+		var count int
+		err := db.Get(&count, `SELECT COUNT(*) FROM query_results WHERE query_id = ? AND host_id = ?`, queryID, hostID)
+		require.NoError(t, err)
+		return count
+	}
+
+	count := getQueryResultsCount(query1GlobalID, host1GlobalID)
+	require.Equal(t, 1, count) // result for global queries are not deleted.
+	count = getQueryResultsCount(query1GlobalID, host2Team1ID)
+	require.Equal(t, 2, count) // result for global queries are not deleted.
+	count = getQueryResultsCount(query1GlobalID, host3Team2ID)
+	require.Equal(t, 3, count) // result for global queries are not deleted.
+
+	count = getQueryResultsCount(query2Team1ID, host1GlobalID)
+	require.Equal(t, 0, count) // query results of a team query different than the host's team are deleted.
+	count = getQueryResultsCount(query2Team1ID, host2Team1ID)
+	require.Equal(t, 5, count) // team query results of the host's team are not deleted.
+	count = getQueryResultsCount(query2Team1ID, host3Team2ID)
+	require.Equal(t, 0, count) // query results of a team query different than the host's team are deleted.
+
+	count = getQueryResultsCount(query3Team2ID, host1GlobalID)
+	require.Equal(t, 0, count) // query results of a team query different than the host's team are deleted.
+	count = getQueryResultsCount(query3Team2ID, host2Team1ID)
+	require.Equal(t, 0, count) // query results of a team query different than the host's team are deleted.
+	count = getQueryResultsCount(query3Team2ID, host3Team2ID)
+	require.Equal(t, 9, count) // team query results of the host's team are not deleted.
+}

server/datastore/mysql/policies.go

@@ -841,6 +841,22 @@ func cleanupPolicyMembershipOnTeamChange(ctx context.Context, tx sqlx.ExtContext
 	return nil
 }
 
+func cleanupQueryResultsOnTeamChange(ctx context.Context, tx sqlx.ExtContext, hostIDs []uint) error {
+	// Similar to cleanupPolicyMembershipOnTeamChange, hosts can belong to one team only, so we just delete all
+	// the query results of the hosts that belong to queries that are not global.
+	const cleanupQuery = `
+		DELETE FROM query_results
+		WHERE query_id IN (SELECT id FROM queries WHERE team_id IS NOT NULL) AND host_id IN (?)`
+	query, args, err := sqlx.In(cleanupQuery, hostIDs)
+	if err != nil {
+		return ctxerr.Wrap(ctx, err, "build cleanup query results query")
+	}
+	if _, err := tx.ExecContext(ctx, query, args...); err != nil {
+		return ctxerr.Wrap(ctx, err, "exec cleanup query results query")
+	}
+	return nil
+}
+
 func cleanupPolicyMembershipOnPolicyUpdate(ctx context.Context, db sqlx.ExecerContext, policyID uint, platforms string) error {
 	if platforms == "" {
 		// all platforms allowed, nothing to clean up