feat(coderd/database): add template_usage_stats table and rollup query

coder · mafredri · Mar 18, 2024 · Mar 19, 2024 · Mar 18, 2024 · Mar 20, 2024
commit 90521174411f29d62a8e6896238c38e2c46bae98
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/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -162,6 +162,7 @@ type data struct {
 	templateVersionParameters     []database.TemplateVersionParameter
 	templateVersionVariables      []database.TemplateVersionVariable
 	templates                     []database.TemplateTable
+	templateUsageStats            []database.TemplateUsageStat
 	workspaceAgents               []database.WorkspaceAgent
 	workspaceAgentMetadata        []database.WorkspaceAgentMetadatum
 	workspaceAgentLogs            []database.WorkspaceAgentLog
@@ -3584,6 +3585,34 @@ func (q *FakeQuerier) GetTemplateParameterInsights(ctx context.Context, arg data
 	return rows, nil
 }
 
+func (q *FakeQuerier) GetTemplateUsageStats(_ context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return nil, err
+	}
+
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	var stats []database.TemplateUsageStat
+	for _, stat := range q.templateUsageStats {
+		// Exclude all chunks that don't fall exactly within the range.
+		if stat.StartTime.Before(arg.StartTime) || stat.EndTime.After(arg.EndTime) {
+			continue
+		}
+		if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, stat.TemplateID) {
+			continue
+		}
+		stats = append(stats, stat)
+	}
+
+	if len(stats) == 0 {
+		return nil, sql.ErrNoRows
+	}
+
+	return stats, nil
+}
+
 func (q *FakeQuerier) GetTemplateVersionByID(ctx context.Context, templateVersionID uuid.UUID) (database.TemplateVersion, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
@@ -7923,6 +7952,14 @@ func (*FakeQuerier) UpsertTailnetTunnel(_ context.Context, arg database.UpsertTa
 	return database.TailnetTunnel{}, ErrUnimplemented
 }
 
+func (q *FakeQuerier) UpsertTemplateUsageStats(_ context.Context) error {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	// TODO(mafredri): Implement.
+	panic("not implemented")
+}
+
 func (q *FakeQuerier) UpsertWorkspaceAgentPortShare(_ context.Context, arg database.UpsertWorkspaceAgentPortShareParams) (database.WorkspaceAgentPortShare, error) {
 	err := validateDatabaseType(arg)
 	if err != nil {

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/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000202_template_usage_stats.down.sql b/coderd/database/migrations/000202_template_usage_stats.down.sql
@@ -0,0 +1 @@
+DROP TABLE template_usage_stats;
diff --git a/coderd/database/migrations/000202_template_usage_stats.up.sql b/coderd/database/migrations/000202_template_usage_stats.up.sql
@@ -0,0 +1,36 @@
+CREATE TABLE template_usage_stats (
+  start_time timestamptz NOT NULL,
+  end_time timestamptz NOT NULL,
+  template_id uuid NOT NULL,
+  user_id uuid NOT NULL,
+  median_latency_ms real NULL,
+  usage_mins smallint NOT NULL,
+  ssh_mins smallint NOT NULL,
+  sftp_mins smallint NOT NULL,
+  reconnecting_pty_mins smallint NOT NULL,
+  vscode_mins smallint NOT NULL,
+  jetbrains_mins smallint NOT NULL,
+  app_usage_mins jsonb NULL,
+
+  PRIMARY KEY (start_time, template_id, user_id)
+);
+
+COMMENT ON TABLE template_usage_stats IS 'Records aggregated usage statistics for templates/users. All usage is rounded up to the nearest minute.';
+COMMENT ON COLUMN template_usage_stats.start_time IS 'Start time of the usage period.';
+COMMENT ON COLUMN template_usage_stats.end_time IS 'End time of the usage period.';
+COMMENT ON COLUMN template_usage_stats.template_id IS 'ID of the template being used.';
+COMMENT ON COLUMN template_usage_stats.user_id IS 'ID of the user using the template.';
+COMMENT ON COLUMN template_usage_stats.median_latency_ms IS 'Median latency the user is experiencing, in milliseconds. Null means no value was recorded.';
+COMMENT ON COLUMN template_usage_stats.usage_mins IS 'Total minutes the user has been using the template.';
+COMMENT ON COLUMN template_usage_stats.ssh_mins IS 'Total minutes the user has been using SSH.';
+COMMENT ON COLUMN template_usage_stats.sftp_mins IS 'Total minutes the user has been using SFTP.';
+COMMENT ON COLUMN template_usage_stats.reconnecting_pty_mins IS 'Total minutes the user has been using the reconnecting PTY.';
+COMMENT ON COLUMN template_usage_stats.vscode_mins IS 'Total minutes the user has been using VSCode.';
+COMMENT ON COLUMN template_usage_stats.jetbrains_mins IS 'Total minutes the user has been using JetBrains.';
+COMMENT ON COLUMN template_usage_stats.app_usage_mins IS 'Object with app names as keys and total minutes used as values. Null means no app usage was recorded.';
+
+CREATE UNIQUE INDEX ON template_usage_stats (start_time, template_id, user_id);
+CREATE INDEX ON template_usage_stats (start_time DESC);
+
+COMMENT ON INDEX template_usage_stats_start_time_template_id_user_id_idx IS 'Index for primary key.';
+COMMENT ON INDEX template_usage_stats_start_time_idx IS 'Index for querying MAX(start_time).';
diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go