feat(coderd): return agent script timings

Add the agent script timings into the `/workspaces/:workspaceId/timings` response.
coder · BrunoQuaresma · Oct 14, 2024 · Oct 1, 2024 · Oct 2, 2024 · Oct 3, 2024
commit 1b62a395bd644b5685fb2eb89304748e2ad031e5
diff --git a/coderd/agentapi/scripts.go b/coderd/agentapi/scripts.go
@@ -47,7 +47,7 @@ func (s *ScriptsAPI) ScriptCompleted(ctx context.Context, req *agentproto.Worksp
 
 	//nolint:gocritic // We need permissions to write to the DB here and we are in the context of the agent.
 	ctx = dbauthz.AsProvisionerd(ctx)
-	err = s.Database.InsertWorkspaceAgentScriptTimings(ctx, database.InsertWorkspaceAgentScriptTimingsParams{
+	_, err = s.Database.InsertWorkspaceAgentScriptTimings(ctx, database.InsertWorkspaceAgentScriptTimingsParams{
 		ScriptID:  scriptID,
 		Stage:     stage,
 		Status:    status,

diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -2421,6 +2421,13 @@ func (q *querier) GetWorkspaceAgentPortShare(ctx context.Context, arg database.G
 	return q.db.GetWorkspaceAgentPortShare(ctx, arg)
 }
 
+func (q *querier) GetWorkspaceAgentScriptTimingsByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]database.GetWorkspaceAgentScriptTimingsByWorkspaceIDRow, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
+		return nil, err
+	}
+	return q.db.GetWorkspaceAgentScriptTimingsByWorkspaceID(ctx, workspaceID)
+}
+
 func (q *querier) GetWorkspaceAgentScriptsByAgentIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceAgentScript, error) {
 	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
 		return nil, err
@@ -3034,9 +3041,9 @@ func (q *querier) InsertWorkspaceAgentMetadata(ctx context.Context, arg database
 	return q.db.InsertWorkspaceAgentMetadata(ctx, arg)
 }
 
-func (q *querier) InsertWorkspaceAgentScriptTimings(ctx context.Context, arg database.InsertWorkspaceAgentScriptTimingsParams) error {
+func (q *querier) InsertWorkspaceAgentScriptTimings(ctx context.Context, arg database.InsertWorkspaceAgentScriptTimingsParams) (database.WorkspaceAgentScriptTiming, error) {
 	if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceSystem); err != nil {
-		return err
+		return database.WorkspaceAgentScriptTiming{}, err
 	}
 	return q.db.InsertWorkspaceAgentScriptTimings(ctx, arg)
 }

diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -189,6 +189,18 @@ func WorkspaceAgent(t testing.TB, db database.Store, orig database.WorkspaceAgen
 	return agt
 }
 
+func WorkspaceAgentScripts(t testing.TB, db database.Store, orig database.InsertWorkspaceAgentScriptsParams) []database.WorkspaceAgentScript {
+	scripts, err := db.InsertWorkspaceAgentScripts(genCtx, orig)
+	require.NoError(t, err, "insert workspace agent script")
+	return scripts
+}
+
+func WorkspaceAgentScriptTiming(t testing.TB, db database.Store, orig database.InsertWorkspaceAgentScriptTimingsParams) database.WorkspaceAgentScriptTiming {
+	timing, err := db.InsertWorkspaceAgentScriptTimings(genCtx, orig)
+	require.NoError(t, err, "insert workspace agent script")
+	return timing
+}
+
 func Workspace(t testing.TB, db database.Store, orig database.Workspace) database.Workspace {
 	t.Helper()
 

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -5793,6 +5793,67 @@
 	return database.WorkspaceAgentPortShare{}, sql.ErrNoRows
 }
 
+func (q *FakeQuerier) GetWorkspaceAgentScriptTimingsByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]database.GetWorkspaceAgentScriptTimingsByWorkspaceIDRow, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	build, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, workspaceID)
+	if err != nil {
+		return nil, err
-		return nil, err
+		return nil, xerrors.Errorf("failed to retrieve workspace build: %w", err)
-		return nil, err
+		return nil, xerrors.Errorf("failed to retrieve workspace build: %w", err)
+	}
+
+	resources, err := q.GetWorkspaceResourcesByJobID(ctx, build.JobID)
+	if err != nil {
+		return nil, err
+	}
+	resourceIDs := make([]uuid.UUID, 0, len(resources))
+	for _, res := range resources {
+		resourceIDs = append(resourceIDs, res.ID)
+	}
+
+	agents, err := q.GetWorkspaceAgentsByResourceIDs(ctx, resourceIDs)
+	if err != nil {
+		return nil, err
+	}
+	agentIDs := make([]uuid.UUID, 0, len(agents))
+	for _, agent := range agents {
+		agentIDs = append(agentIDs, agent.ID)
+	}
+
+	scripts, err := q.GetWorkspaceAgentScriptsByAgentIDs(ctx, agentIDs)
+	if err != nil {
+		return nil, err
+	}
+	scriptIDs := make([]uuid.UUID, 0, len(scripts))
+	for _, script := range scripts {
+		scriptIDs = append(scriptIDs, script.ID)
+	}
+
+	rows := []database.GetWorkspaceAgentScriptTimingsByWorkspaceIDRow{}
+	for _, t := range q.workspaceAgentScriptTimings {
+		if slices.Contains(scriptIDs, t.ScriptID) {
+			var script database.WorkspaceAgentScript
+			for _, s := range scripts {
+				if s.ID == t.ScriptID {
+					script = s
+					break
+				}
+			}
+
+			rows = append(rows, database.GetWorkspaceAgentScriptTimingsByWorkspaceIDRow{
+				ScriptID:    t.ScriptID,
+				StartedAt:   t.StartedAt,
+				EndedAt:     t.EndedAt,
+				ExitCode:    t.ExitCode,
+				Stage:       t.Stage,
+				Status:      t.Status,
+				DisplayName: script.DisplayName,
+			})
+		}
+	}
+	return rows, nil
+}
+
 func (q *FakeQuerier) GetWorkspaceAgentScriptsByAgentIDs(_ context.Context, ids []uuid.UUID) ([]database.WorkspaceAgentScript, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
@@ -7844,28 +7905,26 @@
 	return nil
 }
 
