Skip to content

Commit bb2ad6c

Browse files
SasSwartpull[bot]
SasSwart
authored and
pull[bot]
committed
feat(coderd): update API to allow filtering provisioner daemons by tags (#15448)
This PR provides new parameters to an endpoint that will be necessary for #15048
1 parent a845622 commit bb2ad6c

27 files changed

+389
-80
lines changed

‎coderd/apidoc/docs.go

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

‎coderd/apidoc/swagger.json

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

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1890,7 +1890,7 @@ func (q *querier) GetProvisionerDaemons(ctx context.Context) ([]database.Provisi
18901890
return fetchWithPostFilter(q.auth, policy.ActionRead, fetch)(ctx, nil)
18911891
}
18921892

1893-
func (q *querier) GetProvisionerDaemonsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerDaemon, error) {
1893+
func (q *querier) GetProvisionerDaemonsByOrganization(ctx context.Context, organizationID database.GetProvisionerDaemonsByOrganizationParams) ([]database.ProvisionerDaemon, error) {
18941894
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetProvisionerDaemonsByOrganization)(ctx, organizationID)
18951895
}
18961896

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2066,9 +2066,9 @@ func (s *MethodTestSuite) TestExtraMethods() {
20662066
}),
20672067
})
20682068
s.NoError(err, "insert provisioner daemon")
2069-
ds, err := db.GetProvisionerDaemonsByOrganization(context.Background(), org.ID)
2069+
ds, err := db.GetProvisionerDaemonsByOrganization(context.Background(), database.GetProvisionerDaemonsByOrganizationParams{OrganizationID: org.ID})
20702070
s.NoError(err, "get provisioner daemon by org")
2071-
check.Args(org.ID).Asserts(d, policy.ActionRead).Returns(ds)
2071+
check.Args(database.GetProvisionerDaemonsByOrganizationParams{OrganizationID: org.ID}).Asserts(d, policy.ActionRead).Returns(ds)
20722072
}))
20732073
s.Run("DeleteOldProvisionerDaemons", s.Subtest(func(db database.Store, check *expects) {
20742074
_, err := db.UpsertProvisionerDaemon(context.Background(), database.UpsertProvisionerDaemonParams{
@@ -2560,7 +2560,7 @@ func (s *MethodTestSuite) TestSystemFunctions() {
25602560
j := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{
25612561
StartedAt: sql.NullTime{Valid: false},
25622562
})
2563-
check.Args(database.AcquireProvisionerJobParams{OrganizationID: j.OrganizationID, Types: []database.ProvisionerType{j.Provisioner}, Tags: must(json.Marshal(j.Tags))}).
2563+
check.Args(database.AcquireProvisionerJobParams{OrganizationID: j.OrganizationID, Types: []database.ProvisionerType{j.Provisioner}, ProvisionerTags: must(json.Marshal(j.Tags))}).
25642564
Asserts( /*rbac.ResourceSystem, policy.ActionUpdate*/ )
25652565
}))
25662566
s.Run("UpdateProvisionerJobWithCompleteByID", s.Subtest(func(db database.Store, check *expects) {

‎coderd/database/dbfake/dbfake.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -194,8 +194,8 @@ func (b WorkspaceBuildBuilder) Do() WorkspaceResponse {
194194
UUID: uuid.New(),
195195
Valid: true,
196196
},
197-
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
198-
Tags: []byte(`{"scope": "organization"}`),
197+
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
198+
ProvisionerTags: []byte(`{"scope": "organization"}`),
199199
})
200200
require.NoError(b.t, err, "acquire starting job")
201201
if j.ID == job.ID {

‎coderd/database/dbgen/dbgen.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -531,11 +531,11 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
531531
}
532532
if !orig.StartedAt.Time.IsZero() {
533533
job, err = db.AcquireProvisionerJob(genCtx, database.AcquireProvisionerJobParams{
534-
StartedAt: orig.StartedAt,
535-
OrganizationID: job.OrganizationID,
536-
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
537-
Tags: must(json.Marshal(orig.Tags)),
538-
WorkerID: uuid.NullUUID{},
534+
StartedAt: orig.StartedAt,
535+
OrganizationID: job.OrganizationID,
536+
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
537+
ProvisionerTags: must(json.Marshal(orig.Tags)),
538+
WorkerID: uuid.NullUUID{},
539539
})
540540
require.NoError(t, err)
541541
// There is no easy way to make sure we acquire the correct job.

‎coderd/database/dbmem/dbmem.go

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1194,8 +1194,8 @@ func (q *FakeQuerier) AcquireProvisionerJob(_ context.Context, arg database.Acqu
11941194
continue
11951195
}
11961196
tags := map[string]string{}
1197-
if arg.Tags != nil {
1198-
err := json.Unmarshal(arg.Tags, &tags)
1197+
if arg.ProvisionerTags != nil {
1198+
err := json.Unmarshal(arg.ProvisionerTags, &tags)
11991199
if err != nil {
12001200
return provisionerJob, xerrors.Errorf("unmarshal: %w", err)
12011201
}
@@ -3625,16 +3625,28 @@ func (q *FakeQuerier) GetProvisionerDaemons(_ context.Context) ([]database.Provi
36253625
return out, nil
36263626
}
36273627

3628-
func (q *FakeQuerier) GetProvisionerDaemonsByOrganization(_ context.Context, organizationID uuid.UUID) ([]database.ProvisionerDaemon, error) {
3628+
func (q *FakeQuerier) GetProvisionerDaemonsByOrganization(_ context.Context, arg database.GetProvisionerDaemonsByOrganizationParams) ([]database.ProvisionerDaemon, error) {
36293629
q.mutex.RLock()
36303630
defer q.mutex.RUnlock()
36313631

36323632
daemons := make([]database.ProvisionerDaemon, 0)
36333633
for _, daemon := range q.provisionerDaemons {
3634-
if daemon.OrganizationID == organizationID {
3635-
daemon.Tags = maps.Clone(daemon.Tags)
3636-
daemons = append(daemons, daemon)
3634+
if daemon.OrganizationID != arg.OrganizationID {
3635+
continue
3636+
}
3637+
// Special case for untagged provisioners: only match untagged jobs.
3638+
// Ref: coderd/database/queries/provisionerjobs.sql:24-30
3639+
// CASE WHEN nested.tags :: jsonb = '{"scope": "organization", "owner": ""}' :: jsonb
3640+
// THEN nested.tags :: jsonb = @tags :: jsonb
3641+
if tagsEqual(arg.WantTags, tagsUntagged) && !tagsEqual(arg.WantTags, daemon.Tags) {
3642+
continue
3643+
}
3644+
// ELSE nested.tags :: jsonb <@ @tags :: jsonb
3645+
if !tagsSubset(arg.WantTags, daemon.Tags) {
3646+
continue
36373647
}
3648+
daemon.Tags = maps.Clone(daemon.Tags)
3649+
daemons = append(daemons, daemon)
36383650
}
36393651

36403652
return daemons, nil

‎coderd/database/dbmetrics/querymetrics.go

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

‎coderd/database/dbmock/dbmock.go

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

‎coderd/database/dump.sql

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

‎coderd/database/migrations/000274_check_tags.down.sql

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
DROP FUNCTION IF EXISTS provisioner_tagset_contains(tagset, tagset);
2+
3+
DROP DOMAIN IF EXISTS tagset;

‎coderd/database/migrations/000274_check_tags.up.sql

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
CREATE DOMAIN tagset AS jsonb;
2+
3+
COMMENT ON DOMAIN tagset IS 'A set of tags that match provisioner daemons to provisioner jobs, which can originate from workspaces or templates. tagset is a narrowed type over jsonb. It is expected to be the JSON representation of map[string]string. That is, {"key1": "value1", "key2": "value2"}. We need the narrowed type instead of just using jsonb so that we can give sqlc a type hint, otherwise it defaults to json.RawMessage. json.RawMessage is a suboptimal type to use in the context that we need tagset for.';
4+
5+
CREATE OR REPLACE FUNCTION provisioner_tagset_contains(provisioner_tags tagset, job_tags tagset)
6+
RETURNS boolean AS $$
7+
BEGIN
8+
RETURN CASE
9+
-- Special case for untagged provisioners, where only an exact match should count
10+
WHEN job_tags::jsonb = '{"scope": "organization", "owner": ""}'::jsonb THEN job_tags::jsonb = provisioner_tags::jsonb
11+
-- General case
12+
ELSE job_tags::jsonb <@ provisioner_tags::jsonb
13+
END;
14+
END;
15+
$$ LANGUAGE plpgsql;
16+
17+
COMMENT ON FUNCTION provisioner_tagset_contains(tagset, tagset) IS 'Returns true if the provisioner_tags contains the job_tags, or if the job_tags represents an untagged provisioner and the superset is exactly equal to the subset.';

‎coderd/database/querier.go

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

‎coderd/database/querier_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1020,7 +1020,7 @@ func TestQueuePosition(t *testing.T) {
10201020
UUID: uuid.New(),
10211021
Valid: true,
10221022
},
1023-
Tags: json.RawMessage("{}"),
1023+
ProvisionerTags: json.RawMessage("{}"),
10241024
})
10251025
require.NoError(t, err)
10261026
require.Equal(t, jobs[0].ID, job.ID)

‎coderd/database/queries.sql.go

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

‎coderd/database/queries/provisionerdaemons.sql

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,11 @@ SELECT
1010
FROM
1111
provisioner_daemons
1212
WHERE
13-
organization_id = @organization_id;
13+
-- This is the original search criteria:
14+
organization_id = @organization_id :: uuid
15+
AND
16+
-- adding support for searching by tags:
17+
(@want_tags :: tagset = 'null' :: tagset OR provisioner_tagset_contains(provisioner_daemons.tags::tagset, @want_tags::tagset));
1418

1519
-- name: DeleteOldProvisionerDaemons :exec
1620
-- Delete provisioner daemons that have been created at least a week ago

‎coderd/database/queries/provisionerjobs.sql

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,21 +16,17 @@ WHERE
1616
SELECT
1717
id
1818
FROM
19-
provisioner_jobs AS nested
19+
provisioner_jobs AS potential_job
2020
WHERE
21-
nested.started_at IS NULL
22-
AND nested.organization_id = @organization_id
21+
potential_job.started_at IS NULL
22+
AND potential_job.organization_id = @organization_id
2323
-- Ensure the caller has the correct provisioner.
24-
AND nested.provisioner = ANY(@types :: provisioner_type [ ])
25-
AND CASE
26-
-- Special case for untagged provisioners: only match untagged jobs.
27-
WHEN nested.tags :: jsonb = '{"scope": "organization", "owner": ""}' :: jsonb
28-
THEN nested.tags :: jsonb = @tags :: jsonb
29-
-- Ensure the caller satisfies all job tags.
30-
ELSE nested.tags :: jsonb <@ @tags :: jsonb
31-
END
24+
AND potential_job.provisioner = ANY(@types :: provisioner_type [ ])
25+
-- elsewhere, we use the tagset type, but here we use jsonb for backward compatibility
26+
-- they are aliases and the code that calls this query already relies on a different type
27+
AND provisioner_tagset_contains(@provisioner_tags :: jsonb, potential_job.tags :: jsonb)
3228
ORDER BY
33-
nested.created_at
29+
potential_job.created_at
3430
FOR UPDATE
3531
SKIP LOCKED
3632
LIMIT
@@ -160,4 +156,4 @@ RETURNING *;
160156
-- name: GetProvisionerJobTimingsByJobID :many
161157
SELECT * FROM provisioner_job_timings
162158
WHERE job_id = $1
163-
ORDER BY started_at ASC;
159+
ORDER BY started_at ASC;

‎coderd/database/sqlc.yaml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,9 @@ sql:
3535
- db_type: "name_organization_pair"
3636
go_type:
3737
type: "NameOrganizationPair"
38+
- db_type: "tagset"
39+
go_type:
40+
type: "StringMap"
3841
- column: "custom_roles.site_permissions"
3942
go_type:
4043
type: "CustomRolePermissions"

‎coderd/provisionerdserver/acquirer.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -130,8 +130,8 @@ func (a *Acquirer) AcquireJob(
130130
UUID: worker,
131131
Valid: true,
132132
},
133-
Types: pt,
134-
Tags: dbTags,
133+
Types: pt,
134+
ProvisionerTags: dbTags,
135135
})
136136
if xerrors.Is(err, sql.ErrNoRows) {
137137
logger.Debug(ctx, "no job available")

‎coderd/provisionerdserver/acquirer_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -653,7 +653,7 @@ func (s *fakeTaggedStore) AcquireProvisionerJob(
653653
) {
654654
defer func() { s.params <- params }()
655655
var tags provisionerdserver.Tags
656-
err := json.Unmarshal(params.Tags, &tags)
656+
err := json.Unmarshal(params.ProvisionerTags, &tags)
657657
if !assert.NoError(s.t, err) {
658658
return database.ProvisionerJob{}, err
659659
}

0 commit comments

Comments
 (0)