work

.codegen/_openapi_sha

@@ -1 +1 @@
-a6a317df8327c9b1e5cb59a03a42ffa2aabeef6d
+1668b0db17e23605f8c9d29fb3b674c01590732d

.gitattributes

@@ -77,9 +77,11 @@ experimental/mocks/service/marketplace/mock_provider_provider_analytics_dashboar
 experimental/mocks/service/marketplace/mock_provider_providers_interface.go linguist-generated=true
 experimental/mocks/service/ml/mock_experiments_interface.go linguist-generated=true
 experimental/mocks/service/ml/mock_model_registry_interface.go linguist-generated=true
+experimental/mocks/service/oauth2/mock_account_federation_policy_interface.go linguist-generated=true
 experimental/mocks/service/oauth2/mock_custom_app_integration_interface.go linguist-generated=true
 experimental/mocks/service/oauth2/mock_o_auth_published_apps_interface.go linguist-generated=true
 experimental/mocks/service/oauth2/mock_published_app_integration_interface.go linguist-generated=true
+experimental/mocks/service/oauth2/mock_service_principal_federation_policy_interface.go linguist-generated=true
 experimental/mocks/service/oauth2/mock_service_principal_secrets_interface.go linguist-generated=true
 experimental/mocks/service/pipelines/mock_pipelines_interface.go linguist-generated=true
 experimental/mocks/service/provisioning/mock_credentials_interface.go linguist-generated=true

.github/workflows/codeql.yml

@@ -0,0 +1,42 @@
+name: "CodeQL Advanced"
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+  schedule:
+    - cron: '41 1 * * 0'
+
+jobs:
+  analyze:
+    name: Analyze (${{ matrix.language }})
+    runs-on:
+      group: databricks-deco-testing-runner-group
+      labels: ubuntu-latest-deco
+    permissions:
+      # required for all workflows
+      security-events: write
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+        - language: go
+          build-mode: autobuild
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    # Initializes the CodeQL tools for scanning.
+    - name: Initialize CodeQL
+      uses: github/codeql-action/init@v3
+      with:
+        languages: ${{ matrix.language }}
+        build-mode: ${{ matrix.build-mode }}
+
+    - name: Perform CodeQL Analysis
+      uses: github/codeql-action/analyze@v3
+      with:
+        category: "/language:${{matrix.language}}"

.github/workflows/external-message.yml

@@ -13,7 +13,10 @@ on:
 
 jobs:
   comment-on-pr:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: databricks-deco-testing-runner-group
+      labels: ubuntu-latest-deco
+
     permissions:
       pull-requests: write
 
@@ -44,13 +47,13 @@ jobs:
           gh pr comment ${{ github.event.pull_request.number }} --body \
           "<!-- INTEGRATION_TESTS_MANUAL -->
           If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:
