chore: replace GetManagedAgentCount query with aggregate table (#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 https://github.com/coder/internal/issues/943
This commit is contained in:
Dean Sheather
2025-08-30 03:39:37 +10:00
committed by GitHub
parent 433f9c4a38
commit 39bf3ba628
18 changed files with 488 additions and 116 deletions
+106
View File
@@ -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))
}
}