diff --git a/lib/services/local/generic/generic.go b/lib/services/local/generic/generic.go index 8b64a5f02d036..72a10b8b71c0b 100644 --- a/lib/services/local/generic/generic.go +++ b/lib/services/local/generic/generic.go @@ -78,7 +78,6 @@ func (c *ServiceConfig[T]) CheckAndSetDefaults() error { if c.UnmarshalFunc == nil { return trace.BadParameter("unmarshal func is missing") } - return nil } @@ -167,17 +166,17 @@ func (s *Service[T]) GetResources(ctx context.Context) ([]T, error) { // ListResources returns a paginated list of resources. func (s *Service[T]) ListResources(ctx context.Context, pageSize int, pageToken string) ([]T, string, error) { - resources, next, err := s.ListResourcesReturnNextResource(ctx, pageSize, pageToken) - var nextKey string - if next != nil { - nextKey = backend.GetPaginationKey(*next) - } + resources, _, nextKey, err := s.listResourcesReturnNextResourceWithKey(ctx, pageSize, pageToken) return resources, nextKey, trace.Wrap(err) } // ListResourcesReturnNextResource returns a paginated list of resources. The next resource is returned, which allows consumers to construct // the next pagination key as appropriate. func (s *Service[T]) ListResourcesReturnNextResource(ctx context.Context, pageSize int, pageToken string) ([]T, *T, error) { + resources, next, _, err := s.listResourcesReturnNextResourceWithKey(ctx, pageSize, pageToken) + return resources, next, trace.Wrap(err) +} +func (s *Service[T]) listResourcesReturnNextResourceWithKey(ctx context.Context, pageSize int, pageToken string) ([]T, *T, string, error) { rangeStart := backend.Key(s.backendPrefix, pageToken) rangeEnd := backend.RangeEnd(backend.ExactKey(s.backendPrefix)) @@ -191,26 +190,32 @@ func (s *Service[T]) ListResourcesReturnNextResource(ctx context.Context, pageSi // no filter provided get the range directly result, err := s.backend.GetRange(ctx, rangeStart, rangeEnd, limit) if err != nil { - return nil, nil, trace.Wrap(err) + return nil, nil, "", trace.Wrap(err) } out := make([]T, 0, len(result.Items)) + var lastKey []byte for _, item := range result.Items { resource, err := s.unmarshalFunc(item.Value, services.WithRevision(item.Revision), services.WithResourceID(item.ID)) if err != nil { - return nil, nil, trace.Wrap(err) + return nil, nil, "", trace.Wrap(err) } out = append(out, resource) + lastKey = item.Key } - var next *T + var ( + next *T + nextKey string + ) if len(out) > pageSize { next = &out[pageSize] // Truncate the last item that was used to determine next row existence. out = out[:pageSize] + nextKey = trimLastKey(string(lastKey), s.backendPrefix) } - return out, next, nil + return out, next, nextKey, nil } // ListResourcesWithFilter returns a paginated list of resources that match the given filter. @@ -226,6 +231,7 @@ func (s *Service[T]) ListResourcesWithFilter(ctx context.Context, pageSize int, limit := pageSize + 1 var resources []T + var lastKey []byte if err := backend.IterateRange( ctx, s.backend, @@ -239,6 +245,7 @@ func (s *Service[T]) ListResourcesWithFilter(ctx context.Context, pageSize int, return false, trace.Wrap(err) } if matcher(resource) { + lastKey = item.Key resources = append(resources, resource) } if len(resources) == pageSize { @@ -252,7 +259,7 @@ func (s *Service[T]) ListResourcesWithFilter(ctx context.Context, pageSize int, var nextKey string if len(resources) > pageSize { - nextKey = backend.GetPaginationKey(resources[pageSize]) + nextKey = trimLastKey(string(lastKey), s.backendPrefix) // Truncate the last item that was used to determine next row existence. resources = resources[:pageSize] } @@ -460,3 +467,13 @@ func (s *Service[T]) RunWhileLocked(ctx context.Context, lockName string, ttl ti return fn(ctx, s.backend) })) } + +func trimLastKey(lastKey, prefix string) string { + if !strings.HasSuffix(prefix, string(backend.Separator)) { + prefix += string(backend.Separator) + } + if !strings.HasPrefix(prefix, string(backend.Separator)) { + prefix = string(backend.Separator) + prefix + } + return strings.TrimPrefix(lastKey, prefix) +} diff --git a/lib/services/local/generic/generic_test.go b/lib/services/local/generic/generic_test.go index 2efc4dcc3bbe6..da3c9b66c9279 100644 --- a/lib/services/local/generic/generic_test.go +++ b/lib/services/local/generic/generic_test.go @@ -333,6 +333,65 @@ func TestGenericListResourcesReturnNextResource(t *testing.T) { require.Nil(t, next) } +func TestGenericListResourcesWithMultiplePrefixes(t *testing.T) { + ctx := context.Background() + + memBackend, err := memory.New(memory.Config{ + Context: ctx, + Clock: clockwork.NewFakeClock(), + }) + require.NoError(t, err) + + service, err := NewService(&ServiceConfig[*testResource]{ + Backend: memBackend, + ResourceKind: "generic resource", + PageLimit: 200, + BackendPrefix: "generic_prefix", + UnmarshalFunc: unmarshalResource, + MarshalFunc: marshalResource, + }) + require.NoError(t, err) + + // Create a couple test resources. + r1 := newTestResource("r1") + r2 := newTestResource("r2") + + _, err = service.WithPrefix("a-unique-prefix").UpsertResource(ctx, r1) + require.NoError(t, err) + _, err = service.WithPrefix("another-unique-prefix").UpsertResource(ctx, r2) + require.NoError(t, err) + + page, next, err := service.ListResources(ctx, 1, "") + require.NoError(t, err) + require.Empty(t, cmp.Diff([]*testResource{r1}, page, + cmpopts.IgnoreFields(types.Metadata{}, "Revision"), + )) + require.Equal(t, next, "another-unique-prefix"+string(backend.Separator)+r2.GetName()) + + page, next, err = service.ListResources(ctx, 1, next) + require.NoError(t, err) + require.Empty(t, cmp.Diff([]*testResource{r2}, page, + cmpopts.IgnoreFields(types.Metadata{}, "Revision"), + )) + require.Empty(t, next) + + _, err = service.WithPrefix("a-unique-prefix").UpsertResource(ctx, r2) + require.NoError(t, err) + page, next, err = service.WithPrefix("a-unique-prefix").ListResources(ctx, 1, "") + require.NoError(t, err) + require.Empty(t, cmp.Diff([]*testResource{r1}, page, + cmpopts.IgnoreFields(types.Metadata{}, "Revision"), + )) + require.Equal(t, next, r2.GetName()) + + page, next, err = service.WithPrefix("a-unique-prefix").ListResources(ctx, 1, next) + require.NoError(t, err) + require.Empty(t, cmp.Diff([]*testResource{r2}, page, + cmpopts.IgnoreFields(types.Metadata{}, "Revision"), + )) + require.Empty(t, next) +} + func TestGenericListResourcesWithFilter(t *testing.T) { ctx := context.Background() diff --git a/lib/services/local/generic/generic_wrapper.go b/lib/services/local/generic/generic_wrapper.go index 3fa6fc1c1c418..6ccb9aafb40de 100644 --- a/lib/services/local/generic/generic_wrapper.go +++ b/lib/services/local/generic/generic_wrapper.go @@ -33,8 +33,8 @@ func NewServiceWrapper[T types.ResourceMetadata]( resourceKind string, backendPrefix string, marshalFunc MarshalFunc[T], - unmarshalFunc UnmarshalFunc[T]) (*ServiceWrapper[T], error) { - + unmarshalFunc UnmarshalFunc[T], +) (*ServiceWrapper[T], error) { cfg := &ServiceConfig[resourceMetadataAdapter[T]]{ Backend: backend, ResourceKind: resourceKind,