refactor: add lifecycle fields to prebuild claim database query

ssncferreira · ssncferreira · commit 8653a33ed222 · 2025-08-11T12:00:04.000Z
diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -1438,9 +1438,12 @@ func UserSecret(t testing.TB, db database.Store, seed database.UserSecret) datab
 
 func ClaimPrebuild(t testing.TB, db database.Store, newUserID uuid.UUID, newName string, presetID uuid.UUID) database.ClaimPrebuiltWorkspaceRow {
 	claimedWorkspace, err := db.ClaimPrebuiltWorkspace(genCtx, database.ClaimPrebuiltWorkspaceParams{
-		NewUserID: newUserID,
-		NewName:   newName,
-		PresetID:  presetID,
+		NewUserID:         newUserID,
+		NewName:           newName,
+		PresetID:          presetID,
+		AutostartSchedule: sql.NullString{},
+		NextStartAt:       sql.NullTime{},
+		WorkspaceTtl:      sql.NullInt64{},
 	})
 	require.NoError(t, err, "claim prebuilt workspace")
 
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
@@ -3,6 +3,11 @@ UPDATE workspaces w
 SET owner_id   = @new_user_id::uuid,
 	name       = @new_name::text,
 	updated_at = NOW(),
+	-- Update autostart_schedule, next_start_at and ttl according to template and workspace-level
+	-- configurations, allowing the workspace to be managed by the lifecycle executor as expected.
+	autostart_schedule = @autostart_schedule,
+	next_start_at = @next_start_at,
+	ttl = @workspace_ttl,
 	-- Update last_used_at during claim to ensure the claimed workspace is treated as recently used.
 	-- This avoids unintended dormancy caused by prebuilds having stale usage timestamps.
 	last_used_at = NOW(),
diff --git a/coderd/prebuilds/api.go b/coderd/prebuilds/api.go
@@ -2,6 +2,7 @@ package prebuilds
 
 import (
 	"context"
+	"database/sql"
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
@@ -54,6 +55,14 @@ type StateSnapshotter interface {
 }
 
 type Claimer interface {
-	Claim(ctx context.Context, userID uuid.UUID, name string, presetID uuid.UUID) (*uuid.UUID, error)
+	Claim(
+		ctx context.Context,
+		userID uuid.UUID,
+		name string,
+		presetID uuid.UUID,
+		autostartSchedule sql.NullString,
+		nextStartAt sql.NullTime,
+		ttl sql.NullInt64,
+	) (*uuid.UUID, error)
 	Initiator() uuid.UUID
 }
diff --git a/coderd/prebuilds/noop.go b/coderd/prebuilds/noop.go
@@ -2,6 +2,7 @@ package prebuilds
 
 import (
 	"context"
+	"database/sql"
 
 	"github.com/google/uuid"
 
@@ -28,7 +29,7 @@ var DefaultReconciler ReconciliationOrchestrator = NoopReconciler{}
 
 type NoopClaimer struct{}
 
-func (NoopClaimer) Claim(context.Context, uuid.UUID, string, uuid.UUID) (*uuid.UUID, error) {
+func (NoopClaimer) Claim(context.Context, uuid.UUID, string, uuid.UUID, sql.NullString, sql.NullTime, sql.NullInt64) (*uuid.UUID, error) {
 	// Not entitled to claim prebuilds in AGPL version.
 	return nil, ErrAGPLDoesNotSupportPrebuiltWorkspaces
 }
diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -638,7 +638,11 @@ func createWorkspace(
 		// If a template preset was chosen, try claim a prebuilt workspace.
 		if req.TemplateVersionPresetID != uuid.Nil {
 			// Try and claim an eligible prebuild, if available.
-			claimedWorkspace, err = claimPrebuild(ctx, prebuildsClaimer, db, api.Logger, req, owner)
+			// On successful claim, initialize all lifecycle fields from template and workspace-level config
+			// so the newly claimed workspace is properly managed by the lifecycle executor.
+			claimedWorkspace, err = claimPrebuild(
+				ctx, prebuildsClaimer, db, api.Logger, req, owner,
+				dbAutostartSchedule, nextStartAt, dbTTL)
 			// If claiming fails with an expected error (no claimable prebuilds or AGPL does not support prebuilds),
 			// we fall back to creating a new workspace. Otherwise, propagate the unexpected error.
 			if err != nil {
@@ -690,24 +694,7 @@ func createWorkspace(
 			}
 			workspaceID = minimumWorkspace.ID
 		} else {
-			// Prebuild found! Update lifecycle related parameters
-			err = db.UpdateWorkspaceAutostart(ctx, database.UpdateWorkspaceAutostartParams{
-				ID:                claimedWorkspace.ID,
-				AutostartSchedule: dbAutostartSchedule,
-				NextStartAt:       nextStartAt,
-			})
-			if err != nil {
-				return xerrors.Errorf("update claimed workspace Autostart: %w", err)
-			}
-
-			err = db.UpdateWorkspaceTTL(ctx, database.UpdateWorkspaceTTLParams{
-				ID:  claimedWorkspace.ID,
-				Ttl: dbTTL,
-			})
-			if err != nil {
-				return xerrors.Errorf("update claimed workspace TTL: %w", err)
-			}
-
+			// Prebuild found!
 			workspaceID = claimedWorkspace.ID
 			initiatorID = prebuildsClaimer.Initiator()
 		}
@@ -890,8 +877,17 @@ func requestTemplate(ctx context.Context, rw http.ResponseWriter, req codersdk.C
 	return template, true
 }
 
-func claimPrebuild(ctx context.Context, claimer prebuilds.Claimer, db database.Store, logger slog.Logger, req codersdk.CreateWorkspaceRequest, owner workspaceOwner) (*database.Workspace, error) {
-	claimedID, err := claimer.Claim(ctx, owner.ID, req.Name, req.TemplateVersionPresetID)
+func claimPrebuild(
+	ctx context.Context,
+	claimer prebuilds.Claimer,
+	db database.Store, logger slog.Logger,
+	req codersdk.CreateWorkspaceRequest,
+	owner workspaceOwner,
+	autostartSchedule sql.NullString,
+	nextStartAt sql.NullTime,
+	ttl sql.NullInt64,
+) (*database.Workspace, error) {
+	claimedID, err := claimer.Claim(ctx, owner.ID, req.Name, req.TemplateVersionPresetID, autostartSchedule, nextStartAt, ttl)
 	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)
diff --git a/enterprise/coderd/prebuilds/claim.go b/enterprise/coderd/prebuilds/claim.go
@@ -27,11 +27,17 @@ func (c EnterpriseClaimer) Claim(
 	userID uuid.UUID,
 	name string,
 	presetID uuid.UUID,
+	autostartSchedule sql.NullString,
+	nextStartAt sql.NullTime,
+	ttl sql.NullInt64,
 ) (*uuid.UUID, error) {
 	result, err := c.store.ClaimPrebuiltWorkspace(ctx, database.ClaimPrebuiltWorkspaceParams{
-		NewUserID: userID,
-		NewName:   name,
-		PresetID:  presetID,
+		NewUserID:         userID,
+		NewName:           name,
+		PresetID:          presetID,
+		AutostartSchedule: autostartSchedule,
+		NextStartAt:       nextStartAt,
+		WorkspaceTtl:      ttl,
 	})
 	if err != nil {
 		switch {
