feat: add workload identity federation support

config_schema.json

@@ -126,13 +126,21 @@
       "name": "databricks-client-secret",
       "displayName": "OAuth2 Client Secret",
       "description": "The Databricks service principal's client secret used to connect to the Databricks Account and Workspace API",
-      "isRequired": true,
       "isSecret": true,
-      "stringField": {
-        "rules": {
-          "isRequired": true
-        }
-      }
+      "stringField": {}
+    },
+    {
+      "name": "databricks-token-file",
+      "displayName": "Federation Token File",
+      "description": "Path to a file containing an external JWT for workload identity federation (e.g. a SPIFFE JWT-SVID or Kubernetes projected ServiceAccount token). The file is re-read on each token refresh to support credential rotation.",
+      "stringField": {}
+    },
+    {
+      "name": "databricks-token",
+      "displayName": "Federation Token",
+      "description": "An external JWT for workload identity federation (RFC 8693 token exchange). Use --databricks-token-file for automatic credential rotation.",
+      "isSecret": true,
+      "stringField": {}
     },
     {
       "name": "hostname",
@@ -143,6 +151,24 @@
       }
     }
   ],
+  "constraints": [
+    {
+      "kind": "CONSTRAINT_KIND_MUTUALLY_EXCLUSIVE",
+      "fieldNames": [
+        "databricks-client-secret",
+        "databricks-token-file",
+        "databricks-token"
+      ]
+    },
+    {
+      "kind": "CONSTRAINT_KIND_AT_LEAST_ONE",
+      "fieldNames": [
+        "databricks-client-secret",
+        "databricks-token-file",
+        "databricks-token"
+      ]
+    }
+  ],
   "displayName": "Databricks",
   "helpUrl": "/docs/baton/databricks",
   "iconUrl": "/static/app-icons/databricks.svg"

pkg/config/config.go

@@ -21,9 +21,19 @@ var (
 		"databricks-client-secret",
 		field.WithDescription("The Databricks service principal's client secret used to connect to the Databricks Account and Workspace API"),
 		field.WithIsSecret(true),
-		field.WithRequired(true),
 		field.WithDisplayName("OAuth2 Client Secret"),
 	)
+	DatabricksTokenFileField = field.StringField(
+		"databricks-token-file",
+		field.WithDescription("Path to a file containing an external JWT for workload identity federation (e.g. a SPIFFE JWT-SVID or Kubernetes projected ServiceAccount token). The file is re-read on each token refresh to support credential rotation."),
+		field.WithDisplayName("Federation Token File"),
+	)
+	DatabricksTokenField = field.StringField(
+		"databricks-token",
+		field.WithDescription("An external JWT for workload identity federation (RFC 8693 token exchange). Use --databricks-token-file for automatic credential rotation."),
+		field.WithIsSecret(true),
+		field.WithDisplayName("Federation Token"),
+	)
 	AccountHostnameField = field.StringField(
 		"account-hostname",
 		field.WithDefaultValue("accounts.cloud.databricks.com"),
@@ -47,6 +57,8 @@ var (
 		AccountIdField,
 		DatabricksClientIdField,
 		DatabricksClientSecretField,
+		DatabricksTokenFileField,
+		DatabricksTokenField,
 		HostnameField,
 		BaseURLField,
 	}
@@ -58,4 +70,8 @@ var Config = field.NewConfiguration(
 	field.WithConnectorDisplayName("Databricks"),
 	field.WithHelpUrl("/docs/baton/databricks"),
 	field.WithIconUrl("/static/app-icons/databricks.svg"),
+	field.WithConstraints(
+		field.FieldsMutuallyExclusive(DatabricksClientSecretField, DatabricksTokenFileField, DatabricksTokenField),
+		field.FieldsAtLeastOneUsed(DatabricksClientSecretField, DatabricksTokenFileField, DatabricksTokenField),
+	),
 )

pkg/connector/connector.go

@@ -189,13 +189,23 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.Connect
 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)
 
