mirror of
https://github.com/coder/coder.git
synced 2026-10-16 06:37:58 +08:00
## Stack Part of [CODAGT-1341](https://linear.app/codercom/issue/CODAGT-1341/support-organization-level-skills-in-coder-agents). This stack adds organization skills to Coder Agents: admins publish a skill once per organization, share it with groups or members, and chats in that organization load it next to personal and workspace skills. It replaces the 21-PR draft stack #30584 to #30604 with five PRs, bottom to top: 1. **#30609** data model: one migration, queries, RBAC and dbauthz (this PR) 2. #30610 organization skill API routes and the scope-neutral SDK 3. #30611 sharing ACL routes 4. #30612 settings UI: shared skills table, admin page, member section 5. #30613 chats: chatd loading, `/` menu, docs ## What changes This PR adds the data model and authorization for organization skills. Personal skills keep working unchanged: `git diff origin/main -- codersdk site` shows only generated RBAC and API key scope additions (`apikey_scopes_gen.go`, `rbacresources_gen.go`, `rbacresourcesGenerated.ts`, the RBAC and scope unions in `typesGenerated.ts`) and `codersdk.ResourceTypeOrganizationSkill` with its friendly string. No personal route, SDK type or site component changes, and `coderd/userskills_test.go` changes only by mechanical constant renames. **One migration, `000621_skills`** (please review it closely): - renames `user_skills` to `skills` with its constraints, index and triggers, and recreates the cap function and `delete_deleted_user_resources`, whose PL/pgSQL bodies name the table; - adds `organization_id` and `project_id` (`ON DELETE CASCADE`), `enabled` (default true), and `group_acl` / `user_acl` (`jsonb`, default `{}`), with CHECKs for exactly one owner, ACLs that are JSON objects, and ACLs that stay empty unless an organization owns the row, plus partial unique name indexes per owner kind; - replaces the cap trigger with `enforce_skills_per_owner_limit()`: 100 rows per owner, disabled rows included, locking the owner row `FOR NO KEY UPDATE` so unrelated foreign key checks are not blocked; - adds a `BEFORE DELETE` trigger on `organization_members` that removes a departing member from `user_acl` on that organization's skills, so rejoining doesn't restore a direct grant; - adds the `organization_skill` `resource_type` and `api_key_scope` enum values last; nothing in the migration uses them. The down migration undoes the three old migrations in reverse order: it drops organization and project rows and restores `user_skills`, its names, the cap function and `user_id NOT NULL`. Authorization: a new RBAC resource `organization_skill` (create, read, update, delete, share). Owner and org admin get it; site and org auditors get read; members read only through the skill's ACL, which applies only to members of its organization. `Skill.RBACObject()` returns the organization object with both ACLs; project rows stay denied to everyone. New sqlc queries cover organization insert, get, list (SQL-filtered through `GetAuthorizedOrganizationSkillMetadata` and a `regosql` converter), update and delete, with dbauthz wrappers and the audit table entries. No routes yet. ## Commits The first commit reproduces the reviewed code from the old stack; each later commit is a labeled trim or fix. - `baa3150e3b7` feat: add the organization skills data model. Reproduction: every non-site file this commit changes equals the reviewed old `authz` head `1e16b3004c4`, except that the old stack's migrations 000615, 000616 and 000617 are merged into one migration (then `000615_skills`, now `000621_skills`) with one fixture (`testdata/fixtures/000621_skills.up.sql`) and one migration test (now `testMigration000621Skills`). At that commit, `dump.sql` was byte-identical to the old head's. - `1b84e83463a` test(coderd/database): keep only the SQL paths of the org skill matrix. `TestOrganizationSkillAuthorization` becomes `TestOrganizationSkillQueries` instead of being deleted, because coverage showed it is this PR's only run of the `GetAuthorizedOrganizationSkillMetadata` row scan (80% to 65% without it) and of the organization update and delete queries (100% to 0%). Swapping `group_acl` and `user_acl` in the converter makes `GrantedMember` fail. The removed cases are decided by `TestRolePermissions`, `TestSkillRBACObjectOrganizationACL` and the dbauthz method suite. - `e2931b412e9` test(coderd/database): fold the skill cap cases into existing coverage. The user case is covered by `TestUserSkillLimit` and `TestUserSkillLimitConcurrentCreates`; the organization case folds into `ConcurrentOrganizationInsertsHonorCap`, which now also checks that disabled rows count and another organization is unaffected (mutation checked); the project case stays. - `be62aea9021` test(coderd/rbac): drop `TestOrganizationSkillReadScope`, covered by `TestExternalScopeNames` and the generated scope file. - `0c275ba37de` merges main at `454a2fcc697`, which took 000615 to 000618: it renumbers the migration to 000619 and resolves `dump.sql`, `check_constraint.go` and the generated API docs. `8f2b4b87ee2` merges main at `ad85646164e`, which took 000619: it renumbers the migration to 000620. - `5248d56b64f` merges main at `e7b76872098`, which took 000620 (`000620_audit_chat_project_instructions_resource_type`): it renumbers the migration, fixture and migration test to `000621_skills` (step `Migration000621Skills`, version 621), and resolves `enterprise/audit/table.go` (keeping main's `AuditableChatProjectInstructions` entries and the `Skill` entry), `dump.sql` and `models.go` (the `resource_type` enum gets `chat_project_instructions`, then `organization_skill`) and the regenerated `schemas.md`. The up and down SQL are unchanged by every renumbering. ## Size Production +549/-67, tests +728/-33, generated +881/-314; total +2158/-414. This is above the 500-line reviewer-friendly limit on purpose: Mike asked for 4 to 5 PRs instead of the 21 small ones this stack replaces. ## Replaces #30584 (rename `user_skills` to `skills`), #30585 (skill owners), #30586 (organization skill authorization). ## Validation - Migration test: up keeps the existing personal row and its ID; down restores the `user_skills` names, cap function and `NOT NULL`, and drops organization rows; soft-deleting a user removes their skills after up, after down, and after re-up. - `go test` passed on `be62aea9021`, before the main merges, for `coderd` (skill tests), `coderd/database`, `dbauthz`, `migrations`, `gentest`, `coderd/rbac/...`, `coderd/x/chatd`, `coderd/x/skills`, `enterprise/audit` and `enterprise/coderd`; `make lint/check-scopes` passed; the four personal skills Vitest files pass unchanged. - On `5248d56b64f`: build, vet, `TestMigrate`, `TestMigrationChain`, the `migrations`, `database`, `dbauthz`, `gentest`, `rbac` and `enterprise/audit` packages pass; `make gen` run twice is clean, and `make pre-commit` passes. - Remote UAT round 7 through Coder Agents passed on the whole stack; see #30613's Validation for scope and findings. > Xum opened this PR on behalf of @ibetitsmike.
4480 lines
172 KiB
Go
4480 lines
172 KiB
Go
package migrations_test
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"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/sqlc-dev/pqtype"
|
|
"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/dbgen"
|
|
"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("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()
|
|
|
|
// dbtestutil.Open clones an already migrated template database, but this
|
|
// package tests the migrations themselves, so start from Postgres' stock
|
|
// empty template instead.
|
|
connection, err := dbtestutil.Open(t, dbtestutil.WithDBFrom("template1"))
|
|
require.NoError(t, err)
|
|
|
|
db, err := sql.Open("postgres", connection)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { _ = db.Close() })
|
|
|
|
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/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")
|
|
})
|
|
}
|
|
}
|
|
|
|
// migrationStepper applies the next migration and reports the version it
|
|
// reached, see migrations.Stepper.
|
|
type migrationStepper = func() (version uint, more bool, err error)
|
|
|
|
// TestMigrationChain steps one database through the migration chain once and
|
|
// runs each per-migration test when the chain reaches the migration before
|
|
// the one it covers, so the hundreds of preceding migrations are applied once
|
|
// instead of once per test. Each step must leave the chain clean at its
|
|
// version and must not assert whole-table counts that fixtures committed by
|
|
// earlier steps would change; tests that cannot are standalone.
|
|
//
|
|
//nolint:tparallel,paralleltest // Subtests share one database and step a single migration chain in version order.
|
|
func TestMigrationChain(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
return
|
|
}
|
|
|
|
sqlDB := testSQLDB(t)
|
|
next, err := migrations.Stepper(sqlDB)
|
|
require.NoError(t, err)
|
|
|
|
stepTo := func(target uint) {
|
|
t.Helper()
|
|
for {
|
|
version, more, err := next()
|
|
require.NoError(t, err)
|
|
require.Truef(t, more, "migration %d not found", target)
|
|
if version == target {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
steps := []struct {
|
|
name string
|
|
version uint
|
|
run func(t *testing.T, sqlDB *sql.DB, next migrationStepper)
|
|
}{
|
|
{"Migration000362AggregateUsageEvents", 362, testMigration000362AggregateUsageEvents},
|
|
{"Migration000387MigrateTaskWorkspaces", 387, testMigration000387MigrateTaskWorkspaces},
|
|
{"Migration000457ChatAccessRole", 457, testMigration000457ChatAccessRole},
|
|
{"Migration000475AgentsAccessOrgRole", 475, testMigration000475AgentsAccessOrgRole},
|
|
{"Migration000498SoftDeleteStaleWorkspaceAgents", 498, testMigration000498SoftDeleteStaleWorkspaceAgents},
|
|
{"Migration000546ChatHistoryAPIKeyConstraints", 546, testMigration000546ChatHistoryAPIKeyConstraints},
|
|
{"Migration000556UserSecretsEnabled", 556, testMigration000556UserSecretsEnabled},
|
|
{"Migration000587RemoveAgentsAccessRole", 587, testMigration000587RemoveAgentsAccessRole},
|
|
{"Migration000590WorkspaceAgentSessionCounts", 590, testMigration000590WorkspaceAgentSessionCounts},
|
|
{"Migration000595RemoveTaskPermissions", 595, testMigration000595RemoveTaskPermissions},
|
|
{"Migration000602RestoreAgentsAccessDefaultRole", 602, testMigration000602RestoreAgentsAccessDefaultRole},
|
|
{"Migration000606ChatDiffStatusOriginCredentials", 606, testMigration000606ChatDiffStatusOriginCredentials},
|
|
{"Migration000621Skills", 621, testMigration000621Skills},
|
|
}
|
|
for _, step := range steps {
|
|
stepTo(step.version - 1)
|
|
if !t.Run(step.name, func(t *testing.T) {
|
|
step.run(t, sqlDB, next)
|
|
}) {
|
|
t.Fatalf("%s failed, the chain state is unknown so the remaining steps are skipped", step.name)
|
|
}
|
|
}
|
|
|
|
for {
|
|
_, more, err := next()
|
|
require.NoError(t, err)
|
|
if !more {
|
|
break
|
|
}
|
|
}
|
|
|
|
// These tests need the fully migrated schema.
|
|
migrated := []struct {
|
|
name string
|
|
run func(t *testing.T, sqlDB *sql.DB)
|
|
}{
|
|
{"Migration000543ChatMessageSearchText", testMigration000543ChatMessageSearchText},
|
|
{"Migration000543ChatSearchSchemaIndexes", testMigration000543ChatSearchSchemaIndexes},
|
|
{"Migration000543ChatSearchSchemaBehavior", testMigration000543ChatSearchSchemaBehavior},
|
|
{"Migration000583ChatModelOverrideOrgScope", testMigration000583ChatModelOverrideOrgScope},
|
|
}
|
|
for _, tt := range migrated {
|
|
if !t.Run(tt.name, func(t *testing.T) {
|
|
tt.run(t, sqlDB)
|
|
}) {
|
|
t.Fatalf("%s failed, the database state is unknown so the remaining tests are skipped", tt.name)
|
|
}
|
|
}
|
|
}
|
|
|
|
// testMigration000362AggregateUsageEvents tests the migration that aggregates
|
|
// usage events into daily rows correctly.
|
|
func testMigration000362AggregateUsageEvents(t *testing.T, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 362
|
|
|
|
db := database.New(sqlDB)
|
|
|
|
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, sqlDB *sql.DB, next migrationStepper) {
|
|
// 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
|
|
|
|
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, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 457
|
|
|
|
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, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 475
|
|
|
|
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 TestMigration000593RemoveHasAITaskAndTaskBuildReasons(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
const migrationVersion = 593
|
|
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.WaitLong)
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
orgID, userID, templateID, versionID, templateJobID, workspaceID := uuid.New(), uuid.New(), uuid.New(), uuid.New(), uuid.New(), uuid.New()
|
|
buildIDs := []uuid.UUID{uuid.New(), uuid.New(), uuid.New(), uuid.New(), uuid.New()}
|
|
fixtures := []struct {
|
|
query string
|
|
args []any
|
|
}{
|
|
{
|
|
"INSERT INTO organizations (id, name, display_name, description, created_at, updated_at, default_org_member_roles) VALUES ($1, $2, '', '', $3, $3, '{}')",
|
|
[]any{orgID, orgID.String(), now},
|
|
},
|
|
{
|
|
"INSERT INTO users (id, username, email, hashed_password, created_at, updated_at) VALUES ($1, $2, $3, '', $4, $4)",
|
|
[]any{userID, userID.String(), userID.String() + "@test.com", now},
|
|
},
|
|
{
|
|
"INSERT INTO provisioner_jobs (id, created_at, updated_at, organization_id, initiator_id, provisioner, storage_method, file_id, type, input) VALUES ($1, $2, $2, $3, $4, 'echo', 'file', $1, 'template_version_import', '{}')",
|
|
[]any{templateJobID, now, orgID, userID},
|
|
},
|
|
{
|
|
"INSERT INTO template_versions (id, organization_id, name, readme, created_at, updated_at, job_id, created_by, has_ai_task) VALUES ($1, $2, 'v1', '', $3, $3, $4, $5, true)",
|
|
[]any{versionID, orgID, now, templateJobID, userID},
|
|
},
|
|
{
|
|
"INSERT INTO templates (id, organization_id, name, created_at, updated_at, provisioner, active_version_id, created_by) VALUES ($1, $2, 'build-reasons', $3, $3, 'echo', $4, $5)",
|
|
[]any{templateID, orgID, now, versionID, userID},
|
|
},
|
|
{
|
|
"UPDATE template_versions SET template_id = $1 WHERE id = $2",
|
|
[]any{templateID, versionID},
|
|
},
|
|
{
|
|
"INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, name) VALUES ($1, $2, $2, $3, $4, $5, 'build-reasons')",
|
|
[]any{workspaceID, now, userID, orgID, templateID},
|
|
},
|
|
}
|
|
for _, fixture := range fixtures {
|
|
_, err := sqlDB.ExecContext(ctx, fixture.query, fixture.args...)
|
|
require.NoError(t, err)
|
|
}
|
|
for i, reason := range []string{"task_auto_pause", "task_manual_pause", "task_resume", "initiator", "autostop"} {
|
|
jobID := uuid.New()
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
"INSERT INTO provisioner_jobs (id, created_at, updated_at, organization_id, initiator_id, provisioner, storage_method, file_id, type, input) VALUES ($1, $2, $2, $3, $4, 'echo', 'file', $1, 'workspace_build', '{}')",
|
|
jobID, now, orgID, userID)
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
"INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, job_id, reason, has_ai_task) VALUES ($1, $2, $2, $3, $4, $5, 'start', $6, $7, $8, $9)",
|
|
buildIDs[i], now, workspaceID, versionID, i+1, userID, jobID, reason, i == 0)
|
|
require.NoError(t, err)
|
|
}
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
"INSERT INTO workspace_build_orchestrations (id, created_at, updated_at, workspace_id, parent_build_id, child_transition, child_reason) VALUES ($1, $2, $2, $3, $4, 'start', 'task_auto_pause')",
|
|
uuid.New(), now, workspaceID, buildIDs[0])
|
|
require.NoError(t, err)
|
|
|
|
version, _, err := next()
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, migrationVersion, version)
|
|
var count int
|
|
var reasons pq.StringArray
|
|
const reasonsQuery = "SELECT array_agg(reason::text ORDER BY build_number) FROM workspace_builds WHERE workspace_id = $1"
|
|
wantReasons := pq.StringArray{"autostop", "initiator", "initiator", "initiator", "autostop"}
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, reasonsQuery, workspaceID).Scan(&reasons))
|
|
require.Equal(t, wantReasons, reasons)
|
|
var childReason string
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT child_reason FROM workspace_build_orchestrations WHERE workspace_id = $1", workspaceID).Scan(&childReason))
|
|
require.Equal(t, "autostop", childReason)
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT count(*) FROM template_version_with_user").Scan(&count))
|
|
require.Equal(t, 1, count)
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT count(*) FROM workspace_build_with_user").Scan(&count))
|
|
require.Equal(t, len(buildIDs), count)
|
|
|
|
downSQL, err := os.ReadFile("000593_remove_has_ai_task_and_task_build_reasons.down.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT count(*) FROM template_versions WHERE has_ai_task IS FALSE").Scan(&count))
|
|
require.Equal(t, 1, count)
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT count(*) FROM workspace_builds WHERE has_ai_task IS FALSE").Scan(&count))
|
|
require.Equal(t, len(buildIDs), count)
|
|
// The original reasons cannot be recovered after remapping.
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, reasonsQuery, workspaceID).Scan(&reasons))
|
|
require.Equal(t, wantReasons, reasons)
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT child_reason FROM workspace_build_orchestrations WHERE workspace_id = $1", workspaceID).Scan(&childReason))
|
|
require.Equal(t, "autostop", childReason)
|
|
|
|
upSQL, err := os.ReadFile("000593_remove_has_ai_task_and_task_build_reasons.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
func TestMigration000593TaskBuildReasonRewriteTiming(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
return
|
|
}
|
|
|
|
// The enum replacement rewrites workspace_builds, risking slow upgrades
|
|
// on large deployments.
|
|
const (
|
|
migrationVersion = 593
|
|
workspaceCount = 1_000
|
|
buildCount = 100_000
|
|
buildsPerWorkspace = buildCount / workspaceCount
|
|
orchestrationCount = 5_000
|
|
taskReasonBuilds = 3 * buildCount / 10
|
|
)
|
|
|
|
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)
|
|
|
|
orgID := uuid.New()
|
|
userID := uuid.New()
|
|
templateID := uuid.New()
|
|
templateVersionID := uuid.New()
|
|
templateJobID := uuid.New()
|
|
|
|
// Deterministic IDs let the bulk inserts join rows generated by
|
|
// generate_series without round-tripping IDs through the test.
|
|
const (
|
|
workspaceIDPattern = "('a0000000-0000-4000-8000-' || lpad(%s::text, 12, '0'))::uuid"
|
|
jobIDPattern = "('b0000000-0000-4000-8000-' || lpad(%s::text, 12, '0'))::uuid"
|
|
buildIDPattern = "('c0000000-0000-4000-8000-' || lpad(%s::text, 12, '0'))::uuid"
|
|
)
|
|
|
|
fixtures := []struct {
|
|
query string
|
|
args []any
|
|
}{
|
|
{
|
|
`INSERT INTO organizations (id, name, display_name, description, created_at, updated_at, default_org_member_roles) VALUES ($1, $2, $3, $4, $5, $6, '{}')`,
|
|
[]any{orgID, "timing-org", "Timing 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, "timinguser", "timing@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", "", 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, "timing-template", now, now, "terraform", templateVersionID, userID},
|
|
},
|
|
{
|
|
`UPDATE template_versions SET template_id = $1 WHERE id = $2`,
|
|
[]any{templateID, templateVersionID},
|
|
},
|
|
{
|
|
fmt.Sprintf(`INSERT INTO workspaces (id, created_at, updated_at, owner_id, organization_id, template_id, name, last_used_at)
|
|
SELECT `+workspaceIDPattern+`, $1, $1, $2, $3, $4, 'timing-ws-' || g, $1
|
|
FROM generate_series(0, %d) AS g`, "g", workspaceCount-1),
|
|
[]any{now, userID, orgID, templateID},
|
|
},
|
|
{
|
|
fmt.Sprintf(`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)
|
|
SELECT `+jobIDPattern+`, $1, $1, $1, $1, '', $2, $3, 'terraform', 'file', $4, 'workspace_build', '{}'::jsonb, '{}'::jsonb
|
|
FROM generate_series(0, %d) AS g`, "g", buildCount-1),
|
|
[]any{now, orgID, userID, uuid.New()},
|
|
},
|
|
{
|
|
fmt.Sprintf(`INSERT INTO workspace_builds (id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, job_id, reason, has_ai_task)
|
|
SELECT `+buildIDPattern+`, $1, $1, `+fmt.Sprintf(workspaceIDPattern, "(g / %d)")+`, $2, (g %% %d) + 1,
|
|
CASE WHEN g %% 2 = 0 THEN 'start' ELSE 'stop' END::workspace_transition, $3, `+jobIDPattern+`,
|
|
CASE g %% 10
|
|
WHEN 0 THEN 'task_auto_pause'
|
|
WHEN 1 THEN 'task_manual_pause'
|
|
WHEN 2 THEN 'task_resume'
|
|
WHEN 3 THEN 'autostart'
|
|
WHEN 4 THEN 'autostop'
|
|
ELSE 'initiator'
|
|
END::build_reason,
|
|
CASE WHEN g %% 10 < 3 THEN true END
|
|
FROM generate_series(0, %d) AS g`, "g", buildsPerWorkspace, buildsPerWorkspace, "g", buildCount-1),
|
|
[]any{now, templateVersionID, userID},
|
|
},
|
|
{
|
|
fmt.Sprintf(`INSERT INTO workspace_build_orchestrations (id, created_at, updated_at, workspace_id, parent_build_id, child_transition, child_reason)
|
|
SELECT ('e0000000-0000-4000-8000-' || lpad(g::text, 12, '0'))::uuid, $1, $1, `+fmt.Sprintf(workspaceIDPattern, "(g * 20 / %d)")+`, `+fmt.Sprintf(buildIDPattern, "(g * 20)")+`, 'start',
|
|
CASE g %% 3
|
|
WHEN 0 THEN 'task_auto_pause'
|
|
WHEN 1 THEN 'task_resume'
|
|
ELSE 'initiator'
|
|
END::build_reason
|
|
FROM generate_series(0, %d) AS g`, buildsPerWorkspace, orchestrationCount-1),
|
|
[]any{now},
|
|
},
|
|
}
|
|
for i, fixture := range fixtures {
|
|
_, err := sqlDB.ExecContext(ctx, fixture.query, fixture.args...)
|
|
require.NoErrorf(t, err, "fixture %d", i)
|
|
}
|
|
|
|
countByReason := func(table, column, reason string) int {
|
|
var n int
|
|
//nolint:gosec // Test-only query built from constant table/column names.
|
|
err := sqlDB.QueryRowContext(ctx, fmt.Sprintf("SELECT count(*) FROM %s WHERE %s::text = $1", table, column), reason).Scan(&n)
|
|
require.NoError(t, err)
|
|
return n
|
|
}
|
|
wantAutostop := countByReason("workspace_builds", "reason", "autostop") + countByReason("workspace_builds", "reason", "task_auto_pause")
|
|
wantInitiator := countByReason("workspace_builds", "reason", "initiator") + countByReason("workspace_builds", "reason", "task_manual_pause") + countByReason("workspace_builds", "reason", "task_resume")
|
|
wantChildAutostop := countByReason("workspace_build_orchestrations", "child_reason", "autostop") + countByReason("workspace_build_orchestrations", "child_reason", "task_auto_pause")
|
|
wantChildInitiator := countByReason("workspace_build_orchestrations", "child_reason", "initiator") + countByReason("workspace_build_orchestrations", "child_reason", "task_resume")
|
|
|
|
start := time.Now()
|
|
version, _, err := next()
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, migrationVersion, version)
|
|
elapsed := time.Since(start)
|
|
t.Logf("migration %d over %d workspace builds (%d with task reasons) and %d orchestrations took %s", migrationVersion, buildCount, taskReasonBuilds, orchestrationCount, elapsed)
|
|
// Bound migration time to catch regressions that could prolong upgrade downtime.
|
|
require.Lessf(t, elapsed, 2*time.Minute, "migration %d took %s at %d workspace builds", migrationVersion, elapsed, buildCount)
|
|
|
|
var totalBuilds int
|
|
err = sqlDB.QueryRowContext(ctx, "SELECT count(*) FROM workspace_builds").Scan(&totalBuilds)
|
|
require.NoError(t, err)
|
|
require.Equal(t, buildCount, totalBuilds)
|
|
require.Equal(t, wantAutostop, countByReason("workspace_builds", "reason", "autostop"))
|
|
require.Equal(t, wantInitiator, countByReason("workspace_builds", "reason", "initiator"))
|
|
require.Equal(t, wantChildAutostop, countByReason("workspace_build_orchestrations", "child_reason", "autostop"))
|
|
require.Equal(t, wantChildInitiator, countByReason("workspace_build_orchestrations", "child_reason", "initiator"))
|
|
|
|
var count int
|
|
err = sqlDB.QueryRowContext(ctx, "SELECT count(*) FROM information_schema.columns WHERE table_schema = 'public' AND table_name = 'workspace_builds' AND column_name = 'has_ai_task'").Scan(&count)
|
|
require.NoError(t, err)
|
|
require.Zero(t, count)
|
|
}
|
|
|
|
func testMigration000595RemoveTaskPermissions(t *testing.T, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 595
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
orgID, userID, roleID, appID, taskAppID := uuid.New(), uuid.New(), uuid.New(), uuid.New(), uuid.New()
|
|
const (
|
|
taskPermission = `{"negate": false, "resource_type": "task", "action": "read"}`
|
|
userPermission = `{"negate": false, "resource_type": "user", "action": "read"}`
|
|
)
|
|
permissions := []struct {
|
|
column string
|
|
before string
|
|
want string
|
|
}{
|
|
{"site_permissions", "[" + taskPermission + ", " + userPermission + "]", "[" + userPermission + "]"},
|
|
{"org_permissions", "[" + taskPermission + "]", "[]"},
|
|
{"user_permissions", "[" + userPermission + ", " + taskPermission + "]", "[" + userPermission + "]"},
|
|
{"member_permissions", "[" + userPermission + "]", "[" + userPermission + "]"},
|
|
}
|
|
for _, fixture := range []struct {
|
|
query string
|
|
args []any
|
|
}{
|
|
{
|
|
"INSERT INTO organizations (id, name, display_name, description, created_at, updated_at, default_org_member_roles) VALUES ($1, $2, '', '', $3, $3, '{}')",
|
|
[]any{orgID, orgID.String(), now},
|
|
},
|
|
{
|
|
"INSERT INTO users (id, username, email, hashed_password, created_at, updated_at) VALUES ($1, $2, $3, '', $4, $4)",
|
|
[]any{userID, userID.String(), userID.String() + "@test.com", now},
|
|
},
|
|
{
|
|
"INSERT INTO custom_roles (id, name, display_name, organization_id, site_permissions, org_permissions, user_permissions, member_permissions) VALUES ($1, 'task-role', 'Task Role', $2, $3, $4, $5, $6)",
|
|
[]any{roleID, orgID, permissions[0].before, permissions[1].before, permissions[2].before, permissions[3].before},
|
|
},
|
|
} {
|
|
_, err := sqlDB.ExecContext(ctx, fixture.query, fixture.args...)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
apps := []struct {
|
|
id uuid.UUID
|
|
name string
|
|
scope string
|
|
wantScope string
|
|
}{
|
|
{appID, "mixed-app", "task:read template:read", "template:read"},
|
|
// An empty scope reads as unrestricted, so a whitespace-only one stands in for "nothing allowed".
|
|
{taskAppID, "task-app", "task:read", " "},
|
|
}
|
|
for _, app := range apps {
|
|
_, err := sqlDB.ExecContext(ctx,
|
|
"INSERT INTO oauth2_provider_apps (id, created_at, updated_at, name, icon, callback_url, scope) VALUES ($1, $2, $2, $3, '', 'http://localhost/callback', $4)",
|
|
app.id, now, app.name, app.scope)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
keys := []struct {
|
|
id string
|
|
scopes string
|
|
allowList string
|
|
wantDeleted bool
|
|
wantScopes string
|
|
wantAllowList string
|
|
}{
|
|
{id: "mixed-key", scopes: "{task:read,template:read}", allowList: "{*:*}", wantScopes: "{template:read}", wantAllowList: "{*:*}"},
|
|
{id: "multi-task-scopes-key", scopes: "{task:read,task:create,template:read}", allowList: "{*:*}", wantScopes: "{template:read}", wantAllowList: "{*:*}"},
|
|
{id: "task-key", scopes: "{task:read}", allowList: "{*:*}", wantDeleted: true},
|
|
{id: "mixed-allow-key", scopes: "{template:read}", allowList: "{task:*,user:*}", wantScopes: "{template:read}", wantAllowList: "{user:*}"},
|
|
{id: "multi-task-allows-key", scopes: "{template:read}", allowList: "{task:*,task:3f0c9b2e-5d41-4a7b-9c1e-8d2f6a4b7c5e,user:*}", wantScopes: "{template:read}", wantAllowList: "{user:*}"},
|
|
{id: "task-allow-key", scopes: "{template:read}", allowList: "{task:*}", wantDeleted: true},
|
|
{id: "task-only-key", scopes: "{task:read,task:create}", allowList: "{task:*,task:3f0c9b2e-5d41-4a7b-9c1e-8d2f6a4b7c5e}", wantDeleted: true},
|
|
}
|
|
for _, key := range keys {
|
|
_, err := sqlDB.ExecContext(ctx,
|
|
"INSERT INTO api_keys (id, hashed_secret, user_id, last_used, expires_at, created_at, updated_at, login_type, scopes, allow_list, token_name) VALUES ($1, $2, $3, $4, $4, $4, $4, 'token', $5, $6, $1)",
|
|
key.id, []byte(key.id), userID, now, key.scopes, key.allowList)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Each grant is stored as both an authorization code and a token, keyed by the grant name.
|
|
grants := []struct {
|
|
key string
|
|
scope string
|
|
wantDeleted bool
|
|
wantScope string
|
|
}{
|
|
{key: "mixed-grant", scope: "task:read template:read", wantScope: "template:read"},
|
|
{key: "task-grant", scope: "task:read", wantDeleted: true},
|
|
}
|
|
for _, grant := range grants {
|
|
_, err := sqlDB.ExecContext(ctx,
|
|
"INSERT INTO oauth2_provider_app_codes (id, created_at, expires_at, secret_prefix, hashed_secret, user_id, app_id, scope) VALUES ($1, $2, $2, $3, $3, $4, $5, $6)",
|
|
uuid.New(), now, []byte(grant.key), userID, appID, grant.scope)
|
|
require.NoError(t, err)
|
|
// oauth2_provider_app_tokens.api_key_id is a NOT NULL foreign key, so this key only satisfies it and is not asserted.
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
"INSERT INTO api_keys (id, hashed_secret, user_id, last_used, expires_at, created_at, updated_at, login_type, scopes, allow_list, token_name) VALUES ($1, $2, $3, $4, $4, $4, $4, 'oauth2_provider_app', '{template:read}', '{*:*}', $1)",
|
|
grant.key, []byte(grant.key), userID, now)
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
"INSERT INTO oauth2_provider_app_tokens (id, created_at, expires_at, hash_prefix, refresh_hash, api_key_id, user_id, app_id, scope) VALUES ($1, $2, $2, $3, $3, $4, $5, $6, $7)",
|
|
uuid.New(), now, []byte(grant.key), grant.key, userID, appID, grant.scope)
|
|
require.NoError(t, err)
|
|
}
|
|
grantTables := []struct{ name, prefixColumn string }{
|
|
{"oauth2_provider_app_codes", "secret_prefix"},
|
|
{"oauth2_provider_app_tokens", "hash_prefix"},
|
|
}
|
|
|
|
assertMigrated := func() {
|
|
t.Helper()
|
|
for _, p := range permissions {
|
|
var got string
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT "+p.column+" FROM custom_roles WHERE id = $1", roleID).Scan(&got))
|
|
require.JSONEq(t, p.want, got, p.column)
|
|
}
|
|
for _, app := range apps {
|
|
var scope string
|
|
require.NoError(t, sqlDB.QueryRowContext(ctx, "SELECT scope FROM oauth2_provider_apps WHERE id = $1", app.id).Scan(&scope))
|
|
require.Equal(t, app.wantScope, scope, app.name)
|
|
}
|
|
for _, key := range keys {
|
|
var scopes, allowList string
|
|
err := sqlDB.QueryRowContext(ctx, "SELECT scopes::text, allow_list::text FROM api_keys WHERE id = $1", key.id).Scan(&scopes, &allowList)
|
|
if key.wantDeleted {
|
|
require.ErrorIs(t, err, sql.ErrNoRows, key.id)
|
|
continue
|
|
}
|
|
require.NoError(t, err, key.id)
|
|
require.Equal(t, key.wantScopes, scopes, key.id)
|
|
require.Equal(t, key.wantAllowList, allowList, key.id)
|
|
}
|
|
for _, grant := range grants {
|
|
for _, table := range grantTables {
|
|
var scope string
|
|
err := sqlDB.QueryRowContext(ctx, "SELECT scope FROM "+table.name+" WHERE "+table.prefixColumn+" = $1", []byte(grant.key)).Scan(&scope)
|
|
if grant.wantDeleted {
|
|
require.ErrorIs(t, err, sql.ErrNoRows, "%s %s", table.name, grant.key)
|
|
continue
|
|
}
|
|
require.NoError(t, err, "%s %s", table.name, grant.key)
|
|
require.Equal(t, grant.wantScope, scope, "%s %s", table.name, grant.key)
|
|
}
|
|
}
|
|
}
|
|
|
|
version, _, err := next()
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, migrationVersion, version)
|
|
assertMigrated()
|
|
|
|
downSQL, err := os.ReadFile("000595_remove_task_permissions.down.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
// Downgrading cannot restore revoked grants or stripped permissions.
|
|
assertMigrated()
|
|
|
|
upSQL, err := os.ReadFile("000595_remove_task_permissions.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
assertMigrated()
|
|
}
|
|
|
|
func testMigration000587RemoveAgentsAccessRole(t *testing.T, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 587
|
|
|
|
db := database.New(sqlDB)
|
|
user := dbgen.User(t, db, database.User{
|
|
RBACRoles: []string{"auditor", "agents-access"},
|
|
})
|
|
org := dbgen.Organization(t, db, database.Organization{
|
|
DefaultOrgMemberRoles: []string{"organization-workspace-access", "agents-access"},
|
|
})
|
|
dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
|
OrganizationID: org.ID,
|
|
UserID: user.ID,
|
|
Roles: []string{"organization-auditor", "agents-access"},
|
|
})
|
|
|
|
version, _, err := next()
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, migrationVersion, version)
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
var siteRoles, orgRoles, defaultRoles pq.StringArray
|
|
err = sqlDB.QueryRowContext(ctx, "SELECT rbac_roles FROM users WHERE id = $1", user.ID).Scan(&siteRoles)
|
|
require.NoError(t, err)
|
|
err = sqlDB.QueryRowContext(ctx, "SELECT roles FROM organization_members WHERE organization_id = $1 AND user_id = $2", org.ID, user.ID).Scan(&orgRoles)
|
|
require.NoError(t, err)
|
|
err = sqlDB.QueryRowContext(ctx, "SELECT default_org_member_roles FROM organizations WHERE id = $1", org.ID).Scan(&defaultRoles)
|
|
require.NoError(t, err)
|
|
|
|
require.Equal(t, []string{"auditor"}, []string(siteRoles))
|
|
require.Equal(t, []string{"organization-auditor"}, []string(orgRoles))
|
|
require.Equal(t, []string{"organization-workspace-access"}, []string(defaultRoles))
|
|
}
|
|
|
|
func testMigration000602RestoreAgentsAccessDefaultRole(t *testing.T, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 602
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
db := database.New(sqlDB)
|
|
|
|
orgs := []struct {
|
|
defaults []string
|
|
want []string
|
|
}{
|
|
{[]string{"organization-workspace-access"}, []string{"organization-workspace-access", "agents-access"}},
|
|
{[]string{}, []string{"agents-access"}},
|
|
{[]string{"organization-workspace-access", "organization-auditor"}, []string{"organization-workspace-access", "organization-auditor", "agents-access"}},
|
|
}
|
|
orgIDs := make([]uuid.UUID, len(orgs))
|
|
for i, o := range orgs {
|
|
org := dbgen.Organization(t, db, database.Organization{})
|
|
// dbgen replaces an empty slice with the current defaults.
|
|
_, err := sqlDB.ExecContext(ctx, "UPDATE organizations SET default_org_member_roles = $1 WHERE id = $2", pq.StringArray(o.defaults), org.ID)
|
|
require.NoError(t, err)
|
|
orgIDs[i] = org.ID
|
|
}
|
|
|
|
member := dbgen.User(t, db, database.User{})
|
|
dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
|
OrganizationID: orgIDs[0],
|
|
UserID: member.ID,
|
|
Roles: []string{"organization-auditor"},
|
|
})
|
|
serviceAccount := dbgen.User(t, db, database.User{IsServiceAccount: true})
|
|
dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
|
OrganizationID: orgIDs[0],
|
|
UserID: serviceAccount.ID,
|
|
Roles: []string{},
|
|
})
|
|
legacyRole := dbgen.CustomRole(t, db, database.CustomRole{
|
|
Name: "agents-access",
|
|
OrganizationID: uuid.NullUUID{UUID: orgIDs[1], Valid: true},
|
|
})
|
|
unrelatedRole := dbgen.CustomRole(t, db, database.CustomRole{
|
|
OrganizationID: uuid.NullUUID{UUID: orgIDs[1], Valid: true},
|
|
})
|
|
|
|
memberRoles := func(userID uuid.UUID) []string {
|
|
t.Helper()
|
|
var roles pq.StringArray
|
|
err := sqlDB.QueryRowContext(ctx, "SELECT roles FROM organization_members WHERE organization_id = $1 AND user_id = $2", orgIDs[0], userID).Scan(&roles)
|
|
require.NoError(t, err)
|
|
return roles
|
|
}
|
|
assertDefaults := func(withAgentsAccess bool) {
|
|
t.Helper()
|
|
for i, o := range orgs {
|
|
var roles pq.StringArray
|
|
err := sqlDB.QueryRowContext(ctx, "SELECT default_org_member_roles FROM organizations WHERE id = $1", orgIDs[i]).Scan(&roles)
|
|
require.NoError(t, err)
|
|
want := o.defaults
|
|
if withAgentsAccess {
|
|
want = o.want
|
|
}
|
|
require.Equal(t, want, []string(roles), "org %d", i)
|
|
}
|
|
}
|
|
customRoleExists := func(id uuid.UUID) bool {
|
|
t.Helper()
|
|
var exists bool
|
|
err := sqlDB.QueryRowContext(ctx, "SELECT EXISTS (SELECT 1 FROM custom_roles WHERE id = $1)", id).Scan(&exists)
|
|
require.NoError(t, err)
|
|
return exists
|
|
}
|
|
|
|
version, _, err := next()
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, migrationVersion, version)
|
|
assertDefaults(true)
|
|
require.Equal(t, []string{"organization-auditor"}, memberRoles(member.ID))
|
|
require.Empty(t, memberRoles(serviceAccount.ID))
|
|
require.False(t, customRoleExists(legacyRole.ID))
|
|
require.True(t, customRoleExists(unrelatedRole.ID))
|
|
|
|
// An explicit grant made after the upgrade is removed on downgrade.
|
|
_, err = sqlDB.ExecContext(ctx, "UPDATE organization_members SET roles = array_append(roles, 'agents-access') WHERE user_id = $1", serviceAccount.ID)
|
|
require.NoError(t, err)
|
|
|
|
downSQL, err := os.ReadFile("000602_restore_agents_access_default_role.down.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
assertDefaults(false)
|
|
require.Equal(t, []string{"organization-auditor"}, memberRoles(member.ID))
|
|
require.Empty(t, memberRoles(serviceAccount.ID))
|
|
|
|
upSQL, err := os.ReadFile("000602_restore_agents_access_default_role.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
assertDefaults(true)
|
|
}
|
|
|
|
func testMigration000606ChatDiffStatusOriginCredentials(t *testing.T, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 606
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
ownerID, orgID, providerID, modelConfigID := uuid.New(), uuid.New(), uuid.New(), uuid.New()
|
|
for _, fixture := range []struct {
|
|
query string
|
|
args []any
|
|
}{
|
|
{
|
|
"INSERT INTO users (id, username, email, hashed_password, created_at, updated_at) VALUES ($1, $2, $3, '', NOW(), NOW())",
|
|
[]any{ownerID, ownerID.String(), ownerID.String() + "@test.com"},
|
|
},
|
|
{
|
|
"INSERT INTO organizations (id, name, display_name, description, created_at, updated_at, default_org_member_roles) VALUES ($1, $2, '', '', NOW(), NOW(), '{}')",
|
|
[]any{orgID, orgID.String()},
|
|
},
|
|
{
|
|
"INSERT INTO ai_providers (id, type, name, base_url) VALUES ($1, 'openai', $2, 'https://example.com')",
|
|
[]any{providerID, providerID.String()},
|
|
},
|
|
{
|
|
"INSERT INTO chat_model_configs (id, model, context_limit, compression_threshold, ai_provider_id, organization_id) VALUES ($1, 'model', 1000, 50, $2, $3)",
|
|
[]any{modelConfigID, providerID, orgID},
|
|
},
|
|
} {
|
|
_, err := sqlDB.ExecContext(ctx, fixture.query, fixture.args...)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
origin string
|
|
want string
|
|
}{
|
|
{"strips a token", "https://ghp_abc@github.com/o/r.git", "https://github.com/o/r.git"},
|
|
{"strips a user and password", "https://oauth2:tok@gitlab.com/g/r.git", "https://gitlab.com/g/r.git"},
|
|
{"strips a password with an at sign", "https://user:p@ss@github.com/o/r.git", "https://github.com/o/r.git"},
|
|
{"strips over http", "http://tok@git.example.com/o/r.git", "http://git.example.com/o/r.git"},
|
|
{"strips and keeps the port", "https://tok@github.com:443/o/r", "https://github.com:443/o/r"},
|
|
{"strips without a path", "https://tok@github.com", "https://github.com"},
|
|
{"strips and keeps an at sign in the path", "https://tok@github.com/o/r@v1", "https://github.com/o/r@v1"},
|
|
|
|
{"keeps an origin without credentials", "https://github.com/o/r.git", "https://github.com/o/r.git"},
|
|
{"keeps an at sign in the path", "https://github.com/o/r@v1", "https://github.com/o/r@v1"},
|
|
{"keeps an scp-style SSH user", "git@github.com:o/r.git", "git@github.com:o/r.git"},
|
|
{"keeps an SSH URL user", "ssh://git@github.com/o/r.git", "ssh://git@github.com/o/r.git"},
|
|
{"keeps an empty origin", "", ""},
|
|
}
|
|
|
|
// The old key is chat_id, so each case needs its own chat.
|
|
chatIDs := make([]uuid.UUID, len(tests))
|
|
for i, tt := range tests {
|
|
chatIDs[i] = uuid.New()
|
|
_, err := sqlDB.ExecContext(ctx,
|
|
"INSERT INTO chats (id, owner_id, organization_id, last_model_config_id) VALUES ($1, $2, $3, $4)",
|
|
chatIDs[i], ownerID, orgID, modelConfigID)
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
"INSERT INTO chat_diff_statuses (chat_id, git_remote_origin) VALUES ($1, $2)",
|
|
chatIDs[i], tt.origin)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
version, _, err := next()
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, migrationVersion, version)
|
|
|
|
for i, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var got string
|
|
err := sqlDB.QueryRowContext(ctx, "SELECT git_remote_origin FROM chat_diff_statuses WHERE chat_id = $1", chatIDs[i]).Scan(&got)
|
|
require.NoError(t, err)
|
|
require.Equal(t, tt.want, got)
|
|
})
|
|
}
|
|
}
|
|
|
|
func testMigration000621Skills(t *testing.T, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 621
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
insertUser := func() uuid.UUID {
|
|
t.Helper()
|
|
id := uuid.New()
|
|
_, err := sqlDB.ExecContext(ctx,
|
|
"INSERT INTO users (id, username, email, hashed_password, created_at, updated_at) VALUES ($1, $2, $3, '', NOW(), NOW())",
|
|
id, id.String(), id.String()+"@test.com")
|
|
require.NoError(t, err)
|
|
return id
|
|
}
|
|
type skillTable struct{ name, insert, exists string }
|
|
userSkillsTable := skillTable{
|
|
name: "user_skills",
|
|
insert: "INSERT INTO user_skills (id, user_id, name, description, content) VALUES ($1, $2, $3, 'desc', 'body')",
|
|
exists: "SELECT EXISTS (SELECT 1 FROM user_skills WHERE id = $1)",
|
|
}
|
|
skillsTable := skillTable{
|
|
name: "skills",
|
|
insert: "INSERT INTO skills (id, user_id, name, description, content) VALUES ($1, $2, $3, 'desc', 'body')",
|
|
exists: "SELECT EXISTS (SELECT 1 FROM skills WHERE id = $1)",
|
|
}
|
|
insertSkill := func(table skillTable, userID uuid.UUID) uuid.UUID {
|
|
t.Helper()
|
|
id := uuid.New()
|
|
_, err := sqlDB.ExecContext(ctx, table.insert, id, userID, "skill-"+id.String())
|
|
require.NoError(t, err)
|
|
return id
|
|
}
|
|
skillExists := func(table skillTable, id uuid.UUID) bool {
|
|
t.Helper()
|
|
var exists bool
|
|
err := sqlDB.QueryRowContext(ctx, table.exists, id).Scan(&exists)
|
|
require.NoError(t, err)
|
|
return exists
|
|
}
|
|
// Soft deletion runs delete_deleted_user_resources, whose body must name
|
|
// the current table.
|
|
softDeleteRemovesSkills := func(table skillTable) {
|
|
t.Helper()
|
|
userID := insertUser()
|
|
skillID := insertSkill(table, userID)
|
|
_, err := sqlDB.ExecContext(ctx, "UPDATE users SET deleted = true WHERE id = $1", userID)
|
|
require.NoError(t, err)
|
|
require.False(t, skillExists(table, skillID))
|
|
}
|
|
capFunction := func(table skillTable) string {
|
|
t.Helper()
|
|
var def string
|
|
err := sqlDB.QueryRowContext(ctx, `
|
|
SELECT pg_get_functiondef(tgfoid) FROM pg_trigger
|
|
WHERE tgrelid = $1::regclass AND tgname LIKE 'trigger_%skills_per_%'`, table.name).Scan(&def)
|
|
require.NoError(t, err)
|
|
return def
|
|
}
|
|
userNullable := func(table skillTable) string {
|
|
t.Helper()
|
|
var nullable string
|
|
err := sqlDB.QueryRowContext(ctx,
|
|
"SELECT is_nullable FROM information_schema.columns WHERE table_name = $1 AND column_name = 'user_id'",
|
|
table.name).Scan(&nullable)
|
|
require.NoError(t, err)
|
|
return nullable
|
|
}
|
|
|
|
oldNames := []string{
|
|
"trigger_upsert_user_skills:insert_user_skill_fail_if_user_deleted",
|
|
"trigger_user_skills_per_user_limit:enforce_user_skills_per_user_limit",
|
|
"user_skills_content_size",
|
|
"user_skills_description_size",
|
|
"user_skills_name_format",
|
|
"user_skills_name_size",
|
|
"user_skills_pkey",
|
|
"user_skills_user_id_fkey",
|
|
"user_skills_user_id_name_idx",
|
|
}
|
|
require.Equal(t, oldNames, skillsSchemaNames(ctx, t, sqlDB, userSkillsTable.name))
|
|
oldCapFunction := capFunction(userSkillsTable)
|
|
require.Equal(t, "NO", userNullable(userSkillsTable))
|
|
|
|
userID := insertUser()
|
|
personalSkillID := insertSkill(userSkillsTable, userID)
|
|
|
|
version, _, err := next()
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, migrationVersion, version)
|
|
|
|
require.Equal(t, []string{
|
|
"skills_acl_only_on_organization_skills",
|
|
"skills_content_size",
|
|
"skills_description_size",
|
|
"skills_group_acl_is_object",
|
|
"skills_name_format",
|
|
"skills_name_size",
|
|
"skills_organization_id_fkey",
|
|
"skills_organization_id_name_idx",
|
|
"skills_pkey",
|
|
"skills_project_id_fkey",
|
|
"skills_project_id_name_idx",
|
|
"skills_single_owner",
|
|
"skills_user_acl_is_object",
|
|
"skills_user_id_fkey",
|
|
"skills_user_id_name_idx",
|
|
"trigger_skills_per_owner_limit:enforce_skills_per_owner_limit",
|
|
"trigger_upsert_skills:insert_user_skill_fail_if_user_deleted",
|
|
}, skillsSchemaNames(ctx, t, sqlDB, skillsTable.name))
|
|
require.Equal(t, "YES", userNullable(skillsTable))
|
|
|
|
var (
|
|
gotUserID, orgID, projectID uuid.NullUUID
|
|
name, description, content string
|
|
enabled bool
|
|
groupACL, userACL string
|
|
)
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT user_id, organization_id, project_id, name, description, content, enabled, group_acl::text, user_acl::text FROM skills WHERE id = $1",
|
|
personalSkillID).Scan(&gotUserID, &orgID, &projectID, &name, &description, &content, &enabled, &groupACL, &userACL)
|
|
require.NoError(t, err)
|
|
require.Equal(t, uuid.NullUUID{UUID: userID, Valid: true}, gotUserID)
|
|
require.False(t, orgID.Valid)
|
|
require.False(t, projectID.Valid)
|
|
require.Equal(t, "skill-"+personalSkillID.String(), name)
|
|
require.Equal(t, "desc", description)
|
|
require.Equal(t, "body", content)
|
|
require.True(t, enabled)
|
|
require.Equal(t, "{}", groupACL)
|
|
require.Equal(t, "{}", userACL)
|
|
softDeleteRemovesSkills(skillsTable)
|
|
|
|
org := uuid.New()
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
"INSERT INTO organizations (id, name, display_name, description, created_at, updated_at, default_org_member_roles) VALUES ($1, $2, $2, '', NOW(), NOW(), '{}')",
|
|
org, "org-"+org.String())
|
|
require.NoError(t, err)
|
|
orgSkillID := uuid.New()
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
"INSERT INTO skills (id, organization_id, name, description, content) VALUES ($1, $2, 'org', 'desc', 'body')",
|
|
orgSkillID, org)
|
|
require.NoError(t, err)
|
|
|
|
downSQL, err := os.ReadFile("000621_skills.down.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
require.Equal(t, oldNames, skillsSchemaNames(ctx, t, sqlDB, userSkillsTable.name))
|
|
require.Equal(t, oldCapFunction, capFunction(userSkillsTable))
|
|
require.Equal(t, "NO", userNullable(userSkillsTable))
|
|
require.True(t, skillExists(userSkillsTable, personalSkillID))
|
|
require.False(t, skillExists(userSkillsTable, orgSkillID))
|
|
softDeleteRemovesSkills(userSkillsTable)
|
|
|
|
upSQL, err := os.ReadFile("000621_skills.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
require.True(t, skillExists(skillsTable, personalSkillID))
|
|
softDeleteRemovesSkills(skillsTable)
|
|
}
|
|
|
|
// skillsSchemaNames lists a table's constraints, indexes, and triggers with
|
|
// their functions.
|
|
func skillsSchemaNames(ctx context.Context, t *testing.T, sqlDB *sql.DB, table string) []string {
|
|
t.Helper()
|
|
rows, err := sqlDB.QueryContext(ctx, `
|
|
SELECT conname::text FROM pg_constraint WHERE conrelid = $1::regclass
|
|
UNION SELECT indexname::text FROM pg_indexes WHERE tablename = $1::text
|
|
UNION SELECT tgname::text || ':' || tgfoid::regproc::text FROM pg_trigger WHERE tgrelid = $1::regclass AND NOT tgisinternal
|
|
ORDER BY 1`, table)
|
|
require.NoError(t, err)
|
|
defer rows.Close()
|
|
var names []string
|
|
for rows.Next() {
|
|
var name string
|
|
require.NoError(t, rows.Scan(&name))
|
|
names = append(names, name)
|
|
}
|
|
require.NoError(t, rows.Err())
|
|
return names
|
|
}
|
|
|
|
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 TestMigration000542ChatReasoningEffortBackfill(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
const priorMigrationVersion = 539
|
|
|
|
sqlDB := testSQLDB(t)
|
|
|
|
next, err := migrations.Stepper(sqlDB)
|
|
require.NoError(t, err)
|
|
for {
|
|
version, more, err := next()
|
|
require.NoError(t, err)
|
|
if !more || version == priorMigrationVersion {
|
|
break
|
|
}
|
|
}
|
|
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
|
|
tx, err := sqlDB.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
defer tx.Rollback()
|
|
azureID := uuid.New()
|
|
bedrockID := uuid.New()
|
|
emptyID := uuid.New()
|
|
invalidID := uuid.New()
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO ai_providers (id, type, name, enabled, base_url, created_at, updated_at)
|
|
VALUES
|
|
($1, 'azure', 'test-azure-reasoning', TRUE, '', $3, $3),
|
|
($2, 'bedrock', 'test-bedrock-reasoning', TRUE, '', $3, $3)
|
|
`, azureID, bedrockID, now)
|
|
require.NoError(t, err)
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO chat_model_configs (id, ai_provider_id, model, display_name, enabled, context_limit, compression_threshold, options, created_at, updated_at)
|
|
VALUES
|
|
($3, $1, 'gpt-5.1-azure', 'Azure GPT-5.1', TRUE, 200000, 70, '{"provider_options": {"azure": {"reasoning_effort": " LOW "}}}', $5, $5),
|
|
($4, $2, 'anthropic.claude-opus-4-6', 'Bedrock Claude Opus', TRUE, 200000, 70, '{"provider_options": {"bedrock": {"effort": "minimal"}}}', $5, $5),
|
|
(gen_random_uuid(), $1, 'gpt-5.1-empty-effort', 'Azure Empty Effort', TRUE, 200000, 70, '{"provider_options": {"azure": {"reasoning_effort": ""}}}', $5, $5),
|
|
(gen_random_uuid(), $2, 'anthropic.invalid-effort', 'Bedrock Invalid Effort', TRUE, 200000, 70, '{"provider_options": {"bedrock": {"effort": "extreme"}}}', $5, $5)
|
|
`, azureID, bedrockID, emptyID, invalidID, now)
|
|
require.NoError(t, err)
|
|
require.NoError(t, tx.Commit())
|
|
|
|
migrationSQL, err := os.ReadFile("000542_chat_reasoning_effort.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(migrationSQL))
|
|
require.NoError(t, err)
|
|
|
|
rows, err := sqlDB.QueryContext(ctx, `
|
|
SELECT ap.type::text, cmc.model, cmc.options->'reasoning_effort'->>'default'
|
|
FROM chat_model_configs cmc
|
|
JOIN ai_providers ap ON ap.id = cmc.ai_provider_id
|
|
WHERE cmc.ai_provider_id IN ($1, $2)
|
|
ORDER BY cmc.model
|
|
`, azureID, bedrockID)
|
|
require.NoError(t, err)
|
|
defer rows.Close()
|
|
|
|
got := map[string]sql.NullString{}
|
|
for rows.Next() {
|
|
var provider, model string
|
|
var effort sql.NullString
|
|
require.NoError(t, rows.Scan(&provider, &model, &effort))
|
|
got[provider+":"+model] = effort
|
|
}
|
|
require.NoError(t, rows.Err())
|
|
require.Equal(t, sql.NullString{String: "low", Valid: true}, got["azure:gpt-5.1-azure"])
|
|
require.Equal(t, sql.NullString{}, got["azure:gpt-5.1-empty-effort"])
|
|
require.Equal(t, sql.NullString{String: "minimal", Valid: true}, got["bedrock:anthropic.claude-opus-4-6"])
|
|
require.Equal(t, sql.NullString{}, got["bedrock:anthropic.invalid-effort"])
|
|
}
|
|
|
|
func testMigration000546ChatHistoryAPIKeyConstraints(t *testing.T, sqlDB *sql.DB, _ migrationStepper) {
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
constraintNames := []string{
|
|
"chat_messages_api_key_id_fkey",
|
|
"chat_queued_messages_api_key_id_fkey",
|
|
}
|
|
assertConstraintCount := func(t *testing.T, want int) {
|
|
t.Helper()
|
|
for _, name := range constraintNames {
|
|
var got int
|
|
err := sqlDB.QueryRowContext(ctx, `
|
|
SELECT COUNT(*)
|
|
FROM pg_constraint
|
|
WHERE conname = $1
|
|
`, name).Scan(&got)
|
|
require.NoError(t, err)
|
|
require.Equal(t, want, got, name)
|
|
}
|
|
}
|
|
|
|
upSQL, err := os.ReadFile("000546_drop_chat_history_api_key_fks.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
assertConstraintCount(t, 0)
|
|
|
|
downSQL, err := os.ReadFile("000546_drop_chat_history_api_key_fks.down.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
assertConstraintCount(t, 1)
|
|
|
|
for _, name := range constraintNames {
|
|
var count int
|
|
err := sqlDB.QueryRowContext(ctx, `
|
|
SELECT COUNT(*)
|
|
FROM pg_constraint
|
|
WHERE conname = $1 AND confdeltype = 'n'
|
|
`, name).Scan(&count)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, count, name)
|
|
}
|
|
}
|
|
|
|
func TestMigration000555LegacyNoneLoginToPassword(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
const priorMigrationVersion = 554
|
|
|
|
sqlDB := testSQLDB(t)
|
|
|
|
next, err := migrations.Stepper(sqlDB)
|
|
require.NoError(t, err)
|
|
for {
|
|
version, more, err := next()
|
|
require.NoError(t, err)
|
|
if !more || version == priorMigrationVersion {
|
|
break
|
|
}
|
|
}
|
|
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
|
|
legacyNoneID := uuid.New()
|
|
serviceAccountID := uuid.New()
|
|
systemID := uuid.New()
|
|
passwordID := uuid.New()
|
|
|
|
// A legacy machine user: login_type 'none', not a service account, not a
|
|
// system user. This is the only row the migration should convert.
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
`INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type, is_service_account, is_system)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`,
|
|
legacyNoneID, "legacy-none", "legacy-none@test.com", []byte{}, now, now, "active", pq.StringArray{}, "none", false, false)
|
|
require.NoError(t, err)
|
|
|
|
// A service account must keep login_type 'none' (a CHECK constraint requires
|
|
// service accounts to use 'none' and an empty email).
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
`INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type, is_service_account, is_system)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`,
|
|
serviceAccountID, "service-account", "", []byte{}, now, now, "active", pq.StringArray{}, "none", true, false)
|
|
require.NoError(t, err)
|
|
|
|
// A system user must be left untouched.
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
`INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type, is_service_account, is_system)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`,
|
|
systemID, "system-user", "system@test.com", []byte{}, now, now, "active", pq.StringArray{}, "none", false, true)
|
|
require.NoError(t, err)
|
|
|
|
// An existing password user must be left untouched.
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
`INSERT INTO users (id, username, email, hashed_password, created_at, updated_at, status, rbac_roles, login_type, is_service_account, is_system)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`,
|
|
passwordID, "password-user", "password@test.com", []byte("hashed"), now, now, "active", pq.StringArray{}, "password", false, false)
|
|
require.NoError(t, err)
|
|
|
|
migrationSQL, err := os.ReadFile("000555_legacy_none_login_to_password.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(migrationSQL))
|
|
require.NoError(t, err)
|
|
|
|
getUser := func(t *testing.T, id uuid.UUID) (loginType, email string) {
|
|
t.Helper()
|
|
err := sqlDB.QueryRowContext(ctx,
|
|
`SELECT login_type::text, email FROM users WHERE id = $1`, id).Scan(&loginType, &email)
|
|
require.NoError(t, err)
|
|
return loginType, email
|
|
}
|
|
|
|
// The legacy machine user is converted to password auth with its email
|
|
// preserved.
|
|
gotLoginType, gotEmail := getUser(t, legacyNoneID)
|
|
require.Equal(t, "password", gotLoginType)
|
|
require.Equal(t, "legacy-none@test.com", gotEmail)
|
|
|
|
// Service accounts, system users, and existing password users are unchanged.
|
|
gotLoginType, _ = getUser(t, serviceAccountID)
|
|
require.Equal(t, "none", gotLoginType)
|
|
gotLoginType, _ = getUser(t, systemID)
|
|
require.Equal(t, "none", gotLoginType)
|
|
gotLoginType, _ = getUser(t, passwordID)
|
|
require.Equal(t, "password", gotLoginType)
|
|
}
|
|
|
|
// TestMigration000558AuditOAuth2ProviderSettingsEnumInSingleTxn reproduces
|
|
// the production upgrade path, where every pending migration in a deploy
|
|
// runs inside a single transaction (see pgTxnDriver). 000558 adds
|
|
// 'oauth2_provider_settings' to the resource_type enum via ALTER TYPE ...
|
|
// ADD VALUE. Postgres forbids using an enum value added by ADD VALUE within
|
|
// the same transaction that added it, so this confirms the audit write path
|
|
// (a separate transaction, exactly like commitAudit() performs it for a real
|
|
// PUT to the DCR settings endpoint) can use the new value immediately after
|
|
// the migration transaction commits, and that pre-existing audit data from
|
|
// before the upgrade survives untouched.
|
|
func TestMigration000558AuditOAuth2ProviderSettingsEnumInSingleTxn(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
sqlDB := testSQLDB(t)
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
|
|
// Apply everything through 557 and commit, simulating a deployment
|
|
// that was already running the previous release, with real
|
|
// pre-existing audit data, before the upgrade that adds 558.
|
|
applyMigrationsInTxn(ctx, t, sqlDB, 1, 557)
|
|
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
preUpgradeLogID := uuid.New()
|
|
_, err := sqlDB.ExecContext(ctx, `
|
|
INSERT INTO audit_logs (
|
|
id, time, user_id, organization_id, resource_type, resource_id,
|
|
resource_target, action, diff, status_code, additional_fields,
|
|
request_id, resource_icon
|
|
) VALUES (
|
|
$1, $2, $3, $4, 'oauth2_provider_app', $5, 'pre-upgrade-app',
|
|
'write', '{}', 200, '{}', $6, ''
|
|
)`,
|
|
preUpgradeLogID, now, uuid.New(), uuid.New(), uuid.New(), uuid.New(),
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
// Apply 558 in the same single transaction production uses for the
|
|
// whole pending batch.
|
|
applyMigrationsInTxn(ctx, t, sqlDB, 558, 558)
|
|
|
|
// Pre-existing audit data survives the upgrade untouched.
|
|
var resourceTarget string
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
`SELECT resource_target FROM audit_logs WHERE id = $1`, preUpgradeLogID,
|
|
).Scan(&resourceTarget)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "pre-upgrade-app", resourceTarget)
|
|
|
|
// The new enum value is usable immediately after the migration
|
|
// transaction commits, in a separate transaction, exactly like a real
|
|
// PUT to the DCR settings endpoint would write it.
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
INSERT INTO audit_logs (
|
|
id, time, user_id, organization_id, resource_type, resource_id,
|
|
resource_target, action, diff, status_code, additional_fields,
|
|
request_id, resource_icon
|
|
) VALUES (
|
|
$1, $2, $3, $4, 'oauth2_provider_settings', $5, '', 'write',
|
|
'{}', 200, '{}', $6, ''
|
|
)`,
|
|
uuid.New(), now, uuid.New(), uuid.New(), uuid.New(), uuid.New(),
|
|
)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
//nolint:tparallel,paralleltest // Subtests share one database and exercise sequential migration state.
|
|
func TestMigration000563TemplateAgentsAllowedBackfill(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
sqlDB, ctx, orgID, userID, templateIDs := setupMigration000563Templates(t)
|
|
upSQL, err := os.ReadFile("000563_template_agents_allowed.up.sql")
|
|
require.NoError(t, err)
|
|
downSQL, err := os.ReadFile("000563_template_agents_allowed.down.sql")
|
|
require.NoError(t, err)
|
|
|
|
staleID := uuid.New()
|
|
tests := []struct {
|
|
name string
|
|
value string
|
|
present bool
|
|
checkPostMigrationDefault bool
|
|
want map[uuid.UUID]bool
|
|
}{
|
|
{
|
|
name: "valid nonempty list",
|
|
value: fmt.Sprintf(`[%q]`, templateIDs[0]),
|
|
present: true,
|
|
checkPostMigrationDefault: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: true,
|
|
templateIDs[1]: false,
|
|
},
|
|
},
|
|
{
|
|
name: "stale template ID",
|
|
value: fmt.Sprintf(`[%q,%q]`, templateIDs[0], staleID),
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: true,
|
|
templateIDs[1]: false,
|
|
},
|
|
},
|
|
{
|
|
name: "missing",
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: true,
|
|
templateIDs[1]: true,
|
|
},
|
|
},
|
|
{
|
|
name: "empty string",
|
|
value: "",
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: true,
|
|
templateIDs[1]: true,
|
|
},
|
|
},
|
|
{
|
|
name: "JSON null",
|
|
value: "null",
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: true,
|
|
templateIDs[1]: true,
|
|
},
|
|
},
|
|
{
|
|
name: "invalid JSON",
|
|
value: "{",
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: false,
|
|
templateIDs[1]: false,
|
|
},
|
|
},
|
|
{
|
|
name: "JSON object",
|
|
value: `{}`,
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: false,
|
|
templateIDs[1]: false,
|
|
},
|
|
},
|
|
{
|
|
name: "JSON scalar",
|
|
value: `"value"`,
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: false,
|
|
templateIDs[1]: false,
|
|
},
|
|
},
|
|
{
|
|
name: "invalid UUID element",
|
|
value: `["not-a-uuid"]`,
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: false,
|
|
templateIDs[1]: false,
|
|
},
|
|
},
|
|
{
|
|
name: "null element",
|
|
value: `[null]`,
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: false,
|
|
templateIDs[1]: false,
|
|
},
|
|
},
|
|
{
|
|
name: "mixed valid and invalid elements",
|
|
value: fmt.Sprintf(`[%q,"not-a-uuid"]`, templateIDs[0]),
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: false,
|
|
templateIDs[1]: false,
|
|
},
|
|
},
|
|
{
|
|
name: "empty array",
|
|
value: "[]",
|
|
present: true,
|
|
want: map[uuid.UUID]bool{
|
|
templateIDs[0]: true,
|
|
templateIDs[1]: true,
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
_, err := sqlDB.ExecContext(ctx, `DELETE FROM site_configs WHERE key = 'agents_template_allowlist'`)
|
|
require.NoError(t, err)
|
|
if tt.present {
|
|
_, err = sqlDB.ExecContext(ctx, `INSERT INTO site_configs (key, value) VALUES ('agents_template_allowlist', $1)`, tt.value)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
_, err = sqlDB.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() {
|
|
_, err := sqlDB.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
})
|
|
|
|
rows, err := sqlDB.QueryContext(ctx, `SELECT id, agents_allowed FROM templates`)
|
|
require.NoError(t, err)
|
|
got := make(map[uuid.UUID]bool, len(templateIDs))
|
|
for rows.Next() {
|
|
var id uuid.UUID
|
|
var agentsAllowed bool
|
|
require.NoError(t, rows.Scan(&id, &agentsAllowed))
|
|
got[id] = agentsAllowed
|
|
}
|
|
require.NoError(t, rows.Close())
|
|
require.NoError(t, rows.Err())
|
|
require.Equal(t, tt.want, got)
|
|
|
|
var stored string
|
|
err = sqlDB.QueryRowContext(ctx, `SELECT value FROM site_configs WHERE key = 'agents_template_allowlist'`).Scan(&stored)
|
|
if tt.present {
|
|
require.NoError(t, err)
|
|
require.Equal(t, tt.value, stored)
|
|
} else {
|
|
require.ErrorIs(t, err, sql.ErrNoRows)
|
|
}
|
|
|
|
if tt.checkPostMigrationDefault {
|
|
newTemplateID := uuid.New()
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
INSERT INTO templates (id, organization_id, name, created_at, updated_at, provisioner, active_version_id, created_by)
|
|
VALUES ($1, $2, $3, NOW(), NOW(), 'terraform', $4, $5)
|
|
`, newTemplateID, orgID, "post-migration-template", uuid.New(), userID)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() {
|
|
_, err := sqlDB.ExecContext(ctx, `DELETE FROM templates WHERE id = $1`, newTemplateID)
|
|
require.NoError(t, err)
|
|
})
|
|
|
|
var agentsAllowed bool
|
|
err = sqlDB.QueryRowContext(ctx, `SELECT agents_allowed FROM template_with_names WHERE id = $1`, newTemplateID).Scan(&agentsAllowed)
|
|
require.NoError(t, err)
|
|
require.True(t, agentsAllowed)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func setupMigration000563Templates(t *testing.T) (
|
|
sqlDB *sql.DB,
|
|
ctx context.Context,
|
|
orgID uuid.UUID,
|
|
userID uuid.UUID,
|
|
templateIDs []uuid.UUID,
|
|
) {
|
|
t.Helper()
|
|
|
|
const migrationVersion = 562
|
|
|
|
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)
|
|
orgID = uuid.New()
|
|
userID = uuid.New()
|
|
templateIDs = []uuid.UUID{uuid.New(), uuid.New()}
|
|
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
INSERT INTO organizations (id, name, display_name, description, created_at, updated_at, default_org_member_roles)
|
|
VALUES ($1, $2, $3, $4, $5, $5, '{}')
|
|
`, orgID, "agents-allowed-org", "Agents Allowed Org", "Migration test", now)
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.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, $5, 'active', '{}', 'password')
|
|
`, userID, "agents-allowed-user", "agents-allowed@example.com", []byte{}, now)
|
|
require.NoError(t, err)
|
|
for i, templateID := range templateIDs {
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
INSERT INTO templates (id, organization_id, name, created_at, updated_at, provisioner, active_version_id, created_by)
|
|
VALUES ($1, $2, $3, $4, $4, 'terraform', $5, $6)
|
|
`, templateID, orgID, fmt.Sprintf("agents-allowed-template-%d", i), now, uuid.New(), userID)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
return sqlDB, ctx, orgID, userID, templateIDs
|
|
}
|
|
|
|
func testMigration000498SoftDeleteStaleWorkspaceAgents(t *testing.T, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 498
|
|
|
|
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.
|
|
}
|
|
|
|
func testMigration000543ChatMessageSearchText(t *testing.T, sqlDB *sql.DB) {
|
|
cases := []struct {
|
|
name string
|
|
content sql.NullString
|
|
want sql.NullString
|
|
}{
|
|
{
|
|
name: "SingleTextPart",
|
|
content: sql.NullString{String: `[{"type":"text","text":"hello world"}]`, Valid: true},
|
|
want: sql.NullString{String: "hello world", Valid: true},
|
|
},
|
|
{
|
|
name: "TextInterleavedWithNonText",
|
|
content: sql.NullString{String: `[
|
|
{"type":"text","text":"first"},
|
|
{"type":"reasoning","text":"thinking"},
|
|
{"type":"tool-call","toolName":"execute"},
|
|
{"type":"text","text":"second"}
|
|
]`, Valid: true},
|
|
want: sql.NullString{String: "first second", Valid: true},
|
|
},
|
|
{
|
|
name: "OnlyNonTextParts",
|
|
content: sql.NullString{String: `[{"type":"reasoning","text":"thinking"}]`, Valid: true},
|
|
want: sql.NullString{},
|
|
},
|
|
{
|
|
name: "ScalarContent",
|
|
content: sql.NullString{String: `"hello"`, Valid: true},
|
|
want: sql.NullString{},
|
|
},
|
|
{
|
|
name: "EmptyArray",
|
|
content: sql.NullString{String: `[]`, Valid: true},
|
|
want: sql.NullString{},
|
|
},
|
|
{
|
|
name: "NullInput",
|
|
content: sql.NullString{},
|
|
want: sql.NullString{},
|
|
},
|
|
{
|
|
name: "ElementsMissingTypeOrText",
|
|
content: sql.NullString{String: `[{"text":"no type"},{"type":"text"},{"type":"text","text":"kept"}]`, Valid: true},
|
|
want: sql.NullString{String: "kept", Valid: true},
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
|
var got sql.NullString
|
|
err := sqlDB.QueryRowContext(ctx,
|
|
`SELECT chat_message_search_text($1::jsonb)`, tc.content,
|
|
).Scan(&got)
|
|
require.NoError(t, err)
|
|
require.Equal(t, tc.want, got)
|
|
})
|
|
}
|
|
}
|
|
|
|
// Shared eligibility predicate of the two partial chat_messages search
|
|
// indexes. Queries must repeat it verbatim.
|
|
const eligibilityPredicate = `deleted = false
|
|
AND visibility IN ('user', 'both')
|
|
AND role IN ('user', 'assistant')`
|
|
|
|
func testMigration000543ChatSearchSchemaIndexes(t *testing.T, sqlDB *sql.DB) {
|
|
cases := []struct {
|
|
name string
|
|
table string
|
|
partial bool
|
|
}{
|
|
{name: "idx_chat_messages_search_tsv", table: "chat_messages", partial: true},
|
|
{name: "idx_chat_messages_search_tsv_pending", table: "chat_messages", partial: true},
|
|
{name: "idx_chats_title_fts", table: "chats", partial: false},
|
|
{name: "idx_chat_diff_statuses_pr_title_fts", table: "chat_diff_statuses", partial: false},
|
|
}
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
|
var table string
|
|
var partial bool
|
|
err := sqlDB.QueryRowContext(ctx, `
|
|
SELECT i.tablename, x.indpred IS NOT NULL
|
|
FROM pg_indexes i
|
|
JOIN pg_class c ON c.relname = i.indexname
|
|
JOIN pg_index x ON x.indexrelid = c.oid
|
|
WHERE i.indexname = $1`, tc.name,
|
|
).Scan(&table, &partial)
|
|
require.NoError(t, err, "index %s should exist", tc.name)
|
|
require.Equal(t, tc.table, table, "index %s table", tc.name)
|
|
require.Equal(t, tc.partial, partial, "index %s partial", tc.name)
|
|
})
|
|
}
|
|
}
|
|
|
|
func testMigration000543ChatSearchSchemaBehavior(t *testing.T, sqlDB *sql.DB) {
|
|
db := database.New(sqlDB)
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
|
|
org := dbgen.Organization(t, db, database.Organization{})
|
|
owner := dbgen.User(t, db, database.User{})
|
|
_ = dbgen.ChatProvider(t, db, database.ChatProvider{Provider: "openai", DisplayName: "OpenAI"})
|
|
modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
|
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
|
UpdatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
|
IsDefault: true,
|
|
})
|
|
chat := dbgen.Chat(t, db, database.Chat{
|
|
OrganizationID: org.ID,
|
|
OwnerID: owner.ID,
|
|
LastModelConfigID: modelCfg.ID,
|
|
})
|
|
|
|
newMsg := func(role database.ChatMessageRole, visibility database.ChatMessageVisibility, content string) database.ChatMessage {
|
|
seed := database.ChatMessage{
|
|
ChatID: chat.ID,
|
|
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
|
Role: role,
|
|
Visibility: visibility,
|
|
}
|
|
if content != "" {
|
|
seed.Content = pqtype.NullRawMessage{RawMessage: []byte(content), Valid: true}
|
|
}
|
|
return dbgen.ChatMessage(t, db, seed)
|
|
}
|
|
textContent := func(text string) string {
|
|
return `[{"type":"text","text":"` + text + `"}]`
|
|
}
|
|
|
|
pendingIDs := func(ctx context.Context, limit int) []int64 {
|
|
rows, err := sqlDB.QueryContext(ctx, `
|
|
SELECT id FROM chat_messages
|
|
WHERE search_tsv IS NULL AND `+eligibilityPredicate+`
|
|
ORDER BY id DESC
|
|
LIMIT $1`, limit)
|
|
require.NoError(t, err)
|
|
defer rows.Close()
|
|
var ids []int64
|
|
for rows.Next() {
|
|
var id int64
|
|
require.NoError(t, rows.Scan(&id))
|
|
ids = append(ids, id)
|
|
}
|
|
require.NoError(t, rows.Err())
|
|
return ids
|
|
}
|
|
|
|
// Insert regression: RETURNING * must survive the new column, and new
|
|
// rows must start with search_tsv NULL so they enter the pending queue.
|
|
eligibleText := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("deploy the search feature"))
|
|
var tsvIsNull bool
|
|
err := sqlDB.QueryRowContext(ctx,
|
|
`SELECT search_tsv IS NULL FROM chat_messages WHERE id = $1`, eligibleText.ID,
|
|
).Scan(&tsvIsNull)
|
|
require.NoError(t, err)
|
|
require.True(t, tsvIsNull, "new rows must have search_tsv NULL")
|
|
|
|
eligibleNoText := newMsg(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, `[{"type":"reasoning","text":"thinking"}]`)
|
|
toolMsg := newMsg(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, textContent("tool output about deploy"))
|
|
modelOnly := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, textContent("model-only deploy note"))
|
|
deletedMsg := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("deleted deploy message"))
|
|
_, err = sqlDB.ExecContext(ctx, `UPDATE chat_messages SET deleted = true WHERE id = $1`, deletedMsg.ID)
|
|
require.NoError(t, err)
|
|
|
|
// Only eligible rows appear in the queue, newest first. The tool-role,
|
|
// model-only, and soft-deleted rows are excluded even though their
|
|
// search_tsv is NULL.
|
|
require.Equal(t, []int64{eligibleNoText.ID, eligibleText.ID}, pendingIDs(ctx, 10))
|
|
|
|
// Sweep-style UPDATE. The '' sentinel (not NULL) marks no-text rows as
|
|
// swept; search_tsv_config records the config that produced the vector.
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
UPDATE chat_messages
|
|
SET search_tsv = COALESCE(to_tsvector('english', chat_message_search_text(content)), ''::tsvector),
|
|
search_tsv_config = 'english'
|
|
WHERE id = ANY($1)`, pq.Array([]int64{eligibleText.ID, eligibleNoText.ID}))
|
|
require.NoError(t, err)
|
|
require.Empty(t, pendingIDs(ctx, 10), "swept rows must leave the queue, including no-text rows")
|
|
|
|
// A non-NULL vector with a stale config is not in the NULL queue but
|
|
// must be found by the stale-reindex predicate: this is how rows
|
|
// written by an old binary during a rolling upgrade (which cannot set
|
|
// search_tsv_config) are eventually rewritten.
|
|
staleIDs := func(ctx context.Context, limit int) []int64 {
|
|
rows, err := sqlDB.QueryContext(ctx, `
|
|
SELECT id FROM chat_messages
|
|
WHERE search_tsv IS NOT NULL
|
|
AND search_tsv_config IS DISTINCT FROM 'english'
|
|
AND `+eligibilityPredicate+`
|
|
ORDER BY id DESC
|
|
LIMIT $1`, limit)
|
|
require.NoError(t, err)
|
|
defer rows.Close()
|
|
var ids []int64
|
|
for rows.Next() {
|
|
var id int64
|
|
require.NoError(t, rows.Scan(&id))
|
|
ids = append(ids, id)
|
|
}
|
|
require.NoError(t, rows.Err())
|
|
return ids
|
|
}
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
UPDATE chat_messages
|
|
SET search_tsv = to_tsvector('simple', chat_message_search_text(content)),
|
|
search_tsv_config = NULL
|
|
WHERE id = $1`, eligibleText.ID)
|
|
require.NoError(t, err)
|
|
require.Empty(t, pendingIDs(ctx, 10), "stale rows must not enter the NULL queue")
|
|
require.Equal(t, []int64{eligibleText.ID}, staleIDs(ctx, 10),
|
|
"stale-config vectors must match the stale-reindex predicate")
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
UPDATE chat_messages
|
|
SET search_tsv = COALESCE(to_tsvector('english', chat_message_search_text(content)), ''::tsvector),
|
|
search_tsv_config = 'english'
|
|
WHERE id = $1`, eligibleText.ID)
|
|
require.NoError(t, err)
|
|
require.Empty(t, staleIDs(ctx, 10))
|
|
|
|
// Soft-deleting an unswept row removes it from the queue without a sweep.
|
|
unswept := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("unswept deploy row"))
|
|
require.Equal(t, []int64{unswept.ID}, pendingIDs(ctx, 10))
|
|
_, err = sqlDB.ExecContext(ctx, `UPDATE chat_messages SET deleted = true WHERE id = $1`, unswept.ID)
|
|
require.NoError(t, err)
|
|
require.Empty(t, pendingIDs(ctx, 10))
|
|
|
|
// Search contract: populate search_tsv on every row (including
|
|
// ineligible ones) and assert the search-index predicate filters them.
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
UPDATE chat_messages
|
|
SET search_tsv = COALESCE(to_tsvector('english', chat_message_search_text(content)), ''::tsvector),
|
|
search_tsv_config = 'english'
|
|
WHERE chat_id = $1`, chat.ID)
|
|
require.NoError(t, err)
|
|
|
|
rows, err := sqlDB.QueryContext(ctx, `
|
|
SELECT id FROM chat_messages
|
|
WHERE search_tsv @@ websearch_to_tsquery('english', $1)
|
|
AND search_tsv IS NOT NULL
|
|
AND `+eligibilityPredicate+`
|
|
ORDER BY id`, "deploy")
|
|
require.NoError(t, err)
|
|
defer rows.Close()
|
|
var matched []int64
|
|
for rows.Next() {
|
|
var id int64
|
|
require.NoError(t, rows.Scan(&id))
|
|
matched = append(matched, id)
|
|
}
|
|
require.NoError(t, rows.Err())
|
|
require.Equal(t, []int64{eligibleText.ID}, matched,
|
|
"search must exclude deleted, model-only, and tool-role rows (%d %d %d)",
|
|
toolMsg.ID, modelOnly.ID, deletedMsg.ID)
|
|
}
|
|
|
|
func TestMigration000585ChatSearchEnglishConfigDown(t *testing.T) {
|
|
t.Parallel()
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
}
|
|
|
|
sqlDB := testSQLDB(t)
|
|
require.NoError(t, migrations.Up(sqlDB))
|
|
db := database.New(sqlDB)
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
|
|
org := dbgen.Organization(t, db, database.Organization{})
|
|
owner := dbgen.User(t, db, database.User{})
|
|
_ = dbgen.ChatProvider(t, db, database.ChatProvider{Provider: "openai", DisplayName: "OpenAI"})
|
|
modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
|
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
|
UpdatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
|
IsDefault: true,
|
|
})
|
|
chat := dbgen.Chat(t, db, database.Chat{
|
|
OrganizationID: org.ID,
|
|
OwnerID: owner.ID,
|
|
LastModelConfigID: modelCfg.ID,
|
|
})
|
|
|
|
newMsg := func(text string) database.ChatMessage {
|
|
return dbgen.ChatMessage(t, db, database.ChatMessage{
|
|
ChatID: chat.ID,
|
|
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
|
Role: database.ChatMessageRoleUser,
|
|
Visibility: database.ChatMessageVisibilityBoth,
|
|
Content: pqtype.NullRawMessage{RawMessage: []byte(`[{"type":"text","text":"` + text + `"}]`), Valid: true},
|
|
})
|
|
}
|
|
|
|
// A row the upgraded sweep re-vectorized, a stale row it never
|
|
// reached, and a row that was still pending.
|
|
rewritten := newMsg("rewritten with english")
|
|
stale := newMsg("stale simple vector")
|
|
pending := newMsg("still pending")
|
|
|
|
_, err := sqlDB.ExecContext(ctx, `
|
|
UPDATE chat_messages
|
|
SET search_tsv = to_tsvector('english', chat_message_search_text(content)),
|
|
search_tsv_config = 'english'
|
|
WHERE id = $1`, rewritten.ID)
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
UPDATE chat_messages
|
|
SET search_tsv = to_tsvector('simple', chat_message_search_text(content)),
|
|
search_tsv_config = NULL
|
|
WHERE id = $1`, stale.ID)
|
|
require.NoError(t, err)
|
|
|
|
var revisionBefore int64
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
`SELECT revision FROM chat_messages WHERE id = $1`, rewritten.ID,
|
|
).Scan(&revisionBefore)
|
|
require.NoError(t, err)
|
|
|
|
downSQL, err := os.ReadFile("000585_chat_search_english_config.down.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
|
|
// The 'english' row is handed back to the parent sweep: NULL vector
|
|
// re-enters the parent's search_tsv IS NULL pending queue.
|
|
tsvState := func(id int64) (isNull bool, isSimple bool) {
|
|
err := sqlDB.QueryRowContext(ctx, `
|
|
SELECT search_tsv IS NULL,
|
|
search_tsv IS NOT DISTINCT FROM to_tsvector('simple', chat_message_search_text(content))
|
|
FROM chat_messages WHERE id = $1`, id,
|
|
).Scan(&isNull, &isSimple)
|
|
require.NoError(t, err)
|
|
return isNull, isSimple
|
|
}
|
|
isNull, _ := tsvState(rewritten.ID)
|
|
require.True(t, isNull, "english-stamped rows must be reset so the parent sweep reindexes them")
|
|
isNull, isSimple := tsvState(stale.ID)
|
|
require.False(t, isNull, "rows never re-vectorized must keep their vector")
|
|
require.True(t, isSimple, "rows never re-vectorized must keep their 'simple' lexemes")
|
|
isNull, _ = tsvState(pending.ID)
|
|
require.True(t, isNull, "pending rows stay pending")
|
|
|
|
// The reset ran under the replacement trigger functions, so it must
|
|
// not have advanced the message revision.
|
|
var revisionAfter int64
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
`SELECT revision FROM chat_messages WHERE id = $1`, rewritten.ID,
|
|
).Scan(&revisionAfter)
|
|
require.NoError(t, err)
|
|
require.Equal(t, revisionBefore, revisionAfter, "down migration reset must not bump revisions")
|
|
|
|
var columnCount int
|
|
err = sqlDB.QueryRowContext(ctx, `
|
|
SELECT COUNT(*) FROM information_schema.columns
|
|
WHERE table_name = 'chat_messages' AND column_name = 'search_tsv_config'`,
|
|
).Scan(&columnCount)
|
|
require.NoError(t, err)
|
|
require.Zero(t, columnCount, "search_tsv_config must be dropped")
|
|
}
|
|
|
|
func testMigration000556UserSecretsEnabled(t *testing.T, sqlDB *sql.DB, next migrationStepper) {
|
|
const migrationVersion = 556
|
|
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
|
|
userID := uuid.New()
|
|
envSecretID := uuid.New()
|
|
fileSecretID := uuid.New()
|
|
bothEmptySecretID := 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{userID, "user-secrets-enabled", "user-secrets-enabled@test.com", []byte{}, now, now, "active", pq.StringArray{}, "password"},
|
|
},
|
|
// env-only secret: should remain enabled after migration.
|
|
{
|
|
`INSERT INTO user_secrets (id, user_id, name, description, value, env_name, file_path, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`,
|
|
[]any{envSecretID, userID, "env-secret", "", "v1", "ENV_SECRET", "", now, now},
|
|
},
|
|
// file-only secret: should remain enabled after migration.
|
|
{
|
|
`INSERT INTO user_secrets (id, user_id, name, description, value, env_name, file_path, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`,
|
|
[]any{fileSecretID, userID, "file-secret", "", "v2", "", "/tmp/file-secret", now, now},
|
|
},
|
|
// Both env_name and file_path empty: silently skipped today by
|
|
// the agent manifest layer. Should be flipped to enabled=false
|
|
// by the migration so the behavior is preserved exactly under
|
|
// the new "always inject when enabled" rule.
|
|
{
|
|
`INSERT INTO user_secrets (id, user_id, name, description, value, env_name, file_path, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`,
|
|
[]any{bothEmptySecretID, userID, "both-empty", "", "v3", "", "", 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)
|
|
|
|
getEnabled := func(t *testing.T, id uuid.UUID) bool {
|
|
t.Helper()
|
|
var enabled bool
|
|
err := sqlDB.QueryRowContext(ctx,
|
|
"SELECT enabled FROM user_secrets WHERE id = $1", id,
|
|
).Scan(&enabled)
|
|
require.NoError(t, err)
|
|
return enabled
|
|
}
|
|
|
|
require.True(t, getEnabled(t, envSecretID),
|
|
"env-only secret should remain enabled")
|
|
require.True(t, getEnabled(t, fileSecretID),
|
|
"file-only secret should remain enabled")
|
|
require.False(t, getEnabled(t, bothEmptySecretID),
|
|
"secret with both targets empty should be flipped to disabled "+
|
|
"to preserve the previous implicit-skip behavior")
|
|
}
|
|
|
|
// TestMigration000562OAuth2PublicClientTokensBackfill seeds a pre-migration
|
|
// oauth2_provider_app_tokens row (the only shape that could exist before this
|
|
// migration, since app_secret_id was NOT NULL) and asserts that the new app_id
|
|
// column is backfilled from the existing app_secret_id -> app_id join, and
|
|
// that app_secret_id becomes nullable afterward.
|
|
func TestMigration000562OAuth2PublicClientTokensBackfill(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
const priorMigrationVersion = 561
|
|
|
|
sqlDB := testSQLDB(t)
|
|
|
|
next, err := migrations.Stepper(sqlDB)
|
|
require.NoError(t, err)
|
|
for {
|
|
version, more, err := next()
|
|
require.NoError(t, err)
|
|
if !more || version == priorMigrationVersion {
|
|
break
|
|
}
|
|
}
|
|
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
|
|
userID := uuid.New()
|
|
appID := uuid.New()
|
|
secretID := uuid.New()
|
|
tokenID := uuid.New()
|
|
const apiKeyID = "test562apikeyid"
|
|
|
|
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, 'test-user-562', 'test-562@example.com', ''::bytea, $2, $2, 'active', '{}', 'password')
|
|
`, userID, now)
|
|
require.NoError(t, err)
|
|
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO api_keys (id, hashed_secret, user_id, last_used, expires_at, created_at, updated_at, login_type, scopes, allow_list)
|
|
VALUES ($1, ''::bytea, $2, $3, $3, $3, $3, 'oauth2_provider_app', '{}', '{*}')
|
|
`, apiKeyID, userID, now)
|
|
require.NoError(t, err)
|
|
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_apps (id, created_at, updated_at, name, icon, callback_url)
|
|
VALUES ($1, $2, $2, 'test-app-562', '', 'http://localhost/callback')
|
|
`, appID, now)
|
|
require.NoError(t, err)
|
|
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_app_secrets (id, created_at, hashed_secret, display_secret, app_id, secret_prefix)
|
|
VALUES ($1, $2, ''::bytea, '****1234', $3, 'prefix562'::bytea)
|
|
`, secretID, now, appID)
|
|
require.NoError(t, err)
|
|
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_app_tokens (id, created_at, expires_at, hash_prefix, refresh_hash, app_secret_id, api_key_id, user_id)
|
|
VALUES ($1, $2, $3, 'prefix562'::bytea, ''::bytea, $4, $5, $6)
|
|
`, tokenID, now, now.Add(time.Hour), secretID, apiKeyID, userID)
|
|
require.NoError(t, err)
|
|
|
|
require.NoError(t, tx.Commit())
|
|
|
|
migrationSQL, err := os.ReadFile("000562_oauth2_public_client_tokens.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(migrationSQL))
|
|
require.NoError(t, err)
|
|
|
|
var backfilledAppID uuid.UUID
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
`SELECT app_id FROM oauth2_provider_app_tokens WHERE id = $1`, tokenID,
|
|
).Scan(&backfilledAppID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, appID, backfilledAppID, "app_id should be backfilled from app_secret_id's existing join")
|
|
|
|
var isNullable string
|
|
err = sqlDB.QueryRowContext(ctx, `
|
|
SELECT is_nullable FROM information_schema.columns
|
|
WHERE table_name = 'oauth2_provider_app_tokens' AND column_name = 'app_secret_id'
|
|
`).Scan(&isNullable)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "YES", isNullable, "app_secret_id should be nullable after the migration")
|
|
}
|
|
|
|
// setupMigration000565Apps steps to the migration just before 000565 and seeds
|
|
// one oauth2_provider_apps row per token_endpoint_auth_method shape that 000566
|
|
// repairs or preserves. 000565's NULL client_type case is seeded in its own
|
|
// test, since it is about the column this helper always populates. Returns the
|
|
// app IDs keyed by the shape they were seeded with.
|
|
func setupMigration000565Apps(t *testing.T) (*sql.DB, context.Context, map[string]uuid.UUID) {
|
|
t.Helper()
|
|
|
|
const priorMigrationVersion = 564
|
|
|
|
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", priorMigrationVersion)
|
|
}
|
|
if version == priorMigrationVersion {
|
|
break
|
|
}
|
|
}
|
|
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
|
|
// "legacy" is the row the known bug produced: the client asked for "none",
|
|
// which RFC 7591 section 2 defines as public, but client_type was hardcoded
|
|
// confidential and a secret was issued anyway.
|
|
ids := map[string]uuid.UUID{
|
|
"legacy": uuid.New(), // confidential + none -> auth method must be repaired
|
|
"nullMethod": uuid.New(), // confidential + NULL -> auth method must be repaired
|
|
"publicMismatched": uuid.New(), // public + client_secret_basic -> auth method must be repaired
|
|
"confidentialBasic": uuid.New(), // already consistent -> must be left alone
|
|
"confidentialPost": uuid.New(), // already consistent, and post must not be flattened to basic
|
|
"publicNone": uuid.New(), // already consistent -> must be left alone
|
|
"confidentialEmpty": uuid.New(), // confidential + '' -> repaired only by the widened predicate
|
|
"confidentialJunk": uuid.New(), // confidential + unknown -> repaired only by the widened predicate
|
|
"publicNull": uuid.New(), // public + NULL -> the second UPDATE's IS NULL arm
|
|
}
|
|
|
|
seed := func(id uuid.UUID, name, clientType string, authMethod *string) {
|
|
t.Helper()
|
|
_, err := sqlDB.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_apps
|
|
(id, created_at, updated_at, name, icon, callback_url, client_type, token_endpoint_auth_method)
|
|
VALUES ($1, $2, $2, $3, '', 'http://localhost/callback', $4, $5)
|
|
`, id, now, name, clientType, authMethod)
|
|
require.NoError(t, err)
|
|
}
|
|
seed(ids["legacy"], "test-565-legacy", "confidential", new("none"))
|
|
seed(ids["nullMethod"], "test-565-null", "confidential", nil)
|
|
seed(ids["publicMismatched"], "test-565-public-mismatch", "public", new("client_secret_basic"))
|
|
seed(ids["confidentialBasic"], "test-565-basic", "confidential", new("client_secret_basic"))
|
|
seed(ids["confidentialPost"], "test-565-post", "confidential", new("client_secret_post"))
|
|
seed(ids["publicNone"], "test-565-public", "public", new("none"))
|
|
// Neither of these is producible by today's write path: ApplyDefaults maps
|
|
// "" to client_secret_basic and Valid() rejects unknown methods. They stand
|
|
// in for history the squashed log cannot rule out, and both would satisfy
|
|
// the eventual cross-column constraint while remaining unusable.
|
|
seed(ids["confidentialEmpty"], "test-565-empty", "confidential", new(""))
|
|
seed(ids["confidentialJunk"], "test-565-junk", "confidential", new("client_secret_jwt"))
|
|
seed(ids["publicNull"], "test-565-public-null", "public", nil)
|
|
|
|
return sqlDB, ctx, ids
|
|
}
|
|
|
|
// TestMigration000565OAuth2ClientTypeConstraint covers the constraint migration:
|
|
// a NULL client_type is backfilled, the column becomes NOT NULL, and the CHECK
|
|
// rejects any value other than the two canonical ones.
|
|
func TestMigration000565OAuth2ClientTypeConstraint(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
sqlDB, ctx, _ := setupMigration000565Apps(t)
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
|
|
// A NULL client_type must survive the migration as 'confidential', the
|
|
// fail-closed direction, rather than blocking SET NOT NULL.
|
|
nullTypeID := uuid.New()
|
|
_, err := sqlDB.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_apps
|
|
(id, created_at, updated_at, name, icon, callback_url, client_type)
|
|
VALUES ($1, $2, $2, 'test-565-nulltype', '', 'http://localhost/callback', NULL)
|
|
`, nullTypeID, now)
|
|
require.NoError(t, err)
|
|
|
|
migrationSQL, err := os.ReadFile("000565_oauth2_client_type_constraint.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(migrationSQL))
|
|
require.NoError(t, err)
|
|
|
|
var backfilled string
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
`SELECT client_type FROM oauth2_provider_apps WHERE id = $1`, nullTypeID,
|
|
).Scan(&backfilled)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "confidential", backfilled,
|
|
"a NULL client_type must read as confidential, the direction that keeps requiring a secret")
|
|
|
|
var isNullable string
|
|
err = sqlDB.QueryRowContext(ctx, `
|
|
SELECT is_nullable FROM information_schema.columns
|
|
WHERE table_name = 'oauth2_provider_apps' AND column_name = 'client_type'
|
|
`).Scan(&isNullable)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "NO", isNullable, "client_type should be NOT NULL after the migration")
|
|
|
|
// Prove the CHECK rejects rather than trusting that it exists.
|
|
for _, badValue := range []string{"Public", "PUBLIC", "public ", "bogus", ""} {
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_apps
|
|
(id, created_at, updated_at, name, icon, callback_url, client_type)
|
|
VALUES ($1, $2, $2, $3, '', 'http://localhost/callback', $4)
|
|
`, uuid.New(), now, "test-565-bad-"+badValue, badValue)
|
|
require.Error(t, err, "client_type %q must be rejected by the CHECK constraint", badValue)
|
|
require.True(t, database.IsCheckViolation(err, database.CheckOauth2ProviderAppsClientTypeCheck),
|
|
"client_type %q must fail the check constraint specifically", badValue)
|
|
}
|
|
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_apps
|
|
(id, created_at, updated_at, name, icon, callback_url, client_type)
|
|
VALUES ($1, $2, $2, 'test-565-null-after', '', 'http://localhost/callback', NULL)
|
|
`, uuid.New(), now)
|
|
require.ErrorContains(t, err, "not-null")
|
|
|
|
for _, goodValue := range []string{"confidential", "public"} {
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_apps
|
|
(id, created_at, updated_at, name, icon, callback_url, client_type)
|
|
VALUES ($1, $2, $2, $3, '', 'http://localhost/callback', $4)
|
|
`, uuid.New(), now, "test-565-good-"+goodValue, goodValue)
|
|
require.NoError(t, err, "client_type %q must remain valid", goodValue)
|
|
}
|
|
}
|
|
|
|
func testMigration000590WorkspaceAgentSessionCounts(t *testing.T, sqlDB *sql.DB, _ migrationStepper) {
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
migrationSQL, err := os.ReadFile("000590_workspace_agent_session_counts.up.sql")
|
|
require.NoError(t, err)
|
|
|
|
// Everything the rollup has not consumed converts, no matter how far the
|
|
// watermark lags. Rows the rollup already consumed convert to '{}': the
|
|
// insights queries read them from template_usage_stats instead. wantCounts
|
|
// is ordered oldest row first.
|
|
tests := []struct {
|
|
name string
|
|
statsAgeHours []int
|
|
sessionCount int
|
|
watermarkAgeHours []int
|
|
wantCounts []string
|
|
}{
|
|
{name: "empty database"},
|
|
{name: "no watermark and brief activity", statsAgeHours: []int{0}, sessionCount: 2, wantCounts: []string{`{"vscode": 2}`}},
|
|
{name: "no watermark and idle stats", statsAgeHours: []int{0}, wantCounts: []string{`{}`}},
|
|
{name: "no watermark and a long activity backlog", statsAgeHours: []int{48, 0}, sessionCount: 2, wantCounts: []string{`{"vscode": 2}`, `{"vscode": 2}`}},
|
|
// Without a watermark there is nothing in template_usage_stats to fall
|
|
// back on, so even ancient activity converts rather than roll up as
|
|
// zero minutes later.
|
|
{name: "no watermark and expired activity", statsAgeHours: []int{181 * 24}, sessionCount: 2, wantCounts: []string{`{"vscode": 2}`}},
|
|
// A weekend shutoff records nothing while coderd is down, so the
|
|
// watermark is days old with only the minutes before the shutoff
|
|
// behind it.
|
|
{name: "weekend shutoff", statsAgeHours: []int{63}, sessionCount: 2, watermarkAgeHours: []int{64}, wantCounts: []string{`{"vscode": 2}`}},
|
|
{name: "idle since the last rollup", statsAgeHours: []int{0}, watermarkAgeHours: []int{48}, wantCounts: []string{`{}`}},
|
|
{name: "stalled rollup", statsAgeHours: []int{47, 0}, sessionCount: 2, watermarkAgeHours: []int{48}, wantCounts: []string{`{"vscode": 2}`, `{"vscode": 2}`}},
|
|
{name: "activity already rolled up", statsAgeHours: []int{50}, sessionCount: 2, watermarkAgeHours: []int{25}, wantCounts: []string{`{}`}},
|
|
{name: "fresh watermark", statsAgeHours: []int{0}, sessionCount: 2, watermarkAgeHours: []int{23}, wantCounts: []string{`{"vscode": 2}`}},
|
|
{name: "latest watermark wins", statsAgeHours: []int{0}, sessionCount: 2, watermarkAgeHours: []int{25, 23}, wantCounts: []string{`{"vscode": 2}`}},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
tx, err := sqlDB.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { _ = tx.Rollback() })
|
|
|
|
for _, ageHours := range tt.statsAgeHours {
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO workspace_agent_stats (
|
|
id, created_at, user_id, agent_id, workspace_id, template_id,
|
|
connection_count, session_count_vscode
|
|
) VALUES (
|
|
gen_random_uuid(), statement_timestamp() - $1::bigint * interval '1 hour',
|
|
gen_random_uuid(), gen_random_uuid(), gen_random_uuid(), gen_random_uuid(), 1, $2
|
|
)
|
|
`, ageHours, tt.sessionCount)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
for _, ageHours := range tt.watermarkAgeHours {
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO template_usage_stats (
|
|
start_time, end_time, template_id, user_id, median_latency_ms,
|
|
usage_mins, ssh_mins, sftp_mins, reconnecting_pty_mins,
|
|
vscode_mins, jetbrains_mins, app_usage_mins
|
|
) VALUES (
|
|
statement_timestamp() - $1::bigint * interval '1 hour',
|
|
statement_timestamp() - $1::bigint * interval '1 hour' + interval '30 minutes',
|
|
gen_random_uuid(), gen_random_uuid(), NULL, 0, 0, 0, 0, 0, 0, NULL
|
|
)
|
|
`, ageHours)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
|
require.NoError(t, err)
|
|
|
|
rows, err := tx.QueryContext(ctx, `SELECT session_counts FROM workspace_agent_stats ORDER BY created_at`)
|
|
require.NoError(t, err)
|
|
defer rows.Close()
|
|
var gotCounts []string
|
|
for rows.Next() {
|
|
var sessionCounts []byte
|
|
require.NoError(t, rows.Scan(&sessionCounts))
|
|
gotCounts = append(gotCounts, string(sessionCounts))
|
|
}
|
|
require.NoError(t, rows.Err())
|
|
require.Len(t, gotCounts, len(tt.wantCounts))
|
|
for i, want := range tt.wantCounts {
|
|
require.JSONEq(t, want, gotCounts[i])
|
|
}
|
|
})
|
|
}
|
|
|
|
// The down migration has columns for the four known apps only, so a count
|
|
// under any other name is lost.
|
|
t.Run("down restores known apps only", func(t *testing.T) {
|
|
tx, err := sqlDB.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { _ = tx.Rollback() })
|
|
|
|
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
|
require.NoError(t, err)
|
|
|
|
_, err = tx.ExecContext(ctx, `
|
|
INSERT INTO workspace_agent_stats (
|
|
id, created_at, user_id, agent_id, workspace_id, template_id,
|
|
connection_count, session_counts
|
|
) VALUES (
|
|
gen_random_uuid(), statement_timestamp(), gen_random_uuid(),
|
|
gen_random_uuid(), gen_random_uuid(), gen_random_uuid(), 1,
|
|
'{"vscode": 2, "some_future_ide": 3}'::jsonb
|
|
)
|
|
`)
|
|
require.NoError(t, err)
|
|
|
|
downSQL, err := os.ReadFile("000590_workspace_agent_session_counts.down.sql")
|
|
require.NoError(t, err)
|
|
_, err = tx.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
|
|
var vscode, jetbrains, reconnectingPTY, ssh int64
|
|
err = tx.QueryRowContext(ctx, `
|
|
SELECT session_count_vscode, session_count_jetbrains,
|
|
session_count_reconnecting_pty, session_count_ssh
|
|
FROM workspace_agent_stats
|
|
`).Scan(&vscode, &jetbrains, &reconnectingPTY, &ssh)
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, 2, vscode)
|
|
require.EqualValues(t, 0, jetbrains)
|
|
require.EqualValues(t, 0, reconnectingPTY)
|
|
require.EqualValues(t, 0, ssh)
|
|
})
|
|
}
|
|
|
|
// TestMigration000566OAuth2AuthMethodBackfill covers the repair the backfill
|
|
// exists for, which the testdata/fixtures run does not reach: its only
|
|
// oauth2_provider_apps row seeds no token_endpoint_auth_method at all, so the
|
|
// '= none' branch, the one that fixes the actual bug, matches zero rows in CI.
|
|
func TestMigration000566OAuth2AuthMethodBackfill(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
sqlDB, ctx, ids := setupMigration000565Apps(t)
|
|
|
|
// 000566 runs after 000565, so apply both in order.
|
|
for _, name := range []string{
|
|
"000565_oauth2_client_type_constraint.up.sql",
|
|
"000566_oauth2_auth_method_backfill.up.sql",
|
|
} {
|
|
migrationSQL, err := os.ReadFile(name)
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(migrationSQL))
|
|
require.NoError(t, err, "applying %s", name)
|
|
}
|
|
|
|
tests := []struct {
|
|
key string
|
|
wantMethod string
|
|
reason string
|
|
}{
|
|
{
|
|
key: "legacy",
|
|
wantMethod: "client_secret_basic",
|
|
reason: "a confidential client declaring \"none\" is the bug this migration repairs",
|
|
},
|
|
{
|
|
key: "nullMethod",
|
|
wantMethod: "client_secret_basic",
|
|
reason: "a confidential client with no declaration gets the RFC 7591 default",
|
|
},
|
|
{
|
|
key: "publicMismatched",
|
|
wantMethod: "none",
|
|
reason: "a public client cannot declare a secret-based method",
|
|
},
|
|
{
|
|
key: "confidentialBasic",
|
|
wantMethod: "client_secret_basic",
|
|
reason: "already consistent, must be left alone",
|
|
},
|
|
{
|
|
key: "confidentialPost",
|
|
wantMethod: "client_secret_post",
|
|
reason: "client_secret_post is consistent with confidential and must not be flattened to basic",
|
|
},
|
|
{
|
|
key: "publicNone",
|
|
wantMethod: "none",
|
|
reason: "already consistent, must be left alone",
|
|
},
|
|
{
|
|
key: "confidentialEmpty",
|
|
wantMethod: "client_secret_basic",
|
|
reason: "an empty declaration is not a valid confidential method and must be repaired, not just the known 'none' case",
|
|
},
|
|
{
|
|
key: "confidentialJunk",
|
|
wantMethod: "client_secret_basic",
|
|
reason: "an unrecognized declaration must be repaired too, so the migration is idempotent against history it cannot inspect",
|
|
},
|
|
{
|
|
key: "publicNull",
|
|
wantMethod: "none",
|
|
reason: "exercises the second UPDATE's IS NULL arm, which no other seeded row reaches",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.key, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
// A parallel subtest's own context, not the parent's. The parent's
|
|
// deadline starts when it is created, but a parallel subtest does
|
|
// not run until a -parallel slot frees, so it can spend most of the
|
|
// budget queued behind other tests in the package and then fail with
|
|
// "context deadline exceeded" on a sub-20ms read. That failure names
|
|
// no code and gets worse as the package grows.
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
|
|
var gotMethod, gotClientType string
|
|
err := sqlDB.QueryRowContext(ctx, `
|
|
SELECT token_endpoint_auth_method, client_type
|
|
FROM oauth2_provider_apps WHERE id = $1
|
|
`, ids[tt.key]).Scan(&gotMethod, &gotClientType)
|
|
require.NoError(t, err)
|
|
require.Equal(t, tt.wantMethod, gotMethod, tt.reason)
|
|
})
|
|
}
|
|
|
|
// IS DISTINCT FROM, not <>: with <> a NULL declaration makes the comparison
|
|
// NULL and WHERE drops the row, so an unrepaired NULL would count as
|
|
// consistent. The same trap awaits the permanent cross-column CHECK.
|
|
var contradictions int
|
|
err := sqlDB.QueryRowContext(ctx, `
|
|
SELECT count(*) FROM oauth2_provider_apps
|
|
WHERE (token_endpoint_auth_method = 'none') IS DISTINCT FROM (client_type = 'public')
|
|
`).Scan(&contradictions)
|
|
require.NoError(t, err)
|
|
require.Zero(t, contradictions,
|
|
"after the backfill no row may declare an auth method that contradicts its enforced client type")
|
|
|
|
var stillConfidential string
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
`SELECT client_type FROM oauth2_provider_apps WHERE id = $1`, ids["legacy"],
|
|
).Scan(&stillConfidential)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "confidential", stillConfidential,
|
|
"the backfill aligns the declaration to what is enforced, so the enforced value must be unchanged")
|
|
}
|
|
|
|
func TestMigration000580ChatModelConfigOrganization(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
const migrationVersion = 580
|
|
const previousMigrationVersion = 579
|
|
|
|
sqlDB := testSQLDB(t)
|
|
|
|
// Migrate up to the migration before the org-scoping 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", previousMigrationVersion)
|
|
}
|
|
if version == previousMigrationVersion {
|
|
break
|
|
}
|
|
}
|
|
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
|
|
providerID := uuid.New()
|
|
defaultConfigID := uuid.New()
|
|
plainConfigID := uuid.New()
|
|
deletedConfigID := uuid.New()
|
|
historyOrgID := uuid.New()
|
|
chatID := uuid.New()
|
|
messageID := int64(578001)
|
|
queuedMessageID := int64(578002)
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
|
|
fixtures := []struct {
|
|
query string
|
|
args []any
|
|
}{
|
|
{
|
|
`INSERT INTO ai_providers (id, type, name, enabled, base_url, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7)`,
|
|
[]any{providerID, "openai", "openai-578", true, "https://api.openai.com/v1", now, now},
|
|
},
|
|
// The deployment's single live default config.
|
|
{
|
|
`INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, context_limit, compression_threshold, ai_provider_id, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`,
|
|
[]any{defaultConfigID, "gpt-5.2", "Default 578", true, true, 200000, 70, providerID, now, now},
|
|
},
|
|
// A live non-default config.
|
|
{
|
|
`INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, context_limit, compression_threshold, ai_provider_id, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`,
|
|
[]any{plainConfigID, "gpt-5.2-mini", "Plain 578", true, false, 128000, 70, providerID, now, now},
|
|
},
|
|
// A soft-deleted config is backfilled like any row.
|
|
{
|
|
`INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, deleted, deleted_at, context_limit, compression_threshold, ai_provider_id, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)`,
|
|
[]any{deletedConfigID, "gpt-4-legacy", "Deleted 578", false, false, true, now, 128000, 70, providerID, now, now},
|
|
},
|
|
{
|
|
`INSERT INTO organizations (id, name, description, display_name, default_org_member_roles, created_at, updated_at)
|
|
VALUES ($1, $2, '', '', $3, $4, $4)`,
|
|
[]any{historyOrgID, "history-org-578", pq.StringArray{}, now},
|
|
},
|
|
// A non-default organization chat keeps the pre-migration model ID.
|
|
{
|
|
`INSERT INTO chats (id, owner_id, last_model_config_id, title, organization_id, created_at, updated_at)
|
|
SELECT $1, u.id, $2, $3, $4, $5, $5
|
|
FROM users u
|
|
ORDER BY u.created_at, u.id
|
|
LIMIT 1`,
|
|
[]any{chatID, plainConfigID, "History 578", historyOrgID, now},
|
|
},
|
|
{
|
|
`INSERT INTO chat_messages (id, chat_id, model_config_id, role, content, content_version, created_by, created_at)
|
|
SELECT $1, $2, $3, 'user', '{"type":"text","text":"history"}'::jsonb, 1, owner_id, $4
|
|
FROM chats WHERE id = $2`,
|
|
[]any{messageID, chatID, plainConfigID, now},
|
|
},
|
|
{
|
|
`INSERT INTO chat_queued_messages (id, chat_id, model_config_id, content, created_by, created_at)
|
|
SELECT $1, $2, $3, '{"type":"text","text":"queued"}'::jsonb, owner_id, $4
|
|
FROM chats WHERE id = $2`,
|
|
[]any{queuedMessageID, chatID, plainConfigID, now},
|
|
},
|
|
}
|
|
|
|
tx, err := sqlDB.BeginTx(ctx, nil)
|
|
require.NoError(t, err)
|
|
defer tx.Rollback()
|
|
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)
|
|
|
|
var defaultOrgID uuid.UUID
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT id FROM organizations WHERE is_default = true",
|
|
).Scan(&defaultOrgID)
|
|
require.NoError(t, err)
|
|
|
|
// Every row, including the soft-deleted one, is backfilled to the
|
|
// default org.
|
|
var backfilled int
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT COUNT(*) FROM chat_model_configs WHERE organization_id = $1", defaultOrgID,
|
|
).Scan(&backfilled)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 3, backfilled, "all existing configs should move only to the default org")
|
|
|
|
var totalConfigs int
|
|
err = sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM chat_model_configs").Scan(&totalConfigs)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 3, totalConfigs, "the migration must not copy configs to other organizations")
|
|
|
|
var notNullViolation bool
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT EXISTS(SELECT 1 FROM chat_model_configs WHERE organization_id IS NULL)",
|
|
).Scan(¬NullViolation)
|
|
require.NoError(t, err)
|
|
require.False(t, notNullViolation, "no row may keep a NULL organization_id")
|
|
|
|
var chatModelConfigID, messageModelConfigID, queuedModelConfigID uuid.UUID
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT last_model_config_id FROM chats WHERE id = $1", chatID,
|
|
).Scan(&chatModelConfigID)
|
|
require.NoError(t, err)
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT model_config_id FROM chat_messages WHERE id = $1", messageID,
|
|
).Scan(&messageModelConfigID)
|
|
require.NoError(t, err)
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT model_config_id FROM chat_queued_messages WHERE id = $1", queuedMessageID,
|
|
).Scan(&queuedModelConfigID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, plainConfigID, chatModelConfigID)
|
|
require.Equal(t, plainConfigID, messageModelConfigID)
|
|
require.Equal(t, plainConfigID, queuedModelConfigID)
|
|
|
|
// Every row's group_acl is seeded with the everyone-in-org read
|
|
// entry keyed by the org ID (the Everyone group shares the org ID).
|
|
seededACL := map[string]any{
|
|
defaultOrgID.String(): map[string]any{"permissions": []string{"read"}},
|
|
}
|
|
for _, id := range []uuid.UUID{defaultConfigID, plainConfigID, deletedConfigID} {
|
|
var groupACL []byte
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT group_acl FROM chat_model_configs WHERE id = $1", id,
|
|
).Scan(&groupACL)
|
|
require.NoError(t, err)
|
|
require.JSONEq(t, string(mustJSON(t, seededACL)), string(groupACL),
|
|
"group_acl should carry the everyone-in-org read entry")
|
|
}
|
|
|
|
// The single-default index is now keyed per organization: the index
|
|
// definition references organization_id and keeps its predicate.
|
|
var indexDef string
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT indexdef FROM pg_indexes WHERE indexname = 'idx_chat_model_configs_single_default'",
|
|
).Scan(&indexDef)
|
|
require.NoError(t, err)
|
|
require.Contains(t, indexDef, "organization_id")
|
|
require.Contains(t, indexDef, "is_default = true")
|
|
require.Contains(t, indexDef, "deleted = false")
|
|
|
|
// The org lookup index exists.
|
|
var orgIndexExists bool
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT EXISTS(SELECT 1 FROM pg_indexes WHERE indexname = 'idx_chat_model_configs_organization_id')",
|
|
).Scan(&orgIndexExists)
|
|
require.NoError(t, err)
|
|
require.True(t, orgIndexExists)
|
|
|
|
// The default org can host only one live default: a second insert
|
|
// violates the per-org partial unique index.
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
`INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, context_limit, compression_threshold, ai_provider_id, organization_id, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`,
|
|
uuid.New(), "gpt-5.2-alt", "Second Default 578", true, true, 200000, 70, providerID, defaultOrgID, now, now,
|
|
)
|
|
require.Error(t, err, "a second live default in the same org must be rejected")
|
|
require.Contains(t, err.Error(), "idx_chat_model_configs_single_default")
|
|
|
|
// A second org can host its own live default.
|
|
secondOrgID := uuid.New()
|
|
secondOrgDefaultConfigID := uuid.New()
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
`INSERT INTO organizations (id, name, description, display_name, default_org_member_roles, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7)`,
|
|
secondOrgID, "second-org-578", "", "", pq.StringArray{}, now, now,
|
|
)
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
`INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, context_limit, compression_threshold, ai_provider_id, organization_id, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`,
|
|
secondOrgDefaultConfigID, "gpt-5.2-org2", "Second Org Default 578", true, true, 200000, 70, providerID, secondOrgID, now, now,
|
|
)
|
|
require.NoError(t, err, "each org can host its own live default")
|
|
|
|
// The ACL object CHECKs reject non-object values.
|
|
_, err = sqlDB.ExecContext(ctx,
|
|
`INSERT INTO chat_model_configs (id, model, display_name, enabled, context_limit, compression_threshold, ai_provider_id, organization_id, group_acl, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, '[]'::jsonb, $9, $10)`,
|
|
uuid.New(), "bad-acl", "Bad ACL 578", true, 128000, 70, providerID, secondOrgID, now, now,
|
|
)
|
|
require.Error(t, err, "non-object group_acl must be rejected")
|
|
|
|
var rowCountBeforeDown int
|
|
err = sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM chat_model_configs").Scan(&rowCountBeforeDown)
|
|
require.NoError(t, err)
|
|
|
|
downSQL, err := os.ReadFile("000580_chat_model_config_organization.down.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
|
|
var defaultOrgIsDefault bool
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT is_default FROM chat_model_configs WHERE id = $1", defaultConfigID,
|
|
).Scan(&defaultOrgIsDefault)
|
|
require.NoError(t, err)
|
|
require.True(t, defaultOrgIsDefault)
|
|
|
|
var secondOrgIsDefault bool
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT is_default FROM chat_model_configs WHERE id = $1", secondOrgDefaultConfigID,
|
|
).Scan(&secondOrgIsDefault)
|
|
require.NoError(t, err)
|
|
require.False(t, secondOrgIsDefault)
|
|
|
|
var rowCountAfterDown int
|
|
err = sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM chat_model_configs").Scan(&rowCountAfterDown)
|
|
require.NoError(t, err)
|
|
require.Equal(t, rowCountBeforeDown, rowCountAfterDown)
|
|
|
|
err = sqlDB.QueryRowContext(ctx,
|
|
"SELECT indexdef FROM pg_indexes WHERE indexname = 'idx_chat_model_configs_single_default'",
|
|
).Scan(&indexDef)
|
|
require.NoError(t, err)
|
|
require.Contains(t, indexDef, "((1))")
|
|
|
|
upSQL, err := os.ReadFile("000580_chat_model_config_organization.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
func mustJSON(t *testing.T, v any) []byte {
|
|
t.Helper()
|
|
raw, err := json.Marshal(v)
|
|
require.NoError(t, err)
|
|
return raw
|
|
}
|
|
|
|
func testMigration000583ChatModelOverrideOrgScope(t *testing.T, db *sql.DB) {
|
|
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
|
upSQL, err := os.ReadFile("000583_chat_model_override_org_scope.up.sql")
|
|
require.NoError(t, err)
|
|
downSQL, err := os.ReadFile("000583_chat_model_override_org_scope.down.sql")
|
|
require.NoError(t, err)
|
|
|
|
_, err = db.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
|
|
now := time.Now().UTC().Truncate(time.Microsecond)
|
|
userID := uuid.New()
|
|
providerID := uuid.New()
|
|
modelID := uuid.New()
|
|
|
|
var orgID uuid.UUID
|
|
require.NoError(t, db.QueryRowContext(ctx,
|
|
"SELECT id FROM organizations WHERE is_default = true").Scan(&orgID))
|
|
|
|
_, err = db.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, $5, 'active', '{}', 'password')`,
|
|
userID, "model-override-"+userID.String(), userID.String()+"@example.com", []byte{}, now)
|
|
require.NoError(t, err)
|
|
_, err = db.ExecContext(ctx, `
|
|
INSERT INTO ai_providers (id, type, name, enabled, base_url, created_at, updated_at)
|
|
VALUES ($1, 'openai', $2, true, 'https://example.com', $3, $3)`,
|
|
providerID, "model-override-"+providerID.String(), now)
|
|
require.NoError(t, err)
|
|
_, err = db.ExecContext(ctx, `
|
|
INSERT INTO chat_model_configs (
|
|
id, model, display_name, enabled, is_default, deleted,
|
|
context_limit, compression_threshold, ai_provider_id, organization_id,
|
|
group_acl, user_acl, created_at, updated_at
|
|
) VALUES ($1, $2, $2, true, false, false, 128000, 70, $3, $4, '{}', '{}', $5, $5)`,
|
|
modelID, "model-override-"+modelID.String(), providerID, orgID, now)
|
|
require.NoError(t, err)
|
|
|
|
advisorConfig := fmt.Sprintf(
|
|
`{"enabled":true,"max_uses_per_run":2,"model_config_id":%q,"reasoning_effort":"low"}`, modelID)
|
|
_, err = db.ExecContext(ctx, `
|
|
INSERT INTO site_configs (key, value) VALUES
|
|
('agents_chat_general_model_override', $1),
|
|
('agents_chat_title_generation_model_override', 'not-a-uuid'),
|
|
('agents_advisor_config', $2)
|
|
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value`,
|
|
modelID.String()+":high", advisorConfig)
|
|
require.NoError(t, err)
|
|
_, err = db.ExecContext(ctx, `
|
|
INSERT INTO user_configs (user_id, key, value) VALUES
|
|
($1, 'chat_personal_model_override:root', $2),
|
|
($1, 'chat_personal_model_override:general', 'chat_default')`,
|
|
userID, "model:"+modelID.String()+":max")
|
|
require.NoError(t, err)
|
|
|
|
_, err = db.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
|
|
// Legacy overrides are dropped rather than migrated: the new tables start
|
|
// empty and the serialized keys are gone.
|
|
var count int
|
|
require.NoError(t, db.QueryRowContext(ctx,
|
|
"SELECT COUNT(*) FROM chat_organization_model_overrides").Scan(&count))
|
|
require.Zero(t, count)
|
|
require.NoError(t, db.QueryRowContext(ctx,
|
|
"SELECT COUNT(*) FROM chat_user_model_overrides").Scan(&count))
|
|
require.Zero(t, count)
|
|
require.NoError(t, db.QueryRowContext(ctx, `
|
|
SELECT COUNT(*) FROM site_configs
|
|
WHERE key IN (
|
|
'agents_chat_general_model_override',
|
|
'agents_chat_explore_model_override',
|
|
'agents_chat_title_generation_model_override',
|
|
'agents_chat_compaction_model_override'
|
|
)`).Scan(&count))
|
|
require.Zero(t, count)
|
|
require.NoError(t, db.QueryRowContext(ctx, `
|
|
SELECT COUNT(*) FROM user_configs
|
|
WHERE user_id = $1 AND key LIKE 'chat\_personal\_model\_override:%'`, userID).Scan(&count))
|
|
require.Zero(t, count)
|
|
|
|
// The migration does not parse the advisor runtime config; stale model
|
|
// fields stay in the stored JSON and are ignored by readers.
|
|
var value string
|
|
require.NoError(t, db.QueryRowContext(ctx,
|
|
"SELECT value FROM site_configs WHERE key = 'agents_advisor_config'").Scan(&value))
|
|
require.Equal(t, advisorConfig, value)
|
|
|
|
// The composite foreign key rejects model references from another
|
|
// organization.
|
|
otherOrgID := uuid.New()
|
|
_, err = db.ExecContext(ctx, `
|
|
INSERT INTO organizations (
|
|
id, name, display_name, description, created_at, updated_at,
|
|
is_default, default_org_member_roles
|
|
) VALUES ($1, $2, $2, '', $3, $3, false, '{}')`,
|
|
otherOrgID, "model-override-"+otherOrgID.String(), now)
|
|
require.NoError(t, err)
|
|
_, err = db.ExecContext(ctx, `
|
|
INSERT INTO chat_organization_model_overrides (organization_id, context, model_config_id)
|
|
VALUES ($1, 'explore', $2)`, otherOrgID, modelID)
|
|
require.Error(t, err)
|
|
require.True(t, database.IsForeignKeyViolation(err))
|
|
|
|
_, err = db.ExecContext(ctx, string(downSQL))
|
|
require.NoError(t, err)
|
|
_, err = db.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Migration 000599 makes redirect_uris the source of truth for an OAuth2
|
|
// app's redirect URIs, with the callback as the first entry. Every row shape
|
|
// the column could hold before the migration must come out with the callback
|
|
// first and nothing lost, and running the statement again must change
|
|
// nothing.
|
|
func TestMigration000599OAuth2RedirectURIsPrimary(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
const migrationVersion = 599
|
|
|
|
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)
|
|
|
|
const (
|
|
callback = "https://app.example.com/callback"
|
|
other = "https://other.example.com/callback"
|
|
third = "https://third.example.com/callback"
|
|
)
|
|
rows := []struct {
|
|
name string
|
|
id uuid.UUID
|
|
before pq.StringArray
|
|
after pq.StringArray
|
|
}{
|
|
{name: "null-list", id: uuid.New(), before: nil, after: pq.StringArray{callback}},
|
|
{name: "empty-list", id: uuid.New(), before: pq.StringArray{}, after: pq.StringArray{callback}},
|
|
{name: "callback-first", id: uuid.New(), before: pq.StringArray{callback, other}, after: pq.StringArray{callback, other}},
|
|
{name: "callback-absent", id: uuid.New(), before: pq.StringArray{other, third}, after: pq.StringArray{callback, other, third}},
|
|
{name: "callback-second", id: uuid.New(), before: pq.StringArray{other, callback}, after: pq.StringArray{callback, other}},
|
|
}
|
|
for _, row := range rows {
|
|
_, err = sqlDB.ExecContext(ctx, `
|
|
INSERT INTO oauth2_provider_apps (id, created_at, updated_at, name, icon, callback_url, redirect_uris)
|
|
VALUES ($1, $2, $2, $3, '', $4, $5)
|
|
`, row.id, now, row.name, callback, row.before)
|
|
require.NoError(t, err, row.name)
|
|
}
|
|
|
|
version, more, err := next()
|
|
require.NoError(t, err)
|
|
require.True(t, more)
|
|
require.EqualValues(t, migrationVersion, version)
|
|
|
|
assertRows := func() {
|
|
t.Helper()
|
|
for _, row := range rows {
|
|
var got pq.StringArray
|
|
var callbackURL string
|
|
err := sqlDB.QueryRowContext(ctx, `
|
|
SELECT redirect_uris, callback_url FROM oauth2_provider_apps WHERE id = $1
|
|
`, row.id).Scan(&got, &callbackURL)
|
|
require.NoError(t, err, row.name)
|
|
require.Equal(t, row.after, got, row.name)
|
|
require.Equal(t, callback, callbackURL, row.name)
|
|
}
|
|
}
|
|
assertRows()
|
|
|
|
// Running the statement again must be a no-op.
|
|
upSQL, err := os.ReadFile("000599_oauth2_redirect_uris_primary.up.sql")
|
|
require.NoError(t, err)
|
|
_, err = sqlDB.ExecContext(ctx, string(upSQL))
|
|
require.NoError(t, err)
|
|
assertRows()
|
|
}
|