Refactoring example

after.go

func (ds *Datastore) batchSetMDMAppleDeclarations(ctx context.Context, tx sqlx.ExtContext, tmID *uint,
	incomingDeclarations []*fleet.MDMAppleDeclaration) (updatedDB bool, err error) {

	enrichedDeclarations := enrichDeclarations(incomingDeclarations)

	existingDecls, err := ds.getExistingDeclarations(ctx, tx, enrichedDeclarations, tmID)
	if err != nil {
		return false, ctxerr.Wrap(ctx, err, "load existing declarations")
	}

	// figure out which declarations we should not delete, and put those into keepNames list
	keepNames := make([]string, 0, len(existingDecls)+len(fleetmdm.ListFleetReservedMacOSDeclarationNames()))
	for _, p := range existingDecls {
		if newP := enrichedDeclarations.byName[p.Name]; newP != nil {
			keepNames = append(keepNames, p.Name)
		}
	}
	keepNames = append(keepNames, fleetmdm.ListFleetReservedMacOSDeclarationNames()...)

	deletedDeclarations, err := ds.deleteObsoleteDeclarations(ctx, tx, keepNames, tmID)
	if err != nil {
		return false, ctxerr.Wrap(ctx, err, "delete obsolete declarations")
	}

	insertedOrUpdatedDeclarations, err := ds.insertOrUpdateDeclarations(ctx, tx, enrichedDeclarations, tmID)
	if err != nil {
		return false, ctxerr.Wrap(ctx, err, "insert/update declarations")
	}

	updatedLabels, err := ds.updateDeclarationsLabelAssociations(ctx, tx, enrichedDeclarations, tmID)
	if err != nil {
		return false, ctxerr.Wrap(ctx, err, "update declaration label associations")
	}

	return deletedDeclarations || insertedOrUpdatedDeclarations || updatedLabels, nil
}

before.go

