fix(metrics): speed up top workflows by runs count query

app/controlplane/pkg/biz/orgmetrics_integration_test.go

@@ -18,19 +18,179 @@ package biz_test
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/biz"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/biz/testhelpers"
 	"github.com/chainloop-dev/chainloop/pkg/credentials"
 	creds "github.com/chainloop-dev/chainloop/pkg/credentials/mocks"
+	"github.com/google/uuid"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 )
 
+// OCI credentials used to create the CAS backend required by workflow runs
+const (
+	ociTestRepo     = "repo"
+	ociTestUsername = "username"
+	ociTestPassword = "pass"
+)
+
+var ociTestKeypair = &credentials.OCIKeypair{Repo: ociTestRepo, Username: ociTestUsername, Password: ociTestPassword}
+
 func TestOrgMetricsUseCase(t *testing.T) {
 	suite.Run(t, new(orgMetricsGetLastWorkflowStatusByRunTestSuite))
+	suite.Run(t, new(orgMetricsTopWorkflowsTestSuite))
+}
+
+type orgMetricsTopWorkflowsTestSuite struct {
+	testhelpers.UseCasesEachTestSuite
+	org, otherOrg *biz.Organization
+	p1, p2        uuid.UUID
+}
+
+// SetupTest creates the following dataset
+// org:
+//   - wf-a (project p1): 2 success, 1 error runs
+//   - wf-b (project p2): 2 success runs
+//   - wf-c (project p1): 1 initialized run
+//   - wf-deleted (project p1): 5 success runs, soft-deleted afterwards
+//
+// otherOrg:
+//   - wf-other (project p1): 4 success runs
+func (s *orgMetricsTopWorkflowsTestSuite) SetupTest() {
+	t := s.T()
+	ctx := context.Background()
+
+	credsWriter := creds.NewReaderWriter(t)
+	credsWriter.On("SaveCredentials", mock.Anything, mock.Anything, ociTestKeypair).Return("stored-OCI-secret", nil)
+	s.TestingUseCases = testhelpers.NewTestingUseCases(t, testhelpers.WithCredsReaderWriter(credsWriter))
+
+	var err error
+	s.org, err = s.Organization.CreateWithRandomName(ctx)
+	require.NoError(t, err)
+	s.otherOrg, err = s.Organization.CreateWithRandomName(ctx)
+	require.NoError(t, err)
+
+	wfA := s.createWorkflowWithRuns(s.org, "wf-a", "p1", biz.WorkflowRunSuccess, biz.WorkflowRunSuccess, biz.WorkflowRunError)
+	wfB := s.createWorkflowWithRuns(s.org, "wf-b", "p2", biz.WorkflowRunSuccess, biz.WorkflowRunSuccess)
+	s.createWorkflowWithRuns(s.org, "wf-c", "p1", biz.WorkflowRunInitialized)
+	wfDeleted := s.createWorkflowWithRuns(s.org, "wf-deleted", "p1",
+		biz.WorkflowRunSuccess, biz.WorkflowRunSuccess, biz.WorkflowRunSuccess, biz.WorkflowRunSuccess, biz.WorkflowRunSuccess)
+	s.createWorkflowWithRuns(s.otherOrg, "wf-other", "p1",
+		biz.WorkflowRunSuccess, biz.WorkflowRunSuccess, biz.WorkflowRunSuccess, biz.WorkflowRunSuccess)
+
+	require.NoError(t, s.Workflow.Delete(ctx, s.org.ID, wfDeleted.ID.String()))
+
+	s.p1, s.p2 = wfA.ProjectID, wfB.ProjectID
+}
+
+// createWorkflowWithRuns creates a workflow in the given org and project, and one run per given status
+func (s *orgMetricsTopWorkflowsTestSuite) createWorkflowWithRuns(org *biz.Organization, name, project string, statuses ...biz.WorkflowRunStatus) *biz.Workflow {
+	t := s.T()
+	ctx := context.Background()
+
+	wf, err := s.Workflow.Create(ctx, &biz.WorkflowCreateOpts{Name: name, Project: project, OrgID: org.ID})
+	require.NoError(t, err)
+
+	contractVersion, err := s.WorkflowContract.Describe(ctx, org.ID, wf.ContractID.String(), 0)
+	require.NoError(t, err)
+
+	casBackend, err := s.CASBackend.CreateOrUpdate(ctx, org.ID, ociTestRepo, ociTestUsername, ociTestPassword, backendType, true)
+	require.NoError(t, err)
+
+	for _, status := range statuses {
+		run, err := s.WorkflowRun.Create(ctx, &biz.WorkflowRunCreateOpts{
+			WorkflowID: wf.ID.String(), ContractRevision: contractVersion, CASBackendID: casBackend.ID, ProjectVersion: version1,
+		})
+		require.NoError(t, err)
+
+		if status != biz.WorkflowRunInitialized {
+			require.NoError(t, s.WorkflowRun.MarkAsFinished(ctx, run.ID.String(), status, ""))
+		}
+	}
+
+	return wf
+}
+
+func (s *orgMetricsTopWorkflowsTestSuite) TestTopWorkflowsByRunsCount() {
+	type want struct {
+		name     string
+		project  string
+		total    int32
+		byStatus map[string]int32
+	}
+
+	wantA := want{name: "wf-a", project: "p1", total: 3, byStatus: map[string]int32{"success": 2, "error": 1}}
+	wantB := want{name: "wf-b", project: "p2", total: 2, byStatus: map[string]int32{"success": 2}}
+	wantC := want{name: "wf-c", project: "p1", total: 1, byStatus: map[string]int32{"initialized": 1}}
+
+	currentWindow := &biz.TimeWindow{From: time.Now().Add(-time.Hour), To: time.Now().Add(time.Hour)}
+
+	testCases := []struct {
+		name         string
+		numWorkflows int
+		timeWindow   *biz.TimeWindow
+		projectIDs   func() []uuid.UUID
+		want         []want
+	}{
+		{
+			name:         "all visible workflows sorted by total, excluding deleted workflows and other orgs",
+			numWorkflows: 10,
+			want:         []want{wantA, wantB, wantC},
+		},
+		{
+			name:         "limited to the top N",
+			numWorkflows: 2,
+			want:         []want{wantA, wantB},
+		},
+		{
+			name:         "filtered by visible projects",
+			numWorkflows: 10,
+			projectIDs:   func() []uuid.UUID { return []uuid.UUID{s.p2} },
+			want:         []want{wantB},
+		},
+		{
+			name:         "no visible projects returns nothing",
+			numWorkflows: 10,
+			projectIDs:   func() []uuid.UUID { return []uuid.UUID{} },
+			want:         []want{},
+		},
+		{
+			name:         "time window without runs returns nothing",
+			numWorkflows: 10,
+			timeWindow:   &biz.TimeWindow{From: time.Now().Add(time.Hour), To: time.Now().Add(2 * time.Hour)},
+			want:         []want{},
+		},
+	}
+
+	for _, tc := range testCases {
+		s.Run(tc.name, func() {
+			tw := currentWindow
+			if tc.timeWindow != nil {
+				tw = tc.timeWindow
+			}
+
+			var projectIDs []uuid.UUID
+			if tc.projectIDs != nil {
+				projectIDs = tc.projectIDs()
+			}
+
+			got, err := s.OrgMetrics.TopWorkflowsByRunsCount(context.Background(), s.org.ID, tc.numWorkflows, tw, projectIDs)
+			s.Require().NoError(err)
+			s.Require().Len(got, len(tc.want))
+
+			for i, w := range tc.want {
+				s.Equal(w.name, got[i].Workflow.Name)
+				s.Equal(w.project, got[i].Workflow.Project)
+				s.Equal(w.total, got[i].Total)
+				s.Equal(w.byStatus, got[i].ByStatus)
+			}
+		})
+	}
 }
 
 type orgMetricsGetLastWorkflowStatusByRunTestSuite struct {
@@ -49,7 +209,7 @@ func (s *orgMetricsGetLastWorkflowStatusByRunTestSuite) SetupTest() {
 	assert := assert.New(s.T())
 
 	credsWriter := creds.NewReaderWriter(s.T())
-	credsWriter.On("SaveCredentials", mock.Anything, mock.Anything, &credentials.OCIKeypair{Repo: "repo", Username: "username", Password: "pass"}).Return("stored-OCI-secret", nil)
+	credsWriter.On("SaveCredentials", mock.Anything, mock.Anything, ociTestKeypair).Return("stored-OCI-secret", nil)
 
 	s.TestingUseCases = testhelpers.NewTestingUseCases(s.T(), testhelpers.WithCredsReaderWriter(credsWriter))
 
@@ -86,7 +246,7 @@ func (s *orgMetricsGetLastWorkflowStatusByRunTestSuite) TestGetLastWorkflowStatu
 		assert.NotNil(s.T(), contractVersion)
 
 		// Create a CAS backend
-		casBackend, err := s.TestingUseCases.CASBackend.CreateOrUpdate(ctx, s.org.ID, "repo", "username", "pass", backendType, true)
+		casBackend, err := s.CASBackend.CreateOrUpdate(ctx, s.org.ID, ociTestRepo, ociTestUsername, ociTestPassword, backendType, true)
 		assert.NoError(s.T(), err)
 		assert.NotNil(s.T(), casBackend)
 

app/controlplane/pkg/data/orgmetrics.go

@@ -16,14 +16,19 @@
 package data
 
 import (
+	"bytes"
+	"cmp"
 	"context"
 	"fmt"
+	"maps"
+	"slices"
 	"sort"
 	"time"
 
 	"entgo.io/ent/dialect/sql"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/biz"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/data/ent"
+	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/data/ent/predicate"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/data/ent/workflow"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/data/ent/workflowrun"
 	"github.com/chainloop-dev/chainloop/pkg/otelx"
@@ -146,79 +151,96 @@ func (repo *OrgMetricsRepo) TopWorkflowsByRunsCount(ctx context.Context, orgID u
 	ctx, span := otelx.Start(ctx, orgMetricsRepoTracer, "OrgMetricsRepo.TopWorkflowsByRunsCount")
 	defer span.End()
 
-	var runs []struct {
-		WorkflowID string `json:"workflow_id"`
-		State      string
-		Count      int32
+	// Filter on the denormalized workflow_runs.organization_id so the
+	// (organization_id, created_at) index drives the scan, instead of reaching
+	// the runs through the workflows table. Runs of deleted workflows are excluded.
+	wfPredicates := []predicate.Workflow{workflow.DeletedAtIsNil()}
+	if projectIDs != nil {
+		wfPredicates = append(wfPredicates, workflow.ProjectIDIn(projectIDs...))
 	}
 
-	// Get workflow runs grouped by state and workflowRunID
-	wfQuery := orgScopedQuery(repo.data.DB, orgID).
-		QueryWorkflows()
-
-	if projectIDs != nil {
-		wfQuery = wfQuery.Where(workflow.ProjectIDIn(projectIDs...))
+	var runs []struct {
+		WorkflowID uuid.UUID `json:"workflow_id"`
+		State      string    `json:"state"`
+		Count      int32     `json:"count"`
 	}
 
-	if err := wfQuery.QueryWorkflowruns().
+	if err := repo.data.DB.WorkflowRun.Query().
 		Where(
+			workflowrun.OrganizationID(orgID),
 			workflowrun.CreatedAtGTE(tw.From),
 			workflowrun.CreatedAtLTE(tw.To),
+			workflowrun.HasWorkflowWith(wfPredicates...),
 		).
-		GroupBy(workflowrun.WorkflowColumn, workflowrun.FieldState).
+		GroupBy(workflowrun.FieldWorkflowID, workflowrun.FieldState).
 		Aggregate(ent.Count()).
 		Scan(ctx, &runs); err != nil {
-		return nil, err
+		return nil, fmt.Errorf("aggregating workflow runs: %w", err)
 	}
 
-	// Map resultMap to include totals
-	resultMap := make(map[string]*biz.TopWorkflowsByRunsCountItem)
+	totals := make(map[uuid.UUID]*biz.TopWorkflowsByRunsCountItem)
 	for _, r := range runs {
-		var item *biz.TopWorkflowsByRunsCountItem
-		var found bool
-		if item, found = resultMap[r.WorkflowID]; !found {
-			workflowID, err := uuid.Parse(r.WorkflowID)
-			if err != nil {
-				return nil, err
-			}
-
-			wf, err := orgScopedQuery(repo.data.DB, orgID).QueryWorkflows().WithProject().Where(workflow.ID(workflowID)).First(ctx)
-			if err != nil {
-				if ent.IsNotFound(err) {
-					continue
-				}
-
-				return nil, err
-			}
-
-			wfRes, err := entWFToBizWF(ctx, wf)
-			if err != nil {
-				return nil, fmt.Errorf("converting entity: %w", err)
-			}
-
-			item = &biz.TopWorkflowsByRunsCountItem{ByStatus: make(map[string]int32), Workflow: wfRes}
+		item, ok := totals[r.WorkflowID]
+		if !ok {
+			item = &biz.TopWorkflowsByRunsCountItem{ByStatus: make(map[string]int32)}
+			totals[r.WorkflowID] = item
 		}
 
 		item.ByStatus[r.State] = r.Count
 		item.Total += r.Count
-		resultMap[r.WorkflowID] = item
 	}
 
-	result := make([]*biz.TopWorkflowsByRunsCountItem, 0, len(resultMap))
-	for _, r := range resultMap {
-		result = append(result, r)
+	// Pick the top workflows before loading them, sorted by total runs.
+	// Ties are broken by ID so the result is deterministic.
+	topIDs := slices.Collect(maps.Keys(totals))
+	slices.SortFunc(topIDs, func(a, b uuid.UUID) int {
+		if c := cmp.Compare(totals[b].Total, totals[a].Total); c != 0 {
+			return c
+		}
+		return bytes.Compare(a[:], b[:])
+	})
+
+	if len(topIDs) > numWorkflows {
+		topIDs = topIDs[:numWorkflows]
+	}
+
+	result := make([]*biz.TopWorkflowsByRunsCountItem, 0, len(topIDs))
+	if len(topIDs) == 0 {
+		return result, nil
 	}
 
-	// Sort and limit to numWorkflows
-	sort.SliceStable(result, func(i, j int) bool {
-		return result[i].Total > result[j].Total
-	})
+	// Load all the selected workflows and their projects in a single round-trip
+	wfs, err := repo.data.DB.Workflow.Query().
+		Where(workflow.IDIn(topIDs...), workflow.OrganizationID(orgID), workflow.DeletedAtIsNil()).
+		WithProject().
+		All(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("loading workflows: %w", err)
+	}
 
-	if len(result) < numWorkflows {
-		numWorkflows = len(result)
+	wfByID := make(map[uuid.UUID]*ent.Workflow, len(wfs))
+	for _, wf := range wfs {
+		wfByID[wf.ID] = wf
 	}
 
-	return result[0:numWorkflows], nil
+	for _, id := range topIDs {
+		wf, ok := wfByID[id]
+		if !ok {
+			// deleted between both queries
+			continue
+		}
+
+		wfRes, err := entWFToBizWF(ctx, wf)
+		if err != nil {
+			return nil, fmt.Errorf("converting entity: %w", err)
+		}
+
+		item := totals[id]
+		item.Workflow = wfRes
+		result = append(result, item)
+	}
+
+	return result, nil
 }
 
 func (repo *OrgMetricsRepo) DailyRunsCount(ctx context.Context, orgID, workflowID uuid.UUID, tw *biz.TimeWindow, projectIDs []uuid.UUID) ([]*biz.DayRunsCount, error) {