Expose Provisioning States via Cache interface

gravitational · Oct 30, 2024 · 8abc610 · 8abc610
1 parent acdaa01
commit 8abc610
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 3 deletions.
diff --git a/lib/auth/authclient/api.go b/lib/auth/authclient/api.go
@@ -34,6 +34,7 @@ import (
 	integrationpb "github.com/gravitational/teleport/api/gen/proto/go/teleport/integration/v1"
 	kubewaitingcontainerpb "github.com/gravitational/teleport/api/gen/proto/go/teleport/kubewaitingcontainer/v1"
 	machineidv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/machineid/v1"
+	provisioningv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/provisioning/v1"
 	userprovisioningpb "github.com/gravitational/teleport/api/gen/proto/go/teleport/userprovisioning/v2"
 	userspb "github.com/gravitational/teleport/api/gen/proto/go/teleport/users/v1"
 	usertasksv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/usertasks/v1"
@@ -1231,6 +1232,9 @@ type Cache interface {
 	ListStaticHostUsers(ctx context.Context, pageSize int, startKey string) ([]*userprovisioningpb.StaticHostUser, string, error)
 	// GetStaticHostUser returns a static host user by name.
 	GetStaticHostUser(ctx context.Context, name string) (*userprovisioningpb.StaticHostUser, error)
+
+	// GetProvisioningState gets a specific provisioning state
+	GetProvisioningState(context.Context, services.DownstreamID, services.ProvisioningStateID) (*provisioningv1.PrincipalState, error)
 }
 
 type NodeWrapper struct {

diff --git a/lib/cache/cache.go b/lib/cache/cache.go
@@ -44,6 +44,7 @@ import (
 	dbobjectv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/dbobject/v1"
 	kubewaitingcontainerpb "github.com/gravitational/teleport/api/gen/proto/go/teleport/kubewaitingcontainer/v1"
 	notificationsv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/notifications/v1"
+	provisioningv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/provisioning/v1"
 	userprovisioningpb "github.com/gravitational/teleport/api/gen/proto/go/teleport/userprovisioning/v2"
 	userspb "github.com/gravitational/teleport/api/gen/proto/go/teleport/users/v1"
 	usertasksv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/usertasks/v1"
@@ -3544,3 +3545,16 @@ func (c *Cache) GetAccessGraphSettings(ctx context.Context) (*clusterconfigpb.Ac
 	}
 	return rg.reader.GetAccessGraphSettings(ctx)
 }
+
+func (c *Cache) GetProvisioningState(ctx context.Context, downstream services.DownstreamID, id services.ProvisioningStateID) (*provisioningv1.PrincipalState, error) {
+	ctx, span := c.Tracer.Start(ctx, "cache/GetProvisioningState")
+	defer span.End()
+
+	rg, err := readCollectionCache(c, c.collections.provisioningStates)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	defer rg.Release()
+
+	return rg.reader.GetProvisioningState(ctx, downstream, id)
+}
diff --git a/lib/services/provisioningstates.go b/lib/services/provisioningstates.go
@@ -33,12 +33,16 @@ type DownstreamID string
 // disambiguated by a DownstreamID when queried or updated.
 type ProvisioningStateID string
 
-// DownstreamProvisioningStates defines an interface for managing principal
-// provisioning state records scoped by a target downstream receiver.
-type DownstreamProvisioningStates interface {
+type DownstreamProvisioningStateGetter interface {
 	// GetProvisioningState fetches a single provisioning state record for a given
 	// downstream and principal,
 	GetProvisioningState(context.Context, DownstreamID, ProvisioningStateID) (*provisioningv1.PrincipalState, error)
+}
+
+// DownstreamProvisioningStates defines an interface for managing principal
+// provisioning state records scoped by a target downstream receiver.
+type DownstreamProvisioningStates interface {
+	DownstreamProvisioningStateGetter
 
 	// ListProvisioningStates lists all provisioning state records for a given
 	// downstream receiver.