package migrations_test import ( "context" "database/sql" "fmt" "os" "path/filepath" "slices" "strings" "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" ) func TestMain(m *testing.M) { goleak.VerifyTestMain(m, testutil.GoleakOptions...) } func TestMigrate(t *testing.T) { t.Parallel() if testing.Short() { t.SkipNow() return } t.Run("Once", func(t *testing.T) { t.Parallel() db := testSQLDB(t) err := migrations.Up(db) require.NoError(t, err) }) t.Run("Parallel", func(t *testing.T) { t.Parallel() db := testSQLDB(t) eg := errgroup.Group{} eg.Go(func() error { return migrations.Up(db) }) eg.Go(func() error { return migrations.Up(db) }) require.NoError(t, eg.Wait()) }) t.Run("Twice", func(t *testing.T) { t.Parallel() db := testSQLDB(t) err := migrations.Up(db) require.NoError(t, err) err = migrations.Up(db) require.NoError(t, err) }) t.Run("UpDownUp", func(t *testing.T) { t.Parallel() db := testSQLDB(t) err := migrations.Up(db) require.NoError(t, err) err = migrations.Down(db) require.NoError(t, err) err = migrations.Up(db) require.NoError(t, err) }) } func testSQLDB(t testing.TB) *sql.DB { t.Helper() connection, err := dbtestutil.Open(t) require.NoError(t, err) db, err := sql.Open("postgres", connection) require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) // dbtestutil.Open automatically runs migrations, but we want to actually test // migration behavior in this package. _, err = db.Exec(`DROP SCHEMA public CASCADE`) require.NoError(t, err) _, err = db.Exec(`CREATE SCHEMA public`) require.NoError(t, err) return db } // paralleltest linter doesn't correctly handle table-driven tests (https://github.com/kunwardeep/paralleltest/issues/8) // nolint:paralleltest func TestCheckLatestVersion(t *testing.T) { t.Parallel() type test struct { currentVersion uint existingVersions []uint expectedResult string } tests := []test{ // successful cases {1, []uint{1}, ""}, {3, []uint{1, 2, 3}, ""}, {3, []uint{1, 3}, ""}, // failure cases {1, []uint{1, 2}, "current version is 1, but later version 2 exists"}, {2, []uint{1, 2, 3}, "current version is 2, but later version 3 exists"}, {4, []uint{1, 2, 3}, "get previous migration: prev for version 4 : file does not exist"}, {4, []uint{1, 2, 3, 5}, "get previous migration: prev for version 4 : file does not exist"}, } for i, tc := range tests { t.Run(fmt.Sprintf("entry %d", i), func(t *testing.T) { t.Parallel() driver, _ := stub.WithInstance(nil, &stub.Config{}) stub, ok := driver.(*stub.Stub) require.True(t, ok) for _, version := range tc.existingVersions { stub.Migrations.Append(&source.Migration{ Version: version, Identifier: "", Direction: source.Up, Raw: "", }) } err := migrations.CheckLatestVersion(driver, tc.currentVersion) var errMessage string if err != nil { errMessage = err.Error() } require.Equal(t, tc.expectedResult, errMessage) }) } } func setupMigrate(t *testing.T, db *sql.DB, name, path string) (source.Driver, *migrate.Migrate) { t.Helper() ctx := context.Background() conn, err := db.Conn(ctx) require.NoError(t, err) dbDriver, err := migratepostgres.WithConnection(ctx, conn, &migratepostgres.Config{ MigrationsTable: "test_migrate_" + name, }) require.NoError(t, err) dirFS := os.DirFS(path) d, err := iofs.New(dirFS, ".") require.NoError(t, err) t.Cleanup(func() { d.Close() }) m, err := migrate.NewWithInstance(name, d, "", dbDriver) require.NoError(t, err) t.Cleanup(func() { m.Close() }) return d, m } type tableStats struct { mu sync.Mutex s map[string]int } func (s *tableStats) Add(table string, n int) { s.mu.Lock() defer s.mu.Unlock() s.s[table] += n } func (s *tableStats) Empty() []string { s.mu.Lock() defer s.mu.Unlock() var m []string for table, n := range s.s { if n == 0 { m = append(m, table) } } return m } func TestMigrateUpWithFixtures(t *testing.T) { t.Parallel() if testing.Short() { t.SkipNow() return } type testCase struct { name string path string // For determining if test case table stats // are used to determine test coverage. useStats bool } tests := []testCase{ { name: "fixtures", path: filepath.Join("testdata", "fixtures"), useStats: true, }, // More test cases added via glob below. } // Folders in testdata/full_dumps represent fixtures for a full // deployment of Coder. matches, err := filepath.Glob(filepath.Join("testdata", "full_dumps", "*")) require.NoError(t, err) for _, match := range matches { tests = append(tests, testCase{ name: filepath.Base(match), path: match, useStats: true, }) } // These tables are allowed to have zero rows for now, // but we should eventually add fixtures for them. ignoredTablesForStats := []string{ "audit_logs", "external_auth_links", "group_members", "licenses", "replicas", "template_version_parameters", "workspace_build_parameters", "template_version_variables", "dbcrypt_keys", // having zero rows is a valid state for this table "template_version_workspace_tags", "notification_report_generator_logs", } s := &tableStats{s: make(map[string]int)} // This will run after all subtests have run and fail the test if // new tables have been added without covering them with fixtures. t.Cleanup(func() { emptyTables := s.Empty() slices.Sort(emptyTables) for _, table := range ignoredTablesForStats { i := slices.Index(emptyTables, table) if i >= 0 { emptyTables = slices.Delete(emptyTables, i, i+1) } } if len(emptyTables) > 0 { t.Log("The following tables have zero rows, consider adding fixtures for them or create a full database dump:") t.Errorf("tables have zero rows: %v", emptyTables) t.Log("See https://github.com/coder/coder/blob/main/docs/about/contributing/backend.md#database-fixtures-for-testing-migrations for more information") } }) for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() db := testSQLDB(t) // Prepare database for stepping up. err := migrations.Down(db) require.NoError(t, err) // Initialize migrations for fixtures. fDriver, fMigrate := setupMigrate(t, db, tt.name, tt.path) nextStep, err := migrations.Stepper(db) require.NoError(t, err) var fixtureVer uint nextFixtureVer, err := fDriver.First() require.NoError(t, err) for { version, more, err := nextStep() require.NoError(t, err) if !more { // We reached the end of the migrations. break } if nextFixtureVer == version { err = fMigrate.Steps(1) require.NoError(t, err) fixtureVer = version nv, _ := fDriver.Next(nextFixtureVer) if nv > 0 { nextFixtureVer = nv } } t.Logf("migrated to version %d, fixture version %d", version, fixtureVer) } ctx := testutil.Context(t, testutil.WaitSuperLong) // Gather number of rows for all existing tables // at the end of the migrations and fixtures. var tables pq.StringArray err = db.QueryRowContext(ctx, ` SELECT array_agg(tablename) FROM pg_catalog.pg_tables WHERE schemaname != 'information_schema' AND schemaname != 'pg_catalog' AND tablename NOT LIKE 'test_migrate_%' `).Scan(&tables) require.NoError(t, err) for _, table := range tables { var count int err = db.QueryRowContext(ctx, "SELECT COUNT(*) FROM "+table).Scan(&count) require.NoError(t, err) if tt.useStats { s.Add(table, count) } } // Test that migration down is successful after up. err = migrations.Down(db) require.NoError(t, err, "final migration down should be successful") }) } } // TestMigration000362AggregateUsageEvents tests the migration that aggregates // usage events into daily rows correctly. func TestMigration000362AggregateUsageEvents(t *testing.T) { t.Parallel() const migrationVersion = 362 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}`), }, } ctx := testutil.Context(t, testutil.WaitSuperLong) 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)) } } func TestMigration000387MigrateTaskWorkspaces(t *testing.T) { t.Parallel() // This test verifies the migration of task workspaces to the new tasks data model. // Test cases: // // Task 1 (ws1) - Basic case: // - Single build with has_ai_task=true, prompt, and parameters // - Verifies: all task fields are populated correctly // // Task 2 (ws2) - No AI Prompt parameter: // - Single build with has_ai_task=true but NO AI Prompt parameter // - Verifies: prompt defaults to empty string (tests LEFT JOIN for optional prompt) // // Task 3 (ws3) - Latest build is stop: // - Build 1: start with agents/apps and prompt // - Build 2: stop build (references same app via ai_task_sidebar_app_id) // - Verifies: twa uses latest build number with agents/apps from that build's ai_task_sidebar_app_id // // Antagonists - Should NOT be migrated: // - Regular workspace without has_ai_task flag // - Deleted workspace (w.deleted = true) const migrationVersion = 387 sqlDB := testSQLDB(t) // Migrate up to the migration before the task workspace migration. 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 } } now := time.Now().UTC().Truncate(time.Microsecond) deletingAt := now.Add(24 * time.Hour).Truncate(time.Microsecond) // Define all IDs upfront. orgID := uuid.New() userID := uuid.New() templateID := uuid.New() templateVersionID := uuid.New() templateJobID := uuid.New() // Task workspace 1: basic case with prompt and parameters. ws1ID := uuid.New() ws1Build1JobID := uuid.New() ws1Build1ID := uuid.New() ws1Resource1ID := uuid.New() ws1Agent1ID := uuid.New() ws1App1ID := uuid.New() // Task workspace 2: no AI Prompt parameter. ws2ID := uuid.New() ws2Build1JobID := uuid.New() ws2Build1ID := uuid.New() ws2Resource1ID := uuid.New() ws2Agent1ID := uuid.New() ws2App1ID := uuid.New() // Task workspace 3: has both start and stop builds. ws3ID := uuid.New() ws3Build1JobID := uuid.New() ws3Build1ID := uuid.New() ws3Resource1ID := uuid.New() ws3Agent1ID := uuid.New() ws3App1ID := uuid.New() ws3Build2JobID := uuid.New() ws3Build2ID := uuid.New() ws3Resource2ID := uuid.New() // Antagonist 1: deleted workspace. wsAntDeletedID := uuid.New() wsAntDeletedBuild1JobID := uuid.New() wsAntDeletedBuild1ID := uuid.New() wsAntDeletedResource1ID := uuid.New() wsAntDeletedAgent1ID := uuid.New() wsAntDeletedApp1ID := uuid.New() // Antagonist 2: regular workspace without has_ai_task. wsAntID := uuid.New() wsAntBuild1JobID := uuid.New() wsAntBuild1ID := uuid.New() // Create all fixtures in a single transaction. ctx := testutil.Context(t, testutil.WaitSuperLong) tx, err := sqlDB.BeginTx(ctx, nil) require.NoError(t, err) defer tx.Rollback() // Execute fixture setup as individual statements. fixtures := []struct { query string args []any }{ // Setup organization, user, and template. { `INSERT INTO organizations (id, name, display_name, description, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6)`, []any{orgID, "test-org", "Test Org", "Test Org", now, now}, }, { `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{userID, "testuser", "test@example.com", []byte{}, now, now, "active", []byte("{}"), "password"}, }, { `INSERT INTO provisioner_jobs (id, created_at, updated_at, started_at, completed_at, error, organization_id, initiator_id, provisioner, storage_method, file_id, type, input, tags) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, []any{templateJobID, now, now, now, now, "", orgID, userID, "terraform", "file", uuid.New(), "template_version_import", []byte("{}"), []byte("{}")}, }, { `INSERT INTO template_versions (id, organization_id, name, readme, created_at, updated_at, job_id, created_by) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, []any{templateVersionID, orgID, "v1.0", "Test template", now, now, templateJobID, userID}, }, { `INSERT INTO templates (id, organization_id, name, created_at, updated_at, provisioner, active_version_id, created_by) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, []any{templateID, orgID, "test-template", now, now, "terraform", templateVersionID, userID}, }, { `UPDATE template_versions SET template_id = $1 WHERE id = $2`, []any{templateID, templateVersionID}, }, // Task workspace 1 is a normal start build. { `INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, deleted, name, last_used_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{ws1ID, now, now, userID, orgID, templateID, false, "task-ws-1", now}, }, { `INSERT INTO provisioner_jobs (id, created_at, updated_at, started_at, completed_at, error, organization_id, initiator_id, provisioner, storage_method, file_id, type, input, tags) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, []any{ws1Build1JobID, now, now, now, now, "", orgID, userID, "terraform", "file", uuid.New(), "workspace_build", []byte("{}"), []byte("{}")}, }, { `INSERT INTO workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, daily_cost, instance_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{ws1Resource1ID, now, ws1Build1JobID, "start", "docker_container", "main", false, "", 0, ""}, }, { `INSERT INTO workspace_agents (id, created_at, updated_at, name, resource_id, auth_token, architecture, operating_system, directory, connection_timeout_seconds, lifecycle_state, logs_length, logs_overflowed) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)`, []any{ws1Agent1ID, now, now, "agent1", ws1Resource1ID, uuid.New(), "amd64", "linux", "/home/coder", 120, "ready", 0, false}, }, { `INSERT INTO workspace_apps (id, created_at, agent_id, slug, display_name, icon, command, url, subdomain, external) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{ws1App1ID, now, ws1Agent1ID, "code-server", "Code Server", "", "", "http://localhost:8080", false, false}, }, { `INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, has_ai_task, ai_task_sidebar_app_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)`, []any{ws1Build1ID, now, now, ws1ID, templateVersionID, 1, "start", userID, []byte{}, ws1Build1JobID, now.Add(8 * time.Hour), "initiator", 0, now.Add(8 * time.Hour), true, ws1App1ID}, }, { `INSERT INTO workspace_build_parameters (workspace_build_id, name, value) VALUES ($1, $2, $3)`, []any{ws1Build1ID, "AI Prompt", "Build a web server"}, }, { `INSERT INTO workspace_build_parameters (workspace_build_id, name, value) VALUES ($1, $2, $3)`, []any{ws1Build1ID, "region", "us-east-1"}, }, { `INSERT INTO workspace_build_parameters (workspace_build_id, name, value) VALUES ($1, $2, $3)`, []any{ws1Build1ID, "instance_type", "t2.micro"}, }, // Task workspace 2: no AI Prompt parameter (tests LEFT JOIN). { `INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, deleted, name, last_used_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{ws2ID, now, now, userID, orgID, templateID, false, "task-ws-2-no-prompt", now}, }, { `INSERT INTO provisioner_jobs (id, created_at, updated_at, started_at, completed_at, error, organization_id, initiator_id, provisioner, storage_method, file_id, type, input, tags) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, []any{ws2Build1JobID, now, now, now, now, "", orgID, userID, "terraform", "file", uuid.New(), "workspace_build", []byte("{}"), []byte("{}")}, }, { `INSERT INTO workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, daily_cost, instance_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{ws2Resource1ID, now, ws2Build1JobID, "start", "docker_container", "main", false, "", 0, ""}, }, { `INSERT INTO workspace_agents (id, created_at, updated_at, name, resource_id, auth_token, architecture, operating_system, directory, connection_timeout_seconds, lifecycle_state, logs_length, logs_overflowed) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)`, []any{ws2Agent1ID, now, now, "agent2", ws2Resource1ID, uuid.New(), "amd64", "linux", "/home/coder", 120, "ready", 0, false}, }, { `INSERT INTO workspace_apps (id, created_at, agent_id, slug, display_name, icon, command, url, subdomain, external) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{ws2App1ID, now, ws2Agent1ID, "terminal", "Terminal", "", "", "http://localhost:3000", false, false}, }, { `INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, has_ai_task, ai_task_sidebar_app_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)`, []any{ws2Build1ID, now, now, ws2ID, templateVersionID, 1, "start", userID, []byte{}, ws2Build1JobID, now.Add(8 * time.Hour), "initiator", 0, now.Add(8 * time.Hour), true, ws2App1ID}, }, // Note: No AI Prompt parameter for ws2 - this tests the LEFT JOIN for optional prompt. // Task workspace 3: has both start and stop builds. { `INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, deleted, name, last_used_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{ws3ID, now, now, userID, orgID, templateID, false, "task-ws-3-stop", now}, }, { `INSERT INTO provisioner_jobs (id, created_at, updated_at, started_at, completed_at, error, organization_id, initiator_id, provisioner, storage_method, file_id, type, input, tags) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, []any{ws3Build1JobID, now, now, now, now, "", orgID, userID, "terraform", "file", uuid.New(), "workspace_build", []byte("{}"), []byte("{}")}, }, { `INSERT INTO workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, daily_cost, instance_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{ws3Resource1ID, now, ws3Build1JobID, "start", "docker_container", "main", false, "", 0, ""}, }, { `INSERT INTO workspace_agents (id, created_at, updated_at, name, resource_id, auth_token, architecture, operating_system, directory, connection_timeout_seconds, lifecycle_state, logs_length, logs_overflowed) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)`, []any{ws3Agent1ID, now, now, "agent3", ws3Resource1ID, uuid.New(), "amd64", "linux", "/home/coder", 120, "ready", 0, false}, }, { `INSERT INTO workspace_apps (id, created_at, agent_id, slug, display_name, icon, command, url, subdomain, external) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{ws3App1ID, now, ws3Agent1ID, "app3", "App3", "", "", "http://localhost:5000", false, false}, }, { `INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, has_ai_task, ai_task_sidebar_app_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)`, []any{ws3Build1ID, now, now, ws3ID, templateVersionID, 1, "start", userID, []byte{}, ws3Build1JobID, now.Add(8 * time.Hour), "initiator", 0, now.Add(8 * time.Hour), true, ws3App1ID}, }, { `INSERT INTO workspace_build_parameters (workspace_build_id, name, value) VALUES ($1, $2, $3)`, []any{ws3Build1ID, "AI Prompt", "Task with stop build"}, }, { `INSERT INTO provisioner_jobs (id, created_at, updated_at, started_at, completed_at, error, organization_id, initiator_id, provisioner, storage_method, file_id, type, input, tags) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, []any{ws3Build2JobID, now, now, now, now, "", orgID, userID, "terraform", "file", uuid.New(), "workspace_build", []byte("{}"), []byte("{}")}, }, { `INSERT INTO workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, daily_cost, instance_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{ws3Resource2ID, now, ws3Build2JobID, "stop", "docker_container", "main", false, "", 0, ""}, }, { `INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, has_ai_task, ai_task_sidebar_app_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)`, []any{ws3Build2ID, now, now, ws3ID, templateVersionID, 2, "stop", userID, []byte{}, ws3Build2JobID, now.Add(8 * time.Hour), "initiator", 0, now.Add(8 * time.Hour), true, ws3App1ID}, }, // Antagonist 1: deleted workspace. { `INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, deleted, name, last_used_at, deleting_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{wsAntDeletedID, now, now, userID, orgID, templateID, true, "deleted-task-workspace", now, deletingAt}, }, { `INSERT INTO provisioner_jobs (id, created_at, updated_at, started_at, completed_at, error, organization_id, initiator_id, provisioner, storage_method, file_id, type, input, tags) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, []any{wsAntDeletedBuild1JobID, now, now, now, now, "", orgID, userID, "terraform", "file", uuid.New(), "workspace_build", []byte("{}"), []byte("{}")}, }, { `INSERT INTO workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, daily_cost, instance_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{wsAntDeletedResource1ID, now, wsAntDeletedBuild1JobID, "start", "docker_container", "main", false, "", 0, ""}, }, { `INSERT INTO workspace_agents (id, created_at, updated_at, name, resource_id, auth_token, architecture, operating_system, directory, connection_timeout_seconds, lifecycle_state, logs_length, logs_overflowed) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)`, []any{wsAntDeletedAgent1ID, now, now, "agent-deleted", wsAntDeletedResource1ID, uuid.New(), "amd64", "linux", "/home/coder", 120, "ready", 0, false}, }, { `INSERT INTO workspace_apps (id, created_at, agent_id, slug, display_name, icon, command, url, subdomain, external) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, []any{wsAntDeletedApp1ID, now, wsAntDeletedAgent1ID, "app-deleted", "AppDeleted", "", "", "http://localhost:6000", false, false}, }, { `INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, has_ai_task, ai_task_sidebar_app_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)`, []any{wsAntDeletedBuild1ID, now, now, wsAntDeletedID, templateVersionID, 1, "start", userID, []byte{}, wsAntDeletedBuild1JobID, now.Add(8 * time.Hour), "initiator", 0, now.Add(8 * time.Hour), true, wsAntDeletedApp1ID}, }, { `INSERT INTO workspace_build_parameters (workspace_build_id, name, value) VALUES ($1, $2, $3)`, []any{wsAntDeletedBuild1ID, "AI Prompt", "Should not migrate deleted"}, }, // Antagonist 2: regular workspace without has_ai_task. { `INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, deleted, name, last_used_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{wsAntID, now, now, userID, orgID, templateID, false, "regular-workspace", now}, }, { `INSERT INTO provisioner_jobs (id, created_at, updated_at, started_at, completed_at, error, organization_id, initiator_id, provisioner, storage_method, file_id, type, input, tags) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, []any{wsAntBuild1JobID, now, now, now, now, "", orgID, userID, "terraform", "file", uuid.New(), "workspace_build", []byte("{}"), []byte("{}")}, }, { `INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`, []any{wsAntBuild1ID, now, now, wsAntID, templateVersionID, 1, "start", userID, []byte{}, wsAntBuild1JobID, now.Add(8 * time.Hour), "initiator", 0, now.Add(8 * time.Hour)}, }, } for _, fixture := range fixtures { _, err = tx.ExecContext(ctx, fixture.query, fixture.args...) require.NoError(t, err) } err = tx.Commit() require.NoError(t, err) // Run the migration. version, _, err := next() require.NoError(t, err) require.EqualValues(t, migrationVersion, version) // Should have exactly 3 tasks (not antagonists). var taskCount int err = sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM tasks").Scan(&taskCount) require.NoError(t, err) require.Equal(t, 3, taskCount, "should have created 3 tasks from workspaces") // Verify task 1, normal start build. var task1 struct { id uuid.UUID name string workspaceID uuid.UUID templateVersionID uuid.UUID prompt string templateParameters []byte createdAt time.Time deletedAt *time.Time } err = sqlDB.QueryRowContext(ctx, ` SELECT id, name, workspace_id, template_version_id, prompt, template_parameters, created_at, deleted_at FROM tasks WHERE workspace_id = $1 `, ws1ID).Scan(&task1.id, &task1.name, &task1.workspaceID, &task1.templateVersionID, &task1.prompt, &task1.templateParameters, &task1.createdAt, &task1.deletedAt) require.NoError(t, err) require.Equal(t, "task-ws-1", task1.name) require.Equal(t, "Build a web server", task1.prompt) require.JSONEq(t, `{"region":"us-east-1","instance_type":"t2.micro"}`, string(task1.templateParameters)) require.Nil(t, task1.deletedAt) // Verify task_workspace_apps for task 1. var twa1 struct { buildNumber int32 agentID uuid.UUID appID uuid.UUID } err = sqlDB.QueryRowContext(ctx, ` SELECT workspace_build_number, workspace_agent_id, workspace_app_id FROM task_workspace_apps WHERE task_id = $1 `, task1.id).Scan(&twa1.buildNumber, &twa1.agentID, &twa1.appID) require.NoError(t, err) require.Equal(t, int32(1), twa1.buildNumber) require.Equal(t, ws1Agent1ID, twa1.agentID) require.Equal(t, ws1App1ID, twa1.appID) // Verify task 2, no AI Prompt parameter. var task2 struct { id uuid.UUID name string prompt string templateParameters []byte deletedAt *time.Time } err = sqlDB.QueryRowContext(ctx, ` SELECT id, name, prompt, template_parameters, deleted_at FROM tasks WHERE workspace_id = $1 `, ws2ID).Scan(&task2.id, &task2.name, &task2.prompt, &task2.templateParameters, &task2.deletedAt) require.NoError(t, err) require.Equal(t, "task-ws-2-no-prompt", task2.name) require.Equal(t, "", task2.prompt, "prompt should be empty string when no AI Prompt parameter") require.JSONEq(t, `{}`, string(task2.templateParameters), "no parameters") require.Nil(t, task2.deletedAt) // Verify task_workspace_apps for task 2. var twa2 struct { buildNumber int32 agentID uuid.UUID appID uuid.UUID } err = sqlDB.QueryRowContext(ctx, ` SELECT workspace_build_number, workspace_agent_id, workspace_app_id FROM task_workspace_apps WHERE task_id = $1 `, task2.id).Scan(&twa2.buildNumber, &twa2.agentID, &twa2.appID) require.NoError(t, err) require.Equal(t, int32(1), twa2.buildNumber) require.Equal(t, ws2Agent1ID, twa2.agentID) require.Equal(t, ws2App1ID, twa2.appID) // Verify task 3, has both start and stop builds. var task3 struct { id uuid.UUID name string prompt string templateParameters []byte templateVersionID uuid.UUID deletedAt *time.Time } err = sqlDB.QueryRowContext(ctx, ` SELECT id, name, prompt, template_parameters, template_version_id, deleted_at FROM tasks WHERE workspace_id = $1 `, ws3ID).Scan(&task3.id, &task3.name, &task3.prompt, &task3.templateParameters, &task3.templateVersionID, &task3.deletedAt) require.NoError(t, err) require.Equal(t, "task-ws-3-stop", task3.name) require.Equal(t, "Task with stop build", task3.prompt) require.JSONEq(t, `{}`, string(task3.templateParameters), "no other parameters") require.Equal(t, templateVersionID, task3.templateVersionID) require.Nil(t, task3.deletedAt) // Verify task_workspace_apps for task 3 uses latest build and its ai_task_sidebar_app_id. var twa3 struct { buildNumber int32 agentID uuid.UUID appID uuid.UUID } err = sqlDB.QueryRowContext(ctx, ` SELECT workspace_build_number, workspace_agent_id, workspace_app_id FROM task_workspace_apps WHERE task_id = $1 `, task3.id).Scan(&twa3.buildNumber, &twa3.agentID, &twa3.appID) require.NoError(t, err) require.Equal(t, int32(2), twa3.buildNumber, "should use latest build number") require.Equal(t, ws3Agent1ID, twa3.agentID, "should use agent from latest build's ai_task_sidebar_app_id") require.Equal(t, ws3App1ID, twa3.appID, "should use app from latest build's ai_task_sidebar_app_id") // Verify antagonists should NOT be migrated. var antCount int err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM tasks WHERE workspace_id IN ($1, $2) `, wsAntDeletedID, wsAntID).Scan(&antCount) require.NoError(t, err) require.Equal(t, 0, antCount, "antagonist workspaces (deleted and regular) should not be migrated") } func TestMigration000457ChatAccessRole(t *testing.T) { t.Parallel() const migrationVersion = 457 sqlDB := testSQLDB(t) // Migrate up to the migration before the one that grants // agents-access roles. 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 } } ctx := testutil.Context(t, testutil.WaitSuperLong) // Define test users. userWithChat := uuid.New() // Has a chat, no agents-access role. userAlreadyHasRole := uuid.New() // Has a chat and already has agents-access. userNoChat := uuid.New() // No chat at all. userWithChatAndRoles := uuid.New() // Has a chat and other existing roles. now := time.Now().UTC().Truncate(time.Microsecond) // We need a chat_provider and chat_model_config for the chats FK. providerID := uuid.New() modelConfigID := uuid.New() tx, err := sqlDB.BeginTx(ctx, nil) require.NoError(t, err) defer tx.Rollback() fixtures := []struct { query string args []any }{ // Insert test users with varying rbac_roles. { `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{userWithChat, "user-with-chat", "chat@test.com", []byte{}, now, now, "active", pq.StringArray{}, "password"}, }, { `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{userAlreadyHasRole, "user-already-has-role", "already@test.com", []byte{}, now, now, "active", pq.StringArray{"agents-access"}, "password"}, }, { `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{userNoChat, "user-no-chat", "nochat@test.com", []byte{}, now, now, "active", pq.StringArray{}, "password"}, }, { `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{userWithChatAndRoles, "user-with-roles", "roles@test.com", []byte{}, now, now, "active", pq.StringArray{"template-admin"}, "password"}, }, // Insert a chat provider and model config for the chats FK. { `INSERT INTO chat_providers (id, provider, display_name, api_key, enabled, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, $7)`, []any{providerID, "openai", "OpenAI", "", true, now, now}, }, { `INSERT INTO chat_model_configs (id, provider, model, display_name, enabled, context_limit, compression_threshold, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{modelConfigID, "openai", "gpt-4", "GPT 4", true, 100000, 70, now, now}, }, // Insert chats for users A, B, and D (not C). { `INSERT INTO chats (id, owner_id, last_model_config_id, title, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6)`, []any{uuid.New(), userWithChat, modelConfigID, "Chat A", now, now}, }, { `INSERT INTO chats (id, owner_id, last_model_config_id, title, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6)`, []any{uuid.New(), userAlreadyHasRole, modelConfigID, "Chat B", now, now}, }, { `INSERT INTO chats (id, owner_id, last_model_config_id, title, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6)`, []any{uuid.New(), userWithChatAndRoles, modelConfigID, "Chat D", now, now}, }, } for i, f := range fixtures { _, err := tx.ExecContext(ctx, f.query, f.args...) require.NoError(t, err, "fixture %d", i) } require.NoError(t, tx.Commit()) // Run the migration. version, _, err := next() require.NoError(t, err) require.EqualValues(t, migrationVersion, version) // Helper to get rbac_roles for a user. getRoles := func(t *testing.T, userID uuid.UUID) []string { t.Helper() var roles pq.StringArray err := sqlDB.QueryRowContext(ctx, "SELECT rbac_roles FROM users WHERE id = $1", userID, ).Scan(&roles) require.NoError(t, err) return roles } // Verify: user with chat gets agents-access. roles := getRoles(t, userWithChat) require.Contains(t, roles, "agents-access", "user with chat should get agents-access") // Verify: user who already had agents-access has no duplicate. roles = getRoles(t, userAlreadyHasRole) count := 0 for _, r := range roles { if r == "agents-access" { count++ } } require.Equal(t, 1, count, "user who already had agents-access should not get a duplicate") // Verify: user without chat does NOT get agents-access. roles = getRoles(t, userNoChat) require.NotContains(t, roles, "agents-access", "user without chat should not get agents-access") // Verify: user with chat and existing roles gets agents-access // appended while preserving existing roles. roles = getRoles(t, userWithChatAndRoles) require.Contains(t, roles, "agents-access", "user with chat and other roles should get agents-access") require.Contains(t, roles, "template-admin", "existing roles should be preserved") } func TestMigration000475AgentsAccessOrgRole(t *testing.T) { t.Parallel() const migrationVersion = 475 sqlDB := testSQLDB(t) // Migrate up to the migration before 000475. 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 } } ctx := testutil.Context(t, testutil.WaitSuperLong) // Seed: a user with site-level agents-access who is a member of // two orgs, plus a second user who is a member of one org and // does not have the role. userWithRole := uuid.New() userWithoutRole := uuid.New() org1ID := uuid.New() org2ID := uuid.New() now := time.Now().UTC().Truncate(time.Microsecond) tx, err := sqlDB.BeginTx(ctx, nil) require.NoError(t, err) defer tx.Rollback() fixtures := []struct { query string args []any }{ { `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{userWithRole, "user-with-role", "withrole@test.com", []byte{}, now, now, "active", pq.StringArray{"agents-access"}, "password"}, }, { `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, []any{userWithoutRole, "user-without-role", "withoutrole@test.com", []byte{}, now, now, "active", pq.StringArray{}, "password"}, }, { `INSERT INTO organizations (id, name, display_name, description, icon, created_at, updated_at, is_default) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, []any{org1ID, "org-1", "Org 1", "", "", now, now, false}, }, { `INSERT INTO organizations (id, name, display_name, description, icon, created_at, updated_at, is_default) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, []any{org2ID, "org-2", "Org 2", "", "", now, now, false}, }, { `INSERT INTO organization_members (organization_id, user_id, created_at, updated_at, roles) VALUES ($1, $2, $3, $4, $5)`, []any{org1ID, userWithRole, now, now, pq.StringArray{}}, }, { `INSERT INTO organization_members (organization_id, user_id, created_at, updated_at, roles) VALUES ($1, $2, $3, $4, $5)`, []any{org2ID, userWithRole, now, now, pq.StringArray{}}, }, { `INSERT INTO organization_members (organization_id, user_id, created_at, updated_at, roles) VALUES ($1, $2, $3, $4, $5)`, []any{org1ID, userWithoutRole, now, now, pq.StringArray{}}, }, } for i, f := range fixtures { _, err := tx.ExecContext(ctx, f.query, f.args...) require.NoError(t, err, "fixture %d", i) } require.NoError(t, tx.Commit()) // Run migration 000475. version, _, err := next() require.NoError(t, err) require.EqualValues(t, migrationVersion, version) // Verify: userWithRole no longer has agents-access at site level. var siteRoles pq.StringArray err = sqlDB.QueryRowContext(ctx, "SELECT rbac_roles FROM users WHERE id = $1", userWithRole, ).Scan(&siteRoles) require.NoError(t, err) require.NotContains(t, siteRoles, "agents-access", "agents-access should be removed from users.rbac_roles") // Verify: userWithRole has agents-access in both orgs. for _, orgID := range []uuid.UUID{org1ID, org2ID} { var orgRoles pq.StringArray err = sqlDB.QueryRowContext(ctx, "SELECT roles FROM organization_members WHERE user_id = $1 AND organization_id = $2", userWithRole, orgID, ).Scan(&orgRoles) require.NoError(t, err) require.Contains(t, orgRoles, "agents-access", "agents-access should be granted in org %s", orgID) } // Verify: userWithoutRole did not gain agents-access. var orgRoles pq.StringArray err = sqlDB.QueryRowContext(ctx, "SELECT roles FROM organization_members WHERE user_id = $1 AND organization_id = $2", userWithoutRole, org1ID, ).Scan(&orgRoles) require.NoError(t, err) require.NotContains(t, orgRoles, "agents-access", "agents-access should not be granted to a user who didn't have it") // Verify: no DB row exists for agents-access as a custom_role. // The role is now a builtin, resolved in Go via RoleByName. var customRoleCount int err = sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM custom_roles WHERE name = 'agents-access'", ).Scan(&customRoleCount) require.NoError(t, err) require.Equal(t, 0, customRoleCount, "no custom_roles row should exist for agents-access") // Verify: creating a new organization does NOT insert an // agents-access custom_role via the trigger. It should only // insert organization-member and organization-service-account. newOrgID := uuid.New() _, err = sqlDB.ExecContext(ctx, `INSERT INTO organizations (id, name, display_name, description, icon, created_at, updated_at, is_default) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, newOrgID, "new-org", "New Org", "", "", now, now, false, ) require.NoError(t, err) rows, err := sqlDB.QueryContext(ctx, "SELECT name FROM custom_roles WHERE organization_id = $1 AND is_system = true ORDER BY name", newOrgID, ) require.NoError(t, err) defer rows.Close() var gotRoleNames []string for rows.Next() { var name string require.NoError(t, rows.Scan(&name)) gotRoleNames = append(gotRoleNames, name) } require.NoError(t, rows.Err()) require.ElementsMatch(t, []string{"organization-member", "organization-service-account"}, gotRoleNames, "trigger should only create org-member and org-service-account system roles", ) } func TestMigration000504AIProvidersBackfill(t *testing.T) { t.Parallel() const migrationVersion = 504 sqlDB := testSQLDB(t) 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 } } ctx := testutil.Context(t, testutil.WaitSuperLong) now := time.Now().UTC().Truncate(time.Microsecond) userID := uuid.New() openAIProviderID := uuid.New() anthropicProviderID := uuid.New() openAIUserKeyID := uuid.New() anthropicUserKeyID := uuid.New() openAIModelConfigID := uuid.New() anthropicModelConfigID := uuid.New() tx, err := sqlDB.BeginTx(ctx, nil) require.NoError(t, err) defer tx.Rollback() _, err = tx.ExecContext(ctx, `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, userID, "ai-provider-backfill", "ai-provider-backfill@test.com", []byte{}, now, now, "active", pq.StringArray{}, "password", ) require.NoError(t, err) _, err = tx.ExecContext(ctx, ` INSERT INTO chat_providers (id, provider, display_name, api_key, enabled, base_url, created_at, updated_at) VALUES ($1, 'openai', 'OpenAI', 'sk-provider-openai', TRUE, 'https://api.openai.example.com/v1', $3, $3), ($2, 'anthropic', '', '', FALSE, '', $3, $3) `, openAIProviderID, anthropicProviderID, now) require.NoError(t, err) _, err = tx.ExecContext(ctx, ` INSERT INTO user_chat_provider_keys (id, user_id, chat_provider_id, api_key, created_at, updated_at) VALUES ($1, $3, $4, 'sk-user-openai', $6, $6), ($2, $3, $5, 'sk-user-anthropic', $6, $6) `, openAIUserKeyID, anthropicUserKeyID, userID, openAIProviderID, anthropicProviderID, now) require.NoError(t, err) _, err = tx.ExecContext(ctx, ` INSERT INTO chat_model_configs (id, provider, model, display_name, enabled, context_limit, compression_threshold, created_at, updated_at) VALUES ($1, 'openai', 'gpt-4', 'GPT 4', TRUE, 100000, 70, $3, $3), ($2, 'anthropic', 'claude-3-5-sonnet-latest', 'Claude 3.5 Sonnet', TRUE, 200000, 70, $3, $3) `, openAIModelConfigID, anthropicModelConfigID, now) require.NoError(t, err) require.NoError(t, tx.Commit()) var preBackfillCount int err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM ai_providers WHERE id IN ($1, $2) `, openAIProviderID, anthropicProviderID).Scan(&preBackfillCount) require.NoError(t, err) require.Zero(t, preBackfillCount, "test setup should start before the legacy chat providers are backfilled") var preBackfillModelConfigCount int err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM chat_model_configs WHERE id IN ($1, $2) AND ai_provider_id IS NOT NULL `, openAIModelConfigID, anthropicModelConfigID).Scan(&preBackfillModelConfigCount) require.NoError(t, err) require.Zero(t, preBackfillModelConfigCount, "test setup should start before model configs point at AI providers") version, more, err := next() require.NoError(t, err) require.True(t, more) require.EqualValues(t, migrationVersion, version) assertBackfilledProvider := func(providerID uuid.UUID, providerType, name string, displayName sql.NullString, enabled bool, baseURL string) { t.Helper() var provider struct { Typ string Name string DisplayName sql.NullString Enabled bool BaseURL string } err = sqlDB.QueryRowContext(ctx, ` SELECT type, name, display_name, enabled, base_url FROM ai_providers WHERE id = $1 `, providerID).Scan(&provider.Typ, &provider.Name, &provider.DisplayName, &provider.Enabled, &provider.BaseURL) require.NoError(t, err) require.Equal(t, providerType, provider.Typ) require.Equal(t, name, provider.Name) require.Equal(t, displayName, provider.DisplayName) require.Equal(t, enabled, provider.Enabled) require.Equal(t, baseURL, provider.BaseURL) } assertBackfilledProvider( openAIProviderID, "openai", "agents-openai", sql.NullString{String: "OpenAI", Valid: true}, true, "https://api.openai.example.com/v1", ) assertBackfilledProvider( anthropicProviderID, "anthropic", "agents-anthropic", sql.NullString{}, false, "", ) var providerKeyCount int err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM ai_provider_keys WHERE provider_id = $1 AND api_key = 'sk-provider-openai' `, openAIProviderID).Scan(&providerKeyCount) require.NoError(t, err) require.Equal(t, 1, providerKeyCount, "non-empty legacy provider API key should be copied") err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM ai_provider_keys WHERE provider_id = $1 `, anthropicProviderID).Scan(&providerKeyCount) require.NoError(t, err) require.Zero(t, providerKeyCount, "empty legacy provider API key should not create an AI provider key") assertBackfilledUserKey := func(userKeyID, providerID uuid.UUID, apiKey string) { t.Helper() var userKeyCount int err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM user_ai_provider_keys WHERE id = $1 AND user_id = $2 AND ai_provider_id = $3 AND api_key = $4 `, userKeyID, userID, providerID, apiKey).Scan(&userKeyCount) require.NoError(t, err) require.Equal(t, 1, userKeyCount) } assertBackfilledUserKey(openAIUserKeyID, openAIProviderID, "sk-user-openai") assertBackfilledUserKey(anthropicUserKeyID, anthropicProviderID, "sk-user-anthropic") assertModelConfigProviderID := func(modelConfigID, providerID uuid.UUID) { t.Helper() var aiProviderID sql.NullString err = sqlDB.QueryRowContext(ctx, `SELECT ai_provider_id::text FROM chat_model_configs WHERE id = $1`, modelConfigID, ).Scan(&aiProviderID) require.NoError(t, err) require.Equal(t, sql.NullString{String: providerID.String(), Valid: true}, aiProviderID) } assertModelConfigProviderID(openAIModelConfigID, openAIProviderID) assertModelConfigProviderID(anthropicModelConfigID, anthropicProviderID) var legacyProviderCount int err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM chat_providers WHERE id IN ($1, $2) `, openAIProviderID, anthropicProviderID).Scan(&legacyProviderCount) require.NoError(t, err) require.Equal(t, 2, legacyProviderCount, "backfill should leave legacy rows for the rest of the stack") downSQL, err := os.ReadFile("000504_ai_providers_backfill.down.sql") require.NoError(t, err) _, err = sqlDB.ExecContext(ctx, string(downSQL)) require.NoError(t, err) err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM ai_providers WHERE id IN ($1, $2) `, openAIProviderID, anthropicProviderID).Scan(&providerKeyCount) require.NoError(t, err) require.Zero(t, providerKeyCount, "down migration should remove backfilled AI providers") err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM ai_provider_keys WHERE provider_id IN ($1, $2) `, openAIProviderID, anthropicProviderID).Scan(&providerKeyCount) require.NoError(t, err) require.Zero(t, providerKeyCount, "down migration should remove backfilled provider keys") var userKeyCount int err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM user_ai_provider_keys WHERE id IN ($1, $2) `, openAIUserKeyID, anthropicUserKeyID).Scan(&userKeyCount) require.NoError(t, err) require.Zero(t, userKeyCount, "down migration should remove backfilled user keys") err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM chat_model_configs WHERE id IN ($1, $2) AND ai_provider_id IS NOT NULL `, openAIModelConfigID, anthropicModelConfigID).Scan(&preBackfillModelConfigCount) require.NoError(t, err) require.Zero(t, preBackfillModelConfigCount, "down migration should clear model config AI provider references") err = sqlDB.QueryRowContext(ctx, ` SELECT COUNT(*) FROM chat_providers WHERE id IN ($1, $2) `, openAIProviderID, anthropicProviderID).Scan(&legacyProviderCount) require.NoError(t, err) require.Equal(t, 2, legacyProviderCount, "down migration should leave the legacy source rows intact") } // TestMigration000504AIProvidersBackfillOverridesNameConflict verifies that a // pre-existing live ai_providers row whose name collides with the backfill // (for example, agents-openai) is soft-deleted so the chat_providers-derived // row inserted by the migration becomes authoritative. This scenario should // not occur in practice since no other process writes to ai_providers before // this migration runs, but the migration tolerates it rather than failing. func TestMigration000504AIProvidersBackfillOverridesNameConflict(t *testing.T) { t.Parallel() const migrationVersion = 504 sqlDB := testSQLDB(t) 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 } } ctx := testutil.Context(t, testutil.WaitSuperLong) now := time.Now().UTC().Truncate(time.Microsecond) chatProviderID := uuid.New() staleProviderID := uuid.New() tx, err := sqlDB.BeginTx(ctx, nil) require.NoError(t, err) defer tx.Rollback() // Pre-existing live ai_providers row that collides on name. _, err = tx.ExecContext(ctx, `INSERT INTO ai_providers (id, type, name, display_name, enabled, base_url, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, staleProviderID, "openai", "agents-openai", "Stale OpenAI", true, "https://stale.example.com/v1", now, now, ) require.NoError(t, err) // chat_providers row whose backfill will collide with the stale row above. _, err = tx.ExecContext(ctx, `INSERT INTO chat_providers (id, provider, display_name, api_key, enabled, base_url, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, chatProviderID, "openai", "OpenAI", "sk-provider", true, "https://api.openai.example.com/v1", now, now, ) require.NoError(t, err) require.NoError(t, tx.Commit()) version, more, err := next() require.NoError(t, err) require.True(t, more) require.EqualValues(t, migrationVersion, version) // The stale row must be soft-deleted and disabled so the unique name index // (which is partial WHERE deleted = FALSE) no longer covers it. var stale struct { Deleted bool Enabled bool } err = sqlDB.QueryRowContext(ctx, `SELECT deleted, enabled FROM ai_providers WHERE id = $1`, staleProviderID, ).Scan(&stale.Deleted, &stale.Enabled) require.NoError(t, err) require.True(t, stale.Deleted, "pre-existing conflicting ai_providers row should be soft-deleted") require.False(t, stale.Enabled, "pre-existing conflicting ai_providers row should be disabled") // The new authoritative row must exist with the chat_providers id, the // agents-openai name, and the chat_providers base_url. var fresh struct { Name string BaseURL string Deleted bool Enabled bool } err = sqlDB.QueryRowContext(ctx, `SELECT name, base_url, deleted, enabled FROM ai_providers WHERE id = $1`, chatProviderID, ).Scan(&fresh.Name, &fresh.BaseURL, &fresh.Deleted, &fresh.Enabled) require.NoError(t, err) require.Equal(t, "agents-openai", fresh.Name) require.Equal(t, "https://api.openai.example.com/v1", fresh.BaseURL) require.False(t, fresh.Deleted) require.True(t, fresh.Enabled) } // TestMigration000504AIProvidersBackfillEnumInSingleTxn reproduces the // production migration path, where every pending migration runs inside a // single transaction (see pgTxnDriver). Migration 000499 widens // ai_provider_type with ALTER TYPE ... ADD VALUE, and 000504 casts existing // chat_providers rows to that enum. Postgres forbids using an enum value // added by ADD VALUE within the same transaction, so when a legacy provider // uses one of the new values (for example openai-compat) the batch fails with // "unsafe use of new value". The per-step Stepper used by the other tests // commits each migration separately and cannot surface this. func TestMigration000504AIProvidersBackfillEnumInSingleTxn(t *testing.T) { t.Parallel() sqlDB := testSQLDB(t) ctx := testutil.Context(t, testutil.WaitSuperLong) // Apply everything through 498 and commit, so chat_providers exists and is // populated before the batch under test runs, matching a deployment that // ran an earlier migration batch before this one. applyMigrationsInTxn(ctx, t, sqlDB, 1, 498) now := time.Now().UTC().Truncate(time.Microsecond) providerID := uuid.New() // A legacy provider whose type is one of the values added in 000499. _, err := sqlDB.ExecContext(ctx, ` INSERT INTO chat_providers (id, provider, display_name, api_key, enabled, base_url, created_at, updated_at) VALUES ($1, 'openai-compat', 'OpenAI Compatible', '', TRUE, 'https://api.example.com/v1', $2, $2) `, providerID, now) require.NoError(t, err) // Apply 000499 through 000504 in a single transaction, as production does. applyMigrationsInTxn(ctx, t, sqlDB, 499, 504) var typ string err = sqlDB.QueryRowContext(ctx, `SELECT type FROM ai_providers WHERE id = $1`, providerID, ).Scan(&typ) require.NoError(t, err) require.Equal(t, "openai-compat", typ) } // applyMigrationsInTxn executes the up SQL for every migration whose version is // in [from, to] inside a single transaction, mirroring pgTxnDriver. The whole // batch commits or rolls back together. func applyMigrationsInTxn(ctx context.Context, t *testing.T, sqlDB *sql.DB, from, to int) { t.Helper() entries, err := os.ReadDir(".") require.NoError(t, err) var files []string for _, entry := range entries { name := entry.Name() if !strings.HasSuffix(name, ".up.sql") { continue } var version int if _, err := fmt.Sscanf(name, "%06d_", &version); err != nil { continue } if version >= from && version <= to { files = append(files, name) } } slices.Sort(files) tx, err := sqlDB.BeginTx(ctx, nil) require.NoError(t, err) defer tx.Rollback() for _, name := range files { query, err := os.ReadFile(name) require.NoError(t, err) _, err = tx.ExecContext(ctx, string(query)) require.NoErrorf(t, err, "apply migration %s", name) } require.NoError(t, tx.Commit()) } func TestMigration000498SoftDeleteStaleWorkspaceAgents(t *testing.T) { t.Parallel() const migrationVersion = 498 sqlDB := testSQLDB(t) // Step up to migrationVersion - 1. 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 } } ctx := testutil.Context(t, testutil.WaitSuperLong) now := time.Now().UTC().Truncate(time.Microsecond) // Seed the prerequisite tables. Two workspaces share the same EC2-style // instance id across several builds; a third workspace has a single // build on a different instance (baseline, must not be affected). userID := uuid.New() orgID := uuid.New() templateID := uuid.New() templateVersionID := uuid.New() fileID := uuid.New() wsA := uuid.New() wsB := uuid.New() wsSingle := uuid.New() wsDeleted := uuid.New() instanceAB := "i-shared-ab" instanceSingle := "i-solo" instanceDeleted := "i-deleted" // For workspace A: 3 builds on the same instance. // For workspace B: 2 builds on the same instance (different workspace, // same instance id, exercises the cross-workspace scoping case). // For wsSingle: 1 build, should stay non-deleted after the backfill. // For wsDeleted: 1 build on a soft-deleted workspace. Agent should be // marked deleted even though it's on the latest build. type build struct { id uuid.UUID jobID uuid.UUID resourceID uuid.UUID agentID uuid.UUID buildNum int32 wsID uuid.UUID instanceID string } mkBuild := func(ws uuid.UUID, buildNum int32, instance string) build { return build{ id: uuid.New(), jobID: uuid.New(), resourceID: uuid.New(), agentID: uuid.New(), buildNum: buildNum, wsID: ws, instanceID: instance, } } aBuilds := []build{ mkBuild(wsA, 1, instanceAB), mkBuild(wsA, 2, instanceAB), mkBuild(wsA, 3, instanceAB), } bBuilds := []build{ mkBuild(wsB, 1, instanceAB), mkBuild(wsB, 2, instanceAB), } singleBuilds := []build{ mkBuild(wsSingle, 1, instanceSingle), } deletedBuilds := []build{ mkBuild(wsDeleted, 1, instanceDeleted), } allBuilds := append(append(append(append([]build{}, aBuilds...), bBuilds...), singleBuilds...), deletedBuilds...) tx, err := sqlDB.BeginTx(ctx, nil) require.NoError(t, err) defer tx.Rollback() // Minimal user / org / template / template_version / file. _, err = tx.ExecContext(ctx, `INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`, userID, "seed", "seed@test.com", []byte{}, now, now, "active", pq.StringArray{}, "password", ) require.NoError(t, err) _, err = tx.ExecContext(ctx, `INSERT INTO organizations (id, name, display_name, description, icon, created_at, updated_at, is_default) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, orgID, "seed-org", "Seed Org", "", "", now, now, false, ) require.NoError(t, err) _, err = tx.ExecContext(ctx, `INSERT INTO files (id, hash, created_at, created_by, mimetype, data) VALUES ($1, $2, $3, $4, $5, $6)`, fileID, "hash", now, userID, "application/octet-stream", []byte{}, ) require.NoError(t, err) _, err = tx.ExecContext(ctx, `INSERT INTO templates (id, created_at, updated_at, organization_id, name, provisioner, active_version_id, description, created_by, group_acl, user_acl, display_name) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)`, templateID, now, now, orgID, "tpl", "echo", templateVersionID, "", userID, "{}", "{}", "", ) require.NoError(t, err) _, err = tx.ExecContext(ctx, `INSERT INTO template_versions (id, template_id, organization_id, created_at, updated_at, name, readme, job_id, created_by, message) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, templateVersionID, templateID, orgID, now, now, "v", "", uuid.New(), userID, "", ) require.NoError(t, err) for _, ws := range []uuid.UUID{wsA, wsB, wsSingle} { _, err = tx.ExecContext(ctx, `INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, name, deleted, automatic_updates) VALUES ($1, $2, $3, $4, $5, $6, $7, false, 'never')`, ws, now, now, userID, orgID, templateID, "ws-"+ws.String()[:8], ) require.NoError(t, err) } // wsDeleted is a soft-deleted workspace. Its agent is on the latest // build but must still be soft-deleted by the migration. _, err = tx.ExecContext(ctx, `INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, name, deleted, automatic_updates) VALUES ($1, $2, $3, $4, $5, $6, $7, true, 'never')`, wsDeleted, now, now, userID, orgID, templateID, "ws-"+wsDeleted.String()[:8], ) require.NoError(t, err) // For every build: provisioner_job -> workspace_build -> workspace_resource -> workspace_agent. for _, b := range allBuilds { _, err = tx.ExecContext(ctx, `INSERT INTO provisioner_jobs (id, created_at, updated_at, organization_id, initiator_id, provisioner, storage_method, type, input, file_id) VALUES ($1, $2, $3, $4, $5, 'echo', 'file', 'workspace_build', '{}', $6)`, b.jobID, now, now, orgID, userID, fileID, ) require.NoError(t, err) _, err = tx.ExecContext(ctx, `INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, job_id, reason) VALUES ($1, $2, $3, $4, $5, $6, 'start', $7, $8, 'initiator')`, b.id, now, now, b.wsID, templateVersionID, b.buildNum, userID, b.jobID, ) require.NoError(t, err) _, err = tx.ExecContext(ctx, `INSERT INTO workspace_resources (id, created_at, job_id, transition, type, name) VALUES ($1, $2, $3, 'start', 'aws_instance', 'dev')`, b.resourceID, now, b.jobID, ) require.NoError(t, err) _, err = tx.ExecContext(ctx, `INSERT INTO workspace_agents (id, created_at, updated_at, name, resource_id, auth_token, auth_instance_id, architecture, operating_system, deleted) VALUES ($1, $2, $3, 'main', $4, $5, $6, 'amd64', 'linux', false)`, b.agentID, now, now, b.resourceID, uuid.New(), b.instanceID, ) require.NoError(t, err) } require.NoError(t, tx.Commit()) // Sanity check pre-migration: all agents should be deleted=false. var preDeletedCount int err = sqlDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM workspace_agents WHERE deleted = true`).Scan(&preDeletedCount) require.NoError(t, err) require.Equal(t, 0, preDeletedCount, "no agents should be deleted pre-migration") // Run migration 491. version, more, err := next() require.NoError(t, err) require.True(t, more) require.EqualValues(t, migrationVersion, version) // Backfill assertions: // wsA: builds 1,2,3 → keep agent for build 3, delete for 1 and 2. // wsB: builds 1,2 → keep agent for build 2, delete for 1. // wsSingle: 1 build → keep. // Per workspace, exactly one agent remains deleted=false. check := func(label string, expectDeleted bool, agent uuid.UUID) { var deleted bool err := sqlDB.QueryRowContext(ctx, `SELECT deleted FROM workspace_agents WHERE id = $1`, agent).Scan(&deleted) require.NoError(t, err, label) require.Equal(t, expectDeleted, deleted, label) } check("wsA build 1 (old) should be deleted", true, aBuilds[0].agentID) check("wsA build 2 (old) should be deleted", true, aBuilds[1].agentID) check("wsA build 3 (latest) should be kept", false, aBuilds[2].agentID) check("wsB build 1 (old) should be deleted", true, bBuilds[0].agentID) check("wsB build 2 (latest) should be kept", false, bBuilds[1].agentID) check("wsSingle build 1 (solo latest) should be kept", false, singleBuilds[0].agentID) check("wsDeleted: agent on deleted workspace should be soft-deleted even though it's the latest build", true, deletedBuilds[0].agentID) // The ongoing invariants are enforced by wsbuilder.Builder.Build and // provisionerdserver.CompleteJob via SoftDeletePriorWorkspaceAgents and // SoftDeleteWorkspaceAgentsByWorkspaceID. Those paths are covered by // the querier tests TestSoftDeletePriorWorkspaceAgents and // TestSoftDeleteWorkspaceAgentsByWorkspaceID, plus integration tests // under coderd/coderd_test.go; not retested here. }