chore: replace GetManagedAgentCount query with aggregate table

- 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` 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.
coder · deansheather · Aug 29, 2025 · Aug 29, 2025 · Aug 29, 2025 · Aug 29, 2025
commit 4b05f9743635e2e1a9d40802893d0b84cf4e38a8
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -2252,14 +2252,6 @@ func (q *querier) GetLogoURL(ctx context.Context) (string, error) {
 	return q.db.GetLogoURL(ctx)
 }
 
-func (q *querier) GetManagedAgentCount(ctx context.Context, arg database.GetManagedAgentCountParams) (int64, error) {
-	// Must be able to read all workspaces to check usage.
-	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceWorkspace); err != nil {
-		return 0, xerrors.Errorf("authorize read all workspaces: %w", err)
-	}
-	return q.db.GetManagedAgentCount(ctx, arg)
-}
-
 func (q *querier) GetNotificationMessagesByStatus(ctx context.Context, arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage, error) {
 	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceNotificationMessage); err != nil {
 		return nil, err
@@ -3058,6 +3050,13 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe
 	return q.db.GetAuthorizedTemplates(ctx, arg, prep)
 }
 
+func (q *querier) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil {
+		return 0, err
+	}
+	return q.db.GetTotalUsageDCManagedAgentsV1(ctx, arg)
+}
+
 func (q *querier) GetUnexpiredLicenses(ctx context.Context) ([]database.License, error) {
 	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceLicense); err != nil {
 		return nil, err

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -723,12 +723,6 @@ func (s *MethodTestSuite) TestLicense() {
 		dbm.EXPECT().GetAnnouncementBanners(gomock.Any()).Return("value", nil).AnyTimes()
 		check.Args().Asserts().Returns("value")
 	}))
-	s.Run("GetManagedAgentCount", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
-		start := dbtime.Now()
-		end := start.Add(time.Hour)
-		dbm.EXPECT().GetManagedAgentCount(gomock.Any(), database.GetManagedAgentCountParams{StartTime: start, EndTime: end}).Return(int64(0), nil).AnyTimes()
-		check.Args(database.GetManagedAgentCountParams{StartTime: start, EndTime: end}).Asserts(rbac.ResourceWorkspace, policy.ActionRead).Returns(int64(0))
-	}))
 }
 
 func (s *MethodTestSuite) TestOrganization() {
@@ -4403,4 +4397,12 @@ func (s *MethodTestSuite) TestUsageEvents() {
 		db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), params).Return(nil)
 		check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionUpdate)
 	}))
+
+	s.Run("GetTotalUsageDCManagedAgentsV1", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		db.EXPECT().GetTotalUsageDCManagedAgentsV1(gomock.Any(), gomock.Any()).Return(int64(1), nil)
+		check.Args(database.GetTotalUsageDCManagedAgentsV1Params{
+			StartDate: time.Time{},
+			EndDate:   time.Time{},
+		}).Asserts(rbac.ResourceUsageEvent, policy.ActionRead)
+	}))
 }
diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000362_aggregate_usage_events.down.sql b/coderd/database/migrations/000362_aggregate_usage_events.down.sql
@@ -0,0 +1,3 @@
+DROP TRIGGER IF EXISTS trigger_aggregate_usage_event ON usage_events;
+DROP FUNCTION IF EXISTS aggregate_usage_event();
+DROP TABLE IF EXISTS usage_events_daily;
diff --git a/coderd/database/migrations/000362_aggregate_usage_events.up.sql b/coderd/database/migrations/000362_aggregate_usage_events.up.sql
@@ -0,0 +1,65 @@
+CREATE TABLE usage_events_daily (
+  day date NOT NULL, -- always grouped by day in UTC
+  event_type text NOT NULL,
+  usage_data jsonb NOT NULL,
+  PRIMARY KEY (day, event_type)
+);
+
+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.';
+COMMENT ON COLUMN usage_events_daily.day IS 'The date of the summed usage events, always in UTC.';
+
+-- Function to handle usage event aggregation
+CREATE OR REPLACE FUNCTION aggregate_usage_event()
+RETURNS TRIGGER AS $$
+BEGIN
+    -- Check for supported event types and throw error for unknown types
+    IF NEW.event_type NOT IN ('dc_managed_agents_v1') THEN
+        RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type;
+    END IF;
+
+    INSERT INTO usage_events_daily (day, event_type, usage_data)
+    VALUES (
+        -- Extract the date from the created_at timestamp, always using UTC for
+        -- consistency
+        date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date,
+        NEW.event_type,
+        NEW.event_data
+    )
+    ON CONFLICT (day, event_type) DO UPDATE SET
+        usage_data = CASE
+            -- Handle simple counter events by summing the count
+            WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN
+                jsonb_build_object(
+                    'count',
+                    COALESCE((usage_events_daily.usage_data->>'count')::int, 0) +
+                    COALESCE((NEW.event_data->>'count')::int, 0)
+                )
+        END;
+
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Create trigger to automatically aggregate usage events
+CREATE TRIGGER trigger_aggregate_usage_event
+    AFTER INSERT ON usage_events
+    FOR EACH ROW
+    EXECUTE FUNCTION aggregate_usage_event();
+
+-- Populate usage_events_daily with existing data
+INSERT INTO
+    usage_events_daily (day, event_type, usage_data)
+SELECT
+    date_trunc('day', created_at AT TIME ZONE 'UTC')::date AS day,
+    event_type,
+    jsonb_build_object('count', SUM((event_data->>'count')::int)) AS usage_data
+FROM
+    usage_events
+WHERE
+    -- The only event type we currently support is dc_managed_agents_v1
+    event_type = 'dc_managed_agents_v1'
+GROUP BY
+    date_trunc('day', created_at AT TIME ZONE 'UTC')::date,
+    event_type
+ON CONFLICT (day, event_type) DO UPDATE SET
+    usage_data = EXCLUDED.usage_data;
diff --git a/coderd/database/migrations/migrate_test.go b/coderd/database/migrations/migrate_test.go
@@ -9,17 +9,20 @@ import (
 	"slices"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/golang-migrate/migrate/v4"
 	migratepostgres "github.com/golang-migrate/migrate/v4/database/postgres"
 	"github.com/golang-migrate/migrate/v4/source"
 	"github.com/golang-migrate/migrate/v4/source/iofs"
 	"github.com/golang-migrate/migrate/v4/source/stub"
+	"github.com/google/uuid"
 	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 	"golang.org/x/sync/errgroup"
 
+	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	"github.com/coder/coder/v2/coderd/database/migrations"
 	"github.com/coder/coder/v2/testutil"
@@ -363,3 +366,106 @@ func TestMigrateUpWithFixtures(t *testing.T) {
 		})
 	}
 }
+
+// TestMigration000362AggregateUsageEvents tests the migration that aggregates
+// usage events into daily rows correctly.
+func TestMigration000362AggregateUsageEvents(t *testing.T) {
+	t.Parallel()
+
+	const migrationVersion = 362
+
+	// Similarly to the other test, this test will probably time out in CI.
+	ctx := testutil.Context(t, testutil.WaitSuperLong)
+
+	sqlDB := testSQLDB(t)
+	db := database.New(sqlDB)
+
+	// Migrate up to the migration before the one that aggregates usage events.
+	next, err := migrations.Stepper(sqlDB)
+	require.NoError(t, err)
+	for {
+		version, more, err := next()
+		require.NoError(t, err)
+		if !more {
+			t.Fatalf("migration %d not found", migrationVersion)
+		}
+		if version == migrationVersion-1 {
+			break
+		}
+	}
+
+	locSydney, err := time.LoadLocation("Australia/Sydney")
+	require.NoError(t, err)
+
+	usageEvents := []struct {
+		// The only possible event type is dc_managed_agents_v1 when this
+		// migration gets applied.
+		eventData []byte
+		createdAt time.Time
+	}{
+		{
+			eventData: []byte(`{"count": 41}`),
+			createdAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
+		},
+		{
+			eventData: []byte(`{"count": 1}`),
+			// 2025-01-01 in UTC
+			createdAt: time.Date(2025, 1, 2, 8, 38, 57, 0, locSydney),
+		},
+		{
+			eventData: []byte(`{"count": 1}`),
+			createdAt: time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC),
+		},
+	}
+	expectedDailyRows := []struct {
+		day       time.Time
+		usageData []byte
+	}{
+		{
+			day:       time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
+			usageData: []byte(`{"count": 42}`),
+		},
+		{
+			day:       time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC),
+			usageData: []byte(`{"count": 1}`),
+		},
+	}
+
+	for _, usageEvent := range usageEvents {
+		err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
+			ID:        uuid.New().String(),
+			EventType: "dc_managed_agents_v1",
+			EventData: usageEvent.eventData,
+			CreatedAt: usageEvent.createdAt,
+		})
+		require.NoError(t, err)
+	}
+
+	// Migrate up to the migration that aggregates usage events.
+	version, _, err := next()
+	require.NoError(t, err)
+	require.EqualValues(t, migrationVersion, version)
+
+	// Get all of the newly created daily rows. This query is not exposed in the
+	// querier interface intentionally.
+	rows, err := sqlDB.QueryContext(ctx, "SELECT day, event_type, usage_data FROM usage_events_daily ORDER BY day ASC")
+	require.NoError(t, err, "perform query")
+	defer rows.Close()
+	var out []database.UsageEventsDaily
+	for rows.Next() {
+		var row database.UsageEventsDaily
+		err := rows.Scan(&row.Day, &row.EventType, &row.UsageData)
+		require.NoError(t, err, "scan row")
+		out = append(out, row)
+	}
+
+	// Verify that the daily rows match our expectations.
+	require.Len(t, out, len(expectedDailyRows))
+	for i, row := range out {
+		require.Equal(t, "dc_managed_agents_v1", row.EventType)
+		// The read row might be `+0000` rather than `UTC` specifically, so just
+		// ensure it's within 1 second of the expected time.
+		require.WithinDuration(t, expectedDailyRows[i].day, row.Day, time.Second)
+		require.JSONEq(t, string(expectedDailyRows[i].usageData), string(row.UsageData))
+	}
+}