+	if cfg.DatabricksTokenFile != "" || cfg.DatabricksToken != "" {
+		l.Info("using workload identity federation (token exchange)")
+		return databricks.NewTokenFederation(
+			accountID,
+			databricksClientId,
+			cfg.DatabricksTokenFile,
+			cfg.DatabricksToken,
+			accountHostname,
+		)
+	}
+
 	return databricks.NewOAuth2(
 		accountID,
 		databricksClientId,
-		databricksClientSecret,
+		cfg.DatabricksClientSecret,
 		accountHostname,
 	)
 }

pkg/databricks/auth.go

@@ -2,8 +2,14 @@
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
+	"net/url"
+	"os"
+	"strings"
+	"time"
 
 	"github.com/conductorone/baton-sdk/pkg/uhttp"
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
@@ -51,6 +57,95 @@
 	return httpClient, nil
 }
 
-func (o *OAuth2) Apply(req *http.Request) {
-	// No need to set the Authorization header here, the oauth2 client does it automatically
+func (o *OAuth2) Apply(req *http.Request) {}
+
+// TokenFederation authenticates via RFC 8693 token exchange.
+// It presents an externally-issued JWT (e.g. SPIFFE JWT-SVID, Kubernetes SA token)
+// and exchanges it for a short-lived Databricks OAuth token.
+type TokenFederation struct {
+	clientID    string
+	tokenURL    string
+	tokenFile   string
+	staticToken string
 }
+
+func NewTokenFederation(accId, clientId, tokenFile, token, accountHostname string) *TokenFederation {
+	return &TokenFederation{
+		clientID:    clientId,
+		tokenURL:    fmt.Sprintf("https://%s/oidc/accounts/%s/v1/token", accountHostname, accId),
+		tokenFile:   tokenFile,
+		staticToken: token,
+	}
+}
+
+func (tf *TokenFederation) readSubjectToken() (string, error) {
+	if tf.tokenFile != "" {
+		data, err := os.ReadFile(tf.tokenFile)
+		if err != nil {
+			return "", fmt.Errorf("failed to read token file %s: %w", tf.tokenFile, err)
+		}
+		return strings.TrimSpace(string(data)), nil
+	}
+	return tf.staticToken, nil
+}
+
+// Token performs an RFC 8693 token exchange and returns the resulting OAuth2 token.
+func (tf *TokenFederation) Token() (*oauth2.Token, error) {
+	subjectToken, err := tf.readSubjectToken()
+	if err != nil {
+		return nil, err
+	}
+	if subjectToken == "" {
+		return nil, fmt.Errorf("empty subject token for workload identity federation")
+	}
+
+	form := url.Values{
+		"grant_type":         {"urn:ietf:params:oauth:grant-type:token-exchange"},
+		"subject_token":      {subjectToken},
+		"subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"},
+		"client_id":          {tf.clientID},
+		"scope":              {"all-apis"},
+	}
+
+	resp, err := http.PostForm(tf.tokenURL, form)
+	if err != nil {
+		return nil, fmt.Errorf("token exchange request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read token exchange response: %w", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(body))
+	}
+
+	var tokenResp struct {
+		AccessToken string `json:"access_token"`
+		TokenType   string `json:"token_type"`
+		ExpiresIn   int64  `json:"expires_in"`
+	}
+	if err := json.Unmarshal(body, &tokenResp); err != nil {
+		return nil, fmt.Errorf("failed to decode token exchange response: %w", err)
+	}
+
+	token := &oauth2.Token{
+		AccessToken: tokenResp.AccessToken,
+		TokenType:   tokenResp.TokenType,
+	}
+	if tokenResp.ExpiresIn > 0 {
+		token.Expiry = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
+	}
+
+	return token, nil
+}
+
+func (tf *TokenFederation) GetClient(ctx context.Context) (*http.Client, error) {
+	ts := oauth2.ReuseTokenSource(nil, tf)
+	httpClient := oauth2.NewClient(ctx, ts)
+	return httpClient, nil
+}
+
+func (tf *TokenFederation) Apply(req *http.Request) {}