Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit39bf3ba

Browse files
authored
chore: replace GetManagedAgentCount query with aggregate table (#19636)
- Removes GetManagedAgentCount query- Adds new table `usage_events_daily` which stores aggregated usageevents by the type and UTC day- Adds trigger to update the values in this table when a new row isinserted into `usage_events`- Adds a migration that adds `usage_events_daily` rows for existing datain `usage_events`- Adds tests for the trigger- Adds tests for the backfill query in the migrationSince the `usage_events` table is unreleased currently, this migrationwill do nothing on real deployments and will only affect previewdeployments such as dogfood.Closescoder/internal#943
1 parent433f9c4 commit39bf3ba

18 files changed

+488
-116
lines changed

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2252,14 +2252,6 @@ func (q *querier) GetLogoURL(ctx context.Context) (string, error) {
22522252
returnq.db.GetLogoURL(ctx)
22532253
}
22542254

2255-
func (q*querier)GetManagedAgentCount(ctx context.Context,arg database.GetManagedAgentCountParams) (int64,error) {
2256-
// Must be able to read all workspaces to check usage.
2257-
iferr:=q.authorizeContext(ctx,policy.ActionRead,rbac.ResourceWorkspace);err!=nil {
2258-
return0,xerrors.Errorf("authorize read all workspaces: %w",err)
2259-
}
2260-
returnq.db.GetManagedAgentCount(ctx,arg)
2261-
}
2262-
22632255
func (q*querier)GetNotificationMessagesByStatus(ctx context.Context,arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage,error) {
22642256
iferr:=q.authorizeContext(ctx,policy.ActionRead,rbac.ResourceNotificationMessage);err!=nil {
22652257
returnnil,err
@@ -3058,6 +3050,13 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe
30583050
returnq.db.GetAuthorizedTemplates(ctx,arg,prep)
30593051
}
30603052

3053+
func (q*querier)GetTotalUsageDCManagedAgentsV1(ctx context.Context,arg database.GetTotalUsageDCManagedAgentsV1Params) (int64,error) {
3054+
iferr:=q.authorizeContext(ctx,policy.ActionRead,rbac.ResourceUsageEvent);err!=nil {
3055+
return0,err
3056+
}
3057+
returnq.db.GetTotalUsageDCManagedAgentsV1(ctx,arg)
3058+
}
3059+
30613060
func (q*querier)GetUnexpiredLicenses(ctx context.Context) ([]database.License,error) {
30623061
iferr:=q.authorizeContext(ctx,policy.ActionRead,rbac.ResourceLicense);err!=nil {
30633062
returnnil,err

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -723,12 +723,6 @@ func (s *MethodTestSuite) TestLicense() {
723723
dbm.EXPECT().GetAnnouncementBanners(gomock.Any()).Return("value",nil).AnyTimes()
724724
check.Args().Asserts().Returns("value")
725725
}))
726-
s.Run("GetManagedAgentCount",s.Mocked(func(dbm*dbmock.MockStore,_*gofakeit.Faker,check*expects) {
727-
start:=dbtime.Now()
728-
end:=start.Add(time.Hour)
729-
dbm.EXPECT().GetManagedAgentCount(gomock.Any(), database.GetManagedAgentCountParams{StartTime:start,EndTime:end}).Return(int64(0),nil).AnyTimes()
730-
check.Args(database.GetManagedAgentCountParams{StartTime:start,EndTime:end}).Asserts(rbac.ResourceWorkspace,policy.ActionRead).Returns(int64(0))
731-
}))
732726
}
733727

734728
func (s*MethodTestSuite)TestOrganization() {
@@ -4284,4 +4278,12 @@ func (s *MethodTestSuite) TestUsageEvents() {
42844278
db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(),params).Return(nil)
42854279
check.Args(params).Asserts(rbac.ResourceUsageEvent,policy.ActionUpdate)
42864280
}))
4281+
4282+
s.Run("GetTotalUsageDCManagedAgentsV1",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4283+
db.EXPECT().GetTotalUsageDCManagedAgentsV1(gomock.Any(),gomock.Any()).Return(int64(1),nil)
4284+
check.Args(database.GetTotalUsageDCManagedAgentsV1Params{
4285+
StartDate: time.Time{},
4286+
EndDate: time.Time{},
4287+
}).Asserts(rbac.ResourceUsageEvent,policy.ActionRead)
4288+
}))
42874289
}

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 7 additions & 7 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 15 additions & 15 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/dump.sql‎

