Claim prebuild by preset ID

dannykopping · dannykopping · commit c3f67f8169ae · 2025-02-14T13:38:17.000Z
Signed-off-by: Danny Kopping &lt;danny@coder.com&gt;
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/prebuilds.sql b/coderd/database/queries/prebuilds.sql
@@ -105,6 +105,7 @@ WHERE w.id IN (SELECT p.id
 			   WHERE (b.transition = 'start'::workspace_transition
 				   AND pj.job_status IN ('succeeded'::provisioner_job_status))
 				 AND b.template_version_id = t.active_version_id
+				 AND b.template_version_preset_id = @preset_id::uuid
 			   ORDER BY random()
 			   LIMIT 1 FOR UPDATE OF p SKIP LOCKED)
 RETURNING w.id, w.name;
diff --git a/coderd/prebuilds/claim.go b/coderd/prebuilds/claim.go
@@ -9,7 +9,7 @@ import (
 	"golang.org/x/xerrors"
 )
 
-func Claim(ctx context.Context, store database.Store, userID uuid.UUID, name string) (*uuid.UUID, error) {
+func Claim(ctx context.Context, store database.Store, userID uuid.UUID, name string, presetID uuid.UUID) (*uuid.UUID, error) {
 	var prebuildID *uuid.UUID
 	err := store.InTx(func(db database.Store) error {
 		// TODO: do we need this?
@@ -22,6 +22,7 @@ func Claim(ctx context.Context, store database.Store, userID uuid.UUID, name str
 		result, err := db.ClaimPrebuild(ctx, database.ClaimPrebuildParams{
 			NewUserID: userID,
 			NewName:   name,
+			PresetID:  presetID,
 		})
 		if err != nil {
 			switch {
diff --git a/coderd/presets.go b/coderd/presets.go
@@ -45,6 +45,10 @@ func (api *API) templateVersionPresets(rw http.ResponseWriter, r *http.Request)
 			Name: preset.Name,
 		}
 		for _, presetParam := range presetParams {
+			if presetParam.TemplateVersionPresetID != preset.ID {
+				continue
+			}
+
 			sdkPreset.Parameters = append(sdkPreset.Parameters, codersdk.PresetParameter{
 				Name:  presetParam.Name,
 				Value: presetParam.Value,
diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -633,12 +633,18 @@ func createWorkspace(
 		runningWorkspaceAgentID uuid.UUID
 	)
 	err = api.Database.InTx(func(db database.Store) error {
-		var workspaceID uuid.UUID
+		var (
+			workspaceID      uuid.UUID
+			claimedWorkspace *database.Workspace
+		)
 
-		// Try and claim an eligible prebuild, if available.
-		claimedWorkspace, err := claimPrebuild(ctx, db, api.Logger, req, owner)
-		if err != nil {
-			return xerrors.Errorf("claim prebuild: %w", err)
+		// If a template preset was chosen, try claim a prebuild.
+		if req.TemplateVersionPresetID != uuid.Nil {
+			// Try and claim an eligible prebuild, if available.
+			claimedWorkspace, err = claimPrebuild(ctx, db, api.Logger, req, owner)
+			if err != nil {
+				return xerrors.Errorf("claim prebuild: %w", err)
+			}
 		}
 
 		// No prebuild found; regular flow.
@@ -813,9 +819,7 @@ func claimPrebuild(ctx context.Context, db database.Store, logger slog.Logger, r
 	claimCtx, cancel := context.WithTimeout(ownerCtx, time.Second*10) // TODO: don't use elevated authz context
 	defer cancel()
 
-	// TODO: implement matching logic
-	// TODO: pass down rich params for matching
-	claimedID, err := prebuilds.Claim(claimCtx, db, owner.ID, req.Name)
+	claimedID, err := prebuilds.Claim(claimCtx, db, owner.ID, req.Name, req.TemplateVersionPresetID)
 	if err != nil {
 		// TODO: enhance this by clarifying whether this *specific* prebuild failed or whether there are none to claim.
 		return nil, xerrors.Errorf("claim prebuild: %w", err)
