Skip to content

[BB-1755, BB-1771]Add retry logic for not found errors in connector methods and fix missed grants with oauth #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
FeliLucero1 wants to merge 3 commits into from
Closed

[BB-1755, BB-1771]Add retry logic for not found errors in connector methods and fix missed grants with oauth #32

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8ba9cf1
Add retry logic for not found errors in connector methods
FeliLucero1 Nov 12, 2025
8a16f69
fix group grant missed with oauth
FeliLucero1 Nov 14, 2025
5aaf4be
include rate limiting annotations
FeliLucero1 Nov 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 35 additions & 35 deletions pkg/connector/groups.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,50 +189,48 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, pToken
return nil, "", nil, fmt.Errorf("databricks-connector: failed to parse group resource id: %w", err)
}

groupTrait, err := rs.GetGroupTrait(resource)
Copy link

@laurenleach laurenleach Nov 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@FeliLucero1 was the group trait missing members? or just didn't include them at all?

Copy link
Contributor Author

@FeliLucero1 FeliLucero1 Nov 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Used to lose some

Copy link
Contributor Author

@FeliLucero1 FeliLucero1 Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems that the users It wasn't losing were users who had permissions over the group but weren't members of it, therefore It was losing all members

if err != nil {
return nil, "", nil, err
}

var workspaceId string
isWorkspaceGroup := parentId.ResourceType == workspaceResourceType.Id
if isWorkspaceGroup {
workspaceId = parentId.Resource
}

// membership grants
membersPayload, ok := rs.GetProfileStringValue(groupTrait.Profile, "members")
if ok {
members := strings.Split(membersPayload, ",")

for _, m := range members {
pp := strings.Split(m, "/")
if len(pp) != 2 {
return nil, "", nil, fmt.Errorf("databricks-connector: invalid member format of %s: %w", m, err)
}
// Always fetch the group with members attribute to ensure we get the members
// regardless of authentication type (OAuth vs personal access token)
group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
if err != nil {
return nil, "", nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
}
Comment on lines +199 to +204
Copy link

@coderabbitai coderabbitai bot Nov 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Apply retry wrapper consistently with Grant and Revoke methods.

The GetGroup error is not wrapped with makeRetryableIfNotFound here, but the same calls in Grant (line 318) and Revoke (line 462) do wrap it. Per the PR objectives (adding retry logic for not found errors), this method should also benefit from the retry behavior.

Apply this diff:

 	group, rateLimitData, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
 	if err != nil {
+		err = makeRetryableIfNotFound(err)
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
 	}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// Always fetch the group with members attribute to ensure we get the members
// regardless of authentication type (OAuth vs personal access token)
group, rateLimitData, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
if err != nil {
return nil, "", nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
}
// Always fetch the group with members attribute to ensure we get the members
// regardless of authentication type (OAuth vs personal access token)
group, rateLimitData, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, "", nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
}
🤖 Prompt for AI Agents 
In pkg/connector/groups.go around lines 199 to 204, the call to
g.client.GetGroup returns an error that is not wrapped with
makeRetryableIfNotFound like the same calls in Grant and Revoke; update the
error handling to pass the returned error through makeRetryableIfNotFound before
returning so transient "not found" errors become retryable (i.e., replace the
current fmt.Errorf return with returning makeRetryableIfNotFound(err) wrapped
appropriately to preserve the original error context).


memberType, memberID := pp[0], pp[1]
var resourceId *v2.ResourceId
var anns []protoreflect.ProtoMessage
for _, member := range group.Members {
// member.Ref contains the type and ID separated by "/", e.g., "Users/123" or "Groups/456"
pp := strings.Split(member.Ref, "/")
if len(pp) != 2 {
return nil, "", nil, fmt.Errorf("databricks-connector: invalid member format of %s", member.Ref)
}

switch memberType {
case "Users":
resourceId = &v2.ResourceId{ResourceType: userResourceType.Id, Resource: memberID}
case "Groups":
rid, expandAnnotation, err := groupGrantExpansion(ctx, memberID, parentId)
if err != nil {
return rv, "", nil, err
}
resourceId = rid
anns = append(anns, expandAnnotation)
case "ServicePrincipals":
resourceId = &v2.ResourceId{ResourceType: servicePrincipalResourceType.Id, Resource: memberID}
default:
return nil, "", nil, fmt.Errorf("databricks-connector: invalid member type: %s", memberType)
}
memberType, memberID := pp[0], pp[1]
var resourceId *v2.ResourceId
var anns []protoreflect.ProtoMessage

rv = append(rv, grant.NewGrant(resource, groupMemberEntitlement, resourceId, grant.WithAnnotation(anns...)))
switch memberType {
case "Users":
resourceId = &v2.ResourceId{ResourceType: userResourceType.Id, Resource: memberID}
case "Groups":
rid, expandAnnotation, err := groupGrantExpansion(ctx, memberID, parentId)
if err != nil {
return rv, "", nil, err
}
resourceId = rid
anns = append(anns, expandAnnotation)
case "ServicePrincipals":
resourceId = &v2.ResourceId{ResourceType: servicePrincipalResourceType.Id, Resource: memberID}
default:
return nil, "", nil, fmt.Errorf("databricks-connector: invalid member type: %s", memberType)
}

rv = append(rv, grant.NewGrant(resource, groupMemberEntitlement, resourceId, grant.WithAnnotation(anns...)))
}

// role permissions grants
Expand Down Expand Up @@ -306,8 +304,9 @@ func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitl
membershipEntitlementID := ent.NewEntitlementID(entitlement.Resource, groupMemberEntitlement)
managerEntitlementID := ent.NewEntitlementID(entitlement.Resource, groupManagerEntitlement)
if entitlement.Id == membershipEntitlementID {
group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource)
group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
}