Lines changed: 47 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
DROPTRIGGER IF EXISTS trigger_aggregate_usage_eventON usage_events;
2+
DROPFUNCTION IF EXISTS aggregate_usage_event();
3+
DROPTABLE IF EXISTS usage_events_daily;
Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
CREATETABLEusage_events_daily (
2+
daydateNOT NULL,-- always grouped by day in UTC
3+
event_typetextNOT NULL,
4+
usage_data jsonbNOT NULL,
5+
PRIMARY KEY (day, event_type)
6+
);
7+
8+
COMMENT ON TABLE usage_events_daily IS'usage_events_daily is a daily rollup of usage events. It stores the total usage for each event type by day.';
9+
COMMENT ON COLUMN usage_events_daily.day IS'The date of the summed usage events, always in UTC.';
10+
11+
-- Function to handle usage event aggregation
12+
CREATE OR REPLACEFUNCTIONaggregate_usage_event()
13+
RETURNS TRIGGERAS $$
14+
BEGIN
15+
-- Check for supported event types and throw error for unknown types
16+
IFNEW.event_type NOTIN ('dc_managed_agents_v1') THEN
17+
RAISE EXCEPTION'Unhandled usage event type in aggregate_usage_event: %',NEW.event_type;
18+
END IF;
19+
20+
INSERT INTO usage_events_daily (day, event_type, usage_data)
21+
VALUES (
22+
-- Extract the date from the created_at timestamp, always using UTC for
23+
-- consistency
24+
date_trunc('day',NEW.created_at ATTIME ZONE'UTC')::date,
25+
NEW.event_type,
26+
NEW.event_data
27+
)
28+
ON CONFLICT (day, event_type) DOUPDATESET
29+
usage_data= CASE
30+
-- Handle simple counter events by summing the count
31+
WHENNEW.event_typeIN ('dc_managed_agents_v1') THEN
32+
jsonb_build_object(
33+
'count',
34+
COALESCE((usage_events_daily.usage_data->>'count')::bigint,0)+
35+
COALESCE((NEW.event_data->>'count')::bigint,0)
36+
)
37+
END;
38+
39+
RETURN NEW;
40+
END;
41+
$$ LANGUAGE plpgsql;
42+
43+
-- Create trigger to automatically aggregate usage events
44+
CREATETRIGGERtrigger_aggregate_usage_event
45+
AFTER INSERTON usage_events
46+
FOR EACH ROW
47+
EXECUTE FUNCTION aggregate_usage_event();
48+
49+
-- Populate usage_events_daily with existing data
50+
INSERT INTO
51+
usage_events_daily (day, event_type, usage_data)
52+
SELECT
53+
date_trunc('day', created_at ATTIME ZONE'UTC')::dateAS day,
54+
event_type,
55+
jsonb_build_object('count',SUM((event_data->>'count')::bigint))AS usage_data
56+
FROM
57+
usage_events
58+
WHERE
59+
-- The only event type we currently support is dc_managed_agents_v1
60+
event_type='dc_managed_agents_v1'
61+
GROUP BY
62+
date_trunc('day', created_at ATTIME ZONE'UTC')::date,
63+
event_type
64+
ON CONFLICT (day, event_type) DOUPDATESET
65+
usage_data=EXCLUDED.usage_data;

