feat: support resource selector in FindChecks and added tests

flanksource · Jan 31, 2024 · a70e3dc · a70e3dc
1 parent fa46d79
commit a70e3dc
Show file tree

Hide file tree

Showing 3 changed files with 156 additions and 27 deletions.
diff --git a/getters.go b/getters.go
@@ -147,7 +147,7 @@ func FindChecks(ctx context.Context, resourceSelectors types.ResourceSelectors,
 		}
 	}
 
-	var uniqueChecks []models.Check
+	var allChecks []models.Check
 	for _, resourceSelector := range resourceSelectors {
 		hash := "FindChecks-CachePrefix" + resourceSelector.Hash()
 		cacheToUse := getterCache
@@ -161,22 +161,41 @@ func FindChecks(ctx context.Context, resourceSelectors types.ResourceSelectors,
 				return nil, err
 			}
 
-			uniqueChecks = append(uniqueChecks, checks...)
+			allChecks = append(allChecks, checks...)
 			continue
 		}
 
+		var uniqueChecks []models.Check
+		selectorOpts := opts
+
+		if query := firstResourceSelectorQuery(ctx, resourceSelector); query != nil {
+			var checks []models.Check
+			if err := apply(query, opts...).Find(&checks).Error; err != nil {
+				return nil, fmt.Errorf("error getting checks with selectors[%v]: %w", resourceSelector, err)
+			}
+
+			uniqueChecks = checks
+			selectorOpts = append(selectorOpts, WhereClause("id::text in ?", lo.Map(
+				checks,
+				func(c models.Check, _ int) string { return c.ID.String() }),
+			))
+		}
+
 		if resourceSelector.LabelSelector != "" {
-			checks, err := FindChecksByLabel(ctx, resourceSelector.LabelSelector, opts...)
+			checks, err := FindChecksByLabel(ctx, resourceSelector.LabelSelector, selectorOpts...)
 			if err != nil {
 				return nil, fmt.Errorf("error getting checks with label selectors[%s]: %w", resourceSelector.LabelSelector, err)
 			}
 
-			cacheToUse.SetDefault(hash, lo.Map(checks, func(c models.Check, _ int) string { return c.ID.String() }))
-			uniqueChecks = append(uniqueChecks, checks...)
+			uniqueChecks = checks
 		}
+
+		cacheToUse.SetDefault(hash, lo.Map(uniqueChecks, func(c models.Check, _ int) string { return c.ID.String() }))
+
+		allChecks = append(allChecks, uniqueChecks...)
 	}
 