Expand Down Expand Up @@ -449,8 +448,9 @@ func (g *groupBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations
membershipEntitlementID := ent.NewEntitlementID(entitlement.Resource, groupMemberEntitlement)
managerEntitlementID := ent.NewEntitlementID(entitlement.Resource, groupManagerEntitlement)
if entitlement.Id == membershipEntitlementID {
group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource)
group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
}

Expand Down
27 changes: 27 additions & 0 deletions pkg/connector/helpers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,18 @@ package connector

import (
"context"
"errors"
"fmt"
"net/http"
"slices"
"strings"

"github.com/conductorone/baton-databricks/pkg/databricks"
v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
"github.com/conductorone/baton-sdk/pkg/annotations"
rs "github.com/conductorone/baton-sdk/pkg/types/resource"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/types/known/structpb"
)

Expand Down Expand Up @@ -70,20 +74,23 @@ func prepareResourceId(ctx context.Context, c *databricks.Client, workspaceId st
case UsersType:
userID, _, err := c.FindUserID(ctx, workspaceId, principal)
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, fmt.Errorf("failed to find user %s: %w", principal, err)
}

resourceId = &v2.ResourceId{ResourceType: userResourceType.Id, Resource: userID}
case GroupsType:
groupID, _, err := c.FindGroupID(ctx, workspaceId, principal)
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, fmt.Errorf("failed to find group %s: %w", principal, err)
}

resourceId = &v2.ResourceId{ResourceType: groupResourceType.Id, Resource: groupID}
case ServicePrincipalsType:
servicePrincipalID, _, err := c.FindServicePrincipalID(ctx, workspaceId, principal)
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, fmt.Errorf("failed to find service principal %s: %w", principal, err)
}