func (ds *Datastore) batchSetMDMAppleDeclarations(ctx context.Context, tx sqlx.ExtContext, tmID *uint,
	incomingDeclarations []*fleet.MDMAppleDeclaration,
) (declarations []*fleet.MDMAppleDeclaration, updatedDB bool, err error) {
	const insertStmt = `
INSERT INTO mdm_apple_declarations (
	declaration_uuid,
	identifier,
	name,
	raw_json,
	secrets_updated_at,
	uploaded_at,
	team_id
)
VALUES (
	?,?,?,?,?,CURRENT_TIMESTAMP(),?
)
ON DUPLICATE KEY UPDATE
  uploaded_at = IF(raw_json = VALUES(raw_json) AND name = VALUES(name) AND NULLIF(secrets_updated_at = VALUES(secrets_updated_at), TRUE), uploaded_at, CURRENT_TIMESTAMP()),
  secrets_updated_at = VALUES(secrets_updated_at),
  name = VALUES(name),
  identifier = VALUES(identifier),
  raw_json = VALUES(raw_json)
`

	fmtDeleteStmt := `
DELETE FROM
  mdm_apple_declarations
WHERE
  team_id = ? AND %s
`
	andNameNotInList := "name NOT IN (?)" // added to fmtDeleteStmt if needed

	const loadExistingDecls = `
SELECT
  name,
  declaration_uuid,
  raw_json
FROM
  mdm_apple_declarations
WHERE
  team_id = ? AND
  name IN (?)
`

	var declTeamID uint
	if tmID != nil {
		declTeamID = *tmID
	}

	// build a list of names for the incoming declarations, will keep the
	// existing ones if there's a match and no change
	incomingNames := make([]string, len(incomingDeclarations))
	// at the same time, index the incoming declarations keyed by name for ease
	// or processing
	incomingDecls := make(map[string]*fleet.MDMAppleDeclaration, len(incomingDeclarations))
	for i, p := range incomingDeclarations {
		incomingNames[i] = p.Name
		incomingDecls[p.Name] = p
	}

	var existingDecls []*fleet.MDMAppleDeclaration

	if len(incomingNames) > 0 {
		// load existing declarations that match the incoming declarations by names
		stmt, args, err := sqlx.In(loadExistingDecls, declTeamID, incomingNames)
		if err != nil || strings.HasPrefix(ds.testBatchSetMDMAppleProfilesErr, "inselect") { // TODO(JVE): do we need to create similar errors for testing decls?
			if err == nil {
				err = errors.New(ds.testBatchSetMDMAppleProfilesErr)
			}
			return nil, false, ctxerr.Wrap(ctx, err, "build query to load existing declarations")
		}
		if err := sqlx.SelectContext(ctx, tx, &existingDecls, stmt, args...); err != nil || strings.HasPrefix(ds.testBatchSetMDMAppleProfilesErr, "select") {
			if err == nil {
				err = errors.New(ds.testBatchSetMDMAppleProfilesErr)
			}
			return nil, false, ctxerr.Wrap(ctx, err, "load existing declarations")
		}
	}

	// figure out if we need to delete any declarations
	keepNames := make([]string, 0, len(incomingNames))
	for _, p := range existingDecls {
		if newP := incomingDecls[p.Name]; newP != nil {
			keepNames = append(keepNames, p.Name)
		}
	}
	keepNames = append(keepNames, fleetmdm.ListFleetReservedMacOSDeclarationNames()...)

	var delArgs []any
	var delStmt string
	if len(keepNames) == 0 {
		// delete all declarations for the team
		delStmt = fmt.Sprintf(fmtDeleteStmt, "TRUE")
		delArgs = []any{declTeamID}
	} else {
		// delete the obsolete declarations (all those that are not in keepNames)
		stmt, args, err := sqlx.In(fmt.Sprintf(fmtDeleteStmt, andNameNotInList), declTeamID, keepNames)
		if err != nil {
			return nil, false, ctxerr.Wrap(ctx, err, "build query to delete obsolete profiles")
		}
		delStmt = stmt
		delArgs = args
	}

	var result sql.Result
	if result, err = tx.ExecContext(ctx, delStmt, delArgs...); err != nil || strings.HasPrefix(ds.testBatchSetMDMAppleProfilesErr,
		"delete") {
		if err == nil {
			err = errors.New(ds.testBatchSetMDMAppleProfilesErr)
		}
		return nil, false, ctxerr.Wrap(ctx, err, "delete obsolete declarations")
	}
	if result != nil {
		rows, _ := result.RowsAffected()
		updatedDB = rows > 0
	}

	for _, d := range incomingDeclarations {
		declUUID := fleet.MDMAppleDeclarationUUIDPrefix + uuid.NewString()
		if result, err = tx.ExecContext(ctx, insertStmt,
			declUUID,
			d.Identifier,
			d.Name,
			d.RawJSON,
			d.SecretsUpdatedAt,
			declTeamID); err != nil || strings.HasPrefix(ds.testBatchSetMDMAppleProfilesErr, "insert") {
			if err == nil {
				err = errors.New(ds.testBatchSetMDMAppleProfilesErr)
			}
			return nil, false, ctxerr.Wrapf(ctx, err, "insert new/edited declaration with identifier %q", d.Identifier)
		}
		updatedDB = updatedDB || insertOnDuplicateDidInsertOrUpdate(result)
	}

	incomingLabels := []fleet.ConfigurationProfileLabel{}
	if len(incomingNames) > 0 {
		var newlyInsertedDecls []*fleet.MDMAppleDeclaration
		// load current declarations (again) that match the incoming declarations by name to grab their uuids
		// this is an easy way to grab the identifiers for both the existing declarations and the new ones we generated.
		//
		// TODO(roberto): if we're a bit careful, we can harvest this
		// information without this extra request in the previous DB
		// calls. Due to time constraints, I'm leaving that
		// optimization for a later iteration.
		stmt, args, err := sqlx.In(loadExistingDecls, declTeamID, incomingNames)
		if err != nil {
			return nil, false, ctxerr.Wrap(ctx, err, "build query to load newly inserted declarations")
		}
		if err := sqlx.SelectContext(ctx, tx, &newlyInsertedDecls, stmt, args...); err != nil {
			return nil, false, ctxerr.Wrap(ctx, err, "load newly inserted declarations")
		}

		for _, newlyInsertedDecl := range newlyInsertedDecls {
			incomingDecl, ok := incomingDecls[newlyInsertedDecl.Name]
			if !ok {
				return nil, false, ctxerr.Wrapf(ctx, err, "declaration %q is in the database but was not incoming", newlyInsertedDecl.Name)
			}

			for _, label := range incomingDecl.LabelsIncludeAll {
				label.ProfileUUID = newlyInsertedDecl.DeclarationUUID
				label.Exclude = false
				label.RequireAll = true
				incomingLabels = append(incomingLabels, label)
			}
			for _, label := range incomingDecl.LabelsIncludeAny {
				label.ProfileUUID = newlyInsertedDecl.DeclarationUUID
				label.Exclude = false
				label.RequireAll = false
				incomingLabels = append(incomingLabels, label)
			}
			for _, label := range incomingDecl.LabelsExcludeAny {
				label.ProfileUUID = newlyInsertedDecl.DeclarationUUID
				label.Exclude = true
				label.RequireAll = false
				incomingLabels = append(incomingLabels, label)
			}
		}
	}

	var updatedLabels bool
	if updatedLabels, err = batchSetDeclarationLabelAssociationsDB(ctx, tx,
		incomingLabels); err != nil || strings.HasPrefix(ds.testBatchSetMDMAppleProfilesErr, "labels") {
		if err == nil {
			err = errors.New(ds.testBatchSetMDMAppleProfilesErr)
		}
		return nil, false, ctxerr.Wrap(ctx, err, "inserting apple declaration label associations")
	}

	return incomingDeclarations, updatedDB || updatedLabels, nil
}