fix VPP migration edge case (#22460)

#22415

# Checklist for submitter

If some of the following don't apply, delete the relevant line.

<!-- Note that API documentation changes are now addressed by the
product design team. -->

- [x] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/Committing-Changes.md#changes-files)
for more information.
- [x] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements)
- [x] Added/updated tests
- [x] If database migrations are included, checked table schema to
confirm autoupdate
- For database migrations:
- [x] Checked schema for all modified table for columns that will
auto-update timestamps during migration.
- [x] Confirmed that updating the timestamps is acceptable, and will not
cause unwanted side effects.
- [x] Ensured the correct collation is explicitly set for character
columns (`COLLATE utf8mb4_unicode_ci`).
- [x] Manual QA for all new/changed functionality

changes/22415-fix-vpp-migration

@@ -0,0 +1 @@
+* Fixed an issue with the migration adding support for multiple VPP tokens that would happen if a token is removed prior to upgrading Fleet.

server/datastore/mysql/migrations/tables/20240829170033_AddVPPTokenIDToVppAppsTeams.go

@@ -3,6 +3,8 @@ package tables
 import (
 	"database/sql"
 	"fmt"
+
+	"github.com/pkg/errors"
 )
 
 func init() {
@@ -14,7 +16,11 @@ func Up_20240829170033(tx *sql.Tx) error {
 ALTER TABLE vpp_apps_teams
 	ADD COLUMN vpp_token_id int(10) UNSIGNED NOT NULL`
 
-	stmtAssociate := `UPDATE vpp_apps_teams SET vpp_token_id = (SELECT id FROM vpp_tokens LIMIT 1)`
+	stmtFindToken := `SELECT id FROM vpp_tokens LIMIT 1` //nolint:gosec
+
+	stmtCleanAssociations := `DELETE FROM vpp_apps_teams`
+
+	stmtAssociate := `UPDATE vpp_apps_teams SET vpp_token_id = ?`
 
 	stmtAddConstraint := `
 ALTER TABLE vpp_apps_teams
@@ -27,8 +33,22 @@ ALTER TABLE vpp_apps_teams
 
 	// Associate apps with the first token available. If we're
 	// migrating from single-token VPP this should be correct.
-	if _, err := tx.Exec(stmtAssociate); err != nil {
-		return fmt.Errorf("failed to associate vpp apps with first token: %w", err)
+	var vppTokenID uint
+	err := tx.QueryRow(stmtFindToken).Scan(&vppTokenID)
+	if err != nil {
+		if !errors.Is(err, sql.ErrNoRows) {
+			return fmt.Errorf("get existing VPP token ID: %w", err)
+		}
+	}
+
+	if vppTokenID > 0 {
+		if _, err := tx.Exec(stmtAssociate, vppTokenID); err != nil {
+			return fmt.Errorf("failed to associate vpp apps with first token: %w", err)
+		}
+	} else {
+		if _, err := tx.Exec(stmtCleanAssociations); err != nil {
+			return fmt.Errorf("failed clean orphaned VPP team associations: %w", err)
+		}
 	}
 
 	if _, err := tx.Exec(stmtAddConstraint); err != nil {

server/datastore/mysql/migrations/tables/20240829170033_AddVPPTokenIDToVppAppsTeams_test.go

@@ -0,0 +1,70 @@
+package tables
+
+import (
+	"testing"
+	"time"
+
+	"github.com/jmoiron/sqlx"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUp_20240829170033_Existing(t *testing.T) {
+	db := applyUpToPrev(t)
+
+	// insert a vpp token
+	vppTokenID := execNoErrLastID(t, db, "INSERT INTO vpp_tokens (organization_name, location, renew_at, token) VALUES (?, ?, ?, ?)", "org", "location", time.Now(), "token")
+
+	// create a couple teams
+	tm1 := execNoErrLastID(t, db, "INSERT INTO teams (name) VALUES ('team1')")
+	tm2 := execNoErrLastID(t, db, "INSERT INTO teams (name) VALUES ('team2')")
+
+	// create a couple of vpp apps
+	adamID1 := "123"
+	execNoErr(t, db, `INSERT INTO vpp_apps (adam_id, platform) VALUES (?, "iOS")`, adamID1)
+	adamID2 := "456"
+	execNoErr(t, db, `INSERT INTO vpp_apps (adam_id, platform) VALUES (?, "iOS")`, adamID2)
+
+	// insert some teams with vpp apps
+	execNoErr(t, db, `INSERT INTO vpp_apps_teams (adam_id, team_id, global_or_team_id, platform, self_service) VALUES (?, ?, ?, ?, ?)`, adamID1, tm1, 0, "iOS", 0)
+	execNoErr(t, db, `INSERT INTO vpp_apps_teams (adam_id, team_id, global_or_team_id, platform, self_service) VALUES (?, ?, ?, ?, ?)`, adamID2, tm2, 0, "iOS", 0)
+
+	// apply current migration
+	applyNext(t, db)
+
+	// ensure vpp_token_id is set for all teams
+	var vppTokenIDs []int
+	err := sqlx.Select(db, &vppTokenIDs, `SELECT vpp_token_id FROM vpp_apps_teams`)
+	require.NoError(t, err)
+	require.Len(t, vppTokenIDs, 2)
+	for _, tokenID := range vppTokenIDs {
+		require.Equal(t, int(vppTokenID), tokenID)
+	}
+}
+
+func TestUp_20240829170033_NoTokens(t *testing.T) {
+	db := applyUpToPrev(t)
+
+	// create a couple teams
+	tm1 := execNoErrLastID(t, db, "INSERT INTO teams (name) VALUES ('team1')")
+	tm2 := execNoErrLastID(t, db, "INSERT INTO teams (name) VALUES ('team2')")
+
+	// create a couple of vpp apps
+	adamID1 := "123"
+	execNoErr(t, db, `INSERT INTO vpp_apps (adam_id, platform) VALUES (?, "iOS")`, adamID1)
+	adamID2 := "456"
+	execNoErr(t, db, `INSERT INTO vpp_apps (adam_id, platform) VALUES (?, "iOS")`, adamID2)
+
+	// insert some teams with vpp apps
+	execNoErr(t, db, `INSERT INTO vpp_apps_teams (adam_id, team_id, global_or_team_id, platform, self_service) VALUES (?, ?, ?, ?, ?)`, adamID1, tm1, 0, "iOS", 0)
+	execNoErr(t, db, `INSERT INTO vpp_apps_teams (adam_id, team_id, global_or_team_id, platform, self_service) VALUES (?, ?, ?, ?, ?)`, adamID2, tm2, 0, "iOS", 0)
+
+	// apply current migration
+	applyNext(t, db)
+
+	// ensure no rows are left in vpp_apps_teams (since there are no tokens)
+	var count int
+	err := sqlx.Get(db, &count, `SELECT COUNT(*) FROM vpp_apps_teams`)
+	require.NoError(t, err)
+	require.Zero(t, count)
+
+}