From 82938f55b31e8588ff992ce01775de1dad4a78f0 Mon Sep 17 00:00:00 2001 From: Alejandro Bernal Date: Mon, 27 Jan 2025 10:14:24 -0500 Subject: [PATCH 1/3] fixed azure databricks + docs --- README.md | 9 +++++++++ cmd/baton-databricks/main.go | 2 ++ pkg/config/schema.go | 1 - pkg/databricks/auth.go | 4 ++-- pkg/databricks/client.go | 2 +- 5 files changed, 14 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 41ab1370..76760e25 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,14 @@ to work with only specified workspaces and their respective tokens. You can provide multiple tokens by separating them with a comma. This method requires admin access to each workspace you want to sync. +# Using Azure Databricks + +To work with Azure Databricks, you need to provide the hostname flag. + +```bash +baton-databricks --hostname "azuredatabricks.net" +``` + # Getting Started ## brew @@ -120,6 +128,7 @@ Flags: --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) -f, --file string The path to the c1z file to sync with ($BATON_FILE) (default "sync.c1z") -h, --help help for baton-databricks + --hostname 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") --password string The Databricks password used to connect to the Databricks API ($BATON_PASSWORD) diff --git a/cmd/baton-databricks/main.go b/cmd/baton-databricks/main.go index 7cb88413..87c5ec1a 100644 --- a/cmd/baton-databricks/main.go +++ b/cmd/baton-databricks/main.go @@ -58,6 +58,7 @@ func prepareClientAuth(ctx context.Context, cfg *viper.Viper) databricks.Auth { password := cfg.GetString(config.PasswordField.FieldName) workspaces := cfg.GetStringSlice(config.WorkspacesField.FieldName) tokens := cfg.GetStringSlice(config.TokensField.FieldName) + hostname := cfg.GetString(config.HostnameField.FieldName) switch { case username != "" && password != "": @@ -81,6 +82,7 @@ func prepareClientAuth(ctx context.Context, cfg *viper.Viper) databricks.Auth { accountID, databricksClientId, databricksClientSecret, + hostname, ) return cAuth case AreTokensSet(workspaces, tokens): diff --git a/pkg/config/schema.go b/pkg/config/schema.go index 528ee310..9b70529e 100644 --- a/pkg/config/schema.go +++ b/pkg/config/schema.go @@ -12,7 +12,6 @@ var ( AccountHostnameField = field.StringField( "account-hostname", field.WithDescription("The hostname used to connect to the Databricks account API"), - field.WithDefaultValue("accounts.cloud.databricks.com"), ) AccountIdField = field.StringField( "account-id", diff --git a/pkg/databricks/auth.go b/pkg/databricks/auth.go index 459d3447..8d2dd13c 100644 --- a/pkg/databricks/auth.go +++ b/pkg/databricks/auth.go @@ -93,12 +93,12 @@ type OAuth2 struct { cfg *clientcredentials.Config } -func NewOAuth2(accId, clientId, clientSecret string) *OAuth2 { +func NewOAuth2(accId, clientId, clientSecret, hostname string) *OAuth2 { return &OAuth2{ cfg: &clientcredentials.Config{ ClientID: clientId, ClientSecret: clientSecret, - TokenURL: fmt.Sprintf("https://accounts.cloud.databricks.com/oidc/accounts/%s/v1/token", accId), + TokenURL: fmt.Sprintf("https://accounts.%s/oidc/accounts/%s/v1/token", hostname, accId), Scopes: []string{"all-apis"}, }, } diff --git a/pkg/databricks/client.go b/pkg/databricks/client.go index fbd203bf..b06300b0 100644 --- a/pkg/databricks/client.go +++ b/pkg/databricks/client.go @@ -52,7 +52,7 @@ func NewClient(ctx context.Context, httpClient *http.Client, hostname, accountHo } if accountHostname == "" { - accountHostname = "accounts." + defaultHost + accountHostname = "accounts." + hostname } accountBaseUrl := &url.URL{ Scheme: "https", From 397673d6dce9df08f7443d77790b41afdef4930d Mon Sep 17 00:00:00 2001 From: Alejandro Bernal Date: Tue, 28 Jan 2025 16:45:14 -0500 Subject: [PATCH 2/3] use accountHostname when requesting token --- README.md | 3 ++- cmd/baton-databricks/main.go | 9 ++++----- pkg/databricks/auth.go | 4 ++-- pkg/databricks/client.go | 23 ++++++++++++++++------- 4 files changed, 24 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 76760e25..0cfc8fa5 100644 --- a/README.md +++ b/README.md @@ -121,6 +121,7 @@ Available Commands: help Help about any command Flags: + --account-hostname string The hostname used to connect to the Databricks account API ($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) @@ -128,7 +129,7 @@ Flags: --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) -f, --file string The path to the c1z file to sync with ($BATON_FILE) (default "sync.c1z") -h, --help help for baton-databricks - --hostname The Databricks hostname used to connect to the Databricks API ($BATON_HOSTNAME) (default "cloud.databricks.com") + --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") --password string The Databricks password used to connect to the Databricks API ($BATON_PASSWORD) diff --git a/cmd/baton-databricks/main.go b/cmd/baton-databricks/main.go index 87c5ec1a..1fa18ca8 100644 --- a/cmd/baton-databricks/main.go +++ b/cmd/baton-databricks/main.go @@ -58,7 +58,7 @@ func prepareClientAuth(ctx context.Context, cfg *viper.Viper) databricks.Auth { password := cfg.GetString(config.PasswordField.FieldName) workspaces := cfg.GetStringSlice(config.WorkspacesField.FieldName) tokens := cfg.GetStringSlice(config.TokensField.FieldName) - hostname := cfg.GetString(config.HostnameField.FieldName) + accountHostname := databricks.GetAccountHostname(cfg) switch { case username != "" && password != "": @@ -82,7 +82,7 @@ func prepareClientAuth(ctx context.Context, cfg *viper.Viper) databricks.Auth { accountID, databricksClientId, databricksClientSecret, - hostname, + accountHostname, ) return cAuth case AreTokensSet(workspaces, tokens): @@ -108,9 +108,8 @@ func getConnector(ctx context.Context, cfg *viper.Viper) (types.ConnectorServer, return nil, err } - hostname := cfg.GetString(config.HostnameField.FieldName) - accountHostname := cfg.GetString(config.AccountHostnameField.FieldName) - + hostname := databricks.GetHostname(cfg) + accountHostname := databricks.GetAccountHostname(cfg) auth := prepareClientAuth(ctx, cfg) cb, err := connector.New( ctx, diff --git a/pkg/databricks/auth.go b/pkg/databricks/auth.go index 8d2dd13c..1042d563 100644 --- a/pkg/databricks/auth.go +++ b/pkg/databricks/auth.go @@ -93,12 +93,12 @@ type OAuth2 struct { cfg *clientcredentials.Config } -func NewOAuth2(accId, clientId, clientSecret, hostname string) *OAuth2 { +func NewOAuth2(accId, clientId, clientSecret, accountHostname string) *OAuth2 { return &OAuth2{ cfg: &clientcredentials.Config{ ClientID: clientId, ClientSecret: clientSecret, - TokenURL: fmt.Sprintf("https://accounts.%s/oidc/accounts/%s/v1/token", hostname, accId), + TokenURL: fmt.Sprintf("https://%s/oidc/accounts/%s/v1/token", accountHostname, accId), Scopes: []string{"all-apis"}, }, } diff --git a/pkg/databricks/client.go b/pkg/databricks/client.go index b06300b0..a886534e 100644 --- a/pkg/databricks/client.go +++ b/pkg/databricks/client.go @@ -6,8 +6,10 @@ import ( "net/http" "net/url" + "github.com/conductorone/baton-databricks/pkg/config" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/uhttp" + "github.com/spf13/viper" ) const ( @@ -42,18 +44,25 @@ type Client struct { isWSAPIAvailable bool } -func NewClient(ctx context.Context, httpClient *http.Client, hostname, accountHostname, accountID string, auth Auth) (*Client, error) { - if hostname == "" { - hostname = defaultHost +func GetHostname(cfg *viper.Viper) string { + if cfg.GetString(config.HostnameField.FieldName) == "" { + return defaultHost } + return cfg.GetString(config.HostnameField.FieldName) +} + +func GetAccountHostname(cfg *viper.Viper) string { + if cfg.GetString(config.AccountHostnameField.FieldName) == "" { + return "accounts." + GetHostname(cfg) + } + return cfg.GetString(config.AccountHostnameField.FieldName) +} + +func NewClient(ctx context.Context, httpClient *http.Client, hostname, accountHostname, accountID string, auth Auth) (*Client, error) { baseUrl := &url.URL{ Scheme: "https", Host: hostname, } - - if accountHostname == "" { - accountHostname = "accounts." + hostname - } accountBaseUrl := &url.URL{ Scheme: "https", Host: accountHostname, From 7de479121ec92bcc651e5e763e71995c264dee34 Mon Sep 17 00:00:00 2001 From: Bjorn Tipling Date: Fri, 31 Jan 2025 15:15:05 -0800 Subject: [PATCH 3/3] detect GCP and Azure hostnames and return correct account hostnames --- pkg/config/schema.go | 5 ----- pkg/databricks/client.go | 11 ++++++++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pkg/config/schema.go b/pkg/config/schema.go index 9b70529e..44401390 100644 --- a/pkg/config/schema.go +++ b/pkg/config/schema.go @@ -9,10 +9,6 @@ import ( ) var ( - AccountHostnameField = field.StringField( - "account-hostname", - field.WithDescription("The hostname used to connect to the Databricks account API"), - ) AccountIdField = field.StringField( "account-id", field.WithDescription("The Databricks account ID used to connect to the Databricks Account and Workspace API"), @@ -48,7 +44,6 @@ var ( field.WithDescription("The Databricks access tokens scoped to specific workspaces used to connect to the Databricks Workspace API"), ) configurationFields = []field.SchemaField{ - AccountHostnameField, AccountIdField, DatabricksClientIdField, DatabricksClientSecretField, diff --git a/pkg/databricks/client.go b/pkg/databricks/client.go index a886534e..401dc400 100644 --- a/pkg/databricks/client.go +++ b/pkg/databricks/client.go @@ -5,6 +5,7 @@ import ( "fmt" "net/http" "net/url" + "strings" "github.com/conductorone/baton-databricks/pkg/config" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" @@ -14,6 +15,8 @@ import ( const ( defaultHost = "cloud.databricks.com" + azureHost = "azuredatabricks.net" + gcpHost = "gcp.databricks.net" // Some of these are case sensitive. usersEndpoint = "/api/2.0/preview/scim/v2/Users" @@ -52,10 +55,12 @@ func GetHostname(cfg *viper.Viper) string { } func GetAccountHostname(cfg *viper.Viper) string { - if cfg.GetString(config.AccountHostnameField.FieldName) == "" { - return "accounts." + GetHostname(cfg) + if strings.HasSuffix(GetHostname(cfg), azureHost) { + return "accounts." + azureHost + } else if strings.HasSuffix(GetHostname(cfg), gcpHost) { + return "accounts." + gcpHost } - return cfg.GetString(config.AccountHostnameField.FieldName) + return "accounts." + GetHostname(cfg) } func NewClient(ctx context.Context, httpClient *http.Client, hostname, accountHostname, accountID string, auth Auth) (*Client, error) {