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

@@ -9,11 +9,6 @@ import (
 )
 
 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",
 		field.WithDescription("The Databricks account ID used to connect to the Databricks Account and Workspace API"),
@@ -49,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,

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

@@ -5,13 +5,18 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strings"
 
+	"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 (
 	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"
@@ -42,18 +47,27 @@ 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 strings.HasSuffix(GetHostname(cfg), azureHost) {
+		return "accounts." + azureHost
+	} else if strings.HasSuffix(GetHostname(cfg), gcpHost) {
+		return "accounts." + gcpHost
+	}
+	return "accounts." + GetHostname(cfg)
+}
+
+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,