-func (q *FakeQuerier) InsertWorkspaceAgentScriptTimings(_ context.Context, arg database.InsertWorkspaceAgentScriptTimingsParams) error {
+func (q *FakeQuerier) InsertWorkspaceAgentScriptTimings(_ context.Context, arg database.InsertWorkspaceAgentScriptTimingsParams) (database.WorkspaceAgentScriptTiming, error) {
 	err := validateDatabaseType(arg)
 	if err != nil {
-		return err
+		return database.WorkspaceAgentScriptTiming{}, err
 	}
 
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 
-	q.workspaceAgentScriptTimings = append(q.workspaceAgentScriptTimings,
-		//nolint:gosimple // Stop the linter complaining about changing the type of `arg`.
-		database.WorkspaceAgentScriptTiming{
-			ScriptID:  arg.ScriptID,
-			StartedAt: arg.StartedAt,
-			EndedAt:   arg.EndedAt,
-			ExitCode:  arg.ExitCode,
-			Stage:     arg.Stage,
-			Status:    arg.Status,
-		},
-	)
+	timing := database.WorkspaceAgentScriptTiming{
+		ScriptID:  arg.ScriptID,
+		StartedAt: arg.StartedAt,
+		EndedAt:   arg.EndedAt,
+		ExitCode:  arg.ExitCode,
+		Stage:     arg.Stage,
+		Status:    arg.Status,
+	}
+	q.workspaceAgentScriptTimings = append(q.workspaceAgentScriptTimings, timing)
 
-	return nil
+	return timing, nil
 }
 
 func (q *FakeQuerier) InsertWorkspaceAgentScripts(_ context.Context, arg database.InsertWorkspaceAgentScriptsParams) ([]database.WorkspaceAgentScript, error) {

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go