From 8abc610973ca82a92f41f04c03e3fce328339b17 Mon Sep 17 00:00:00 2001
From: Trent Clarke <trent@goteleport.com>
Date: Wed, 30 Oct 2024 19:07:32 +1100
Subject: [PATCH] Expose Provisioning States via Cache interface

---
 lib/auth/authclient/api.go         |  4 ++++
 lib/cache/cache.go                 | 14 ++++++++++++++
 lib/services/provisioningstates.go | 10 +++++++---
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/lib/auth/authclient/api.go b/lib/auth/authclient/api.go
index efc4eab13b54d..2d4ebe3204e6a 100644
--- 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
index e9273cfb0e8a8..2c3013e9b0d37 100644
--- 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
index e51e4b9b782fc..8d1e1c0de5497 100644
--- 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.