feat: Add workspace api support

README.md

@@ -10,9 +10,11 @@ Check out [Baton](https://github.com/conductorone/baton) to learn more about the
 
 To work with the connector, you can choose from multiple ways to run it, but the main requirement is to have a Databricks account and its ID. You can find the ID of an account, after you log into account platform and click on your username in right top corner that will open a dropdown menu with the account ID along other options.
 
-Another requirement is to have valid credentials to run the connector with. You can use either OAuth client credentials flow or basic auth flow (username and password). To use the OAuth, you need to create a service principal and add OAuth secret (client id and secret) to it. You can do that by going to the user management tab and clicking on the Service Principals tab. Then click on the Add Service principal button and name it. You then need to add OAuth secret to it by clicking on the Generate secret button. You can use this secret to authenticate across all workspaces that service principal has access to. To use basic auth, you just need to provide a username and password of a user that has access to the Databricks API. 
+Another requirement is to have valid credentials to run the connector with. This will decide how connector will be executed. You can use either OAuth client credentials flow or Basic auth flow (username and password) or Bearer auth flow. Both OAuth and Basic can be used across account and all workspaces you have access to. Bearer auth can be used only for a specific workspace.
 
-To get more data about workspaces, you will need to provide a Databricks workspace access token for each workspace you want to sync or already mentioned OAuth secret. You can create a new token by logging into the workspace and going into user settings. Then go to Developer tab and create a new token. You can also use a basic auth across all workspaces you have access to, but this is not recommended as the token is more secure and must be scoped to specific workspaces.
+To use the OAuth, you need to create a service principal and add OAuth secret (client id and secret) to it. You can do that by going to the user management tab and clicking on the Service Principals tab. Then click on the Add Service principal button and name it. You then need to add OAuth secret to it by clicking on the Generate secret button. You can use this secret to authenticate across all workspaces that service principal has access to. To use basic auth, you just need to provide a username and password of a user that has access to the Databricks API. Both methods require admin access to the Databricks account and each workspace you want to sync.
+
+To use bearer auth, you need to provide a Databricks workspace access token. You can create a new token by logging into the workspace and going into user settings. Then go to Developer tab and create a new access token.
 
 # Getting Started
 
@@ -53,6 +55,8 @@ baton resources
 - Users
 - Roles
 
+By default, connector will fetch all resources from the account and all workspaces. You can limit the scope of the sync by providing a list of workspaces to sync with. You can do that by providing a comma-separated list of workspace hostnames to the `--workspaces` flag. Connector expects admin permissions in all workspaces you want to sync with.
+
 # Contributing, Support and Issues
 
 We started Baton because we were tired of taking screenshots and manually building spreadsheets. We welcome contributions, and ideas, no matter how small -- our goal is to make identity and permissions sprawl less painful for everyone. If you have questions, problems, or ideas: Please open a Github Issue!

cmd/baton-databricks/config.go

@@ -29,19 +29,11 @@ func (c *config) IsOauth() bool {
 	return c.DatabricksClientId != "" && c.DatabricksClientSecret != ""
 }
 
-func (c *config) AreWorkspacesSet() bool {
-	return len(c.Workspaces) > 0
-}
-
 func (c *config) AreTokensSet() bool {
-	return len(c.Workspaces) == len(c.Tokens)
-}
-
-func (c *config) IsAccAuthReady() bool {
-	return c.IsOauth() || c.IsBasicAuth()
+	return (len(c.Tokens) > 0) && (len(c.Workspaces) == len(c.Tokens))
 }
 
-func (c *config) IsWorkspaceAuthReady() bool {
+func (c *config) IsAuthReady() bool {
 	return c.AreTokensSet() || c.IsOauth() || c.IsBasicAuth()
 }
 
@@ -51,12 +43,8 @@ func validateConfig(ctx context.Context, cfg *config) error {
 		return fmt.Errorf("account ID must be provided, use --help for more information")
 	}
 
-	if !cfg.IsAccAuthReady() {
-		return fmt.Errorf("username and password must be provided, use --help for more information")
-	}
-
-	if cfg.AreWorkspacesSet() && !cfg.IsWorkspaceAuthReady() {
-		return fmt.Errorf("either access token along workspaces or username and password must be provided, use --help for more information")
+	if !cfg.IsAuthReady() {
+		return fmt.Errorf("either access token along workspaces or username and password or client id and client secret must be provided, use --help for more information")
 	}
 
 	return nil

cmd/baton-databricks/main.go

@@ -40,25 +40,28 @@ func main() {
 func prepareClientAuth(ctx context.Context, cfg *config) databricks.Auth {
 	l := ctxzap.Extract(ctx)
 
-	if cfg.IsBasicAuth() {
+	switch {
+	case cfg.IsBasicAuth():
 		l.Info("using basic auth", zap.String("account-id", cfg.AccountId), zap.String("username", cfg.Username))
 		cAuth := databricks.NewBasicAuth(cfg.Username, cfg.Password)
-
 		return cAuth
-	} else if cfg.IsOauth() {
+	case cfg.IsOauth():
 		l.Info("using oauth", zap.String("account-id", cfg.AccountId), zap.String("client-id", cfg.DatabricksClientId))
 		cAuth := databricks.NewOAuth2(cfg.AccountId, cfg.DatabricksClientId, cfg.DatabricksClientSecret)
-
 		return cAuth
+	case cfg.AreTokensSet():
+		l.Info("using access token", zap.String("account-id", cfg.AccountId))
+		cAuth := databricks.NewTokenAuth(cfg.Workspaces, cfg.Tokens)
+		return cAuth
+	default:
+		return nil
 	}
-
-	return nil
 }
 
 func getConnector(ctx context.Context, cfg *config) (types.ConnectorServer, error) {
 	l := ctxzap.Extract(ctx)
 	auth := prepareClientAuth(ctx, cfg)
-	cb, err := connector.New(ctx, cfg.AccountId, auth)
+	cb, err := connector.New(ctx, cfg.AccountId, auth, cfg.Workspaces)
 	if err != nil {
 		l.Error("error creating connector", zap.Error(err))
 		return nil, err

pkg/connector/account.go

@@ -12,6 +12,7 @@ import (
 	ent "github.com/conductorone/baton-sdk/pkg/types/entitlement"
 	"github.com/conductorone/baton-sdk/pkg/types/grant"
 	rs "github.com/conductorone/baton-sdk/pkg/types/resource"
+	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
 const (
@@ -39,18 +40,25 @@ func (a *accountBuilder) ResourceType(ctx context.Context) *v2.ResourceType {
 	return accountResourceType
 }
 
-func accountResource(ctx context.Context, accID string) (*v2.Resource, error) {
-	resource, err := rs.NewResource(
-		accID,
-		accountResourceType,
-		accID,
-		rs.WithAnnotation(
+func accountResource(ctx context.Context, accID string, accAPIAvailable bool) (*v2.Resource, error) {
+	children := []protoreflect.ProtoMessage{
+		&v2.ChildResourceType{ResourceTypeId: workspaceResourceType.Id},
+	}
+
+	if accAPIAvailable {
+		children = append(children,
 			&v2.ChildResourceType{ResourceTypeId: userResourceType.Id},
 			&v2.ChildResourceType{ResourceTypeId: groupResourceType.Id},
 			&v2.ChildResourceType{ResourceTypeId: servicePrincipalResourceType.Id},
-			&v2.ChildResourceType{ResourceTypeId: workspaceResourceType.Id},
 			&v2.ChildResourceType{ResourceTypeId: roleResourceType.Id},
-		),
+		)
+	}
+
+	resource, err := rs.NewResource(
+		accID,
+		accountResourceType,
+		accID,
+		rs.WithAnnotation(children...),
 	)
 
 	if err != nil {
@@ -63,7 +71,7 @@ func accountResource(ctx context.Context, accID string) (*v2.Resource, error) {
 func (a *accountBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId, pToken *pagination.Token) ([]*v2.Resource, string, annotations.Annotations, error) {
 	var rv []*v2.Resource
 
-	ur, err := accountResource(ctx, a.client.GetAccountId())
+	ur, err := accountResource(ctx, a.client.GetAccountId(), a.client.IsAccountAPIAvailable())
 	if err != nil {
 		return nil, "", nil, err
 	}
@@ -75,6 +83,10 @@ func (a *accountBuilder) List(ctx context.Context, parentResourceID *v2.Resource
 
 // Entitlements returns slice of entitlements for marketplace admins under account.
 func (a *accountBuilder) Entitlements(_ context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
+	if !a.client.IsAccountAPIAvailable() {
+		return nil, "", nil, nil
+	}
+
 	var rv []*v2.Entitlement
 
 	permissiongOptions := []ent.EntitlementOption{
@@ -89,7 +101,14 @@ func (a *accountBuilder) Entitlements(_ context.Context, resource *v2.Resource,
 }
 
 // Grants returns grants for marketplace admins under account.
+// To get marketplace admins, we can only use the account API.
 func (a *accountBuilder) Grants(ctx context.Context, resource *v2.Resource, pToken *pagination.Token) ([]*v2.Grant, string, annotations.Annotations, error) {
+	if !a.client.IsAccountAPIAvailable() {
+		return nil, "", nil, nil
+	}
+
+	a.client.SetAccountConfig()
+
 	var rv []*v2.Grant
 
 	// list rule sets for the account
@@ -102,12 +121,17 @@ func (a *accountBuilder) Grants(ctx context.Context, resource *v2.Resource, pTok
 		// rule set contains role and its principals, each one with resource type and resource id seperated by "/"
 		if strings.Contains(ruleSet.Role, MarketplaceAdminRole) {
 			for _, p := range ruleSet.Principals {
-				resourceId, anns, err := prepareResourceID(ctx, a.client, p)
+				resourceId, err := prepareResourceID(ctx, a.client, p)
 				if err != nil {
 					return nil, "", nil, fmt.Errorf("databricks-connector: failed to prepare resource id for principal %s: %w", p, err)
 				}
 
-				rv = append(rv, grant.NewGrant(resource, MarketplaceAdminRole, resourceId, grant.WithAnnotation(anns...)))
+				var annotations []protoreflect.ProtoMessage
+				if resourceId.ResourceType == groupResourceType.Id {
+					annotations = append(annotations, expandGrantForGroup(resourceId.Resource))
+				}
+
+				rv = append(rv, grant.NewGrant(resource, MarketplaceAdminRole, resourceId, grant.WithAnnotation(annotations...)))
 			}
 		}
 	}

pkg/connector/connector.go

@@ -12,7 +12,8 @@ import (
 )
 
 type Databricks struct {
-	client *databricks.Client
+	client     *databricks.Client
+	workspaces []string
 }
 
 // ResourceSyncers returns a ResourceSyncer for each resource type that should be synced from the upstream service.
@@ -22,8 +23,7 @@ func (d *Databricks) ResourceSyncers(ctx context.Context) []connectorbuilder.Res
 		newGroupBuilder(d.client),
 		newServicePrincipalBuilder(d.client),
 		newUserBuilder(d.client),
-		// TODO: implement workspace builder that will be able to use different clients for different workspaces
-		newWorkspaceBuilder(d.client),
+		newWorkspaceBuilder(d.client, d.workspaces),
 		newRoleBuilder(d.client),
 	}
 }
@@ -43,24 +43,83 @@ func (d *Databricks) Metadata(ctx context.Context) (*v2.ConnectorMetadata, error
 }
 
 // Validate is called to ensure that the connector is properly configured. It should exercise any API credentials
-// to be sure that they are valid.
+// to be sure that they are valid. Since this connector works with two APIs and can have different types of credentials
+// it is important to validate that the connector is properly configured before attempting to sync.
 func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, error) {
-	_, _, err := d.client.ListUsers(ctx, &databricks.PaginationVars{Count: 1})
-	if err != nil {
-		return nil, fmt.Errorf("databricks-connector: failed to validate client: %w", err)
+	cfg := d.client.GetCurrentConfig()
+	isAccAPIAvailable := false
+	isWSAPIAvailable := false
+
+	// Check if we can list users from Account API (unless we are using token auth specific to a single workspace).
+	if !d.client.IsTokenAuth() {
+		_, err := d.client.ListRoles(ctx, "", "")
+		if err == nil {
+			isAccAPIAvailable = true
+		}
 	}
 
+	// Validate that credentials are valid for each targetted workspace.
+	if len(d.workspaces) > 0 {
+		for _, workspace := range d.workspaces {
+			d.client.SetWorkspaceConfig(workspace)
+
+			_, err := d.client.ListRoles(ctx, "", "")
+			if err != nil && !isAccAPIAvailable {
+				return nil, fmt.Errorf("databricks-connector: failed to validate credentials for workspace %s: %w", workspace, err)
+			}
+
+			isWSAPIAvailable = true
+		}
+	}
+
+	// Validate that credentials are valid for every workspace.
+	if len(d.workspaces) == 0 {
+		workspaces, err := d.client.ListWorkspaces(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err)
+		}
+
+		for _, workspace := range workspaces {
+			d.client.SetWorkspaceConfig(workspace.Host)
+
+			_, err := d.client.ListRoles(ctx, "", "")
+			if err != nil && !isAccAPIAvailable {
+				return nil, fmt.Errorf("databricks-connector: failed to validate credentials for workspace %s: %w", workspace.Host, err)
+			}
+
+			isWSAPIAvailable = true
+		}
+	}
+
+	// Resolve the result.
+	if !isAccAPIAvailable && !isWSAPIAvailable {
+		return nil, fmt.Errorf("databricks-connector: failed to validate credentials")
+	}
+
+	// Restore the original config.
+	d.client.UpdateConfig(cfg)
+	d.client.UpdateAvailability(isAccAPIAvailable, isWSAPIAvailable)
+
 	return nil, nil
 }
 
 // New returns a new instance of the connector.
-func New(ctx context.Context, acc string, auth databricks.Auth) (*Databricks, error) {
+func New(ctx context.Context, acc string, auth databricks.Auth, workspaces []string) (*Databricks, error) {
 	httpClient, err := auth.GetClient(ctx)
 	if err != nil {
 		return nil, err
 	}
 
 	client := databricks.NewClient(httpClient, acc, auth)
 
-	return &Databricks{client}, nil
+	if client.IsTokenAuth() {
+		client.SetWorkspaceConfig(workspaces[0])
+	} else {
+		client.SetAccountConfig()
+	}
+
+	return &Databricks{
+		client,
+		workspaces,
+	}, nil
 }