diff --git a/cli/delete_test.go b/cli/delete_test.go index c01893419f..2e550d7484 100644 --- a/cli/delete_test.go +++ b/cli/delete_test.go @@ -111,7 +111,6 @@ func TestDelete(t *testing.T) { // The API checks if the user has any workspaces, so we cannot delete a user // this way. ctx := testutil.Context(t, testutil.WaitShort) - // nolint:gocritic // Unit test err := api.Database.UpdateUserDeletedByID(dbauthz.AsSystemRestricted(ctx), deleteMeUser.ID) require.NoError(t, err) diff --git a/cli/provisioners_test.go b/cli/provisioners_test.go index 30a89714ff..0c3fe5ae2f 100644 --- a/cli/provisioners_test.go +++ b/cli/provisioners_test.go @@ -31,7 +31,6 @@ func TestProvisioners_Golden(t *testing.T) { // Replace UUIDs with predictable values for golden files. replace := make(map[string]string) updateReplaceUUIDs := func(coderdAPI *coderd.API) { - //nolint:gocritic // This is a test. systemCtx := dbauthz.AsSystemRestricted(context.Background()) provisioners, err := coderdAPI.Database.GetProvisionerDaemons(systemCtx) require.NoError(t, err) diff --git a/coderd/agentapi/subagent_test.go b/coderd/agentapi/subagent_test.go index 0a95a70e52..1b6eef936f 100644 --- a/coderd/agentapi/subagent_test.go +++ b/coderd/agentapi/subagent_test.go @@ -163,7 +163,7 @@ func TestSubAgentAPI(t *testing.T) { agentID, err := uuid.FromBytes(createResp.Agent.Id) require.NoError(t, err) - agent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test. + agent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) require.NoError(t, err) assert.Equal(t, tt.agentName, agent.Name) @@ -621,7 +621,7 @@ func TestSubAgentAPI(t *testing.T) { agentID, err := uuid.FromBytes(createResp.Agent.Id) require.NoError(t, err) - apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test. + apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) require.NoError(t, err) // Sort the apps for determinism @@ -751,7 +751,7 @@ func TestSubAgentAPI(t *testing.T) { agentID, err := uuid.FromBytes(createResp.Agent.Id) require.NoError(t, err) - apps, err := db.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test. + apps, err := db.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) require.NoError(t, err) require.Len(t, apps, 1) require.Equal(t, "k5jd7a99-duplicate-slug", apps[0].Slug) @@ -789,7 +789,7 @@ func TestSubAgentAPI(t *testing.T) { require.NoError(t, err) // Then: It is deleted. - _, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgent.ID) //nolint:gocritic // this is a test. + _, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgent.ID) require.ErrorIs(t, err, sql.ErrNoRows) }) @@ -830,10 +830,10 @@ func TestSubAgentAPI(t *testing.T) { require.NoError(t, err) // Then: The correct one is deleted. - _, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentOne.ID) //nolint:gocritic // this is a test. + _, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentOne.ID) require.ErrorIs(t, err, sql.ErrNoRows) - _, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentTwo.ID) //nolint:gocritic // this is a test. + _, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentTwo.ID) require.NoError(t, err) }) @@ -871,7 +871,7 @@ func TestSubAgentAPI(t *testing.T) { var notAuthorizedError dbauthz.NotAuthorizedError require.ErrorAs(t, err, ¬AuthorizedError) - _, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentOne.ID) //nolint:gocritic // this is a test. + _, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentOne.ID) require.NoError(t, err) }) @@ -912,7 +912,7 @@ func TestSubAgentAPI(t *testing.T) { require.NoError(t, err) // Verify that the apps were created - apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), subAgentID) //nolint:gocritic // this is a test. + apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), subAgentID) require.NoError(t, err) require.Len(t, apps, 2) @@ -923,7 +923,7 @@ func TestSubAgentAPI(t *testing.T) { require.NoError(t, err) // Then: The agent is deleted - _, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), subAgentID) //nolint:gocritic // this is a test. + _, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), subAgentID) require.ErrorIs(t, err, sql.ErrNoRows) // And: The apps are *retained* to avoid causing issues @@ -1068,7 +1068,7 @@ func TestSubAgentAPI(t *testing.T) { agentID, err := uuid.FromBytes(createResp.Agent.Id) require.NoError(t, err) - subAgent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test. + subAgent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) require.NoError(t, err) require.Equal(t, len(tt.expectedApps), len(subAgent.DisplayApps), "display apps count mismatch") @@ -1118,14 +1118,14 @@ func TestSubAgentAPI(t *testing.T) { require.NoError(t, err) // Verify display apps - subAgent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test. + subAgent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) require.NoError(t, err) require.Len(t, subAgent.DisplayApps, 2) require.Equal(t, database.DisplayAppVscode, subAgent.DisplayApps[0]) require.Equal(t, database.DisplayAppWebTerminal, subAgent.DisplayApps[1]) // Verify regular apps - apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test. + apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) require.NoError(t, err) require.Len(t, apps, 1) require.Equal(t, "v4qhkq17-custom-app", apps[0].Slug) @@ -1190,7 +1190,7 @@ func TestSubAgentAPI(t *testing.T) { }) // When: We list the sub agents. - listResp, err := api.ListSubAgents(ctx, &proto.ListSubAgentsRequest{}) //nolint:gocritic // this is a test. + listResp, err := api.ListSubAgents(ctx, &proto.ListSubAgentsRequest{}) require.NoError(t, err) listedChildAgents := listResp.Agents diff --git a/coderd/coderd.go b/coderd/coderd.go index 8ab204f8a3..5debc13d21 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -23,6 +23,7 @@ import ( "github.com/coder/coder/v2/coderd/oauth2provider" "github.com/coder/coder/v2/coderd/pproflabel" "github.com/coder/coder/v2/coderd/prebuilds" + "github.com/coder/coder/v2/coderd/usage" "github.com/coder/coder/v2/coderd/wsbuilder" "github.com/andybalholm/brotli" @@ -200,6 +201,7 @@ type Options struct { TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore] UserQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore] AccessControlStore *atomic.Pointer[dbauthz.AccessControlStore] + UsageInserter *atomic.Pointer[usage.Inserter] // CoordinatorResumeTokenProvider is used to provide and validate resume // tokens issued by and passed to the coordinator DRPC API. CoordinatorResumeTokenProvider tailnet.ResumeTokenProvider @@ -428,6 +430,13 @@ func New(options *Options) *API { v := schedule.NewAGPLUserQuietHoursScheduleStore() options.UserQuietHoursScheduleStore.Store(&v) } + if options.UsageInserter == nil { + options.UsageInserter = &atomic.Pointer[usage.Inserter]{} + } + if options.UsageInserter.Load() == nil { + inserter := usage.NewAGPLInserter() + options.UsageInserter.Store(&inserter) + } if options.OneTimePasscodeValidityPeriod == 0 { options.OneTimePasscodeValidityPeriod = 20 * time.Minute } @@ -590,6 +599,7 @@ func New(options *Options) *API { UserQuietHoursScheduleStore: options.UserQuietHoursScheduleStore, AccessControlStore: options.AccessControlStore, BuildUsageChecker: &buildUsageChecker, + UsageInserter: options.UsageInserter, FileCache: files.New(options.PrometheusRegistry, options.Authorizer), Experiments: experiments, WebpushDispatcher: options.WebPushDispatcher, @@ -1690,6 +1700,9 @@ type API struct { // BuildUsageChecker is a pointer as it's passed around to multiple // components. BuildUsageChecker *atomic.Pointer[wsbuilder.UsageChecker] + // UsageInserter is a pointer to an atomic pointer because it is passed to + // multiple components. + UsageInserter *atomic.Pointer[usage.Inserter] UpdatesProvider tailnet.WorkspaceUpdatesProvider @@ -1905,6 +1918,7 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n &api.Auditor, api.TemplateScheduleStore, api.UserQuietHoursScheduleStore, + api.UsageInserter, api.DeploymentValues, provisionerdserver.Options{ OIDCConfig: api.OIDCConfig, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index a716c04adc..94e60db47c 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -213,6 +213,8 @@ var ( // Provisionerd creates workspaces resources monitor rbac.ResourceWorkspaceAgentResourceMonitor.Type: {policy.ActionCreate}, rbac.ResourceWorkspaceAgentDevcontainers.Type: {policy.ActionCreate}, + // Provisionerd creates usage events + rbac.ResourceUsageEvent.Type: {policy.ActionCreate}, }), Org: map[string][]rbac.Permission{}, User: []rbac.Permission{}, @@ -510,17 +512,19 @@ var ( Scope: rbac.ScopeAll, }.WithCachedASTValue() - subjectUsageTracker = rbac.Subject{ - Type: rbac.SubjectTypeUsageTracker, - FriendlyName: "Usage Tracker", + subjectUsagePublisher = rbac.Subject{ + Type: rbac.SubjectTypeUsagePublisher, + FriendlyName: "Usage Publisher", ID: uuid.Nil.String(), Roles: rbac.Roles([]rbac.Role{ { - Identifier: rbac.RoleIdentifier{Name: "usage-tracker"}, - DisplayName: "Usage Tracker", + Identifier: rbac.RoleIdentifier{Name: "usage-publisher"}, + DisplayName: "Usage Publisher", Site: rbac.Permissions(map[string][]policy.Action{ - rbac.ResourceLicense.Type: {policy.ActionRead}, - rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate}, + rbac.ResourceLicense.Type: {policy.ActionRead}, + // The usage publisher doesn't create events, just + // reads/processes them. + rbac.ResourceUsageEvent.Type: {policy.ActionRead, policy.ActionUpdate}, }), Org: map[string][]rbac.Permission{}, User: []rbac.Permission{}, @@ -604,10 +608,10 @@ func AsFileReader(ctx context.Context) context.Context { return As(ctx, subjectFileReader) } -// AsUsageTracker returns a context with an actor that has permissions required -// for creating, reading, and updating usage events. -func AsUsageTracker(ctx context.Context) context.Context { - return As(ctx, subjectUsageTracker) +// AsUsagePublisher returns a context with an actor that has permissions +// required for creating, reading, and updating usage events. +func AsUsagePublisher(ctx context.Context) context.Context { + return As(ctx, subjectUsagePublisher) } var AsRemoveActor = rbac.Subject{ @@ -3038,7 +3042,7 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe } func (q *querier) GetUnexpiredLicenses(ctx context.Context) ([]database.License, error) { - if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceLicense); err != nil { return nil, err } return q.db.GetUnexpiredLicenses(ctx) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index ce70a9b1f1..ad444d1025 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -758,6 +758,18 @@ func (s *MethodTestSuite) TestLicense() { check.Args().Asserts(l, policy.ActionRead). Returns([]database.License{l}) })) + s.Run("GetUnexpiredLicenses", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + l := database.License{ + ID: 1, + Exp: time.Now().Add(time.Hour * 24 * 30), + UUID: uuid.New(), + } + db.EXPECT().GetUnexpiredLicenses(gomock.Any()). + Return([]database.License{l}, nil). + AnyTimes() + check.Args().Asserts(rbac.ResourceLicense, policy.ActionRead). + Returns([]database.License{l}) + })) s.Run("InsertLicense", s.Subtest(func(db database.Store, check *expects) { check.Args(database.InsertLicenseParams{}). Asserts(rbac.ResourceLicense, policy.ActionCreate) @@ -3770,9 +3782,6 @@ func (s *MethodTestSuite) TestSystemFunctions() { s.Run("GetActiveUserCount", s.Subtest(func(db database.Store, check *expects) { check.Args(false).Asserts(rbac.ResourceSystem, policy.ActionRead).Returns(int64(0)) })) - s.Run("GetUnexpiredLicenses", s.Subtest(func(db database.Store, check *expects) { - check.Args().Asserts(rbac.ResourceSystem, policy.ActionRead) - })) s.Run("GetAuthorizationUserRoles", s.Subtest(func(db database.Store, check *expects) { u := dbgen.User(s.T(), db, database.User{}) check.Args(u.ID).Asserts(rbac.ResourceSystem, policy.ActionRead) diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go index 484d59beab..8e46566ed2 100644 --- a/coderd/externalauth/externalauth_test.go +++ b/coderd/externalauth/externalauth_test.go @@ -337,7 +337,6 @@ func TestRefreshToken(t *testing.T) { require.Equal(t, 1, validateCalls, "token is validated") require.Equal(t, 1, refreshCalls, "token is refreshed") require.NotEqualf(t, link.OAuthAccessToken, updated.OAuthAccessToken, "token is updated") - //nolint:gocritic // testing dbLink, err := db.GetExternalAuthLink(dbauthz.AsSystemRestricted(context.Background()), database.GetExternalAuthLinkParams{ ProviderID: link.ProviderID, UserID: link.UserID, diff --git a/coderd/files/cache_test.go b/coderd/files/cache_test.go index 6f8f74e74f..b81deae5d9 100644 --- a/coderd/files/cache_test.go +++ b/coderd/files/cache_test.go @@ -45,7 +45,6 @@ func TestCancelledFetch(t *testing.T) { cache := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) // Cancel the context for the first call; should fail. - //nolint:gocritic // Unit testing ctx, cancel := context.WithCancel(dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort))) cancel() _, err := cache.Acquire(ctx, dbM, fileID) @@ -71,7 +70,6 @@ func TestCancelledConcurrentFetch(t *testing.T) { cache := files.LeakCache{Cache: files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})} - //nolint:gocritic // Unit testing ctx := dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort)) // Cancel the context for the first call; should fail. @@ -99,7 +97,6 @@ func TestConcurrentFetch(t *testing.T) { }) cache := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{}) - //nolint:gocritic // Unit testing ctx := dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort)) // Expect 2 calls to Acquire before we continue the test @@ -151,7 +148,6 @@ func TestCacheRBAC(t *testing.T) { Scope: rbac.ScopeAll, }) - //nolint:gocritic // Unit testing cacheReader := dbauthz.AsFileReader(ctx) t.Run("NoRolesOpen", func(t *testing.T) { @@ -207,7 +203,6 @@ func cachePromMetricName(metric string) string { func TestConcurrency(t *testing.T) { t.Parallel() - //nolint:gocritic // Unit testing ctx := dbauthz.AsFileReader(t.Context()) const fileSize = 10 @@ -268,7 +263,6 @@ func TestConcurrency(t *testing.T) { func TestRelease(t *testing.T) { t.Parallel() - //nolint:gocritic // Unit testing ctx := dbauthz.AsFileReader(t.Context()) const fileSize = 10 diff --git a/coderd/idpsync/group_test.go b/coderd/idpsync/group_test.go index 478d6557de..7f4ee9f435 100644 --- a/coderd/idpsync/group_test.go +++ b/coderd/idpsync/group_test.go @@ -328,7 +328,6 @@ func TestGroupSyncTable(t *testing.T) { }, } - //nolint:gocritic // testing defOrg, err := db.GetDefaultOrganization(dbauthz.AsSystemRestricted(ctx)) require.NoError(t, err) SetupOrganization(t, s, db, user, defOrg.ID, def) @@ -527,7 +526,6 @@ func TestApplyGroupDifference(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitMedium) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) org := dbgen.Organization(t, db, database.Organization{}) diff --git a/coderd/idpsync/role_test.go b/coderd/idpsync/role_test.go index 6df091097b..db172e0ee4 100644 --- a/coderd/idpsync/role_test.go +++ b/coderd/idpsync/role_test.go @@ -273,7 +273,6 @@ func TestRoleSyncTable(t *testing.T) { } // Also assert site wide roles - //nolint:gocritic // unit testing assertions allRoles, err := db.GetAuthorizationUserRoles(dbauthz.AsSystemRestricted(ctx), user.ID) require.NoError(t, err) diff --git a/coderd/insights_test.go b/coderd/insights_test.go index d916b20fea..cf5f63065d 100644 --- a/coderd/insights_test.go +++ b/coderd/insights_test.go @@ -754,7 +754,6 @@ func TestTemplateInsights_Golden(t *testing.T) { Database: db, AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize, }) - //nolint:gocritic // This is a test. err = reporter.ReportAppStats(dbauthz.AsSystemRestricted(ctx), stats) require.NoError(t, err, "want no error inserting app stats") @@ -1646,7 +1645,6 @@ func TestUserActivityInsights_Golden(t *testing.T) { Database: db, AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize, }) - //nolint:gocritic // This is a test. err = reporter.ReportAppStats(dbauthz.AsSystemRestricted(ctx), stats) require.NoError(t, err, "want no error inserting app stats") diff --git a/coderd/notifications/manager_test.go b/coderd/notifications/manager_test.go index e9c309f0a0..30af0c88b8 100644 --- a/coderd/notifications/manager_test.go +++ b/coderd/notifications/manager_test.go @@ -31,7 +31,6 @@ func TestBufferedUpdates(t *testing.T) { // setup - // nolint:gocritic // Unit test. ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitSuperLong)) store, ps := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -108,7 +107,6 @@ func TestBuildPayload(t *testing.T) { // SETUP - // nolint:gocritic // Unit test. ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitSuperLong)) store, _ := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -166,7 +164,6 @@ func TestStopBeforeRun(t *testing.T) { // SETUP - // nolint:gocritic // Unit test. ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitSuperLong)) store, ps := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -187,7 +184,6 @@ func TestRunStopRace(t *testing.T) { // SETUP - // nolint:gocritic // Unit test. ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitMedium)) store, ps := dbtestutil.NewDB(t) logger := testutil.Logger(t) diff --git a/coderd/notifications/metrics_test.go b/coderd/notifications/metrics_test.go index 5517f86061..6ba6635a50 100644 --- a/coderd/notifications/metrics_test.go +++ b/coderd/notifications/metrics_test.go @@ -37,7 +37,6 @@ func TestMetrics(t *testing.T) { t.Skip("This test requires postgres; it relies on business-logic only implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -226,7 +225,6 @@ func TestPendingUpdatesMetric(t *testing.T) { t.Parallel() // SETUP - // nolint:gocritic // Unit test. ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -320,7 +318,6 @@ func TestInflightDispatchesMetric(t *testing.T) { t.Parallel() // SETUP - // nolint:gocritic // Unit test. ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -400,7 +397,6 @@ func TestCustomMethodMetricCollection(t *testing.T) { t.Skip("This test requires postgres; it relies on business-logic only implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) diff --git a/coderd/notifications/notifications_test.go b/coderd/notifications/notifications_test.go index e213a62df9..f5e72a8327 100644 --- a/coderd/notifications/notifications_test.go +++ b/coderd/notifications/notifications_test.go @@ -70,7 +70,6 @@ func TestBasicNotificationRoundtrip(t *testing.T) { t.Skip("This test requires postgres; it relies on business-logic only implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -137,7 +136,6 @@ func TestSMTPDispatch(t *testing.T) { // SETUP - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -203,7 +201,6 @@ func TestWebhookDispatch(t *testing.T) { // SETUP - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -287,7 +284,6 @@ func TestBackpressure(t *testing.T) { store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitShort)) const method = database.NotificationMethodWebhook @@ -416,7 +412,6 @@ func TestRetries(t *testing.T) { } const maxAttempts = 3 - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -516,7 +511,6 @@ func TestExpiredLeaseIsRequeued(t *testing.T) { t.Skip("This test requires postgres; it relies on business-logic only implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -536,7 +530,6 @@ func TestExpiredLeaseIsRequeued(t *testing.T) { noopInterceptor := newNoopStoreSyncer(store) - // nolint:gocritic // Unit test. mgrCtx, cancelManagerCtx := context.WithCancel(dbauthz.AsNotifier(context.Background())) t.Cleanup(cancelManagerCtx) @@ -645,7 +638,6 @@ func TestNotifierPaused(t *testing.T) { // Setup. - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -1323,7 +1315,6 @@ func TestNotificationTemplates_Golden(t *testing.T) { return &db, &api.Logger, &user }() - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) _, pubsub := dbtestutil.NewDB(t) @@ -1406,13 +1397,11 @@ func TestNotificationTemplates_Golden(t *testing.T) { // as appearance changes are enterprise features and we do not want to mix those // can't use the api if tc.appName != "" { - // nolint:gocritic // Unit test. err = (*db).UpsertApplicationName(dbauthz.AsSystemRestricted(ctx), "Custom Application") require.NoError(t, err) } if tc.logoURL != "" { - // nolint:gocritic // Unit test. err = (*db).UpsertLogoURL(dbauthz.AsSystemRestricted(ctx), "https://custom.application/logo.png") require.NoError(t, err) } @@ -1510,7 +1499,6 @@ func TestNotificationTemplates_Golden(t *testing.T) { }() _, pubsub := dbtestutil.NewDB(t) - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) // Spin up the mock webhook server @@ -1650,7 +1638,6 @@ func TestDisabledByDefaultBeforeEnqueue(t *testing.T) { t.Skip("This test requires postgres; it is testing business-logic implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, _ := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -1676,7 +1663,6 @@ func TestDisabledBeforeEnqueue(t *testing.T) { t.Skip("This test requires postgres; it is testing business-logic implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, _ := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -1712,7 +1698,6 @@ func TestDisabledAfterEnqueue(t *testing.T) { t.Skip("This test requires postgres; it is testing business-logic implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -1769,7 +1754,6 @@ func TestCustomNotificationMethod(t *testing.T) { t.Skip("This test requires postgres; it relies on business-logic only implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -1873,7 +1857,6 @@ func TestNotificationsTemplates(t *testing.T) { t.Skip("This test requires postgres; it relies on business-logic only implemented in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) api := coderdtest.New(t, createOpts(t)) @@ -1910,7 +1893,6 @@ func TestNotificationDuplicates(t *testing.T) { t.Skip("This test requires postgres; it is testing the dedupe hash trigger in the database") } - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -2007,7 +1989,6 @@ func TestNotificationTargetMatrix(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, pubsub := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -2051,7 +2032,6 @@ func TestNotificationOneTimePasswordDeliveryTargets(t *testing.T) { t.Run("Inbox", func(t *testing.T) { t.Parallel() - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, _ := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -2076,7 +2056,6 @@ func TestNotificationOneTimePasswordDeliveryTargets(t *testing.T) { t.Run("SMTP", func(t *testing.T) { t.Parallel() - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, _ := dbtestutil.NewDB(t) logger := testutil.Logger(t) @@ -2100,7 +2079,6 @@ func TestNotificationOneTimePasswordDeliveryTargets(t *testing.T) { t.Run("Webhook", func(t *testing.T) { t.Parallel() - // nolint:gocritic // Unit test. ctx := dbauthz.AsNotifier(testutil.Context(t, testutil.WaitSuperLong)) store, _ := dbtestutil.NewDB(t) logger := testutil.Logger(t) diff --git a/coderd/notifications/reports/generator_internal_test.go b/coderd/notifications/reports/generator_internal_test.go index f61064c4e0..6dcff17311 100644 --- a/coderd/notifications/reports/generator_internal_test.go +++ b/coderd/notifications/reports/generator_internal_test.go @@ -505,7 +505,6 @@ func TestReportFailedWorkspaceBuilds(t *testing.T) { func setup(t *testing.T) (context.Context, slog.Logger, database.Store, pubsub.Pubsub, *notificationstest.FakeEnqueuer, *quartz.Mock) { t.Helper() - // nolint:gocritic // reportFailedWorkspaceBuilds is called by system. ctx := dbauthz.AsSystemRestricted(context.Background()) logger := slogtest.Make(t, &slogtest.Options{}) db, ps := dbtestutil.NewDB(t) diff --git a/coderd/prometheusmetrics/insights/metricscollector_test.go b/coderd/prometheusmetrics/insights/metricscollector_test.go index 9382fa5013..5c18ec6d1a 100644 --- a/coderd/prometheusmetrics/insights/metricscollector_test.go +++ b/coderd/prometheusmetrics/insights/metricscollector_test.go @@ -128,7 +128,6 @@ func TestCollectInsights(t *testing.T) { AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize, }) refTime := time.Now().Add(-3 * time.Minute).Truncate(time.Minute) - //nolint:gocritic // This is a test. err = reporter.ReportAppStats(dbauthz.AsSystemRestricted(context.Background()), []workspaceapps.StatsReport{ { UserID: user.ID, diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go index f0091ca63e..93573131a0 100644 --- a/coderd/provisionerdserver/provisionerdserver.go +++ b/coderd/provisionerdserver/provisionerdserver.go @@ -29,6 +29,7 @@ import ( "cdr.dev/slog" + "github.com/coder/coder/v2/coderd/usage" "github.com/coder/coder/v2/coderd/util/slice" "github.com/coder/coder/v2/codersdk/drpcsdk" @@ -121,6 +122,7 @@ type server struct { DeploymentValues *codersdk.DeploymentValues NotificationsEnqueuer notifications.Enqueuer PrebuildsOrchestrator *atomic.Pointer[prebuilds.ReconciliationOrchestrator] + UsageInserter *atomic.Pointer[usage.Inserter] OIDCConfig promoauth.OAuth2Config @@ -174,6 +176,7 @@ func NewServer( auditor *atomic.Pointer[audit.Auditor], templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore], userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore], + usageInserter *atomic.Pointer[usage.Inserter], deploymentValues *codersdk.DeploymentValues, options Options, enqueuer notifications.Enqueuer, @@ -195,6 +198,9 @@ func NewServer( if userQuietHoursScheduleStore == nil { return nil, xerrors.New("userQuietHoursScheduleStore is nil") } + if usageInserter == nil { + return nil, xerrors.New("usageCollector is nil") + } if deploymentValues == nil { return nil, xerrors.New("deploymentValues is nil") } @@ -244,6 +250,7 @@ func NewServer( heartbeatInterval: options.HeartbeatInterval, heartbeatFn: options.HeartbeatFn, PrebuildsOrchestrator: prebuildsOrchestrator, + UsageInserter: usageInserter, } if s.heartbeatFn == nil { @@ -2030,6 +2037,20 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro sidebarAppID = uuid.NullUUID{} } + if hasAITask && workspaceBuild.Transition == database.WorkspaceTransitionStart { + // Insert usage event for managed agents. + usageInserter := s.UsageInserter.Load() + if usageInserter != nil { + event := usage.DCManagedAgentsV1{ + Count: 1, + } + err = (*usageInserter).InsertDiscreteUsageEvent(ctx, db, event) + if err != nil { + return xerrors.Errorf("insert %q event: %w", event.EventType(), err) + } + } + } + hasExternalAgent := false for _, resource := range jobType.WorkspaceBuild.Resources { if resource.Type == "coder_external_agent" { diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go index 7fb351bf0c..8bb06eb52c 100644 --- a/coderd/provisionerdserver/provisionerdserver_test.go +++ b/coderd/provisionerdserver/provisionerdserver_test.go @@ -16,6 +16,7 @@ import ( "time" "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/trace" @@ -30,7 +31,9 @@ import ( "github.com/coder/coder/v2/buildinfo" "github.com/coder/coder/v2/coderd/audit" + "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" @@ -44,6 +47,7 @@ import ( "github.com/coder/coder/v2/coderd/schedule" "github.com/coder/coder/v2/coderd/schedule/cron" "github.com/coder/coder/v2/coderd/telemetry" + "github.com/coder/coder/v2/coderd/usage" "github.com/coder/coder/v2/coderd/wspubsub" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" @@ -67,6 +71,13 @@ func testUserQuietHoursScheduleStore() *atomic.Pointer[schedule.UserQuietHoursSc return ptr } +func testUsageInserter() *atomic.Pointer[usage.Inserter] { + ptr := &atomic.Pointer[usage.Inserter]{} + inserter := usage.NewAGPLInserter() + ptr.Store(&inserter) + return ptr +} + func TestAcquireJob_LongPoll(t *testing.T) { t.Parallel() //nolint:dogsled @@ -681,12 +692,20 @@ func TestUpdateJob(t *testing.T) { t.Run("NotRunning", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, nil) + user := dbgen.User(t, db, database.User{}) + version := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + CreatedBy: user.ID, + OrganizationID: pd.OrganizationID, + JobID: uuid.New(), + }) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: uuid.New(), - Provisioner: database.ProvisionerTypeEcho, - StorageMethod: database.ProvisionerStorageMethodFile, - Type: database.ProvisionerJobTypeTemplateVersionDryRun, - Input: json.RawMessage("{}"), + ID: version.JobID, + Provisioner: database.ProvisionerTypeEcho, + StorageMethod: database.ProvisionerStorageMethodFile, + Type: database.ProvisionerJobTypeTemplateVersionDryRun, + Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{ + TemplateVersionID: version.ID, + })), OrganizationID: pd.OrganizationID, Tags: pd.Tags, }) @@ -700,12 +719,20 @@ func TestUpdateJob(t *testing.T) { t.Run("NotOwner", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, nil) + user := dbgen.User(t, db, database.User{}) + version := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + CreatedBy: user.ID, + OrganizationID: pd.OrganizationID, + JobID: uuid.New(), + }) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: uuid.New(), - Provisioner: database.ProvisionerTypeEcho, - StorageMethod: database.ProvisionerStorageMethodFile, - Type: database.ProvisionerJobTypeTemplateVersionDryRun, - Input: json.RawMessage("{}"), + ID: version.JobID, + Provisioner: database.ProvisionerTypeEcho, + StorageMethod: database.ProvisionerStorageMethodFile, + Type: database.ProvisionerJobTypeTemplateVersionDryRun, + Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{ + TemplateVersionID: version.ID, + })), OrganizationID: pd.OrganizationID, Tags: pd.Tags, }) @@ -730,38 +757,57 @@ func TestUpdateJob(t *testing.T) { require.ErrorContains(t, err, "you don't own this job") }) - setupJob := func(t *testing.T, db database.Store, srvID, orgID uuid.UUID, tags database.StringMap) uuid.UUID { - job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: uuid.New(), - OrganizationID: orgID, - Provisioner: database.ProvisionerTypeEcho, - Type: database.ProvisionerJobTypeTemplateVersionImport, - StorageMethod: database.ProvisionerStorageMethodFile, - Input: json.RawMessage("{}"), - Tags: tags, - }) + setupJob := func(t *testing.T, db database.Store, srvID, orgID uuid.UUID, tags database.StringMap) (templateVersionID, jobID uuid.UUID) { + templateVersionID = uuid.New() + jobID = uuid.New() + err := db.InTx(func(db database.Store) error { + user := dbgen.User(t, db, database.User{}) + version := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + ID: templateVersionID, + CreatedBy: user.ID, + OrganizationID: orgID, + JobID: jobID, + }) + job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ + ID: version.JobID, + OrganizationID: orgID, + Provisioner: database.ProvisionerTypeEcho, + Type: database.ProvisionerJobTypeTemplateVersionImport, + StorageMethod: database.ProvisionerStorageMethodFile, + Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{ + TemplateVersionID: version.ID, + })), + Tags: tags, + }) + if err != nil { + return xerrors.Errorf("insert provisioner job: %w", err) + } + _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ + WorkerID: uuid.NullUUID{ + UUID: srvID, + Valid: true, + }, + Types: []database.ProvisionerType{database.ProvisionerTypeEcho}, + StartedAt: sql.NullTime{ + Time: dbtime.Now(), + Valid: true, + }, + OrganizationID: orgID, + ProvisionerTags: must(json.Marshal(job.Tags)), + }) + if err != nil { + return xerrors.Errorf("acquire provisioner job: %w", err) + } + return nil + }, nil) require.NoError(t, err) - _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ - WorkerID: uuid.NullUUID{ - UUID: srvID, - Valid: true, - }, - Types: []database.ProvisionerType{database.ProvisionerTypeEcho}, - StartedAt: sql.NullTime{ - Time: dbtime.Now(), - Valid: true, - }, - OrganizationID: orgID, - ProvisionerTags: must(json.Marshal(job.Tags)), - }) - require.NoError(t, err) - return job.ID + return templateVersionID, jobID } t.Run("Success", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) + _, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) _, err := srv.UpdateJob(ctx, &proto.UpdateJobRequest{ JobId: job.String(), }) @@ -771,7 +817,7 @@ func TestUpdateJob(t *testing.T) { t.Run("Logs", func(t *testing.T) { t.Parallel() srv, db, ps, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) + _, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) published := make(chan struct{}) @@ -796,23 +842,14 @@ func TestUpdateJob(t *testing.T) { t.Run("Readme", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) - versionID := uuid.New() - user := dbgen.User(t, db, database.User{}) - err := db.InsertTemplateVersion(ctx, database.InsertTemplateVersionParams{ - ID: versionID, - CreatedBy: user.ID, - OrganizationID: pd.OrganizationID, - JobID: job, - }) - require.NoError(t, err) - _, err = srv.UpdateJob(ctx, &proto.UpdateJobRequest{ + templateVersionID, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) + _, err := srv.UpdateJob(ctx, &proto.UpdateJobRequest{ JobId: job.String(), Readme: []byte("# hello world"), }) require.NoError(t, err) - version, err := db.GetTemplateVersionByID(ctx, versionID) + version, err := db.GetTemplateVersionByID(ctx, templateVersionID) require.NoError(t, err) require.Equal(t, "# hello world", version.Readme) }) @@ -825,16 +862,7 @@ func TestUpdateJob(t *testing.T) { defer cancel() srv, db, _, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) - versionID := uuid.New() - user := dbgen.User(t, db, database.User{}) - err := db.InsertTemplateVersion(ctx, database.InsertTemplateVersionParams{ - ID: versionID, - CreatedBy: user.ID, - JobID: job, - OrganizationID: pd.OrganizationID, - }) - require.NoError(t, err) + templateVersionID, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) firstTemplateVariable := &sdkproto.TemplateVariable{ Name: "first", Type: "string", @@ -863,7 +891,7 @@ func TestUpdateJob(t *testing.T) { require.NoError(t, err) require.Len(t, response.VariableValues, 2) - templateVariables, err := db.GetTemplateVersionVariables(ctx, versionID) + templateVariables, err := db.GetTemplateVersionVariables(ctx, templateVersionID) require.NoError(t, err) require.Len(t, templateVariables, 2) require.Equal(t, templateVariables[0].Value, firstTemplateVariable.DefaultValue) @@ -875,16 +903,7 @@ func TestUpdateJob(t *testing.T) { defer cancel() srv, db, _, pd := setup(t, false, &overrides{}) - user := dbgen.User(t, db, database.User{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) - versionID := uuid.New() - err := db.InsertTemplateVersion(ctx, database.InsertTemplateVersionParams{ - CreatedBy: user.ID, - ID: versionID, - JobID: job, - OrganizationID: pd.OrganizationID, - }) - require.NoError(t, err) + templateVersionID, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) firstTemplateVariable := &sdkproto.TemplateVariable{ Name: "first", Type: "string", @@ -909,7 +928,7 @@ func TestUpdateJob(t *testing.T) { // Even though there is an error returned, variables are stored in the database // to show the schema in the site UI. - templateVariables, err := db.GetTemplateVersionVariables(ctx, versionID) + templateVariables, err := db.GetTemplateVersionVariables(ctx, templateVersionID) require.NoError(t, err) require.Len(t, templateVariables, 2) require.Equal(t, templateVariables[0].Value, firstTemplateVariable.DefaultValue) @@ -923,18 +942,9 @@ func TestUpdateJob(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - srv, db, _, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) - versionID := uuid.New() - user := dbgen.User(t, db, database.User{}) - err := db.InsertTemplateVersion(ctx, database.InsertTemplateVersionParams{ - ID: versionID, - CreatedBy: user.ID, - JobID: job, - OrganizationID: pd.OrganizationID, - }) - require.NoError(t, err) - _, err = srv.UpdateJob(ctx, &proto.UpdateJobRequest{ + srv, db, _, pd := setup(t, false, nil) + templateVersionID, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) + _, err := srv.UpdateJob(ctx, &proto.UpdateJobRequest{ JobId: job.String(), WorkspaceTags: map[string]string{ "bird": "tweety", @@ -943,7 +953,7 @@ func TestUpdateJob(t *testing.T) { }) require.NoError(t, err) - workspaceTags, err := db.GetTemplateVersionWorkspaceTags(ctx, versionID) + workspaceTags, err := db.GetTemplateVersionWorkspaceTags(ctx, templateVersionID) require.NoError(t, err) require.Len(t, workspaceTags, 2) require.Equal(t, workspaceTags[0].Key, "bird") @@ -955,7 +965,7 @@ func TestUpdateJob(t *testing.T) { t.Run("LogSizeLimit", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) + _, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) // Create a log message that exceeds the 1MB limit largeOutput := strings.Repeat("a", 1048577) // 1MB + 1 byte @@ -979,7 +989,7 @@ func TestUpdateJob(t *testing.T) { t.Run("IncrementalLogSizeOverflow", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) + _, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) // Send logs that together exceed the limit mediumOutput := strings.Repeat("b", 524289) // Half a MB + 1 byte @@ -1020,7 +1030,7 @@ func TestUpdateJob(t *testing.T) { t.Run("LogSizeTracking", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) + _, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) logOutput := "test log message" expectedSize := int32(len(logOutput)) // #nosec G115 - Log length is 16. @@ -1045,7 +1055,7 @@ func TestUpdateJob(t *testing.T) { t.Run("LogOverflowStopsProcessing", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) - job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) + _, job := setupJob(t, db, pd.ID, pd.OrganizationID, pd.Tags) // First: trigger overflow largeOutput := strings.Repeat("a", 1048577) // 1MB + 1 byte @@ -1108,12 +1118,20 @@ func TestFailJob(t *testing.T) { t.Run("NotOwner", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, nil) + user := dbgen.User(t, db, database.User{}) + version := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + CreatedBy: user.ID, + OrganizationID: pd.OrganizationID, + JobID: uuid.New(), + }) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: uuid.New(), - Provisioner: database.ProvisionerTypeEcho, - StorageMethod: database.ProvisionerStorageMethodFile, - Type: database.ProvisionerJobTypeTemplateVersionImport, - Input: json.RawMessage("{}"), + ID: version.JobID, + Provisioner: database.ProvisionerTypeEcho, + StorageMethod: database.ProvisionerStorageMethodFile, + Type: database.ProvisionerJobTypeTemplateVersionImport, + Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{ + TemplateVersionID: version.ID, + })), OrganizationID: pd.OrganizationID, Tags: pd.Tags, }) @@ -1139,13 +1157,21 @@ func TestFailJob(t *testing.T) { }) t.Run("AlreadyCompleted", func(t *testing.T) { t.Parallel() - srv, db, _, pd := setup(t, false, &overrides{}) + srv, db, _, pd := setup(t, false, nil) + user := dbgen.User(t, db, database.User{}) + version := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + CreatedBy: user.ID, + OrganizationID: pd.OrganizationID, + JobID: uuid.New(), + }) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: uuid.New(), - Provisioner: database.ProvisionerTypeEcho, - Type: database.ProvisionerJobTypeTemplateVersionImport, - StorageMethod: database.ProvisionerStorageMethodFile, - Input: json.RawMessage("{}"), + ID: version.JobID, + Provisioner: database.ProvisionerTypeEcho, + Type: database.ProvisionerJobTypeTemplateVersionImport, + StorageMethod: database.ProvisionerStorageMethodFile, + Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{ + TemplateVersionID: version.ID, + })), OrganizationID: pd.OrganizationID, Tags: pd.Tags, }) @@ -1310,14 +1336,22 @@ func TestCompleteJob(t *testing.T) { t.Run("NotOwner", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, nil) + user := dbgen.User(t, db, database.User{}) + version := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + CreatedBy: user.ID, + OrganizationID: pd.OrganizationID, + JobID: uuid.New(), + }) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: uuid.New(), + ID: version.JobID, Provisioner: database.ProvisionerTypeEcho, StorageMethod: database.ProvisionerStorageMethodFile, - Type: database.ProvisionerJobTypeWorkspaceBuild, + Type: database.ProvisionerJobTypeTemplateVersionImport, OrganizationID: pd.OrganizationID, - Input: json.RawMessage("{}"), - Tags: pd.Tags, + Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{ + TemplateVersionID: version.ID, + })), + Tags: pd.Tags, }) require.NoError(t, err) _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ @@ -1361,10 +1395,12 @@ func TestCompleteJob(t *testing.T) { OrganizationID: pd.OrganizationID, ID: jobID, Provisioner: database.ProvisionerTypeEcho, - Input: []byte(`{"template_version_id": "` + versionID.String() + `"}`), - StorageMethod: database.ProvisionerStorageMethodFile, - Type: database.ProvisionerJobTypeTemplateVersionImport, - Tags: pd.Tags, + Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{ + TemplateVersionID: versionID, + })), + StorageMethod: database.ProvisionerStorageMethodFile, + Type: database.ProvisionerJobTypeTemplateVersionImport, + Tags: pd.Tags, }) require.NoError(t, err) _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ @@ -1410,14 +1446,22 @@ func TestCompleteJob(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + version := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + CreatedBy: user.ID, + OrganizationID: org.ID, + JobID: uuid.New(), + }) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ ID: uuid.New(), OrganizationID: org.ID, Provisioner: database.ProvisionerTypeEcho, Type: database.ProvisionerJobTypeTemplateVersionDryRun, StorageMethod: database.ProvisionerStorageMethodFile, - Input: json.RawMessage("{}"), - Tags: pd.Tags, + Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{ + TemplateVersionID: version.ID, + })), + Tags: pd.Tags, }) require.NoError(t, err) _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ @@ -1628,25 +1672,49 @@ func TestCompleteJob(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) jobID := uuid.New() - versionID := uuid.New() user := dbgen.User(t, db, database.User{}) - err := db.InsertTemplateVersion(ctx, database.InsertTemplateVersionParams{ + tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ CreatedBy: user.ID, - ID: versionID, - JobID: jobID, OrganizationID: pd.OrganizationID, + JobID: jobID, + }) + template := dbgen.Template(t, db, database.Template{ + CreatedBy: user.ID, + OrganizationID: pd.OrganizationID, + ActiveVersionID: tv.ID, + }) + err := db.UpdateTemplateVersionByID(ctx, database.UpdateTemplateVersionByIDParams{ + ID: tv.ID, + TemplateID: uuid.NullUUID{ + UUID: template.ID, + Valid: true, + }, + UpdatedAt: dbtime.Now(), + Name: tv.Name, + Message: tv.Message, }) require.NoError(t, err) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: user.ID, + OrganizationID: pd.OrganizationID, + TemplateID: template.ID, + }) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ ID: jobID, Provisioner: database.ProvisionerTypeEcho, - Input: []byte(`{"template_version_id": "` + versionID.String() + `"}`), + Input: json.RawMessage("{}"), StorageMethod: database.ProvisionerStorageMethodFile, Type: database.ProvisionerJobTypeWorkspaceBuild, OrganizationID: pd.OrganizationID, Tags: pd.Tags, }) require.NoError(t, err) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: tv.ID, + InitiatorID: user.ID, + JobID: jobID, + }) _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ OrganizationID: pd.OrganizationID, WorkerID: uuid.NullUUID{ @@ -1697,11 +1765,13 @@ func TestCompleteJob(t *testing.T) { }) require.NoError(t, err) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: jobID, - Provisioner: database.ProvisionerTypeEcho, - Input: []byte(`{"template_version_id": "` + versionID.String() + `"}`), + ID: jobID, + Provisioner: database.ProvisionerTypeEcho, + Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{ + TemplateVersionID: versionID, + })), StorageMethod: database.ProvisionerStorageMethodFile, - Type: database.ProvisionerJobTypeWorkspaceBuild, + Type: database.ProvisionerJobTypeTemplateVersionImport, OrganizationID: pd.OrganizationID, Tags: pd.Tags, }) @@ -1766,10 +1836,12 @@ func TestCompleteJob(t *testing.T) { OrganizationID: pd.OrganizationID, ID: jobID, Provisioner: database.ProvisionerTypeEcho, - Input: []byte(`{"template_version_id": "` + versionID.String() + `"}`), - StorageMethod: database.ProvisionerStorageMethodFile, - Type: database.ProvisionerJobTypeWorkspaceBuild, - Tags: pd.Tags, + Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{ + TemplateVersionID: versionID, + })), + StorageMethod: database.ProvisionerStorageMethodFile, + Type: database.ProvisionerJobTypeTemplateVersionImport, + Tags: pd.Tags, }) require.NoError(t, err) _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ @@ -2091,12 +2163,20 @@ func TestCompleteJob(t *testing.T) { t.Run("TemplateDryRun", func(t *testing.T) { t.Parallel() srv, db, _, pd := setup(t, false, &overrides{}) + user := dbgen.User(t, db, database.User{}) + version := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + CreatedBy: user.ID, + OrganizationID: pd.OrganizationID, + JobID: uuid.New(), + }) job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{ - ID: uuid.New(), - Provisioner: database.ProvisionerTypeEcho, - Type: database.ProvisionerJobTypeTemplateVersionDryRun, - StorageMethod: database.ProvisionerStorageMethodFile, - Input: json.RawMessage("{}"), + ID: version.JobID, + Provisioner: database.ProvisionerTypeEcho, + Type: database.ProvisionerJobTypeTemplateVersionDryRun, + StorageMethod: database.ProvisionerStorageMethodFile, + Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{ + TemplateVersionID: version.ID, + })), OrganizationID: pd.OrganizationID, Tags: pd.Tags, }) @@ -2191,8 +2271,10 @@ func TestCompleteJob(t *testing.T) { Transition: database.WorkspaceTransitionStart, }}, provisionerJobParams: database.InsertProvisionerJobParams{ - Type: database.ProvisionerJobTypeTemplateVersionDryRun, - Input: json.RawMessage("{}"), + Type: database.ProvisionerJobTypeTemplateVersionDryRun, + Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{ + TemplateVersionID: templateVersionID, + })), }, }, { @@ -2349,22 +2431,26 @@ func TestCompleteJob(t *testing.T) { OrganizationID: pd.OrganizationID, }) tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + ID: templateVersionID, CreatedBy: user.ID, OrganizationID: pd.OrganizationID, TemplateID: uuid.NullUUID{UUID: tpl.ID, Valid: true}, JobID: job.ID, }) - workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ - TemplateID: tpl.ID, - OrganizationID: pd.OrganizationID, - OwnerID: user.ID, - }) - _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ - ID: workspaceBuildID, - JobID: job.ID, - WorkspaceID: workspace.ID, - TemplateVersionID: tv.ID, - }) + + if jobParams.Type == database.ProvisionerJobTypeWorkspaceBuild { + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + TemplateID: tpl.ID, + OrganizationID: pd.OrganizationID, + OwnerID: user.ID, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + ID: workspaceBuildID, + JobID: job.ID, + WorkspaceID: workspace.ID, + TemplateVersionID: tv.ID, + }) + } require.NoError(t, err) _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ @@ -2672,7 +2758,10 @@ func TestCompleteJob(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() - srv, db, _, pd := setup(t, false, &overrides{}) + fakeUsageInserter, usageInserterPtr := newFakeUsageInserter() + srv, db, _, pd := setup(t, false, &overrides{ + usageInserter: usageInserterPtr, + }) importJobID := uuid.New() tvID := uuid.New() @@ -2741,6 +2830,10 @@ func TestCompleteJob(t *testing.T) { require.NoError(t, err) require.True(t, version.HasAITask.Valid) // We ALWAYS expect a value to be set, therefore not nil, i.e. valid = true. require.Equal(t, tc.expected, version.HasAITask.Bool) + + // We never expect a usage event to be collected for + // template imports. + require.Empty(t, fakeUsageInserter.collectedEvents) }) } }) @@ -2750,22 +2843,27 @@ func TestCompleteJob(t *testing.T) { // will be set as well in that case. t.Run("WorkspaceBuild", func(t *testing.T) { type testcase struct { - name string - input *proto.CompletedJob_WorkspaceBuild - expected bool + name string + transition database.WorkspaceTransition + input *proto.CompletedJob_WorkspaceBuild + expectHasAiTask bool + expectUsageEvent bool } sidebarAppID := uuid.NewString() for _, tc := range []testcase{ { - name: "has_ai_task is false by default", - input: &proto.CompletedJob_WorkspaceBuild{ + name: "has_ai_task is false by default", + transition: database.WorkspaceTransitionStart, + input: &proto.CompletedJob_WorkspaceBuild{ // No AiTasks defined. }, - expected: false, + expectHasAiTask: false, + expectUsageEvent: false, }, { - name: "has_ai_task is set to true", + name: "has_ai_task is set to true", + transition: database.WorkspaceTransitionStart, input: &proto.CompletedJob_WorkspaceBuild{ AiTasks: []*sdkproto.AITask{ { @@ -2792,11 +2890,13 @@ func TestCompleteJob(t *testing.T) { }, }, }, - expected: true, + expectHasAiTask: true, + expectUsageEvent: true, }, // Checks regression for https://github.com/coder/coder/issues/18776 { - name: "non-existing app", + name: "non-existing app", + transition: database.WorkspaceTransitionStart, input: &proto.CompletedJob_WorkspaceBuild{ AiTasks: []*sdkproto.AITask{ { @@ -2808,13 +2908,49 @@ func TestCompleteJob(t *testing.T) { }, }, }, - expected: false, + expectHasAiTask: false, + expectUsageEvent: false, + }, + { + name: "has_ai_task is set to true, but transition is not start", + transition: database.WorkspaceTransitionStop, + input: &proto.CompletedJob_WorkspaceBuild{ + AiTasks: []*sdkproto.AITask{ + { + Id: uuid.NewString(), + SidebarApp: &sdkproto.AITaskSidebarApp{ + Id: sidebarAppID, + }, + }, + }, + Resources: []*sdkproto.Resource{ + { + Agents: []*sdkproto.Agent{ + { + Id: uuid.NewString(), + Name: "a", + Apps: []*sdkproto.App{ + { + Id: sidebarAppID, + Slug: "test-app", + }, + }, + }, + }, + }, + }, + }, + expectHasAiTask: true, + expectUsageEvent: false, }, } { t.Run(tc.name, func(t *testing.T) { t.Parallel() - srv, db, _, pd := setup(t, false, &overrides{}) + fakeUsageInserter, usageInserterPtr := newFakeUsageInserter() + srv, db, _, pd := setup(t, false, &overrides{ + usageInserter: usageInserterPtr, + }) importJobID := uuid.New() tvID := uuid.New() @@ -2868,7 +3004,7 @@ func TestCompleteJob(t *testing.T) { WorkspaceID: workspaceTable.ID, TemplateVersionID: version.ID, InitiatorID: user.ID, - Transition: database.WorkspaceTransitionStart, + Transition: tc.transition, }) _, err = db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ @@ -2899,11 +3035,22 @@ func TestCompleteJob(t *testing.T) { build, err = db.GetWorkspaceBuildByID(ctx, build.ID) require.NoError(t, err) require.True(t, build.HasAITask.Valid) // We ALWAYS expect a value to be set, therefore not nil, i.e. valid = true. - require.Equal(t, tc.expected, build.HasAITask.Bool) + require.Equal(t, tc.expectHasAiTask, build.HasAITask.Bool) - if tc.expected { + if tc.expectHasAiTask { require.Equal(t, sidebarAppID, build.AITaskSidebarAppID.UUID.String()) } + + if tc.expectUsageEvent { + // Check that a usage event was collected. + require.Len(t, fakeUsageInserter.collectedEvents, 1) + require.Equal(t, usage.DCManagedAgentsV1{ + Count: 1, + }, fakeUsageInserter.collectedEvents[0]) + } else { + // Check that no usage event was collected. + require.Empty(t, fakeUsageInserter.collectedEvents) + } }) } }) @@ -3835,6 +3982,7 @@ type overrides struct { externalAuthConfigs []*externalauth.Config templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore] userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore] + usageInserter *atomic.Pointer[usage.Inserter] clock *quartz.Mock acquireJobLongPollDuration time.Duration heartbeatFn func(ctx context.Context) error @@ -3855,13 +4003,14 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi var externalAuthConfigs []*externalauth.Config tss := testTemplateScheduleStore() uqhss := testUserQuietHoursScheduleStore() + usageInserter := testUsageInserter() clock := quartz.NewReal() pollDur := time.Duration(0) if ov == nil { ov = &overrides{} } if ov.ctx == nil { - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(dbauthz.AsProvisionerd(context.Background())) t.Cleanup(cancel) ov.ctx = ctx } @@ -3892,6 +4041,15 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi require.True(t, swapped) } } + if ov.usageInserter != nil { + tUsageInserter := usageInserter.Load() + // keep the initial test value if the override hasn't set the atomic pointer. + usageInserter = ov.usageInserter + if usageInserter.Load() == nil { + swapped := usageInserter.CompareAndSwap(nil, tUsageInserter) + require.True(t, swapped) + } + } if ov.clock != nil { clock = ov.clock } @@ -3929,6 +4087,10 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi var op atomic.Pointer[agplprebuilds.ReconciliationOrchestrator] op.Store(&prebuildsOrchestrator) + // Use an authz wrapped database for the server to ensure permission checks + // work. + authorizer := rbac.NewStrictCachingAuthorizer(prometheus.NewRegistry()) + serverDB := dbauthz.New(db, authorizer, logger, coderdtest.AccessControlStorePointer()) srv, err := provisionerdserver.NewServer( ov.ctx, proto.CurrentVersion.String(), @@ -3938,7 +4100,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi slogtest.Make(t, &slogtest.Options{IgnoreErrors: ignoreLogErrors}), []database.ProvisionerType{database.ProvisionerTypeEcho}, provisionerdserver.Tags(daemon.Tags), - db, + serverDB, ps, provisionerdserver.NewAcquirer(ov.ctx, logger.Named("acquirer"), db, ps), telemetry.NewNoop(), @@ -3947,6 +4109,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi auditPtr, tss, uqhss, + usageInserter, deploymentValues, provisionerdserver.Options{ ExternalAuthConfigs: externalAuthConfigs, @@ -4061,3 +4224,22 @@ func (s *fakeStream) cancel() { s.canceled = true s.c.Broadcast() } + +type fakeUsageInserter struct { + collectedEvents []usage.Event +} + +var _ usage.Inserter = &fakeUsageInserter{} + +func newFakeUsageInserter() (*fakeUsageInserter, *atomic.Pointer[usage.Inserter]) { + ptr := &atomic.Pointer[usage.Inserter]{} + fake := &fakeUsageInserter{} + var inserter usage.Inserter = fake + ptr.Store(&inserter) + return fake, ptr +} + +func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usage.DiscreteEvent) error { + f.collectedEvents = append(f.collectedEvents, event) + return nil +} diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go index a8130bea17..0b48a24aeb 100644 --- a/coderd/rbac/authz.go +++ b/coderd/rbac/authz.go @@ -76,7 +76,7 @@ const ( SubjectTypeNotifier SubjectType = "notifier" SubjectTypeSubAgentAPI SubjectType = "sub_agent_api" SubjectTypeFileReader SubjectType = "file_reader" - SubjectTypeUsageTracker SubjectType = "usage_tracker" + SubjectTypeUsagePublisher SubjectType = "usage_publisher" ) const ( diff --git a/coderd/usage/inserter.go b/coderd/usage/inserter.go index 08ca8dec3e..3a0e85f273 100644 --- a/coderd/usage/inserter.go +++ b/coderd/usage/inserter.go @@ -10,6 +10,8 @@ import ( type Inserter interface { // InsertDiscreteUsageEvent writes a discrete usage event to the database // within the given transaction. + // The caller context must be authorized to create usage events in the + // database. InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event DiscreteEvent) error } diff --git a/coderd/userauth_test.go b/coderd/userauth_test.go index 4c9412fda3..504b102e9e 100644 --- a/coderd/userauth_test.go +++ b/coderd/userauth_test.go @@ -335,7 +335,6 @@ func TestUserOAuth2Github(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) - // nolint:gocritic // Unit test count, err := db.GetUserCount(dbauthz.AsSystemRestricted(ctx), false) require.NoError(t, err) require.Equal(t, int64(1), count) @@ -897,7 +896,6 @@ func TestUserOAuth2Github(t *testing.T) { require.Empty(t, links) // Make sure a user_link cannot be created with a deleted user. - // nolint:gocritic // Unit test _, err = db.InsertUserLink(dbauthz.AsSystemRestricted(ctx), database.InsertUserLinkParams{ UserID: deleted.ID, LoginType: "github", diff --git a/coderd/users_test.go b/coderd/users_test.go index 5928fc6486..22c9fad5ee 100644 --- a/coderd/users_test.go +++ b/coderd/users_test.go @@ -1544,7 +1544,6 @@ func TestUsersFilter(t *testing.T) { } userClient, userData := coderdtest.CreateAnotherUser(t, client, first.OrganizationID, roles...) // Set the last seen for each user to a unique day - // nolint:gocritic // Unit test _, err := api.Database.UpdateUserLastSeenAt(dbauthz.AsSystemRestricted(ctx), database.UpdateUserLastSeenAtParams{ ID: userData.ID, LastSeenAt: lastSeenNow.Add(-1 * time.Hour * 24 * time.Duration(i)), @@ -1572,7 +1571,6 @@ func TestUsersFilter(t *testing.T) { // Add users with different creation dates for testing date filters for i := 0; i < 3; i++ { - // nolint:gocritic // Using system context is necessary to seed data in tests user1, err := api.Database.InsertUser(dbauthz.AsSystemRestricted(ctx), database.InsertUserParams{ ID: uuid.New(), Email: fmt.Sprintf("before%d@coder.com", i), @@ -1594,7 +1592,6 @@ func TestUsersFilter(t *testing.T) { require.NoError(t, err) users = append(users, sdkUser1) - // nolint:gocritic //Using system context is necessary to seed data in tests user2, err := api.Database.InsertUser(dbauthz.AsSystemRestricted(ctx), database.InsertUserParams{ ID: uuid.New(), Email: fmt.Sprintf("during%d@coder.com", i), @@ -1615,7 +1612,6 @@ func TestUsersFilter(t *testing.T) { require.NoError(t, err) users = append(users, sdkUser2) - // nolint:gocritic // Using system context is necessary to seed data in tests user3, err := api.Database.InsertUser(dbauthz.AsSystemRestricted(ctx), database.InsertUserParams{ ID: uuid.New(), Email: fmt.Sprintf("after%d@coder.com", i), @@ -1912,7 +1908,6 @@ func TestGetUsers(t *testing.T) { Email: "test2@coder.com", Username: "test2", }) - // nolint:gocritic // Unit test err := db.UpdateUserGithubComUserID(dbauthz.AsSystemRestricted(ctx), database.UpdateUserGithubComUserIDParams{ ID: first.UserID, GithubComUserID: sql.NullInt64{ diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index ac58df1b77..6f28b12af5 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -562,7 +562,6 @@ func TestWorkspaceAgentConnectRPC(t *testing.T) { seed := database.WorkspaceTable{OrganizationID: user.OrganizationID, OwnerID: user.UserID} wsb := dbfake.WorkspaceBuild(t, db, seed).WithAgent().Do() // When: the workspace is marked as soft-deleted - // nolint:gocritic // this is a test err := db.UpdateWorkspaceDeletedByID( dbauthz.AsProvisionerd(ctx), database.UpdateWorkspaceDeletedByIDParams{ID: wsb.Workspace.ID, Deleted: true}, @@ -633,7 +632,6 @@ func TestWorkspaceAgentClientCoordinate_BadVersion(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) agentToken, err := uuid.Parse(r.AgentToken) require.NoError(t, err) - //nolint: gocritic // testing ao, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentToken) require.NoError(t, err) @@ -724,7 +722,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) { agentTokenUUID, err := uuid.Parse(r.AgentToken) require.NoError(t, err) ctx := testutil.Context(t, testutil.WaitLong) - agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) //nolint + agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) require.NoError(t, err) // Connect with no resume token, and ensure that the peer ID is set to a @@ -796,7 +794,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) { agentTokenUUID, err := uuid.Parse(r.AgentToken) require.NoError(t, err) ctx := testutil.Context(t, testutil.WaitLong) - agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) //nolint + agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) require.NoError(t, err) // Connect with no resume token, and ensure that the peer ID is set to a diff --git a/coderd/workspacebuilds_test.go b/coderd/workspacebuilds_test.go index 633acae328..e888115093 100644 --- a/coderd/workspacebuilds_test.go +++ b/coderd/workspacebuilds_test.go @@ -55,7 +55,6 @@ func TestWorkspaceBuild(t *testing.T) { Auditor: auditor, }) user := coderdtest.CreateFirstUser(t, client) - //nolint:gocritic // testing up, err := db.UpdateUserProfile(dbauthz.AsSystemRestricted(ctx), database.UpdateUserProfileParams{ ID: user.UserID, Email: coderdtest.FirstUserParams.Email, @@ -518,7 +517,6 @@ func TestWorkspaceBuildsProvisionerState(t *testing.T) { OrganizationID: first.OrganizationID, }).Do() - // nolint:gocritic // For testing daemons, err := store.GetProvisionerDaemons(dbauthz.AsSystemReadProvisionerDaemons(ctx)) require.NoError(t, err) require.Empty(t, daemons, "Provisioner daemons should be empty for this test") diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go index 8fc11ef6c8..4df83114c6 100644 --- a/coderd/workspaces_test.go +++ b/coderd/workspaces_test.go @@ -1427,7 +1427,6 @@ func TestWorkspaceFilterAllStatus(t *testing.T) { t.Parallel() // For this test, we do not care about permissions. - // nolint:gocritic // unit testing ctx := dbauthz.AsSystemRestricted(context.Background()) db, pubsub := dbtestutil.NewDB(t) client := coderdtest.New(t, &coderdtest.Options{ @@ -2215,15 +2214,12 @@ func TestWorkspaceFilterManual(t *testing.T) { after := coderdtest.CreateWorkspace(t, client, template.ID) _ = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, after.LatestBuild.ID) - //nolint:gocritic // Unit testing context err := api.Database.UpdateWorkspaceLastUsedAt(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceLastUsedAtParams{ ID: before.ID, LastUsedAt: now.UTC().Add(time.Hour * -1), }) require.NoError(t, err) - // Unit testing context - //nolint:gocritic // Unit testing context err = api.Database.UpdateWorkspaceLastUsedAt(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceLastUsedAtParams{ ID: after.ID, LastUsedAt: now.UTC().Add(time.Hour * 1), @@ -2916,14 +2912,14 @@ func TestWorkspaceUpdateTTL(t *testing.T) { // This is a hack, but the max_deadline isn't precisely configurable // without a lot of unnecessary hassle. - dbBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), build.ID) //nolint:gocritic // test + dbBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), build.ID) require.NoError(t, err) - dbJob, err := db.GetProvisionerJobByID(dbauthz.AsSystemRestricted(ctx), dbBuild.JobID) //nolint:gocritic // test + dbJob, err := db.GetProvisionerJobByID(dbauthz.AsSystemRestricted(ctx), dbBuild.JobID) require.NoError(t, err) require.True(t, dbJob.CompletedAt.Valid) initialDeadline := dbJob.CompletedAt.Time.Add(deadline) expectedMaxDeadline := dbJob.CompletedAt.Time.Add(maxDeadline) - err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{ //nolint:gocritic // test + err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{ ID: build.ID, Deadline: initialDeadline, MaxDeadline: expectedMaxDeadline, @@ -4507,14 +4503,12 @@ func TestOIDCRemoved(t *testing.T) { user, userData := coderdtest.CreateAnotherUser(t, owner, first.OrganizationID, rbac.ScopedRoleOrgAdmin(first.OrganizationID)) ctx := testutil.Context(t, testutil.WaitMedium) - //nolint:gocritic // unit test _, err := db.UpdateUserLoginType(dbauthz.AsSystemRestricted(ctx), database.UpdateUserLoginTypeParams{ NewLoginType: database.LoginTypeOIDC, UserID: userData.ID, }) require.NoError(t, err) - //nolint:gocritic // unit test _, err = db.InsertUserLink(dbauthz.AsSystemRestricted(ctx), database.InsertUserLinkParams{ UserID: userData.ID, LoginType: database.LoginTypeOIDC, @@ -4603,7 +4597,6 @@ func TestWorkspaceFilterHasAITask(t *testing.T) { }) if aiTaskPrompt != nil { - //nolint:gocritic // unit test err := db.InsertWorkspaceBuildParameters(dbauthz.AsSystemRestricted(ctx), database.InsertWorkspaceBuildParametersParams{ WorkspaceBuildID: build.ID, Name: []string{provider.TaskPromptParameterName}, @@ -4806,7 +4799,6 @@ func TestMultipleAITasksDisallowed(t *testing.T) { ws := coderdtest.CreateWorkspace(t, client, template.ID) coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID) - //nolint: gocritic // testing ctx := dbauthz.AsSystemRestricted(t.Context()) pj, err := db.GetProvisionerJobByID(ctx, ws.LatestBuild.Job.ID) require.NoError(t, err) diff --git a/enterprise/cli/prebuilds_test.go b/enterprise/cli/prebuilds_test.go index 76d11a41d6..cf0c741050 100644 --- a/enterprise/cli/prebuilds_test.go +++ b/enterprise/cli/prebuilds_test.go @@ -434,7 +434,6 @@ func TestSchedulePrebuilds(t *testing.T) { }).Do() // Mark the prebuilt workspace's agent as ready so the prebuild can be claimed - // nolint:gocritic ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong)) agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken)) require.NoError(t, err) diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index 3b1fd63ab1..f58ec86b58 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -20,6 +20,7 @@ import ( "github.com/coder/coder/v2/enterprise/audit/backends" "github.com/coder/coder/v2/enterprise/coderd" "github.com/coder/coder/v2/enterprise/coderd/dormancy" + "github.com/coder/coder/v2/enterprise/coderd/usage" "github.com/coder/coder/v2/enterprise/dbcrypt" "github.com/coder/coder/v2/enterprise/trialer" "github.com/coder/coder/v2/tailnet" @@ -116,11 +117,33 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { o.ExternalTokenEncryption = cs } + if o.LicenseKeys == nil { + o.LicenseKeys = coderd.Keys + } + + closers := &multiCloser{} + + // Create the enterprise API. api, err := coderd.New(ctx, o) if err != nil { return nil, nil, err } - return api.AGPL, api, nil + closers.Add(api) + + // Start the enterprise usage publisher routine. This won't do anything + // unless the deployment is licensed and one of the licenses has usage + // publishing enabled. + publisher := usage.NewTallymanPublisher(ctx, options.Logger, options.Database, o.LicenseKeys, + usage.PublisherWithHTTPClient(api.HTTPClient), + ) + err = publisher.Start() + if err != nil { + _ = closers.Close() + return nil, nil, xerrors.Errorf("start usage publisher: %w", err) + } + closers.Add(publisher) + + return api.AGPL, closers, nil }) cmd.AddSubcommands( @@ -128,3 +151,23 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { ) return cmd } + +type multiCloser struct { + closers []io.Closer +} + +var _ io.Closer = &multiCloser{} + +func (m *multiCloser) Add(closer io.Closer) { + m.closers = append(m.closers, closer) +} + +func (m *multiCloser) Close() error { + var errs []error + for _, closer := range m.closers { + if err := closer.Close(); err != nil { + errs = append(errs, xerrors.Errorf("close %T: %w", closer, err)) + } + } + return errors.Join(errs...) +} diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 8190de103c..a81e165854 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -10,6 +10,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/coder/coder/v2/buildinfo" @@ -21,10 +22,12 @@ import ( "github.com/coder/coder/v2/coderd/pproflabel" agplprebuilds "github.com/coder/coder/v2/coderd/prebuilds" "github.com/coder/coder/v2/coderd/rbac/policy" + agplusage "github.com/coder/coder/v2/coderd/usage" "github.com/coder/coder/v2/coderd/wsbuilder" "github.com/coder/coder/v2/enterprise/coderd/connectionlog" "github.com/coder/coder/v2/enterprise/coderd/enidpsync" "github.com/coder/coder/v2/enterprise/coderd/portsharing" + "github.com/coder/coder/v2/enterprise/coderd/usage" "github.com/coder/quartz" "golang.org/x/xerrors" @@ -90,6 +93,13 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { if options.Entitlements == nil { options.Entitlements = entitlements.New() } + if options.Options.UsageInserter == nil { + options.Options.UsageInserter = &atomic.Pointer[agplusage.Inserter]{} + } + if options.Options.UsageInserter.Load() == nil { + collector := usage.NewDBInserter() + options.Options.UsageInserter.Store(&collector) + } ctx, cancelFunc := context.WithCancel(ctx) diff --git a/enterprise/coderd/coderd_test.go b/enterprise/coderd/coderd_test.go index 94d9e4fda2..302b367c30 100644 --- a/enterprise/coderd/coderd_test.go +++ b/enterprise/coderd/coderd_test.go @@ -154,7 +154,6 @@ func TestEntitlements(t *testing.T) { entitlements, err := anotherClient.Entitlements(context.Background()) require.NoError(t, err) require.False(t, entitlements.HasLicense) - //nolint:gocritic // unit test ctx := testDBAuthzRole(context.Background()) _, err = api.Database.InsertLicense(ctx, database.InsertLicenseParams{ UploadedAt: dbtime.Now(), @@ -186,7 +185,6 @@ func TestEntitlements(t *testing.T) { require.False(t, entitlements.HasLicense) // Valid ctx := context.Background() - //nolint:gocritic // unit test _, err = api.Database.InsertLicense(testDBAuthzRole(ctx), database.InsertLicenseParams{ UploadedAt: dbtime.Now(), Exp: dbtime.Now().AddDate(1, 0, 0), @@ -198,7 +196,6 @@ func TestEntitlements(t *testing.T) { }) require.NoError(t, err) // Expired - //nolint:gocritic // unit test _, err = api.Database.InsertLicense(testDBAuthzRole(ctx), database.InsertLicenseParams{ UploadedAt: dbtime.Now(), Exp: dbtime.Now().AddDate(-1, 0, 0), @@ -208,7 +205,6 @@ func TestEntitlements(t *testing.T) { }) require.NoError(t, err) // Invalid - //nolint:gocritic // unit test _, err = api.Database.InsertLicense(testDBAuthzRole(ctx), database.InsertLicenseParams{ UploadedAt: dbtime.Now(), Exp: dbtime.Now().AddDate(1, 0, 0), diff --git a/enterprise/coderd/enidpsync/organizations_test.go b/enterprise/coderd/enidpsync/organizations_test.go index 13a9bd69ed..c3bae7cd1d 100644 --- a/enterprise/coderd/enidpsync/organizations_test.go +++ b/enterprise/coderd/enidpsync/organizations_test.go @@ -56,7 +56,6 @@ func TestOrganizationSync(t *testing.T) { requireUserOrgs := func(t *testing.T, db database.Store, user database.User, expected []uuid.UUID) { t.Helper() - // nolint:gocritic // in testing members, err := db.OrganizationMembers(dbauthz.AsSystemRestricted(context.Background()), database.OrganizationMembersParams{ UserID: user.ID, }) diff --git a/enterprise/coderd/idpsync_test.go b/enterprise/coderd/idpsync_test.go index d34701c3f6..49d83a6268 100644 --- a/enterprise/coderd/idpsync_test.go +++ b/enterprise/coderd/idpsync_test.go @@ -39,7 +39,6 @@ func TestGetGroupSyncSettings(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) dbresv := runtimeconfig.OrganizationResolver(user.OrganizationID, runtimeconfig.NewStoreResolver(db)) entry := runtimeconfig.MustNew[*idpsync.GroupSyncSettings]("group-sync-settings") - //nolint:gocritic // Requires system context to set runtime config err := entry.SetRuntimeValue(dbauthz.AsSystemRestricted(ctx), dbresv, &idpsync.GroupSyncSettings{Field: "august"}) require.NoError(t, err) diff --git a/enterprise/coderd/prebuilds/metricscollector_test.go b/enterprise/coderd/prebuilds/metricscollector_test.go index 1e9f3f5082..b852079beb 100644 --- a/enterprise/coderd/prebuilds/metricscollector_test.go +++ b/enterprise/coderd/prebuilds/metricscollector_test.go @@ -231,7 +231,6 @@ func TestMetricsCollector(t *testing.T) { } // Force an update to the metrics state to allow the collector to collect fresh metrics. - // nolint:gocritic // Authz context needed to retrieve state. require.NoError(t, collector.UpdateState(dbauthz.AsPrebuildsOrchestrator(ctx), testutil.WaitLong)) metricsFamilies, err := registry.Gather() @@ -367,7 +366,6 @@ func TestMetricsCollector_DuplicateTemplateNames(t *testing.T) { "organization_name": defaultOrg.Name, } - // nolint:gocritic // Authz context needed to retrieve state. ctx = dbauthz.AsPrebuildsOrchestrator(ctx) // Then: metrics collect successfully. diff --git a/enterprise/coderd/provisionerdaemons.go b/enterprise/coderd/provisionerdaemons.go index c830495278..65b03a7d6b 100644 --- a/enterprise/coderd/provisionerdaemons.go +++ b/enterprise/coderd/provisionerdaemons.go @@ -352,6 +352,7 @@ func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request) &api.AGPL.Auditor, api.AGPL.TemplateScheduleStore, api.AGPL.UserQuietHoursScheduleStore, + api.AGPL.UsageInserter, api.DeploymentValues, provisionerdserver.Options{ ExternalAuthConfigs: api.ExternalAuthConfigs, diff --git a/enterprise/coderd/provisionerdaemons_test.go b/enterprise/coderd/provisionerdaemons_test.go index a94a60ffff..5797e978fa 100644 --- a/enterprise/coderd/provisionerdaemons_test.go +++ b/enterprise/coderd/provisionerdaemons_test.go @@ -682,7 +682,6 @@ func TestProvisionerDaemonServe(t *testing.T) { if tc.insertParams.Name != "" { tc.insertParams.OrganizationID = user.OrganizationID - // nolint:gocritic // test _, err := db.InsertProvisionerKey(dbauthz.AsSystemRestricted(ctx), tc.insertParams) require.NoError(t, err) } @@ -945,7 +944,6 @@ func TestGetProvisionerDaemons(t *testing.T) { daemonCreatedAt := time.Now() - //nolint:gocritic // We're not testing auth on the following in this test provisionerKey, err := db.InsertProvisionerKey(dbauthz.AsSystemRestricted(ctx), database.InsertProvisionerKeyParams{ Name: "Test Provisioner Key", ID: uuid.New(), @@ -956,7 +954,6 @@ func TestGetProvisionerDaemons(t *testing.T) { }) require.NoError(t, err, "should be able to create a provisioner key") - //nolint:gocritic // We're not testing auth on the following in this test pd, err := db.UpsertProvisionerDaemon(dbauthz.AsSystemRestricted(ctx), database.UpsertProvisionerDaemonParams{ CreatedAt: daemonCreatedAt, Name: "Test Provisioner Daemon", diff --git a/enterprise/coderd/schedule/template_test.go b/enterprise/coderd/schedule/template_test.go index 70dc308489..e764826f76 100644 --- a/enterprise/coderd/schedule/template_test.go +++ b/enterprise/coderd/schedule/template_test.go @@ -719,7 +719,6 @@ func TestNotifications(t *testing.T) { // Lower the dormancy TTL to ensure the schedule recalculates deadlines and // triggers notifications. - // nolint:gocritic // Need an actor in the context. _, err = templateScheduleStore.Set(dbauthz.AsNotifier(ctx), db, template, agplschedule.TemplateScheduleOptions{ TimeTilDormant: timeTilDormant / 2, TimeTilDormantAutoDelete: timeTilDormant / 2, diff --git a/enterprise/coderd/usage/inserter.go b/enterprise/coderd/usage/inserter.go index 3320c25d45..e07b536e75 100644 --- a/enterprise/coderd/usage/inserter.go +++ b/enterprise/coderd/usage/inserter.go @@ -13,16 +13,17 @@ import ( "github.com/coder/quartz" ) -// Inserter accepts usage events and stores them in the database for publishing. -type Inserter struct { +// dbCollector collects usage events and stores them in the database for +// publishing. +type dbCollector struct { clock quartz.Clock } -var _ agplusage.Inserter = &Inserter{} +var _ agplusage.Inserter = &dbCollector{} -// NewInserter creates a new database-backed usage event inserter. -func NewInserter(opts ...InserterOptions) *Inserter { - c := &Inserter{ +// NewDBInserter creates a new database-backed usage event inserter. +func NewDBInserter(opts ...InserterOption) agplusage.Inserter { + c := &dbCollector{ clock: quartz.NewReal(), } for _, opt := range opts { @@ -31,17 +32,17 @@ func NewInserter(opts ...InserterOptions) *Inserter { return c } -type InserterOptions func(*Inserter) +type InserterOption func(*dbCollector) // InserterWithClock sets the quartz clock to use for the inserter. -func InserterWithClock(clock quartz.Clock) InserterOptions { - return func(c *Inserter) { +func InserterWithClock(clock quartz.Clock) InserterOption { + return func(c *dbCollector) { c.clock = clock } } // InsertDiscreteUsageEvent implements agplusage.Inserter. -func (c *Inserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event agplusage.DiscreteEvent) error { +func (i *dbCollector) InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event agplusage.DiscreteEvent) error { if !event.EventType().IsDiscrete() { return xerrors.Errorf("event type %q is not a discrete event", event.EventType()) } @@ -61,6 +62,6 @@ func (c *Inserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.Sto ID: uuid.New().String(), EventType: string(event.EventType()), EventData: jsonData, - CreatedAt: dbtime.Time(c.clock.Now()), + CreatedAt: dbtime.Time(i.clock.Now()), }) } diff --git a/enterprise/coderd/usage/inserter_test.go b/enterprise/coderd/usage/inserter_test.go index c5abd931cf..b7ced536ae 100644 --- a/enterprise/coderd/usage/inserter_test.go +++ b/enterprise/coderd/usage/inserter_test.go @@ -28,7 +28,7 @@ func TestInserter(t *testing.T) { ctrl := gomock.NewController(t) db := dbmock.NewMockStore(ctrl) clock := quartz.NewMock(t) - inserter := usage.NewInserter(usage.InserterWithClock(clock)) + inserter := usage.NewDBInserter(usage.InserterWithClock(clock)) now := dbtime.Now() events := []struct { @@ -51,8 +51,8 @@ func TestInserter(t *testing.T) { for _, event := range events { eventJSON := jsoninate(t, event.event) - db.EXPECT().InsertUsageEvent(ctx, gomock.Any()).DoAndReturn( - func(ctx interface{}, params database.InsertUsageEventParams) error { + db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx any, params database.InsertUsageEventParams) error { _, err := uuid.Parse(params.ID) assert.NoError(t, err) assert.Equal(t, string(event.event.EventType()), params.EventType) @@ -76,7 +76,7 @@ func TestInserter(t *testing.T) { db := dbmock.NewMockStore(ctrl) // We should get an error if the event is invalid. - inserter := usage.NewInserter() + inserter := usage.NewDBInserter() err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{ Count: 0, // invalid }) diff --git a/enterprise/coderd/usage/publisher.go b/enterprise/coderd/usage/publisher.go index e872284116..8c0811c772 100644 --- a/enterprise/coderd/usage/publisher.go +++ b/enterprise/coderd/usage/publisher.go @@ -15,11 +15,11 @@ import ( "cdr.dev/slog" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/pproflabel" agplusage "github.com/coder/coder/v2/coderd/usage" "github.com/coder/coder/v2/cryptorand" - "github.com/coder/coder/v2/enterprise/coderd" "github.com/coder/coder/v2/enterprise/coderd/license" "github.com/coder/quartz" ) @@ -49,17 +49,17 @@ type Publisher interface { } type tallymanPublisher struct { - ctx context.Context - ctxCancel context.CancelFunc - log slog.Logger - db database.Store - done chan struct{} + ctx context.Context + ctxCancel context.CancelFunc + log slog.Logger + db database.Store + licenseKeys map[string]ed25519.PublicKey + done chan struct{} // Configured with options: ingestURL string httpClient *http.Client clock quartz.Clock - licenseKeys map[string]ed25519.PublicKey initialDelay time.Duration } @@ -67,19 +67,21 @@ var _ Publisher = &tallymanPublisher{} // NewTallymanPublisher creates a Publisher that publishes usage events to // Coder's Tallyman service. -func NewTallymanPublisher(ctx context.Context, log slog.Logger, db database.Store, opts ...TallymanPublisherOption) Publisher { +func NewTallymanPublisher(ctx context.Context, log slog.Logger, db database.Store, keys map[string]ed25519.PublicKey, opts ...TallymanPublisherOption) Publisher { ctx, cancel := context.WithCancel(ctx) - publisher := &tallymanPublisher{ - ctx: ctx, - ctxCancel: cancel, - log: log, - db: db, - done: make(chan struct{}), + ctx = dbauthz.AsUsagePublisher(ctx) //nolint:gocritic // we intentionally want to be able to process usage events - ingestURL: tallymanIngestURLV1, - httpClient: http.DefaultClient, - clock: quartz.NewReal(), - licenseKeys: coderd.Keys, + publisher := &tallymanPublisher{ + ctx: ctx, + ctxCancel: cancel, + log: log, + db: db, + licenseKeys: keys, + done: make(chan struct{}), + + ingestURL: tallymanIngestURLV1, + httpClient: http.DefaultClient, + clock: quartz.NewReal(), } for _, opt := range opts { opt(publisher) @@ -92,6 +94,9 @@ type TallymanPublisherOption func(*tallymanPublisher) // PublisherWithHTTPClient sets the HTTP client to use for publishing usage events. func PublisherWithHTTPClient(httpClient *http.Client) TallymanPublisherOption { return func(p *tallymanPublisher) { + if httpClient == nil { + httpClient = http.DefaultClient + } p.httpClient = httpClient } } @@ -103,14 +108,6 @@ func PublisherWithClock(clock quartz.Clock) TallymanPublisherOption { } } -// PublisherWithLicenseKeys sets the license public keys to use for license -// validation. -func PublisherWithLicenseKeys(keys map[string]ed25519.PublicKey) TallymanPublisherOption { - return func(p *tallymanPublisher) { - p.licenseKeys = keys - } -} - // PublisherWithIngestURL sets the ingest URL to use for publishing usage // events. func PublisherWithIngestURL(ingestURL string) TallymanPublisherOption { @@ -149,6 +146,10 @@ func (p *tallymanPublisher) Start() error { p.initialDelay = tallymanPublishInitialMinimumDelay + time.Duration(plusDelay) } + if len(p.licenseKeys) == 0 { + return xerrors.New("no license keys provided") + } + pproflabel.Go(ctx, pproflabel.Service(pproflabel.ServiceTallymanPublisher), func(ctx context.Context) { p.publishLoop(ctx, deploymentUUID) }) diff --git a/enterprise/coderd/usage/publisher_test.go b/enterprise/coderd/usage/publisher_test.go index a2a997b032..7a17935a64 100644 --- a/enterprise/coderd/usage/publisher_test.go +++ b/enterprise/coderd/usage/publisher_test.go @@ -10,16 +10,20 @@ import ( "time" "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/goleak" "go.uber.org/mock/gomock" "cdr.dev/slog/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/rbac" agplusage "github.com/coder/coder/v2/coderd/usage" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" "github.com/coder/coder/v2/enterprise/coderd/usage" @@ -40,6 +44,7 @@ func TestIntegration(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) log := slogtest.Make(t, nil) db, _ := dbtestutil.NewDB(t) + clock := quartz.NewMock(t) deploymentID, licenseJWT := configureDeployment(ctx, t, db) now := time.Now() @@ -60,7 +65,7 @@ func TestIntegration(t *testing.T) { return handler(req) })) - inserter := usage.NewInserter( + inserter := usage.NewDBInserter( usage.InserterWithClock(clock), ) // Insert an old event that should never be published. @@ -80,10 +85,12 @@ func TestIntegration(t *testing.T) { require.NoErrorf(t, err, "collecting event %d", i) } - publisher := usage.NewTallymanPublisher(ctx, log, db, + // Wrap the publisher's DB in a dbauthz to ensure that the publisher has + // enough permissions. + authzDB := dbauthz.New(db, rbac.NewAuthorizer(prometheus.NewRegistry()), log, coderdtest.AccessControlStorePointer()) + publisher := usage.NewTallymanPublisher(ctx, log, authzDB, coderdenttest.Keys, usage.PublisherWithClock(clock), usage.PublisherWithIngestURL(ingestURL), - usage.PublisherWithLicenseKeys(coderdenttest.Keys), ) defer publisher.Close() @@ -212,10 +219,9 @@ func TestPublisherNoEligibleLicenses(t *testing.T) { } })) - publisher := usage.NewTallymanPublisher(ctx, log, db, + publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys, usage.PublisherWithClock(clock), usage.PublisherWithIngestURL(ingestURL), - usage.PublisherWithLicenseKeys(coderdenttest.Keys), ) defer publisher.Close() @@ -283,14 +289,13 @@ func TestPublisherClaimExpiry(t *testing.T) { return tallymanAcceptAllHandler(req) })) - inserter := usage.NewInserter( + inserter := usage.NewDBInserter( usage.InserterWithClock(clock), ) - publisher := usage.NewTallymanPublisher(ctx, log, db, + publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys, usage.PublisherWithClock(clock), usage.PublisherWithIngestURL(ingestURL), - usage.PublisherWithLicenseKeys(coderdenttest.Keys), usage.PublisherWithInitialDelay(17*time.Minute), ) defer publisher.Close() @@ -367,10 +372,9 @@ func TestPublisherMissingEvents(t *testing.T) { } })) - publisher := usage.NewTallymanPublisher(ctx, log, db, + publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys, usage.PublisherWithClock(clock), usage.PublisherWithIngestURL(ingestURL), - usage.PublisherWithLicenseKeys(coderdenttest.Keys), ) // Expect the publisher to call SelectUsageEventsForPublishing, followed by @@ -510,10 +514,9 @@ func TestPublisherLicenseSelection(t *testing.T) { return tallymanAcceptAllHandler(req) })) - publisher := usage.NewTallymanPublisher(ctx, log, db, + publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys, usage.PublisherWithClock(clock), usage.PublisherWithIngestURL(ingestURL), - usage.PublisherWithLicenseKeys(coderdenttest.Keys), ) defer publisher.Close() @@ -579,10 +582,9 @@ func TestPublisherTallymanError(t *testing.T) { } })) - publisher := usage.NewTallymanPublisher(ctx, log, db, + publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys, usage.PublisherWithClock(clock), usage.PublisherWithIngestURL(ingestURL), - usage.PublisherWithLicenseKeys(coderdenttest.Keys), ) defer publisher.Close() diff --git a/enterprise/coderd/userauth_test.go b/enterprise/coderd/userauth_test.go index 46207f319d..fd4706a25e 100644 --- a/enterprise/coderd/userauth_test.go +++ b/enterprise/coderd/userauth_test.go @@ -941,7 +941,6 @@ func TestGroupSync(t *testing.T) { require.NoError(t, err) } - // nolint:gocritic _, err := runner.API.Database.UpdateUserLoginType(dbauthz.AsSystemRestricted(ctx), database.UpdateUserLoginTypeParams{ NewLoginType: database.LoginTypeOIDC, UserID: user.ID, diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index f49e135ad5..f39b090ca2 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -462,7 +462,6 @@ func TestWorkspaceSerialization(t *testing.T) { // +------------------------------+------------------+ // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -520,7 +519,6 @@ func TestWorkspaceSerialization(t *testing.T) { // +------------------------------+------------------+ // Works! ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -589,7 +587,6 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // pq: could not serialize access due to concurrent update ctx := testutil.Context(t, testutil.WaitShort) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -642,7 +639,6 @@ func TestWorkspaceSerialization(t *testing.T) { // | CommitTx() | | // +---------------------+----------------------------------+ ctx := testutil.Context(t, testutil.WaitShort) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -686,7 +682,6 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+----------------------------------+ // Works! ctx := testutil.Context(t, testutil.WaitShort) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) var err error @@ -741,7 +736,6 @@ func TestWorkspaceSerialization(t *testing.T) { // | | CommitTx() | // +---------------------+---------------------+ ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -799,7 +793,6 @@ func TestWorkspaceSerialization(t *testing.T) { // | | CommitTx() | // +---------------------+---------------------+ ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ @@ -860,7 +853,6 @@ func TestWorkspaceSerialization(t *testing.T) { // +---------------------+---------------------+ // pq: could not serialize access due to read/write dependencies among transactions ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // testing ctx = dbauthz.AsSystemRestricted(ctx) myWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ diff --git a/enterprise/coderd/workspaces_test.go b/enterprise/coderd/workspaces_test.go index dc44a8794e..1cdcd9fb43 100644 --- a/enterprise/coderd/workspaces_test.go +++ b/enterprise/coderd/workspaces_test.go @@ -570,7 +570,6 @@ func TestCreateUserWorkspace(t *testing.T) { return a }).Do() - // nolint:gocritic // this is a test ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong)) agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(r.AgentToken)) require.NoError(t, err) @@ -1708,7 +1707,6 @@ func TestWorkspaceAutobuild(t *testing.T) { // We want to test the database nullifies the NextStartAt so we // make a raw DB call here. We pass in NextStartAt here so we // can test the database will nullify it and not us. - //nolint: gocritic // We need system context to modify this. err = db.UpdateWorkspaceAutostart(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceAutostartParams{ ID: ws.ID, AutostartSchedule: sql.NullString{Valid: true, String: sched.String()}, @@ -2720,7 +2718,6 @@ func TestPrebuildUpdateLifecycleParams(t *testing.T) { }).Do() // Mark the prebuilt workspace's agent as ready so the prebuild can be claimed - // nolint:gocritic ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong)) agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken)) require.NoError(t, err) @@ -3722,7 +3719,6 @@ func TestWorkspaceByOwnerAndName(t *testing.T) { require.Equal(t, ws.LatestBuild.MatchedProvisioners.Available, 0) // Verify that the provisioner daemon is registered in the database - //nolint:gocritic // unit testing daemons, err := db.GetProvisionerDaemons(dbauthz.AsSystemRestricted(ctx)) require.NoError(t, err) require.Equal(t, 1, len(daemons)) @@ -3758,7 +3754,6 @@ func TestWorkspaceByOwnerAndName(t *testing.T) { ctx = testutil.Context(t, testutil.WaitLong) // Reset the context to avoid timeouts. - // nolint:gocritic // unit testing daemons, err := db.GetProvisionerDaemons(dbauthz.AsSystemRestricted(ctx)) require.NoError(t, err) require.Equal(t, len(daemons), 1) @@ -3768,8 +3763,6 @@ func TestWorkspaceByOwnerAndName(t *testing.T) { require.NoError(t, err) // Simulate it's subsequent deletion from the database: - - // nolint:gocritic // unit testing _, err = db.UpsertProvisionerDaemon(dbauthz.AsSystemRestricted(ctx), database.UpsertProvisionerDaemonParams{ Name: daemons[0].Name, OrganizationID: daemons[0].OrganizationID, @@ -3787,7 +3780,6 @@ func TestWorkspaceByOwnerAndName(t *testing.T) { }, }) require.NoError(t, err) - // nolint:gocritic // unit testing err = db.DeleteOldProvisionerDaemons(dbauthz.AsSystemRestricted(ctx)) require.NoError(t, err) @@ -3798,7 +3790,6 @@ func TestWorkspaceByOwnerAndName(t *testing.T) { require.Equal(t, workspace.LatestBuild.MatchedProvisioners.Count, 0) require.Equal(t, workspace.LatestBuild.MatchedProvisioners.Available, 0) - // nolint:gocritic // unit testing _, err = client.WorkspaceByOwnerAndName(dbauthz.As(ctx, userSubject), username, workspace.Name, codersdk.WorkspaceOptions{}) require.NoError(t, err) require.Equal(t, workspace.LatestBuild.Status, codersdk.WorkspaceStatusPending) @@ -3835,7 +3826,6 @@ func TestWorkspaceByOwnerAndName(t *testing.T) { ctx = testutil.Context(t, testutil.WaitLong) // Reset the context to avoid timeouts. - // nolint:gocritic // unit testing daemons, err := db.GetProvisionerDaemons(dbauthz.AsSystemRestricted(ctx)) require.NoError(t, err) require.Equal(t, len(daemons), 1) @@ -3844,7 +3834,6 @@ func TestWorkspaceByOwnerAndName(t *testing.T) { err = closer.Close() require.NoError(t, err) - // nolint:gocritic // unit testing _, err = db.UpsertProvisionerDaemon(dbauthz.AsSystemRestricted(ctx), database.UpsertProvisionerDaemonParams{ Name: daemons[0].Name, OrganizationID: daemons[0].OrganizationID, diff --git a/scripts/rules.go b/scripts/rules.go index f15582d12a..dce029a102 100644 --- a/scripts/rules.go +++ b/scripts/rules.go @@ -37,7 +37,9 @@ func dbauthzAuthorizationContext(m dsl.Matcher) { Where( m["c"].Type.Implements("context.Context") && // Only report on functions that start with "As". - m["f"].Text.Matches("^As"), + m["f"].Text.Matches("^As") && + // Ignore test usages of dbauthz contexts. + !m.File().Name.Matches(`_test\.go$`), ). // Instructions for fixing the lint error should be included on the dangerous function. Report("Using '$f' is dangerous and should be accompanied by a comment explaining why it's ok and a nolint.")