Skip to content

Commit 7d73585

Browse files
mafredrimafredri
committed
fix(coderd): mark sub agent deletion via boolean instead of delete
Fixes coder/internal#685
1 parent 44d4646 commit 7d73585

File tree

12 files changed

+260
-37
lines changed

12 files changed

+260
-37
lines changed

‎coderd/agentapi/subagent_test.go

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -875,14 +875,9 @@ func TestSubAgentAPI(t *testing.T) {
875875
require.NoError(t, err)
876876
})
877877

878-
t.Run("DeletesWorkspaceApps", func(t *testing.T) {
878+
t.Run("DeleteRetainsWorkspaceApps", func(t *testing.T) {
879879
t.Parallel()
880880

881-
// Skip test on in-memory database since CASCADE DELETE is not implemented
882-
if !dbtestutil.WillUsePostgres() {
883-
t.Skip("CASCADE DELETE behavior requires PostgreSQL")
884-
}
885-
886881
log := testutil.Logger(t)
887882
ctx := testutil.Context(t, testutil.WaitShort)
888883
clock := quartz.NewMock(t)
@@ -931,11 +926,11 @@ func TestSubAgentAPI(t *testing.T) {
931926
_, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), subAgentID) //nolint:gocritic // this is a test.
932927
require.ErrorIs(t, err, sql.ErrNoRows)
933928

934-
// And: The apps are also deleted (due to CASCADE DELETE)
935-
// Use raw database since authorization layer requires agent to exist
929+
// And: The apps are *retained* to avoid causing issues
930+
// where the resources are expected to be present.
936931
appsAfterDeletion, err := db.GetWorkspaceAppsByAgentID(ctx, subAgentID)
937932
require.NoError(t, err)
938-
require.Empty(t, appsAfterDeletion)
933+
require.NotEmpty(t, appsAfterDeletion)
939934
})
940935
})
941936

‎coderd/database/dbgen/dbgen.go

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ func WorkspaceAgent(t testing.TB, db database.Store, orig database.WorkspaceAgen
209209
},
210210
ConnectionTimeoutSeconds: takeFirst(orig.ConnectionTimeoutSeconds, 3600),
211211
TroubleshootingURL: takeFirst(orig.TroubleshootingURL, "https://example.com"),
212-
MOTDFile: takeFirst(orig.TroubleshootingURL, ""),
212+
MOTDFile: takeFirst(orig.MOTDFile, ""),
213213
DisplayApps: append([]database.DisplayApp{}, orig.DisplayApps...),
214214
DisplayOrder: takeFirst(orig.DisplayOrder, 1),
215215
APIKeyScope: takeFirst(orig.APIKeyScope, database.AgentKeyScopeEnumAll),
@@ -226,6 +226,35 @@ func WorkspaceAgent(t testing.TB, db database.Store, orig database.WorkspaceAgen
226226
})
227227
require.NoError(t, err, "update workspace agent first connected at")
228228
}
229+
230+
// Add a test antagonist. For every agent we add a deleted sub agent
231+
// to discover cases where deletion should be handled.
232+
subAgt, err := db.InsertWorkspaceAgent(genCtx, database.InsertWorkspaceAgentParams{
233+
ID: uuid.New(),
234+
ParentID: uuid.NullUUID{UUID: agt.ID, Valid: true},
235+
CreatedAt: dbtime.Now(),
236+
UpdatedAt: dbtime.Now(),
237+
Name: testutil.GetRandomName(t),
238+
ResourceID: agt.ResourceID,
239+
AuthToken: uuid.New(),
240+
AuthInstanceID: sql.NullString{},
241+
Architecture: agt.Architecture,
242+
EnvironmentVariables: pqtype.NullRawMessage{},
243+
OperatingSystem: agt.OperatingSystem,
244+
Directory: agt.Directory,
245+
InstanceMetadata: pqtype.NullRawMessage{},
246+
ResourceMetadata: pqtype.NullRawMessage{},
247+
ConnectionTimeoutSeconds: agt.ConnectionTimeoutSeconds,
248+
TroubleshootingURL: agt.TroubleshootingURL,
249+
MOTDFile: "",
250+
DisplayApps: nil,
251+
DisplayOrder: agt.DisplayOrder,
252+
APIKeyScope: agt.APIKeyScope,
253+
})
254+
require.NoError(t, err, "insert workspace agent subagent antagonist")
255+
err = db.DeleteWorkspaceSubAgentByID(genCtx, subAgt.ID)
256+
require.NoError(t, err, "delete workspace agent subagent antagonist")
257+
229258
return agt
230259
}
231260

