perf(db): optimize in-progress builds queries by filtering on b.team_id (#1987)

Author: dobrac

packages/api/internal/handlers/admin_cancel_team_builds.go

@@ -23,7 +23,7 @@ func (a *APIStore) PostAdminTeamsTeamIDBuildsCancel(c *gin.Context, teamID uuid.
 
 	logger.L().Info(ctx, "Admin cancelling all builds for team", logger.WithTeamID(teamID.String()))
 
-	builds, err := a.sqlcDB.GetCancellableTemplateBuildsByTeam(ctx, teamID)
+	builds, err := a.sqlcDB.GetCancellableTemplateBuildsByTeam(ctx, &teamID)
 	if err != nil {
 		a.sendAPIStoreError(c, http.StatusInternalServerError, "Failed to get builds")
 

packages/api/internal/template/register_build.go

@@ -60,9 +60,19 @@ func RegisterBuild(
 	ctx, span := tracer.Start(ctx, "register build")
 	defer span.End()
 
+	// Add default tag if no tags are present
+	tags := data.Tags
+	if len(tags) == 0 {
+		tags = []string{id.DefaultTag}
+	}
+
 	// This is a simple implementation of concurrency limit
 	// It does not guarantee that the limit is not exceeded, but it should be good enough for now (considering overall low number of total builds)
-	templateIDs, err := db.GetInProgressTemplateBuildsByTeam(ctx, data.Team.ID)
+	otherBuildCount, err := db.GetInProgressTemplateBuildsByTeam(ctx, queries.GetInProgressTemplateBuildsByTeamParams{
+		TeamID:            &data.Team.ID,
+		ExcludeTemplateID: data.TemplateID,
+		ExcludeTags:       tags,
+	})
 	if err != nil {
 		return nil, &api.APIError{
 			Err:       err,
@@ -71,13 +81,8 @@ func RegisterBuild(
 		}
 	}
 
-	// Exclude the current build if it's a rebuild (it will be canceled)
-	teamBuildsExcludingCurrent := gutils.Filter(templateIDs, func(templateID string) bool {
-		return templateID != data.TemplateID
-	})
-
 	totalConcurrentTemplateBuilds := data.Team.Limits.BuildConcurrency
-	if len(teamBuildsExcludingCurrent) >= int(totalConcurrentTemplateBuilds) {
+	if otherBuildCount >= totalConcurrentTemplateBuilds {
 		telemetry.ReportError(ctx, "team has reached max concurrent template builds", nil, telemetry.WithTeamID(data.Team.ID.String()), attribute.Int64("total.concurrent_template_builds", totalConcurrentTemplateBuilds))
 
 		return nil, &api.APIError{
@@ -101,12 +106,6 @@ func RegisterBuild(
 		}
 	}
 
-	// Add default tag if no tags are present
-	tags := data.Tags
-	if len(tags) == 0 {
-		tags = []string{id.DefaultTag}
-	}
-
 	telemetry.SetAttributes(ctx,
 		attribute.String("env.team.id", data.Team.ID.String()),
 		attribute.String("env.team.name", data.Team.Name),

packages/db/migrations/20260225120000_add_env_builds_team_status_group_index.sql

@@ -0,0 +1,10 @@
+-- +goose Up
+-- +goose NO TRANSACTION
+
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_env_builds_team_status_group
+  ON public.env_builds (team_id, status_group);
+
+-- +goose Down
+-- +goose NO TRANSACTION
+
+DROP INDEX CONCURRENTLY IF EXISTS idx_env_builds_team_status_group;

packages/db/queries/builds/get_inprogress_builds.sql

@@ -8,19 +8,27 @@ WHERE b.status_group IN ('pending', 'in_progress')
   AND e.source = 'template'
 ORDER BY b.id, b.created_at DESC;
 
--- name: GetInProgressTemplateBuildsByTeam :many
-SELECT DISTINCT ON (b.id) e.id as template_id
+-- name: GetInProgressTemplateBuildsByTeam :one
+SELECT COUNT(DISTINCT b.id) as build_count
 FROM public.env_builds b
 JOIN public.env_build_assignments eba ON eba.build_id = b.id
 JOIN public.envs e ON e.id = eba.env_id
-WHERE e.team_id = $1 AND b.status_group IN ('pending', 'in_progress') AND e.source = 'template'
-ORDER BY b.id, b.created_at DESC;
+WHERE b.team_id = @team_id
+  AND b.status_group IN ('pending', 'in_progress')
+  AND e.source = 'template'
+  AND NOT EXISTS (
+    SELECT 1 FROM public.env_build_assignments exc
+    WHERE exc.build_id = b.id
+      AND exc.env_id = @exclude_template_id
+      AND exc.tag = ANY(@exclude_tags::text[])
+  );
 
 -- name: GetCancellableTemplateBuildsByTeam :many
 SELECT DISTINCT ON (b.id) b.id as build_id, e.id as template_id, e.cluster_id, b.cluster_node_id
 FROM public.env_builds b
 JOIN public.env_build_assignments eba ON eba.build_id = b.id
 JOIN public.envs e ON e.id = eba.env_id
-WHERE e.team_id = $1 AND b.status_group IN ('pending', 'in_progress')
+WHERE b.team_id = $1
+  AND b.status_group IN ('pending', 'in_progress')
   AND e.source = 'template'
-ORDER BY b.id, b.created_at DESC;
+ORDER BY b.id;