Expand Down Expand Up @@ -119,20 +126,23 @@ func preparePrincipalId(ctx context.Context, c *databricks.Client, workspaceId,
case userResourceType.Id:
username, _, err := c.FindUsername(ctx, workspaceId, principalId)
if err != nil {
err = makeRetryableIfNotFound(err)
return "", fmt.Errorf("failed to find user %s/%s: %w", workspaceId, principalId, err)
}

result = fmt.Sprintf("%s/%s", UsersType, username)
case groupResourceType.Id:
displayName, _, err := c.FindGroupDisplayName(ctx, workspaceId, principalId)
if err != nil {
err = makeRetryableIfNotFound(err)
return "", fmt.Errorf("failed to find group %s/%s: %w", workspaceId, principalId, err)
}

result = fmt.Sprintf("%s/%s", GroupsType, displayName)
case servicePrincipalResourceType.Id:
appID, _, err := c.FindServicePrincipalAppID(ctx, workspaceId, principalId)
if err != nil {
err = makeRetryableIfNotFound(err)
return "", fmt.Errorf("failed to find service principal %s/%s: %w", workspaceId, principalId, err)
}

Expand Down Expand Up @@ -202,3 +212,20 @@ func prepareWorkspaceRole(entitlement string) string {

return parts[1]
}

func makeRetryableIfNotFound(err error) error {
if err == nil {
return nil
}

var apiErr *databricks.APIError
if errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound {
return status.Error(codes.Unavailable, fmt.Sprintf("resource not found yet, will retry: %v", err))
}

if status.Code(err) == codes.NotFound {
return status.Error(codes.Unavailable, fmt.Sprintf("resource not found yet, will retry: %v", err))
}

return err
}
7 changes: 5 additions & 2 deletions pkg/connector/roles.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,7 @@ func (r *roleBuilder) Grant(ctx context.Context, principal *v2.Resource, entitle
case userResourceType.Id:
u, _, err := r.client.GetUser(ctx, workspaceId, principal.Id.Resource)
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, fmt.Errorf("databricks-connector: failed to get user: %w", err)
}

Expand All @@ -347,8 +348,9 @@ func (r *roleBuilder) Grant(ctx context.Context, principal *v2.Resource, entitle
}

case groupResourceType.Id:
g, _, err := r.client.GetGroup(ctx, workspaceId, principal.Id.Resource)
g, _, err := r.client.GetGroup(ctx, workspaceId, principal.Id.Resource, databricks.NewGroupRolesAttrVars())
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, fmt.Errorf("databricks-connector: failed to get group: %w", err)
}

Expand All @@ -362,6 +364,7 @@ func (r *roleBuilder) Grant(ctx context.Context, principal *v2.Resource, entitle
case servicePrincipalResourceType.Id:
sp, _, err := r.client.GetServicePrincipal(ctx, workspaceId, principal.Id.Resource)
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, fmt.Errorf("databricks-connector: failed to get service principal: %w", err)
}

Expand Down Expand Up @@ -428,7 +431,7 @@ func (r *roleBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations.
}