‎coderd/database/dbmem/dbmem.go

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -792,7 +792,7 @@ func (q *FakeQuerier) getWorkspaceAgentByIDNoLock(_ context.Context, id uuid.UUI
792792
// The schema sorts this by created at, so we iterate the array backwards.
793793
for i := len(q.workspaceAgents) - 1; i >= 0; i-- {
794794
agent := q.workspaceAgents[i]
795-
if agent.ID == id {
795+
if !agent.Deleted && agent.ID == id {
796796
return agent, nil
797797
}
798798
}
@@ -802,6 +802,9 @@ func (q *FakeQuerier) getWorkspaceAgentByIDNoLock(_ context.Context, id uuid.UUI
802802
func (q *FakeQuerier) getWorkspaceAgentsByResourceIDsNoLock(_ context.Context, resourceIDs []uuid.UUID) ([]database.WorkspaceAgent, error) {
803803
workspaceAgents := make([]database.WorkspaceAgent, 0)
804804
for _, agent := range q.workspaceAgents {
805+
if agent.Deleted {
806+
continue
807+
}
805808
for _, resourceID := range resourceIDs {
806809
if agent.ResourceID != resourceID {
807810
continue
@@ -2543,13 +2546,13 @@ func (q *FakeQuerier) DeleteWorkspaceAgentPortSharesByTemplate(_ context.Context
25432546
return nil
25442547
}
25452548

2546-
func (q *FakeQuerier) DeleteWorkspaceSubAgentByID(ctx context.Context, id uuid.UUID) error {
2549+
func (q *FakeQuerier) DeleteWorkspaceSubAgentByID(_ context.Context, id uuid.UUID) error {
25472550
q.mutex.Lock()
25482551
defer q.mutex.Unlock()
25492552

25502553
for i, agent := range q.workspaceAgents {
25512554
if agent.ID == id && agent.ParentID.Valid {
2552-
q.workspaceAgents = slices.Delete(q.workspaceAgents, i, i+1)
2555+
q.workspaceAgents[i].Deleted = true
25532556
return nil
25542557
}
25552558
}
@@ -7066,6 +7069,10 @@ func (q *FakeQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(_ context.Conte
70667069
latestBuildNumber := make(map[uuid.UUID]int32)
70677070

70687071
for _, agt := range q.workspaceAgents {
7072+
if agt.Deleted {
7073+
continue
7074+
}
7075+
70697076
// get the related workspace and user
70707077
for _, res := range q.workspaceResources {
70717078
if agt.ResourceID != res.ID {
@@ -7135,7 +7142,7 @@ func (q *FakeQuerier) GetWorkspaceAgentByInstanceID(_ context.Context, instanceI
71357142
// The schema sorts this by created at, so we iterate the array backwards.
71367143
for i := len(q.workspaceAgents) - 1; i >= 0; i-- {
71377144
agent := q.workspaceAgents[i]
7138-
if agent.AuthInstanceID.Valid && agent.AuthInstanceID.String == instanceID {
7145+
if !agent.Deleted && agent.AuthInstanceID.Valid && agent.AuthInstanceID.String == instanceID {
71397146
return agent, nil
71407147
}
71417148
}
@@ -7695,13 +7702,13 @@ func (q *FakeQuerier) GetWorkspaceAgentUsageStatsAndLabels(_ context.Context, cr
76957702
return stats, nil
76967703
}
76977704

7698-
func (q *FakeQuerier) GetWorkspaceAgentsByParentID(ctx context.Context, parentID uuid.UUID) ([]database.WorkspaceAgent, error) {
7705+
func (q *FakeQuerier) GetWorkspaceAgentsByParentID(_ context.Context, parentID uuid.UUID) ([]database.WorkspaceAgent, error) {
76997706
q.mutex.RLock()
77007707
defer q.mutex.RUnlock()
77017708

77027709
workspaceAgents := make([]database.WorkspaceAgent, 0)
77037710
for _, agent := range q.workspaceAgents {
7704-
if !agent.ParentID.Valid || agent.ParentID.UUID != parentID {
7711+
if !agent.ParentID.Valid || agent.ParentID.UUID != parentID || agent.Deleted {
77057712
continue
77067713
}
77077714

@@ -7748,6 +7755,9 @@ func (q *FakeQuerier) GetWorkspaceAgentsCreatedAfter(_ context.Context, after ti
77487755

77497756
workspaceAgents := make([]database.WorkspaceAgent, 0)
77507757
for _, agent := range q.workspaceAgents {
7758+
if agent.Deleted {
7759+
continue
7760+
}
77517761
if agent.CreatedAt.After(after) {
77527762
workspaceAgents = append(workspaceAgents, agent)
77537763
}

‎coderd/database/dump.sql

Lines changed: 5 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000338_use_deleted_boolean_for_subagents.down.sql

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

‎coderd/database/migrations/000338_use_deleted_boolean_for_subagents.up.sql

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
ALTER TABLE workspace_agents
2+
ADD COLUMN deleted BOOLEAN NOT NULL DEFAULT FALSE;
3+
4+
COMMENT ON COLUMN workspace_agents.deleted IS 'Indicates whether or not the agent has been deleted. This is currently only applicable to sub agents.';
5+
6+
-- Recreate the trigger with deleted check.
7+
DROP TRIGGER IF EXISTS workspace_agent_name_unique_trigger ON workspace_agents;
8+
DROP FUNCTION IF EXISTS check_workspace_agent_name_unique();
9+
10+
CREATE OR REPLACE FUNCTION check_workspace_agent_name_unique()
11+
RETURNS TRIGGER AS $$
12+
DECLARE
13+
workspace_build_id uuid;
14+
agents_with_name int;
15+
BEGIN
16+
-- Find the workspace build the workspace agent is being inserted into.
17+
SELECT workspace_builds.id INTO workspace_build_id
18+
FROM workspace_resources
19+
JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id
20+
WHERE workspace_resources.id = NEW.resource_id;
21+
22+
-- If the agent doesn't have a workspace build, we'll allow the insert.
23+
IF workspace_build_id IS NULL THEN
24+
RETURN NEW;
25+
END IF;
26+
27+
-- Count how many agents in this workspace build already have the given agent name.
28+
SELECT COUNT(*) INTO agents_with_name
29+
FROM workspace_agents
30+
JOIN workspace_resources ON workspace_resources.id = workspace_agents.resource_id
31+
JOIN workspace_builds ON workspace_builds.job_id = workspace_resources.job_id
32+
WHERE workspace_builds.id = workspace_build_id
33+
AND workspace_agents.name = NEW.name
34+
AND workspace_agents.id != NEW.id
35+
AND workspace_agents.deleted = FALSE; -- Ensure we only count non-deleted agents.
36+
37+
-- If there's already an agent with this name, raise an error
38+
IF agents_with_name > 0 THEN
39+
RAISE EXCEPTION 'workspace agent name "%" already exists in this workspace build', NEW.name
40+
USING ERRCODE = 'unique_violation';
41+
END IF;
42+
43+
RETURN NEW;
44+
END;
45+
$$ LANGUAGE plpgsql;
46+
47+
CREATE TRIGGER workspace_agent_name_unique_trigger
48+
BEFORE INSERT OR UPDATE OF name, resource_id ON workspace_agents
49+
FOR EACH ROW
50+
EXECUTE FUNCTION check_workspace_agent_name_unique();
51+
52+
COMMENT ON TRIGGER workspace_agent_name_unique_trigger ON workspace_agents IS
53+
'Use a trigger instead of a unique constraint because existing data may violate
54+
the uniqueness requirement. A trigger allows us to enforce uniqueness going
55+
forward without requiring a migration to clean up historical data.';

‎coderd/database/models.go

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)