[BB-1409] Add caching to reduce API calls

pkg/connector/api_token.go

@@ -95,7 +95,7 @@ func (o *apiTokenResourceType) List(
 	})
 	if err != nil {
 		if isRatelimited(resp) {
-			return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+			return nil, "", nil, uhttp.WrapErrorsWithRateLimitInfo(codes.Unavailable, resp.Response, err)
 		}
 		return nil, "", nil, err
 	}

pkg/connector/connector.go

@@ -97,21 +97,24 @@ type GitHub struct {
 	graphqlClient  *githubv4.Client
 	hasSAMLEnabled *bool
 	orgCache       *orgNameCache
+	userCache      *userDataCache
+	teamCache      *teamDataCache
 	syncSecrets    bool
 	enterprises    []string
 }
 
 func (gh *GitHub) ResourceSyncers(ctx context.Context) []connectorbuilder.ResourceSyncer {
 	resourceSyncers := []connectorbuilder.ResourceSyncer{
-		orgBuilder(gh.client, gh.appClient, gh.orgCache, gh.orgs, gh.syncSecrets),
-		teamBuilder(gh.client, gh.orgCache),
-		userBuilder(gh.client, gh.hasSAMLEnabled, gh.graphqlClient, gh.orgCache, gh.orgs),
-		repositoryBuilder(gh.client, gh.orgCache),
-		orgRoleBuilder(gh.client, gh.orgCache),
+		orgBuilder(gh.client, gh.appClient, gh.orgCache, gh.userCache, gh.orgs, gh.syncSecrets),
+		teamBuilder(gh.client, gh.orgCache, gh.userCache, gh.teamCache),
+		userBuilder(gh.client, gh.hasSAMLEnabled, gh.graphqlClient, gh.orgCache, gh.userCache, gh.orgs),
+		repositoryBuilder(gh.client, gh.orgCache, gh.userCache, gh.teamCache),
+		orgRoleBuilder(gh.client, gh.orgCache, gh.userCache),
 		invitationBuilder(invitationBuilderParams{
-			client:   gh.client,
-			orgCache: gh.orgCache,
-			orgs:     gh.orgs,
+			client:    gh.client,
+			orgCache:  gh.orgCache,
+			userCache: gh.userCache,
+			orgs:      gh.orgs,
 		}),
 	}
 
@@ -120,7 +123,7 @@ func (gh *GitHub) ResourceSyncers(ctx context.Context) []connectorbuilder.Resour
 	}
 
 	if len(gh.enterprises) > 0 {
-		resourceSyncers = append(resourceSyncers, enterpriseRoleBuilder(gh.client, gh.customClient, gh.enterprises))
+		resourceSyncers = append(resourceSyncers, enterpriseRoleBuilder(gh.client, gh.customClient, gh.userCache, gh.enterprises))
 	}
 	return resourceSyncers
 }
@@ -314,6 +317,8 @@ func New(ctx context.Context, ghc *cfg.Github, appKey string) (*GitHub, error) {
 		enterprises:   ghc.Enterprises,
 		graphqlClient: graphqlClient,
 		orgCache:      newOrgNameCache(ghClient),
+		userCache:     newUserDataCache(ghClient),
+		teamCache:     newTeamDataCache(ghClient),
 		syncSecrets:   ghc.SyncSecrets,
 	}
 	return gh, nil
