Remove unnecessary nested transactions in batch-set of MDM profiles (#16449)

This commit is contained in:
Martin Angers
2024-02-05 10:51:32 -05:00
committed by GitHub
parent fa89dff66f
commit 9c37db2a05
6 changed files with 284 additions and 145 deletions
+1
View File
@@ -0,0 +1 @@
* Removed unnecessary nested database transactions in batch-setting of MDM profiles.
+102 -68
View File
@@ -1166,11 +1166,23 @@ func (ds *Datastore) GetNanoMDMEnrollment(ctx context.Context, id string) (*flee
}
func (ds *Datastore) BatchSetMDMAppleProfiles(ctx context.Context, tmID *uint, profiles []*fleet.MDMAppleConfigProfile) error {
return ds.withTx(ctx, func(tx sqlx.ExtContext) error {
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
return ds.batchSetMDMAppleProfilesDB(ctx, tx, tmID, profiles)
})
}
// set this in tests to simulate an error at various stages in the
// batchSetMDMAppleProfilesDB execution: if the string starts with "insert", it
// will be in the insert/upsert stage, "delete" for deletion, "select" to load
// existing ones, "reselect" to reload existing ones after insert, and "labels"
// to simulate an error in batch setting the profile label associations.
// "inselect", "inreselect", "indelete", etc. can also be used to fail the
// sqlx.In before the corresponding statement.
//
// e.g.: testBatchSetMDMAppleProfilesErr = "insert:fail"
var testBatchSetMDMAppleProfilesErr string
// batchSetMDMAppleProfilesDB must be called from inside a transaction.
func (ds *Datastore) batchSetMDMAppleProfilesDB(
ctx context.Context,
tx sqlx.ExtContext,
@@ -1229,90 +1241,112 @@ ON DUPLICATE KEY UPDATE
incomingProfs[p.Identifier] = p
}
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var existingProfiles []*fleet.MDMAppleConfigProfile
var existingProfiles []*fleet.MDMAppleConfigProfile
if len(incomingIdents) > 0 {
// load existing profiles that match the incoming profiles by identifiers
stmt, args, err := sqlx.In(loadExistingProfiles, profTeamID, incomingIdents)
if err != nil {
return ctxerr.Wrap(ctx, err, "build query to load existing profiles")
if len(incomingIdents) > 0 {
// load existing profiles that match the incoming profiles by identifiers
stmt, args, err := sqlx.In(loadExistingProfiles, profTeamID, incomingIdents)
if err != nil || strings.HasPrefix(testBatchSetMDMAppleProfilesErr, "inselect") {
if err == nil {
err = errors.New(testBatchSetMDMAppleProfilesErr)
}
if err := sqlx.SelectContext(ctx, tx, &existingProfiles, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "load existing profiles")
return ctxerr.Wrap(ctx, err, "build query to load existing profiles")
}
if err := sqlx.SelectContext(ctx, tx, &existingProfiles, stmt, args...); err != nil || strings.HasPrefix(testBatchSetMDMAppleProfilesErr, "select") {
if err == nil {
err = errors.New(testBatchSetMDMAppleProfilesErr)
}
return ctxerr.Wrap(ctx, err, "load existing profiles")
}
}
// figure out if we need to delete any profiles
keepIdents := make([]string, 0, len(incomingIdents))
for _, p := range existingProfiles {
if newP := incomingProfs[p.Identifier]; newP != nil {
keepIdents = append(keepIdents, p.Identifier)
// figure out if we need to delete any profiles
keepIdents := make([]string, 0, len(incomingIdents))
for _, p := range existingProfiles {
if newP := incomingProfs[p.Identifier]; newP != nil {
keepIdents = append(keepIdents, p.Identifier)
}
}
// profiles that are managed and delivered by Fleet
fleetIdents := []string{}
for ident := range mobileconfig.FleetPayloadIdentifiers() {
fleetIdents = append(fleetIdents, ident)
}
var (
stmt string
args []interface{}
err error
)
// delete the obsolete profiles (all those that are not in keepIdents or delivered by Fleet)
stmt, args, err = sqlx.In(deleteProfilesNotInList, profTeamID, append(keepIdents, fleetIdents...))
if err != nil || strings.HasPrefix(testBatchSetMDMAppleProfilesErr, "indelete") {
if err == nil {
err = errors.New(testBatchSetMDMAppleProfilesErr)
}
return ctxerr.Wrap(ctx, err, "build statement to delete obsolete profiles")
}
if _, err := tx.ExecContext(ctx, stmt, args...); err != nil || strings.HasPrefix(testBatchSetMDMAppleProfilesErr, "delete") {
if err == nil {
err = errors.New(testBatchSetMDMAppleProfilesErr)
}
return ctxerr.Wrap(ctx, err, "delete obsolete profiles")
}
// insert the new profiles and the ones that have changed
for _, p := range incomingProfs {
if _, err := tx.ExecContext(ctx, insertNewOrEditedProfile, profTeamID, p.Identifier, p.Name, p.Mobileconfig); err != nil || strings.HasPrefix(testBatchSetMDMAppleProfilesErr, "insert") {
if err == nil {
err = errors.New(testBatchSetMDMAppleProfilesErr)
}
return ctxerr.Wrapf(ctx, err, "insert new/edited profile with identifier %q", p.Identifier)
}
}
// profiles that are managed and delivered by Fleet
fleetIdents := []string{}
for ident := range mobileconfig.FleetPayloadIdentifiers() {
fleetIdents = append(fleetIdents, ident)
}
var (
stmt string
args []interface{}
err error
)
// delete the obsolete profiles (all those that are not in keepIdents or delivered by Fleet)
stmt, args, err = sqlx.In(deleteProfilesNotInList, profTeamID, append(keepIdents, fleetIdents...))
if err != nil {
return ctxerr.Wrap(ctx, err, "build statement to delete obsolete profiles")
}
if _, err := tx.ExecContext(ctx, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "delete obsolete profiles")
}
// insert the new profiles and the ones that have changed
for _, p := range incomingProfs {
if _, err := tx.ExecContext(ctx, insertNewOrEditedProfile, profTeamID, p.Identifier, p.Name, p.Mobileconfig); err != nil {
return ctxerr.Wrapf(ctx, err, "insert new/edited profile with identifier %q", p.Identifier)
// build a list of labels so the associations can be batch-set all at once
// TODO: with minor changes this chunk of code could be shared
// between macOS and Windows, but at the time of this
// implementation we're under tight time constraints.
incomingLabels := []fleet.ConfigurationProfileLabel{}
if len(incomingIdents) > 0 {
var newlyInsertedProfs []*fleet.MDMAppleConfigProfile
// load current profiles (again) that match the incoming profiles by name to grab their uuids
stmt, args, err := sqlx.In(loadExistingProfiles, profTeamID, incomingIdents)
if err != nil || strings.HasPrefix(testBatchSetMDMAppleProfilesErr, "inreselect") {
if err == nil {
err = errors.New(testBatchSetMDMAppleProfilesErr)
}
return ctxerr.Wrap(ctx, err, "build query to load newly inserted profiles")
}
if err := sqlx.SelectContext(ctx, tx, &newlyInsertedProfs, stmt, args...); err != nil || strings.HasPrefix(testBatchSetMDMAppleProfilesErr, "reselect") {
if err == nil {
err = errors.New(testBatchSetMDMAppleProfilesErr)
}
return ctxerr.Wrap(ctx, err, "load newly inserted profiles")
}
// build a list of labels so the associations can be batch-set all at once
// TODO: with minor changes this chunk of code could be shared
// between macOS and Windows, but at the time of this
// implementation we're under tight time constraints.
incomingLabels := []fleet.ConfigurationProfileLabel{}
if len(incomingIdents) > 0 {
var newlyInsertedProfs []*fleet.MDMAppleConfigProfile
// load current profiles (again) that match the incoming profiles by name to grab their uuids
stmt, args, err := sqlx.In(loadExistingProfiles, profTeamID, incomingIdents)
if err != nil {
return ctxerr.Wrap(ctx, err, "build query to load newly inserted profiles")
}
if err := sqlx.SelectContext(ctx, tx, &newlyInsertedProfs, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "load newly inserted profiles")
for _, newlyInsertedProf := range newlyInsertedProfs {
incomingProf, ok := incomingProfs[newlyInsertedProf.Identifier]
if !ok {
return ctxerr.Wrapf(ctx, err, "profile %q is in the database but was not incoming", newlyInsertedProf.Identifier)
}
for _, newlyInsertedProf := range newlyInsertedProfs {
incomingProf, ok := incomingProfs[newlyInsertedProf.Identifier]
if !ok {
return ctxerr.Wrapf(ctx, err, "profile %q is in the database but was not incoming", newlyInsertedProf.Identifier)
}
for _, label := range incomingProf.Labels {
label.ProfileUUID = newlyInsertedProf.ProfileUUID
incomingLabels = append(incomingLabels, label)
}
for _, label := range incomingProf.Labels {
label.ProfileUUID = newlyInsertedProf.ProfileUUID
incomingLabels = append(incomingLabels, label)
}
}
}
// insert label associations
if err := batchSetProfileLabelAssociationsDB(ctx, tx, incomingLabels, "darwin"); err != nil {
return ctxerr.Wrap(ctx, err, "inserting apple profile label associations")
// insert label associations
if err := batchSetProfileLabelAssociationsDB(ctx, tx, incomingLabels, "darwin"); err != nil || strings.HasPrefix(testBatchSetMDMAppleProfilesErr, "labels") {
if err == nil {
err = errors.New(testBatchSetMDMAppleProfilesErr)
}
return nil
})
return ctxerr.Wrap(ctx, err, "inserting apple profile label associations")
}
return nil
}
func (ds *Datastore) BulkDeleteMDMAppleHostsConfigProfiles(ctx context.Context, profs []*fleet.MDMAppleProfilePayload) error {
+1 -1
View File
@@ -88,7 +88,7 @@ INNER JOIN hosts h ON h.uuid = mwe.host_uuid
}
func (ds *Datastore) BatchSetMDMProfiles(ctx context.Context, tmID *uint, macProfiles []*fleet.MDMAppleConfigProfile, winProfiles []*fleet.MDMWindowsConfigProfile) error {
return ds.withTx(ctx, func(tx sqlx.ExtContext) error {
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
if err := ds.batchSetMDMWindowsProfilesDB(ctx, tx, tmID, winProfiles); err != nil {
return ctxerr.Wrap(ctx, err, "batch set windows profiles")
}
+67
View File
@@ -33,6 +33,7 @@ func TestMDMShared(t *testing.T) {
{"TestBulkSetPendingMDMHostProfilesBatch3", testBulkSetPendingMDMHostProfilesBatch3},
{"TestGetHostMDMAppleProfilesExpectedForVerification", testGetHostMDMAppleProfilesExpectedForVerification},
{"TestBatchSetProfileLabelAssociations", testBatchSetProfileLabelAssociations},
{"TestBatchSetProfilesTransactionError", testBatchSetMDMProfilesTransactionError},
}
for _, c := range cases {
@@ -3021,3 +3022,69 @@ func testBatchSetProfileLabelAssociations(t *testing.T, ds *Datastore) {
require.Error(t, err)
})
}
func testBatchSetMDMProfilesTransactionError(t *testing.T, ds *Datastore) {
ctx := context.Background()
lbl, err := ds.NewLabel(ctx, &fleet.Label{Name: "label", Query: "select 1"})
require.NoError(t, err)
cases := []struct {
windowsErr string
appleErr string
wantErr string
}{
{"select:a", "", "batch set windows profiles: load existing profiles: select:a"},
{"insert:b", "", ": insert:b"},
{"delete:c", "", "batch set windows profiles: delete obsolete profiles: delete:c"},
{"reselect:d", "", "batch set windows profiles: load newly inserted profiles: reselect:d"},
{"labels:e", "", "batch set windows profiles: inserting windows profile label associations: labels:e"},
{"inselect:k", "", "batch set windows profiles: build query to load existing profiles: inselect:k"},
{"indelete:l", "", "batch set windows profiles: build statement to delete obsolete profiles: indelete:l"},
{"inreselect:m", "", "batch set windows profiles: build query to load newly inserted profiles: inreselect:m"},
{"", "select:f", "batch set apple profiles: load existing profiles: select:f"},
{"", "insert:g", ": insert:g"},
{"", "delete:h", "batch set apple profiles: delete obsolete profiles: delete:h"},
{"", "reselect:i", "batch set apple profiles: load newly inserted profiles: reselect:i"},
{"", "labels:j", "batch set apple profiles: inserting apple profile label associations: labels:j"},
{"", "inselect:n", "batch set apple profiles: build query to load existing profiles: inselect:n"},
{"", "indelete:o", "batch set apple profiles: build statement to delete obsolete profiles: indelete:o"},
{"", "inreselect:p", "batch set apple profiles: build query to load newly inserted profiles: inreselect:p"},
}
for _, c := range cases {
t.Run(c.windowsErr+" "+c.appleErr, func(t *testing.T) {
t.Cleanup(func() {
testBatchSetMDMAppleProfilesErr = ""
testBatchSetMDMWindowsProfilesErr = ""
})
appleProfs := []*fleet.MDMAppleConfigProfile{
configProfileForTest(t, "N1", "I1", "a"),
configProfileForTest(t, "N2", "I2", "b"),
}
winProfs := []*fleet.MDMWindowsConfigProfile{
windowsConfigProfileForTest(t, "W1", "l1"),
windowsConfigProfileForTest(t, "W2", "l2"),
}
// set the initial profiles without error
err := ds.BatchSetMDMProfiles(ctx, nil, appleProfs, winProfs)
require.NoError(t, err)
// now ensure all steps are required (add a profile, delete a profile, set labels)
appleProfs = []*fleet.MDMAppleConfigProfile{
configProfileForTest(t, "N1", "I1", "aa"),
configProfileForTest(t, "N3", "I3", "c", lbl),
}
winProfs = []*fleet.MDMWindowsConfigProfile{
windowsConfigProfileForTest(t, "W1", "l11"),
windowsConfigProfileForTest(t, "W3", "l3", lbl),
}
// setup the expected errors
testBatchSetMDMAppleProfilesErr = c.appleErr
testBatchSetMDMWindowsProfilesErr = c.windowsErr
err = ds.BatchSetMDMProfiles(ctx, nil, appleProfs, winProfs)
require.ErrorContains(t, err, c.wantErr)
})
}
}
+112 -75
View File
@@ -1561,6 +1561,18 @@ ON DUPLICATE KEY UPDATE
return nil
}
// set this in tests to simulate an error at various stages in the
// batchSetMDMWindowsProfilesDB execution: if the string starts with "insert",
// it will be in the insert/upsert stage, "delete" for deletion, "select" to
// load existing ones, "reselect" to reload existing ones after insert, and
// "labels" to simulate an error in batch setting the profile label
// associations. "inselect", "inreselect", "indelete", etc. can also be used to
// fail the sqlx.In before the corresponding statement.
//
// e.g.: testBatchSetMDMWindowsProfilesErr = "insert:fail"
var testBatchSetMDMWindowsProfilesErr string
// batchSetMDMWindowsProfilesDB must be called from inside a transaction.
func (ds *Datastore) batchSetMDMWindowsProfilesDB(
ctx context.Context,
tx sqlx.ExtContext,
@@ -1625,91 +1637,116 @@ ON DUPLICATE KEY UPDATE
incomingProfs[p.Name] = p
}
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var existingProfiles []*fleet.MDMWindowsConfigProfile
var existingProfiles []*fleet.MDMWindowsConfigProfile
if len(incomingNames) > 0 {
// load existing profiles that match the incoming profiles by name
stmt, args, err := sqlx.In(loadExistingProfiles, profTeamID, incomingNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build query to load existing profiles")
if len(incomingNames) > 0 {
// load existing profiles that match the incoming profiles by name
stmt, args, err := sqlx.In(loadExistingProfiles, profTeamID, incomingNames)
if err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "inselect") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
if err := sqlx.SelectContext(ctx, tx, &existingProfiles, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "load existing profiles")
return ctxerr.Wrap(ctx, err, "build query to load existing profiles")
}
if err := sqlx.SelectContext(ctx, tx, &existingProfiles, stmt, args...); err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "select") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
return ctxerr.Wrap(ctx, err, "load existing profiles")
}
}
// figure out if we need to delete any profiles
keepNames := make([]string, 0, len(incomingNames))
for _, p := range existingProfiles {
if newP := incomingProfs[p.Name]; newP != nil {
keepNames = append(keepNames, p.Name)
}
}
var (
stmt string
args []interface{}
err error
)
// delete the obsolete profiles (all those that are not in keepNames)
if len(keepNames) > 0 {
stmt, args, err = sqlx.In(deleteProfilesNotInList, profTeamID, keepNames)
if err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "indelete") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
return ctxerr.Wrap(ctx, err, "build statement to delete obsolete profiles")
}
if _, err := tx.ExecContext(ctx, stmt, args...); err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "delete") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
return ctxerr.Wrap(ctx, err, "delete obsolete profiles")
}
} else {
if _, err := tx.ExecContext(ctx, deleteAllProfilesForTeam, profTeamID); err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "delete") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
return ctxerr.Wrap(ctx, err, "delete all profiles for team")
}
}
// insert the new profiles and the ones that have changed
for _, p := range incomingProfs {
if _, err := tx.ExecContext(ctx, insertNewOrEditedProfile, profTeamID, p.Name, p.SyncML); err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "insert") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
return ctxerr.Wrapf(ctx, err, "insert new/edited profile with name %q", p.Name)
}
}
// build a list of labels so the associations can be batch-set all at once
// TODO: with minor changes this chunk of code could be shared
// between macOS and Windows, but at the time of this
// implementation we're under tight time constraints.
incomingLabels := []fleet.ConfigurationProfileLabel{}
if len(incomingNames) > 0 {
var newlyInsertedProfs []*fleet.MDMWindowsConfigProfile
// load current profiles (again) that match the incoming profiles by name to grab their uuids
stmt, args, err := sqlx.In(loadExistingProfiles, profTeamID, incomingNames)
if err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "inreselect") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
return ctxerr.Wrap(ctx, err, "build query to load newly inserted profiles")
}
if err := sqlx.SelectContext(ctx, tx, &newlyInsertedProfs, stmt, args...); err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "reselect") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
return ctxerr.Wrap(ctx, err, "load newly inserted profiles")
}
// figure out if we need to delete any profiles
keepNames := make([]string, 0, len(incomingNames))
for _, p := range existingProfiles {
if newP := incomingProfs[p.Name]; newP != nil {
keepNames = append(keepNames, p.Name)
for _, newlyInsertedProf := range newlyInsertedProfs {
incomingProf, ok := incomingProfs[newlyInsertedProf.Name]
if !ok {
return ctxerr.Wrapf(ctx, err, "profile %q is in the database but was not incoming", newlyInsertedProf.Name)
}
for _, label := range incomingProf.Labels {
label.ProfileUUID = newlyInsertedProf.ProfileUUID
incomingLabels = append(incomingLabels, label)
}
}
}
var (
stmt string
args []interface{}
err error
)
// delete the obsolete profiles (all those that are not in keepNames)
if len(keepNames) > 0 {
stmt, args, err = sqlx.In(deleteProfilesNotInList, profTeamID, keepNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build statement to delete obsolete profiles")
}
if _, err := tx.ExecContext(ctx, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "delete obsolete profiles")
}
} else {
if _, err := tx.ExecContext(ctx, deleteAllProfilesForTeam, profTeamID); err != nil {
return ctxerr.Wrap(ctx, err, "delete all profiles for team")
}
// insert/delete the label associations
if err := batchSetProfileLabelAssociationsDB(ctx, tx, incomingLabels, "windows"); err != nil || strings.HasPrefix(testBatchSetMDMWindowsProfilesErr, "labels") {
if err == nil {
err = errors.New(testBatchSetMDMWindowsProfilesErr)
}
return ctxerr.Wrap(ctx, err, "inserting windows profile label associations")
}
// insert the new profiles and the ones that have changed
for _, p := range incomingProfs {
if _, err := tx.ExecContext(ctx, insertNewOrEditedProfile, profTeamID, p.Name, p.SyncML); err != nil {
return ctxerr.Wrapf(ctx, err, "insert new/edited profile with name %q", p.Name)
}
}
// build a list of labels so the associations can be batch-set all at once
// TODO: with minor changes this chunk of code could be shared
// between macOS and Windows, but at the time of this
// implementation we're under tight time constraints.
incomingLabels := []fleet.ConfigurationProfileLabel{}
if len(incomingNames) > 0 {
var newlyInsertedProfs []*fleet.MDMWindowsConfigProfile
// load current profiles (again) that match the incoming profiles by name to grab their uuids
stmt, args, err := sqlx.In(loadExistingProfiles, profTeamID, incomingNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build query to load newly inserted profiles")
}
if err := sqlx.SelectContext(ctx, tx, &newlyInsertedProfs, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "load newly inserted profiles")
}
for _, newlyInsertedProf := range newlyInsertedProfs {
incomingProf, ok := incomingProfs[newlyInsertedProf.Name]
if !ok {
return ctxerr.Wrapf(ctx, err, "profile %q is in the database but was not incoming", newlyInsertedProf.Name)
}
for _, label := range incomingProf.Labels {
label.ProfileUUID = newlyInsertedProf.ProfileUUID
incomingLabels = append(incomingLabels, label)
}
}
}
// insert/delete the label associations
if err := batchSetProfileLabelAssociationsDB(ctx, tx, incomingLabels, "windows"); err != nil {
return ctxerr.Wrap(ctx, err, "inserting windows profile label associations")
}
return nil
})
return nil
}
func (ds *Datastore) bulkSetPendingMDMWindowsHostProfilesDB(
+1 -1
View File
@@ -2007,7 +2007,7 @@ func testBatchSetMDMWindowsProfiles(t *testing.T, ds *Datastore) {
ctx := context.Background()
applyAndExpect := func(newSet []*fleet.MDMWindowsConfigProfile, tmID *uint, want []*fleet.MDMWindowsConfigProfile) map[string]string {
err := ds.withTx(ctx, func(tx sqlx.ExtContext) error {
err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
return ds.batchSetMDMWindowsProfilesDB(ctx, tx, tmID, newSet)
})
require.NoError(t, err)