From 6e255c72c6cb5e2e5c7e0c66833cd1b8b4a6de0c Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Wed, 28 May 2025 14:21:17 +0100 Subject: [PATCH] chore(coderd/database): enforce agent name unique within workspace build (#18052) Adds a database trigger that runs on insert and update of the `workspace_agents` table. The trigger ensures that the agent name is unique within the context of the workspace build it is being inserted into. --- coderd/database/dump.sql | 43 ++++ ...rkspace_agent_name_unique_trigger.down.sql | 2 + ...workspace_agent_name_unique_trigger.up.sql | 45 ++++ coderd/database/querier_test.go | 234 ++++++++++++++++++ coderd/insights_test.go | 8 +- 5 files changed, 328 insertions(+), 4 deletions(-) create mode 100644 coderd/database/migrations/000332_workspace_agent_name_unique_trigger.down.sql create mode 100644 coderd/database/migrations/000332_workspace_agent_name_unique_trigger.up.sql diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index a50abf13c5..acb9780b82 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -316,6 +316,43 @@ CREATE TYPE workspace_transition AS ENUM ( 'delete' ); +CREATE FUNCTION check_workspace_agent_name_unique() RETURNS trigger + LANGUAGE plpgsql + AS $$ +DECLARE + workspace_build_id uuid; + agents_with_name int; +BEGIN + -- Find the workspace build the workspace agent is being inserted into. + SELECT workspace_builds.id INTO workspace_build_id + FROM workspace_resources + JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id + WHERE workspace_resources.id = NEW.resource_id; + + -- If the agent doesn't have a workspace build, we'll allow the insert. + IF workspace_build_id IS NULL THEN + RETURN NEW; + END IF; + + -- Count how many agents in this workspace build already have the given agent name. + SELECT COUNT(*) INTO agents_with_name + FROM workspace_agents + JOIN workspace_resources ON workspace_resources.id = workspace_agents.resource_id + JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id + WHERE workspace_builds.id = workspace_build_id + AND workspace_agents.name = NEW.name + AND workspace_agents.id != NEW.id; + + -- If there's already an agent with this name, raise an error + IF agents_with_name > 0 THEN + RAISE EXCEPTION 'workspace agent name "%" already exists in this workspace build', NEW.name + USING ERRCODE = 'unique_violation'; + END IF; + + RETURN NEW; +END; +$$; + CREATE FUNCTION compute_notification_message_dedupe_hash() RETURNS trigger LANGUAGE plpgsql AS $$ @@ -2773,6 +2810,12 @@ CREATE TRIGGER update_notification_message_dedupe_hash BEFORE INSERT OR UPDATE O CREATE TRIGGER user_status_change_trigger AFTER INSERT OR UPDATE ON users FOR EACH ROW EXECUTE FUNCTION record_user_status_change(); +CREATE TRIGGER workspace_agent_name_unique_trigger BEFORE INSERT OR UPDATE OF name, resource_id ON workspace_agents FOR EACH ROW EXECUTE FUNCTION check_workspace_agent_name_unique(); + +COMMENT ON TRIGGER workspace_agent_name_unique_trigger ON workspace_agents IS 'Use a trigger instead of a unique constraint because existing data may violate +the uniqueness requirement. A trigger allows us to enforce uniqueness going +forward without requiring a migration to clean up historical data.'; + ALTER TABLE ONLY api_keys ADD CONSTRAINT api_keys_user_id_uuid_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; diff --git a/coderd/database/migrations/000332_workspace_agent_name_unique_trigger.down.sql b/coderd/database/migrations/000332_workspace_agent_name_unique_trigger.down.sql new file mode 100644 index 0000000000..916e1d469e --- /dev/null +++ b/coderd/database/migrations/000332_workspace_agent_name_unique_trigger.down.sql @@ -0,0 +1,2 @@ +DROP TRIGGER IF EXISTS workspace_agent_name_unique_trigger ON workspace_agents; +DROP FUNCTION IF EXISTS check_workspace_agent_name_unique(); diff --git a/coderd/database/migrations/000332_workspace_agent_name_unique_trigger.up.sql b/coderd/database/migrations/000332_workspace_agent_name_unique_trigger.up.sql new file mode 100644 index 0000000000..7b10fcdc1d --- /dev/null +++ b/coderd/database/migrations/000332_workspace_agent_name_unique_trigger.up.sql @@ -0,0 +1,45 @@ +CREATE OR REPLACE FUNCTION check_workspace_agent_name_unique() +RETURNS TRIGGER AS $$ +DECLARE + workspace_build_id uuid; + agents_with_name int; +BEGIN + -- Find the workspace build the workspace agent is being inserted into. + SELECT workspace_builds.id INTO workspace_build_id + FROM workspace_resources + JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id + WHERE workspace_resources.id = NEW.resource_id; + + -- If the agent doesn't have a workspace build, we'll allow the insert. + IF workspace_build_id IS NULL THEN + RETURN NEW; + END IF; + + -- Count how many agents in this workspace build already have the given agent name. + SELECT COUNT(*) INTO agents_with_name + FROM workspace_agents + JOIN workspace_resources ON workspace_resources.id = workspace_agents.resource_id + JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id + WHERE workspace_builds.id = workspace_build_id + AND workspace_agents.name = NEW.name + AND workspace_agents.id != NEW.id; + + -- If there's already an agent with this name, raise an error + IF agents_with_name > 0 THEN + RAISE EXCEPTION 'workspace agent name "%" already exists in this workspace build', NEW.name + USING ERRCODE = 'unique_violation'; + END IF; + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER workspace_agent_name_unique_trigger + BEFORE INSERT OR UPDATE OF name, resource_id ON workspace_agents + FOR EACH ROW + EXECUTE FUNCTION check_workspace_agent_name_unique(); + +COMMENT ON TRIGGER workspace_agent_name_unique_trigger ON workspace_agents IS +'Use a trigger instead of a unique constraint because existing data may violate +the uniqueness requirement. A trigger allows us to enforce uniqueness going +forward without requiring a migration to clean up historical data.'; diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index b5ed8b019c..6f7ee8fe88 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -4,12 +4,14 @@ import ( "context" "database/sql" "encoding/json" + "errors" "fmt" "sort" "testing" "time" "github.com/google/uuid" + "github.com/lib/pq" "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -4720,6 +4722,238 @@ func TestGetPresetsAtFailureLimit(t *testing.T) { }) } +func TestWorkspaceAgentNameUniqueTrigger(t *testing.T) { + t.Parallel() + + if !dbtestutil.WillUsePostgres() { + t.Skip("This test makes use of a database trigger not implemented in dbmem") + } + + createWorkspaceWithAgent := func(t *testing.T, db database.Store, org database.Organization, agentName string) (database.WorkspaceBuild, database.WorkspaceResource, database.WorkspaceAgent) { + t.Helper() + + user := dbgen.User(t, db, database.User{}) + template := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + TemplateID: uuid.NullUUID{Valid: true, UUID: template.ID}, + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OrganizationID: org.ID, + TemplateID: template.ID, + OwnerID: user.ID, + }) + job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + OrganizationID: org.ID, + }) + build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + BuildNumber: 1, + JobID: job.ID, + WorkspaceID: workspace.ID, + TemplateVersionID: templateVersion.ID, + }) + resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: build.JobID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: resource.ID, + Name: agentName, + }) + + return build, resource, agent + } + + t.Run("DuplicateNamesInSameWorkspaceResource", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + org := dbgen.Organization(t, db, database.Organization{}) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: A workspace with an agent + _, resource, _ := createWorkspaceWithAgent(t, db, org, "duplicate-agent") + + // When: Another agent is created for that workspace with the same name. + _, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{ + ID: uuid.New(), + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + Name: "duplicate-agent", // Same name as agent1 + ResourceID: resource.ID, + AuthToken: uuid.New(), + Architecture: "amd64", + OperatingSystem: "linux", + APIKeyScope: database.AgentKeyScopeEnumAll, + }) + + // Then: We expect it to fail. + require.Error(t, err) + var pqErr *pq.Error + require.True(t, errors.As(err, &pqErr)) + require.Equal(t, pq.ErrorCode("23505"), pqErr.Code) // unique_violation + require.Contains(t, pqErr.Message, `workspace agent name "duplicate-agent" already exists in this workspace build`) + }) + + t.Run("DuplicateNamesInSameProvisionerJob", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + org := dbgen.Organization(t, db, database.Organization{}) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: A workspace with an agent + _, resource, agent := createWorkspaceWithAgent(t, db, org, "duplicate-agent") + + // When: A child agent is created for that workspace with the same name. + _, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{ + ID: uuid.New(), + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + Name: agent.Name, + ResourceID: resource.ID, + AuthToken: uuid.New(), + Architecture: "amd64", + OperatingSystem: "linux", + APIKeyScope: database.AgentKeyScopeEnumAll, + }) + + // Then: We expect it to fail. + require.Error(t, err) + var pqErr *pq.Error + require.True(t, errors.As(err, &pqErr)) + require.Equal(t, pq.ErrorCode("23505"), pqErr.Code) // unique_violation + require.Contains(t, pqErr.Message, `workspace agent name "duplicate-agent" already exists in this workspace build`) + }) + + t.Run("DuplicateChildNamesOverMultipleResources", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + org := dbgen.Organization(t, db, database.Organization{}) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: A workspace with two agents + _, resource1, agent1 := createWorkspaceWithAgent(t, db, org, "parent-agent-1") + + resource2 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{JobID: resource1.JobID}) + agent2 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: resource2.ID, + Name: "parent-agent-2", + }) + + // Given: One agent has a child agent + agent1Child := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ParentID: uuid.NullUUID{Valid: true, UUID: agent1.ID}, + Name: "child-agent", + ResourceID: resource1.ID, + }) + + // When: A child agent is inserted for the other parent. + _, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{ + ID: uuid.New(), + ParentID: uuid.NullUUID{Valid: true, UUID: agent2.ID}, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + Name: agent1Child.Name, + ResourceID: resource2.ID, + AuthToken: uuid.New(), + Architecture: "amd64", + OperatingSystem: "linux", + APIKeyScope: database.AgentKeyScopeEnumAll, + }) + + // Then: We expect it to fail. + require.Error(t, err) + var pqErr *pq.Error + require.True(t, errors.As(err, &pqErr)) + require.Equal(t, pq.ErrorCode("23505"), pqErr.Code) // unique_violation + require.Contains(t, pqErr.Message, `workspace agent name "child-agent" already exists in this workspace build`) + }) + + t.Run("SameNamesInDifferentWorkspaces", func(t *testing.T) { + t.Parallel() + + agentName := "same-name-different-workspace" + + db, _ := dbtestutil.NewDB(t) + org := dbgen.Organization(t, db, database.Organization{}) + + // Given: A workspace with an agent + _, _, agent1 := createWorkspaceWithAgent(t, db, org, agentName) + require.Equal(t, agentName, agent1.Name) + + // When: A second workspace is created with an agent having the same name + _, _, agent2 := createWorkspaceWithAgent(t, db, org, agentName) + require.Equal(t, agentName, agent2.Name) + + // Then: We expect there to be different agents with the same name. + require.NotEqual(t, agent1.ID, agent2.ID) + require.Equal(t, agent1.Name, agent2.Name) + }) + + t.Run("NullWorkspaceID", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + org := dbgen.Organization(t, db, database.Organization{}) + ctx := testutil.Context(t, testutil.WaitShort) + + // Given: A resource that does not belong to a workspace build (simulating template import) + orphanJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + Type: database.ProvisionerJobTypeTemplateVersionImport, + OrganizationID: org.ID, + }) + orphanResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: orphanJob.ID, + }) + + // And this resource has a workspace agent. + agent1, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{ + ID: uuid.New(), + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + Name: "orphan-agent", + ResourceID: orphanResource.ID, + AuthToken: uuid.New(), + Architecture: "amd64", + OperatingSystem: "linux", + APIKeyScope: database.AgentKeyScopeEnumAll, + }) + require.NoError(t, err) + require.Equal(t, "orphan-agent", agent1.Name) + + // When: We created another resource that does not belong to a workspace build. + orphanJob2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + Type: database.ProvisionerJobTypeTemplateVersionImport, + OrganizationID: org.ID, + }) + orphanResource2 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: orphanJob2.ID, + }) + + // Then: We expect to be able to create an agent in this new resource that has the same name. + agent2, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{ + ID: uuid.New(), + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + Name: "orphan-agent", // Same name as agent1 + ResourceID: orphanResource2.ID, + AuthToken: uuid.New(), + Architecture: "amd64", + OperatingSystem: "linux", + APIKeyScope: database.AgentKeyScopeEnumAll, + }) + require.NoError(t, err) + require.Equal(t, "orphan-agent", agent2.Name) + require.NotEqual(t, agent1.ID, agent2.ID) + }) +} + func requireUsersMatch(t testing.TB, expected []database.User, found []database.GetUsersRow, msg string) { t.Helper() require.ElementsMatch(t, expected, database.ConvertUserRows(found), msg) diff --git a/coderd/insights_test.go b/coderd/insights_test.go index 47a80df528..693bb48811 100644 --- a/coderd/insights_test.go +++ b/coderd/insights_test.go @@ -609,8 +609,8 @@ func TestTemplateInsights_Golden(t *testing.T) { Name: "example", Type: "aws_instance", Agents: []*proto.Agent{{ - Id: uuid.NewString(), // Doesn't matter, not used in DB. - Name: "dev", + Id: uuid.NewString(), // Doesn't matter, not used in DB. + Name: fmt.Sprintf("dev-%d", len(resources)), // Ensure unique name per agent Auth: &proto.Agent_Token{ Token: authToken.String(), }, @@ -1525,8 +1525,8 @@ func TestUserActivityInsights_Golden(t *testing.T) { Name: "example", Type: "aws_instance", Agents: []*proto.Agent{{ - Id: uuid.NewString(), // Doesn't matter, not used in DB. - Name: "dev", + Id: uuid.NewString(), // Doesn't matter, not used in DB. + Name: fmt.Sprintf("dev-%d", len(resources)), // Ensure unique name per agent Auth: &proto.Agent_Token{ Token: authToken.String(), },