@@ -454,7 +459,7 @@ func getOrgs(ctx context.Context, client *github.Client, orgs []string) ([]strin
 		orgs, resp, err := client.Organizations.List(ctx, "", &github.ListOptions{Page: page, PerPage: maxPageSize})
 		if err != nil {
 			if isRatelimited(resp) {
-				return nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+				return nil, uhttp.WrapErrorsWithRateLimitInfo(codes.Unavailable, resp.Response, err)
 			}
 			if isAuthError(resp) {
 				return nil, uhttp.WrapErrors(codes.Unauthenticated, "github-connector: failed to retrieve org", err)

pkg/connector/enterprise_role.go

@@ -22,6 +22,7 @@ type enterpriseRoleResourceType struct {
 	resourceType   *v2.ResourceType
 	client         *github.Client
 	customClient   *customclient.Client
+	userCache      *userDataCache
 	enterprises    []string
 	roleUsersCache map[string][]string
 	mu             *sync.Mutex
@@ -139,10 +140,10 @@ func (o *enterpriseRoleResourceType) Grants(
 
 	ret := []*v2.Grant{}
 	for _, userLogin := range cache[resource.Id.Resource] {
-		user, resp, err := o.client.Users.Get(ctx, userLogin)
+		user, resp, err := o.userCache.GetUserByLogin(ctx, userLogin)
 		if err != nil {
 			if isRatelimited(resp) {
-				return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+				return nil, "", nil, uhttp.WrapErrorsWithRateLimitInfo(codes.Unavailable, resp.Response, err)
 			}
 			return nil, "", nil, fmt.Errorf("baton-github: error getting user %s: %w", userLogin, err)
 		}
@@ -162,11 +163,12 @@ func (o *enterpriseRoleResourceType) Grants(
 	return ret, "", nil, nil
 }
 
-func enterpriseRoleBuilder(client *github.Client, customClient *customclient.Client, enterprises []string) *enterpriseRoleResourceType {
+func enterpriseRoleBuilder(client *github.Client, customClient *customclient.Client, userCache *userDataCache, enterprises []string) *enterpriseRoleResourceType {
 	return &enterpriseRoleResourceType{
 		resourceType:   resourceTypeEnterpriseRole,
 		client:         client,
 		customClient:   customClient,
+		userCache:      userCache,
 		enterprises:    enterprises,
 		roleUsersCache: make(map[string][]string),
 		mu:             &sync.Mutex{},

pkg/connector/helpers.go

@@ -67,6 +67,115 @@ func newOrgNameCache(c *github.Client) *orgNameCache {
 	}
 }
 
+type userDataCache struct {
+	sync.RWMutex
+	c            *github.Client
+	usersById    map[int64]*github.User
+	usersByLogin map[string]*github.User
+}
+
+func (u *userDataCache) GetUser(ctx context.Context, userID int64) (*github.User, *github.Response, error) {
+	u.RLock()
+	if user, ok := u.usersById[userID]; ok {
+		u.RUnlock()
+		return user, nil, nil
+	}
+	u.RUnlock()
+
+	u.Lock()
+	defer u.Unlock()
+
+	if user, ok := u.usersById[userID]; ok {
+		return user, nil, nil
+	}
+
+	user, resp, err := u.c.Users.GetByID(ctx, userID)
+	if err != nil {
+		return nil, resp, err
+	}
+
+	u.usersById[userID] = user
+	if user.Login != nil {
+		u.usersByLogin[*user.Login] = user
+	}
+
+	return user, resp, nil
+}
+
+func (u *userDataCache) GetUserByLogin(ctx context.Context, login string) (*github.User, *github.Response, error) {
+	u.RLock()
+	if user, ok := u.usersByLogin[login]; ok {
+		u.RUnlock()
+		return user, nil, nil
+	}
+	u.RUnlock()
+
+	u.Lock()
+	defer u.Unlock()
+
+	if user, ok := u.usersByLogin[login]; ok {
+		return user, nil, nil
+	}
+
+	user, resp, err := u.c.Users.Get(ctx, login)
+	if err != nil {
+		return nil, resp, err
+	}
+
+	u.usersByLogin[login] = user
+	if user.ID != nil {
+		u.usersById[*user.ID] = user
+	}
+
+	return user, resp, nil
+}
+
+func newUserDataCache(c *github.Client) *userDataCache {
+	return &userDataCache{
+		c:            c,
+		usersById:    make(map[int64]*github.User),
+		usersByLogin: make(map[string]*github.User),
+	}
+}
+
+type teamDataCache struct {
+	sync.RWMutex
+	c     *github.Client
+	teams map[int64]*github.Team
+}
+
+func (t *teamDataCache) GetTeam(ctx context.Context, orgID int64, teamID int64) (*github.Team, *github.Response, error) {
+	t.RLock()
+	if team, ok := t.teams[teamID]; ok {
+		t.RUnlock()
+		return team, nil, nil
+	}
+	t.RUnlock()
+
+	t.Lock()
+	defer t.Unlock()
+
+	if team, ok := t.teams[teamID]; ok {
+		return team, nil, nil
+	}
+
+	team, resp, err := t.c.Teams.GetTeamByID(ctx, orgID, teamID) //nolint:staticcheck // TODO: migrate to GetTeamBySlug
+	if err != nil {
+		return nil, resp, err
+	}
+
+	t.teams[teamID] = team
+
+	return team, resp, nil
+}
+
+func newTeamDataCache(c *github.Client) *teamDataCache {
+	return &teamDataCache{
+		c:     c,
+		teams: make(map[int64]*github.Team),
+	}
+}
+
 func v1AnnotationsForResourceType(resourceTypeID string) annotations.Annotations {
 	annos := annotations.Annotations{}
 	annos.Update(&v2.V1Identifier{

pkg/connector/invitation.go

@@ -40,9 +40,10 @@ func invitationToUserResource(invitation *github.Invitation) (*v2.Resource, erro
 }
 
 type invitationResourceType struct {
-	client   *github.Client
-	orgCache *orgNameCache
-	orgs     []string
+	client    *github.Client
+	orgCache  *orgNameCache
+	userCache *userDataCache
+	orgs      []string
 }
 
 var _ connectorbuilder.AccountManager = &invitationResourceType{}
@@ -225,15 +226,17 @@ func getCreateUserParams(accountInfo *v2.AccountInfo) (*createUserParams, error)
 }
 
 type invitationBuilderParams struct {
-	client   *github.Client
-	orgCache *orgNameCache
-	orgs     []string
+	client    *github.Client
+	orgCache  *orgNameCache
+	userCache *userDataCache
+	orgs      []string
 }
 
 func invitationBuilder(p invitationBuilderParams) *invitationResourceType {
 	return &invitationResourceType{
-		client:   p.client,
-		orgCache: p.orgCache,
-		orgs:     p.orgs,
+		client:    p.client,
+		orgCache:  p.orgCache,
+		userCache: p.userCache,
+		orgs:      p.orgs,
 	}
 }

pkg/connector/org.go

@@ -38,6 +38,7 @@ type orgResourceType struct {
 	appClient    *github.Client
 	orgs         map[string]struct{}
 	orgCache     *orgNameCache
+	userCache    *userDataCache
 	syncSecrets  bool
 }
 
@@ -103,7 +104,7 @@ func (o *orgResourceType) List(
 	orgs, resp, err := o.client.Organizations.List(ctx, "", opts)
 	if err != nil {
 		if isRatelimited(resp) {
-			return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+			return nil, "", nil, uhttp.WrapErrorsWithRateLimitInfo(codes.Unavailable, resp.Response, err)
 		}
 		return nil, "", nil, fmt.Errorf("github-connector: failed to fetch org: %w", err)
 	}
@@ -131,7 +132,7 @@ func (o *orgResourceType) List(
 			}
 
 			if isRatelimited(resp) {
-				return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+				return nil, "", nil, uhttp.WrapErrorsWithRateLimitInfo(codes.Unavailable, resp.Response, err)
 			}
 			return nil, "", nil, err
 		}
@@ -229,7 +230,7 @@ func (o *orgResourceType) Grants(
 			}
 			errMsg := "github-connectorv2: failed to list org members"
 			if isRatelimited(resp) {
-				return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+				return nil, "", nil, uhttp.WrapErrorsWithRateLimitInfo(codes.Unavailable, resp.Response, err)
 			}
 			return nil, "", nil, fmt.Errorf("%s: %w", errMsg, err)
 		}
@@ -294,7 +295,7 @@ func (o *orgResourceType) Grant(ctx context.Context, principal *v2.Resource, en
 		return nil, err
 	}
 
-	user, _, err := o.client.Users.GetByID(ctx, principalID)
+	user, _, err := o.userCache.GetUser(ctx, principalID)
 	if err != nil {
 		return nil, fmt.Errorf("github-connectorv2: failed to get user: %w", err)
 	}
@@ -386,7 +387,7 @@ func (o *orgResourceType) Revoke(ctx context.Context, grant *v2.Grant) (annotati
 		return nil, err
 	}
 
-	user, _, err := o.client.Users.GetByID(ctx, principalID)
+	user, _, err := o.userCache.GetUser(ctx, principalID)
 	if err != nil {
 		return nil, fmt.Errorf("github-connectorv2: failed to get user: %w", err)
 	}
@@ -416,7 +417,7 @@ func (o *orgResourceType) Revoke(ctx context.Context, grant *v2.Grant) (annotati
 	return nil, nil
 }
 
-func orgBuilder(client, appClient *github.Client, orgCache *orgNameCache, orgs []string, syncSecrets bool) *orgResourceType {
+func orgBuilder(client, appClient *github.Client, orgCache *orgNameCache, userCache *userDataCache, orgs []string, syncSecrets bool) *orgResourceType {
 	orgMap := make(map[string]struct{})
 
 	for _, o := range orgs {
@@ -429,6 +430,7 @@ func orgBuilder(client, appClient *github.Client, orgCache *orgNameCache, orgs [
 		client:       client,
 		appClient:    appClient,
 		orgCache:     orgCache,
+		userCache:    userCache,
 		syncSecrets:  syncSecrets,
 	}
 }

pkg/connector/org_role.go

@@ -40,6 +40,7 @@ type orgRoleResourceType struct {
 	resourceType *v2.ResourceType
 	client       *github.Client
 	orgCache     *orgNameCache
+	userCache    *userDataCache
 }
 
 func orgRoleResource(
@@ -91,7 +92,7 @@ func (o *orgRoleResourceType) List(
 			return nil, "", nil, nil
 		}
 		if isRatelimited(resp) {
-			return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+			return nil, "", nil, uhttp.WrapErrorsWithRateLimitInfo(codes.Unavailable, resp.Response, err)
 		}
 		return nil, "", nil, fmt.Errorf("failed to list organization roles: %w", err)
 	}
@@ -228,7 +229,7 @@ func (o *orgRoleResourceType) Grants(
 				return nil, pageToken, nil, nil
 			}
 			if isRatelimited(resp) {
-				return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+				return nil, "", nil, uhttp.WrapErrorsWithRateLimitInfo(codes.Unavailable, resp.Response, err)
 			}
 			return nil, "", nil, fmt.Errorf("failed to list role teams: %w", err)
 		}
@@ -325,7 +326,7 @@ func (o *orgRoleResourceType) Grant(ctx context.Context, principal *v2.Resource,
 		return nil, fmt.Errorf("github-connectorv2: invalid entitlement ID: %s", entitlement.Id)
 	}
 
-	user, _, err := o.client.Users.GetByID(ctx, userID)
+	user, _, err := o.userCache.GetUser(ctx, userID)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get user: %w", err)
 	}
@@ -400,7 +401,7 @@ func (o *orgRoleResourceType) Revoke(ctx context.Context, grant *v2.Grant) (anno
 		return nil, fmt.Errorf("invalid user ID: %w", err)
 	}
 
-	user, _, err := o.client.Users.GetByID(ctx, userID)
+	user, _, err := o.userCache.GetUser(ctx, userID)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get user: %w", err)
 	}
@@ -423,10 +424,11 @@ func (o *orgRoleResourceType) Revoke(ctx context.Context, grant *v2.Grant) (anno
 	return nil, nil
 }
 
-func orgRoleBuilder(client *github.Client, orgCache *orgNameCache) *orgRoleResourceType {
+func orgRoleBuilder(client *github.Client, orgCache *orgNameCache, userCache *userDataCache) *orgRoleResourceType {
 	return &orgRoleResourceType{
 		resourceType: resourceTypeOrgRole,
 		client:       client,
 		orgCache:     orgCache,
+		userCache:    userCache,
 	}
 }