Skip to content

Commit a25d856

Browse files
deansheatherdeansheather
authored
chore: add usage tracking package (#19095)
Not used in coderd yet, see stack. Adds two new packages: - `coderd/usage`: provides an interface for the "Collector" as well as a stub implementation for AGPL - `enterprise/coderd/usage`: provides an interface for the "Publisher" as well as a Tallyman implementation Relates to coder/internal#814
1 parent e92af2b commit a25d856

36 files changed

+2069
-17
lines changed

‎CODEOWNERS

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ tailnet/proto/ @spikecurtis @johnstcn
77
vpn/vpn.proto @spikecurtis @johnstcn
88
vpn/version.go @spikecurtis @johnstcn
99

10-
1110
# This caching code is particularly tricky, and one must be very careful when
1211
# altering it.
1312
coderd/files/ @aslilac
@@ -34,3 +33,8 @@ site/CLAUDE.md
3433
# requires elite ball knowledge of most of the scheduling code to make changes
3534
# without inadvertently affecting other parts of the codebase.
3635
coderd/schedule/autostop.go @deansheather @DanielleMaywood
36+
37+
# Usage tracking code requires intimate knowledge of Tallyman and Metronome, as
38+
# well as guidance from revenue.
39+
coderd/usage/ @deansheather @spikecurtis
40+
enterprise/coderd/usage/ @deansheather @spikecurtis

‎coderd/apidoc/docs.go

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/check_constraint.go

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -509,6 +509,25 @@ var (
509509
}),
510510
Scope: rbac.ScopeAll,
511511
}.WithCachedASTValue()
512+
513+
subjectUsageTracker = rbac.Subject{
514+
Type: rbac.SubjectTypeUsageTracker,
515+
FriendlyName: "Usage Tracker",
516+
ID: uuid.Nil.String(),
517+
Roles: rbac.Roles([]rbac.Role{
518+
{
519+
Identifier: rbac.RoleIdentifier{Name: "usage-tracker"},
520+
DisplayName: "Usage Tracker",
521+
Site: rbac.Permissions(map[string][]policy.Action{
522+
rbac.ResourceLicense.Type: {policy.ActionRead},
523+
rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
524+
}),
525+
Org: map[string][]rbac.Permission{},
526+
User: []rbac.Permission{},
527+
},
528+
}),
529+
Scope: rbac.ScopeAll,
530+
}.WithCachedASTValue()
512531
)
513532

514533
// AsProvisionerd returns a context with an actor that has permissions required
@@ -579,10 +598,18 @@ func AsPrebuildsOrchestrator(ctx context.Context) context.Context {
579598
return As(ctx, subjectPrebuildsOrchestrator)
580599
}
581600

601+
// AsFileReader returns a context with an actor that has permissions required
602+
// for reading all files.
582603
func AsFileReader(ctx context.Context) context.Context {
583604
return As(ctx, subjectFileReader)
584605
}
585606

607+
// AsUsageTracker returns a context with an actor that has permissions required
608+
// for creating, reading, and updating usage events.
609+
func AsUsageTracker(ctx context.Context) context.Context {
610+
return As(ctx, subjectUsageTracker)
611+
}
612+
586613
var AsRemoveActor = rbac.Subject{
587614
ID: "remove-actor",
588615
}
@@ -3951,6 +3978,13 @@ func (q *querier) InsertTemplateVersionWorkspaceTag(ctx context.Context, arg dat
39513978
return q.db.InsertTemplateVersionWorkspaceTag(ctx, arg)
39523979
}
39533980

3981+
func (q *querier) InsertUsageEvent(ctx context.Context, arg database.InsertUsageEventParams) error {
3982+
if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceUsageEvent); err != nil {
3983+
return err
3984+
}
3985+
return q.db.InsertUsageEvent(ctx, arg)
3986+
}
3987+
39543988
func (q *querier) InsertUser(ctx context.Context, arg database.InsertUserParams) (database.User, error) {
39553989
// Always check if the assigned roles can actually be assigned by this actor.
39563990
impliedRoles := append([]rbac.RoleIdentifier{rbac.RoleMember()}, q.convertToDeploymentRoles(arg.RBACRoles)...)
@@ -4306,6 +4340,14 @@ func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string)
43064340
return q.db.RevokeDBCryptKey(ctx, activeKeyDigest)
43074341
}
43084342

4343+
func (q *querier) SelectUsageEventsForPublishing(ctx context.Context, arg time.Time) ([]database.UsageEvent, error) {
4344+
// ActionUpdate because we're updating the publish_started_at column.
4345+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceUsageEvent); err != nil {
4346+
return nil, err
4347+
}
4348+
return q.db.SelectUsageEventsForPublishing(ctx, arg)
4349+
}
4350+
43094351
func (q *querier) TryAcquireLock(ctx context.Context, id int64) (bool, error) {
43104352
return q.db.TryAcquireLock(ctx, id)
43114353
}
@@ -4787,6 +4829,13 @@ func (q *querier) UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, arg da
47874829
return fetchAndExec(q.log, q.auth, policy.ActionUpdate, fetch, q.db.UpdateTemplateWorkspacesLastUsedAt)(ctx, arg)
47884830
}
47894831

4832+
func (q *querier) UpdateUsageEventsPostPublish(ctx context.Context, arg database.UpdateUsageEventsPostPublishParams) error {
4833+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceUsageEvent); err != nil {
4834+
return err
4835+
}
4836+
return q.db.UpdateUsageEventsPostPublish(ctx, arg)
4837+
}
4838+
47904839
func (q *querier) UpdateUserDeletedByID(ctx context.Context, id uuid.UUID) error {
47914840
return deleteQ(q.log, q.auth, q.db.GetUserByID, q.db.UpdateUserDeletedByID)(ctx, id)
47924841
}

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5666,3 +5666,34 @@ func (s *MethodTestSuite) TestUserSecrets() {
56665666
Asserts(userSecret, policy.ActionRead, userSecret, policy.ActionDelete)
56675667
}))
56685668
}
5669+
5670+
func (s *MethodTestSuite) TestUsageEvents() {
5671+
s.Run("InsertUsageEvent", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
5672+
params := database.InsertUsageEventParams{
5673+
ID: "1",
5674+
EventType: "dc_managed_agents_v1",
5675+
EventData: []byte("{}"),
5676+
CreatedAt: dbtime.Now(),
5677+
}
5678+
db.EXPECT().InsertUsageEvent(gomock.Any(), params).Return(nil)
5679+
check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionCreate)
5680+
}))
5681+
5682+
s.Run("SelectUsageEventsForPublishing", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
5683+
now := dbtime.Now()
5684+
db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), now).Return([]database.UsageEvent{}, nil)
5685+
check.Args(now).Asserts(rbac.ResourceUsageEvent, policy.ActionUpdate)
5686+
}))
5687+
5688+
s.Run("UpdateUsageEventsPostPublish", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
5689+
now := dbtime.Now()
5690+
params := database.UpdateUsageEventsPostPublishParams{
5691+
Now: now,
5692+
IDs: []string{"1", "2"},
5693+
FailureMessages: []string{"error", "error"},
5694+
SetPublishedAts: []bool{false, false},
5695+
}
5696+
db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), params).Return(nil)
5697+
check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionUpdate)
5698+
}))
5699+
}

‎coderd/database/dbmetrics/querymetrics.go

Lines changed: 21 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go

Lines changed: 43 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dump.sql

Lines changed: 30 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000359_create_usage_events_table.down.sql

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
DROP TABLE usage_events;

0 commit comments

Comments
 (0)