Extend the generic storage service to support custom names for names (#52826)

* Extend the generic storage service to support custom names for names

* Drop the custom per-resource name as it's unused

Author: espadolini

lib/services/local/access_graph.go

@@ -48,7 +48,7 @@ type AccessGraphSecretsService struct {
 // SSH Keys. Future implementations might extend them.
 func NewAccessGraphSecretsService(b backend.Backend) (*AccessGraphSecretsService, error) {
 	authorizedKeysSvc, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*accessgraphsecretspb.AuthorizedKey]{
+		generic.ServiceConfig[*accessgraphsecretspb.AuthorizedKey]{
 			Backend:       b,
 			ResourceKind:  types.KindAccessGraphSecretAuthorizedKey,
 			BackendPrefix: backend.NewKey(authorizedKeysPrefix),
@@ -60,7 +60,7 @@ func NewAccessGraphSecretsService(b backend.Backend) (*AccessGraphSecretsService
 	}
 
 	privateKeysSvc, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*accessgraphsecretspb.PrivateKey]{
+		generic.ServiceConfig[*accessgraphsecretspb.PrivateKey]{
 			Backend:       b,
 			ResourceKind:  types.KindAccessGraphSecretPrivateKey,
 			BackendPrefix: backend.NewKey(privateKeysPrefix),

lib/services/local/access_list_test.go

@@ -407,7 +407,7 @@ func TestAccessListDedupeOwnersBackwardsCompat(t *testing.T) {
 	accessListDuplicateOwners.Spec.Owners = append(accessListDuplicateOwners.Spec.Owners, accessListDuplicateOwners.Spec.Owners[0])
 	require.Len(t, accessListDuplicateOwners.Spec.Owners, 3)
 
-	item, err := service.service.MakeBackendItem(accessListDuplicateOwners, accessListDuplicateOwners.GetName())
+	item, err := service.service.MakeBackendItem(accessListDuplicateOwners)
 	require.NoError(t, err)
 	_, err = mem.Put(ctx, item)
 	require.NoError(t, err)
@@ -490,7 +490,6 @@ func TestAccessListUpsertWithMembers(t *testing.T) {
 		require.NoError(t, err)
 		require.Empty(t, members)
 	})
-
 }
 
 func TestAccessListMembersCRUD(t *testing.T) {

lib/services/local/access_monitoring_rules.go

@@ -41,7 +41,7 @@ type AccessMonitoringRulesService struct {
 // NewAccessMonitoringRulesService creates a new AccessMonitoringRulesService.
 func NewAccessMonitoringRulesService(b backend.Backend) (*AccessMonitoringRulesService, error) {
 	service, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*accessmonitoringrulesv1.AccessMonitoringRule]{
+		generic.ServiceConfig[*accessmonitoringrulesv1.AccessMonitoringRule]{
 			Backend:       b,
 			ResourceKind:  types.KindAccessMonitoringRule,
 			BackendPrefix: backend.NewKey(accessMonitoringRulesPrefix),

lib/services/local/autoupdate.go

@@ -47,44 +47,44 @@ type AutoUpdateService struct {
 // NewAutoUpdateService returns a new AutoUpdateService.
 func NewAutoUpdateService(b backend.Backend) (*AutoUpdateService, error) {
 	config, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*autoupdate.AutoUpdateConfig]{
+		generic.ServiceConfig[*autoupdate.AutoUpdateConfig]{
 			Backend:       b,
 			ResourceKind:  types.KindAutoUpdateConfig,
 			BackendPrefix: backend.NewKey(autoUpdateConfigPrefix),
 			MarshalFunc:   services.MarshalProtoResource[*autoupdate.AutoUpdateConfig],
 			UnmarshalFunc: services.UnmarshalProtoResource[*autoupdate.AutoUpdateConfig],
 			ValidateFunc:  update.ValidateAutoUpdateConfig,
-			KeyFunc: func(*autoupdate.AutoUpdateConfig) string {
+			NameKeyFunc: func(string) string {
 				return types.MetaNameAutoUpdateConfig
 			},
 		})
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 	version, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*autoupdate.AutoUpdateVersion]{
+		generic.ServiceConfig[*autoupdate.AutoUpdateVersion]{
 			Backend:       b,
 			ResourceKind:  types.KindAutoUpdateVersion,
 			BackendPrefix: backend.NewKey(autoUpdateVersionPrefix),
 			MarshalFunc:   services.MarshalProtoResource[*autoupdate.AutoUpdateVersion],
 			UnmarshalFunc: services.UnmarshalProtoResource[*autoupdate.AutoUpdateVersion],
 			ValidateFunc:  update.ValidateAutoUpdateVersion,
-			KeyFunc: func(version *autoupdate.AutoUpdateVersion) string {
+			NameKeyFunc: func(string) string {
 				return types.MetaNameAutoUpdateVersion
 			},
 		})
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 	rollout, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*autoupdate.AutoUpdateAgentRollout]{
+		generic.ServiceConfig[*autoupdate.AutoUpdateAgentRollout]{
 			Backend:       b,
 			ResourceKind:  types.KindAutoUpdateAgentRollout,
 			BackendPrefix: backend.NewKey(autoUpdateAgentRolloutPrefix),
 			MarshalFunc:   services.MarshalProtoResource[*autoupdate.AutoUpdateAgentRollout],
 			UnmarshalFunc: services.UnmarshalProtoResource[*autoupdate.AutoUpdateAgentRollout],
 			ValidateFunc:  update.ValidateAutoUpdateAgentRollout,
-			KeyFunc: func(_ *autoupdate.AutoUpdateAgentRollout) string {
+			NameKeyFunc: func(string) string {
 				return types.MetaNameAutoUpdateAgentRollout
 			},
 		})
@@ -222,7 +222,7 @@ func itemFromAutoUpdateConfig(config *autoupdate.AutoUpdateConfig) (*backend.Ite
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
-	value, err := services.MarshalProtoResource[*autoupdate.AutoUpdateConfig](config)
+	value, err := services.MarshalProtoResource(config)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -248,7 +248,7 @@ func itemFromAutoUpdateVersion(version *autoupdate.AutoUpdateVersion) (*backend.
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
-	value, err := services.MarshalProtoResource[*autoupdate.AutoUpdateVersion](version)
+	value, err := services.MarshalProtoResource(version)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}

lib/services/local/bot_instance.go

@@ -45,7 +45,7 @@ type BotInstanceService struct {
 // NewBotInstanceService creates a new BotInstanceService with the given backend.
 func NewBotInstanceService(b backend.Backend, clock clockwork.Clock) (*BotInstanceService, error) {
 	service, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*machineidv1.BotInstance]{
+		generic.ServiceConfig[*machineidv1.BotInstance]{
 			Backend:       b,
 			ResourceKind:  types.KindBotInstance,
 			BackendPrefix: backend.NewKey(botInstancePrefix),

lib/services/local/crown_jewels.go

@@ -39,7 +39,7 @@ const crownJewelsKey = "crown_jewels"
 // NewCrownJewelsService creates a new CrownJewelsService.
 func NewCrownJewelsService(b backend.Backend) (*CrownJewelsService, error) {
 	service, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*crownjewelv1.CrownJewel]{
+		generic.ServiceConfig[*crownjewelv1.CrownJewel]{
 			Backend:       b,
 			ResourceKind:  types.KindCrownJewel,
 			BackendPrefix: backend.NewKey(crownJewelsKey),

lib/services/local/databaseobject.go

@@ -74,7 +74,7 @@ const (
 
 func NewDatabaseObjectService(b backend.Backend) (*DatabaseObjectService, error) {
 	service, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*dbobjectv1.DatabaseObject]{
+		generic.ServiceConfig[*dbobjectv1.DatabaseObject]{
 			Backend:       b,
 			ResourceKind:  types.KindDatabaseObject,
 			BackendPrefix: backend.NewKey(databaseObjectPrefix),

lib/services/local/databaseobjectimportrule.go

@@ -72,7 +72,7 @@ const (
 
 func NewDatabaseObjectImportRuleService(b backend.Backend) (services.DatabaseObjectImportRules, error) {
 	service, err := generic.NewServiceWrapper(
-		generic.ServiceWrapperConfig[*databaseobjectimportrulev1.DatabaseObjectImportRule]{
+		generic.ServiceConfig[*databaseobjectimportrulev1.DatabaseObjectImportRule]{
 			Backend:       b,
 			ResourceKind:  types.KindDatabaseObjectImportRule,
 			BackendPrefix: backend.NewKey(databaseObjectImportRulePrefix),

lib/services/local/generic/generic.go

@@ -44,7 +44,7 @@ type MarshalFunc[T any] func(T, ...services.MarshalOption) ([]byte, error)
 type UnmarshalFunc[T any] func([]byte, ...services.MarshalOption) (T, error)
 
 // ServiceConfig is the configuration for the service configuration.
-type ServiceConfig[T Resource] struct {
+type ServiceConfig[T any] struct {
 	// Backend used to persist the resource.
 	Backend backend.Backend
 	// ResourceKind is the friendly name of the resource.
@@ -64,9 +64,10 @@ type ServiceConfig[T Resource] struct {
 	// If set to 0, the default interval of 250ms will be used.
 	// WARNING: If set to a negative value, the RunWhileLocked function will retry immediately.
 	RunWhileLockedRetryInterval time.Duration
-	// KeyFunc optionally allows resource to have a custom key. If not provided the
-	// name of the resource will be used.
-	KeyFunc func(T) string
+	// NameKeyFunc optionally allows resources to have a custom key suffix, by
+	// transforming the name of the resource or the input given to methods that
+	// take a resource name. If unset, the name is used without changes.
+	NameKeyFunc func(name string) string
 }
 
 func (c *ServiceConfig[T]) CheckAndSetDefaults() error {
@@ -95,10 +96,6 @@ func (c *ServiceConfig[T]) CheckAndSetDefaults() error {
 		c.ValidateFunc = func(t T) error { return nil }
 	}
 
-	if c.KeyFunc == nil {
-		c.KeyFunc = func(t T) string { return t.GetName() }
-	}
-
 	return nil
 }
 
@@ -112,7 +109,7 @@ type Service[T Resource] struct {
 	unmarshalFunc               UnmarshalFunc[T]
 	validateFunc                func(T) error
 	runWhileLockedRetryInterval time.Duration
-	keyFunc                     func(T) string
+	nameKeyFunc                 func(name string) string
 }
 
 // NewService will return a new generic service with the given config. This will
@@ -131,7 +128,7 @@ func NewService[T Resource](cfg *ServiceConfig[T]) (*Service[T], error) {
 		unmarshalFunc:               cfg.UnmarshalFunc,
 		validateFunc:                cfg.ValidateFunc,
 		runWhileLockedRetryInterval: cfg.RunWhileLockedRetryInterval,
-		keyFunc:                     cfg.KeyFunc,
+		nameKeyFunc:                 cfg.NameKeyFunc,
 	}, nil
 }
 
@@ -140,18 +137,16 @@ func (s *Service[T]) WithPrefix(parts ...string) *Service[T] {
 	if len(parts) == 0 {
 		return s
 	}
+	s2 := *s
+	s2.backendPrefix = s2.backendPrefix.AppendKey(backend.NewKey(parts...))
+	return &s2
+}
 
-	return &Service[T]{
-		backend:                     s.backend,
-		resourceKind:                s.resourceKind,
-		pageLimit:                   s.pageLimit,
-		backendPrefix:               s.backendPrefix.AppendKey(backend.NewKey(parts...)),
-		marshalFunc:                 s.marshalFunc,
-		unmarshalFunc:               s.unmarshalFunc,
-		validateFunc:                s.validateFunc,
-		runWhileLockedRetryInterval: s.runWhileLockedRetryInterval,
-		keyFunc:                     s.keyFunc,
+func (s *Service[T]) nameKey(name string) string {
+	if s.nameKeyFunc != nil {
+		return s.nameKeyFunc(name)
 	}
+	return name
 }
 
 // CountResources will return a count of all resources in the prefix range.
@@ -204,6 +199,7 @@ func (s *Service[T]) ListResourcesReturnNextResource(ctx context.Context, pageSi
 	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 := s.backendPrefix.AppendKey(backend.KeyFromString(pageToken))
 	rangeEnd := backend.RangeEnd(s.backendPrefix.ExactKey())
@@ -291,12 +287,11 @@ func (s *Service[T]) ListResourcesWithFilter(ctx context.Context, pageSize int,
 	}
 
 	return resources, nextKey, nil
-
 }
 
 // GetResource returns the specified resource.
 func (s *Service[T]) GetResource(ctx context.Context, name string) (resource T, err error) {
-	item, err := s.backend.Get(ctx, s.MakeKey(backend.NewKey(name)))
+	item, err := s.backend.Get(ctx, s.MakeKey(backend.NewKey(s.nameKey(name))))
 	if err != nil {
 		if trace.IsNotFound(err) {
 			return resource, trace.NotFound("%s %q doesn't exist", s.resourceKind, name)
@@ -315,7 +310,7 @@ func (s *Service[T]) CreateResource(ctx context.Context, resource T) (T, error)
 		return t, trace.Wrap(err)
 	}
 
-	item, err := s.MakeBackendItem(resource, s.keyFunc(resource))
+	item, err := s.MakeBackendItem(resource)
 	if err != nil {
 		return t, trace.Wrap(err)
 	}
@@ -340,7 +335,7 @@ func (s *Service[T]) UpdateResource(ctx context.Context, resource T) (T, error)
 		return t, trace.Wrap(err)
 	}
 
-	item, err := s.MakeBackendItem(resource, s.keyFunc(resource))
+	item, err := s.MakeBackendItem(resource)
 	if err != nil {
 		return t, trace.Wrap(err)
 	}
@@ -365,7 +360,7 @@ func (s *Service[T]) ConditionalUpdateResource(ctx context.Context, resource T)
 		return t, trace.Wrap(err)
 	}
 
-	item, err := s.MakeBackendItem(resource, s.keyFunc(resource))
+	item, err := s.MakeBackendItem(resource)
 	if err != nil {
 		return t, trace.Wrap(err)
 	}
@@ -390,7 +385,7 @@ func (s *Service[T]) UpsertResource(ctx context.Context, resource T) (T, error)
 		return t, trace.Wrap(err)
 	}
 
-	item, err := s.MakeBackendItem(resource, s.keyFunc(resource))
+	item, err := s.MakeBackendItem(resource)
 	if err != nil {
 		return t, trace.Wrap(err)
 	}
@@ -406,7 +401,7 @@ func (s *Service[T]) UpsertResource(ctx context.Context, resource T) (T, error)
 
 // DeleteResource removes the specified resource.
 func (s *Service[T]) DeleteResource(ctx context.Context, name string) error {
-	err := s.backend.Delete(ctx, s.MakeKey(backend.NewKey(name)))
+	err := s.backend.Delete(ctx, s.MakeKey(backend.NewKey(s.nameKey(name))))
 	if err != nil {
 		if trace.IsNotFound(err) {
 			return trace.NotFound("%s %q doesn't exist", s.resourceKind, name)
@@ -425,7 +420,7 @@ func (s *Service[T]) DeleteAllResources(ctx context.Context) error {
 // UpdateAndSwapResource will get the resource from the backend, modify it, and swap the new value into the backend.
 func (s *Service[T]) UpdateAndSwapResource(ctx context.Context, name string, modify func(T) error) (T, error) {
 	var t T
-	existingItem, err := s.backend.Get(ctx, s.MakeKey(backend.NewKey(name)))
+	existingItem, err := s.backend.Get(ctx, s.MakeKey(backend.NewKey(s.nameKey(name))))
 	if err != nil {
 		if trace.IsNotFound(err) {
 			return t, trace.NotFound("%s %q doesn't exist", s.resourceKind, name)
@@ -447,7 +442,7 @@ func (s *Service[T]) UpdateAndSwapResource(ctx context.Context, name string, mod
 		return t, trace.Wrap(err)
 	}
 
-	replacementItem, err := s.MakeBackendItem(resource, name)
+	replacementItem, err := s.MakeBackendItem(resource)
 	if err != nil {
 		return t, trace.Wrap(err)
 	}
@@ -462,7 +457,8 @@ func (s *Service[T]) UpdateAndSwapResource(ctx context.Context, name string, mod
 }
 
 // MakeBackendItem will check and make the backend item.
-func (s *Service[T]) MakeBackendItem(resource T, name string) (backend.Item, error) {
+func (s *Service[T]) MakeBackendItem(resource T, _ ...any) (backend.Item, error) {
+	// TODO(espadolini): clean up unused variadic after teleport.e is updated
 	if err := services.CheckAndSetDefaults(resource); err != nil {
 		return backend.Item{}, trace.Wrap(err)
 	}
@@ -477,7 +473,7 @@ func (s *Service[T]) MakeBackendItem(resource T, name string) (backend.Item, err
 		return backend.Item{}, trace.Wrap(err)
 	}
 	item := backend.Item{
-		Key:      s.MakeKey(backend.NewKey(name)),
+		Key:      s.MakeKey(backend.NewKey(s.nameKey(resource.GetName()))),
 		Value:    value,
 		Revision: rev,
 	}

lib/services/local/generic/generic_test.go

@@ -540,7 +540,6 @@ func TestGenericListResourcesWithFilterForScale(t *testing.T) {
 }
 
 func TestGenericValidation(t *testing.T) {
-
 	ctx := context.Background()
 
 	memBackend, err := memory.New(memory.Config{
@@ -571,16 +570,18 @@ func TestGenericValidation(t *testing.T) {
 
 	_, err = service.UpsertResource(ctx, r1)
 	require.ErrorIs(t, err, validationErr)
-
 }
 
 func TestGenericKeyOverride(t *testing.T) {
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	memBackend, err := memory.New(memory.Config{
 		Context: ctx,
 		Clock:   clockwork.NewFakeClock(),
 	})
 	require.NoError(t, err)
+	defer memBackend.Close()
 
 	service, err := NewService(&ServiceConfig[*testResource]{
 		Backend:       memBackend,
@@ -589,7 +590,7 @@ func TestGenericKeyOverride(t *testing.T) {
 		BackendPrefix: backend.NewKey("generic_prefix"),
 		UnmarshalFunc: unmarshalResource,
 		MarshalFunc:   marshalResource,
-		KeyFunc:       func(tr *testResource) string { return "llama" },
+		NameKeyFunc:   func(string) string { return "llama" },
 	})
 	require.NoError(t, err)
 
@@ -654,4 +655,18 @@ func TestGenericKeyOverride(t *testing.T) {
 	item, err = memBackend.Get(ctx, backend.NewKey("generic_prefix", r1.GetName()))
 	require.Error(t, err)
 	require.Nil(t, item)
+
+	// Validate that getting the resource through the service uses the overridden name
+	_, err = service.GetResource(ctx, r1.GetName())
+	require.NoError(t, err)
+	_, err = service.GetResource(ctx, "llama")
+	require.NoError(t, err)
+	_, err = service.GetResource(ctx, "notllama")
+	require.NoError(t, err)
+
+	// Validate that deleting the resource also uses the overridden name
+	err = service.DeleteResource(ctx, "notllama")
+	require.NoError(t, err)
+	_, err = memBackend.Get(ctx, backend.NewKey("generic_prefix", "llama"))
+	require.ErrorAs(t, err, new(*trace.NotFoundError))
 }