feat(coderd/database): allow filtering provisioner daemons by tags

SasSwart · SasSwart · commit 3d45f553f804 · 2024-11-08T10:32:34.000Z
diff --git a/coderd/database/migrations/000274_check_tags.down.sql b/coderd/database/migrations/000274_check_tags.down.sql
@@ -0,0 +1 @@
+DROP FUNCTION IF EXISTS tags_compatible(jsonb, jsonb);
diff --git a/coderd/database/migrations/000274_check_tags.up.sql b/coderd/database/migrations/000274_check_tags.up.sql
@@ -0,0 +1,10 @@
+CREATE OR REPLACE FUNCTION tags_compatible(provisioner_tags jsonb, required_tags jsonb)
+RETURNS boolean as $$
+BEGIN
+	RETURN CASE
+		WHEN provisioner_tags = '{"scope": "organization", "owner": ""}' :: jsonb
+		THEN provisioner_tags = required_tags
+		ELSE required_tags <@ provisioner_tags
+	END;
+END;
+$$ LANGUAGE plpgsql;
diff --git a/coderd/database/queries/provisionerdaemons.sql b/coderd/database/queries/provisionerdaemons.sql
@@ -10,7 +10,11 @@ SELECT
 FROM
 	provisioner_daemons
 WHERE
-	organization_id = @organization_id;
+	-- If organization_id is provided, filter by it; otherwise, allow all.
+	(@organization_id IS NULL OR organization_id = @organization_id)
+	AND
+	-- If tags are provided, check compatibility; otherwise, skip tags check.
+	(@tags IS NULL OR tags_compatible(@tags :: jsonb, provisioner_daemons.tags :: jsonb));
 
 -- name: DeleteOldProvisionerDaemons :exec
 -- Delete provisioner daemons that have been created at least a week ago
diff --git a/coderd/database/queries/provisionerjobs.sql b/coderd/database/queries/provisionerjobs.sql
@@ -22,13 +22,7 @@ WHERE
 			AND nested.organization_id = @organization_id
 			-- Ensure the caller has the correct provisioner.
 			AND nested.provisioner = ANY(@types :: provisioner_type [ ])
-			AND CASE
-				-- Special case for untagged provisioners: only match untagged jobs.
-				WHEN nested.tags :: jsonb = '{"scope": "organization", "owner": ""}' :: jsonb
-				THEN nested.tags :: jsonb = @tags :: jsonb
-				-- Ensure the caller satisfies all job tags.
-				ELSE nested.tags :: jsonb <@ @tags :: jsonb
-			END
+			AND tags_compatible(nested.tags, @tags)
 		ORDER BY
 			nested.created_at
 		FOR UPDATE
@@ -160,4 +154,4 @@ RETURNING *;
 -- name: GetProvisionerJobTimingsByJobID :many
 SELECT * FROM provisioner_job_timings
 WHERE job_id = $1
-ORDER BY started_at ASC;
+ORDER BY started_at ASC;

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+DROP FUNCTION IF EXISTS tags_compatible(jsonb, jsonb);`