From 359cc84c4c13f727fa32b2f4238da369481da88d Mon Sep 17 00:00:00 2001
From: Martin Angers <martin.n.angers@gmail.com>
Date: Tue, 7 Jan 2025 16:51:27 -0500
Subject: [PATCH] Support setting priority for setup experience-related scripts

---
 ee/server/service/setup_experience.go | 10 ++++++----
 server/datastore/mysql/scripts.go     | 21 ++++++++++++++++-----
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/ee/server/service/setup_experience.go b/ee/server/service/setup_experience.go
index a8c51c5f6dd7..79698a67e28d 100644
--- a/ee/server/service/setup_experience.go
+++ b/ee/server/service/setup_experience.go
@@ -224,12 +224,14 @@ func (svc *Service) SetupExperienceNextStep(ctx context.Context, hostUUID string
 				return false, ctxerr.Errorf(ctx, "setup experience script missing content id: %d", *script.SetupExperienceScriptID)
 			}
 			req := &fleet.HostScriptRequestPayload{
-				HostID:                  host.ID,
-				ScriptName:              script.Name,
-				ScriptContentID:         *script.ScriptContentID,
+				HostID:          host.ID,
+				ScriptName:      script.Name,
+				ScriptContentID: *script.ScriptContentID,
+				// because the script execution request is associated with setup experience,
+				// it will be enqueued with a higher priority and will run before other
+				// items in the queue.
 				SetupExperienceScriptID: script.SetupExperienceScriptID,
 			}
-			// TODO(mna): setup experience scripts go to the unified queue, but must be higher priority.
 			res, err := svc.ds.NewHostScriptExecutionRequest(ctx, req)
 			if err != nil {
 				return false, ctxerr.Wrap(ctx, err, "queueing setup experience script execution request")
diff --git a/server/datastore/mysql/scripts.go b/server/datastore/mysql/scripts.go
index 5c8c2c963330..aa23668b10fb 100644
--- a/server/datastore/mysql/scripts.go
+++ b/server/datastore/mysql/scripts.go
@@ -64,9 +64,12 @@ func newHostScriptExecutionRequest(ctx context.Context, tx sqlx.ExtContext, requ
 	const (
 		insStmt = `
 INSERT INTO upcoming_activities
-	(host_id, user_id, activity_type, execution_id, script_id, script_content_id, policy_id, setup_experience_script_id, payload)
+	(
+		host_id, user_id, activity_type, execution_id, script_id, script_content_id,
+		policy_id, setup_experience_script_id, priority, payload
+	)
 VALUES
-	(?, ?, 'script', ?, ?, ?, ?, ?,
+	(?, ?, 'script', ?, ?, ?, ?, ?, ?,
 		JSON_OBJECT(
 			'sync_request', ?,
 			'is_internal', ?,
@@ -88,6 +91,13 @@ WHERE
 `
 	)
 
+	var priority int
+	if request.SetupExperienceScriptID != nil {
+		// a bit naive/simplistic for now, but we'll support user-provided
+		// priorities in a future story and we can improve on how we manage those.
+		priority = 100
+	}
+
 	execID := uuid.New().String()
 	result, err := tx.ExecContext(ctx, insStmt,
 		request.HostID,
@@ -97,6 +107,7 @@ WHERE
 		request.ScriptContentID,
 		request.PolicyID,
 		request.SetupExperienceScriptID,
+		priority,
 		request.SyncRequest,
 		isInternal,
 		request.UserID,
@@ -582,10 +593,10 @@ func (ds *Datastore) deletePendingHostScriptExecutionsForPolicy(ctx context.Cont
 	}
 
 	deleteUAStmt := `
-		DELETE FROM 
+		DELETE FROM
 			upcoming_activities
-		WHERE 
-			policy_id = ? AND 
+		WHERE
+			policy_id = ? AND
 			activity_type = 'script' AND
 			script_id IN (
 				SELECT id FROM scripts WHERE scripts.global_or_team_id = ?