Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion .github/workflows/ci.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ jobs:
- name: Test grant/revoking entitlements
env:
BATON_DATABRICKS: ./baton-databricks
uses: ConductorOne/github-workflows/actions/sync-test@v2
uses: ConductorOne/github-workflows/actions/sync-test@v4

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟠 Bug: sync-test@v4 adds an "Auth Error" step that runs the connector with invalidated credentials and requires exit code 16 (Unauthenticated) or 7 (PermissionDenied); this connector exits 2 (Unknown), so the bump makes the test job fail deterministically (run log: Connector exited 2 with invalid credentials; expected Unauthenticated (16) or PermissionDenied (7)).

Root cause is pkg/connector/connector.go:126 (and the sibling error paths at :138/:146) returning a bare fmt.Errorf, which gRPC maps to Unknown. Wrap credential failures with a status code — e.g. uhttp.WrapErrors(codes.Unauthenticated, "databricks-connector: failed to list workspaces", err) — or hold the bump until that's done.

with:
connector: ./baton-databricks
baton-entitlement: 'group:account/8c6f99ec-78ce-4654-8f92-e716b3dd67a7/group/66271399072731:member'
Expand Down
29 changes: 19 additions & 10 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -85,20 +85,29 @@ 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. You can also provide a list of
workspace access tokens to the `--workspace-tokens` flag. This will limit the
sync to only workspaces that are associated with those tokens. You can also use
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.
By default (OAuth), the connector fetches all resources from the account and all
workspaces. To limit the scope, pass a comma-separated list of workspace
deployment names to the `--workspaces` flag.

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.
the OAuth fields by default and rejects the config. In this mode `--workspaces`
is required and pairs with `--workspace-tokens` by position: the Nth token
authenticates the Nth workspace, so the two lists must have the same length and
order. `--workspaces` and `--workspace-tokens` are not unioned, and a
length mismatch is rejected at startup.

Workspace-token auth cannot reach the Databricks account API, so it syncs less
than OAuth: no account-level entitlements or grants and no workspace-membership
entitlements, and users, groups, service principals and roles are parented under
their workspace instead of the account. The connector logs a warning at startup
when the account API is unreachable. Use OAuth for full account and
cross-workspace coverage.

OAuth requires a reachable account API. If the account API check fails at
startup, the connector fails validation instead of falling back to a
workspace-only sync, even when `--workspaces` is set.

