diff --git a/README.md b/README.md index 289d0208..91a837ec 100644 --- a/README.md +++ b/README.md @@ -18,20 +18,18 @@ 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. 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. +will decide how connector will be executed. You can use either the OAuth client +credentials flow or the Bearer auth flow. OAuth can be used across account and +all workspaces you have access to. Bearer auth can be used only for a specific +workspace. 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. +across all workspaces that service principal has access to. This requires 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 @@ -55,14 +53,14 @@ baton-databricks --hostname "azuredatabricks.net" ``` brew install conductorone/baton/baton conductorone/baton/baton-databricks -BATON_ACCOUNT_ID=account_id BATON_USERNAME=username BATON_PASSWORD=password baton-databricks +BATON_ACCOUNT_ID=account_id BATON_DATABRICKS_CLIENT_ID=client_id BATON_DATABRICKS_CLIENT_SECRET=client_secret baton-databricks baton resources ``` ## docker ``` -docker run --rm -v $(pwd):/out -e BATON_ACCOUNT_ID=account_id BATON_USERNAME=username BATON_PASSWORD=password ghcr.io/conductorone/baton-databricks:latest -f "/out/sync.c1z" +docker run --rm -v $(pwd):/out -e BATON_ACCOUNT_ID=account_id BATON_DATABRICKS_CLIENT_ID=client_id BATON_DATABRICKS_CLIENT_SECRET=client_secret ghcr.io/conductorone/baton-databricks:latest -f "/out/sync.c1z" docker run --rm -v $(pwd):/out ghcr.io/conductorone/baton:latest -f "/out/sync.c1z" resources ``` @@ -72,7 +70,7 @@ docker run --rm -v $(pwd):/out ghcr.io/conductorone/baton:latest -f "/out/sync.c go install github.com/conductorone/baton/cmd/baton@main go install github.com/conductorone/baton-databricks/cmd/baton-databricks@main -BATON_ACCOUNT_ID=account_id BATON_USERNAME=username BATON_PASSWORD=password baton-databricks +BATON_ACCOUNT_ID=account_id BATON_DATABRICKS_CLIENT_ID=client_id BATON_DATABRICKS_CLIENT_SECRET=client_secret baton-databricks baton resources ``` @@ -97,6 +95,11 @@ both flags at the same time. If you do that, connector will sync with all workspaces that are associated with provided tokens and all workspaces that are in the list of workspaces. +When authenticating with `--workspace-tokens` instead of the OAuth client ID and +secret, also pass `--auth-method workspace-token` (or set +`BATON_AUTH_METHOD=workspace-token`), otherwise the connector validates against +the OAuth fields by default and rejects the config. + To instead exclude specific workspaces from the sync, pass them to the `--databricks-exclude-workspaces` flag (or the `BATON_DATABRICKS_EXCLUDE_WORKSPACES` environment variable) as a comma-separated @@ -129,25 +132,47 @@ Usage: Available Commands: capabilities Get connector capabilities completion Generate the autocompletion script for the specified shell + config Get the connector config schema + health-check Check the health of a running connector help Help about any command Flags: - --account-hostname string The hostname used to connect to the Databricks account API. If not set, it will be calculated from the hostname field. ($BATON_ACCOUNT_HOSTNAME) - --account-id string required: The Databricks account ID used to connect to the Databricks Account and Workspace API ($BATON_ACCOUNT_ID) - --client-id string The client ID used to authenticate with ConductorOne ($BATON_CLIENT_ID) - --client-secret string The client secret used to authenticate with ConductorOne ($BATON_CLIENT_SECRET) - --databricks-client-id string The Databricks service principal's client ID used to connect to the Databricks Account and Workspace API ($BATON_DATABRICKS_CLIENT_ID) - --databricks-client-secret string The Databricks service principal's client secret used to connect to the Databricks Account and Workspace API ($BATON_DATABRICKS_CLIENT_SECRET) - --databricks-exclude-workspaces strings Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID ($BATON_DATABRICKS_EXCLUDE_WORKSPACES) - -f, --file string The path to the c1z file to sync with ($BATON_FILE) (default "sync.c1z") - -h, --help help for baton-databricks - --hostname string The Databricks hostname used to connect to the Databricks API ($BATON_HOSTNAME) (default "cloud.databricks.com") - --log-format string The output format for logs: json, console ($BATON_LOG_FORMAT) (default "json") - --log-level string The log level: debug, info, warn, error ($BATON_LOG_LEVEL) (default "info") - -p, --provisioning This must be set in order for provisioning actions to be enabled ($BATON_PROVISIONING) - --skip-full-sync This must be set to skip a full sync ($BATON_SKIP_FULL_SYNC) - --ticketing This must be set to enable ticketing support ($BATON_TICKETING) - -v, --version version for baton-databricks + --account-hostname string The hostname used to connect to the Databricks account API. If not set, it will be calculated from the hostname field. ($BATON_ACCOUNT_HOSTNAME) + --account-id string required: The Databricks account ID used to connect to the Databricks Account and Workspace API ($BATON_ACCOUNT_ID) + --auth-method string ($BATON_AUTH_METHOD) + --client-id string The client ID used to authenticate with ConductorOne ($BATON_CLIENT_ID) + --client-secret string The client secret used to authenticate with ConductorOne ($BATON_CLIENT_SECRET) + --databricks-client-id string required: The Databricks service principal's client ID used to connect to the Databricks Account and Workspace API ($BATON_DATABRICKS_CLIENT_ID) + --databricks-client-secret string required: The Databricks service principal's client secret used to connect to the Databricks Account and Workspace API ($BATON_DATABRICKS_CLIENT_SECRET) + --databricks-exclude-workspaces strings Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID ($BATON_DATABRICKS_EXCLUDE_WORKSPACES) + --external-resource-c1z string The path to the c1z file to sync external baton resources with ($BATON_EXTERNAL_RESOURCE_C1Z) + --external-resource-entitlement-id-filter string The entitlement that external users, groups must have access to sync external baton resources ($BATON_EXTERNAL_RESOURCE_ENTITLEMENT_ID_FILTER) + --external-resource-traits strings Resource type traits (e.g. "user", "group", "app") to sync and match from the external resource c1z. When unset the matcher falls back to user and group; passing this flag replaces the full set rather than adding to it. ($BATON_EXTERNAL_RESOURCE_TRAITS) + -f, --file string The path to the c1z file to sync with ($BATON_FILE) (default "sync.c1z") + --health-check Enable the HTTP health check endpoint ($BATON_HEALTH_CHECK) + --health-check-port int Port for the HTTP health check endpoint ($BATON_HEALTH_CHECK_PORT) (default 8081) + -h, --help help for baton-databricks + --hostname string The Databricks hostname used to connect to the Databricks API ($BATON_HOSTNAME) (default "cloud.databricks.com") + --http-timeout-seconds int HTTP client timeout in seconds (max 1800) ($BATON_HTTP_TIMEOUT_SECONDS) (default 300) + --keep-previous-sync-c1z Keep the previously synced c1z on disk to enable ETag replay across service-mode syncs (requires a connector that supports ETag replay; costs one c1z of local disk) ($BATON_KEEP_PREVIOUS_SYNC_C1Z) + --log-format string The output format for logs: json, console ($BATON_LOG_FORMAT) (default "json") + --log-level string The log level: debug, info, warn, error ($BATON_LOG_LEVEL) (default "info") + --log-level-debug-expires-at string The timestamp indicating when debug-level logging should expire ($BATON_LOG_LEVEL_DEBUG_EXPIRES_AT) + --log-path strings The file path to write logs to ($BATON_LOG_PATH) + --otel-collector-endpoint string The endpoint of the OpenTelemetry collector to send observability data to (used for both tracing and logging if specific endpoints are not provided) ($BATON_OTEL_COLLECTOR_ENDPOINT) + --parallel-sync Deprecated: use --workers instead. ($BATON_PARALLEL_SYNC) + -p, --provisioning This must be set in order for provisioning actions to be enabled ($BATON_PROVISIONING) + --skip-entitlements-and-grants This must be set to skip syncing of entitlements and grants ($BATON_SKIP_ENTITLEMENTS_AND_GRANTS) + --skip-full-sync This must be set to skip a full sync ($BATON_SKIP_FULL_SYNC) + --storage-engine string The storage engine to use when opening the sync c1z file: sqlite or pebble. Leave unset to use the baton-sdk default. ($BATON_STORAGE_ENGINE) + --sync-resource-types strings The resource type IDs to sync ($BATON_SYNC_RESOURCE_TYPES) + --sync-resources strings The resource IDs to sync ($BATON_SYNC_RESOURCES) + --task-concurrency int The number of Baton tasks to run concurrently in service mode. Tasks may include sync, grant, revoke, and more. Minimum value is 1, maximum value is 100. ($BATON_TASK_CONCURRENCY) (default 3) + --ticketing This must be set to enable ticketing support ($BATON_TICKETING) + -v, --version version for baton-databricks + --workers int The number of sync workers to use. -1 for auto-detect, 0 for sequential, >0 for parallel ($BATON_WORKERS) + --workspace-tokens strings required: The Databricks personal access tokens scoped to specific workspaces used to connect to the Databricks Workspace API ($BATON_WORKSPACE_TOKENS) + --workspaces strings Limit syncing to the specified workspaces, by deployment name, not workspace ID. Required when using workspace tokens, in the same order as workspace-tokens. ($BATON_WORKSPACES) Use "baton-databricks [command] --help" for more information about a command. ``` diff --git a/config_schema.json b/config_schema.json index 194a841b..b38554d0 100644 --- a/config_schema.json +++ b/config_schema.json @@ -140,14 +140,82 @@ "defaultValue": "cloud.databricks.com" } }, + { + "name": "workspaces", + "displayName": "Workspaces", + "description": "Limit syncing to the specified workspaces, by deployment name, not workspace ID. Required when using workspace tokens, in the same order as workspace-tokens. Mutually exclusive with databricks-exclude-workspaces.", + "stringSliceField": {} + }, + { + "name": "workspace-tokens", + "displayName": "Workspace Tokens", + "description": "The Databricks personal access tokens scoped to specific workspaces used to connect to the Databricks Workspace API", + "isRequired": true, + "isSecret": true, + "stringSliceField": { + "rules": { + "isRequired": true + } + } + }, { "name": "databricks-exclude-workspaces", "displayName": "Exclude Workspaces", - "description": "Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID", + "description": "Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces.", "stringSliceField": {} } ], + "constraints": [ + { + "kind": "CONSTRAINT_KIND_MUTUALLY_EXCLUSIVE", + "fieldNames": [ + "workspaces", + "databricks-exclude-workspaces" + ] + }, + { + "kind": "CONSTRAINT_KIND_DEPENDENT_ON", + "fieldNames": [ + "workspace-tokens" + ], + "secondaryFieldNames": [ + "workspaces" + ] + } + ], "displayName": "Databricks", "helpUrl": "/docs/baton/databricks", - "iconUrl": "/static/app-icons/databricks.svg" + "iconUrl": "/static/app-icons/databricks.svg", + "fieldGroups": [ + { + "name": "oauth2", + "displayName": "OAuth2", + "helpText": "Authenticate as a service principal using an OAuth2 client ID and secret.", + "fields": [ + "account-id", + "databricks-client-id", + "databricks-client-secret", + "hostname", + "account-hostname", + "workspaces", + "base-url", + "databricks-exclude-workspaces" + ], + "default": true + }, + { + "name": "workspace-token", + "displayName": "Workspace token", + "helpText": "Authenticate with a personal access token scoped to each workspace.", + "fields": [ + "account-id", + "workspaces", + "workspace-tokens", + "hostname", + "account-hostname", + "base-url", + "databricks-exclude-workspaces" + ] + } + ] } \ No newline at end of file diff --git a/docs/connector.mdx b/docs/connector.mdx index 88ce529e..e76e8503 100644 --- a/docs/connector.mdx +++ b/docs/connector.mdx @@ -42,7 +42,7 @@ A user with the **Account admin** role in each Databricks workspace you want to ### Generate Databricks credentials -You have three authentication choices when setting up the Databricks connector: +You have two authentication choices when setting up the Databricks connector: - **OAuth** (syncs info from all Databricks workspaces) @@ -69,10 +69,6 @@ You have three authentication choices when setting up the Databricks connector: -- **Username and password** (syncs info from all Databricks workspaces) - - You do not need to generate any additional credentials to use this method. - **Done.** Here's the set of credentials you'll need when setting up the connector: - Account ID @@ -83,13 +79,7 @@ OR - Account ID - Personal access token -- Workspace ID for the Databricks workspace you're syncing - -OR - -- Account ID -- Username -- Password +- Deployment name of the Databricks workspace you're syncing (the subdomain in the workspace URL, not the workspace ID) Next, move on to the instructions for your chosen setup method. @@ -132,13 +122,13 @@ To complete this task, you'll need: Find the **Settings** area of the page and click **Edit**. - Select whether you're authenticating with **OAuth**, a **Personal access token**, or your **Username and password**. + Select whether you're authenticating with **OAuth** or a **Personal access token**. Paste the account ID you looked up in Step 1 into the **Account ID** field. - Enter the required OAuth, token, or username and password credentials into the other two fields. + Enter the required OAuth or token credentials into the other two fields. **Google Cloud Platform and Azure Databricks customers only:** Enter your Databricks account hostname and hostname in the relevant fields. @@ -226,13 +216,9 @@ stringData: # Databricks credentials, option 2 BATON_ACCOUNT_ID: + BATON_AUTH_METHOD: workspace-token BATON_WORKSPACE_TOKENS: - BATON_WORKSPACES: - - # Databricks credentials, option 3 - BATON_ACCOUNT_ID: - BATON_USERNAME: - BATON_PASSWORD: + BATON_WORKSPACES: # Optional: comma-separated workspaces to exclude from sync (workspace name, deployment name, or numeric ID) BATON_DATABRICKS_EXCLUDE_WORKSPACES: diff --git a/pkg/config/conf.gen.go b/pkg/config/conf.gen.go index 80ebf5b1..6ca5facf 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -9,6 +9,8 @@ type Databricks struct { DatabricksClientId string `mapstructure:"databricks-client-id"` DatabricksClientSecret string `mapstructure:"databricks-client-secret"` Hostname string `mapstructure:"hostname"` + Workspaces []string `mapstructure:"workspaces"` + WorkspaceTokens []string `mapstructure:"workspace-tokens"` BaseUrl string `mapstructure:"base-url"` DatabricksExcludeWorkspaces []string `mapstructure:"databricks-exclude-workspaces"` } diff --git a/pkg/config/config.go b/pkg/config/config.go index fb664cb3..2f27e771 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -1,9 +1,17 @@ package config import ( + "context" + "fmt" + "github.com/conductorone/baton-sdk/pkg/field" ) +const ( + DatabricksOAuth2Group = "oauth2" + DatabricksWorkspaceTokenGroup = "workspace-token" +) + var ( AccountIdField = field.StringField( "account-id", @@ -14,8 +22,8 @@ var ( DatabricksClientIdField = field.StringField( "databricks-client-id", field.WithDescription("The Databricks service principal's client ID used to connect to the Databricks Account and Workspace API"), - field.WithDisplayName("OAuth2 Client ID"), field.WithRequired(true), + field.WithDisplayName("OAuth2 Client ID"), ) DatabricksClientSecretField = field.StringField( "databricks-client-secret", @@ -24,6 +32,22 @@ var ( field.WithRequired(true), field.WithDisplayName("OAuth2 Client Secret"), ) + WorkspacesField = field.StringSliceField( + "workspaces", + field.WithDescription( + "Limit syncing to the specified workspaces, by deployment name, not workspace ID. "+ + "Required when using workspace tokens, in the same order as workspace-tokens. "+ + "Mutually exclusive with databricks-exclude-workspaces.", + ), + field.WithDisplayName("Workspaces"), + ) + WorkspaceTokensField = field.StringSliceField( + "workspace-tokens", + field.WithDescription("The Databricks personal access tokens scoped to specific workspaces used to connect to the Databricks Workspace API"), + field.WithIsSecret(true), + field.WithRequired(true), + field.WithDisplayName("Workspace Tokens"), + ) AccountHostnameField = field.StringField( "account-hostname", field.WithDescription("The hostname used to connect to the Databricks account API. If not set, it will be calculated from the hostname field."), @@ -43,7 +67,7 @@ var ( ) ExcludeWorkspacesField = field.StringSliceField( "databricks-exclude-workspaces", - field.WithDescription("Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID"), + field.WithDescription("Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces."), field.WithDisplayName("Exclude Workspaces"), ) configFields = []field.SchemaField{ @@ -52,6 +76,8 @@ var ( DatabricksClientIdField, DatabricksClientSecretField, HostnameField, + WorkspacesField, + WorkspaceTokensField, BaseURLField, ExcludeWorkspacesField, } @@ -63,4 +89,47 @@ var Config = field.NewConfiguration( field.WithConnectorDisplayName("Databricks"), field.WithHelpUrl("/docs/baton/databricks"), field.WithIconUrl("/static/app-icons/databricks.svg"), + field.WithConstraints( + field.FieldsMutuallyExclusive(WorkspacesField, ExcludeWorkspacesField), + field.FieldsDependentOn([]field.SchemaField{WorkspaceTokensField}, []field.SchemaField{WorkspacesField}), + ), + field.WithFieldGroups([]field.SchemaFieldGroup{ + { + Name: DatabricksOAuth2Group, + DisplayName: "OAuth2", + HelpText: "Authenticate as a service principal using an OAuth2 client ID and secret.", + Fields: []field.SchemaField{ + AccountIdField, DatabricksClientIdField, DatabricksClientSecretField, + HostnameField, AccountHostnameField, WorkspacesField, BaseURLField, ExcludeWorkspacesField, + }, + Default: true, + }, + { + Name: DatabricksWorkspaceTokenGroup, + DisplayName: "Workspace token", + HelpText: "Authenticate with a personal access token scoped to each workspace.", + Fields: []field.SchemaField{AccountIdField, WorkspacesField, WorkspaceTokensField, HostnameField, AccountHostnameField, BaseURLField, ExcludeWorkspacesField}, + Default: false, + }, + }), ) + +// ValidateConfig enforces what field groups can't: OAuth/token exclusion when no +// auth method is set, and equal-length workspaces/workspace-tokens. +func ValidateConfig(ctx context.Context, cfg *Databricks, authMethod string) error { + // A merged/stored config can carry both groups' fields; once authMethod picks one, + // prepareClientAuth only reads that group, so the other group's leftovers are inert. + if authMethod == "" && len(cfg.WorkspaceTokens) > 0 && (cfg.DatabricksClientId != "" || cfg.DatabricksClientSecret != "") { + return fmt.Errorf("databricks-connector: databricks-client-id/databricks-client-secret and workspace-tokens are mutually exclusive") + } + + if authMethod == DatabricksWorkspaceTokenGroup && len(cfg.Workspaces) != len(cfg.WorkspaceTokens) { + return fmt.Errorf( + "databricks-connector: workspaces and workspace-tokens must be the same length, got %d workspaces and %d tokens", + len(cfg.Workspaces), + len(cfg.WorkspaceTokens), + ) + } + + return nil +} diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go new file mode 100644 index 00000000..c88c848f --- /dev/null +++ b/pkg/config/config_test.go @@ -0,0 +1,82 @@ +package config + +import ( + "context" + "testing" +) + +func TestValidateConfig(t *testing.T) { + cases := []struct { + name string + workspaces []string + tokens []string + authMethod string + wantErr bool + }{ + {"no tokens", nil, nil, DatabricksWorkspaceTokenGroup, false}, + {"equal length", []string{"ws-1", "ws-2"}, []string{"tok-1", "tok-2"}, DatabricksWorkspaceTokenGroup, false}, + {"more workspaces than tokens", []string{"ws-1", "ws-2"}, []string{"tok-1"}, DatabricksWorkspaceTokenGroup, true}, + {"tokens without workspaces", nil, []string{"tok-1"}, DatabricksWorkspaceTokenGroup, true}, + {"mismatched lengths ignored outside workspace-token method", []string{"ws-1", "ws-2"}, []string{"tok-1"}, DatabricksOAuth2Group, false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cfg := &Databricks{Workspaces: tc.workspaces, WorkspaceTokens: tc.tokens} + err := ValidateConfig(context.Background(), cfg, tc.authMethod) + if tc.wantErr && err == nil { + t.Fatal("expected error, got nil") + } + if !tc.wantErr && err != nil { + t.Fatalf("expected no error, got %v", err) + } + }) + } +} + +// Both auth modes' fields live in the same struct; with no auth method selected, +// ValidateConfig can't tell which credentials would actually be used, so it must +// reject having both set. +func TestValidateConfigRejectsBothAuthModesWhenAmbiguous(t *testing.T) { + cfg := &Databricks{ + DatabricksClientId: "client-id", + Workspaces: []string{"ws-1"}, + WorkspaceTokens: []string{"tok-1"}, + } + + if err := ValidateConfig(context.Background(), cfg, ""); err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestValidateConfigRejectsClientSecretWithTokensWhenAmbiguous(t *testing.T) { + cfg := &Databricks{ + DatabricksClientSecret: "client-secret", + Workspaces: []string{"ws-1"}, + WorkspaceTokens: []string{"tok-1"}, + } + + if err := ValidateConfig(context.Background(), cfg, ""); err == nil { + t.Fatal("expected error, got nil") + } +} + +// Once an auth method is explicitly selected, prepareClientAuth only reads that +// group's fields, so leftover values from the other group (e.g. stale OAuth +// creds after switching to workspace tokens, or vice versa) must not block startup. +func TestValidateConfigTrustsExplicitAuthMethod(t *testing.T) { + cfg := &Databricks{ + DatabricksClientId: "client-id", + DatabricksClientSecret: "client-secret", + Workspaces: []string{"ws-1"}, + WorkspaceTokens: []string{"tok-1"}, + } + + if err := ValidateConfig(context.Background(), cfg, DatabricksWorkspaceTokenGroup); err != nil { + t.Fatalf("expected no error, got %v", err) + } + + if err := ValidateConfig(context.Background(), cfg, DatabricksOAuth2Group); err != nil { + t.Fatalf("expected no error, got %v", err) + } +} diff --git a/pkg/connector/account.go b/pkg/connector/account.go index 14dc3d80..96aa8f2c 100644 --- a/pkg/connector/account.go +++ b/pkg/connector/account.go @@ -42,6 +42,7 @@ func (a *accountBuilder) ResourceType(ctx context.Context) *v2.ResourceType { return accountResourceType } +// The Account API check below mirrors groupGrantParent (helpers.go); keep both in sync. func (a *accountBuilder) accountResource(_ context.Context) (*v2.Resource, error) { accountId := a.client.GetAccountId() children := []protoreflect.ProtoMessage{ @@ -130,7 +131,13 @@ func (a *accountBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs var annotations []protoreflect.ProtoMessage if resourceId.ResourceType == groupResourceType.Id { - rid, expandAnnotation, err := groupGrantExpansion(ctx, resourceId.Resource, resource.ParentResourceId) + // Grants already returned early above when the account API is unavailable, + // so groups reaching this point are always account-parented. + groupParentResourceId, err := rs.NewResourceID(accountResourceType, a.client.GetAccountId()) + if err != nil { + return rv, nil, err + } + rid, expandAnnotation, err := groupGrantExpansion(ctx, resourceId.Resource, groupParentResourceId) if err != nil { return rv, nil, err } diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 617e60a0..491dd6bd 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -16,7 +16,8 @@ import ( ) type Databricks struct { - client *databricks.Client + client *databricks.Client + workspaces []string } // ResourceSyncers returns a ResourceSyncerV2 for each resource type that should be synced from the upstream service. @@ -26,7 +27,7 @@ func (d *Databricks) ResourceSyncers(ctx context.Context) []connectorbuilder.Res newGroupBuilder(d.client), newServicePrincipalBuilder(d.client), newUserBuilder(d.client), - newWorkspaceBuilder(d.client), + newWorkspaceBuilder(d.client, d.workspaces), newRoleBuilder(d.client), } @@ -108,22 +109,33 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err isAccAPIAvailable := false isWSAPIAvailable := false - // Check if we can list users from Account API. - _, _, err := d.client.ListRoles(ctx, "", "", "") - if err == nil { - isAccAPIAvailable = true + // The Account API is unreachable with workspace tokens, so only probe it for OAuth. + if !d.client.IsTokenAuth() { + _, _, err := d.client.ListRoles(ctx, "", "", "") + if err == nil { + isAccAPIAvailable = true + } } - // Validate that credentials are valid for every workspace. - workspaces, _, err := d.client.ListWorkspaces(ctx) - if err != nil { - return nil, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err) + // With an explicit workspace list (always the case for token auth), validate each + // configured workspace. Otherwise discover every workspace from the Account API. + workspaceNames := d.workspaces + if len(workspaceNames) == 0 { + workspaces, _, err := d.client.ListWorkspaces(ctx) + if err != nil { + return nil, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err) + } + + workspaceNames = make([]string, 0, len(workspaces)) + for _, workspace := range workspaces { + workspaceNames = append(workspaceNames, workspace.DeploymentName) + } } - for _, workspace := range workspaces { - _, _, err := d.client.ListRoles(ctx, workspace.DeploymentName, "", "") + for _, workspace := range workspaceNames { + _, _, err := d.client.ListRoles(ctx, workspace, "", "") if err != nil && !isAccAPIAvailable { - return nil, fmt.Errorf("databricks-connector: failed to validate credentials for workspace %s: %w", workspace.DeploymentName, err) + return nil, fmt.Errorf("databricks-connector: failed to validate credentials for workspace %s: %w", workspace, err) } isWSAPIAvailable = true @@ -148,6 +160,7 @@ func New( baseURL string, auth databricks.Auth, excludeWorkspaces []string, + workspaces []string, ) (*Databricks, error) { httpClient, err := auth.GetClient(ctx) if err != nil { @@ -160,7 +173,8 @@ func New( } return &Databricks{ - client: client, + client: client, + workspaces: workspaces, }, nil } @@ -168,8 +182,17 @@ func New( func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.ConnectorOpts) (connectorbuilder.ConnectorBuilderV2, []connectorbuilder.Opt, error) { l := ctxzap.Extract(ctx) + authMethod := "" + if opts != nil { + authMethod = opts.SelectedAuthMethod + } + + if err := config.ValidateConfig(ctx, cfg, authMethod); err != nil { + return nil, nil, err + } + accountHostname := getAccountHostname(cfg, cfg.Hostname) - auth := prepareClientAuth(ctx, cfg, l) + auth := prepareClientAuth(ctx, cfg, authMethod, l) cb, err := New( ctx, @@ -179,26 +202,27 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.Connect cfg.BaseUrl, auth, cfg.DatabricksExcludeWorkspaces, + cfg.Workspaces, ) if err != nil { - l.Warn("error creating connector", zap.Error(err)) return nil, nil, err } return cb, nil, nil } -func prepareClientAuth(_ context.Context, cfg *config.Databricks, l *zap.Logger) databricks.Auth { - accountID := cfg.AccountId - databricksClientId := cfg.DatabricksClientId - databricksClientSecret := cfg.DatabricksClientSecret - accountHostname := getAccountHostname(cfg, cfg.Hostname) +func prepareClientAuth(_ context.Context, cfg *config.Databricks, authMethod string, l *zap.Logger) databricks.Auth { + if authMethod == config.DatabricksWorkspaceTokenGroup { + l.Debug("using workspace token auth", zap.String("account-id", cfg.AccountId)) + return databricks.NewTokenAuth(cfg.Workspaces, cfg.WorkspaceTokens) + } + l.Debug("using oauth", zap.String("account-id", cfg.AccountId)) return databricks.NewOAuth2( - accountID, - databricksClientId, - databricksClientSecret, - accountHostname, + cfg.AccountId, + cfg.DatabricksClientId, + cfg.DatabricksClientSecret, + getAccountHostname(cfg, cfg.Hostname), ) } diff --git a/pkg/connector/groups.go b/pkg/connector/groups.go index 4091292d..7d223f8b 100644 --- a/pkg/connector/groups.go +++ b/pkg/connector/groups.go @@ -149,6 +149,12 @@ func (g *groupBuilder) Entitlements(ctx context.Context, resource *v2.Resource, // get all assignable roles for this specific group resource roles, _, err := g.client.ListRoles(ctx, workspaceId, GroupsType, groupId.Resource) if err != nil { + if workspaceId != "" && isGroupNotFoundError(err) { + ctxzap.Extract(ctx).Warn("databricks-connector: skipping roles for group not recognized by the rule-sets API", + zap.String("group_id", groupId.Resource), + ) + return rv, nil, nil + } return nil, nil, fmt.Errorf("databricks-connector: failed to list roles for group %s: %w", groupId.Resource, err) } @@ -230,6 +236,12 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.S // role permissions grants ruleSets, rateLimitDataRuleSets, err := g.client.ListRuleSets(ctx, workspaceId, GroupsType, groupId.Resource) if err != nil { + if isWorkspaceGroup && isGroupNotFoundError(err) { + l.Warn("databricks-connector: skipping role rule sets for group not recognized by the rule-sets API", + zap.String("group_id", groupId.Resource), + ) + return rv, &rs.SyncOpResults{Annotations: annos}, nil + } return nil, nil, fmt.Errorf("databricks-connector: failed to list role rule sets for group %s: %w", resource.Id.Resource, err) } diff --git a/pkg/connector/helpers.go b/pkg/connector/helpers.go index 54874a57..be1db76f 100644 --- a/pkg/connector/helpers.go +++ b/pkg/connector/helpers.go @@ -2,7 +2,9 @@ package connector import ( "context" + "errors" "fmt" + "net/http" "slices" "strings" @@ -35,6 +37,17 @@ func parseResourceId(resourceId string) (*v2.ResourceId, *v2.ResourceId, error) return nil, nil, fmt.Errorf("invalid resource ID: %s", resourceId) } +// Mirrors how groupBuilder parents synced groups: account when its API is +// reachable, otherwise the workspace (token auth). accountResource() in +// account.go encodes the same condition for its child-resource-type list; +// keep both in sync. +func groupGrantParent(accountAPIAvailable bool, accountId, workspaceId string) (*v2.ResourceId, error) { + if accountAPIAvailable { + return rs.NewResourceID(accountResourceType, accountId) + } + return rs.NewResourceID(workspaceResourceType, workspaceId) +} + func groupGrantExpansion(ctx context.Context, groupId string, parentResource *v2.ResourceId) (*v2.ResourceId, *v2.GrantExpandable, error) { groupResourceStr := groupResourceId(ctx, groupId, parentResource) resourceId, err := rs.NewResourceID(groupResourceType, groupResourceStr) @@ -144,6 +157,20 @@ func preparePrincipalId(ctx context.Context, c *databricks.Client, workspaceId, return result, nil } +// isGroupNotFoundError matches the rule-sets/roles API's response for a group ID +// it doesn't recognize (e.g. an orphaned or stale workspace SCIM group), distinct +// from other 400s. +func isGroupNotFoundError(err error) bool { + var apiErr *databricks.APIError + if !errors.As(err, &apiErr) { + return false + } + msg := strings.ToLower(apiErr.Message) + return apiErr.StatusCode == http.StatusBadRequest && + strings.Contains(msg, "not found") && + strings.Contains(msg, "group") +} + func isValidPrincipal(principal *v2.ResourceId) bool { return principal.ResourceType == userResourceType.Id || principal.ResourceType == groupResourceType.Id || diff --git a/pkg/connector/helpers_test.go b/pkg/connector/helpers_test.go new file mode 100644 index 00000000..54aaa341 --- /dev/null +++ b/pkg/connector/helpers_test.go @@ -0,0 +1,99 @@ +package connector + +import ( + "context" + "errors" + "net/http" + "testing" + + "github.com/conductorone/baton-databricks/pkg/databricks" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" +) + +// CXH-2166 regression: under token auth (no account API), groups sync parented +// under the workspace. Role grants to those groups must use the same parent, or +// the grant's principal ID references a group resource that was never synced. +func TestGroupGrantParentMatchesSyncedGroupId(t *testing.T) { + ctx := context.Background() + + t.Run("token auth uses workspace parent", func(t *testing.T) { + parent, err := groupGrantParent(false, "acc-1", "dbc-abc") + if err != nil { + t.Fatalf("groupGrantParent: %v", err) + } + + gotResourceId, _, err := groupGrantExpansion(ctx, "group-1", parent) + if err != nil { + t.Fatalf("groupGrantExpansion: %v", err) + } + + wantId := groupResourceId(ctx, "group-1", &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: "dbc-abc"}) + if gotResourceId.Resource != wantId { + t.Errorf("principal ID = %q, want %q (the ID groupBuilder emits for a workspace-parented group)", gotResourceId.Resource, wantId) + } + }) + + t.Run("account API available uses account parent", func(t *testing.T) { + parent, err := groupGrantParent(true, "acc-1", "dbc-abc") + if err != nil { + t.Fatalf("groupGrantParent: %v", err) + } + + gotResourceId, _, err := groupGrantExpansion(ctx, "group-1", parent) + if err != nil { + t.Fatalf("groupGrantExpansion: %v", err) + } + + wantId := groupResourceId(ctx, "group-1", &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: "acc-1"}) + if gotResourceId.Resource != wantId { + t.Errorf("principal ID = %q, want %q (the ID groupBuilder emits for an account-parented group)", gotResourceId.Resource, wantId) + } + }) +} + +func TestIsGroupNotFoundError(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + { + name: "matching group not found", + err: &databricks.APIError{StatusCode: http.StatusBadRequest, Message: "Group 12345 not found"}, + want: true, + }, + { + name: "non-matching 400", + err: &databricks.APIError{StatusCode: http.StatusBadRequest, Message: "invalid role name"}, + want: false, + }, + { + name: "unrelated not-found 400 without group in message", + err: &databricks.APIError{StatusCode: http.StatusBadRequest, Message: "workspace not found"}, + want: false, + }, + { + name: "404 status code", + err: &databricks.APIError{StatusCode: http.StatusNotFound, Message: "Group 12345 not found"}, + want: false, + }, + { + name: "mixed case still matches", + err: &databricks.APIError{StatusCode: http.StatusBadRequest, Message: "GROUP 12345 Not Found"}, + want: true, + }, + { + name: "non-APIError", + err: errors.New("connection reset"), + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isGroupNotFoundError(tt.err); got != tt.want { + t.Errorf("isGroupNotFoundError() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/pkg/connector/roles.go b/pkg/connector/roles.go index abed3569..8fb0ec1f 100644 --- a/pkg/connector/roles.go +++ b/pkg/connector/roles.go @@ -226,12 +226,11 @@ func (r *roleBuilder) Grants(ctx context.Context, resource *v2.Resource, attr rs } if (!isWorkspaceRole && g.HaveRole(roleName)) || (isWorkspaceRole && g.HaveEntitlement(roleName)) { - accountId := r.client.GetAccountId() - accountResourceId, err := rs.NewResourceID(accountResourceType, accountId) + groupParentResourceId, err := groupGrantParent(r.client.IsAccountAPIAvailable(), r.client.GetAccountId(), workspaceId) if err != nil { return rv, nil, err } - resourceId, expandAnnotation, err := groupGrantExpansion(ctx, g.ID, accountResourceId) + resourceId, expandAnnotation, err := groupGrantExpansion(ctx, g.ID, groupParentResourceId) if err != nil { return rv, nil, err } diff --git a/pkg/connector/service-principals.go b/pkg/connector/service-principals.go index 62b92868..a51d383c 100644 --- a/pkg/connector/service-principals.go +++ b/pkg/connector/service-principals.go @@ -197,7 +197,8 @@ func (s *servicePrincipalBuilder) Grants(ctx context.Context, resource *v2.Resou var annotations []protoreflect.ProtoMessage if resourceId.ResourceType == groupResourceType.Id { - groupResourceStr := groupResourceId(ctx, resourceId.Resource, resource.ParentResourceId) + groupParentResourceId := &v2.ResourceId{ResourceType: parentType, Resource: parentID} + groupResourceStr := groupResourceId(ctx, resourceId.Resource, groupParentResourceId) annotations = append(annotations, &v2.GrantExpandable{ EntitlementIds: []string{fmt.Sprintf("group:%s:%s", groupResourceStr, groupMemberEntitlement)}, }) diff --git a/pkg/connector/workspaces.go b/pkg/connector/workspaces.go index 52ed5d92..8a594192 100644 --- a/pkg/connector/workspaces.go +++ b/pkg/connector/workspaces.go @@ -24,12 +24,34 @@ const workspaceMemberEntitlement = "member" type workspaceBuilder struct { client *databricks.Client resourceType *v2.ResourceType + workspaces map[string]struct{} } func (w *workspaceBuilder) ResourceType(ctx context.Context) *v2.ResourceType { return workspaceResourceType } +// minimalWorkspaceResource builds a workspace from just its deployment name, for +// token auth where the Account API (and its numeric workspace IDs) is unreachable. +// Deployment names are unique per Databricks cloud (they form the workspace's +// canonical hostname), so they're safe as the resource ID here. +// Users, groups and service principals hang off the workspace here instead of the account. +func minimalWorkspaceResource(_ context.Context, workspace *databricks.Workspace, parent *v2.ResourceId) (*v2.Resource, error) { + return rs.NewGroupResource( + workspace.DeploymentName, + workspaceResourceType, + workspace.DeploymentName, + nil, + rs.WithParentResourceID(parent), + rs.WithAnnotation( + &v2.ChildResourceType{ResourceTypeId: userResourceType.Id}, + &v2.ChildResourceType{ResourceTypeId: groupResourceType.Id}, + &v2.ChildResourceType{ResourceTypeId: servicePrincipalResourceType.Id}, + &v2.ChildResourceType{ResourceTypeId: roleResourceType.Id}, + ), + ) +} + func workspaceResource(_ context.Context, workspace *databricks.Workspace, parent *v2.ResourceId) (*v2.Resource, error) { profile := map[string]interface{}{ "workspace_id": workspace.ID, @@ -62,12 +84,47 @@ func (w *workspaceBuilder) List(ctx context.Context, parentResourceID *v2.Resour var rv []*v2.Resource + if w.client.IsTokenAuth() { + for workspace := range w.workspaces { + if w.client.IsWorkspaceNameExcluded(workspace) { + continue + } + + ws := &databricks.Workspace{DeploymentName: workspace} + + wr, err := minimalWorkspaceResource(ctx, ws, parentResourceID) + if err != nil { + return nil, nil, err + } + + rv = append(rv, wr) + } + + if len(w.workspaces) > 0 && len(rv) == 0 { + ctxzap.Extract(ctx).Warn("databricks-connector: all configured workspaces are excluded, sync will be empty", + zap.Strings("workspaces", configuredWorkspaceNames(w.workspaces)), + ) + } + + return rv, nil, nil + } + workspaces, _, err := w.client.ListWorkspaces(ctx) if err != nil { return nil, nil, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err) } + matchedConfigured := make(map[string]struct{}, len(w.workspaces)) for _, workspace := range workspaces { + // Skip workspaces outside the configured set when one was provided. + if len(w.workspaces) > 0 { + cfg, ok := matchConfiguredWorkspace(w.workspaces, workspace.DeploymentName, workspace.Name, strconv.Itoa(workspace.ID)) + if !ok { + continue + } + matchedConfigured[cfg] = struct{}{} + } + wCopy := workspace wr, err := workspaceResource(ctx, &wCopy, parentResourceID) @@ -78,9 +135,58 @@ func (w *workspaceBuilder) List(ctx context.Context, parentResourceID *v2.Resour rv = append(rv, wr) } + l := ctxzap.Extract(ctx) + if len(w.workspaces) > 0 && len(matchedConfigured) == 0 { + l.Warn("databricks-connector: none of the configured workspaces matched any account workspace, sync will be empty", + zap.Strings("workspaces", configuredWorkspaceNames(w.workspaces)), + ) + } + for workspace := range w.workspaces { + if _, ok := matchedConfigured[workspace]; ok { + continue + } + if w.client.IsWorkspaceNameExcluded(workspace) { + l.Debug("databricks-connector: configured workspace was excluded from sync", + zap.String("workspace", workspace), + ) + continue + } + l.Debug("databricks-connector: configured workspace not found among account workspaces", + zap.String("workspace", workspace), + ) + } + return rv, nil, nil } +func configuredWorkspaceNames(configured map[string]struct{}) []string { + names := make([]string, 0, len(configured)) + for name := range configured { + names = append(names, name) + } + return names +} + +// matchConfiguredWorkspace looks up a workspace by deployment name, name, or numeric +// ID case-insensitively (mirroring Client.IsWorkspaceNameExcluded), returning the matched +// key so warnings can report the value the user configured. +func matchConfiguredWorkspace(configured map[string]struct{}, candidates ...string) (string, bool) { + for _, candidate := range candidates { + if candidate == "" { + continue + } + if _, ok := configured[candidate]; ok { + return candidate, true + } + for cfg := range configured { + if strings.EqualFold(cfg, candidate) { + return cfg, true + } + } + } + return "", false +} + // Entitlements returns slice of entitlements representing workspace members. // To get workspace members, we can only use the account API. func (w *workspaceBuilder) Entitlements(_ context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Entitlement, *rs.SyncOpResults, error) { @@ -239,9 +345,15 @@ func (w *workspaceBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotat return nil, nil } -func newWorkspaceBuilder(client *databricks.Client) *workspaceBuilder { +func newWorkspaceBuilder(client *databricks.Client, workspaces []string) *workspaceBuilder { + wMap := make(map[string]struct{}, len(workspaces)) + for _, w := range workspaces { + wMap[w] = struct{}{} + } + return &workspaceBuilder{ client: client, resourceType: workspaceResourceType, + workspaces: wMap, } } diff --git a/pkg/databricks/auth.go b/pkg/databricks/auth.go index 63cda5bc..1ef846cf 100644 --- a/pkg/databricks/auth.go +++ b/pkg/databricks/auth.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "strings" "github.com/conductorone/baton-sdk/pkg/uhttp" "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" @@ -29,6 +30,52 @@ func (n *NoAuth) GetClient(ctx context.Context) (*http.Client, error) { return httpClient, nil } +// TokenAuth authenticates each request with the workspace-scoped personal access +// token for the workspace it targets. Account-level requests match no token. +type TokenAuth struct { + tokens map[string]string +} + +func NewTokenAuth(workspaces, tokens []string) *TokenAuth { + tokensMap := make(map[string]string, len(workspaces)) + for i, workspace := range workspaces { + if i >= len(tokens) { + break + } + tokensMap[workspace] = tokens[i] + } + + return &TokenAuth{tokens: tokensMap} +} + +func (t *TokenAuth) Apply(req *http.Request) { + // A workspace request host is ".". A shorter + // deployment name can be a false prefix of a longer one (Azure names + // contain a dot, e.g. "adb-123" of "adb-123.1"), so match the longest one. + host := req.URL.Host + var bestWorkspace, bestToken string + for workspace, token := range t.tokens { + if host != workspace && !strings.HasPrefix(host, workspace+".") { + continue + } + if len(workspace) > len(bestWorkspace) { + bestWorkspace, bestToken = workspace, token + } + } + if bestToken != "" { + req.Header.Set("Authorization", "Bearer "+bestToken) + } +} + +func (t *TokenAuth) GetClient(ctx context.Context) (*http.Client, error) { + httpClient, err := uhttp.NewClient(ctx, uhttp.WithLogger(true, ctxzap.Extract(ctx))) + if err != nil { + return nil, err + } + + return httpClient, nil +} + type OAuth2 struct { cfg *clientcredentials.Config } diff --git a/pkg/databricks/auth_test.go b/pkg/databricks/auth_test.go new file mode 100644 index 00000000..3c19b8d1 --- /dev/null +++ b/pkg/databricks/auth_test.go @@ -0,0 +1,87 @@ +package databricks + +import ( + "net/http" + "net/url" + "testing" +) + +func mustURL(t *testing.T, raw string) *url.URL { + t.Helper() + u, err := url.Parse(raw) + if err != nil { + t.Fatalf("parse %q: %v", raw, err) + } + return u +} + +func TestTokenAuthApply(t *testing.T) { + auth := NewTokenAuth( + []string{"dbc-abc123", "adb-2531901403506481.1"}, + []string{"aws-token", "azure-token"}, + ) + + cases := []struct { + name string + host string + wantToken string + }{ + {"aws deployment name (no dot)", "dbc-abc123.cloud.databricks.com", "aws-token"}, + {"azure deployment name (dotted)", "adb-2531901403506481.1.azuredatabricks.net", "azure-token"}, + {"account host matches nothing", "accounts.azuredatabricks.net", ""}, + {"unknown workspace matches nothing", "dbc-other.cloud.databricks.com", ""}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + req := &http.Request{URL: mustURL(t, "https://"+tc.host+"/api/2.0/preview/scim/v2/Users"), Header: http.Header{}} + auth.Apply(req) + + got := req.Header.Get("Authorization") + want := "" + if tc.wantToken != "" { + want = "Bearer " + tc.wantToken + } + if got != want { + t.Fatalf("Authorization = %q, want %q", got, want) + } + }) + } +} + +// A workspace name that prefixes another must not steal the longer one's token. +func TestTokenAuthApplyPrefixCollision(t *testing.T) { + auth := NewTokenAuth([]string{"dbc-1", "dbc-12"}, []string{"token-1", "token-12"}) + + req := &http.Request{URL: mustURL(t, "https://dbc-12.cloud.databricks.com/x"), Header: http.Header{}} + auth.Apply(req) + + if got := req.Header.Get("Authorization"); got != "Bearer token-12" { + t.Fatalf("Authorization = %q, want %q", got, "Bearer token-12") + } +} + +// An Azure deployment name's own dot must not let a shorter workspace name +// falsely prefix a longer one that embeds it (e.g. "adb-123" of "adb-123.1"). +func TestTokenAuthApplyNestedDottedPrefix(t *testing.T) { + auth := NewTokenAuth([]string{"adb-123", "adb-123.1"}, []string{"token-short", "token-long"}) + + req := &http.Request{URL: mustURL(t, "https://adb-123.1.azuredatabricks.net/x"), Header: http.Header{}} + auth.Apply(req) + + if got := req.Header.Get("Authorization"); got != "Bearer token-long" { + t.Fatalf("Authorization = %q, want %q", got, "Bearer token-long") + } +} + +// Fewer tokens than workspaces must not panic; unmatched workspaces just get no token. +func TestNewTokenAuthFewerTokensThanWorkspaces(t *testing.T) { + auth := NewTokenAuth([]string{"dbc-1", "dbc-2"}, []string{"token-1"}) + + req := &http.Request{URL: mustURL(t, "https://dbc-2.cloud.databricks.com/x"), Header: http.Header{}} + auth.Apply(req) + + if got := req.Header.Get("Authorization"); got != "" { + t.Fatalf("Authorization = %q, want empty", got) + } +} diff --git a/pkg/databricks/client.go b/pkg/databricks/client.go index fcdd1890..76710023 100644 --- a/pkg/databricks/client.go +++ b/pkg/databricks/client.go @@ -118,6 +118,18 @@ func (c *Client) isWorkspaceExcluded(w Workspace) ([]string, bool) { return keys, len(keys) > 0 } +// IsWorkspaceNameExcluded reports whether deploymentName matches the +// databricks-exclude-workspaces set. Checks the name only, not via +// isWorkspaceExcluded: that also matches on ID, and a zero-value ID here would +// let an exclude entry of "0" match every workspace. +func (c *Client) IsWorkspaceNameExcluded(deploymentName string) bool { + if len(c.excludeWorkspaces) == 0 { + return false + } + _, ok := c.excludeWorkspaces[strings.ToLower(deploymentName)] + return ok +} + func (c *Client) workspaceUrl(workspaceId string) *url.URL { return &url.URL{ Scheme: "https", @@ -138,6 +150,11 @@ func (c *Client) UpdateAvailability(accAPI, wsAPI bool) { c.isWSAPIAvailable = wsAPI } +func (c *Client) IsTokenAuth() bool { + _, ok := c.auth.(*TokenAuth) + return ok +} + func (c *Client) UpdateEtag(etag string) { c.etag = etag }