‎coderd/database/migrations/migrate_test.go‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,17 +9,20 @@ import (
99
"slices"
1010
"sync"
1111
"testing"
12+
"time"
1213

1314
"github.com/golang-migrate/migrate/v4"
1415
migratepostgres"github.com/golang-migrate/migrate/v4/database/postgres"
1516
"github.com/golang-migrate/migrate/v4/source"
1617
"github.com/golang-migrate/migrate/v4/source/iofs"
1718
"github.com/golang-migrate/migrate/v4/source/stub"
19+
"github.com/google/uuid"
1820
"github.com/lib/pq"
1921
"github.com/stretchr/testify/require"
2022
"go.uber.org/goleak"
2123
"golang.org/x/sync/errgroup"
2224

25+
"github.com/coder/coder/v2/coderd/database"
2326
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2427
"github.com/coder/coder/v2/coderd/database/migrations"
2528
"github.com/coder/coder/v2/testutil"
@@ -363,3 +366,106 @@ func TestMigrateUpWithFixtures(t *testing.T) {
363366
})
364367
}
365368
}
369+
370+
// TestMigration000362AggregateUsageEvents tests the migration that aggregates
371+
// usage events into daily rows correctly.
372+
funcTestMigration000362AggregateUsageEvents(t*testing.T) {
373+
t.Parallel()
374+
375+
constmigrationVersion=362
376+
377+
// Similarly to the other test, this test will probably time out in CI.
378+
ctx:=testutil.Context(t,testutil.WaitSuperLong)
379+
380+
sqlDB:=testSQLDB(t)
381+
db:=database.New(sqlDB)
382+
383+
// Migrate up to the migration before the one that aggregates usage events.
384+
next,err:=migrations.Stepper(sqlDB)
385+
require.NoError(t,err)
386+
for {
387+
version,more,err:=next()
388+
require.NoError(t,err)
389+
if!more {
390+
t.Fatalf("migration %d not found",migrationVersion)
391+
}
392+
ifversion==migrationVersion-1 {
393+
break
394+
}
395+
}
396+
397+
locSydney,err:=time.LoadLocation("Australia/Sydney")
398+
require.NoError(t,err)
399+
400+
usageEvents:= []struct {
401+
// The only possible event type is dc_managed_agents_v1 when this
402+
// migration gets applied.
403+
eventData []byte
404+
createdAt time.Time
405+
}{
406+
{
407+
eventData: []byte(`{"count": 41}`),
408+
createdAt:time.Date(2025,1,1,0,0,0,0,time.UTC),
409+
},
410+
{
411+
eventData: []byte(`{"count": 1}`),
412+
// 2025-01-01 in UTC
413+
createdAt:time.Date(2025,1,2,8,38,57,0,locSydney),
414+
},
415+
{
416+
eventData: []byte(`{"count": 1}`),
417+
createdAt:time.Date(2025,1,2,0,0,0,0,time.UTC),
418+
},
419+
}
420+
expectedDailyRows:= []struct {
421+
day time.Time
422+
usageData []byte
423+
}{
424+
{
425+
day:time.Date(2025,1,1,0,0,0,0,time.UTC),
426+
usageData: []byte(`{"count": 42}`),
427+
},
428+
{
429+
day:time.Date(2025,1,2,0,0,0,0,time.UTC),
430+
usageData: []byte(`{"count": 1}`),
431+
},
432+
}
433+
434+
for_,usageEvent:=rangeusageEvents {
435+
err:=db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
436+
ID:uuid.New().String(),
437+
EventType:"dc_managed_agents_v1",
438+
EventData:usageEvent.eventData,
439+
CreatedAt:usageEvent.createdAt,
440+
})
441+
require.NoError(t,err)
442+
}
443+
444+
// Migrate up to the migration that aggregates usage events.
445+
version,_,err:=next()
446+
require.NoError(t,err)
447+
require.EqualValues(t,migrationVersion,version)
448+
449+
// Get all of the newly created daily rows. This query is not exposed in the
450+
// querier interface intentionally.
451+
rows,err:=sqlDB.QueryContext(ctx,"SELECT day, event_type, usage_data FROM usage_events_daily ORDER BY day ASC")
452+
require.NoError(t,err,"perform query")
453+
deferrows.Close()
454+
varout []database.UsageEventsDaily
455+
forrows.Next() {
456+
varrow database.UsageEventsDaily
457+
err:=rows.Scan(&row.Day,&row.EventType,&row.UsageData)
458+
require.NoError(t,err,"scan row")
459+
out=append(out,row)
460+
}
461+
462+
// Verify that the daily rows match our expectations.
463+
require.Len(t,out,len(expectedDailyRows))
464+
fori,row:=rangeout {
465+
require.Equal(t,"dc_managed_agents_v1",row.EventType)
466+
// The read row might be `+0000` rather than `UTC` specifically, so just
467+
// ensure it's within 1 second of the expected time.
468+
require.WithinDuration(t,expectedDailyRows[i].day,row.Day,time.Second)
469+
require.JSONEq(t,string(expectedDailyRows[i].usageData),string(row.UsageData))
470+
}
471+
}

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp