Skip to content

Commit 39bf3ba

Browse files
deansheatherdeansheather
authored
chore: replace GetManagedAgentCount query with aggregate table (coder#19636)
- Removes GetManagedAgentCount query - Adds new table `usage_events_daily` which stores aggregated usage events by the type and UTC day - Adds trigger to update the values in this table when a new row is inserted into `usage_events` - Adds a migration that adds `usage_events_daily` rows for existing data in `usage_events` - Adds tests for the trigger - Adds tests for the backfill query in the migration Since the `usage_events` table is unreleased currently, this migration will do nothing on real deployments and will only affect preview deployments such as dogfood. Closes coder/internal#943
1 parent 433f9c4 commit 39bf3ba

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
return q.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-
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceWorkspace); err != nil {
2258-
return 0, xerrors.Errorf("authorize read all workspaces: %w", err)
2259-
}
2260-
return q.db.GetManagedAgentCount(ctx, arg)
2261-
}
2262-
22632255
func (q *querier) GetNotificationMessagesByStatus(ctx context.Context, arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage, error) {
22642256
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceNotificationMessage); err != nil {
22652257
return nil, err
@@ -3058,6 +3050,13 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe
30583050
return q.db.GetAuthorizedTemplates(ctx, arg, prep)
30593051
}
30603052

3053+
func (q *querier) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) {
3054+
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil {
3055+
return 0, err
3056+
}
3057+
return q.db.GetTotalUsageDCManagedAgentsV1(ctx, arg)
3058+
}
3059+
30613060
func (q *querier) GetUnexpiredLicenses(ctx context.Context) ([]database.License, error) {
30623061
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceLicense); err != nil {
30633062
return nil, 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 about customizing 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 about customizing 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 about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000362_aggregate_usage_events.down.sql

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
DROP TRIGGER IF EXISTS trigger_aggregate_usage_event ON usage_events;
2+
DROP FUNCTION IF EXISTS aggregate_usage_event();
3+
DROP TABLE IF EXISTS usage_events_daily;

‎coderd/database/migrations/000362_aggregate_usage_events.up.sql

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
CREATE TABLE usage_events_daily (
2+
day date NOT NULL, -- always grouped by day in UTC
3+
event_type text NOT NULL,
4+
usage_data jsonb NOT 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 REPLACE FUNCTION aggregate_usage_event()
13+
RETURNS TRIGGER AS $$
14+
BEGIN
15+
-- Check for supported event types and throw error for unknown types
16+
IF NEW.event_type NOT IN ('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 AT TIME ZONE 'UTC')::date,
25+
NEW.event_type,
26+
NEW.event_data
27+
)
28+
ON CONFLICT (day, event_type) DO UPDATE SET
29+
usage_data = CASE
30+
-- Handle simple counter events by summing the count
31+
WHEN NEW.event_type IN ('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+
CREATE TRIGGER trigger_aggregate_usage_event
45+
AFTER INSERT ON 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 AT TIME ZONE 'UTC')::date AS 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 AT TIME ZONE 'UTC')::date,
63+
event_type
64+
ON CONFLICT (day, event_type) DO UPDATE SET
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+
func TestMigration000362AggregateUsageEvents(t *testing.T) {
373+
t.Parallel()
374+
375+
const migrationVersion = 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+
if version == 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 := range usageEvents {
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+
defer rows.Close()
454+
var out []database.UsageEventsDaily
455+
for rows.Next() {
456+
var row 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+
for i, row := range out {
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)