feat: limit on resource selector search

flanksource · Sep 5, 2024 · e3655c3 · e3655c3
1 parent 809ff8c
commit e3655c3
Show file tree

Hide file tree

Showing 7 changed files with 113 additions and 27 deletions.
diff --git a/query/checks.go b/query/checks.go
@@ -9,16 +9,16 @@ import (
 	"github.com/google/uuid"
 )
 
-func FindChecks(ctx context.Context, resourceSelectors ...types.ResourceSelector) ([]models.Check, error) {
-	ids, err := FindCheckIDs(ctx, resourceSelectors...)
+func FindChecks(ctx context.Context, limit int, resourceSelectors ...types.ResourceSelector) ([]models.Check, error) {
+	ids, err := FindCheckIDs(ctx, limit, resourceSelectors...)
 	if err != nil {
 		return nil, err
 	}
 
 	return GetChecksByIDs(ctx, ids)
 }
 
-func FindCheckIDs(ctx context.Context, resourceSelectors ...types.ResourceSelector) ([]uuid.UUID, error) {
+func FindCheckIDs(ctx context.Context, limit int, resourceSelectors ...types.ResourceSelector) ([]uuid.UUID, error) {
 	for _, rs := range resourceSelectors {
 		if rs.FieldSelector != "" {
 			return nil, fmt.Errorf("field selector is not supported for checks (%s)", rs.FieldSelector)
@@ -27,12 +27,15 @@ func FindCheckIDs(ctx context.Context, resourceSelectors ...types.ResourceSelect
 
 	var allChecks []uuid.UUID
 	for _, resourceSelector := range resourceSelectors {
-		items, err := queryResourceSelector(ctx, resourceSelector, "checks", nil)
+		items, err := queryResourceSelector(ctx, limit, resourceSelector, "checks", nil)
 		if err != nil {
 			return nil, err
 		}
 
 		allChecks = append(allChecks, items...)
+		if limit > 0 && len(allChecks) >= limit {
+			return allChecks[:limit], nil
+		}
 	}
 
 	return allChecks, nil

diff --git a/query/components.go b/query/components.go
@@ -21,24 +21,27 @@ func GetComponentsByIDs(ctx context.Context, ids []uuid.UUID) ([]models.Componen
 	return components, nil
 }
 
-func FindComponents(ctx context.Context, resourceSelectors ...types.ResourceSelector) ([]models.Component, error) {
-	items, err := FindComponentIDs(ctx, resourceSelectors...)
+func FindComponents(ctx context.Context, limit int, resourceSelectors ...types.ResourceSelector) ([]models.Component, error) {
+	items, err := FindComponentIDs(ctx, limit, resourceSelectors...)
 	if err != nil {
 		return nil, err
 	}
 
 	return GetComponentsByIDs(ctx, items)
 }
 
-func FindComponentIDs(ctx context.Context, resourceSelectors ...types.ResourceSelector) ([]uuid.UUID, error) {
+func FindComponentIDs(ctx context.Context, limit int, resourceSelectors ...types.ResourceSelector) ([]uuid.UUID, error) {
 	var allComponents []uuid.UUID
 	for _, resourceSelector := range resourceSelectors {
-		items, err := queryResourceSelector(ctx, resourceSelector, "components", models.AllowedColumnFieldsInComponents)
+		items, err := queryResourceSelector(ctx, limit, resourceSelector, "components", models.AllowedColumnFieldsInComponents)
 		if err != nil {
 			return nil, err
 		}
 
 		allComponents = append(allComponents, items...)
+		if limit > 0 && len(allComponents) >= limit {
+			return allComponents[:limit], nil
+		}
 	}
 
 	return allComponents, nil

diff --git a/query/config.go b/query/config.go
@@ -336,7 +336,7 @@ func GetConfigsByIDs(ctx context.Context, ids []uuid.UUID) ([]models.ConfigItem,
 }
 
 func FindConfig(ctx context.Context, query types.ConfigQuery) (*models.ConfigItem, error) {
-	res, err := FindConfigsByResourceSelector(ctx, query.ToResourceSelector())
+	res, err := FindConfigsByResourceSelector(ctx, -1, query.ToResourceSelector())
 	if err != nil {
 		return nil, err
 	}
@@ -348,33 +348,36 @@ func FindConfig(ctx context.Context, query types.ConfigQuery) (*models.ConfigIte
 	return &res[0], nil
 }
 
-func FindConfigs(ctx context.Context, config types.ConfigQuery) ([]models.ConfigItem, error) {
-	return FindConfigsByResourceSelector(ctx, config.ToResourceSelector())
+func FindConfigs(ctx context.Context, limit int, config types.ConfigQuery) ([]models.ConfigItem, error) {
+	return FindConfigsByResourceSelector(ctx, limit, config.ToResourceSelector())
 }
 
-func FindConfigIDs(ctx context.Context, config types.ConfigQuery) ([]uuid.UUID, error) {
-	return FindConfigIDsByResourceSelector(ctx, config.ToResourceSelector())
+func FindConfigIDs(ctx context.Context, limit int, config types.ConfigQuery) ([]uuid.UUID, error) {
+	return FindConfigIDsByResourceSelector(ctx, limit, config.ToResourceSelector())
 }
 
-func FindConfigsByResourceSelector(ctx context.Context, resourceSelectors ...types.ResourceSelector) ([]models.ConfigItem, error) {
-	items, err := FindConfigIDsByResourceSelector(ctx, resourceSelectors...)
+func FindConfigsByResourceSelector(ctx context.Context, limit int, resourceSelectors ...types.ResourceSelector) ([]models.ConfigItem, error) {
+	items, err := FindConfigIDsByResourceSelector(ctx, limit, resourceSelectors...)
 	if err != nil {
 		return nil, err
 	}
 
 	return GetConfigsByIDs(ctx, items)
 }
 
-func FindConfigIDsByResourceSelector(ctx context.Context, resourceSelectors ...types.ResourceSelector) ([]uuid.UUID, error) {
+func FindConfigIDsByResourceSelector(ctx context.Context, limit int, resourceSelectors ...types.ResourceSelector) ([]uuid.UUID, error) {
 	var allConfigs []uuid.UUID
 
 	for _, resourceSelector := range resourceSelectors {
-		items, err := queryResourceSelector(ctx, resourceSelector, "config_items", models.AllowedColumnFieldsInConfigs)
+		items, err := queryResourceSelector(ctx, limit, resourceSelector, "config_items", models.AllowedColumnFieldsInConfigs)
 		if err != nil {
 			return nil, err
 		}
 
 		allConfigs = append(allConfigs, items...)
+		if limit > 0 && len(allConfigs) >= limit {
+			return allConfigs[:limit], nil
+		}
 	}
 
 	return allConfigs, nil

diff --git a/query/resource_selector.go b/query/resource_selector.go
@@ -23,6 +23,9 @@ import (
 )
 
 type SearchResourcesRequest struct {
+	// Limit the number of results returned per resource type
+	Limit int `json:"limit"`
+
 	Checks     []types.ResourceSelector `json:"checks"`
 	Components []types.ResourceSelector `json:"components"`
 	Configs    []types.ResourceSelector `json:"configs"`
@@ -55,9 +58,13 @@ type SelectedResource struct {
 func SearchResources(ctx context.Context, req SearchResourcesRequest) (*SearchResourcesResponse, error) {
 	var output SearchResourcesResponse
 
+	if req.Limit <= 0 {
+		req.Limit = 100
+	}
+
 	eg, _ := errgroup.WithContext(ctx)
 	eg.Go(func() error {
-		if items, err := FindConfigsByResourceSelector(ctx, req.Configs...); err != nil {
+		if items, err := FindConfigsByResourceSelector(ctx, req.Limit, req.Configs...); err != nil {
 			return err
 		} else {
 			for i := range items {
@@ -76,7 +83,7 @@ func SearchResources(ctx context.Context, req SearchResourcesRequest) (*SearchRe
 	})
 
 	eg.Go(func() error {
-		if items, err := FindChecks(ctx, req.Checks...); err != nil {
+		if items, err := FindChecks(ctx, req.Limit, req.Checks...); err != nil {
 			return err
 		} else {
 			for i := range items {
@@ -96,7 +103,7 @@ func SearchResources(ctx context.Context, req SearchResourcesRequest) (*SearchRe
 	})
 
 	eg.Go(func() error {
-		if items, err := FindComponents(ctx, req.Components...); err != nil {
+		if items, err := FindComponents(ctx, req.Limit, req.Components...); err != nil {
 			return err
 		} else {
 			for i := range items {
@@ -236,12 +243,12 @@ func SetResourceSelectorClause(ctx context.Context, resourceSelector types.Resou
 }
 
 // queryResourceSelector runs the given resourceSelector and returns the resource ids
-func queryResourceSelector(ctx context.Context, resourceSelector types.ResourceSelector, table string, allowedColumnsAsFields []string) ([]uuid.UUID, error) {
+func queryResourceSelector(ctx context.Context, limit int, resourceSelector types.ResourceSelector, table string, allowedColumnsAsFields []string) ([]uuid.UUID, error) {
 	if resourceSelector.IsEmpty() {
 		return nil, nil
 	}
 
-	hash := fmt.Sprintf("%s-%s", table, resourceSelector.Hash())
+	hash := fmt.Sprintf("%s-%s-%d", table, resourceSelector.Hash(), limit)
 	cacheToUse := getterCache
 	if resourceSelector.Immutable() {
 		cacheToUse = immutableCache
@@ -254,6 +261,10 @@ func queryResourceSelector(ctx context.Context, resourceSelector types.ResourceS
 	}
 
 	query := ctx.DB().Select("id").Table(table)
+	if limit > 0 {
+		query = query.Limit(limit)
+	}
+
 	query, err := SetResourceSelectorClause(ctx, resourceSelector, query, table, allowedColumnsAsFields)
 	if err != nil {
 		return nil, err

diff --git a/relationship_selector.go b/relationship_selector.go
@@ -208,7 +208,7 @@ func LookupComponents(ctx context.Context, lookup RelationshipSelectorTemplate,
 		logger.Tracef("finding all components (%s)", lookupResult)
 	}
 
-	return query.FindComponentIDs(ctx, lookupResult.ToResourceSelector())
+	return query.FindComponentIDs(ctx, 0, lookupResult.ToResourceSelector())
 }
 
 func LookupConfigs(ctx context.Context, lookup RelationshipSelectorTemplate, labels map[string]string, env map[string]any) ([]uuid.UUID, error) {
@@ -223,5 +223,5 @@ func LookupConfigs(ctx context.Context, lookup RelationshipSelectorTemplate, lab
 		logger.Tracef("finding all config items (%s)", lookupResult)
 	}
 
-	return query.FindConfigIDsByResourceSelector(ctx, lookupResult.ToResourceSelector())
+	return query.FindConfigIDsByResourceSelector(ctx, 0, lookupResult.ToResourceSelector())
 }
diff --git a/tests/getters_test.go b/tests/getters_test.go
@@ -69,7 +69,7 @@ var _ = ginkgo.Describe("FindChecks", func() {
 		td := testData[i]
 
 		ginkgo.It(td.Name, func() {
-			components, err := query.FindCheckIDs(DefaultContext, td.Selectors...)
+			components, err := query.FindCheckIDs(DefaultContext, 0, td.Selectors...)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(len(components)).To(Equal(td.Results))
 		})
@@ -123,7 +123,7 @@ var _ = ginkgo.Describe("FindConfigs", func() {
 		td := testData[i]
 
 		ginkgo.It(td.Name, func() {
-			components, err := query.FindConfigIDsByResourceSelector(DefaultContext, td.Selectors...)
+			components, err := query.FindConfigIDsByResourceSelector(DefaultContext, 0, td.Selectors...)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(components).To(ContainElements(testData[i].Results))
 		})
@@ -194,7 +194,7 @@ var _ = ginkgo.Describe("FindComponent", func() {
 		td := testData[i]
 
 		ginkgo.It(td.Name, func() {
-			components, err := query.FindComponentIDs(DefaultContext, td.Selectors...)
+			components, err := query.FindComponentIDs(DefaultContext, 0, td.Selectors...)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(len(components)).To(Equal(td.Results))
 		})

diff --git a/tests/query_resource_selector_test.go b/tests/query_resource_selector_test.go
@@ -215,3 +215,69 @@ var _ = ginkgo.Describe("SearchResourceSelectors", func() {
 		})
 	})
 })
+
+var _ = ginkgo.Describe("Resoure Selector limits", ginkgo.Ordered, func() {
+	ginkgo.BeforeAll(func() {
+		_ = query.SyncConfigCache(DefaultContext)
+	})
+
+	ginkgo.Context("It should return the fixed page size for configs", func() {
+		for limit := 1; limit < 3; limit++ {
+			ginkgo.It(fmt.Sprintf("should work with %d page size", limit), func() {
+				items, err := query.SearchResources(DefaultContext, query.SearchResourcesRequest{
+					Limit:   limit,
+					Configs: []types.ResourceSelector{{FieldSelector: fmt.Sprintf("config_class=%s", models.ConfigClassNode)}},
+				})
+
+				Expect(err).To(BeNil())
+				Expect(limit).To(Equal(len(items.Configs)))
+			})
+		}
+	})
+
+	ginkgo.Context("It should return the fixed page size for components", func() {
+		for limit := 1; limit < 5; limit++ {
+			ginkgo.It(fmt.Sprintf("should work with %d page size", limit), func() {
+				items, err := query.SearchResources(DefaultContext, query.SearchResourcesRequest{
+					Limit:      limit,
+					Components: []types.ResourceSelector{{Types: []string{"Application"}}},
+				})
+
+				Expect(err).To(BeNil())
+				Expect(limit).To(Equal(len(items.Components)))
+			})
+		}
+	})
+
+	ginkgo.Context("It should return the fixed page size for checks", func() {
+		for limit := 1; limit < 3; limit++ {
+			ginkgo.It(fmt.Sprintf("should work with %d page size", limit), func() {
+				items, err := query.SearchResources(DefaultContext, query.SearchResourcesRequest{
+					Limit:  limit,
+					Checks: []types.ResourceSelector{{Types: []string{"http"}, Agent: "all"}},
+				})
+
+				Expect(err).To(BeNil())
+				Expect(limit).To(Equal(len(items.Checks)))
+			})
+		}
+	})
+
+	ginkgo.Context("It should return the fixed page size for all types", func() {
+		for pageSize := 1; pageSize < 3; pageSize++ {
+			ginkgo.It(fmt.Sprintf("should work with %d page size", pageSize), func() {
+				items, err := query.SearchResources(DefaultContext, query.SearchResourcesRequest{
+					Limit:      pageSize,
+					Configs:    []types.ResourceSelector{{FieldSelector: fmt.Sprintf("config_class=%s", models.ConfigClassNode)}},
+					Components: []types.ResourceSelector{{Types: []string{"Application"}}},
+					Checks:     []types.ResourceSelector{{Types: []string{"http"}, Agent: "all"}},
+				})
+
+				Expect(err).To(BeNil())
+				Expect(pageSize).To(Equal(len(items.Configs)))
+				Expect(pageSize).To(Equal(len(items.Components)))
+				Expect(pageSize).To(Equal(len(items.Checks)))
+			})
+		}
+	})
+})