case groupResourceType.Id:
g, _, err := r.client.GetGroup(ctx, workspaceId, principal.Id.Resource)
g, _, err := r.client.GetGroup(ctx, workspaceId, principal.Id.Resource, databricks.NewGroupRolesAttrVars())
if err != nil {
return nil, fmt.Errorf("databricks-connector: failed to get group: %w", err)
}
Expand Down
1 change: 1 addition & 0 deletions pkg/connector/users.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,7 @@ func (o *userBuilder) CreateAccount(ctx context.Context, accountInfo *v2.Account

user, _, err := o.client.GetUser(ctx, "", res.Id)
if err != nil {
err = makeRetryableIfNotFound(err)
return nil, nil, nil, fmt.Errorf("baton-databricks: failed to get user after creation: %w", err)
}

Expand Down
18 changes: 10 additions & 8 deletions pkg/databricks/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ import (

v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
"github.com/conductorone/baton-sdk/pkg/uhttp"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
Comment on lines +12 to +13
Copy link

@coderabbitai coderabbitai bot Nov 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

Verify the gRPC library version.

The imports are necessary for the new error handling, but the library context indicates google.golang.org/grpc version v1.73.0-dev is being used. Development versions can be unstable and are generally not recommended for production use.

Run the following script to check the gRPC version in use:

🏁 Script executed:

#!/bin/bash
# Check the version of google.golang.org/grpc in go.mod
cat go.mod | grep "google.golang.org/grpc"

Length of output: 144





**No actionable concerns with the imports; gRPC v1 is stable, but verify dependency management.**

When a project reaches major version v1 it is considered stable. The imports for `google.golang.org/grpc/codes` and `google.golang.org/grpc/status` are appropriate for error handling. While the codebase currently pulls in v1.73.0-dev (a pre-release version), Version 1.73.0 (stable) should be used when available. Since this is marked as an indirect dependency, it's being pulled in by another package—consider verifying your dependency management to use stable releases where possible rather than development versions in production builds.

[resolve_review_comment]

---

🌐 Web query:

```web
google.golang.org/grpc v1.73.0-dev stability production use

💡 Result:

Short answer: don't use a -dev module version in production. The "-dev" suffix denotes a development (unpublished/unstable) build; use an official tagged release (for example v1.73.0 or a later non-dev release) and pin that version after testing. [1][2][3]

Details and sources:

  • grpc-go has an official v1.73.0 release (use the tagged release, not a dev build). [1]
  • A v1.74.0 tagging issue was reported (the project advised using v1.73.0 until a fixed 1.74.x), illustrating that unstable tags/dev builds can be transient or incorrect. [2]
  • The grpc-go repo documentation and releases are the authoritative source for which versions are stable for production; prefer published semver tags, run your tests, and pin the module in go.mod. [3]

If you want, I can look up the latest non-dev grpc-go release right now and suggest the exact module line to put in go.mod.

🌐 Web query:

grpc-go v1.73.0 release date stable version

💡 Result:

grpc-go v1.73.0 (stable) was released on June 5, 2025. [1][2]

Sources:
[1] GitHub Releases — grpc/grpc-go (v1.73.0).
[2] SourceForge mirror — grpc-go v1.73.0.

Address the development version dependency before production use.

The gRPC library version v1.73.0-dev is a development (unstable) build and must not be used in production. The stable release v1.73.0 has been available since June 5, 2025. Update go.mod to pin the stable release instead. Since this is an indirect dependency, investigate which direct dependency is pulling in the -dev version and either update that dependency or explicitly pin google.golang.org/grpc to v1.73.0 in your go.mod.

🤖 Prompt for AI Agents 
pkg/databricks/client.go around lines 12-13: your module is pulling an indirect
gRPC dev release (v1.73.0-dev); locate which direct dependency requires it (use
go mod graph or go list -m -json all) and either update that direct dependency
to a version that depends on google.golang.org/grpc v1.73.0 or explicitly pin
google.golang.org/grpc to v1.73.0 in go.mod (add a require directive and run go
mod tidy) so the stable v1.73.0 is used instead of the -dev build.

)

const (
Expand Down Expand Up @@ -195,7 +197,7 @@ func (c *Client) FindUserID(
}

if len(users) == 0 {
return "", ratelimitData, err
return "", ratelimitData, status.Error(codes.NotFound, fmt.Sprintf("user not found: %s", username))
}

return users[0].ID, ratelimitData, nil
Expand All @@ -222,7 +224,7 @@ func (c *Client) FindUsername(
}

if len(users) == 0 {
return "", ratelimitData, err
return "", ratelimitData, status.Error(codes.NotFound, fmt.Sprintf("user not found: %s", userID))
}

return users[0].UserName, ratelimitData, nil
Expand Down Expand Up @@ -254,7 +256,7 @@ func (c *Client) ListGroups(
return res.Resources, res.Total, ratelimitData, nil
}

func (c *Client) GetGroup(ctx context.Context, workspaceId, groupId string) (
func (c *Client) GetGroup(ctx context.Context, workspaceId, groupId string, vars ...Vars) (
*Group,
*v2.RateLimitDescription,
error,
Expand All @@ -267,7 +269,7 @@ func (c *Client) GetGroup(ctx context.Context, workspaceId, groupId string) (
}

var res *Group
ratelimitData, err := c.Get(ctx, u, &res)
ratelimitData, err := c.Get(ctx, u, &res, vars...)
if err != nil {
return nil, ratelimitData, err
}
Expand Down Expand Up @@ -307,7 +309,7 @@ func (c *Client) FindGroupID(
}

if len(groups) == 0 {
return "", ratelimitData, err
return "", ratelimitData, status.Error(codes.NotFound, fmt.Sprintf("group not found: %s", displayName))
}

return groups[0].ID, ratelimitData, nil
Expand All @@ -334,7 +336,7 @@ func (c *Client) FindGroupDisplayName(
}

if len(groups) == 0 {
return "", ratelimitData, err
return "", ratelimitData, status.Error(codes.NotFound, fmt.Sprintf("group not found: %s", groupID))
}

return groups[0].DisplayName, ratelimitData, nil
Expand Down Expand Up @@ -425,7 +427,7 @@ func (c *Client) FindServicePrincipalID(
}

if len(servicePrincipals) == 0 {
return "", ratelimitData, err
return "", ratelimitData, status.Error(codes.NotFound, fmt.Sprintf("service principal not found: %s", appID))
}

return servicePrincipals[0].ID, ratelimitData, nil
Expand All @@ -452,7 +454,7 @@ func (c *Client) FindServicePrincipalAppID(
}

if len(servicePrincipals) == 0 {
return "", ratelimitData, err
return "", ratelimitData, status.Error(codes.NotFound, fmt.Sprintf("service principal not found: %s", servicePrincipalID))
}

return servicePrincipals[0].ApplicationID, ratelimitData, nil
Expand Down