diff --git a/pkg/config/conf.gen.go b/pkg/config/conf.gen.go index 4a963e2f..eb3f1676 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -8,11 +8,11 @@ type Github struct { Orgs []string `mapstructure:"orgs"` Enterprises []string `mapstructure:"enterprises"` InstanceUrl string `mapstructure:"instance-url"` - SyncSecrets bool `mapstructure:"sync-secrets"` - OmitArchivedRepositories bool `mapstructure:"omit-archived-repositories"` AppId string `mapstructure:"app-id"` AppPrivatekeyPath []byte `mapstructure:"app-privatekey-path"` Org string `mapstructure:"org"` + SyncSecrets bool `mapstructure:"sync-secrets"` + OmitArchivedRepositories bool `mapstructure:"omit-archived-repositories"` } func (c *Github) findFieldByTag(tagValue string) (any, bool) { diff --git a/pkg/connector/helpers.go b/pkg/connector/helpers.go index 2659668b..308fe25b 100644 --- a/pkg/connector/helpers.go +++ b/pkg/connector/helpers.go @@ -255,6 +255,48 @@ type listUsersQuery struct { } } +// batchSAMLQuery fetches all SAML identities for an organization with pagination. +// This is used to avoid N+1 queries when syncing users. +type batchSAMLQuery struct { + Organization struct { + SamlIdentityProvider struct { + ExternalIdentities struct { + PageInfo struct { + HasNextPage bool + EndCursor githubv4.String + } + Edges []struct { + Node struct { + SamlIdentity struct { + NameId string + Emails []struct { + Value string + } + } + User struct { + Login string + Name string + } + } + } + } `graphql:"externalIdentities(first: 100, after: $cursor)"` + } + } `graphql:"organization(login: $orgLoginName)"` + RateLimit struct { + Limit int + Cost int + Remaining int + ResetAt githubv4.DateTime + } +} + +// SAMLIdentity holds the parsed SAML identity data for a user. +type SAMLIdentity struct { + PrimaryEmail string `json:"primary_email"` + ExtraEmails []string `json:"extra_emails,omitempty"` + Name string `json:"name,omitempty"` +} + type hasSAMLQuery struct { Organization struct { SamlIdentityProvider struct { diff --git a/pkg/connector/org_role.go b/pkg/connector/org_role.go index 99fcf031..a70997ed 100644 --- a/pkg/connector/org_role.go +++ b/pkg/connector/org_role.go @@ -219,6 +219,11 @@ func (o *orgRoleResourceType) Grants( rv = append(rv, grant) } case resourceTypeTeam.Id: + orgID, err := parseResourceToGitHub(resource.ParentResourceId) + if err != nil { + return nil, nil, err + } + listOpts := &github.ListOptions{ Page: page, PerPage: maxPageSize, @@ -256,21 +261,21 @@ func (o *orgRoleResourceType) Grants( // Create expandable grants for teams. To show inherited roles, we need to show the teams that have the role. for _, team := range teams { - teamResource, err := teamResource(team, resource.ParentResourceId) + tr, err := teamResource(team, orgID, resource.ParentResourceId) if err != nil { return nil, nil, err } rv = append(rv, grant.NewGrant( resource, "assigned", - teamResource.Id, + tr.Id, grant.WithAnnotation(&v2.V1Identifier{ Id: fmt.Sprintf("org-role-grant:%s:%d:%s", resource.Id.Resource, team.GetID(), "assigned"), }, &v2.GrantExpandable{ EntitlementIds: []string{ - entitlement.NewEntitlementID(teamResource, teamRoleMaintainer), - entitlement.NewEntitlementID(teamResource, teamRoleMember), + entitlement.NewEntitlementID(tr, teamRoleMaintainer), + entitlement.NewEntitlementID(tr, teamRoleMember), }, Shallow: true, }, diff --git a/pkg/connector/repository.go b/pkg/connector/repository.go index 571bb358..e944a8cd 100644 --- a/pkg/connector/repository.go +++ b/pkg/connector/repository.go @@ -226,6 +226,11 @@ func (o *repositoryResourceType) Grants( } case resourceTypeTeam.Id: + orgID, err := parseResourceToGitHub(resource.ParentResourceId) + if err != nil { + return nil, nil, err + } + listOpts := &github.ListOptions{ Page: page, PerPage: maxPageSize, @@ -269,7 +274,7 @@ func (o *repositoryResourceType) Grants( continue } - tr, err := teamResource(team, resource.ParentResourceId) + tr, err := teamResource(team, orgID, resource.ParentResourceId) if err != nil { return nil, nil, err } diff --git a/pkg/connector/team.go b/pkg/connector/team.go index e6eff4c7..8bb350cc 100644 --- a/pkg/connector/team.go +++ b/pkg/connector/team.go @@ -30,12 +30,15 @@ var teamAccessLevels = []string{ } // teamResource creates a new connector resource for a GitHub Team. It is possible that the team has a parent resource. -func teamResource(team *github.Team, parentResourceID *v2.ResourceId) (*v2.Resource, error) { +// orgID must be passed explicitly since ListTeams() doesn't return the full organization object. +func teamResource(team *github.Team, orgID int64, parentResourceID *v2.ResourceId) (*v2.Resource, error) { profile := map[string]interface{}{ + // Note: members_count and repos_count are only populated when fetching + // individual teams. We skip the per-team API call to avoid N+1 requests. "members_count": team.GetMembersCount(), "repos_count": team.GetReposCount(), // Store the org ID in the profile so that we can reference it when calculating grants - "orgID": team.GetOrganization().GetID(), + "orgID": orgID, } ret, err := rType.NewGroupResource( @@ -104,12 +107,10 @@ func (o *teamResourceType) List(ctx context.Context, parentID *v2.ResourceId, op } for _, team := range teams { - fullTeam, resp, err := o.client.Teams.GetTeamByID(ctx, orgID, team.GetID()) //nolint:staticcheck // TODO: migrate to GetTeamBySlug - if err != nil { - return nil, nil, wrapGitHubError(err, resp, "github-connector: failed to get team details") - } - - tr, err := teamResource(fullTeam, &v2.ResourceId{ResourceType: resourceTypeOrg.Id, Resource: fmt.Sprintf("%d", orgID)}) + // Use team data directly from ListTeams() to avoid N+1 API calls. + // Note: members_count and repos_count won't be populated, but orgID + // is passed explicitly since we already have it from the parent resource. + tr, err := teamResource(team, orgID, &v2.ResourceId{ResourceType: resourceTypeOrg.Id, Resource: fmt.Sprintf("%d", orgID)}) if err != nil { return nil, nil, err } diff --git a/pkg/connector/team_test.go b/pkg/connector/team_test.go index 5700f14c..f6dc0328 100644 --- a/pkg/connector/team_test.go +++ b/pkg/connector/team_test.go @@ -28,7 +28,7 @@ func TestTeam(t *testing.T) { client := teamBuilder(githubClient, cache) organization, _ := organizationResource(ctx, githubOrganization, nil, false) - team, _ := teamResource(githubTeam, organization.Id) + team, _ := teamResource(githubTeam, githubOrganization.GetID(), organization.Id) user, _ := userResource(ctx, githubUser, *githubUser.Email, nil) entitlement := v2.Entitlement{ diff --git a/pkg/connector/user.go b/pkg/connector/user.go index a67df7ba..96149c39 100644 --- a/pkg/connector/user.go +++ b/pkg/connector/user.go @@ -2,6 +2,7 @@ package connector import ( "context" + "encoding/json" "fmt" "net/mail" "strconv" @@ -105,6 +106,14 @@ const ( // enterprise_saml:* keys. This allows bulk-reading SAML mappings with // GetManyJSON without scanning the entire session store. enterpriseSAMLKeysIndex = "enterprise_saml_keys" + + // orgSAMLKeyPrefix is prepended to each GitHub login to form + // individual session keys for org-level SAML, e.g. "org_saml:myorg:octocat". + orgSAMLKeyPrefix = "org_saml:" + + // orgSAMLKeysIndexPrefix is the session key prefix that stores the list of all + // org_saml:* keys for a specific org. Format: "org_saml_keys:myorg" + orgSAMLKeysIndexPrefix = "org_saml_keys:" ) type userResourceType struct { @@ -147,7 +156,7 @@ func (u *userResourceType) List(ctx context.Context, parentID *v2.ResourceId, op // For enterprise SAML: on the first page, fetch from the API and store in // session. On every page, bulk-read the mappings into a local map so the // user loop can do plain map lookups with no session calls. - var enterpriseSAMLEmails map[string]string + var enterpriseSAMLIdentities map[string]SAMLIdentity if currentSAMLState == samlStateEnterprise { _, alreadyFetched, err := session.GetJSON[[]string](ctx, opts.Session, enterpriseSAMLKeysIndex) if err != nil { @@ -166,7 +175,40 @@ func (u *userResourceType) List(ctx context.Context, parentID *v2.ResourceId, op } } if currentSAMLState == samlStateEnterprise { - enterpriseSAMLEmails, err = loadEnterpriseSAMLEmails(ctx, opts.Session) + enterpriseSAMLIdentities, err = loadEnterpriseSAMLIdentities(ctx, opts.Session) + if err != nil { + return nil, nil, err + } + } + } + + // For org-level SAML: batch-fetch all SAML identities on first page and cache them. + // This avoids N+1 GraphQL queries (one per user) during user sync. + var orgSAMLIdentities map[string]SAMLIdentity + if currentSAMLState == samlStateOrgEnabled { + keyIndex := orgSAMLKeysIndexPrefix + orgName + _, alreadyFetched, err := session.GetJSON[[]string](ctx, opts.Session, keyIndex) + if err != nil { + return nil, nil, fmt.Errorf("baton-github: error checking org SAML session: %w", err) + } + if !alreadyFetched { + graphqlRateLimit, fetchErr := u.fetchAndStoreOrgSAML(ctx, opts.Session, orgName) + if fetchErr != nil { + l.Debug("failed to fetch org SAML identities, falling back to REST API emails", + zap.Error(fetchErr), + zap.String("org", orgName)) + // Write empty sentinel so we don't retry + if setErr := session.SetJSON(ctx, opts.Session, keyIndex, []string{}); setErr != nil { + l.Debug("failed to write empty org SAML sentinel to session", zap.Error(setErr)) + } + u.samlStates[orgName] = samlStateDisabled + currentSAMLState = samlStateDisabled + } else if graphqlRateLimit != nil { + annotations.WithRateLimiting(graphqlRateLimit) + } + } + if currentSAMLState == samlStateOrgEnabled { + orgSAMLIdentities, err = loadOrgSAMLIdentities(ctx, opts.Session, orgName) if err != nil { return nil, nil, err } @@ -202,80 +244,59 @@ func (u *userResourceType) List(ctx context.Context, parentID *v2.ResourceId, op return nil, nil, err } - var lastGraphQLRateLimit *struct { - Limit int - Remaining int - ResetAt githubv4.DateTime - } rv := make([]*v2.Resource, 0, len(users)) for _, user := range users { - ghUser, res, err := u.client.Users.GetByID(ctx, user.GetID()) - if err != nil { - // This undocumented API can return 404 for some users. If this fails it means we won't get some of their details like email - if isNotFoundError(res) { - l.Warn("error fetching user by id", zap.Error(err), zap.Int64("user_id", user.GetID())) - ghUser = user - } else { - return nil, nil, wrapGitHubError(err, res, "github-connector: failed to get user by id") - } - } - userEmail := ghUser.GetEmail() + var userEmail string var extraEmails []string + var ghUser *github.User switch currentSAMLState { case samlStateUnknown: return nil, nil, fmt.Errorf("baton-github: unexpected unknown SAML state for org %s", orgName) - case samlStateOrgEnabled: - q := listUsersQuery{} - variables := map[string]interface{}{ - "orgLoginName": githubv4.String(orgName), - "userName": githubv4.String(ghUser.GetLogin()), - } - err = u.graphqlClient.Query(ctx, &q, variables) - if err != nil { - return nil, nil, err - } - if len(q.Organization.SamlIdentityProvider.ExternalIdentities.Edges) == 1 { - samlIdent := q.Organization.SamlIdentityProvider.ExternalIdentities.Edges[0].Node.SamlIdentity - userEmail = samlIdent.NameId - setUserEmail := false - - if userEmail != "" { - setUserEmail = true - } - for _, email := range samlIdent.Emails { - ok := isEmail(email.Value) - if !ok { - continue - } - - if !setUserEmail { - userEmail = email.Value - setUserEmail = true - } else { - extraEmails = append(extraEmails, email.Value) - } + case samlStateOrgEnabled: + // Use cached SAML identity instead of per-user GraphQL query. + // Skip GetByID - we have the identity data we need from SAML cache, + // including the user's Name from the GraphQL batch query. + ghUser = user + login := strings.ToLower(user.GetLogin()) + key := orgSAMLKeyPrefix + orgName + ":" + login + if ident, ok := orgSAMLIdentities[key]; ok { + userEmail = ident.PrimaryEmail + extraEmails = ident.ExtraEmails + if ident.Name != "" && ghUser.Name == nil { + ghUser.Name = &ident.Name } } - lastGraphQLRateLimit = &struct { - Limit int - Remaining int - ResetAt githubv4.DateTime - }{ - Limit: q.RateLimit.Limit, - Remaining: q.RateLimit.Remaining, - ResetAt: q.RateLimit.ResetAt, - } case samlStateEnterprise: - key := enterpriseSAMLKeyPrefix + strings.ToLower(ghUser.GetLogin()) - if samlEmail, ok := enterpriseSAMLEmails[key]; ok && isEmail(samlEmail) { - userEmail = samlEmail + // Use cached SAML identity from consumed licenses (email + name). + // No GetByID needed — name comes from the licenses API. + ghUser = user + key := enterpriseSAMLKeyPrefix + strings.ToLower(user.GetLogin()) + if ident, ok := enterpriseSAMLIdentities[key]; ok { + if isEmail(ident.PrimaryEmail) { + userEmail = ident.PrimaryEmail + } + if ident.Name != "" && ghUser.Name == nil { + ghUser.Name = &ident.Name + } } case samlStateDisabled: - // no SAML enrichment + // No SAML - need to call GetByID to get email + var res *github.Response + ghUser, res, err = u.client.Users.GetByID(ctx, user.GetID()) + if err != nil { + // This undocumented API can return 404 for some users. If this fails it means we won't get some of their details like email + if isNotFoundError(res) { + l.Warn("error fetching user by id", zap.Error(err), zap.Int64("user_id", user.GetID())) + ghUser = user + } else { + return nil, nil, wrapGitHubError(err, res, "github-connector: failed to get user by id") + } + } + userEmail = ghUser.GetEmail() } ur, err := userResource(ctx, ghUser, userEmail, extraEmails) @@ -286,14 +307,6 @@ func (u *userResourceType) List(ctx context.Context, parentID *v2.ResourceId, op rv = append(rv, ur) } annotations.WithRateLimiting(restApiRateLimit) - if lastGraphQLRateLimit != nil && int64(lastGraphQLRateLimit.Remaining) < restApiRateLimit.Remaining { - graphqlRateLimit := &v2.RateLimitDescription{ - Limit: int64(lastGraphQLRateLimit.Limit), - Remaining: int64(lastGraphQLRateLimit.Remaining), - ResetAt: timestamppb.New(lastGraphQLRateLimit.ResetAt.Time), - } - annotations.WithRateLimiting(graphqlRateLimit) - } return rv, &resource.SyncOpResults{ NextPageToken: pageToken, @@ -440,13 +453,13 @@ func (u *userResourceType) checkOrgSAML(ctx context.Context, orgName string) (sa } // fetchAndStoreEnterpriseSAML pages through the consumed licenses API for all -// configured enterprises, aggregates the login-to-SAML-email mappings, and +// configured enterprises, aggregates the login-to-SAML-identity mappings, and // writes them to the session store in a single batch. It also stores the list -// of keys under enterpriseSAMLKeysIndex so that loadEnterpriseSAMLEmails can +// of keys under enterpriseSAMLKeysIndex so that loadEnterpriseSAMLIdentities can // bulk-read them back on subsequent List pages. func (u *userResourceType) fetchAndStoreEnterpriseSAML(ctx context.Context, ss sessions.SessionStore) error { l := ctxzap.Extract(ctx) - samlByLogin := make(map[string]string) + samlByLogin := make(map[string]SAMLIdentity) for _, enterprise := range u.enterprises { // GitHub's consumed-licenses API is 1-indexed; page 0 is undocumented @@ -464,7 +477,10 @@ func (u *userResourceType) fetchAndStoreEnterpriseSAML(ctx context.Context, ss s for _, user := range consumedLicenses.Users { if user.GitHubComSAMLNameID != nil && *user.GitHubComSAMLNameID != "" && user.GitHubComLogin != "" { key := enterpriseSAMLKeyPrefix + strings.ToLower(user.GitHubComLogin) - samlByLogin[key] = *user.GitHubComSAMLNameID + samlByLogin[key] = SAMLIdentity{ + PrimaryEmail: *user.GitHubComSAMLNameID, + Name: user.GitHubComName, + } } } page++ @@ -472,11 +488,17 @@ func (u *userResourceType) fetchAndStoreEnterpriseSAML(ctx context.Context, ss s } keys := make([]string, 0, len(samlByLogin)) - for k := range samlByLogin { + stringMap := make(map[string]string, len(samlByLogin)) + for k, v := range samlByLogin { keys = append(keys, k) + data, err := json.Marshal(v) + if err != nil { + return fmt.Errorf("baton-github: error serializing enterprise SAML identity: %w", err) + } + stringMap[k] = string(data) } - if len(samlByLogin) > 0 { - if err := session.SetManyJSON(ctx, ss, samlByLogin); err != nil { + if len(stringMap) > 0 { + if err := session.SetManyJSON(ctx, ss, stringMap); err != nil { return fmt.Errorf("baton-github: error storing enterprise SAML mappings: %w", err) } } @@ -491,11 +513,11 @@ func (u *userResourceType) fetchAndStoreEnterpriseSAML(ctx context.Context, ss s return nil } -// loadEnterpriseSAMLEmails bulk-reads all enterprise SAML mappings from the +// loadEnterpriseSAMLIdentities bulk-reads all enterprise SAML mappings from the // session store in two calls: one to get the key index, one to get the values. -// Returns a map of "enterprise_saml:" -> SAML email for use as a local +// Returns a map of "enterprise_saml:" -> SAMLIdentity for use as a local // lookup table in the List loop (no session calls needed per user). -func loadEnterpriseSAMLEmails(ctx context.Context, ss sessions.SessionStore) (map[string]string, error) { +func loadEnterpriseSAMLIdentities(ctx context.Context, ss sessions.SessionStore) (map[string]SAMLIdentity, error) { keys, found, err := session.GetJSON[[]string](ctx, ss, enterpriseSAMLKeysIndex) if err != nil { return nil, fmt.Errorf("baton-github: error reading enterprise SAML key index: %w", err) @@ -504,9 +526,138 @@ func loadEnterpriseSAMLEmails(ctx context.Context, ss sessions.SessionStore) (ma return nil, nil } - samlByLogin, err := session.GetManyJSON[string](ctx, ss, keys) + stringMap, err := session.GetManyJSON[string](ctx, ss, keys) if err != nil { return nil, fmt.Errorf("baton-github: error reading enterprise SAML mappings: %w", err) } - return samlByLogin, nil + + result := make(map[string]SAMLIdentity, len(stringMap)) + for k, v := range stringMap { + var ident SAMLIdentity + if err := json.Unmarshal([]byte(v), &ident); err != nil { + continue + } + result[k] = ident + } + return result, nil +} + +// fetchAndStoreOrgSAML pages through the org's SAML external identities via GraphQL, +// aggregates the login-to-SAML-identity mappings, and writes them to the session store. +// This avoids N+1 GraphQL queries when syncing users for orgs with SAML enabled. +func (u *userResourceType) fetchAndStoreOrgSAML(ctx context.Context, ss sessions.SessionStore, orgName string) (*v2.RateLimitDescription, error) { + l := ctxzap.Extract(ctx) + samlByLogin := make(map[string]SAMLIdentity) + + var lastRateLimit *v2.RateLimitDescription + var cursor *githubv4.String + for { + q := batchSAMLQuery{} + variables := map[string]interface{}{ + "orgLoginName": githubv4.String(orgName), + "cursor": cursor, + } + err := u.graphqlClient.Query(ctx, &q, variables) + if err != nil { + return nil, fmt.Errorf("baton-github: error fetching org SAML identities for %s: %w", orgName, err) + } + lastRateLimit = &v2.RateLimitDescription{ + Limit: int64(q.RateLimit.Limit), + Remaining: int64(q.RateLimit.Remaining), + ResetAt: timestamppb.New(q.RateLimit.ResetAt.Time), + } + + for _, edge := range q.Organization.SamlIdentityProvider.ExternalIdentities.Edges { + if edge.Node.User.Login == "" { + continue + } + login := strings.ToLower(edge.Node.User.Login) + key := orgSAMLKeyPrefix + orgName + ":" + login + + ident := SAMLIdentity{ + Name: edge.Node.User.Name, + } + // Extract primary email from NameId or first email + if edge.Node.SamlIdentity.NameId != "" && isEmail(edge.Node.SamlIdentity.NameId) { + ident.PrimaryEmail = edge.Node.SamlIdentity.NameId + } + for _, email := range edge.Node.SamlIdentity.Emails { + if !isEmail(email.Value) { + continue + } + if ident.PrimaryEmail == "" { + ident.PrimaryEmail = email.Value + } else if email.Value != ident.PrimaryEmail { + ident.ExtraEmails = append(ident.ExtraEmails, email.Value) + } + } + if ident.PrimaryEmail != "" { + samlByLogin[key] = ident + } + } + + if !q.Organization.SamlIdentityProvider.ExternalIdentities.PageInfo.HasNextPage { + break + } + cursor = &q.Organization.SamlIdentityProvider.ExternalIdentities.PageInfo.EndCursor + } + + // Store identities in session + keyIndex := orgSAMLKeysIndexPrefix + orgName + keys := make([]string, 0, len(samlByLogin)) + + // Convert map to string values for SetManyJSON + stringMap := make(map[string]string, len(samlByLogin)) + for k, v := range samlByLogin { + keys = append(keys, k) + // Serialize SAMLIdentity to JSON string for storage + data, err := json.Marshal(v) + if err != nil { + return nil, fmt.Errorf("baton-github: error serializing SAML identity: %w", err) + } + stringMap[k] = string(data) + } + + if len(stringMap) > 0 { + if err := session.SetManyJSON(ctx, ss, stringMap); err != nil { + return nil, fmt.Errorf("baton-github: error storing org SAML mappings: %w", err) + } + } + // Always write the key index + if err := session.SetJSON(ctx, ss, keyIndex, keys); err != nil { + return nil, fmt.Errorf("baton-github: error storing org SAML key index: %w", err) + } + + l.Debug("stored org SAML mappings in session", zap.String("org", orgName), zap.Int("count", len(samlByLogin))) + return lastRateLimit, nil +} + +// loadOrgSAMLIdentities bulk-reads all org SAML mappings from the session store. +// Returns a map of "org_saml::" -> SAMLIdentity for local lookups. +func loadOrgSAMLIdentities(ctx context.Context, ss sessions.SessionStore, orgName string) (map[string]SAMLIdentity, error) { + keyIndex := orgSAMLKeysIndexPrefix + orgName + keys, found, err := session.GetJSON[[]string](ctx, ss, keyIndex) + if err != nil { + return nil, fmt.Errorf("baton-github: error reading org SAML key index: %w", err) + } + if !found || len(keys) == 0 { + return nil, nil + } + + stringMap, err := session.GetManyJSON[string](ctx, ss, keys) + if err != nil { + return nil, fmt.Errorf("baton-github: error reading org SAML mappings: %w", err) + } + + // Deserialize JSON strings back to SAMLIdentity + result := make(map[string]SAMLIdentity, len(stringMap)) + for k, v := range stringMap { + var ident SAMLIdentity + if err := json.Unmarshal([]byte(v), &ident); err != nil { + // Skip malformed entries + continue + } + result[k] = ident + } + return result, nil }