Skip to content

Commit 6e255c7

Browse files
chore(coderd/database): enforce agent name unique within workspace build (coder#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.
1 parent 110102a commit 6e255c7

5 files changed

+328
-4
lines changed

coderd/database/dump.sql

Lines changed: 43 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
DROP TRIGGER IF EXISTS workspace_agent_name_unique_trigger ON workspace_agents;
2+
DROP FUNCTION IF EXISTS check_workspace_agent_name_unique();
Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
CREATE OR REPLACE FUNCTION check_workspace_agent_name_unique()
2+
RETURNS TRIGGER AS $$
3+
DECLARE
4+
workspace_build_id uuid;
5+
agents_with_name int;
6+
BEGIN
7+
-- Find the workspace build the workspace agent is being inserted into.
8+
SELECT workspace_builds.id INTO workspace_build_id
9+
FROM workspace_resources
10+
JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id
11+
WHERE workspace_resources.id = NEW.resource_id;
12+
13+
-- If the agent doesn't have a workspace build, we'll allow the insert.
14+
IF workspace_build_id IS NULL THEN
15+
RETURN NEW;
16+
END IF;
17+
18+
-- Count how many agents in this workspace build already have the given agent name.
19+
SELECT COUNT(*) INTO agents_with_name
20+
FROM workspace_agents
21+
JOIN workspace_resources ON workspace_resources.id = workspace_agents.resource_id
22+
JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id
23+
WHERE workspace_builds.id = workspace_build_id
24+
AND workspace_agents.name = NEW.name
25+
AND workspace_agents.id != NEW.id;
26+
27+
-- If there's already an agent with this name, raise an error
28+
IF agents_with_name > 0 THEN
29+
RAISE EXCEPTION 'workspace agent name "%" already exists in this workspace build', NEW.name
30+
USING ERRCODE = 'unique_violation';
31+
END IF;
32+
33+
RETURN NEW;
34+
END;
35+
$$ LANGUAGE plpgsql;
36+
37+
CREATE TRIGGER workspace_agent_name_unique_trigger
38+
BEFORE INSERT OR UPDATE OF name, resource_id ON workspace_agents
39+
FOR EACH ROW
40+
EXECUTE FUNCTION check_workspace_agent_name_unique();
41+
42+
COMMENT ON TRIGGER workspace_agent_name_unique_trigger ON workspace_agents IS
43+
'Use a trigger instead of a unique constraint because existing data may violate
44+
the uniqueness requirement. A trigger allows us to enforce uniqueness going
45+
forward without requiring a migration to clean up historical data.';

coderd/database/querier_test.go

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@ import (
44
"context"
55
"database/sql"
66
"encoding/json"
7+
"errors"
78
"fmt"
89
"sort"
910
"testing"
1011
"time"
1112

1213
"github.com/google/uuid"
14+
"github.com/lib/pq"
1315
"github.com/prometheus/client_golang/prometheus"
1416
"github.com/stretchr/testify/assert"
1517
"github.com/stretchr/testify/require"
@@ -4720,6 +4722,238 @@ func TestGetPresetsAtFailureLimit(t *testing.T) {
47204722
})
47214723
}
47224724