To instead exclude specific workspaces from the sync, pass them to the
`--databricks-exclude-workspaces` flag (or the
Expand Down
29 changes: 21 additions & 8 deletions docs/connector.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -214,24 +214,37 @@ stringData:
BATON_CLIENT_ID: <C1 client ID>
BATON_CLIENT_SECRET: <C1 client secret>

# Databricks credentials, option 1
# Databricks account ID (required for both credential options below)
BATON_ACCOUNT_ID: <Databricks account ID>

# Databricks credentials, option 1: OAuth (syncs the account and all workspaces)
BATON_DATABRICKS_CLIENT_ID: <OAuth client ID>
BATON_DATABRICKS_CLIENT_SECRET: <OAuth client secret>
# Optional (OAuth only): comma-separated workspaces to exclude from sync
# (workspace name, deployment name, or numeric ID). Cannot be combined with BATON_WORKSPACES.
BATON_DATABRICKS_EXCLUDE_WORKSPACES: <workspace-a,workspace-b>

# Databricks credentials, option 2
BATON_ACCOUNT_ID: <Databricks account ID>
# Databricks credentials, option 2: workspace token (syncs only the listed workspaces)
BATON_AUTH_METHOD: workspace-token
BATON_WORKSPACE_TOKENS: <Personal access token>
BATON_WORKSPACES: <Deployment name of the Databricks workspace you're syncing, not the workspace ID>

# Optional: comma-separated workspaces to exclude from sync (workspace name, deployment name, or numeric ID)
BATON_DATABRICKS_EXCLUDE_WORKSPACES: <workspace-a,workspace-b>
# BATON_WORKSPACE_TOKENS and BATON_WORKSPACES pair up by position: the Nth token
# authenticates the Nth workspace, so both lists must have the same length and order.
BATON_WORKSPACE_TOKENS: <token for workspace-a,token for workspace-b>
BATON_WORKSPACES: <deployment name of workspace-a,deployment name of workspace-b>

# Optional: include if you want C1 to provision access using this connector
BATON_PROVISIONING: true
```

<Note>
**Choose one credential option, not both.** Use OAuth (option 1) or workspace tokens (option 2). `BATON_DATABRICKS_EXCLUDE_WORKSPACES` is mutually exclusive with `BATON_WORKSPACES`, so it only applies to OAuth.
</Note>

<Warning>
**Workspace-token auth syncs less than OAuth.** A workspace token cannot reach the Databricks account API, so the sync is scoped to the listed workspaces only: it emits no account-level entitlements or grants and no workspace-membership entitlements, and users, groups, service principals and roles are parented under their workspace instead of the account. For full account and cross-workspace coverage, use OAuth.

OAuth requires a reachable account API. If the account API check fails at startup, the connector fails validation instead of falling back to a workspace-only sync, even when `BATON_WORKSPACES` is set.
</Warning>

See the connector's README or run `--help` to see all available configuration flags and environment variables.

#### Deployment configuration
Expand Down
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ require (
github.com/quasilyte/go-ruleguard/dsl v0.3.23
go.uber.org/zap v1.28.0
golang.org/x/oauth2 v0.36.0
google.golang.org/grpc v1.83.0
google.golang.org/protobuf v1.36.11
)

Expand Down Expand Up @@ -136,7 +137,6 @@ require (
golang.org/x/text v0.40.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20260729162451-8efbd57d26e0 // indirect
google.golang.org/grpc v1.83.0 // indirect
gopkg.in/yaml.v2 v2.4.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
modernc.org/libc v1.72.0 // indirect
Expand Down
8 changes: 5 additions & 3 deletions pkg/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -107,9 +107,11 @@ var Config = field.NewConfiguration(
{
Name: DatabricksWorkspaceTokenGroup,
DisplayName: "Workspace token",
HelpText: "Authenticate with a personal access token scoped to each workspace.",
Fields: []field.SchemaField{AccountIdField, WorkspacesField, WorkspaceTokensField, HostnameField, AccountHostnameField, ExcludeWorkspacesField},
Default: false,
HelpText: "Authenticate with a personal access token scoped to each workspace. " +
"Does not sync account-level data (account entitlements and grants, and " +
"workspace-membership entitlements); use OAuth for full account coverage.",
Fields: []field.SchemaField{AccountIdField, WorkspacesField, WorkspaceTokensField, HostnameField, AccountHostnameField},
Default: false,
},
}),
)
Expand Down
21 changes: 17 additions & 4 deletions pkg/connector/connector.go
Original file line number Diff line number Diff line change
Expand Up @@ -109,12 +109,14 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err
isAccAPIAvailable := false
isWSAPIAvailable := false

// The Account API is unreachable with workspace tokens, so only probe it for OAuth.
// OAuth must reach the account API; a failed check here is a fixable misconfiguration,
// so fail instead of silently dropping account-level data (token auth, handled below,
// can't reach it by design).
if !d.client.IsTokenAuth() {
_, _, err := d.client.ListRoles(ctx, "", "", "")
if err == nil {
isAccAPIAvailable = true
if _, _, err := d.client.ListRoles(ctx, "", "", ""); err != nil {
return nil, fmt.Errorf("databricks-connector: account API validation failed: %w", err)
}
isAccAPIAvailable = true
}
Comment thread
al-conductorone marked this conversation as resolved.

// With an explicit workspace list (always the case for token auth), validate each
Expand Down Expand Up @@ -148,6 +150,17 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err

d.client.UpdateAvailability(isAccAPIAvailable, isWSAPIAvailable)

// Token auth can't reach the account plane: account entitlements/grants and
// workspace-membership entitlements go unsynced and identities re-parent onto the
// workspace. Warn, not Debug (invisible at info level), so this drop isn't silent.
if !isAccAPIAvailable && isWSAPIAvailable {
ctxzap.Extract(ctx).Warn(
Comment thread
al-conductorone marked this conversation as resolved.
Comment thread
al-conductorone marked this conversation as resolved.
Comment thread
al-conductorone marked this conversation as resolved.
"databricks-connector: account API unreachable under workspace-token auth; syncing workspace-scoped data only. " +
"Account entitlements and grants, and workspace-membership entitlements, will not be synced, " +
"and identities are parented under their workspace instead of the account",
)
}

return nil, nil
}

Expand Down
126 changes: 126 additions & 0 deletions pkg/connector/validate_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
package connector

import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"strings"
"testing"

"github.com/conductorone/baton-databricks/pkg/databricks"
"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
)

// rolesTransport answers the assignable-roles calls Validate makes. failAccount
// makes the account-plane check (host "accounts.*") fail so isAccAPIAvailable stays
// false while the workspace check still succeeds.
type rolesTransport struct{ failAccount bool }

func (t rolesTransport) RoundTrip(req *http.Request) (*http.Response, error) {
if t.failAccount && strings.HasPrefix(req.URL.Host, "accounts.") {
return &http.Response{
StatusCode: http.StatusInternalServerError,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"message":"boom"}`)),
Request: req,
}, nil
}
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"roles":[]}`)),
Request: req,
}, nil
}

func captureLogs(ctx context.Context, buf *bytes.Buffer) context.Context {
core := zapcore.NewCore(
zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
zapcore.AddSync(buf),
zapcore.DebugLevel,
)
return ctxzap.ToContext(ctx, zap.New(core))
}

func newValidateConnector(t *testing.T, auth databricks.Auth, tr http.RoundTripper) *Databricks {
t.Helper()
client, err := databricks.NewClient(
context.Background(), &http.Client{Transport: tr},
"cloud.databricks.com", "accounts.cloud.databricks.com",
"acct-1", "", auth, nil,
)
if err != nil {
t.Fatalf("NewClient: %v", err)
}
return &Databricks{client: client, workspaces: []string{"ws1"}}
}

// levelFor scans the captured JSON log lines for the first entry whose message
// contains want and returns its level. Empty string means no such entry.
func levelFor(t *testing.T, buf *bytes.Buffer, want string) string {
t.Helper()
for _, line := range strings.Split(buf.String(), "\n") {
if line == "" {
continue
}
var entry struct {
Level string `json:"level"`
Msg string `json:"msg"`
}
if err := json.Unmarshal([]byte(line), &entry); err != nil {
continue
}
if strings.Contains(entry.Msg, want) {
return entry.Level
}
}
return ""
}

const accountUnreachableMsg = "account API unreachable"

// CXH-2350: dropping the whole account plane is a customer-visible degradation, so the
// startup notice must be visible. It logs at warn, not debug (a debug line is invisible at
// the default info level, which is the silent degradation the ticket was filed to fix).
func TestValidateWorkspaceTokenLogsAtWarn(t *testing.T) {
buf := &bytes.Buffer{}
ctx := captureLogs(context.Background(), buf)
d := newValidateConnector(t, databricks.NewTokenAuth([]string{"ws1"}, []string{"tok"}), rolesTransport{})

if _, err := d.Validate(ctx); err != nil {
t.Fatalf("Validate: %v", err)
}

if got := levelFor(t, buf, accountUnreachableMsg); got != "warn" {
t.Errorf("token-auth notice logged at %q, want %q", got, "warn")
}
}

// Under OAuth a failed account check is a fixable misconfiguration, so Validate
// fails rather than silently dropping account-level data.
func TestValidateOAuthAccountCheckFailureReturnsError(t *testing.T) {
d := newValidateConnector(t, &databricks.NoAuth{}, rolesTransport{failAccount: true})

if _, err := d.Validate(context.Background()); err == nil {
t.Fatal("Validate: want error on OAuth account check failure, got nil")
}
}

// When the account API is reachable (non-token auth), the notice must not fire at all.
func TestValidateAccountReachableNoNotice(t *testing.T) {
buf := &bytes.Buffer{}
ctx := captureLogs(context.Background(), buf)
d := newValidateConnector(t, &databricks.NoAuth{}, rolesTransport{})

if _, err := d.Validate(ctx); err != nil {
t.Fatalf("Validate: %v", err)
}

if got := levelFor(t, buf, accountUnreachableMsg); got != "" {
t.Errorf("notice fired (level %q) when account API was reachable", got)
}
}
33 changes: 31 additions & 2 deletions pkg/databricks/request.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,41 @@ import (
"github.com/conductorone/baton-sdk/pkg/uhttp"
"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
"go.uber.org/zap"
"golang.org/x/oauth2"
"google.golang.org/grpc/codes"
)

const (
AlreadyExists = "AlreadyExists"
)

// wrapTransportAuthError maps an OAuth2 token-retrieval failure to a gRPC
// status. It happens in the oauth2 transport before any API response, so uhttp
// never sees an HTTP status and the error would otherwise surface as Unknown.
// oauth2 returns RetrieveError for any non-2xx from the token endpoint, so a
// transient 5xx must stay retryable (Unavailable) rather than look like bad
// credentials. A missing Response falls back to Unauthenticated.
func wrapTransportAuthError(err error) error {
var retrieveErr *oauth2.RetrieveError
if !errors.As(err, &retrieveErr) {
return err
}

code := codes.Unauthenticated
if retrieveErr.Response != nil {
switch status := retrieveErr.Response.StatusCode; {
case status == http.StatusForbidden:
code = codes.PermissionDenied
case status == http.StatusTooManyRequests:
code = codes.Unavailable
Comment thread
al-conductorone marked this conversation as resolved.
case status >= 500:
code = codes.Unavailable
}
Comment thread
al-conductorone marked this conversation as resolved.
}

return uhttp.WrapErrors(code, "databricks-connector: authentication failed", err)
}

// APIError represents an error response from the Databricks API.
type APIError struct {
StatusCode int
Expand Down Expand Up @@ -182,7 +211,7 @@ func (c *Client) doRequest(
uhttp.WithRatelimitData(ratelimitData),
)
if resp == nil {
return ratelimitData, err
return ratelimitData, wrapTransportAuthError(err)
}

defer resp.Body.Close()
Expand Down Expand Up @@ -256,7 +285,7 @@ func (c *Client) doRequestNoResponse(
uhttp.WithRatelimitData(ratelimitData),
)
if resp == nil {
return ratelimitData, err
return ratelimitData, wrapTransportAuthError(err)
}

defer resp.Body.Close()
Expand Down
36 changes: 36 additions & 0 deletions pkg/databricks/request_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,47 @@ package databricks

import (
"errors"
"fmt"
"net/http"
"strings"
"testing"

"golang.org/x/oauth2"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)

func TestWrapTransportAuthError(t *testing.T) {
retrieveErr := &oauth2.RetrieveError{ErrorCode: "invalid_client"}
retrieveWithStatus := func(status int) *oauth2.RetrieveError {
return &oauth2.RetrieveError{Response: &http.Response{StatusCode: status}}
}

tests := []struct {
name string
in error
want codes.Code
}{
{"nil stays nil", nil, codes.OK},
{"oauth2 retrieve error without response maps to unauthenticated", retrieveErr, codes.Unauthenticated},
{"wrapped oauth2 retrieve error maps to unauthenticated", fmt.Errorf("get workspaces: %w", retrieveErr), codes.Unauthenticated},
{"401 maps to unauthenticated", retrieveWithStatus(http.StatusUnauthorized), codes.Unauthenticated},
{"403 maps to permission denied", retrieveWithStatus(http.StatusForbidden), codes.PermissionDenied},
{"429 stays retryable as unavailable", retrieveWithStatus(http.StatusTooManyRequests), codes.Unavailable},
{"503 stays retryable as unavailable", retrieveWithStatus(http.StatusServiceUnavailable), codes.Unavailable},
{"unrelated error is left unknown", errors.New("boom"), codes.Unknown},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := status.Code(wrapTransportAuthError(tt.in))
if got != tt.want {
t.Errorf("wrapTransportAuthError(%v) grpc code = %v, want %v", tt.in, got, tt.want)
}
})
}
}

func TestNameWorkspace403Remedy(t *testing.T) {
forbidden := &APIError{StatusCode: http.StatusForbidden, Message: "Unauthorized access to Org: 123"}
badRequest := &APIError{StatusCode: http.StatusBadRequest, Message: "bad"}
Expand Down
Loading