-	return lo.UniqBy(uniqueChecks, models.CheckID), nil
+	return lo.UniqBy(allChecks, models.CheckID), nil
 }
 
 func FindComponents(ctx context.Context, resourceSelectors types.ResourceSelectors, opts ...FindOption) (components []models.Component, err error) {
@@ -201,32 +220,15 @@ func FindComponents(ctx context.Context, resourceSelectors types.ResourceSelecto
 		var uniqueComponents []models.Component
 		selectorOpts := opts
 
-		if resourceSelector.Name != "" || resourceSelector.Namespace != "" || resourceSelector.AgentID != "" || len(resourceSelector.Types) != 0 || len(resourceSelector.Statuses) != 0 {
-			query := ctx.DB()
-			if resourceSelector.Name != "" {
-				query = query.Where("name = ?", resourceSelector.Name)
-			}
-			if resourceSelector.Namespace != "" {
-				query = query.Where("namespace = ?", resourceSelector.Namespace)
-			}
-			if resourceSelector.AgentID != "" {
-				query = query.Where("agent_id = ?", resourceSelector.AgentID)
-			}
-			if len(resourceSelector.Types) != 0 {
-				query = query.Where("type IN ?", resourceSelector.Types)
-			}
-			if len(resourceSelector.Statuses) != 0 {
-				query = query.Where("status IN ?", resourceSelector.Statuses)
-			}
-
-			var nameComponents []models.Component
+		if query := firstResourceSelectorQuery(ctx, resourceSelector); query != nil {
+			var components []models.Component
 			if err := apply(query, opts...).Find(&components).Error; err != nil {
 				return nil, fmt.Errorf("error getting components with selectors[%v]: %w", resourceSelector, err)
 			}
 
-			uniqueComponents = nameComponents
+			uniqueComponents = components
 			selectorOpts = append(selectorOpts, WhereClause("id::text in ?", lo.Map(
-				nameComponents,
+				components,
 				func(c models.Component, _ int) string { return c.ID.String() }),
 			))
 		}
@@ -259,6 +261,33 @@ func FindComponents(ctx context.Context, resourceSelectors types.ResourceSelecto
 	return lo.UniqBy(allComponents, models.ComponentID), nil
 }
 
+// firstResourceSelectorQuery returns an ANDed query from all the fields except the
+// label selectors & field selectors.
+func firstResourceSelectorQuery(ctx DBContext, resourceSelector types.ResourceSelector) *gorm.DB {
+	if resourceSelector.Name == "" && resourceSelector.Namespace == "" && resourceSelector.AgentID == "" && len(resourceSelector.Types) == 0 && len(resourceSelector.Statuses) == 0 {
+		return nil
+	}
+
+	query := ctx.DB()
+	if resourceSelector.Name != "" {
+		query = query.Where("name = ?", resourceSelector.Name)
+	}
+	if resourceSelector.Namespace != "" {
+		query = query.Where("namespace = ?", resourceSelector.Namespace)
+	}
+	if resourceSelector.AgentID != "" {
+		query = query.Where("agent_id = ?", resourceSelector.AgentID)
+	}
+	if len(resourceSelector.Types) != 0 {
+		query = query.Where("type IN ?", resourceSelector.Types)
+	}
+	if len(resourceSelector.Statuses) != 0 {
+		query = query.Where("status IN ?", resourceSelector.Statuses)
+	}
+
+	return query
+}
+
 func FindComponentsByLabel(ctx context.Context, labelSelector string, opts ...FindOption) (components []models.Component, err error) {
 	if labelSelector == "" {
 		return nil, nil

diff --git a/tests/fixtures/dummy/checks.go b/tests/fixtures/dummy/checks.go
@@ -15,6 +15,9 @@ var LogisticsAPIHealthHTTPCheck = models.Check{
 	Name:     "logistics-api-health-check",
 	Type:     "http",
 	Status:   "healthy",
+	Labels: map[string]string{
+		"app": "logistics",
+	},
 }
 
 var LogisticsAPIHomeHTTPCheck = models.Check{
@@ -23,6 +26,9 @@ var LogisticsAPIHomeHTTPCheck = models.Check{
 	Name:     "logistics-api-home-check",
 	Type:     "http",
 	Status:   "healthy",
+	Labels: map[string]string{
+		"app": "logistics",
+	},
 }
 
 var LogisticsDBCheck = models.Check{
@@ -31,6 +37,9 @@ var LogisticsDBCheck = models.Check{
 	Name:     "logistics-db-check",
 	Type:     "postgres",
 	Status:   "unhealthy",
+	Labels: map[string]string{
+		"app": "logistics",
+	},
 }
 
 var CartAPIHeathCheckAgent = models.Check{

diff --git a/tests/getters_test.go b/tests/getters_test.go
@@ -8,6 +8,72 @@ import (
 	. "github.com/onsi/gomega"
 )
 
+var _ = ginkgo.Describe("FindChecks", func() {
+	type testRecord struct {
+		Name      string
+		Selectors []types.ResourceSelector
+		Results   int
+	}
+
+	testData := []testRecord{
+		{
+			Name:      "empty",
+			Selectors: []types.ResourceSelector{},
+			Results:   0,
+		},
+		{
+			Name:      "name",
+			Selectors: []types.ResourceSelector{{Name: dummy.LogisticsAPIHealthHTTPCheck.Name}},
+			Results:   1,
+		},
+		{
+			Name:      "names",
+			Selectors: []types.ResourceSelector{{Name: dummy.LogisticsAPIHealthHTTPCheck.Name}, {Name: dummy.LogisticsAPIHomeHTTPCheck.Name}, {Name: dummy.LogisticsDBCheck.Name}},
+			Results:   3,
+		},
+		{
+			Name:      "names but different namespace",
+			Selectors: []types.ResourceSelector{{Namespace: "kube-system", Name: dummy.LogisticsAPIHealthHTTPCheck.Name}, {Namespace: "kube-system", Name: dummy.LogisticsAPIHomeHTTPCheck.Name}},
+			Results:   0,
+		},
+		{
+			Name:      "types",
+			Selectors: []types.ResourceSelector{{Types: []string{dummy.LogisticsDBCheck.Type}}},
+			Results:   1,
+		},
+		{
+			Name:      "repeated (types) to test cache",
+			Selectors: []types.ResourceSelector{{Types: []string{dummy.LogisticsDBCheck.Type}}},
+			Results:   1,
+		},
+		{
+			Name:      "agentID",
+			Selectors: []types.ResourceSelector{{AgentID: dummy.CartAPIHeathCheckAgent.AgentID.String()}},
+			Results:   4,
+		},
+		{
+			Name:      "type & statuses",
+			Selectors: []types.ResourceSelector{{Types: []string{dummy.LogisticsDBCheck.Type}, Statuses: []string{string(dummy.LogisticsDBCheck.Status)}}},
+			Results:   1,
+		},
+		{
+			Name:      "label selector",
+			Selectors: []types.ResourceSelector{{LabelSelector: "app=logistics"}},
+			Results:   3,
+		},
+	}
+
+	for i := range testData {
+		td := testData[i]
+
+		ginkgo.It(td.Name, func() {
+			components, err := duty.FindChecks(DefaultContext, td.Selectors, duty.PickColumns("id"))
+			Expect(err).ToNot(HaveOccurred())
+			Expect(len(components)).To(Equal(td.Results))
+		})
+	}
+})
+
 var _ = ginkgo.Describe("FindComponent", func() {
 	type testRecord struct {
 		Name      string
@@ -26,6 +92,31 @@ var _ = ginkgo.Describe("FindComponent", func() {
 			Selectors: []types.ResourceSelector{{Name: dummy.Logistics.Name}, {Name: dummy.LogisticsAPI.Name}},
 			Results:   2,
 		},
+		{
+			Name:      "names but different namespace",
+			Selectors: []types.ResourceSelector{{Namespace: "kube-system", Name: dummy.Logistics.Name}, {Namespace: "kube-system", Name: dummy.LogisticsAPI.Name}},
+			Results:   0,
+		},
+		{
+			Name:      "types",
+			Selectors: []types.ResourceSelector{{Types: []string{dummy.Logistics.Type}}},
+			Results:   1,
+		},
+		{
+			Name:      "repeated (types) to test cache",
+			Selectors: []types.ResourceSelector{{Types: []string{dummy.Logistics.Type}}},
+			Results:   1,
+		},
+		{
+			Name:      "agentID",
+			Selectors: []types.ResourceSelector{{AgentID: dummy.PaymentsAPI.AgentID.String()}},
+			Results:   1,
+		},
+		{
+			Name:      "type & statuses",
+			Selectors: []types.ResourceSelector{{Types: []string{"KubernetesPod"}, Statuses: []string{string(types.ComponentStatusHealthy)}}},
+			Results:   3,
+		},
 		{
 			Name:      "empty",
 			Selectors: []types.ResourceSelector{},