use session storage as cache to get system users

Author: luisina-santos

cmd/baton-jumpcloud/main.go

@@ -6,11 +6,17 @@ import (
 	cfg "github.com/conductorone/baton-jumpcloud/pkg/config"
 	"github.com/conductorone/baton-jumpcloud/pkg/connector"
 	"github.com/conductorone/baton-sdk/pkg/config"
+	"github.com/conductorone/baton-sdk/pkg/connectorrunner"
 )
 
 var version = "dev"
 
 func main() {
 	ctx := context.Background()
-	config.RunConnector(ctx, "baton-jumpcloud", version, cfg.ConfigurationSchema, connector.NewLambdaConnector)
+	config.RunConnector(ctx,
+		"baton-jumpcloud",
+		version,
+		cfg.ConfigurationSchema,
+		connector.NewLambdaConnector,
+		connectorrunner.WithSessionStoreEnabled())
 }

pkg/client/client.go

@@ -160,38 +160,6 @@ func (jc *Client) ListAdminUsers(ctx context.Context, opts *Options) ([]jcapi1.U
 	return users, resp, pageToken, nil
 }
 
-func (jc *Client) FetchUserByEmail(ctx context.Context, email string) (*jcapi1.Systemuserreturn, error) {
-	ctx, client := jc.client1(ctx)
-
-	if email == "" {
-		return nil, fmt.Errorf("email parameter cannot be empty when fetching user by email")
-	}
-
-	users, resp, err := client.SystemusersApi.SystemusersList(ctx).Filter(fmt.Sprintf("email:$eq:%s", email)).Execute()
-	if err != nil {
-		return nil, wrapSDKError(err, resp, "failed to list users")
-	}
-	defer resp.Body.Close()
-
-	if len(users.Results) == 0 {
-		return nil, uhttp.WrapErrors(
-			codes.NotFound,
-			fmt.Sprintf("user not found for email: %s", email),
-			nil,
-		)
-	}
-
-	if len(users.Results) != 1 {
-		return nil, uhttp.WrapErrors(
-			codes.InvalidArgument,
-			fmt.Sprintf("multiple users found for email: %s", email),
-			nil,
-		)
-	}
-
-	return &users.Results[0], nil
-}
-
 func (jc *Client) GetSystemUserByID(ctx context.Context, userID string) (*jcapi1.Systemuserreturn, error) {
 	ctx, client := jc.client1(ctx)
 

pkg/connector/apps.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"strconv"
-	"strings"
 
 	"github.com/conductorone/baton-jumpcloud/pkg/client"
 	"github.com/conductorone/baton-jumpcloud/pkg/client/jcapi1"
@@ -13,7 +12,6 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 	sdkResources "github.com/conductorone/baton-sdk/pkg/types/resource"
-	"google.golang.org/grpc/codes"
 )
 
 const (
@@ -25,6 +23,7 @@ const (
 type appResourceType struct {
 	resourceType *v2.ResourceType
 	client       *client.Client
+	usersCache   *usersCache
 }
 
 func (o *appResourceType) ResourceType(_ context.Context) *v2.ResourceType {
@@ -148,17 +147,25 @@ func (o *appResourceType) adminGrants(ctx context.Context, resource *v2.Resource
 	defer resp.Body.Close()
 
 	var rv []*v2.Grant
+
+	adminEmails := make([]string, len(users))
+	for i := range users {
+		adminEmails[i] = users[i].GetEmail()
+	}
+	systemUsersMap, err := o.usersCache.GetSystemUsersByEmailList(ctx, opts.Session, adminEmails)
+	if err != nil {
+		return nil, nil, err
+	}
+
 	for i := range users {
 		adminUser := &users[i]
 		var adminPrincipal appAdminPrincipal = adminUser
 
 		// If the user is a system user, we need to fetch the user by email to get the ID
-		systemUser, err := o.client.FetchUserByEmail(ctx, adminUser.GetEmail())
-		if err != nil && !strings.Contains(err.Error(), codes.NotFound.String()) {
-			return nil, nil, fmt.Errorf("failed to fetch system user by email during admin grants processing: %w", err)
-		}
-		if systemUser != nil {
-			adminPrincipal = systemUser
+		if systemUser, ok := systemUsersMap[adminUser.GetEmail()]; ok {
+			if systemUser != nil {
+				adminPrincipal = systemUser
+			}
 		}
 
 		ur := &v2.Resource{
@@ -303,5 +310,6 @@ func newAppBuilder(c *client.Client) *appResourceType {
 	return &appResourceType{
 		resourceType: resourceTypeApp,
 		client:       c,
+		usersCache:   newUsersCache(c),
 	}
 }

pkg/connector/helper.go

@@ -1,19 +1,11 @@
 package connector
 
 import (
-	"errors"
 	"fmt"
-	"sync"
 
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 )
 
-var (
-	userCache                sync.Map
-	errUserNotFoundForEmail  = errors.New("user not found for email")
-	errMultipleUsersForEmail = errors.New("multiple users found for email")
-)
-
 func fmtResourceId(resourceTypeID string, id string) *v2.ResourceId {
 	return &v2.ResourceId{
 		ResourceType: resourceTypeID,

pkg/connector/users.go

@@ -20,6 +20,7 @@ import (
 type userResourceType struct {
 	resourceType *v2.ResourceType
 	client       *client.Client
+	usersCache   *usersCache
 	managers     map[string]*jcapi1.Systemuserreturn
 }
 
@@ -32,6 +33,7 @@ func newUserBuilder(client *client.Client) *userResourceType {
 		resourceType: resourceTypeUser,
 		client:       client,
 		managers:     make(map[string]*jcapi1.Systemuserreturn),
+		usersCache:   newUsersCache(client),
 	}
 }
 
@@ -103,6 +105,12 @@ func (o *userResourceType) List(ctx context.Context, parentResourceID *v2.Resour
 			return nil, nil, fmt.Errorf("failed to list system users: %w", err)
 		}
 
+		// populate the users cache with the system users
+		err = o.usersCache.SetSystemUsers(ctx, opts.Session, systemUsers)
+		if err != nil {
+			return nil, nil, err
+		}
+
 		for i := range systemUsers {
 			ur, err := o.userResource(ctx, &systemUsers[i])
 			if err != nil {
@@ -121,6 +129,14 @@ func (o *userResourceType) List(ctx context.Context, parentResourceID *v2.Resour
 			return nil, nil, fmt.Errorf("failed to list admin users: %w", err)
 		}
 
+		adminEmails := make([]string, len(adminUsers))
+		for i := range adminUsers {
+			adminEmails[i] = adminUsers[i].GetEmail()
+		}
+		systemUsersMap, err := o.usersCache.GetSystemUsersByEmailList(ctx, opts.Session, adminEmails)
+		if err != nil {
+			return nil, nil, err
+		}
 		for i := range adminUsers {
 			adminEmail := adminUsers[i].GetEmail()
 			adminUser, err := o.adminUserResource(&adminUsers[i])
@@ -129,14 +145,12 @@ func (o *userResourceType) List(ctx context.Context, parentResourceID *v2.Resour
 			}
 
 			// Check if the admin user is also a system user, if so we'll use that user instead
-			// TODO (luisina) > use a cache to store user by email, to avoid making multiple requests to the API
-			systemUser, err := o.client.FetchUserByEmail(ctx, adminEmail)
-			if err != nil && !strings.Contains(err.Error(), codes.NotFound.String()) {
-				return nil, nil, err
-			}
-
-			if systemUser != nil {
-				continue
+			if systemUser, ok := systemUsersMap[adminEmail]; ok {
+				if systemUser != nil {
+					continue
+				} else {
+					return nil, nil, err
+				}
 			}
 
 			l.Debug("admin user not found as system user, creating", zap.String("email", adminEmail))

pkg/connector/users_cache.go

@@ -0,0 +1,59 @@
+package connector
+
+import (
+	"context"
+
+	"github.com/conductorone/baton-jumpcloud/pkg/client"
+	"github.com/conductorone/baton-jumpcloud/pkg/client/jcapi1"
+	"github.com/conductorone/baton-sdk/pkg/session"
+	"github.com/conductorone/baton-sdk/pkg/types/sessions"
+)
+
+type usersCache struct {
+	client *client.Client
+}
+
+func newUsersCache(jumpcloudClient *client.Client) *usersCache {
+	return &usersCache{
+		client: jumpcloudClient,
+	}
+}
+
+func (uc *usersCache) GetSystemUsersByEmailList(ctx context.Context, sessionStorage sessions.SessionStore, emails []string) (map[string]*jcapi1.Systemuserreturn, error) {
+	systemUsers, err := session.GetManyJSON[*jcapi1.Systemuserreturn](ctx, sessionStorage, emails)
+	if err != nil {
+		return nil, err
+	}
+	return systemUsers, nil
+}
+
+func (uc *usersCache) SetSystemUsers(ctx context.Context, sessionStorage sessions.SessionStore, systemUsers []jcapi1.Systemuserreturn) error {
+	if len(systemUsers) == 0 {
+		return nil
+	}
+
+	const batchSize = 100
+
+	// Process in batches of 100
+	for i := 0; i < len(systemUsers); i += batchSize {
+		end := i + batchSize
+		if end > len(systemUsers) {
+			end = len(systemUsers)
+		}
+
+		batch := systemUsers[i:end]
+		systemUsersMap := make(map[string]*jcapi1.Systemuserreturn, len(batch))
+
+		// Create map with email as key and pointer to user
+		// Note: we need to take address of the slice element, not the loop variable
+		for j := range batch {
+			systemUsersMap[batch[j].GetEmail()] = &batch[j]
+		}
+
+		if err := session.SetManyJSON(ctx, sessionStorage, systemUsersMap); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}