coder · mafredri · Mar 18, 2024 · Mar 19, 2024 · Mar 18, 2024 · Mar 20, 2024
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -47,6 +47,7 @@ import (
 	"github.com/coder/coder/v2/coderd/batchstats"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/database/dbrollup"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/externalauth"
@@ -180,6 +181,7 @@ type Options struct {
 
 	UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
 	StatsBatcher       *batchstats.Batcher
+	DBRollupInterval   time.Duration
 
 	WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
 
@@ -342,6 +344,9 @@ func New(options *Options) *API {
 	if options.StatsBatcher == nil {
 		panic("developer error: options.StatsBatcher is nil")
 	}
+	if options.DBRollupInterval == 0 {
+		options.DBRollupInterval = dbrollup.DefaultInterval
+	}
 
 	siteCacheDir := options.CacheDir
 	if siteCacheDir != "" {
@@ -414,7 +419,13 @@ func New(options *Options) *API {
 			ctx,
 			options.Logger.Named("acquirer"),
 			options.Database,
-			options.Pubsub),
+			options.Pubsub,
+		),
+		rolluper: dbrollup.New(
+			options.Logger,
+			options.Database,
+			options.DBRollupInterval,
+		),
 		workspaceUsageTracker: options.WorkspaceUsageTracker,
 	}
 
@@ -1190,7 +1201,9 @@ type API struct {
 	statsBatcher *batchstats.Batcher
 
 	Acquirer *provisionerdserver.Acquirer
-
+	// rolluper rolls up template usage stats from raw agent and app
+	// stats. This is used to provide insights in the WebUI.
+	rolluper              *dbrollup.Rolluper
 	workspaceUsageTracker *workspaceusage.Tracker
 }
 
@@ -1203,6 +1216,7 @@ func (api *API) Close() error {
 	api.WebsocketWaitGroup.Wait()
 	api.WebsocketWaitMutex.Unlock()
 
+	api.rolluper.Close()
 	api.metricsCache.Close()
 	if api.updateChecker != nil {
 		api.updateChecker.Close()

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -147,6 +147,7 @@ type Options struct {
 	WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
 	AllowWorkspaceRenames              bool
 	NewTicker                          func(duration time.Duration) (<-chan time.Time, func())
+	DBRollupInterval                   time.Duration
 	WorkspaceUsageTrackerFlush         chan int
 	WorkspaceUsageTrackerTick          chan time.Time
 }
@@ -487,6 +488,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			WorkspaceAppsStatsCollectorOptions: options.WorkspaceAppsStatsCollectorOptions,
 			AllowWorkspaceRenames:              options.AllowWorkspaceRenames,
 			NewTicker:                          options.NewTicker,
+			DBRollupInterval:                   options.DBRollupInterval,
 			WorkspaceUsageTracker:              wuTracker,
 		}
 }

diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -1617,6 +1617,29 @@ func (q *querier) GetTemplateParameterInsights(ctx context.Context, arg database
 	return q.db.GetTemplateParameterInsights(ctx, arg)
 }
 
+func (q *querier) GetTemplateUsageStats(ctx context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) {
+	// Used by dbrollup tests, use same safe-guard as other insights endpoints.
+	// For auditors, check read template_insights, and fall back to update template.
+	if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceTemplateInsights); err != nil {
+		for _, templateID := range arg.TemplateIDs {
+			template, err := q.db.GetTemplateByID(ctx, templateID)
+			if err != nil {
+				return nil, err
+			}
+
+			if err := q.authorizeContext(ctx, rbac.ActionUpdate, template); err != nil {
+				return nil, err
+			}
+		}
+		if len(arg.TemplateIDs) == 0 {
+			if err := q.authorizeContext(ctx, rbac.ActionUpdate, rbac.ResourceTemplate.All()); err != nil {
+				return nil, err
+			}
+		}
+	}
+	return q.db.GetTemplateUsageStats(ctx, arg)
+}
+
 func (q *querier) GetTemplateVersionByID(ctx context.Context, tvid uuid.UUID) (database.TemplateVersion, error) {
 	tv, err := q.db.GetTemplateVersionByID(ctx, tvid)
 	if err != nil {
@@ -3413,6 +3436,13 @@ func (q *querier) UpsertTailnetTunnel(ctx context.Context, arg database.UpsertTa
 	return q.db.UpsertTailnetTunnel(ctx, arg)
 }
 
+func (q *querier) UpsertTemplateUsageStats(ctx context.Context) error {
+	if err := q.authorizeContext(ctx, rbac.ActionUpdate, rbac.ResourceSystem); err != nil {
+		return err
+	}
+	return q.db.UpsertTemplateUsageStats(ctx)
+}
+
 func (q *querier) UpsertWorkspaceAgentPortShare(ctx context.Context, arg database.UpsertWorkspaceAgentPortShareParams) (database.WorkspaceAgentPortShare, error) {
 	workspace, err := q.db.GetWorkspaceByID(ctx, arg.WorkspaceID)
 	if err != nil {

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -956,6 +956,12 @@ func (s *MethodTestSuite) TestTemplate() {
 	s.Run("GetTemplateAppInsightsByTemplate", s.Subtest(func(db database.Store, check *expects) {
 		check.Args(database.GetTemplateAppInsightsByTemplateParams{}).Asserts(rbac.ResourceTemplateInsights, rbac.ActionRead)
 	}))
+	s.Run("GetTemplateUsageStats", s.Subtest(func(db database.Store, check *expects) {
+		check.Args(database.GetTemplateUsageStatsParams{}).Asserts(rbac.ResourceTemplateInsights, rbac.ActionRead).Errors(sql.ErrNoRows)
+	}))
+	s.Run("UpsertTemplateUsageStats", s.Subtest(func(db database.Store, check *expects) {
+		check.Asserts(rbac.ResourceSystem, rbac.ActionUpdate)
+	}))
 }
 
 func (s *MethodTestSuite) TestUser() {

diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -489,6 +489,38 @@ func WorkspaceApp(t testing.TB, db database.Store, orig database.WorkspaceApp) d
 	return resource
 }
 
+func WorkspaceAppStat(t testing.TB, db database.Store, orig database.WorkspaceAppStat) database.WorkspaceAppStat {
+	// This is not going to be correct, but our query doesn't return the ID.
+	id, err := cryptorand.Int63()
+	require.NoError(t, err, "generate id")
+
+	scheme := database.WorkspaceAppStat{
+		ID:               takeFirst(orig.ID, id),
+		UserID:           takeFirst(orig.UserID, uuid.New()),
+		WorkspaceID:      takeFirst(orig.WorkspaceID, uuid.New()),
+		AgentID:          takeFirst(orig.AgentID, uuid.New()),
+		AccessMethod:     takeFirst(orig.AccessMethod, ""),
+		SlugOrPort:       takeFirst(orig.SlugOrPort, ""),
+		SessionID:        takeFirst(orig.SessionID, uuid.New()),
+		SessionStartedAt: takeFirst(orig.SessionStartedAt, dbtime.Now().Add(-time.Minute)),
+		SessionEndedAt:   takeFirst(orig.SessionEndedAt, dbtime.Now()),
+		Requests:         takeFirst(orig.Requests, 1),
+	}
+	err = db.InsertWorkspaceAppStats(genCtx, database.InsertWorkspaceAppStatsParams{
+		UserID:           []uuid.UUID{scheme.UserID},
+		WorkspaceID:      []uuid.UUID{scheme.WorkspaceID},
+		AgentID:          []uuid.UUID{scheme.AgentID},
+		AccessMethod:     []string{scheme.AccessMethod},
+		SlugOrPort:       []string{scheme.SlugOrPort},
+		SessionID:        []uuid.UUID{scheme.SessionID},
+		SessionStartedAt: []time.Time{scheme.SessionStartedAt},
+		SessionEndedAt:   []time.Time{scheme.SessionEndedAt},
+		Requests:         []int32{scheme.Requests},
+	})
+	require.NoError(t, err, "insert workspace agent stat")
+	return scheme
+}
+
 func WorkspaceResource(t testing.TB, db database.Store, orig database.WorkspaceResource) database.WorkspaceResource {
 	resource, err := db.InsertWorkspaceResource(genCtx, database.InsertWorkspaceResourceParams{
 		ID:         takeFirst(orig.ID, uuid.New()),