Add support for Azure & GCP Databricks

.github/workflows/capabilities.yaml

@@ -16,7 +16,7 @@ jobs:
           token: ${{ secrets.RELENG_GITHUB_TOKEN }}
 
       - name: Setup Go
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version-file: 'go.mod'
 

.github/workflows/ci.yaml

@@ -5,11 +5,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Install Go
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version: 1.21.x
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Run linters
         uses: golangci/golangci-lint-action@v6
         with:
@@ -24,11 +24,11 @@ jobs:
     steps:
       - name: Install Go
         if: success()
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version: ${{ matrix.go-version }}
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: go tests
         run: (set -o pipefail && go test -v -covermode=count -json ./... | tee test.json)
       - name: annotate go tests
@@ -54,7 +54,7 @@ jobs:
       CONNECTOR_PRINCIPAL: '5346803201281760'
     steps:
       - name: Install Go
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version: 1.22.x
       - name: Checkout code
@@ -64,4 +64,4 @@ jobs:
       - name: Build baton-databricks
         run: go build ./cmd/baton-databricks
       - name: Run baton-databricks
-        run: ./baton-databricks --account-id ${{ env.BATON_ACCOUNT_ID }} --workspaces ${{ env.BATON_WORKSPACES }} --workspace-tokens ${{ env.BATON_WORKSPACE_TOKENS }}
+        run: ./baton-databricks

.github/workflows/main.yaml

@@ -8,11 +8,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Install Go
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version: 1.21.x
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Run linters
         uses: golangci/golangci-lint-action@v6
         with:
@@ -27,11 +27,11 @@ jobs:
     steps:
       - name: Install Go
         if: success()
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version: ${{ matrix.go-version }}
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: go tests
         run: (set -o pipefail && go test -v -covermode=count -json ./... | tee test.json)
       - name: annotate go tests

.github/workflows/release.yaml

@@ -10,11 +10,11 @@ jobs:
     runs-on: macos-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Set up Go
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version: 1.21.x
       - name: Set up Gon
@@ -37,11 +37,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Set up Go
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version: 1.21.x
       - name: Docker Login

.gon-arm64.json