-          
+
           Trigger:
           [go/deco-tests-run/sdk-go](https://go/deco-tests-run/sdk-go)
 
           Inputs:
           * PR number: ${{github.event.pull_request.number}}
           * Commit SHA: \`${{ env.COMMIT_SHA }}\`
-          
+
           Checks will be approved automatically on success.
           "

.github/workflows/integration-tests.yml

@@ -6,12 +6,16 @@ on:
     types: [opened, synchronize]
 
   merge_group:
-  
+
 
 jobs:
   check-token:
     name: Check secrets access
-    runs-on: ubuntu-latest
+
+    runs-on:
+      group: databricks-deco-testing-runner-group
+      labels: ubuntu-latest-deco
+
     environment: "test-trigger-is"
     outputs:
       has_token: ${{ steps.set-token-status.outputs.has_token }}
@@ -29,11 +33,15 @@ jobs:
 
   trigger-tests:
     name: Trigger Tests
-    runs-on: ubuntu-latest
+
+    runs-on:
+      group: databricks-deco-testing-runner-group
+      labels: ubuntu-latest-deco
+
     needs: check-token
     if: github.event_name == 'pull_request'  && needs.check-token.outputs.has_token == 'true'
     environment: "test-trigger-is"
-    
+
     steps:
     - uses: actions/checkout@v3
 
@@ -45,22 +53,26 @@ jobs:
         private-key: ${{ secrets.DECO_WORKFLOW_TRIGGER_PRIVATE_KEY }}
         owner: ${{ secrets.ORG_NAME }}
         repositories: ${{secrets.REPO_NAME}}
-    
+
     - name: Trigger Workflow in Another Repo
       env:
         GH_TOKEN: ${{ steps.generate-token.outputs.token }}
       run: |
         gh workflow run sdk-go-isolated-pr.yml -R ${{ secrets.ORG_NAME }}/${{secrets.REPO_NAME}} \
         --ref main \
         -f pull_request_number=${{ github.event.pull_request.number }} \
-        -f commit_sha=${{ github.event.pull_request.head.sha }} 
+        -f commit_sha=${{ github.event.pull_request.head.sha }}
+
 
-    
   # The hash for the merge queue may not be the same as the hash for the PR.
   # Auto approve the check for the merge queue to avoid running integration tests twice.
   auto-approve:
     if: github.event_name == 'merge_group'
-    runs-on: ubuntu-latest
+
+    runs-on:
+      group: databricks-deco-testing-runner-group
+      labels: ubuntu-latest-deco
+
     steps:
       - name: Mark Check
         env:
@@ -71,4 +83,4 @@ jobs:
               -H "X-GitHub-Api-Version: 2022-11-28" \
               /repos/${{ github.repository }}/statuses/${{ github.sha }} \
               -f 'state=success' \
-              -f 'context=Integration Tests Check'
+              -f 'context=Integration Tests Check'

config/api_client.go

@@ -9,7 +9,7 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sdk-go/apierr"
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/useragent"
 )

config/auth_azure_cli.go

@@ -11,7 +11,7 @@ import (
 
 	"golang.org/x/oauth2"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/logger"
 )
 

config/auth_azure_client_secret.go

@@ -8,7 +8,7 @@ import (
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/clientcredentials"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/logger"
 )
 

config/auth_azure_github_oidc.go

@@ -6,7 +6,7 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/logger"
 	"golang.org/x/oauth2"

config/auth_azure_msi.go

@@ -8,7 +8,7 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/logger"
 	"golang.org/x/oauth2"

config/auth_basic.go

@@ -6,7 +6,7 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 )
 
 type BasicCredentials struct {

config/auth_default.go

@@ -5,7 +5,7 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/logger"
 )
 

config/auth_gcp_google_credentials.go

@@ -6,7 +6,7 @@ import (
 	"io/ioutil"
 	"os"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/logger"
 	"golang.org/x/oauth2/google"
 	"google.golang.org/api/idtoken"

config/auth_gcp_google_id.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/logger"
 	"golang.org/x/oauth2"
 	"google.golang.org/api/impersonate"

config/auth_m2m.go

@@ -7,7 +7,7 @@ import (
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/clientcredentials"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/logger"
 )
 

config/auth_metadata_service.go

@@ -8,7 +8,7 @@ import (
 	"net/url"
 	"time"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/logger"
 	"golang.org/x/oauth2"

config/auth_pat.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 )
 
 type PatCredentials struct {

config/auth_u2m.go

@@ -7,7 +7,7 @@ import (
 	"net/http"
 	"strings"
 
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/credentials/cache"
 	"github.com/databricks/databricks-sdk-go/credentials/oauth"
 	"github.com/databricks/databricks-sdk-go/logger"

config/config.go

@@ -13,7 +13,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/common"
 	"github.com/databricks/databricks-sdk-go/common/environment"
-	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/credentials/oauth"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/logger"

config/credentials/credentials.go

@@ -0,0 +1,68 @@
+package credentials
+
+import (
+	"net/http"
+
+	"golang.org/x/oauth2"
+)
+
+// CredentialsProvider is an interface for providing credentials to the client.
+// Implementations of this interface should set the necessary headers on the request.
+type CredentialsProvider interface {
+	// SetHeaders sets the necessary headers on the request.
+	SetHeaders(r *http.Request) error
+}
+
+type credentialsProvider struct {
+	setHeaders func(r *http.Request) error
+}
+
+func (c *credentialsProvider) SetHeaders(r *http.Request) error {
+	return c.setHeaders(r)
+}
+
+func NewCredentialsProvider(visitor func(r *http.Request) error) CredentialsProvider {
+	return &credentialsProvider{
+		setHeaders: visitor,
+	}
+}
+
+// OAuthCredentialsProvider is a specialized CredentialsProvider uses and provides an OAuth token.
+type OAuthCredentialsProvider interface {
+	CredentialsProvider
+	// Token returns the OAuth token generated by the provider.
+	Token() (*oauth2.Token, error)
+}
+
+type oauthCredentialsProvider struct {
+	setHeaders func(r *http.Request) error
+	token      func() (*oauth2.Token, error)
+}
+
+func (c *oauthCredentialsProvider) SetHeaders(r *http.Request) error {
+	return c.setHeaders(r)
+}
+
+func (c *oauthCredentialsProvider) Token() (*oauth2.Token, error) {
+	return c.token()
+}
+
+func NewOAuthCredentialsProvider(visitor func(r *http.Request) error, tokenProvider func() (*oauth2.Token, error)) OAuthCredentialsProvider {
+	return &oauthCredentialsProvider{
+		setHeaders: visitor,
+		token:      tokenProvider,
+	}
+}
+
+// OAuthToken represents an OAuth token as defined by the OAuth 2.0 Authorization Framework.
+// https://datatracker.ietf.org/doc/html/rfc6749
+type OAuthToken struct {
+	// The access token issued by the authorization server. This is the token that will be used to authenticate requests.
+	AccessToken string `json:"access_token"  auth:",sensitive"`
+	// Time in seconds until the token expires.
+	ExpiresIn int `json:"expires_in"`
+	// The scope of the token. This is a space-separated list of strings that represent the permissions granted by the token.
+	Scope string `json:"scope"`
+	// The type of token that was issued.
+	TokenType string `json:"token_type"`
+}