mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user