diff --git a/query/checks.go b/query/checks.go index 75c56668..323e4cbb 100644 --- a/query/checks.go +++ b/query/checks.go @@ -9,8 +9,8 @@ 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 } @@ -18,7 +18,7 @@ func FindChecks(ctx context.Context, resourceSelectors ...types.ResourceSelector 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 index bb1e6eed..fecf346e 100644 --- a/query/components.go +++ b/query/components.go @@ -21,8 +21,8 @@ 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 } @@ -30,15 +30,18 @@ func FindComponents(ctx context.Context, resourceSelectors ...types.ResourceSele 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 index 19193ff8..b2c6a4e0 100644 --- 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,16 +348,16 @@ 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 } @@ -365,16 +365,19 @@ func FindConfigsByResourceSelector(ctx context.Context, resourceSelectors ...typ 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 index e352b006..fc7cba0d 100644 --- 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 index cc39fa9d..fd8914bb 100644 --- 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 index 19ee46ff..8ccc6eeb 100644 --- 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 index ef74b775..31175b65 100644 --- 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))) + }) + } + }) +})