4725+
func TestWorkspaceAgentNameUniqueTrigger(t *testing.T) {
4726+
t.Parallel()
4727+
4728+
if !dbtestutil.WillUsePostgres() {
4729+
t.Skip("This test makes use of a database trigger not implemented in dbmem")
4730+
}
4731+
4732+
createWorkspaceWithAgent := func(t *testing.T, db database.Store, org database.Organization, agentName string) (database.WorkspaceBuild, database.WorkspaceResource, database.WorkspaceAgent) {
4733+
t.Helper()
4734+
4735+
user := dbgen.User(t, db, database.User{})
4736+
template := dbgen.Template(t, db, database.Template{
4737+
OrganizationID: org.ID,
4738+
CreatedBy: user.ID,
4739+
})
4740+
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
4741+
TemplateID: uuid.NullUUID{Valid: true, UUID: template.ID},
4742+
OrganizationID: org.ID,
4743+
CreatedBy: user.ID,
4744+
})
4745+
workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
4746+
OrganizationID: org.ID,
4747+
TemplateID: template.ID,
4748+
OwnerID: user.ID,
4749+
})
4750+
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
4751+
Type: database.ProvisionerJobTypeWorkspaceBuild,
4752+
OrganizationID: org.ID,
4753+
})
4754+
build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
4755+
BuildNumber: 1,
4756+
JobID: job.ID,
4757+
WorkspaceID: workspace.ID,
4758+
TemplateVersionID: templateVersion.ID,
4759+
})
4760+
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
4761+
JobID: build.JobID,
4762+
})
4763+
agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
4764+
ResourceID: resource.ID,
4765+
Name: agentName,
4766+
})
4767+
4768+
return build, resource, agent
4769+
}
4770+
4771+
t.Run("DuplicateNamesInSameWorkspaceResource", func(t *testing.T) {
4772+
t.Parallel()
4773+
4774+
db, _ := dbtestutil.NewDB(t)
4775+
org := dbgen.Organization(t, db, database.Organization{})
4776+
ctx := testutil.Context(t, testutil.WaitShort)
4777+
4778+
// Given: A workspace with an agent
4779+
_, resource, _ := createWorkspaceWithAgent(t, db, org, "duplicate-agent")
4780+
4781+
// When: Another agent is created for that workspace with the same name.
4782+
_, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{
4783+
ID: uuid.New(),
4784+
CreatedAt: time.Now(),
4785+
UpdatedAt: time.Now(),
4786+
Name: "duplicate-agent", // Same name as agent1
4787+
ResourceID: resource.ID,
4788+
AuthToken: uuid.New(),
4789+
Architecture: "amd64",
4790+
OperatingSystem: "linux",
4791+
APIKeyScope: database.AgentKeyScopeEnumAll,
4792+
})
4793+
4794+
// Then: We expect it to fail.
4795+
require.Error(t, err)
4796+
var pqErr *pq.Error
4797+
require.True(t, errors.As(err, &pqErr))
4798+
require.Equal(t, pq.ErrorCode("23505"), pqErr.Code) // unique_violation
4799+
require.Contains(t, pqErr.Message, `workspace agent name "duplicate-agent" already exists in this workspace build`)
4800+
})
4801+
4802+
t.Run("DuplicateNamesInSameProvisionerJob", func(t *testing.T) {
4803+
t.Parallel()
4804+
4805+
db, _ := dbtestutil.NewDB(t)
4806+
org := dbgen.Organization(t, db, database.Organization{})
4807+
ctx := testutil.Context(t, testutil.WaitShort)
4808+
4809+
// Given: A workspace with an agent
4810+
_, resource, agent := createWorkspaceWithAgent(t, db, org, "duplicate-agent")
4811+
4812+
// When: A child agent is created for that workspace with the same name.
4813+
_, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{
4814+
ID: uuid.New(),
4815+
CreatedAt: time.Now(),
4816+
UpdatedAt: time.Now(),
4817+
Name: agent.Name,
4818+
ResourceID: resource.ID,
4819+
AuthToken: uuid.New(),
4820+
Architecture: "amd64",
4821+
OperatingSystem: "linux",
4822+
APIKeyScope: database.AgentKeyScopeEnumAll,
4823+
})
4824+
4825+
// Then: We expect it to fail.
4826+
require.Error(t, err)
4827+
var pqErr *pq.Error
4828+
require.True(t, errors.As(err, &pqErr))
4829+
require.Equal(t, pq.ErrorCode("23505"), pqErr.Code) // unique_violation
4830+
require.Contains(t, pqErr.Message, `workspace agent name "duplicate-agent" already exists in this workspace build`)
4831+
})
4832+
4833+
t.Run("DuplicateChildNamesOverMultipleResources", func(t *testing.T) {
4834+
t.Parallel()
4835+
4836+
db, _ := dbtestutil.NewDB(t)
4837+
org := dbgen.Organization(t, db, database.Organization{})
4838+
ctx := testutil.Context(t, testutil.WaitShort)
4839+
4840+
// Given: A workspace with two agents
4841+
_, resource1, agent1 := createWorkspaceWithAgent(t, db, org, "parent-agent-1")
4842+
4843+
resource2 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{JobID: resource1.JobID})
4844+
agent2 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
4845+
ResourceID: resource2.ID,
4846+
Name: "parent-agent-2",
4847+
})
4848+
4849+
// Given: One agent has a child agent
4850+
agent1Child := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
4851+
ParentID: uuid.NullUUID{Valid: true, UUID: agent1.ID},
4852+
Name: "child-agent",
4853+
ResourceID: resource1.ID,
4854+
})
4855+
4856+
// When: A child agent is inserted for the other parent.
4857+
_, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{
4858+
ID: uuid.New(),
4859+
ParentID: uuid.NullUUID{Valid: true, UUID: agent2.ID},
4860+
CreatedAt: time.Now(),
4861+
UpdatedAt: time.Now(),
4862+
Name: agent1Child.Name,
4863+
ResourceID: resource2.ID,
4864+
AuthToken: uuid.New(),
4865+
Architecture: "amd64",
4866+
OperatingSystem: "linux",
4867+
APIKeyScope: database.AgentKeyScopeEnumAll,
4868+
})
4869+
4870+
// Then: We expect it to fail.
4871+
require.Error(t, err)
4872+
var pqErr *pq.Error
4873+
require.True(t, errors.As(err, &pqErr))
4874+
require.Equal(t, pq.ErrorCode("23505"), pqErr.Code) // unique_violation
4875+
require.Contains(t, pqErr.Message, `workspace agent name "child-agent" already exists in this workspace build`)
4876+
})
4877+
4878+
t.Run("SameNamesInDifferentWorkspaces", func(t *testing.T) {
4879+
t.Parallel()
4880+
4881+
agentName := "same-name-different-workspace"
4882+
4883+
db, _ := dbtestutil.NewDB(t)
4884+
org := dbgen.Organization(t, db, database.Organization{})
4885+
4886+
// Given: A workspace with an agent
4887+
_, _, agent1 := createWorkspaceWithAgent(t, db, org, agentName)
4888+
require.Equal(t, agentName, agent1.Name)
4889+
4890+
// When: A second workspace is created with an agent having the same name
4891+
_, _, agent2 := createWorkspaceWithAgent(t, db, org, agentName)
4892+
require.Equal(t, agentName, agent2.Name)
4893+
4894+
// Then: We expect there to be different agents with the same name.
4895+
require.NotEqual(t, agent1.ID, agent2.ID)
4896+
require.Equal(t, agent1.Name, agent2.Name)
4897+
})
4898+
4899+
t.Run("NullWorkspaceID", func(t *testing.T) {
4900+
t.Parallel()
4901+
4902+
db, _ := dbtestutil.NewDB(t)
4903+
org := dbgen.Organization(t, db, database.Organization{})
4904+
ctx := testutil.Context(t, testutil.WaitShort)
4905+
4906+
// Given: A resource that does not belong to a workspace build (simulating template import)
4907+
orphanJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
4908+
Type: database.ProvisionerJobTypeTemplateVersionImport,
4909+
OrganizationID: org.ID,
4910+
})
4911+
orphanResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
4912+
JobID: orphanJob.ID,
4913+
})
4914+
4915+
// And this resource has a workspace agent.
4916+
agent1, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{
4917+
ID: uuid.New(),
4918+
CreatedAt: time.Now(),
4919+
UpdatedAt: time.Now(),
4920+
Name: "orphan-agent",
4921+
ResourceID: orphanResource.ID,
4922+
AuthToken: uuid.New(),
4923+
Architecture: "amd64",
4924+
OperatingSystem: "linux",
4925+
APIKeyScope: database.AgentKeyScopeEnumAll,
4926+
})
4927+
require.NoError(t, err)
4928+
require.Equal(t, "orphan-agent", agent1.Name)
4929+
4930+
// When: We created another resource that does not belong to a workspace build.
4931+
orphanJob2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
4932+
Type: database.ProvisionerJobTypeTemplateVersionImport,
4933+
OrganizationID: org.ID,
4934+
})
4935+
orphanResource2 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
4936+
JobID: orphanJob2.ID,
4937+
})
4938+
4939+
// Then: We expect to be able to create an agent in this new resource that has the same name.
4940+
agent2, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{
4941+
ID: uuid.New(),
4942+
CreatedAt: time.Now(),
4943+
UpdatedAt: time.Now(),
4944+
Name: "orphan-agent", // Same name as agent1
4945+
ResourceID: orphanResource2.ID,
4946+
AuthToken: uuid.New(),
4947+
Architecture: "amd64",
4948+
OperatingSystem: "linux",
4949+
APIKeyScope: database.AgentKeyScopeEnumAll,
4950+
})
4951+
require.NoError(t, err)
4952+
require.Equal(t, "orphan-agent", agent2.Name)
4953+
require.NotEqual(t, agent1.ID, agent2.ID)
4954+
})
4955+
}
4956+
47234957
func requireUsersMatch(t testing.TB, expected []database.User, found []database.GetUsersRow, msg string) {
47244958
t.Helper()
47254959
require.ElementsMatch(t, expected, database.ConvertUserRows(found), msg)

coderd/insights_test.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -609,8 +609,8 @@ func TestTemplateInsights_Golden(t *testing.T) {
609609
Name: "example",
610610
Type: "aws_instance",
611611
Agents: []*proto.Agent{{
612-
Id: uuid.NewString(), // Doesn't matter, not used in DB.
613-
Name: "dev",
612+
Id: uuid.NewString(), // Doesn't matter, not used in DB.
613+
Name: fmt.Sprintf("dev-%d", len(resources)), // Ensure unique name per agent
614614
Auth: &proto.Agent_Token{
615615
Token: authToken.String(),
616616
},
@@ -1525,8 +1525,8 @@ func TestUserActivityInsights_Golden(t *testing.T) {
15251525
Name: "example",
15261526
Type: "aws_instance",
15271527
Agents: []*proto.Agent{{
1528-
Id: uuid.NewString(), // Doesn't matter, not used in DB.
1529-
Name: "dev",
1528+
Id: uuid.NewString(), // Doesn't matter, not used in DB.
1529+
Name: fmt.Sprintf("dev-%d", len(resources)), // Ensure unique name per agent
15301530
Auth: &proto.Agent_Token{
15311531
Token: authToken.String(),
15321532
},

0 commit comments

Comments
 (0)