improve query

sreya · sreya · commit 75c54dbe4f67 · 2023-07-21T00:55:06.000Z
diff --git a/coderd/autobuild/lifecycle_executor.go b/coderd/autobuild/lifecycle_executor.go
@@ -184,7 +184,6 @@ func (e *Executor) runOnce(t time.Time) Stats {
 							Time:  database.Now(),
 							Valid: true,
 						},
-						LockedTtlMs: templateSchedule.LockedTTL.Milliseconds(),
 					})
 					if err != nil {
 						log.Error(e.ctx, "unable to lock workspace",
diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -4847,10 +4847,22 @@ func (q *FakeQuerier) UpdateWorkspaceLockedDeletingAt(_ context.Context, arg dat
 			workspace.LastUsedAt = database.Now()
 			workspace.DeletingAt = sql.NullTime{}
 		}
-		if arg.LockedTtlMs > 0 && !workspace.LockedAt.Time.IsZero() {
-			workspace.DeletingAt = sql.NullTime{
-				Valid: true,
-				Time:  workspace.LockedAt.Time.Add(time.Duration(arg.LockedTtlMs) * time.Millisecond),
+		if !workspace.LockedAt.Time.IsZero() {
+			var template database.TemplateTable
+			for _, t := range q.templates {
+				if t.ID == workspace.TemplateID {
+					template = t
+					break
+				}
+			}
+			if template.ID == uuid.Nil {
+				return xerrors.Errorf("unable to find workspace template")
+			}
+			if template.LockedTTL > 0 {
+				workspace.DeletingAt = sql.NullTime{
+					Valid: true,
+					Time:  workspace.LockedAt.Time.Add(time.Duration(template.LockedTTL)),
+				}
 			}
 		}
 		q.workspaces[index] = workspace
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -489,17 +489,22 @@ WHERE
 	) AND workspaces.deleted = 'false';
 
 -- name: UpdateWorkspaceLockedDeletingAt :exec
-UPDATE workspaces
+UPDATE
+	workspaces
 SET
 	locked_at = $2,
 	-- When a workspace is unlocked we want to update the last_used_at to avoid the workspace getting re-locked.
 	-- if we're locking the workspace then we leave it alone.
 	last_used_at = CASE WHEN $2::timestamptz IS NULL THEN now() at time zone 'utc' ELSE last_used_at END,
 	-- If locked_at is null (meaning unlocked) or the template-defined locked_ttl is 0 we should set
 	-- deleting_at to NULL else set it to the locked_at + locked_ttl duration.
-	deleting_at = CASE WHEN $2::timestamptz IS NULL OR @locked_ttl_ms::bigint = 0 THEN NULL ELSE $2::timestamptz + interval '1 millisecond' * @locked_ttl_ms::bigint END
+	deleting_at = CASE WHEN $2::timestamptz IS NULL OR templates.locked_ttl = 0 THEN NULL ELSE $2::timestamptz + INTERVAL '1 milliseconds' * templates.locked_ttl / 1000000 END
+FROM
+	templates
 WHERE
-	id = $1;
+	workspaces.template_id = templates.id
+AND
+	workspaces.id = $1;
 
 -- name: UpdateWorkspacesDeletingAtByTemplateID :exec
 UPDATE
diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -784,22 +784,10 @@ func (api *API) putWorkspaceLock(rw http.ResponseWriter, r *http.Request) {
 		lockedAt.Time = database.Now()
 	}
 
-	err := api.Database.InTx(func(tx database.Store) error {
-		template, err := tx.GetTemplateByID(ctx, workspace.TemplateID)
-		if err != nil {
-			return xerrors.Errorf("get template by id: %w", err)
-		}
-
-		err = tx.UpdateWorkspaceLockedDeletingAt(ctx, database.UpdateWorkspaceLockedDeletingAtParams{
-			ID:          workspace.ID,
-			LockedAt:    lockedAt,
-			LockedTtlMs: time.Duration(template.LockedTTL).Milliseconds(),
-		})
-		if err != nil {
-			return xerrors.Errorf("update workspace locked at: %w", err)
-		}
-		return nil
-	}, nil)
+	err := api.Database.UpdateWorkspaceLockedDeletingAt(ctx, database.UpdateWorkspaceLockedDeletingAtParams{
+		ID:       workspace.ID,
+		LockedAt: lockedAt,
+	})
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error updating workspace locked status.",
