fix azure databricks + docs

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
@@ -113,13 +121,15 @@ 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)
       --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)
   -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")
       --password string                   The Databricks password used to connect to the Databricks API ($BATON_PASSWORD)

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)
+	accountHostname := databricks.GetAccountHostname(cfg)
 
 	switch {
 	case username != "" && password != "":
@@ -81,6 +82,7 @@ func prepareClientAuth(ctx context.Context, cfg *viper.Viper) databricks.Auth {
 			accountID,
 			databricksClientId,
 			databricksClientSecret,
+			accountHostname,
 		)
 		return cAuth
 	case AreTokensSet(workspaces, tokens):
@@ -106,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,

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",

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, accountHostname 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://%s/oidc/accounts/%s/v1/token", accountHostname, accId),
 			Scopes:       []string{"all-apis"},
 		},
 	}

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." + defaultHost
-	}
 	accountBaseUrl := &url.URL{
 		Scheme: "https",
 		Host:   accountHostname,