@@ -1,5 +1,5 @@
 {
-  "source": ["./dist/macos-arm64_darwin_arm64/baton-databricks"],
+  "source": ["./dist/macos-arm64_darwin_arm64_v8.0/baton-databricks"],
   "bundle_id": "com.conductorone.baton-databricks",
   "apple_id": {
     "username" : "[email protected]"

.goreleaser.yaml

@@ -36,7 +36,7 @@ builds:
     hooks:
       post:
         - gon .gon-arm64.json
-        - mv dist/baton-databricks-darwin-arm64.signed.zip dist/macos-arm64_darwin_arm64/baton-databricks
+        - mv dist/baton-databricks-darwin-arm64.signed.zip dist/macos-arm64_darwin_arm64_v8.0/baton-databricks
 archives:
   - id: linux-archive
     builds:

cmd/baton-databricks/main.go

@@ -106,9 +106,14 @@ 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)
+
 	auth := prepareClientAuth(ctx, cfg)
 	cb, err := connector.New(
 		ctx,
+		hostname,
+		accountHostname,
 		cfg.GetString(config.AccountIdField.FieldName),
 		auth,
 		cfg.GetStringSlice(config.WorkspacesField.FieldName),

go.mod

@@ -3,7 +3,7 @@ module github.com/conductorone/baton-databricks
 go 1.21
 
 require (
-	github.com/conductorone/baton-sdk v0.2.35
+	github.com/conductorone/baton-sdk v0.2.58
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
 	github.com/spf13/viper v1.19.0
 	go.uber.org/zap v1.27.0
@@ -14,7 +14,6 @@ require (
 require (
 	filippo.io/age v1.1.1 // indirect
 	filippo.io/edwards25519 v1.1.0 // indirect
-	github.com/allegro/bigcache/v3 v3.1.0 // indirect
 	github.com/deckarep/golang-set/v2 v2.6.0 // indirect
 	github.com/ncruces/go-strftime v0.1.9 // indirect
 	go.opentelemetry.io/otel v1.27.0 // indirect
@@ -43,10 +42,12 @@ require (
 	github.com/aws/smithy-go v1.20.2 // indirect
 	github.com/benbjohnson/clock v1.3.5 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/dolthub/maphash v0.1.0 // indirect
 	github.com/doug-martin/goqu/v9 v9.19.0 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/envoyproxy/protoc-gen-validate v1.0.4 // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
+	github.com/gammazero/deque v0.2.1 // indirect
 	github.com/glebarez/go-sqlite v1.22.0 // indirect
 	github.com/go-jose/go-jose/v3 v3.0.3 // indirect
 	github.com/go-ole/go-ole v1.3.0 // indirect
@@ -59,6 +60,7 @@ require (
 	github.com/lufia/plan9stats v0.0.0-20240408141607-282e7b5d6b74 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/maypok86/otter v1.2.4 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect

go.sum

@@ -6,8 +6,6 @@ filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/DATA-DOG/go-sqlmock v1.5.0 h1:Shsta01QNfFxHCfpW6YH2STWB0MudeXXEWMr20OEh60=
 github.com/DATA-DOG/go-sqlmock v1.5.0/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM=
-github.com/allegro/bigcache/v3 v3.1.0 h1:H2Vp8VOvxcrB91o86fUSVJFqeuz8kpyyB02eH3bSzwk=
-github.com/allegro/bigcache/v3 v3.1.0/go.mod h1:aPyh7jEvrog9zAwx5N7+JUQX5dZTSGpxF1LAR4dr35I=
 github.com/aws/aws-sdk-go-v2 v1.26.1 h1:5554eUqIYVWpU0YmeeYZ0wU64H2VLBs8TlhRB2L+EkA=
 github.com/aws/aws-sdk-go-v2 v1.26.1/go.mod h1:ffIFB97e2yNsv4aTSGkqtHnppsIJzw7G7BReUZ3jCXM=
 github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.2 h1:x6xsQXGSmW6frevwDA+vi/wqhp1ct18mVXYN08/93to=
@@ -52,8 +50,8 @@ github.com/benbjohnson/clock v1.3.5/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZx
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
 github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc=
-github.com/conductorone/baton-sdk v0.2.35 h1:aSdNvlM5HMti8WdhotrXTHWs+b+BmSqMxtGwsSUFxjY=
-github.com/conductorone/baton-sdk v0.2.35/go.mod h1:hmd/Oz3DPIKD+9QmkusZaA18ZoiinnTDdrxh2skcdUc=
+github.com/conductorone/baton-sdk v0.2.58 h1:LBZ12JeEQar2GKA9Tll5iSt7cW6gCkrPbQq3Zs4UeDk=
+github.com/conductorone/baton-sdk v0.2.58/go.mod h1:s8Cj7lufoHVNz4cFM6etem0yjJgorl42jMhHrfM2sFA=
 github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -62,6 +60,8 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8Yc
 github.com/deckarep/golang-set/v2 v2.6.0 h1:XfcQbWM1LlMB8BsJ8N9vW5ehnnPVIw0je80NsVHagjM=
 github.com/deckarep/golang-set/v2 v2.6.0/go.mod h1:VAky9rY/yGXJOLEDv3OMci+7wtDpOF4IN+y82NBOac4=
 github.com/denisenkom/go-mssqldb v0.10.0/go.mod h1:xbL0rPBG9cCiLr28tMa8zpbdarY27NDyej4t/EjAShU=
+github.com/dolthub/maphash v0.1.0 h1:bsQ7JsF4FkkWyrP3oCnFJgrCUAFbFf3kOl4L/QxPDyQ=
+github.com/dolthub/maphash v0.1.0/go.mod h1:gkg4Ch4CdCDu5h6PMriVLawB7koZ+5ijb9puGMV50a4=
 github.com/doug-martin/goqu/v9 v9.19.0 h1:PD7t1X3tRcUiSdc5TEyOFKujZA5gs3VSA7wxSvBx7qo=
 github.com/doug-martin/goqu/v9 v9.19.0/go.mod h1:nf0Wc2/hV3gYK9LiyqIrzBEVGlI8qW3GuDCEobC4wBQ=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
@@ -76,6 +76,8 @@ github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHk
 github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
 github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=
 github.com/fsnotify/fsnotify v1.7.0/go.mod h1:40Bi/Hjc2AVfZrqy+aj+yEI+/bRxZnMJyTJwOpGvigM=
+github.com/gammazero/deque v0.2.1 h1:qSdsbG6pgp6nL7A0+K/B7s12mcCY/5l5SIUpMOl+dC0=
+github.com/gammazero/deque v0.2.1/go.mod h1:LFroj8x4cMYCukHJDbxFCkT+r9AndaJnFMuZDV34tuU=
 github.com/glebarez/go-sqlite v1.22.0 h1:uAcMJhaA6r3LHMTFgP0SifzgXg46yJkgxqyuyec+ruQ=
 github.com/glebarez/go-sqlite v1.22.0/go.mod h1:PlBIdHe0+aUEFn+r2/uthrWq4FxbzugL0L8Li6yQJbc=
 github.com/go-jose/go-jose/v3 v3.0.3 h1:fFKWeig/irsp7XD2zBxvnmA/XaRWp5V3CBsZXJF7G7k=
@@ -141,8 +143,11 @@ github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0V
 github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
-github.com/mattn/go-sqlite3 v1.14.7 h1:fxWBnXkxfM6sRiuH3bqJ4CfzZojMOLVc0UTsTglEghA=
 github.com/mattn/go-sqlite3 v1.14.7/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
+github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU=
+github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
+github.com/maypok86/otter v1.2.4 h1:HhW1Pq6VdJkmWwcZZq19BlEQkHtI8xgsQzBVXJU0nfc=
+github.com/maypok86/otter v1.2.4/go.mod h1:mKLfoI7v1HOmQMwFgX4QkRk23mX6ge3RDvjdHOWG4R4=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
 github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
 github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4=

pkg/config/schema.go

@@ -9,11 +9,21 @@ 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"),
 		field.WithRequired(true),
 	)
+	HostnameField = field.StringField(
+		"hostname",
+		field.WithDescription("The Databricks hostname used to connect to the Databricks API"),
+		field.WithDefaultValue("cloud.databricks.com"),
+	)
 	DatabricksClientIdField = field.StringField(
 		"databricks-client-id",
 		field.WithDescription("The Databricks service principal's client ID used to connect to the Databricks Account and Workspace API"),
@@ -39,9 +49,11 @@ 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,
+		HostnameField,
 		PasswordField,
 		TokensField,
 		UsernameField,

pkg/connector/account.go

@@ -43,12 +43,13 @@ func (a *accountBuilder) ResourceType(ctx context.Context) *v2.ResourceType {
 	return accountResourceType
 }
 
-func accountResource(ctx context.Context, accID string, accAPIAvailable bool) (*v2.Resource, error) {
+func (a *accountBuilder) accountResource(_ context.Context) (*v2.Resource, error) {
+	accountId := a.client.GetAccountId()
 	children := []protoreflect.ProtoMessage{
 		&v2.ChildResourceType{ResourceTypeId: workspaceResourceType.Id},
 	}
 
-	if accAPIAvailable {
+	if a.client.IsAccountAPIAvailable() {
 		children = append(children,
 			&v2.ChildResourceType{ResourceTypeId: userResourceType.Id},
 			&v2.ChildResourceType{ResourceTypeId: groupResourceType.Id},
@@ -58,9 +59,9 @@ func accountResource(ctx context.Context, accID string, accAPIAvailable bool) (*
 	}
 
 	resource, err := rs.NewResource(
-		accID,
+		accountId,
 		accountResourceType,
-		accID,
+		accountId,
 		rs.WithAnnotation(children...),
 	)
 
@@ -72,16 +73,12 @@ func accountResource(ctx context.Context, accID string, accAPIAvailable bool) (*
 }
 
 func (a *accountBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId, pToken *pagination.Token) ([]*v2.Resource, string, annotations.Annotations, error) {
-	var rv []*v2.Resource
-
-	ur, err := accountResource(ctx, a.client.GetAccountId(), a.client.IsAccountAPIAvailable())
+	ur, err := a.accountResource(ctx)
 	if err != nil {
 		return nil, "", nil, err
 	}
 
-	rv = append(rv, ur)
-
-	return rv, "", nil, nil
+	return []*v2.Resource{ur}, "", nil, nil
 }
 
 // Entitlements returns slice of entitlements for marketplace admins under account.
@@ -116,12 +113,10 @@ func (a *accountBuilder) Grants(ctx context.Context, resource *v2.Resource, pTok
 		return nil, "", nil, nil
 	}
 
-	a.client.SetAccountConfig()
-
 	var rv []*v2.Grant
 
 	// list rule sets for the account
-	ruleSets, _, err := a.client.ListRuleSets(ctx, "", "")
+	ruleSets, _, err := a.client.ListRuleSets(ctx, "", "", "")
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to list rule sets for account %s: %w", resource.Id.Resource, err)
 	}
@@ -130,7 +125,7 @@ func (a *accountBuilder) Grants(ctx context.Context, resource *v2.Resource, pTok
 		// rule set contains role and its principals, each one with resource type and resource id seperated by "/"
 		if strings.Contains(ruleSet.Role, MarketplaceAdminRole) {
 			for _, p := range ruleSet.Principals {
-				resourceId, err := prepareResourceID(ctx, a.client, p)
+				resourceId, err := prepareResourceId(ctx, a.client, "", p)
 				if err != nil {
 					return nil, "", nil, fmt.Errorf("databricks-connector: failed to prepare resource id for principal %s: %w", p, err)
 				}
@@ -167,12 +162,12 @@ func (a *accountBuilder) Grant(ctx context.Context, principal *v2.Resource, enti
 	}
 
 	accID := entitlement.Resource.Id.Resource
-	ruleSets, _, err := a.client.ListRuleSets(ctx, "", "")
+	ruleSets, _, err := a.client.ListRuleSets(ctx, "", "", "")
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to list rule sets for account %s: %w", accID, err)
 	}
 
-	principalID, err := preparePrincipalID(ctx, a.client, principal.Id.ResourceType, principal.Id.Resource)
+	principalID, err := preparePrincipalId(ctx, a.client, "", principal.Id.ResourceType, principal.Id.Resource)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to prepare principal id for principal %s: %w", principal.Id.Resource, err)
 	}
@@ -208,7 +203,7 @@ func (a *accountBuilder) Grant(ctx context.Context, principal *v2.Resource, enti
 		})
 	}
 
-	_, err = a.client.UpdateRuleSets(ctx, "", "", ruleSets)
+	_, err = a.client.UpdateRuleSets(ctx, "", "", "", ruleSets)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to update rule sets for account %s: %w", accID, err)
 	}
@@ -233,7 +228,7 @@ func (a *accountBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotatio
 	}
 
 	accID := entitlement.Resource.Id.Resource
-	ruleSets, _, err := a.client.ListRuleSets(ctx, "", "")
+	ruleSets, _, err := a.client.ListRuleSets(ctx, "", "", "")
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to list rule sets for account %s: %w", accID, err)
 	}
@@ -248,7 +243,7 @@ func (a *accountBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotatio
 		return nil, nil
 	}
 
-	principalID, err := preparePrincipalID(ctx, a.client, principal.Id.ResourceType, principal.Id.Resource)
+	principalID, err := preparePrincipalId(ctx, a.client, "", principal.Id.ResourceType, principal.Id.Resource)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to prepare principal id: %w", err)
 	}
@@ -278,7 +273,7 @@ func (a *accountBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotatio
 		}
 	}
 
-	_, err = a.client.UpdateRuleSets(ctx, "", "", ruleSets)
+	_, err = a.client.UpdateRuleSets(ctx, "", "", "", ruleSets)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to update rule sets for account %s: %w", accID, err)
 	}