├── .package.json ├── staticcheck.conf ├── .codegen └── _openapi_sha ├── httpclient ├── fixtures │ ├── testdata │ │ └── some.json │ ├── map_transport.go │ └── slice_transport.go ├── timeout_context_test.go ├── errors_test.go ├── traceparent │ └── traceparent_test.go └── errors.go ├── .release_metadata.json ├── .vscode ├── extensions.json ├── settings.json └── testing.code-snippets ├── version └── version.go ├── config ├── testdata │ ├── corrupt │ │ └── .databrickscfg │ ├── empty_default │ │ └── .databrickscfg │ ├── azure │ │ └── .databrickscfg │ └── .databrickscfg ├── command.go ├── visitors.go ├── config_file_test.go ├── environments_test.go ├── in_memory_test.go ├── experimental │ └── auth │ │ ├── credentials │ │ ├── pat.go │ │ └── pat_test.go │ │ ├── authconv │ │ ├── authconv_test.go │ │ └── authconv.go │ │ └── auth.go ├── auth_basic.go ├── oauth_visitors_test.go ├── auth_pat.go ├── reflect.go ├── auth_m2m.go ├── visitors_test.go ├── auth_default_test.go ├── token_source_strategy.go ├── azure_test.go ├── environments.go ├── auth_azure_cli_federated_token_test.go └── token_source_strategy_test.go ├── qa ├── lock │ ├── core │ │ ├── lockable.go │ │ └── backend.go │ ├── lock_options.go │ ├── lockables.go │ └── internal │ │ ├── lockable.go │ │ └── lockable_test.go ├── debug.go └── poll │ ├── poll.go │ └── poll_test.go ├── NEXT_CHANGELOG.md ├── examples ├── http-proxy │ └── proxy │ │ ├── certutil.go │ │ ├── certutil_windows.go │ │ ├── certutil_linux.go │ │ └── certutil_darwin.go ├── default-auth │ └── main.go ├── useragent │ └── main.go ├── aad-explicit-spn │ └── main.go ├── zerolog │ ├── README.md │ ├── main.go │ ├── go.mod │ └── zerolog.go ├── slog │ ├── main.go │ ├── README.md │ ├── slog.go │ └── go.mod ├── custom-auth │ └── main.go └── long-running │ └── main.go ├── .github ├── dependabot.yml ├── ISSUE_TEMPLATE │ ├── feature_request.md │ └── sdk-issue.md ├── workflows │ ├── codeql.yml │ └── tagging.yml └── PULL_REQUEST_TEMPLATE.md ├── credentials └── u2m │ ├── error.go │ ├── oauth_argument.go │ ├── workspace_oauth_argument_test.go │ ├── cache │ ├── cache.go │ └── file_test.go │ ├── unified_oauth_argument.go │ ├── account_oauth_argument.go │ └── doc.go ├── common ├── environment │ ├── environments_test.go │ └── azure.go ├── response_test.go ├── types │ ├── fieldmask │ │ └── fieldmask.go │ ├── time │ │ └── time.go │ └── duration │ │ └── duration.go └── request.go ├── experimental ├── api │ ├── limiter.go │ └── options.go └── mocks │ └── mocks_test.go ├── service ├── compute │ ├── ext_leading_whitespace_test.go │ ├── ext_sort.go │ ├── libraries_usage_test.go │ ├── ext_leading_whitespace.go │ ├── policy_families_usage_test.go │ ├── ext_results_test.go │ ├── example_test.go │ └── instance_profiles_usage_test.go ├── iam │ ├── ext_model.go │ └── current_user_usage_test.go ├── common │ └── lro │ │ └── options.go ├── provisioning │ ├── ext_azure.go │ └── ext_azure_test.go ├── sql │ ├── dashboards_usage_test.go │ ├── query_history_usage_test.go │ └── data_sources_usage_test.go ├── workspace │ └── examples_test.go ├── settings │ └── workspace_conf_usage_test.go ├── billing │ └── billable_usage_usage_test.go ├── catalog │ └── metastore_assignments_usage_test.go ├── qualitymonitorv2 │ └── interface.go ├── agentbricks │ ├── interface.go │ └── api.go ├── tags │ └── interface.go └── files │ └── doc.go ├── SECURITY.md ├── openapi ├── code │ ├── preview.go │ └── named_sort.go ├── model_test.go ├── roll │ └── compare.go ├── extensions.go └── gen │ └── main.go ├── .mockery.yaml ├── .codegen.json ├── .goreleaser.yml ├── logger ├── context_test.go ├── context.go ├── simple.go ├── logger_test.go └── logger.go ├── internal ├── workspace_client_test.go ├── metastore_test.go ├── sql_test.go ├── account_client_test.go ├── workspaceconf_test.go ├── testspecs │ └── service │ │ ├── jsonmarshallv2 │ │ ├── api.go │ │ └── impl.go │ │ ├── idempotencytesting │ │ ├── api.go │ │ ├── model.go │ │ └── impl.go │ │ ├── httpcallv2 │ │ └── api.go │ │ └── basicv2 │ │ └── api.go ├── libraries_test.go ├── env │ └── env.go ├── tokens_test.go ├── globalinitscripts_test.go ├── tokenmanagement_test.go ├── ipaccesslists_test.go ├── instancepools_test.go ├── warehouses_test.go ├── repos_test.go ├── connections_test.go ├── clusterpolicies_test.go └── account_test.go ├── .gitignore ├── useragent ├── runtime_test.go ├── runtime.go ├── patterns_test.go └── patterns.go ├── workspace_functions.go ├── marshal ├── types.go ├── types_pint_test.go ├── types_pstring_test.go ├── types_pbool_test.go ├── types_pfloat_test.go ├── types_childnofs_test.go ├── types_basemap_test.go ├── types_int_test.go ├── types_maptointerface_test.go ├── types_bool_test.go ├── types_float_test.go ├── types_string_test.go ├── types_interface_test.go ├── types_maptoarray_test.go ├── types_slice_test.go ├── types_pchildnofs_test.go ├── types_childfs_test.go ├── types_maptocustom_test.go ├── reflect_cache.go └── types_pchildfs_test.go ├── alias.go ├── CONTRIBUTING.md ├── apierr └── unwrap.go ├── DCO ├── Makefile ├── test-config.yaml ├── account_functions.go └── go.mod /.package.json: -------------------------------------------------------------------------------- 1 | {} -------------------------------------------------------------------------------- /staticcheck.conf: -------------------------------------------------------------------------------- 1 | checks = ["inherit", "-SA1019"] -------------------------------------------------------------------------------- /.codegen/_openapi_sha: -------------------------------------------------------------------------------- 1 | 59c4c0f3d5f0ef00cd5350b5674e941a7606d91a -------------------------------------------------------------------------------- /httpclient/fixtures/testdata/some.json: -------------------------------------------------------------------------------- 1 | { 2 | "some": "data" 3 | } -------------------------------------------------------------------------------- /.release_metadata.json: -------------------------------------------------------------------------------- 1 | { 2 | "timestamp": "2025-11-13 14:18:22+0000" 3 | } -------------------------------------------------------------------------------- /.vscode/extensions.json: -------------------------------------------------------------------------------- 1 | { 2 | "recommendations": [ 3 | "jinliming2.vscode-go-template" 4 | ] 5 | } -------------------------------------------------------------------------------- /version/version.go: -------------------------------------------------------------------------------- 1 | package version 2 | 3 | // Version of the SDK, updated manually before every tag 4 | const Version = "0.91.0" 5 | -------------------------------------------------------------------------------- /config/testdata/corrupt/.databrickscfg: -------------------------------------------------------------------------------- 1 | [nohost] 2 | token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ 3 | 4 | [notoken] 5 | host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/ -------------------------------------------------------------------------------- /qa/lock/core/lockable.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | // Lockable is the interface for a lockable resource. 4 | type Lockable interface { 5 | // GetLockId returns the lock id for the resource. 6 | GetLockId() string 7 | } 8 | -------------------------------------------------------------------------------- /NEXT_CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # NEXT CHANGELOG 2 | 3 | ## Release v0.92.0 4 | 5 | ### New Features and Improvements 6 | 7 | ### Bug Fixes 8 | 9 | ### Documentation 10 | 11 | ### Internal Changes 12 | 13 | ### API Changes 14 | -------------------------------------------------------------------------------- /config/testdata/empty_default/.databrickscfg: -------------------------------------------------------------------------------- 1 | ; The profile defined in the DEFAULT section is be used as fallback when no profile is explicitly specified. 2 | [DEFAULT] 3 | 4 | [abc] 5 | host = https://foo 6 | token = xyz 7 | -------------------------------------------------------------------------------- /examples/http-proxy/proxy/certutil.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | type certUtil interface { 4 | GetRegisterCertificateCommands(certPath string) [][]string 5 | GetDeregisterCertificateCommands(certPath string) [][]string 6 | } 7 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "makefile.extensionOutputFolder": "./.vscode", 3 | "editor.rulers": [ 4 | 80, 5 | 125 6 | ], 7 | "[go]": { 8 | "editor.formatOnSave": true 9 | } 10 | } -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | - package-ecosystem: "gomod" 4 | directory: "/" 5 | schedule: 6 | interval: "daily" 7 | commit-message: 8 | prefix: "[Internal] " 9 | open-pull-requests-limit: 0 # disable version updates 10 | -------------------------------------------------------------------------------- /credentials/u2m/error.go: -------------------------------------------------------------------------------- 1 | package u2m 2 | 3 | // InvalidRefreshTokenError is returned from PersistentAuth's Load() method 4 | // if the access token has expired and the refresh token in the token cache 5 | // is invalid. 6 | type InvalidRefreshTokenError struct { 7 | error 8 | } 9 | -------------------------------------------------------------------------------- /qa/debug.go: -------------------------------------------------------------------------------- 1 | package qa 2 | 3 | import ( 4 | "os" 5 | "path" 6 | "strings" 7 | ) 8 | 9 | // detects if test is run from "debug test" feature in VSCode 10 | func IsInDebug() bool { 11 | ex, _ := os.Executable() 12 | return strings.HasPrefix(path.Base(ex), "__debug_bin") 13 | } 14 | -------------------------------------------------------------------------------- /common/environment/environments_test.go: -------------------------------------------------------------------------------- 1 | package environment 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestDefaultEnvironmentIsReturned(t *testing.T) { 10 | assert.Equal(t, ".cloud.databricks.com", GetEnvironmentForHostname("").DnsZone) 11 | } 12 | -------------------------------------------------------------------------------- /experimental/api/limiter.go: -------------------------------------------------------------------------------- 1 | package api 2 | 3 | import "context" 4 | 5 | // Limiter is anything that can wait. It is typically used to implement 6 | // client-side rate limiting. Implementation of this interface must be 7 | // thread-safe. 8 | type Limiter interface { 9 | Wait(context.Context) error 10 | } 11 | -------------------------------------------------------------------------------- /service/compute/ext_leading_whitespace_test.go: -------------------------------------------------------------------------------- 1 | package compute 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestTrimLeadingWhitespace(t *testing.T) { 10 | assert.Equal(t, "foo\nbar\n", TrimLeadingWhitespace(` 11 | 12 | foo 13 | bar 14 | 15 | `)) 16 | } 17 | -------------------------------------------------------------------------------- /SECURITY.md: -------------------------------------------------------------------------------- 1 | ## Reporting a Vulnerability 2 | 3 | We appreciate any security concerns brought to our attention and encourage you to notify us of any potential vulnerabilities discovered in our systems or products. 4 | If you believe you have found a security vulnerability, please report it to us at [security@databricks.com](mailto:security@databricks.com). -------------------------------------------------------------------------------- /openapi/code/preview.go: -------------------------------------------------------------------------------- 1 | package code 2 | 3 | import ( 4 | "strings" 5 | 6 | "github.com/databricks/databricks-sdk-go/openapi" 7 | ) 8 | 9 | func isPrivatePreview(n *openapi.Node) bool { 10 | return strings.ToLower(n.Preview) == "private" 11 | } 12 | 13 | func isPublicPreview(n *openapi.Node) bool { 14 | return strings.ToLower(n.Preview) == "public" 15 | } 16 | -------------------------------------------------------------------------------- /.mockery.yaml: -------------------------------------------------------------------------------- 1 | with-expecter: true 2 | filename: "mock_{{.InterfaceName | snakecase}}.go" 3 | dir: "experimental/mocks/service/{{.PackageName | snakecase}}" 4 | mockname: "Mock{{.InterfaceName}}" 5 | outpkg: "{{.PackageName}}" 6 | packages: 7 | github.com/databricks/databricks-sdk-go/service: 8 | config: 9 | recursive: true 10 | include-regex: ".*Interface" -------------------------------------------------------------------------------- /service/iam/ext_model.go: -------------------------------------------------------------------------------- 1 | package iam 2 | 3 | type PartialUpdate struct { 4 | // Unique ID in the Databricks workspace. 5 | Id string `json:"-" url:"-"` 6 | 7 | Operations []Patch `json:"Operations,omitempty"` 8 | // The schema of the patch request. Must be 9 | // ["urn:ietf:params:scim:api:messages:2.0:PatchOp"]. 10 | Schemas []PatchSchema `json:"schemas,omitempty"` 11 | } 12 | -------------------------------------------------------------------------------- /.codegen.json: -------------------------------------------------------------------------------- 1 | { 2 | "mode": "go_v0", 3 | "api_changelog": true, 4 | "version": { 5 | "version/version.go": "const Version = \"$VERSION\"" 6 | }, 7 | "toolchain": { 8 | "required": [ 9 | "go" 10 | ], 11 | "post_generate": [ 12 | "make fmt", 13 | "go run github.com/vektra/mockery/v2@bfd46e35b15c2689ced221299bdcdeeff8aa0be3" 14 | ] 15 | } 16 | } -------------------------------------------------------------------------------- /openapi/model_test.go: -------------------------------------------------------------------------------- 1 | package openapi 2 | 3 | import ( 4 | "os" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestLoadFromJson(t *testing.T) { 11 | f, err := os.Open("testdata/spec.json") 12 | assert.NoError(t, err) 13 | spec, err := NewFromReader(f) 14 | assert.NoError(t, err) 15 | assert.Equal(t, "Command Execution", spec.Tags[0].Name) 16 | } 17 | -------------------------------------------------------------------------------- /qa/poll/poll.go: -------------------------------------------------------------------------------- 1 | package poll 2 | 3 | import "time" 4 | 5 | type PollFunc[R any] func(time.Duration, func(*R)) (*R, error) 6 | 7 | func Simple[R any](r R) PollFunc[R] { 8 | return func(_ time.Duration, _ func(*R)) (*R, error) { 9 | return &r, nil 10 | } 11 | } 12 | 13 | func SimpleError[R any](err error) PollFunc[R] { 14 | return func(_ time.Duration, _ func(*R)) (*R, error) { 15 | return nil, err 16 | } 17 | } 18 | -------------------------------------------------------------------------------- /.goreleaser.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | 3 | before: 4 | hooks: 5 | - go mod download 6 | 7 | builds: 8 | - skip: true 9 | 10 | # Snapshot configuration for non-release builds 11 | snapshot: 12 | version_template: "{{ .Tag }}" 13 | 14 | # Changelog configuration (overridden by workflow with tag message) 15 | changelog: 16 | sort: asc 17 | filters: 18 | exclude: 19 | - '^docs:' 20 | - '^test:' 21 | 22 | release: {} -------------------------------------------------------------------------------- /config/command.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "os/exec" 6 | "strings" 7 | 8 | "github.com/databricks/databricks-sdk-go/logger" 9 | ) 10 | 11 | // Run a command and return the output. 12 | func runCommand(ctx context.Context, cmd string, args []string) ([]byte, error) { 13 | logger.Debugf(ctx, "Running command: %s %v", cmd, strings.Join(args, " ")) 14 | return exec.Command(cmd, args...).Output() 15 | } 16 | -------------------------------------------------------------------------------- /service/common/lro/options.go: -------------------------------------------------------------------------------- 1 | package lro 2 | 3 | import "time" 4 | 5 | // LroOptions is the options for the Long Running Operations. 6 | // DO NOT USE THIS OPTION. This option is still under development 7 | // and can be updated in the future without notice. 8 | type LroOptions struct { 9 | // Timeout is the timeout for the Long Running Operations. 10 | // If not set, the default timeout is 20 minutes. 11 | Timeout time.Duration 12 | } 13 | -------------------------------------------------------------------------------- /logger/context_test.go: -------------------------------------------------------------------------------- 1 | package logger 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestLoggerContext(t *testing.T) { 11 | logger, ok := FromContext(context.Background()) 12 | assert.Nil(t, logger) 13 | assert.False(t, ok) 14 | 15 | ctx := NewContext(context.Background(), &SimpleLogger{}) 16 | logger, ok = FromContext(ctx) 17 | assert.NotNil(t, logger) 18 | assert.True(t, ok) 19 | } 20 | -------------------------------------------------------------------------------- /examples/default-auth/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/databricks/databricks-sdk-go" 7 | "github.com/databricks/databricks-sdk-go/service/compute" 8 | ) 9 | 10 | func main() { 11 | w := databricks.Must(databricks.NewWorkspaceClient()) 12 | all, err := w.Clusters.ListAll(context.Background(), compute.ListClustersRequest{}) 13 | if err != nil { 14 | panic(err) 15 | } 16 | println("Found %d clusters.", len(all)) 17 | } 18 | -------------------------------------------------------------------------------- /internal/workspace_client_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestUcAccWorkspaceClient_CurrentWorkspaceId_NoTranspile(t *testing.T) { 10 | ctx, w := ucwsTest(t) 11 | expectedWorkspaceId := MustParseInt64(GetEnvOrSkipTest(t, "THIS_WORKSPACE_ID")) 12 | workspaceId, err := w.CurrentWorkspaceID(ctx) 13 | assert.NoError(t, err) 14 | assert.Equal(t, expectedWorkspaceId, workspaceId) 15 | } 16 | -------------------------------------------------------------------------------- /examples/http-proxy/proxy/certutil_windows.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | type certUtilImpl struct{} 4 | 5 | func (w certUtilImpl) GetRegisterCertificateCommands(certPath string) [][]string { 6 | return [][]string{{"certutil", "-addstore", "-f", "ROOT", certPath}} 7 | } 8 | 9 | func (w certUtilImpl) GetDeregisterCertificateCommands(certPath string) [][]string { 10 | return [][]string{{"certutil", "-delstore", "ROOT", certPath}} 11 | } 12 | 13 | func newCertUtil() certUtil { 14 | return certUtilImpl{} 15 | } 16 | -------------------------------------------------------------------------------- /credentials/u2m/oauth_argument.go: -------------------------------------------------------------------------------- 1 | package u2m 2 | 3 | // OAuthArgument is an interface that provides the necessary information to 4 | // authenticate with PersistentAuth. Implementations of this interface must 5 | // implement either the WorkspaceOAuthArgument or AccountOAuthArgument 6 | // interface. 7 | type OAuthArgument interface { 8 | // GetCacheKey returns a unique key for the OAuthArgument. This key is used 9 | // to store and retrieve the token from the token cache. 10 | GetCacheKey() string 11 | } 12 | -------------------------------------------------------------------------------- /service/provisioning/ext_azure.go: -------------------------------------------------------------------------------- 1 | package provisioning 2 | 3 | import "fmt" 4 | 5 | const resourceIdFormat = "/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Databricks/workspaces/%s" 6 | 7 | // Return the AzureResourceID for the workspace, if it is an Azure workspace. 8 | func (w Workspace) AzureResourceId() string { 9 | if w.AzureWorkspaceInfo == nil { 10 | return "" 11 | } 12 | return fmt.Sprintf(resourceIdFormat, w.AzureWorkspaceInfo.SubscriptionId, w.AzureWorkspaceInfo.ResourceGroup, w.WorkspaceName) 13 | } 14 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, built with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | 14 | # Dependency directories (remove the comment below to include it) 15 | vendor/ 16 | 17 | coverage.txt 18 | coverage 19 | 20 | .vscode/*.lograw-specs/** 21 | .vscode/*.log 22 | .vscode/launch.json 23 | 24 | raw-specs/** 25 | out/** 26 | *.log 27 | .databricks/ 28 | .DS_Store 29 | -------------------------------------------------------------------------------- /logger/context.go: -------------------------------------------------------------------------------- 1 | package logger 2 | 3 | import "context" 4 | 5 | type logger int 6 | 7 | var loggerKey logger 8 | 9 | // NewContext returns a new Context that carries the specified logger. 10 | func NewContext(ctx context.Context, logger Logger) context.Context { 11 | return context.WithValue(ctx, loggerKey, logger) 12 | } 13 | 14 | // FromContext returns the Logger value stored in ctx, if any. 15 | func FromContext(ctx context.Context) (Logger, bool) { 16 | u, ok := ctx.Value(loggerKey).(Logger) 17 | return u, ok 18 | } 19 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature Request 3 | about: Suggest a new feature or improvement for the Go SDK for Databricks. 4 | title: "[FEATURE] " 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Problem Statement** 11 | A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] 12 | 13 | **Proposed Solution** 14 | A clear and concise description of what you want to happen. 15 | 16 | **Additional Context** 17 | Add any other context, references or screenshots about the feature request here. 18 | -------------------------------------------------------------------------------- /internal/metastore_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/catalog" 7 | "github.com/stretchr/testify/assert" 8 | "github.com/stretchr/testify/require" 9 | ) 10 | 11 | func TestUcAccMetastoreAssignments(t *testing.T) { 12 | ctx, a := ucacctTest(t) 13 | ws, err := a.MetastoreAssignments.ListAll(ctx, catalog.ListAccountMetastoreAssignmentsRequest{ 14 | MetastoreId: GetEnvOrSkipTest(t, "TEST_METASTORE_ID"), 15 | }) 16 | require.NoError(t, err) 17 | assert.NotEmpty(t, ws) 18 | } 19 | -------------------------------------------------------------------------------- /config/visitors.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "net/http" 5 | 6 | "github.com/databricks/databricks-sdk-go/useragent" 7 | ) 8 | 9 | type AuthType struct{} 10 | 11 | func authInUserAgentVisitor(c *Config) func(*http.Request) error { 12 | return func(r *http.Request) error { 13 | authType := c.AuthType 14 | if t := r.Context().Value(AuthType{}); t != nil { 15 | authType = t.(string) 16 | } 17 | ctx := useragent.InContext(r.Context(), useragent.AuthKey, authType) 18 | *r = *r.WithContext(ctx) // replace request 19 | return nil 20 | } 21 | } 22 | -------------------------------------------------------------------------------- /useragent/runtime_test.go: -------------------------------------------------------------------------------- 1 | package useragent 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestUserAgentRuntime(t *testing.T) { 10 | for _, tc := range []string{ 11 | "1.2.3", 12 | "0.0.0-dev+2e014739024a", 13 | "client.0", 14 | "foo", 15 | "15.0", 16 | "13.3", 17 | } { 18 | t.Run(tc, func(t *testing.T) { 19 | t.Setenv("DATABRICKS_RUNTIME_VERSION", tc) 20 | 21 | v := getRuntimeVersion() 22 | assert.Equal(t, tc, v) 23 | assert.NoError(t, matchAlphanumOrSemVer(v)) 24 | }) 25 | } 26 | } 27 | -------------------------------------------------------------------------------- /internal/sql_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/sql" 7 | "github.com/stretchr/testify/require" 8 | ) 9 | 10 | func TestAccSqlQueryHistory(t *testing.T) { 11 | ctx, w := workspaceTest(t) 12 | _, err := w.QueryHistory.List(ctx, sql.ListQueryHistoryRequest{ 13 | FilterBy: &sql.QueryFilter{ 14 | QueryStartTimeRange: &sql.TimeRange{ 15 | StartTimeMs: 1690243200000, // 25-07-2023 16 | EndTimeMs: 1690329600000, // 26-07-2023 17 | }, 18 | }, 19 | }) 20 | require.NoError(t, err) 21 | } 22 | -------------------------------------------------------------------------------- /config/testdata/azure/.databrickscfg: -------------------------------------------------------------------------------- 1 | [nohost] 2 | token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ 3 | 4 | [justhost] 5 | host = https://adb-123.4.azuredatabricks.net/ 6 | 7 | [pat] 8 | host = https://adb-123.4.azuredatabricks.net/ 9 | token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ 10 | 11 | [spn] 12 | host = https://adb-123.4.azuredatabricks.net/ 13 | azure_workspace_resource_id = "/a/b/c" 14 | azure_tenant_id = 12345678 15 | azure_client_id = 87654321 16 | azure_client_secret = 10000002 17 | 18 | [msi] 19 | azure_use_msi = true 20 | azure_workspace_resource_id = "/a/b/c" 21 | host = https://adb-123.4.azuredatabricks.net/ 22 | -------------------------------------------------------------------------------- /examples/http-proxy/proxy/certutil_linux.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | type certUtilImpl struct{} 4 | 5 | func (l certUtilImpl) GetRegisterCertificateCommands(certPath string) [][]string { 6 | return [][]string{{"sudo", "cp", certPath, "/usr/local/share/ca-certificates/proxy.crt"}, {"sudo", "update-ca-certificates"}} 7 | } 8 | 9 | func (l certUtilImpl) GetDeregisterCertificateCommands(certPath string) [][]string { 10 | return [][]string{{"sudo", "rm", "/usr/local/share/ca-certificates/proxy.crt"}, {"sudo", "update-ca-certificates"}} 11 | } 12 | 13 | func newCertUtil() certUtil { 14 | return certUtilImpl{} 15 | } 16 | -------------------------------------------------------------------------------- /config/config_file_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | "github.com/stretchr/testify/require" 8 | ) 9 | 10 | func TestConfigFileLoad(t *testing.T) { 11 | f, err := LoadFile("testdata/.databrickscfg") 12 | require.NoError(t, err) 13 | assert.NotNil(t, f) 14 | 15 | for _, name := range []string{ 16 | "password-with-double-quotes", 17 | "password-with-single-quotes", 18 | "password-without-quotes", 19 | } { 20 | section := f.Section(name) 21 | require.NotNil(t, section) 22 | assert.Equal(t, "%Y#X$Z", section.Key("password").String()) 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /internal/account_client_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | "github.com/stretchr/testify/require" 8 | ) 9 | 10 | func TestUcAccAccountClient_GetWorkspaceClient_NoTranspile(t *testing.T) { 11 | ctx, a := ucacctTest(t) 12 | workspaceId := MustParseInt64(GetEnvOrSkipTest(t, "TEST_WORKSPACE_ID")) 13 | ws, err := a.Workspaces.GetByWorkspaceId(ctx, int64(workspaceId)) 14 | require.NoError(t, err) 15 | w, err := a.GetWorkspaceClient(*ws) 16 | assert.NoError(t, err) 17 | me, err := w.CurrentUser.Me(ctx) 18 | assert.NoError(t, err) 19 | assert.True(t, me.Active) 20 | } 21 | -------------------------------------------------------------------------------- /examples/http-proxy/proxy/certutil_darwin.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import "os" 4 | 5 | type certUtilImpl struct { 6 | homeDir string 7 | } 8 | 9 | func (d certUtilImpl) GetRegisterCertificateCommands(certPath string) [][]string { 10 | return [][]string{{"security", "add-trusted-cert", "-r", "trustRoot", "-k", d.homeDir + "/Library/Keychains/login.keychain-db", certPath}} 11 | } 12 | 13 | func (d certUtilImpl) GetDeregisterCertificateCommands(certPath string) [][]string { 14 | return [][]string{{"security", "remove-trusted-cert", certPath}} 15 | } 16 | 17 | func newCertUtil() certUtil { 18 | homeDir := os.Getenv("HOME") 19 | return certUtilImpl{homeDir: homeDir} 20 | } 21 | -------------------------------------------------------------------------------- /config/environments_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/common/environment" 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestDefaultEnvironmentIsReturned(t *testing.T) { 11 | c := &Config{} 12 | assert.Equal(t, ".cloud.databricks.com", c.Environment().DnsZone) 13 | } 14 | 15 | func TestOverriddenEnvironmentIsReturned(t *testing.T) { 16 | c := &Config{ 17 | Host: "something.else", 18 | DatabricksEnvironment: &environment.DatabricksEnvironment{ 19 | Cloud: environment.CloudAzure, 20 | DnsZone: "holla", 21 | }, 22 | } 23 | assert.Equal(t, "holla", c.Environment().DnsZone) 24 | } 25 | -------------------------------------------------------------------------------- /service/provisioning/ext_azure_test.go: -------------------------------------------------------------------------------- 1 | package provisioning 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestAzureResourceId_AzureWorkspace(t *testing.T) { 10 | w := Workspace{ 11 | WorkspaceName: "test", 12 | AzureWorkspaceInfo: &AzureWorkspaceInfo{ 13 | SubscriptionId: "sub", 14 | ResourceGroup: "rg", 15 | }, 16 | } 17 | assert.Equal(t, "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/test", w.AzureResourceId()) 18 | } 19 | 20 | func TestAzureResourceId_NonAzureWorkspace(t *testing.T) { 21 | w := Workspace{ 22 | WorkspaceName: "test", 23 | } 24 | assert.Equal(t, "", w.AzureResourceId()) 25 | } 26 | -------------------------------------------------------------------------------- /internal/workspaceconf_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/settings" 7 | "github.com/stretchr/testify/assert" 8 | "github.com/stretchr/testify/require" 9 | ) 10 | 11 | func TestAccWorkspaceConf(t *testing.T) { 12 | ctx, w := workspaceTest(t) 13 | conf, err := w.WorkspaceConf.GetStatus(ctx, settings.GetStatusRequest{ 14 | Keys: "maxTokenLifetimeDays,enableIpAccessLists,enableWorkspaceFilesystem", 15 | }) 16 | require.NoError(t, err) 17 | assert.Equal(t, 3, len(*conf)) 18 | 19 | err = w.WorkspaceConf.SetStatus(ctx, settings.WorkspaceConf{ 20 | "enableWebTerminal": "true", 21 | }) 22 | require.NoError(t, err) 23 | } 24 | -------------------------------------------------------------------------------- /config/in_memory_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" 5 | "golang.org/x/oauth2" 6 | ) 7 | 8 | type InMemoryTokenCache struct { 9 | Tokens map[string]*oauth2.Token 10 | } 11 | 12 | // Lookup implements TokenCache. 13 | func (i *InMemoryTokenCache) Lookup(key string) (*oauth2.Token, error) { 14 | token, ok := i.Tokens[key] 15 | if !ok { 16 | return nil, cache.ErrNotFound 17 | } 18 | return token, nil 19 | } 20 | 21 | // Store implements TokenCache. 22 | func (i *InMemoryTokenCache) Store(key string, t *oauth2.Token) error { 23 | i.Tokens[key] = t 24 | return nil 25 | } 26 | 27 | var _ cache.TokenCache = (*InMemoryTokenCache)(nil) 28 | -------------------------------------------------------------------------------- /internal/testspecs/service/jsonmarshallv2/api.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | // Lorem Ipsum 4 | package jsonmarshallv2 5 | 6 | import ( 7 | "context" 8 | 9 | "github.com/databricks/databricks-sdk-go/client" 10 | ) 11 | 12 | type JsonMarshallV2Interface interface { 13 | GetResource(ctx context.Context, request GetResourceRequest) (*Resource, error) 14 | } 15 | 16 | func NewJsonMarshallV2(client *client.DatabricksClient) *JsonMarshallV2API { 17 | return &JsonMarshallV2API{ 18 | jsonMarshallV2Impl: jsonMarshallV2Impl{ 19 | client: client, 20 | }, 21 | } 22 | } 23 | 24 | // Lorem Ipsum 25 | type JsonMarshallV2API struct { 26 | jsonMarshallV2Impl 27 | } 28 | -------------------------------------------------------------------------------- /config/experimental/auth/credentials/pat.go: -------------------------------------------------------------------------------- 1 | package credentials 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | 8 | "github.com/databricks/databricks-sdk-go/config/experimental/auth" 9 | ) 10 | 11 | var errTokenRequired = errors.New("token is required") 12 | 13 | // NewPATCredentials returns a Credentials that can be used to authenticate with 14 | // a Personal Access Token. 15 | func NewPATCredentials(token string) (auth.Credentials, error) { 16 | if token == "" { 17 | return nil, errTokenRequired 18 | } 19 | return auth.CredentialsFn(func(_ context.Context) (map[string]string, error) { 20 | return map[string]string{ 21 | "Authorization": fmt.Sprintf("Bearer %s", token), 22 | }, nil 23 | }), nil 24 | } 25 | -------------------------------------------------------------------------------- /workspace_functions.go: -------------------------------------------------------------------------------- 1 | package databricks 2 | 3 | import ( 4 | "context" 5 | "strconv" 6 | 7 | "github.com/databricks/databricks-sdk-go/httpclient" 8 | ) 9 | 10 | // CurrentWorkspaceID returns the workspace ID of the workspace that this client is 11 | // connected to. 12 | func (w *WorkspaceClient) CurrentWorkspaceID(ctx context.Context) (int64, error) { 13 | var workspaceIdStr string 14 | err := w.apiClient.Do(ctx, "GET", "/api/2.0/preview/scim/v2/Me", httpclient.WithResponseHeader("X-Databricks-Org-Id", &workspaceIdStr)) 15 | if err != nil { 16 | return 0, err 17 | } 18 | workspaceId, err := strconv.ParseInt(workspaceIdStr, 10, 64) 19 | if err != nil { 20 | return 0, err 21 | } 22 | return workspaceId, nil 23 | } 24 | -------------------------------------------------------------------------------- /service/sql/dashboards_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package sql_test 4 | 5 | import ( 6 | "context" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | "github.com/databricks/databricks-sdk-go/logger" 10 | 11 | "github.com/databricks/databricks-sdk-go/service/sql" 12 | ) 13 | 14 | func ExampleDashboardsAPI_List_dashboards() { 15 | ctx := context.Background() 16 | w, err := databricks.NewWorkspaceClient() 17 | if err != nil { 18 | panic(err) 19 | } 20 | 21 | all, err := w.Dashboards.ListAll(ctx, sql.ListDashboardsRequest{}) 22 | if err != nil { 23 | panic(err) 24 | } 25 | logger.Infof(ctx, "found %v", all) 26 | } 27 | -------------------------------------------------------------------------------- /openapi/roll/compare.go: -------------------------------------------------------------------------------- 1 | package roll 2 | 3 | // This inlines the cmp package from Go 1.21 while we still support Go 1.19 and 1.20. 4 | // Once we drop support for Go 1.19 and 1.20, we can remove this file and use the cmp package directly. 5 | 6 | type Ordered interface { 7 | ~int | ~int8 | ~int16 | ~int32 | ~int64 | 8 | ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr | 9 | ~float32 | ~float64 | 10 | ~string 11 | } 12 | 13 | func isNaN[T Ordered](x T) bool { 14 | return x != x 15 | } 16 | func Compare[T Ordered](x, y T) int { 17 | xNaN := isNaN(x) 18 | yNaN := isNaN(y) 19 | if xNaN && yNaN { 20 | return 0 21 | } 22 | if xNaN || x < y { 23 | return -1 24 | } 25 | if yNaN || x > y { 26 | return +1 27 | } 28 | return 0 29 | } 30 | -------------------------------------------------------------------------------- /qa/lock/lock_options.go: -------------------------------------------------------------------------------- 1 | package lock 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/databricks/databricks-sdk-go/qa/lock/core" 8 | ) 9 | 10 | type LockOption func(*LockOptions) 11 | 12 | func WithBackend(backend core.LockBackend) LockOption { 13 | return func(opts *LockOptions) { 14 | opts.Backend = backend 15 | } 16 | } 17 | 18 | func WithLeaseDuration(duration time.Duration) LockOption { 19 | return func(opts *LockOptions) { 20 | opts.LeaseDuration = duration 21 | } 22 | } 23 | 24 | func InTest(t *testing.T) LockOption { 25 | return func(opts *LockOptions) { 26 | opts.T = t 27 | } 28 | } 29 | 30 | type LockOptions struct { 31 | LeaseDuration time.Duration 32 | Backend core.LockBackend 33 | T *testing.T 34 | } 35 | -------------------------------------------------------------------------------- /marshal/types.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import "reflect" 4 | 5 | var basicTypes = map[reflect.Kind]bool{ 6 | 7 | reflect.Bool: true, 8 | 9 | reflect.String: true, 10 | 11 | // Integers 12 | reflect.Int: true, 13 | reflect.Int8: true, 14 | reflect.Int16: true, 15 | reflect.Int32: true, 16 | reflect.Int64: true, 17 | reflect.Uint: true, 18 | reflect.Uint8: true, 19 | reflect.Uint16: true, 20 | reflect.Uint32: true, 21 | reflect.Uint64: true, 22 | reflect.Uintptr: true, 23 | 24 | // Floats 25 | reflect.Float32: true, 26 | reflect.Float64: true, 27 | 28 | // Complex 29 | reflect.Complex64: true, 30 | reflect.Complex128: true, 31 | } 32 | 33 | func isBasicType(v reflect.Type) bool { 34 | b, ok := basicTypes[v.Kind()] 35 | return b && ok 36 | } 37 | -------------------------------------------------------------------------------- /service/workspace/examples_test.go: -------------------------------------------------------------------------------- 1 | package workspace_test 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/databricks/databricks-sdk-go" 7 | "github.com/databricks/databricks-sdk-go/service/workspace" 8 | ) 9 | 10 | func ExampleReposAPI_GetByPath_checkoutBranchByPath() { 11 | ctx := context.Background() 12 | w, err := databricks.NewWorkspaceClient() 13 | if err != nil { 14 | panic(err) 15 | } 16 | // shortcut for getting RepoInfo by path 17 | repo, err := w.Repos.GetByPath(ctx, "/Repos/path/to/prod") 18 | if err != nil { 19 | panic(err) 20 | } 21 | // because you can update repo only by ID, not by path 22 | err = w.Repos.Update(ctx, workspace.UpdateRepoRequest{ 23 | RepoId: repo.Id, 24 | Branch: "v1.4.18", 25 | }) 26 | if err != nil { 27 | panic(err) 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /useragent/runtime.go: -------------------------------------------------------------------------------- 1 | package useragent 2 | 3 | import ( 4 | "os" 5 | "sync" 6 | ) 7 | 8 | // Clear cached user agent values like the DBR version or the CI/CD provider 9 | // being used. This is useful for testing. 10 | func ClearCache() { 11 | // Reset the sync.Once to their default values. This will recompute the 12 | // values on the next call to Runtime() or CiCdProvider(). 13 | runtimeOnce = sync.Once{} 14 | providerOnce = sync.Once{} 15 | } 16 | 17 | var runtimeOnce sync.Once 18 | var runtimeVersion string 19 | 20 | func getRuntimeVersion() string { 21 | v := os.Getenv("DATABRICKS_RUNTIME_VERSION") 22 | return Sanitize(v) 23 | } 24 | 25 | func Runtime() string { 26 | runtimeOnce.Do(func() { 27 | runtimeVersion = getRuntimeVersion() 28 | }) 29 | return runtimeVersion 30 | } 31 | -------------------------------------------------------------------------------- /examples/useragent/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/databricks/databricks-sdk-go" 7 | "github.com/databricks/databricks-sdk-go/useragent" 8 | ) 9 | 10 | // Run this example by running `go run ./examples/useragent` from the root of this repository. 11 | // 12 | // To run this example, ensure that you have a DEFAULT profile configured in your `.databrickscfg`. 13 | func main() { 14 | useragent.WithProduct("databricks-sdk-example", "0.0.1") 15 | useragent.WithPartner("databricks-sdk-example-abc") 16 | useragent.WithPartner("databricks-sdk-example-def") 17 | useragent.WithUserAgentExtra("test-key", "test-value") 18 | w := databricks.Must(databricks.NewWorkspaceClient()) 19 | _, err := w.CurrentUser.Me(context.Background()) 20 | if err != nil { 21 | panic(err) 22 | } 23 | } 24 | -------------------------------------------------------------------------------- /service/settings/workspace_conf_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package settings_test 4 | 5 | import ( 6 | "context" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | "github.com/databricks/databricks-sdk-go/logger" 10 | 11 | "github.com/databricks/databricks-sdk-go/service/settings" 12 | ) 13 | 14 | func ExampleWorkspaceConfAPI_GetStatus_repos() { 15 | ctx := context.Background() 16 | w, err := databricks.NewWorkspaceClient() 17 | if err != nil { 18 | panic(err) 19 | } 20 | 21 | conf, err := w.WorkspaceConf.GetStatus(ctx, settings.GetStatusRequest{ 22 | Keys: "enableWorkspaceFilesystem", 23 | }) 24 | if err != nil { 25 | panic(err) 26 | } 27 | logger.Infof(ctx, "found %v", conf) 28 | 29 | } 30 | -------------------------------------------------------------------------------- /qa/poll/poll_test.go: -------------------------------------------------------------------------------- 1 | package poll_test 2 | 3 | import ( 4 | "errors" 5 | "testing" 6 | 7 | "github.com/databricks/databricks-sdk-go/qa/poll" 8 | "github.com/databricks/databricks-sdk-go/service/sql" 9 | "github.com/stretchr/testify/assert" 10 | ) 11 | 12 | func TestSimple(t *testing.T) { 13 | waiter := sql.WaitGetWarehouseRunning[int]{ 14 | Poll: poll.Simple(sql.GetWarehouseResponse{ 15 | Id: "test", 16 | }), 17 | } 18 | res, err := waiter.Get() 19 | assert.NoError(t, err) 20 | assert.Equal(t, "test", res.Id) 21 | } 22 | 23 | func TestSimpleError(t *testing.T) { 24 | waiter := sql.WaitGetWarehouseRunning[int]{ 25 | Poll: poll.SimpleError[sql.GetWarehouseResponse](errors.New("test")), 26 | } 27 | _, err := waiter.Get() 28 | if assert.Error(t, err) { 29 | assert.Equal(t, errors.New("test"), err) 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /internal/testspecs/service/idempotencytesting/api.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | // Test service for Idempotency of Operations 4 | package idempotencytesting 5 | 6 | import ( 7 | "context" 8 | 9 | "github.com/databricks/databricks-sdk-go/client" 10 | ) 11 | 12 | type IdempotencyTestingInterface interface { 13 | CreateTestResource(ctx context.Context, request CreateTestResourceRequest) (*TestResource, error) 14 | } 15 | 16 | func NewIdempotencyTesting(client *client.DatabricksClient) *IdempotencyTestingAPI { 17 | return &IdempotencyTestingAPI{ 18 | idempotencyTestingImpl: idempotencyTestingImpl{ 19 | client: client, 20 | }, 21 | } 22 | } 23 | 24 | // Test service for Idempotency of Operations 25 | type IdempotencyTestingAPI struct { 26 | idempotencyTestingImpl 27 | } 28 | -------------------------------------------------------------------------------- /service/billing/billable_usage_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package billing_test 4 | 5 | import ( 6 | "context" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | "github.com/databricks/databricks-sdk-go/logger" 10 | 11 | "github.com/databricks/databricks-sdk-go/service/billing" 12 | ) 13 | 14 | func ExampleBillableUsageAPI_Download_usageDownload() { 15 | ctx := context.Background() 16 | a, err := databricks.NewAccountClient() 17 | if err != nil { 18 | panic(err) 19 | } 20 | 21 | resp, err := a.BillableUsage.Download(ctx, billing.DownloadRequest{ 22 | StartMonth: "2024-08", 23 | EndMonth: "2024-09", 24 | }) 25 | if err != nil { 26 | panic(err) 27 | } 28 | logger.Infof(ctx, "found %v", resp) 29 | 30 | } 31 | -------------------------------------------------------------------------------- /service/sql/query_history_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package sql_test 4 | 5 | import ( 6 | "context" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | 10 | "github.com/databricks/databricks-sdk-go/service/sql" 11 | ) 12 | 13 | func ExampleQueryHistoryAPI_List_sqlQueryHistory() { 14 | ctx := context.Background() 15 | w, err := databricks.NewWorkspaceClient() 16 | if err != nil { 17 | panic(err) 18 | } 19 | 20 | _, err = w.QueryHistory.List(ctx, sql.ListQueryHistoryRequest{ 21 | FilterBy: &sql.QueryFilter{ 22 | QueryStartTimeRange: &sql.TimeRange{ 23 | StartTimeMs: 1690243200000, 24 | EndTimeMs: 1690329600000, 25 | }, 26 | }, 27 | }) 28 | if err != nil { 29 | panic(err) 30 | } 31 | 32 | } 33 | -------------------------------------------------------------------------------- /config/experimental/auth/authconv/authconv_test.go: -------------------------------------------------------------------------------- 1 | package authconv 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | 7 | "golang.org/x/oauth2" 8 | ) 9 | 10 | type mockOauth2TokenSource func() (*oauth2.Token, error) 11 | 12 | func (t mockOauth2TokenSource) Token() (*oauth2.Token, error) { 13 | return t() 14 | } 15 | 16 | func TestIndepotency(t *testing.T) { 17 | wantErr := fmt.Errorf("test error") 18 | wantToken := &oauth2.Token{AccessToken: "test token"} 19 | ts := mockOauth2TokenSource(func() (*oauth2.Token, error) { 20 | return wantToken, wantErr 21 | }) 22 | 23 | gotToken, gotErr := OAuth2TokenSource(AuthTokenSource(ts)).Token() 24 | 25 | if gotErr != wantErr { 26 | t.Errorf("Token() = %v, want %v", gotErr, wantErr) 27 | } 28 | if gotToken != wantToken { 29 | t.Errorf("Token() = %v, want %v", gotToken, wantToken) 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /config/auth_basic.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "encoding/base64" 6 | "fmt" 7 | "net/http" 8 | 9 | "github.com/databricks/databricks-sdk-go/config/credentials" 10 | ) 11 | 12 | type BasicCredentials struct { 13 | } 14 | 15 | func (c BasicCredentials) Name() string { 16 | return "basic" 17 | } 18 | 19 | func (c BasicCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { 20 | if cfg.Username == "" || cfg.Password == "" || cfg.Host == "" { 21 | return nil, nil 22 | } 23 | tokenUnB64 := fmt.Sprintf("%s:%s", cfg.Username, cfg.Password) 24 | b64 := base64.StdEncoding.EncodeToString([]byte(tokenUnB64)) 25 | visitor := func(r *http.Request) error { 26 | r.Header.Set("Authorization", fmt.Sprintf("Basic %s", b64)) 27 | return nil 28 | } 29 | return credentials.CredentialsProviderFn(visitor), nil 30 | } 31 | -------------------------------------------------------------------------------- /config/oauth_visitors_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/stretchr/testify/assert" 8 | "golang.org/x/oauth2" 9 | ) 10 | 11 | type mockTokenSource struct { 12 | mockedTokenFunc func() (*oauth2.Token, error) 13 | } 14 | 15 | func (m mockTokenSource) Token() (*oauth2.Token, error) { 16 | return m.mockedTokenFunc() 17 | } 18 | 19 | func TestAzureReuseTokenSource(t *testing.T) { 20 | mockSource := mockTokenSource{ 21 | mockedTokenFunc: func() (*oauth2.Token, error) { 22 | return &oauth2.Token{ 23 | Expiry: time.Now().Add(35 * time.Second), 24 | }, nil 25 | }, 26 | } 27 | 28 | // Assert the token is not valid if it expires in 35 seconds. 29 | adjustedSource := azureReuseTokenSource(nil, mockSource) 30 | token, err := adjustedSource.Token() 31 | assert.NoError(t, err) 32 | assert.False(t, token.Valid()) 33 | } 34 | -------------------------------------------------------------------------------- /internal/libraries_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/compute" 7 | "github.com/stretchr/testify/require" 8 | ) 9 | 10 | func TestAccLibraries(t *testing.T) { 11 | ctx, w := workspaceTest(t) 12 | clusterId := sharedRunningCluster(t, ctx, w) 13 | 14 | err := w.Libraries.UpdateAndWait(ctx, compute.Update{ 15 | ClusterId: clusterId, 16 | Install: []compute.Library{ 17 | { 18 | Pypi: &compute.PythonPyPiLibrary{ 19 | Package: "dbl-tempo", 20 | }, 21 | }, 22 | }, 23 | }) 24 | require.NoError(t, err) 25 | 26 | err = w.Libraries.UpdateAndWait(ctx, compute.Update{ 27 | ClusterId: clusterId, 28 | Uninstall: []compute.Library{ 29 | { 30 | Pypi: &compute.PythonPyPiLibrary{ 31 | Package: "dbl-tempo", 32 | }, 33 | }, 34 | }, 35 | }) 36 | require.NoError(t, err) 37 | } 38 | -------------------------------------------------------------------------------- /service/catalog/metastore_assignments_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package catalog_test 4 | 5 | import ( 6 | "context" 7 | "os" 8 | 9 | "github.com/databricks/databricks-sdk-go" 10 | "github.com/databricks/databricks-sdk-go/logger" 11 | 12 | "github.com/databricks/databricks-sdk-go/service/catalog" 13 | ) 14 | 15 | func ExampleAccountMetastoreAssignmentsAPI_ListAll_metastoreAssignments() { 16 | ctx := context.Background() 17 | a, err := databricks.NewAccountClient() 18 | if err != nil { 19 | panic(err) 20 | } 21 | 22 | ws, err := a.MetastoreAssignments.ListAll(ctx, catalog.ListAccountMetastoreAssignmentsRequest{ 23 | MetastoreId: os.Getenv("TEST_METASTORE_ID"), 24 | }) 25 | if err != nil { 26 | panic(err) 27 | } 28 | logger.Infof(ctx, "found %v", ws) 29 | 30 | } 31 | -------------------------------------------------------------------------------- /internal/testspecs/service/jsonmarshallv2/impl.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | package jsonmarshallv2 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "net/http" 9 | 10 | "github.com/databricks/databricks-sdk-go/client" 11 | ) 12 | 13 | // unexported type that holds implementations of just JsonMarshallV2 API methods 14 | type jsonMarshallV2Impl struct { 15 | client *client.DatabricksClient 16 | } 17 | 18 | func (a *jsonMarshallV2Impl) GetResource(ctx context.Context, request GetResourceRequest) (*Resource, error) { 19 | var resource Resource 20 | path := fmt.Sprintf("/api/2.0/json-marshall/%v", request.Name) 21 | queryParams := make(map[string]any) 22 | headers := make(map[string]string) 23 | headers["Accept"] = "application/json" 24 | err := a.client.Do(ctx, http.MethodGet, path, headers, queryParams, request, &resource) 25 | return &resource, err 26 | } 27 | -------------------------------------------------------------------------------- /alias.go: -------------------------------------------------------------------------------- 1 | package databricks 2 | 3 | import ( 4 | "github.com/databricks/databricks-sdk-go/config" 5 | "github.com/databricks/databricks-sdk-go/useragent" 6 | "github.com/databricks/databricks-sdk-go/version" 7 | ) 8 | 9 | type Config config.Config 10 | 11 | // Must panics if error is not nil. It's intended to be used with 12 | // [databricks.NewWorkspaceClient] and [databricks.NewAccountClient]. 13 | func Must[T any](c T, err error) T { 14 | if err != nil { 15 | panic(err) 16 | } 17 | return c 18 | } 19 | 20 | // Version of this SDK 21 | func Version() string { 22 | return version.Version 23 | } 24 | 25 | // WithProduct is expected to be set by developers to differentiate their app from others. 26 | // 27 | // Example setting is: 28 | // 29 | // func init() { 30 | // databricks.WithProduct("your-product", "0.0.1") 31 | // } 32 | func WithProduct(name, version string) { 33 | useragent.WithProduct(name, version) 34 | } 35 | -------------------------------------------------------------------------------- /marshal/types_pint_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestPIntDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | PInt: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestPIntValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | PInt: Ptr(0), 24 | }, 25 | jsonString: `{"pint":0,"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | }, 28 | ) 29 | } 30 | 31 | func TestPIntForce(t *testing.T) { 32 | executeBasicMarshalTest(t, 33 | basicMarshalTest{ 34 | st: customStruct{ 35 | PInt: nil, 36 | ForceSendFields: []string{"PInt"}, 37 | }, 38 | jsonString: `{"childfs":{},"childnofs":{}}`, 39 | matchClassic: true, 40 | }, 41 | ) 42 | } 43 | -------------------------------------------------------------------------------- /marshal/types_pstring_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestPStrDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | PStr: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestPStrValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | PStr: Ptr(""), 24 | }, 25 | jsonString: `{"pstring":"","childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | }, 28 | ) 29 | } 30 | 31 | func TestPStrForce(t *testing.T) { 32 | executeBasicMarshalTest(t, 33 | basicMarshalTest{ 34 | st: customStruct{ 35 | PStr: nil, 36 | ForceSendFields: []string{"PStr"}, 37 | }, 38 | jsonString: `{"childfs":{},"childnofs":{}}`, 39 | matchClassic: true, 40 | }, 41 | ) 42 | } 43 | -------------------------------------------------------------------------------- /marshal/types_pbool_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestPBoolDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | PBool: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestPBoolValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | PBool: Ptr(false), 24 | }, 25 | jsonString: `{"pbool":false,"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | }, 28 | ) 29 | } 30 | 31 | func TestPBoolForce(t *testing.T) { 32 | executeBasicMarshalTest(t, 33 | basicMarshalTest{ 34 | st: customStruct{ 35 | PBool: nil, 36 | ForceSendFields: []string{"PBool"}, 37 | }, 38 | jsonString: `{"childfs":{},"childnofs":{}}`, 39 | matchClassic: true, 40 | }, 41 | ) 42 | } 43 | -------------------------------------------------------------------------------- /marshal/types_pfloat_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestPFloatDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | PFloat: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestPFloatValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | PFloat: Ptr(0.0), 24 | }, 25 | jsonString: `{"pfloat":0,"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | }, 28 | ) 29 | } 30 | 31 | func TestPFloatForce(t *testing.T) { 32 | executeBasicMarshalTest(t, 33 | basicMarshalTest{ 34 | st: customStruct{ 35 | PFloat: nil, 36 | ForceSendFields: []string{"PFloat"}, 37 | }, 38 | jsonString: `{"childfs":{},"childnofs":{}}`, 39 | matchClassic: true, 40 | }, 41 | ) 42 | } 43 | -------------------------------------------------------------------------------- /qa/lock/core/backend.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "context" 5 | "time" 6 | ) 7 | 8 | // LockBackend is the interface for a lock backend. 9 | type LockBackend interface { 10 | // PrepareBackend prepares the backend for use. This is called once per lock. 11 | // Locking is aborted if this fails. 12 | PrepareBackend(ctx context.Context, lockId string) error 13 | 14 | // AcquireLock acquires the lock. This is called once per lock. Locking is aborted 15 | // if this fails. 16 | AcquireLock(ctx context.Context, contents *LockState) error 17 | 18 | // RenewLock renews the lock. This is called periodically while the lock is held. 19 | RenewLock(ctx context.Context, leaseId string) error 20 | 21 | // ReleaseLock releases the lock. This is called once per lock. 22 | ReleaseLock(ctx context.Context, leaseId string) error 23 | 24 | // RefreshDuration returns the duration between calls to RenewLock. 25 | RefreshDuration() time.Duration 26 | } 27 | -------------------------------------------------------------------------------- /service/iam/current_user_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package iam_test 4 | 5 | import ( 6 | "context" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | "github.com/databricks/databricks-sdk-go/logger" 10 | ) 11 | 12 | func ExampleCurrentUserAPI_Me_currentUser() { 13 | ctx := context.Background() 14 | w, err := databricks.NewWorkspaceClient() 15 | if err != nil { 16 | panic(err) 17 | } 18 | 19 | me, err := w.CurrentUser.Me(ctx) 20 | if err != nil { 21 | panic(err) 22 | } 23 | logger.Infof(ctx, "found %v", me) 24 | 25 | } 26 | 27 | func ExampleCurrentUserAPI_Me_tokens() { 28 | ctx := context.Background() 29 | w, err := databricks.NewWorkspaceClient() 30 | if err != nil { 31 | panic(err) 32 | } 33 | 34 | me2, err := w.CurrentUser.Me(ctx) 35 | if err != nil { 36 | panic(err) 37 | } 38 | logger.Infof(ctx, "found %v", me2) 39 | 40 | } 41 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | Required development tools: 2 | 3 | * `go install golang.org/x/tools/cmd/goimports@latest` 4 | * `go install honnef.co/go/tools/cmd/staticcheck@v0.3.2` 5 | * `go install gotest.tools/gotestsum@latest` 6 | 7 | ## Developer Certificate of Origin 8 | 9 | To contribute to this repository, you must sign off your commits to certify 10 | that you have the right to contribute the code and that it complies with the 11 | open source license. The rules are pretty simple, if you can certify the 12 | content of [DCO](./DCO), then simply add a "Signed-off-by" line to your 13 | commit message to certify your compliance. Please use your real name as 14 | pseudonymous/anonymous contributions are not accepted. 15 | 16 | ``` 17 | Signed-off-by: Joe Smith 18 | ``` 19 | 20 | If you set your `user.name` and `user.email` git configs, you can sign your 21 | commit automatically with `git commit -s`: 22 | 23 | ``` 24 | git commit -s -m "Your commit message" 25 | ``` -------------------------------------------------------------------------------- /service/sql/data_sources_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package sql_test 4 | 5 | import ( 6 | "context" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | "github.com/databricks/databricks-sdk-go/logger" 10 | ) 11 | 12 | func ExampleDataSourcesAPI_List_queries() { 13 | ctx := context.Background() 14 | w, err := databricks.NewWorkspaceClient() 15 | if err != nil { 16 | panic(err) 17 | } 18 | 19 | srcs, err := w.DataSources.List(ctx) 20 | if err != nil { 21 | panic(err) 22 | } 23 | logger.Infof(ctx, "found %v", srcs) 24 | 25 | } 26 | 27 | func ExampleDataSourcesAPI_List_alerts() { 28 | ctx := context.Background() 29 | w, err := databricks.NewWorkspaceClient() 30 | if err != nil { 31 | panic(err) 32 | } 33 | 34 | srcs, err := w.DataSources.List(ctx) 35 | if err != nil { 36 | panic(err) 37 | } 38 | logger.Infof(ctx, "found %v", srcs) 39 | 40 | } 41 | -------------------------------------------------------------------------------- /config/experimental/auth/authconv/authconv.go: -------------------------------------------------------------------------------- 1 | package authconv 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/databricks/databricks-sdk-go/config/experimental/auth" 7 | "golang.org/x/oauth2" 8 | ) 9 | 10 | // AuthTokenSource converts an oauth2.TokenSource to an auth.TokenSource. 11 | func AuthTokenSource(ts oauth2.TokenSource) auth.TokenSource { 12 | return &authTokenSource{ts: ts} 13 | } 14 | 15 | type authTokenSource struct { 16 | ts oauth2.TokenSource 17 | } 18 | 19 | func (t *authTokenSource) Token(_ context.Context) (*oauth2.Token, error) { 20 | return t.ts.Token() 21 | } 22 | 23 | // OAuth2TokenSource converts an auth.TokenSource to an oauth2.TokenSource. 24 | func OAuth2TokenSource(ts auth.TokenSource) oauth2.TokenSource { 25 | return &oauth2TokenSource{ts: ts} 26 | } 27 | 28 | type oauth2TokenSource struct { 29 | ts auth.TokenSource 30 | } 31 | 32 | func (t *oauth2TokenSource) Token() (*oauth2.Token, error) { 33 | return t.ts.Token(context.Background()) 34 | } 35 | -------------------------------------------------------------------------------- /service/compute/ext_sort.go: -------------------------------------------------------------------------------- 1 | package compute 2 | 3 | import ( 4 | "sort" 5 | ) 6 | 7 | // readable chained sorting helper 8 | func sortByChain(s interface{}, fn func(int) sortCmp) { 9 | sort.Slice(s, func(i, j int) bool { 10 | return fn(i).Less(fn(j)) 11 | }) 12 | } 13 | 14 | type sortCmp interface { 15 | Less(o sortCmp) bool 16 | } 17 | 18 | type boolAsc bool 19 | 20 | func (b boolAsc) Less(o sortCmp) bool { 21 | return bool(b) != bool(o.(boolAsc)) && !bool(b) 22 | } 23 | 24 | type intAsc int 25 | 26 | func (ia intAsc) Less(o sortCmp) bool { 27 | return int(ia) < int(o.(intAsc)) 28 | } 29 | 30 | type strAsc string 31 | 32 | func (s strAsc) Less(o sortCmp) bool { 33 | return string(s) < string(o.(strAsc)) 34 | } 35 | 36 | type sortChain []sortCmp 37 | 38 | func (c sortChain) Less(other sortCmp) bool { 39 | o := other.(sortChain) 40 | for i := range c { 41 | if c[i].Less(o[i]) { 42 | return true 43 | } 44 | if o[i].Less(c[i]) { 45 | break 46 | } 47 | } 48 | return false 49 | } 50 | -------------------------------------------------------------------------------- /config/auth_pat.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | 7 | "github.com/databricks/databricks-sdk-go/config/credentials" 8 | authcred "github.com/databricks/databricks-sdk-go/config/experimental/auth/credentials" 9 | ) 10 | 11 | type PatCredentials struct{} 12 | 13 | func (c PatCredentials) Name() string { 14 | return "pat" 15 | } 16 | 17 | func (c PatCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { 18 | // Host inference is not supported for PAT authentication. This is 19 | // arguably a redundant check as requests will fail anyway if no 20 | // host is provided. This check exists for backward compatibility 21 | // with the previous PAT credentials implementation. 22 | if cfg.Host == "" { 23 | return nil, fmt.Errorf("host is required for PAT authentication") 24 | } 25 | creds, err := authcred.NewPATCredentials(cfg.Token) 26 | if err != nil { 27 | return nil, err 28 | } 29 | return credentials.FromCredentials(creds), nil 30 | } 31 | -------------------------------------------------------------------------------- /marshal/types_childnofs_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestChildNoFSDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | ChildNoFS: structNoFS{}, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestChildNoFSValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | ChildNoFS: structNoFS{ 24 | Int: 3, 25 | }, 26 | }, 27 | jsonString: `{"childnofs":{"childint":3},"childfs":{}}`, 28 | matchClassic: true, 29 | }, 30 | ) 31 | } 32 | 33 | func TestChildNoFSForce(t *testing.T) { 34 | executeBasicMarshalTest(t, 35 | basicMarshalTest{ 36 | st: customStruct{ 37 | ChildNoFS: structNoFS{}, 38 | ForceSendFields: []string{"ChildNoFS"}, 39 | }, 40 | jsonString: `{"childnofs":{},"childfs":{}}`, 41 | matchClassic: true, 42 | }, 43 | ) 44 | } 45 | -------------------------------------------------------------------------------- /internal/testspecs/service/httpcallv2/api.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | // Lorem Ipsum 4 | package httpcallv2 5 | 6 | import ( 7 | "context" 8 | 9 | "github.com/databricks/databricks-sdk-go/client" 10 | ) 11 | 12 | type HttpCallV2Interface interface { 13 | 14 | // This mimics "old" style post requests which have the resource inlined. 15 | CreateResource(ctx context.Context, request CreateResourceRequest) (*Resource, error) 16 | 17 | GetResource(ctx context.Context, request GetResourceRequest) (*Resource, error) 18 | 19 | // This mimics "new" style post requests which have a body field. 20 | UpdateResource(ctx context.Context, request UpdateResourceRequest) (*Resource, error) 21 | } 22 | 23 | func NewHttpCallV2(client *client.DatabricksClient) *HttpCallV2API { 24 | return &HttpCallV2API{ 25 | httpCallV2Impl: httpCallV2Impl{ 26 | client: client, 27 | }, 28 | } 29 | } 30 | 31 | // Lorem Ipsum 32 | type HttpCallV2API struct { 33 | httpCallV2Impl 34 | } 35 | -------------------------------------------------------------------------------- /marshal/types_basemap_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestBasicMapDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | BasicMap: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestBasicMapValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | BasicMap: map[string]string{ 24 | "key": "value", 25 | }, 26 | }, 27 | jsonString: `{"basicmap":{"key":"value"},"childfs":{},"childnofs":{}}`, 28 | matchClassic: true, 29 | }, 30 | ) 31 | } 32 | 33 | func TestBasicMapForce(t *testing.T) { 34 | executeBasicMarshalTest(t, 35 | basicMarshalTest{ 36 | st: customStruct{ 37 | BasicMap: nil, 38 | ForceSendFields: []string{"BasicMap"}, 39 | }, 40 | jsonString: `{"childfs":{},"childnofs":{}}`, 41 | matchClassic: true, 42 | }, 43 | ) 44 | } 45 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/sdk-issue.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: SDK Issue 3 | about: Use this to report an issue with the Go SDK for Databricks. 4 | title: "[ISSUE] " 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Description** 11 | A clear and concise description of what the bug is. 12 | 13 | **Reproduction** 14 | A minimal code sample demonstrating the bug. 15 | 16 | **Expected behavior** 17 | A clear and concise description of what you expected to happen. 18 | 19 | **Is it a regression?** 20 | Did this work in a previous version of the SDK? If so, which versions did you try? 21 | 22 | **Debug Logs** 23 | The SDK logs helpful debugging information when debug logging is enabled. Set the log level to Trace by configuring the default logger to log at trace (for example: add `logger.DefaultLogger = &logger.SimpleLogger{Level: logger.LevelTrace}` to your program), and include the logs here. 24 | 25 | **Other Information** 26 | - OS: [e.g. macOS] 27 | - Version: [e.g. 0.1.0] 28 | 29 | **Additional context** 30 | Add any other context about the problem here. 31 | -------------------------------------------------------------------------------- /marshal/types_int_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestIntDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | Int: 0, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestIntValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | Int: 2, 24 | }, 25 | jsonString: `{"int":2,"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | unmarshalForceSendField: []string{"Int"}, 28 | }, 29 | ) 30 | } 31 | 32 | func TestIntForce(t *testing.T) { 33 | executeBasicMarshalTest(t, 34 | basicMarshalTest{ 35 | st: customStruct{ 36 | Int: 0, 37 | ForceSendFields: []string{"Int"}, 38 | }, 39 | jsonString: `{"int":0,"childfs":{},"childnofs":{}}`, 40 | matchClassic: false, 41 | unmarshalForceSendField: []string{"Int"}, 42 | }, 43 | ) 44 | } 45 | -------------------------------------------------------------------------------- /marshal/types_maptointerface_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestMapToInterfaceDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | MapToInt: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestMapToInterfaceValue(t *testing.T) { 20 | element := map[string]interface{}{ 21 | "key": "3", 22 | } 23 | executeBasicMarshalTest(t, 24 | basicMarshalTest{ 25 | st: customStruct{ 26 | MapToInt: element, 27 | }, 28 | jsonString: `{"maptoint":{"key":"3"},"childfs":{},"childnofs":{}}`, 29 | matchClassic: true, 30 | }, 31 | ) 32 | } 33 | 34 | func TestMapToInterfaceForce(t *testing.T) { 35 | executeBasicMarshalTest(t, 36 | basicMarshalTest{ 37 | st: customStruct{ 38 | MapToInt: nil, 39 | ForceSendFields: []string{"MapToInt"}, 40 | }, 41 | jsonString: `{"childfs":{},"childnofs":{}}`, 42 | matchClassic: true, 43 | }, 44 | ) 45 | } 46 | -------------------------------------------------------------------------------- /qa/lock/lockables.go: -------------------------------------------------------------------------------- 1 | package lock 2 | 3 | import "github.com/databricks/databricks-sdk-go/qa/lock/internal" 4 | 5 | // GitCredentials are unique to the user and workspace. 6 | type GitCredentials struct { 7 | WorkspaceHost string 8 | Username string 9 | } 10 | 11 | func (g GitCredentials) GetLockId() string { 12 | return internal.GenerateBlobName(g) 13 | } 14 | 15 | // Some operations require locking the entire workspace. 16 | type Workspace struct { 17 | WorkspaceId string 18 | } 19 | 20 | func (w Workspace) GetLockId() string { 21 | return internal.GenerateBlobName(w) 22 | } 23 | 24 | // LockableImpl is a default implementation of Lockable. If r is a struct, it uses 25 | // reflection to generate a unique lock ID based on the name and fields of the struct. 26 | // If r is a string, it uses the string as the lock ID. 27 | type LockableImpl[R any] struct { 28 | r R 29 | } 30 | 31 | func NewLockable[R any](r R) LockableImpl[R] { 32 | return LockableImpl[R]{r: r} 33 | } 34 | func (l LockableImpl[R]) GetLockId() string { 35 | return internal.GenerateBlobName(l.r) 36 | } 37 | -------------------------------------------------------------------------------- /internal/env/env.go: -------------------------------------------------------------------------------- 1 | package env 2 | 3 | import ( 4 | "os" 5 | "strings" 6 | "testing" 7 | ) 8 | 9 | // CleanupEnvironment sets up a pristine environment containing only $PATH and $HOME. 10 | // The original environment is restored upon test completion. 11 | // Note: use of this function is incompatible with parallel execution. 12 | func CleanupEnvironment(t *testing.T) { 13 | // Restore environment when test finishes. 14 | environ := os.Environ() 15 | t.Cleanup(func() { 16 | // Restore original environment. 17 | for _, kv := range environ { 18 | kvs := strings.SplitN(kv, "=", 2) 19 | os.Setenv(kvs[0], kvs[1]) 20 | } 21 | }) 22 | 23 | path := os.Getenv("PATH") 24 | pwd := os.Getenv("PWD") 25 | os.Clearenv() 26 | 27 | // We use t.Setenv instead of os.Setenv because the former actively 28 | // prevents a test being run with t.Parallel. Modifying the environment 29 | // within a test is not compatible with running tests in parallel 30 | // because of isolation; the environment is scoped to the process. 31 | t.Setenv("PATH", path) 32 | t.Setenv("HOME", pwd) 33 | } 34 | -------------------------------------------------------------------------------- /service/compute/libraries_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package compute_test 4 | 5 | import ( 6 | "context" 7 | "os" 8 | 9 | "github.com/databricks/databricks-sdk-go" 10 | 11 | "github.com/databricks/databricks-sdk-go/service/compute" 12 | ) 13 | 14 | func ExampleLibrariesAPI_UpdateAndWait_libraries() { 15 | ctx := context.Background() 16 | w, err := databricks.NewWorkspaceClient() 17 | if err != nil { 18 | panic(err) 19 | } 20 | 21 | clusterId := func() string { 22 | clusterId := os.Getenv("DATABRICKS_CLUSTER_ID") 23 | err := w.Clusters.EnsureClusterIsRunning(ctx, clusterId) 24 | if err != nil { 25 | panic(err) 26 | } 27 | return clusterId 28 | }() 29 | 30 | err = w.Libraries.UpdateAndWait(ctx, compute.Update{ 31 | ClusterId: clusterId, 32 | Install: []compute.Library{compute.Library{ 33 | Pypi: &compute.PythonPyPiLibrary{ 34 | Package: "dbl-tempo", 35 | }, 36 | }}, 37 | }) 38 | if err != nil { 39 | panic(err) 40 | } 41 | 42 | } 43 | -------------------------------------------------------------------------------- /marshal/types_bool_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestBoolDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | Bool: false, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestBoolValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | Bool: true, 24 | }, 25 | jsonString: `{"bool":true,"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | unmarshalForceSendField: []string{"Bool"}, 28 | }, 29 | ) 30 | } 31 | 32 | func TestBoolForce(t *testing.T) { 33 | executeBasicMarshalTest(t, 34 | basicMarshalTest{ 35 | st: customStruct{ 36 | Bool: false, 37 | ForceSendFields: []string{"Bool"}, 38 | }, 39 | jsonString: `{"bool":false,"childfs":{},"childnofs":{}}`, 40 | matchClassic: false, 41 | unmarshalForceSendField: []string{"Bool"}, 42 | }, 43 | ) 44 | } 45 | -------------------------------------------------------------------------------- /marshal/types_float_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestFloatDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | Float: 0.0, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestFloatValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | Float: 2.3, 24 | }, 25 | jsonString: `{"float":2.3,"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | unmarshalForceSendField: []string{"Float"}, 28 | }, 29 | ) 30 | } 31 | 32 | func TestFloatForce(t *testing.T) { 33 | executeBasicMarshalTest(t, 34 | basicMarshalTest{ 35 | st: customStruct{ 36 | Float: 0, 37 | ForceSendFields: []string{"Float"}, 38 | }, 39 | jsonString: `{"float":0,"childfs":{},"childnofs":{}}`, 40 | matchClassic: false, 41 | unmarshalForceSendField: []string{"Float"}, 42 | }, 43 | ) 44 | } 45 | -------------------------------------------------------------------------------- /marshal/types_string_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestStringDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | Str: "", 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestStringValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | Str: "val", 24 | }, 25 | jsonString: `{"string":"val","childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | unmarshalForceSendField: []string{"Str"}, 28 | }, 29 | ) 30 | } 31 | 32 | func TestStringForce(t *testing.T) { 33 | executeBasicMarshalTest(t, 34 | basicMarshalTest{ 35 | st: customStruct{ 36 | Str: "", 37 | ForceSendFields: []string{"Str"}, 38 | }, 39 | jsonString: `{"string":"","childfs":{},"childnofs":{}}`, 40 | matchClassic: false, 41 | unmarshalForceSendField: []string{"Str"}, 42 | }, 43 | ) 44 | } 45 | -------------------------------------------------------------------------------- /service/compute/ext_leading_whitespace.go: -------------------------------------------------------------------------------- 1 | package compute 2 | 3 | import ( 4 | "strings" 5 | ) 6 | 7 | // TrimLeadingWhitespace removes leading whitespace, so that Python code blocks 8 | // that are embedded into Go code still could be interpreted properly. 9 | func TrimLeadingWhitespace(commandStr string) (newCommand string) { 10 | lines := strings.Split(strings.ReplaceAll(commandStr, "\t", " "), "\n") 11 | leadingWhitespace := 1<<31 - 1 12 | for _, line := range lines { 13 | for pos, char := range line { 14 | if char == ' ' || char == '\t' { 15 | continue 16 | } 17 | // first non-whitespace character 18 | if pos < leadingWhitespace { 19 | leadingWhitespace = pos 20 | } 21 | // is not needed further 22 | break 23 | } 24 | } 25 | for i := 0; i < len(lines); i++ { 26 | if lines[i] == "" || strings.Trim(lines[i], " \t") == "" { 27 | continue 28 | } 29 | if len(lines[i]) < leadingWhitespace { 30 | newCommand += lines[i] + "\n" // or not.. 31 | } else { 32 | newCommand += lines[i][leadingWhitespace:] + "\n" 33 | } 34 | } 35 | return 36 | } 37 | -------------------------------------------------------------------------------- /internal/testspecs/service/idempotencytesting/model.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | package idempotencytesting 4 | 5 | import ( 6 | "github.com/databricks/databricks-sdk-go/marshal" 7 | ) 8 | 9 | type CreateTestResourceRequest struct { 10 | RequestId string `json:"-" url:"request_id,omitempty"` 11 | 12 | TestResource TestResource `json:"test_resource"` 13 | 14 | ForceSendFields []string `json:"-" url:"-"` 15 | } 16 | 17 | func (s *CreateTestResourceRequest) UnmarshalJSON(b []byte) error { 18 | return marshal.Unmarshal(b, s) 19 | } 20 | 21 | func (s CreateTestResourceRequest) MarshalJSON() ([]byte, error) { 22 | return marshal.Marshal(s) 23 | } 24 | 25 | type TestResource struct { 26 | Id string `json:"id,omitempty"` 27 | 28 | Name string `json:"name,omitempty"` 29 | 30 | ForceSendFields []string `json:"-" url:"-"` 31 | } 32 | 33 | func (s *TestResource) UnmarshalJSON(b []byte) error { 34 | return marshal.Unmarshal(b, s) 35 | } 36 | 37 | func (s TestResource) MarshalJSON() ([]byte, error) { 38 | return marshal.Marshal(s) 39 | } 40 | -------------------------------------------------------------------------------- /apierr/unwrap.go: -------------------------------------------------------------------------------- 1 | package apierr 2 | 3 | func inheritErr(err error, message string) *wrapError { 4 | return &wrapError{ 5 | wrap: err, 6 | message: message, 7 | } 8 | } 9 | 10 | type wrapError struct { 11 | wrap error 12 | message string 13 | } 14 | 15 | func (e *wrapError) Error() string { 16 | return e.message 17 | } 18 | 19 | func (e *wrapError) Unwrap() error { 20 | return e.wrap 21 | } 22 | 23 | func ByStatusCode(statusCode int) (error, bool) { 24 | err, ok := statusCodeMapping[statusCode] 25 | return err, ok 26 | } 27 | 28 | // Unwrap error for easier client code checking 29 | // 30 | // See https://pkg.go.dev/errors#example-Unwrap 31 | func (apiError *APIError) Unwrap() error { 32 | if apiError.unwrap != nil { 33 | return apiError.unwrap 34 | } 35 | byErrorCode, ok := errorCodeMapping[apiError.ErrorCode] 36 | if ok { 37 | return byErrorCode 38 | } 39 | byStatusCode, ok := ByStatusCode(apiError.StatusCode) 40 | if ok { 41 | return byStatusCode 42 | } 43 | // A nil error returned from e.Unwrap() indicates that e does not wrap 44 | // any error. 45 | return nil 46 | } 47 | -------------------------------------------------------------------------------- /.github/workflows/codeql.yml: -------------------------------------------------------------------------------- 1 | name: "CodeQL Advanced" 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | schedule: 9 | - cron: '41 1 * * 0' 10 | 11 | jobs: 12 | analyze: 13 | name: Analyze (${{ matrix.language }}) 14 | runs-on: 15 | group: databricks-deco-testing-runner-group 16 | labels: ubuntu-latest-deco 17 | permissions: 18 | # required for all workflows 19 | security-events: write 20 | 21 | strategy: 22 | fail-fast: false 23 | matrix: 24 | include: 25 | - language: go 26 | build-mode: autobuild 27 | 28 | steps: 29 | - name: Checkout repository 30 | uses: actions/checkout@v4 31 | 32 | # Initializes the CodeQL tools for scanning. 33 | - name: Initialize CodeQL 34 | uses: github/codeql-action/init@v3 35 | with: 36 | languages: ${{ matrix.language }} 37 | build-mode: ${{ matrix.build-mode }} 38 | 39 | - name: Perform CodeQL Analysis 40 | uses: github/codeql-action/analyze@v3 41 | with: 42 | category: "/language:${{matrix.language}}" 43 | -------------------------------------------------------------------------------- /httpclient/timeout_context_test.go: -------------------------------------------------------------------------------- 1 | package httpclient 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | "time" 7 | 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | func TestTimeoutContextTimeout(t *testing.T) { 12 | ctx := context.Background() 13 | ctx, _ = newTimeoutContext(ctx, time.Millisecond*50) 14 | time.Sleep(time.Millisecond * 100) 15 | 16 | // The context should have timed out. 17 | assert.Equal(t, context.Canceled, ctx.Err()) 18 | } 19 | 20 | func TestTimeoutContextTick(t *testing.T) { 21 | ctx := context.Background() 22 | ctx, ticker := newTimeoutContext(ctx, time.Millisecond*50) 23 | 24 | // Extend the deadline a couple of times. 25 | for i := 0; i < 5; i++ { 26 | ticker.Tick() 27 | time.Sleep(time.Millisecond * 25) 28 | } 29 | 30 | // The context should not have timed out. 31 | assert.Nil(t, ctx.Err()) 32 | } 33 | 34 | func TestTimeoutContextCancel(t *testing.T) { 35 | ctx := context.Background() 36 | ctx, ticker := newTimeoutContext(ctx, time.Millisecond*50) 37 | 38 | // Cancel the context. 39 | ticker.Cancel() 40 | 41 | // The context should have timed out. 42 | assert.Equal(t, context.Canceled, ctx.Err()) 43 | } 44 | -------------------------------------------------------------------------------- /httpclient/fixtures/map_transport.go: -------------------------------------------------------------------------------- 1 | package fixtures 2 | 3 | import ( 4 | "fmt" 5 | "net/http" 6 | ) 7 | 8 | // MappingTransport stubs are 2 lines shorter than [SliceTransport] 9 | type MappingTransport map[string]HTTPFixture 10 | 11 | func (fixtures MappingTransport) SkipRetryOnIO() bool { 12 | return true 13 | } 14 | 15 | func (fixtures MappingTransport) RoundTrip(req *http.Request) (*http.Response, error) { 16 | key := fmt.Sprintf("%s %s", req.Method, resourceFromRequest(req)) 17 | f, ok := fixtures[key] 18 | if ok { 19 | err := f.AssertHeaders(req) 20 | if err != nil { 21 | return nil, fmt.Errorf("headers: %w", err) 22 | } 23 | err = f.AssertRequest(req) 24 | if err != nil { 25 | return nil, fmt.Errorf("body: %w", err) 26 | } 27 | return f.Reply(req) 28 | } 29 | expectedRequest, err := bodyStub(req) 30 | if err != nil { 31 | return nil, fmt.Errorf("marshal: %w", err) 32 | } 33 | // whitespace in this string is very important for unit tests 34 | stub := fmt.Sprintf(`"%s": { 35 | %sResponse: XXX { 36 | // fill in specific fields... 37 | }, 38 | },`, key, expectedRequest) 39 | return nil, fmt.Errorf("missing stub, please add: %s", stub) 40 | } 41 | -------------------------------------------------------------------------------- /config/reflect.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import "reflect" 4 | 5 | var kindMap = map[reflect.Kind]string{ 6 | reflect.Bool: "Bool", 7 | reflect.Int: "Int", 8 | reflect.Int8: "Int8", 9 | reflect.Int16: "Int16", 10 | reflect.Int32: "Int32", 11 | reflect.Int64: "Int64", 12 | reflect.Uint: "Uint", 13 | reflect.Uint8: "Uint8", 14 | reflect.Uint16: "Uint16", 15 | reflect.Uint32: "Uint32", 16 | reflect.Uint64: "Uint64", 17 | reflect.Uintptr: "Uintptr", 18 | reflect.Float32: "Float32", 19 | reflect.Float64: "Float64", 20 | reflect.Complex64: "Complex64", 21 | reflect.Complex128: "Complex128", 22 | reflect.Array: "Array", 23 | reflect.Chan: "Chan", 24 | reflect.Func: "Func", 25 | reflect.Interface: "Interface", 26 | reflect.Ptr: "Ptr", 27 | reflect.Slice: "Slice", 28 | reflect.String: "String", 29 | reflect.Struct: "Struct", 30 | reflect.UnsafePointer: "UnsafePointer", 31 | } 32 | 33 | func reflectKind(k reflect.Kind) string { 34 | n, ok := kindMap[k] 35 | if !ok { 36 | return "other" 37 | } 38 | return n 39 | } 40 | -------------------------------------------------------------------------------- /credentials/u2m/workspace_oauth_argument_test.go: -------------------------------------------------------------------------------- 1 | package u2m 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestValidateHost(t *testing.T) { 10 | tests := []struct { 11 | host string 12 | want string 13 | }{ 14 | // Valid hosts 15 | {"https://some-host.com", ""}, 16 | {"http://127.0.0.1", ""}, 17 | {"http://127.0.0.1:5656", ""}, 18 | 19 | // Invalid hosts 20 | {"http://some-host.com", "host must start with 'https://': http://some-host.com"}, 21 | {"https://some-host.com/", "host must not have a trailing slash: https://some-host.com/"}, 22 | } 23 | 24 | for _, test := range tests { 25 | err := validateHost(test.host) 26 | if test.want == "" { 27 | assert.NoError(t, err) 28 | } else { 29 | assert.EqualError(t, err, test.want) 30 | } 31 | 32 | _, err = NewBasicWorkspaceOAuthArgument(test.host) 33 | if test.want == "" { 34 | assert.NoError(t, err) 35 | } else { 36 | assert.EqualError(t, err, test.want) 37 | } 38 | 39 | _, err = NewBasicAccountOAuthArgument(test.host, "123") 40 | if test.want == "" { 41 | assert.NoError(t, err) 42 | } else { 43 | assert.EqualError(t, err, test.want) 44 | } 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /service/qualitymonitorv2/interface.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | package qualitymonitorv2 4 | 5 | import ( 6 | "context" 7 | ) 8 | 9 | // Manage data quality of UC objects (currently support `schema`) 10 | // 11 | // Deprecated: Do not use this interface, it will be removed in a future version of the SDK. 12 | type QualityMonitorV2Service interface { 13 | 14 | // Create a quality monitor on UC object 15 | CreateQualityMonitor(ctx context.Context, request CreateQualityMonitorRequest) (*QualityMonitor, error) 16 | 17 | // Delete a quality monitor on UC object 18 | DeleteQualityMonitor(ctx context.Context, request DeleteQualityMonitorRequest) error 19 | 20 | // Read a quality monitor on UC object 21 | GetQualityMonitor(ctx context.Context, request GetQualityMonitorRequest) (*QualityMonitor, error) 22 | 23 | // (Unimplemented) List quality monitors 24 | ListQualityMonitor(ctx context.Context, request ListQualityMonitorRequest) (*ListQualityMonitorResponse, error) 25 | 26 | // (Unimplemented) Update a quality monitor on UC object 27 | UpdateQualityMonitor(ctx context.Context, request UpdateQualityMonitorRequest) (*QualityMonitor, error) 28 | } 29 | -------------------------------------------------------------------------------- /credentials/u2m/cache/cache.go: -------------------------------------------------------------------------------- 1 | /* 2 | Package cache provides an interface for storing and looking up OAuth tokens. 3 | 4 | The cache should be primarily used for user-to-machine (U2M) OAuth flows. In U2M 5 | OAuth flows, the application needs to store the token for later use, such as in 6 | a separate process, and the cache provides a way to do so without requiring the 7 | user to follow the OAuth flow again. 8 | 9 | In machine-to-machine (M2M) OAuth flows, the application is configured with a 10 | secret and can fetch a new token on demand without user interaction, so the 11 | token cache is not necessary. 12 | */ 13 | package cache 14 | 15 | import ( 16 | "errors" 17 | 18 | "golang.org/x/oauth2" 19 | ) 20 | 21 | var ErrNotFound = errors.New("token not found") 22 | 23 | // TokenCache is an interface for storing and looking up OAuth tokens. 24 | type TokenCache interface { 25 | // Store stores the token with the given key, replacing any existing token. 26 | // If t is nil, it deletes the token. 27 | Store(key string, t *oauth2.Token) error 28 | 29 | // Lookup looks up the token with the given key. If the token is not found, it 30 | // returns ErrNotFound. 31 | Lookup(key string) (*oauth2.Token, error) 32 | } 33 | -------------------------------------------------------------------------------- /DCO: -------------------------------------------------------------------------------- 1 | Developer's Certificate of Origin 1.1 2 | 3 | By making a contribution to this project, I certify that: 4 | 5 | (a) The contribution was created in whole or in part by me and I 6 | have the right to submit it under the open source license 7 | indicated in the file; or 8 | 9 | (b) The contribution is based upon previous work that, to the best 10 | of my knowledge, is covered under an appropriate open source 11 | license and I have the right under that license to submit that 12 | work with modifications, whether created in whole or in part 13 | by me, under the same open source license (unless I am 14 | permitted to submit under a different license), as indicated 15 | in the file; or 16 | 17 | (c) The contribution was provided directly to me by some other 18 | person who certified (a), (b) or (c) and I have not modified 19 | it. 20 | 21 | (d) I understand and agree that this project and the contribution 22 | are public and that a record of the contribution (including all 23 | personal information I submit with it, including my sign-off) is 24 | maintained indefinitely and may be redistributed consistent with 25 | this project or the open source license(s) involved. 26 | -------------------------------------------------------------------------------- /logger/simple.go: -------------------------------------------------------------------------------- 1 | package logger 2 | 3 | import ( 4 | "context" 5 | "log" 6 | ) 7 | 8 | type SimpleLogger struct { 9 | Level Level 10 | } 11 | 12 | func (l *SimpleLogger) Enabled(_ context.Context, level Level) bool { 13 | return level >= l.Level 14 | } 15 | 16 | func (l *SimpleLogger) Tracef(ctx context.Context, format string, v ...any) { 17 | if !l.Enabled(ctx, LevelTrace) { 18 | return 19 | } 20 | log.Printf("[TRACE] "+format, v...) 21 | } 22 | 23 | func (l *SimpleLogger) Debugf(ctx context.Context, format string, v ...any) { 24 | if !l.Enabled(ctx, LevelDebug) { 25 | return 26 | } 27 | log.Printf("[DEBUG] "+format, v...) 28 | } 29 | 30 | func (l *SimpleLogger) Infof(ctx context.Context, format string, v ...any) { 31 | if !l.Enabled(ctx, LevelInfo) { 32 | return 33 | } 34 | log.Printf("[INFO] "+format, v...) 35 | } 36 | 37 | func (l *SimpleLogger) Warnf(ctx context.Context, format string, v ...any) { 38 | if !l.Enabled(ctx, LevelWarn) { 39 | return 40 | } 41 | log.Printf("[WARN] "+format, v...) 42 | } 43 | 44 | func (l *SimpleLogger) Errorf(ctx context.Context, format string, v ...any) { 45 | if !l.Enabled(ctx, LevelError) { 46 | return 47 | } 48 | log.Printf("[ERROR] "+format, v...) 49 | } 50 | -------------------------------------------------------------------------------- /service/agentbricks/interface.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | package agentbricks 4 | 5 | import ( 6 | "context" 7 | ) 8 | 9 | // The Custom LLMs service manages state and powers the UI for the Custom LLM 10 | // product. 11 | // 12 | // Deprecated: Do not use this interface, it will be removed in a future version of the SDK. 13 | type AgentBricksService interface { 14 | 15 | // Cancel a Custom LLM Optimization Run. 16 | CancelOptimize(ctx context.Context, request CancelCustomLlmOptimizationRunRequest) error 17 | 18 | // Create a Custom LLM. 19 | CreateCustomLlm(ctx context.Context, request CreateCustomLlmRequest) (*CustomLlm, error) 20 | 21 | // Delete a Custom LLM. 22 | DeleteCustomLlm(ctx context.Context, request DeleteCustomLlmRequest) error 23 | 24 | // Get a Custom LLM. 25 | GetCustomLlm(ctx context.Context, request GetCustomLlmRequest) (*CustomLlm, error) 26 | 27 | // Start a Custom LLM Optimization Run. 28 | StartOptimize(ctx context.Context, request StartCustomLlmOptimizationRunRequest) (*CustomLlm, error) 29 | 30 | // Update a Custom LLM. 31 | UpdateCustomLlm(ctx context.Context, request UpdateCustomLlmRequest) (*CustomLlm, error) 32 | } 33 | -------------------------------------------------------------------------------- /examples/aad-explicit-spn/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bufio" 5 | "context" 6 | "fmt" 7 | "os" 8 | "strings" 9 | 10 | "github.com/databricks/databricks-sdk-go" 11 | "github.com/databricks/databricks-sdk-go/config" 12 | "github.com/databricks/databricks-sdk-go/service/compute" 13 | ) 14 | 15 | func main() { 16 | w := databricks.Must(databricks.NewWorkspaceClient(&databricks.Config{ 17 | Host: askFor("Host:"), 18 | AzureResourceID: askFor("Azure Resource ID:"), 19 | AzureTenantID: askFor("AAD Tenant ID:"), 20 | AzureClientID: askFor("AAD Client ID:"), 21 | AzureClientSecret: askFor("AAD Client Secret:"), 22 | Credentials: config.AzureClientSecretCredentials{}, 23 | })) 24 | all, err := w.Clusters.ListAll(context.Background(), compute.ListClustersRequest{}) 25 | if err != nil { 26 | panic(err) 27 | } 28 | for _, c := range all { 29 | println(c.ClusterName) 30 | } 31 | } 32 | 33 | func askFor(prompt string) string { 34 | var s string 35 | r := bufio.NewReader(os.Stdin) 36 | for { 37 | fmt.Fprint(os.Stdout, prompt+" ") 38 | s, _ = r.ReadString('\n') 39 | s = strings.TrimSpace(s) 40 | if s != "" { 41 | break 42 | } 43 | } 44 | return s 45 | } 46 | -------------------------------------------------------------------------------- /marshal/types_interface_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestInterfaceDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | Interface: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestInterfaceValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | Interface: structNoFS{ 24 | Int: 3, 25 | }, 26 | }, 27 | jsonString: `{"interface":{"childint":3},"childfs":{},"childnofs":{}}`, 28 | matchClassic: true, 29 | unmarshalResult: &customStruct{ //Reflect can't determine the original type, so unmarshal returns this instead. 30 | Interface: map[string]interface{}{ 31 | "childint": 3.0, 32 | }, 33 | }, 34 | }, 35 | ) 36 | } 37 | 38 | func TestInterfaceForce(t *testing.T) { 39 | executeBasicMarshalTest(t, 40 | basicMarshalTest{ 41 | st: customStruct{ 42 | Interface: nil, 43 | ForceSendFields: []string{"Interface"}, 44 | }, 45 | jsonString: `{"childfs":{},"childnofs":{}}`, 46 | matchClassic: true, 47 | }, 48 | ) 49 | } 50 | -------------------------------------------------------------------------------- /config/auth_m2m.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | 7 | "golang.org/x/oauth2" 8 | "golang.org/x/oauth2/clientcredentials" 9 | 10 | "github.com/databricks/databricks-sdk-go/config/credentials" 11 | "github.com/databricks/databricks-sdk-go/logger" 12 | ) 13 | 14 | type M2mCredentials struct{} 15 | 16 | func (c M2mCredentials) Name() string { 17 | return "oauth-m2m" 18 | } 19 | 20 | func (c M2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { 21 | if cfg.ClientID == "" || cfg.ClientSecret == "" { 22 | return nil, nil 23 | } 24 | endpoints, err := cfg.getOidcEndpoints(ctx) 25 | if err != nil { 26 | return nil, fmt.Errorf("oidc: %w", err) 27 | } 28 | logger.Debugf(ctx, "Generating Databricks OAuth token for Service Principal (%s)", cfg.ClientID) 29 | ts := (&clientcredentials.Config{ 30 | ClientID: cfg.ClientID, 31 | ClientSecret: cfg.ClientSecret, 32 | AuthStyle: oauth2.AuthStyleInHeader, 33 | TokenURL: endpoints.TokenEndpoint, 34 | Scopes: []string{"all-apis"}, 35 | }).TokenSource(ctx) 36 | 37 | visitor := refreshableVisitor(ts) 38 | return credentials.NewOAuthCredentialsProvider(visitor, ts.Token), nil 39 | } 40 | -------------------------------------------------------------------------------- /examples/zerolog/README.md: -------------------------------------------------------------------------------- 1 | # zerolog 2 | 3 | This example shows how to use [zerolog][zerolog] with this SDK. 4 | 5 | [zerolog]: https://pkg.go.dev/github.com/rs/zerolog 6 | 7 | ## Output 8 | 9 | With JSON output: 10 | 11 | ```text 12 | {"level":"trace","global":true,"time":"2023-03-10T08:34:56+01:00","message":"Loading config via environment"} 13 | {"level":"trace","global":true,"time":"2023-03-10T08:34:56+01:00","message":"Loading config via config-file"} 14 | {"level":"debug","global":true,"time":"2023-03-10T08:34:56+01:00","message":"Loading DEFAULT profile from [...]/.databrickscfg"} 15 | {"level":"trace","global":false,"time":"2023-03-10T08:34:56+01:00","message":"Attempting to configure auth: pat"} 16 | {"level":"debug","global":false,"time":"2023-03-10T08:34:57+01:00","message":"GET /api/2.0/preview/scim/v2/Me\n< HTTP/2.0 200 OK\n... 17 | ``` 18 | 19 | With text output: 20 | ```text 21 | 8:35AM TRC Loading config via environment global=true 22 | 8:35AM TRC Loading config via config-file global=true 23 | 8:35AM DBG Loading DEFAULT profile from [...]/.databrickscfg global=true 24 | 8:35AM TRC Attempting to configure auth: pat global=false 25 | 8:35AM DBG GET /api/2.0/preview/scim/v2/Me 26 | < HTTP/2.0 200 OK 27 | ... 28 | ``` 29 | -------------------------------------------------------------------------------- /config/visitors_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "net/http" 6 | "strings" 7 | "testing" 8 | 9 | "github.com/databricks/databricks-sdk-go/useragent" 10 | "github.com/stretchr/testify/assert" 11 | ) 12 | 13 | func TestAuthInUserAgentVisitorDefault(t *testing.T) { 14 | c := &Config{ 15 | AuthType: "PAT", 16 | } 17 | request := &http.Request{} 18 | err := authInUserAgentVisitor(c)(request) 19 | assert.NoError(t, err) 20 | // tests may be developed and run on different versions of different things 21 | uac := strings.Split(useragent.FromContext(request.Context()), " ") 22 | assert.Contains(t, uac, "auth/PAT") 23 | } 24 | 25 | func TestAuthInUserAgentVisitorCustom(t *testing.T) { 26 | c := &Config{ 27 | AuthType: "PAT", 28 | } 29 | request := &http.Request{} 30 | ctx := context.Background() 31 | ctx = context.WithValue(ctx, AuthType{}, "oath") 32 | request = request.WithContext(ctx) 33 | err := authInUserAgentVisitor(c)(request) 34 | assert.NoError(t, err) 35 | // tests may be developed and run on different versions of different things 36 | uac := strings.Split(useragent.FromContext(request.Context()), " ") 37 | assert.NotContains(t, uac, "auth/PAT") 38 | assert.Contains(t, uac, "auth/oath") 39 | } 40 | -------------------------------------------------------------------------------- /logger/logger_test.go: -------------------------------------------------------------------------------- 1 | package logger 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestGetLogger(t *testing.T) { 11 | t.Cleanup(func() { 12 | DefaultLogger = &SimpleLogger{} 13 | }) 14 | 15 | t1 := &SimpleLogger{} 16 | t2 := &SimpleLogger{} 17 | DefaultLogger = t1 18 | 19 | logger := Get(context.Background()) 20 | assert.Equal(t, logger, t1) 21 | 22 | ctx := NewContext(context.Background(), t2) 23 | logger = Get(ctx) 24 | assert.Equal(t, logger, t2) 25 | } 26 | 27 | func TestWhenInfoLevelThenDebugDisabled(t *testing.T) { 28 | t.Cleanup(func() { 29 | DefaultLogger = &SimpleLogger{} 30 | }) 31 | 32 | infoLevelLogger := &SimpleLogger{ 33 | Level: LevelInfo, 34 | } 35 | debugEnabled := infoLevelLogger.Enabled(context.Background(), LevelDebug) 36 | assert.False(t, debugEnabled) 37 | } 38 | 39 | func TestWhenInfoLevelThenErrorEnabled(t *testing.T) { 40 | infoLevelLogger := &SimpleLogger{ 41 | Level: LevelInfo, 42 | } 43 | 44 | errorEnabled := infoLevelLogger.Enabled(context.Background(), LevelError) 45 | assert.True(t, errorEnabled) 46 | } 47 | 48 | func TestDefaultLevelInfo(t *testing.T) { 49 | logger := &SimpleLogger{} 50 | assert.EqualValues(t, LevelInfo, logger.Level) 51 | } 52 | -------------------------------------------------------------------------------- /internal/tokens_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go" 7 | "github.com/databricks/databricks-sdk-go/service/settings" 8 | "github.com/stretchr/testify/assert" 9 | "github.com/stretchr/testify/require" 10 | ) 11 | 12 | func TestAccTokens(t *testing.T) { 13 | ctx, w := workspaceTest(t) 14 | 15 | token, err := w.Tokens.Create(ctx, settings.CreateTokenRequest{ 16 | Comment: RandomName("go-sdk-"), 17 | LifetimeSeconds: 300, 18 | }) 19 | require.NoError(t, err) 20 | t.Cleanup(func() { 21 | err = w.Tokens.DeleteByTokenId(ctx, token.TokenInfo.TokenId) 22 | require.NoError(t, err) 23 | }) 24 | 25 | all, err := w.Tokens.ListAll(ctx) 26 | require.NoError(t, err) 27 | assert.True(t, len(all) >= 1) 28 | 29 | byName, err := w.Tokens.GetByComment(ctx, token.TokenInfo.Comment) 30 | require.NoError(t, err) 31 | assert.Equal(t, token.TokenInfo.TokenId, byName.TokenId) 32 | 33 | wscInner := databricks.Must(databricks.NewWorkspaceClient(&databricks.Config{ 34 | Host: w.Config.Host, 35 | Token: token.TokenValue, 36 | AuthType: "pat", 37 | })) 38 | 39 | me2, err := wscInner.CurrentUser.Me(ctx) 40 | require.NoError(t, err) 41 | assert.Equal(t, me2.UserName, me(t, w).UserName) 42 | } 43 | -------------------------------------------------------------------------------- /internal/testspecs/service/basicv2/api.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | // Lorem Ipsum 4 | package basicv2 5 | 6 | import ( 7 | "context" 8 | 9 | "github.com/databricks/databricks-sdk-go/client" 10 | "github.com/databricks/databricks-sdk-go/listing" 11 | ) 12 | 13 | type BasicV2Interface interface { 14 | CreateBasic(ctx context.Context, request CreateBasicRequest) (*Basic, error) 15 | 16 | DeleteBasic(ctx context.Context, request DeleteBasicRequest) error 17 | 18 | GetBasic(ctx context.Context, request GetBasicRequest) (*Basic, error) 19 | 20 | // 21 | // This method is generated by Databricks SDK Code Generator. 22 | ListBasics(ctx context.Context, request ListBasicsRequest) listing.Iterator[Basic] 23 | 24 | // 25 | // This method is generated by Databricks SDK Code Generator. 26 | ListBasicsAll(ctx context.Context, request ListBasicsRequest) ([]Basic, error) 27 | 28 | UpdateBasic(ctx context.Context, request UpdateBasicRequest) (*Basic, error) 29 | } 30 | 31 | func NewBasicV2(client *client.DatabricksClient) *BasicV2API { 32 | return &BasicV2API{ 33 | basicV2Impl: basicV2Impl{ 34 | client: client, 35 | }, 36 | } 37 | } 38 | 39 | // Lorem Ipsum 40 | type BasicV2API struct { 41 | basicV2Impl 42 | } 43 | -------------------------------------------------------------------------------- /marshal/types_maptoarray_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestMapToArrayDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | MapToArray: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestMapToArrayValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | MapToArray: map[string][]interface{}{ 24 | "key": {structNoFS{ 25 | Int: 3, 26 | }}, 27 | }, 28 | }, 29 | jsonString: `{"maptoarray":{"key":[{"childint":3}]},"childfs":{},"childnofs":{}}`, 30 | matchClassic: true, 31 | unmarshalResult: &customStruct{ 32 | MapToArray: map[string][]interface{}{ 33 | "key": {map[string]interface{}{ 34 | "childint": 3.0, 35 | }}, 36 | }, 37 | }, 38 | }, 39 | ) 40 | } 41 | 42 | func TestMapToArrayForce(t *testing.T) { 43 | executeBasicMarshalTest(t, 44 | basicMarshalTest{ 45 | st: customStruct{ 46 | MapToArray: nil, 47 | ForceSendFields: []string{"MapToArray"}, 48 | }, 49 | jsonString: `{"childfs":{},"childnofs":{}}`, 50 | matchClassic: true, 51 | }, 52 | ) 53 | } 54 | -------------------------------------------------------------------------------- /qa/lock/internal/lockable.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "reflect" 7 | "strings" 8 | ) 9 | 10 | func GenerateBlobName[R any](r R) string { 11 | rv := reflect.ValueOf(r) 12 | if rv.Kind() == reflect.Ptr { 13 | rv = rv.Elem() 14 | } 15 | switch rv.Kind() { 16 | case reflect.String: 17 | return rv.String() 18 | case reflect.Struct: 19 | id := bytes.Buffer{} 20 | t := rv.Type() 21 | if rv.NumField() == 0 { 22 | return t.Name() 23 | } 24 | id.WriteString(fmt.Sprintf("%s:", t.Name())) 25 | for i := 0; i < rv.NumField(); i++ { 26 | fieldName := t.Field(i).Name 27 | id.WriteString(fmt.Sprintf("%s=", fieldName)) 28 | switch v := rv.Field(i).Interface().(type) { 29 | case string: 30 | noSlashes := strings.ReplaceAll(v, "/", "_") 31 | id.WriteString(noSlashes) 32 | case int, int8, int16, int32, int64: 33 | id.WriteString(fmt.Sprintf("%d", v)) 34 | default: 35 | panic(fmt.Errorf("unsupported type %T for field %s, must be string or int", v, fieldName)) 36 | } 37 | if i < rv.NumField()-1 { 38 | id.WriteString(";") 39 | } 40 | } 41 | return id.String() 42 | default: 43 | panic(fmt.Errorf("unsupported type %T, must be string or struct (or pointer to string or pointer to struct)", r)) 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /marshal/types_slice_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestSliceDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | Slice: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestSliceValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | Slice: []int{3, 4}, 24 | }, 25 | jsonString: `{"slice":[3,4],"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | }, 28 | ) 29 | } 30 | 31 | func TestSliceEmptyValue(t *testing.T) { 32 | executeBasicMarshalTest(t, 33 | basicMarshalTest{ 34 | st: customStruct{ 35 | Slice: []int{}, 36 | }, 37 | jsonString: `{"childfs":{},"childnofs":{}}`, 38 | matchClassic: true, 39 | unmarshalResult: &customStruct{ 40 | Slice: nil, 41 | }, 42 | }, 43 | ) 44 | } 45 | 46 | func TestSliceForce(t *testing.T) { 47 | executeBasicMarshalTest(t, 48 | basicMarshalTest{ 49 | st: customStruct{ 50 | Slice: nil, 51 | ForceSendFields: []string{"Slice"}, 52 | }, 53 | jsonString: `{"childfs":{},"childnofs":{}}`, 54 | matchClassic: true, 55 | }, 56 | ) 57 | } 58 | -------------------------------------------------------------------------------- /internal/testspecs/service/idempotencytesting/impl.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | package idempotencytesting 4 | 5 | import ( 6 | "context" 7 | "net/http" 8 | 9 | "github.com/databricks/databricks-sdk-go/client" 10 | "github.com/google/uuid" 11 | "golang.org/x/exp/slices" 12 | ) 13 | 14 | // unexported type that holds implementations of just IdempotencyTesting API methods 15 | type idempotencyTestingImpl struct { 16 | client *client.DatabricksClient 17 | } 18 | 19 | func (a *idempotencyTestingImpl) CreateTestResource(ctx context.Context, request CreateTestResourceRequest) (*TestResource, error) { 20 | var testResource TestResource 21 | if request.RequestId == "" { 22 | request.RequestId = uuid.New().String() 23 | } 24 | path := "/api/2.0/idempotency-testing/resources" 25 | queryParams := make(map[string]any) 26 | 27 | if request.RequestId != "" || slices.Contains(request.ForceSendFields, "RequestId") { 28 | queryParams["request_id"] = request.RequestId 29 | } 30 | headers := make(map[string]string) 31 | headers["Accept"] = "application/json" 32 | headers["Content-Type"] = "application/json" 33 | err := a.client.Do(ctx, http.MethodPost, path, headers, queryParams, request.TestResource, &testResource) 34 | return &testResource, err 35 | } 36 | -------------------------------------------------------------------------------- /openapi/code/named_sort.go: -------------------------------------------------------------------------------- 1 | package code 2 | 3 | import "sort" 4 | 5 | // github.com/databricks/databricks-sdk-go/httpclient/fixtures stub generator uses Named.PascalName() method to come up with 6 | // the best possible field name for generated copy-pastable stubs, though, when this library is attempted to be used together 7 | // with github.com/spf13/viper, we immediately get the following error related to a change in 8 | // golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e as: 9 | // .../entity.go:185:32: type func(a *Field, b *Field) bool of func(a, b *Field) bool {…} does not match inferred type 10 | // ... func(a *Field, b *Field) int for func(a E, b E) int, 11 | // because github.com/spf13/viper v0.17+ transitively depends on golang.org/x/exp v0.0.0-20230905200255-921286631fa9 12 | 13 | // sorts slice predictably by NamesLikeThis 14 | func pascalNameSort[E interface{ PascalName() string }](things []E) { 15 | sort.Slice(things, func(i, j int) bool { 16 | return things[i].PascalName() < things[j].PascalName() 17 | }) 18 | } 19 | 20 | // sorts slice predictably by package_names_and.ClassNamesLikeThis 21 | func fullNameSort[E interface{ FullName() string }](things []E) { 22 | sort.Slice(things, func(i, j int) bool { 23 | return things[i].FullName() < things[j].FullName() 24 | }) 25 | } 26 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | default: build 2 | 3 | build: vendor 4 | @echo "✓ Building source code with go build ..." 5 | @go build -mod vendor -v 6 | 7 | fmt: 8 | @echo "✓ Formatting source code with goimports ..." 9 | @go run golang.org/x/tools/cmd/goimports@v0.34.0 -w $(shell find . -type f -name '*.go' -not -path "./vendor/*") 10 | @echo "✓ Formatting source code with gofmt ..." 11 | @gofmt -w $(shell find . -type f -name '*.go' -not -path "./vendor/*") 12 | 13 | lint: vendor 14 | @echo "✓ Linting source code with https://staticcheck.io/ ..." 15 | @go run honnef.co/go/tools/cmd/staticcheck@v0.5.1 ./... 16 | 17 | test: vendor 18 | @echo "✓ Running tests ..." 19 | @go run gotest.tools/gotestsum@v1.12.2 --format pkgname-and-test-fails \ 20 | --no-summary=skipped --raw-command go test -v \ 21 | -json -short -coverprofile=coverage.txt ./... 22 | 23 | coverage: test 24 | @echo "✓ Opening coverage for unit tests ..." 25 | @go tool cover -html=coverage.txt 26 | 27 | vendor: 28 | @echo "✓ Filling vendor folder with library code ..." 29 | @go mod vendor 30 | 31 | download: 32 | @echo "✓ Downloading dependencies ..." 33 | @go mod download 34 | 35 | doc: 36 | @echo "Open http://localhost:6060" 37 | @go run golang.org/x/tools/cmd/godoc@v0.34.0 -http=localhost:6060 38 | 39 | .PHONY: fmt vendor fmt coverage test lint doc download 40 | -------------------------------------------------------------------------------- /examples/zerolog/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "os" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | sdk "github.com/databricks/databricks-sdk-go/logger" 10 | "github.com/rs/zerolog" 11 | "github.com/rs/zerolog/log" 12 | ) 13 | 14 | var json bool 15 | 16 | func init() { 17 | flag.BoolVar(&json, "json", false, "log json messages") 18 | } 19 | 20 | func main() { 21 | flag.Parse() 22 | 23 | var l zerolog.Logger 24 | if json { 25 | l = log.Output(os.Stderr) 26 | } else { 27 | l = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) 28 | } 29 | 30 | // Define global logger for the SDK. 31 | // This instance is used when no context is present (e.g. upon construction), 32 | // or when a context is present but doesn't hold a logger to use. 33 | sdk.DefaultLogger = &zerologAdapter{l.With().Bool("global", true).Logger()} 34 | 35 | // Define logger on context for the SDK to use. 36 | ctx := sdk.NewContext(context.Background(), &zerologAdapter{ 37 | l.With().Bool("global", false).Logger(), 38 | }) 39 | 40 | // Construct client and make a request. 41 | // Observe a mix of global=true and global=false in the log output. 42 | w := databricks.Must(databricks.NewWorkspaceClient()) 43 | _, err := w.CurrentUser.Me(ctx) 44 | if err != nil { 45 | panic(err) 46 | } 47 | } 48 | -------------------------------------------------------------------------------- /marshal/types_pchildnofs_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestPChildNoFSDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | PChildNoFS: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestPChildNoFSValueEmpty(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | PChildNoFS: &structNoFS{}, 24 | }, 25 | jsonString: `{"pchildnofs":{},"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | }, 28 | ) 29 | } 30 | 31 | func TestPChildNoFSValue(t *testing.T) { 32 | executeBasicMarshalTest(t, 33 | basicMarshalTest{ 34 | st: customStruct{ 35 | PChildNoFS: &structNoFS{ 36 | Int: 3, 37 | }, 38 | }, 39 | jsonString: `{"pchildnofs":{"childint":3},"childfs":{},"childnofs":{}}`, 40 | matchClassic: true, 41 | }, 42 | ) 43 | } 44 | 45 | func TestPChildNoFSForce(t *testing.T) { 46 | executeBasicMarshalTest(t, 47 | basicMarshalTest{ 48 | st: customStruct{ 49 | PChildNoFS: nil, 50 | ForceSendFields: []string{"PChildNoFS"}, 51 | }, 52 | jsonString: `{"childfs":{},"childnofs":{}}`, 53 | matchClassic: true, 54 | }, 55 | ) 56 | } 57 | -------------------------------------------------------------------------------- /config/auth_default_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | "testing" 7 | ) 8 | 9 | func TestDefaultCredentials_Configure(t *testing.T) { 10 | testCases := []struct { 11 | desc string 12 | authType string 13 | wantErr string 14 | }{ 15 | { 16 | desc: "unknown auth type", 17 | authType: "unknown-auth-type-1337", 18 | wantErr: "auth type \"unknown-auth-type-1337\" not found", 19 | }, 20 | { 21 | desc: "not valid auth", 22 | authType: "", 23 | wantErr: "cannot configure default credentials", 24 | }, 25 | } 26 | 27 | for _, tc := range testCases { 28 | t.Run(tc.desc, func(t *testing.T) { 29 | ctx := context.Background() 30 | cfg := &Config{ 31 | AuthType: tc.authType, 32 | resolved: true, // avoid calling EnsureResolved 33 | } 34 | 35 | dc := DefaultCredentials{} 36 | got, gotErr := dc.Configure(ctx, cfg) 37 | 38 | if got != nil { 39 | t.Errorf("DefaultCredentials.Configure: got %v, want nil", got) 40 | } 41 | if gotErr == nil { 42 | t.Errorf("DefaultCredentials.Configure: got error %v, want non-nil", gotErr) 43 | } 44 | if !strings.Contains(gotErr.Error(), tc.wantErr) { 45 | t.Errorf("DefaultCredentials.Configure: got error %v, want error containing %q", gotErr, tc.wantErr) 46 | } 47 | }) 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /examples/slog/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "os" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | sdk "github.com/databricks/databricks-sdk-go/logger" 10 | "golang.org/x/exp/slog" 11 | ) 12 | 13 | var json bool 14 | 15 | func init() { 16 | flag.BoolVar(&json, "json", false, "log json messages") 17 | } 18 | 19 | func main() { 20 | flag.Parse() 21 | 22 | opts := slog.HandlerOptions{ 23 | Level: slog.LevelDebug, 24 | } 25 | 26 | var handler slog.Handler 27 | if json { 28 | handler = opts.NewJSONHandler(os.Stderr) 29 | } else { 30 | handler = opts.NewTextHandler(os.Stderr) 31 | } 32 | 33 | // Define global logger for the SDK. 34 | // This instance is used when no context is present (e.g. upon construction), 35 | // or when a context is present but doesn't hold a logger to use. 36 | sdk.DefaultLogger = &slogAdapter{ 37 | slog.New(handler).With("global", true), 38 | } 39 | 40 | // Define logger on context for the SDK to use. 41 | ctx := sdk.NewContext(context.Background(), &slogAdapter{ 42 | slog.New(handler).With("global", false), 43 | }) 44 | 45 | // Construct client and make a request. 46 | // Observe a mix of global=true and global=false in the log output. 47 | w := databricks.Must(databricks.NewWorkspaceClient()) 48 | _, err := w.CurrentUser.Me(ctx) 49 | if err != nil { 50 | panic(err) 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /config/token_source_strategy.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/databricks/databricks-sdk-go/config/credentials" 7 | "github.com/databricks/databricks-sdk-go/config/experimental/auth" 8 | ) 9 | 10 | // Creates a CredentialsStrategy from a TokenSource. 11 | func NewTokenSourceStrategy(name string, ts auth.TokenSource) CredentialsStrategy { 12 | return &tokenSourceStrategy{name: name, ts: ts} 13 | } 14 | 15 | // tokenSourceStrategy is wrapper on a auth.TokenSource which converts it into 16 | // a CredentialsStrategy. 17 | type tokenSourceStrategy struct { 18 | name string 19 | ts auth.TokenSource 20 | } 21 | 22 | // Configure implements [CredentialsStrategy.Configure]. 23 | func (tss *tokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { 24 | cp := credentials.NewOAuthCredentialsProviderFromTokenSource(auth.NewCachedTokenSource(tss.ts)) 25 | 26 | // Sanity check that a token can be obtained. 27 | // 28 | // TODO: Move this outside of this function. If credentials providers have 29 | // to be tested, this should be done in the main default loop, not here. 30 | if _, err := cp.Token(ctx); err != nil { 31 | return nil, err 32 | } 33 | 34 | return cp, nil 35 | } 36 | 37 | // Name implements [CredentialsStrategy.Name]. 38 | func (t *tokenSourceStrategy) Name() string { 39 | return t.name 40 | } 41 | -------------------------------------------------------------------------------- /common/environment/azure.go: -------------------------------------------------------------------------------- 1 | package environment 2 | 3 | type azureEnvironment struct { 4 | Name string `json:"name"` 5 | ServiceManagementEndpoint string `json:"serviceManagementEndpoint"` 6 | ResourceManagerEndpoint string `json:"resourceManagerEndpoint"` 7 | ActiveDirectoryEndpoint string `json:"activeDirectoryEndpoint"` 8 | } 9 | 10 | // based on github.com/Azure/go-autorest/autorest/azure/azureEnvironments.go 11 | var ( 12 | AzurePublicCloud = azureEnvironment{ 13 | Name: "PUBLIC", 14 | ServiceManagementEndpoint: "https://management.core.windows.net/", 15 | ResourceManagerEndpoint: "https://management.azure.com/", 16 | ActiveDirectoryEndpoint: "https://login.microsoftonline.com/", 17 | } 18 | 19 | AzureUsGovernmentCloud = azureEnvironment{ 20 | Name: "USGOVERNMENT", 21 | ServiceManagementEndpoint: "https://management.core.usgovcloudapi.net/", 22 | ResourceManagerEndpoint: "https://management.usgovcloudapi.net/", 23 | ActiveDirectoryEndpoint: "https://login.microsoftonline.us/", 24 | } 25 | 26 | AzureChinaCloud = azureEnvironment{ 27 | Name: "CHINA", 28 | ServiceManagementEndpoint: "https://management.core.chinacloudapi.cn/", 29 | ResourceManagerEndpoint: "https://management.chinacloudapi.cn/", 30 | ActiveDirectoryEndpoint: "https://login.chinacloudapi.cn/", 31 | } 32 | ) 33 | -------------------------------------------------------------------------------- /test-config.yaml: -------------------------------------------------------------------------------- 1 | ignored_tests: 2 | go: 3 | - package: github.com/databricks/databricks-sdk-go/internal 4 | test_name: TestAccWorkspaceUsers 5 | comment: Failures due to read-after-write inconsistency. Tracked in https://databricks.atlassian.net/browse/ES-1100061 6 | - package: github.com/databricks/databricks-sdk-go/internal 7 | test_name: TestAccGroups 8 | comment: Failures due to read-after-write inconsistency. Tracked in https://databricks.atlassian.net/browse/ES-1100061 9 | - package: github.com/databricks/databricks-sdk-go/internal 10 | test_name: TestAccServicePrincipalsOnAWS 11 | comment: Failures due to read-after-write inconsistency. Tracked in https://databricks.atlassian.net/browse/ES-1100061 12 | - package: github.com/databricks/databricks-sdk-go/internal 13 | test_name: TestMwsAccBudgets 14 | comment: Oauth issues in 2.1 REST APIs. Tracked in https://databricks.atlassian.net/browse/ES-1197977 15 | - package: github.com/databricks/databricks-sdk-go/internal 16 | test_name: TestUcAccMetastores 17 | comment: Metastores are limited in new accounts. Limit increase in https://databricks.atlassian.net/browse/ES-1242706 18 | - package: github.com/databricks/databricks-sdk-go/internal 19 | test_name: TestAccCreateOboTokenOnAws 20 | comment: Flaky test. Tracked in https://databricks.atlassian.net/browse/ES-1243720 -------------------------------------------------------------------------------- /examples/slog/README.md: -------------------------------------------------------------------------------- 1 | # slog 2 | 3 | This example shows how to use [slog][slog] with this library. 4 | 5 | [slog]: https://pkg.go.dev/golang.org/x/exp/slog 6 | 7 | ## Output 8 | 9 | With JSON output: 10 | 11 | ```text 12 | {"time":"2023-03-10T08:33:11.318122+01:00","level":"DEBUG","msg":"Loading config via environment","global":true} 13 | {"time":"2023-03-10T08:33:11.318411+01:00","level":"DEBUG","msg":"Loading config via config-file","global":true} 14 | {"time":"2023-03-10T08:33:11.318552+01:00","level":"DEBUG","msg":"Loading DEFAULT profile from [...]/.databrickscfg","global":true} 15 | {"time":"2023-03-10T08:33:11.31862+01:00","level":"DEBUG","msg":"Attempting to configure auth: pat","global":false} 16 | {"time":"2023-03-10T08:33:12.170168+01:00","level":"DEBUG","msg":"GET /api/2.0/preview/scim/v2/Me\n< HTTP/2.0 200 OK\n 17 | ``` 18 | 19 | With text output: 20 | ```text 21 | time=2023-03-10T08:33:44.219+01:00 level=DEBUG msg="Loading config via environment" global=true 22 | time=2023-03-10T08:33:44.219+01:00 level=DEBUG msg="Loading config via config-file" global=true 23 | time=2023-03-10T08:33:44.219+01:00 level=DEBUG msg="Loading DEFAULT profile from [...]/.databrickscfg" global=true 24 | time=2023-03-10T08:33:44.220+01:00 level=DEBUG msg="Attempting to configure auth: pat" global=false 25 | time=2023-03-10T08:33:45.077+01:00 level=DEBUG msg="GET /api/2.0/preview/scim/v2/Me\n< HTTP/2.0 200 OK\n 26 | ``` 27 | -------------------------------------------------------------------------------- /service/compute/policy_families_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package compute_test 4 | 5 | import ( 6 | "context" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | "github.com/databricks/databricks-sdk-go/logger" 10 | 11 | "github.com/databricks/databricks-sdk-go/service/compute" 12 | ) 13 | 14 | func ExamplePolicyFamiliesAPI_Get_clusterPolicyFamilies() { 15 | ctx := context.Background() 16 | w, err := databricks.NewWorkspaceClient() 17 | if err != nil { 18 | panic(err) 19 | } 20 | 21 | all, err := w.PolicyFamilies.ListAll(ctx, compute.ListPolicyFamiliesRequest{}) 22 | if err != nil { 23 | panic(err) 24 | } 25 | logger.Infof(ctx, "found %v", all) 26 | 27 | firstFamily, err := w.PolicyFamilies.Get(ctx, compute.GetPolicyFamilyRequest{ 28 | PolicyFamilyId: all[0].PolicyFamilyId, 29 | }) 30 | if err != nil { 31 | panic(err) 32 | } 33 | logger.Infof(ctx, "found %v", firstFamily) 34 | 35 | } 36 | 37 | func ExamplePolicyFamiliesAPI_List_clusterPolicyFamilies() { 38 | ctx := context.Background() 39 | w, err := databricks.NewWorkspaceClient() 40 | if err != nil { 41 | panic(err) 42 | } 43 | 44 | all, err := w.PolicyFamilies.ListAll(ctx, compute.ListPolicyFamiliesRequest{}) 45 | if err != nil { 46 | panic(err) 47 | } 48 | logger.Infof(ctx, "found %v", all) 49 | 50 | } 51 | -------------------------------------------------------------------------------- /.github/PULL_REQUEST_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | ## What changes are proposed in this pull request? 2 | 3 | Provide the readers and reviewers with the information they need to understand 4 | this PR in a comprehensive manner. 5 | 6 | Specifically, try to answer the two following questions: 7 | 8 | - **WHAT** changes are being made in the PR? This should be a summary of the 9 | major changes to allow the reader to quickly understand the PR without having 10 | to look at the code. 11 | - **WHY** are these changes needed? This should provide the context that the 12 | reader might be missing. For example, were there any decisions behind the 13 | change that are not reflected in the code itself? 14 | 15 | The “why part” is the most important of the two as it usually cannot be 16 | inferred from the code itself. A well-written PR description will help future 17 | developers (including your future self) to know how to interact and update your 18 | code. 19 | 20 | ## How is this tested? 21 | 22 | Describe any tests you have done; especially if test tests are not part of 23 | the unit tests (e.g. local tests). 24 | 25 | **ALWAYS ANSWER THIS QUESTION:** Answer with "N/A" if tests are not applicable 26 | to your PR (e.g. if the PR only modifies comments). Do not be afraid of 27 | answering "Not tested" if the PR has not been tested. Being clear about what 28 | has been done and not done provides important context to the reviewers. -------------------------------------------------------------------------------- /service/compute/ext_results_test.go: -------------------------------------------------------------------------------- 1 | package compute 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestResults_Error(t *testing.T) { 10 | cr := Results{} 11 | assert.NoError(t, cr.Err()) 12 | cr.ResultType = "error" 13 | assert.EqualError(t, cr.Err(), "") 14 | 15 | cr.Summary = "NotFoundException: Things are going wrong; nested exception is: with something" 16 | assert.Equal(t, "Things are going wrong with something", cr.Error()) 17 | 18 | cr.Summary = "" 19 | cr.Cause = "ExecutionError: \nStatusCode=400\nStatusDescription=ABC\nSomething else" 20 | assert.Equal(t, "\nStatusCode=400\nStatusDescription=ABC", cr.Error()) 21 | 22 | cr.Cause = "ErrorMessage=Error was here\n" 23 | assert.Equal(t, "Error was here", cr.Error()) 24 | 25 | assert.False(t, cr.Scan()) 26 | } 27 | 28 | func TestResults_Scan(t *testing.T) { 29 | cr := Results{ 30 | ResultType: "table", 31 | Data: []interface{}{ 32 | []interface{}{"foo", 1, true}, 33 | []interface{}{"bar", 2, false}, 34 | }, 35 | } 36 | a := "" 37 | b := 0 38 | c := false 39 | assert.True(t, cr.Scan(&a, &b, &c)) 40 | assert.Equal(t, "foo", a) 41 | assert.Equal(t, 1, b) 42 | assert.Equal(t, true, c) 43 | 44 | assert.True(t, cr.Scan(&a, &b, &c)) 45 | assert.Equal(t, "bar", a) 46 | assert.Equal(t, 2, b) 47 | assert.Equal(t, false, c) 48 | 49 | assert.False(t, cr.Scan(&a, &b, &c)) 50 | } 51 | -------------------------------------------------------------------------------- /account_functions.go: -------------------------------------------------------------------------------- 1 | package databricks 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/provisioning" 7 | ) 8 | 9 | // GetWorkspaceClient returns a WorkspaceClient for the given workspace. The 10 | // workspace can be fetched by calling w.Workspaces.Get() or w.Workspaces.List(). 11 | // 12 | // The config used for the workspace is identical to that used for the account, 13 | // except that the host is set to the workspace host, and the account ID is 14 | // not set. 15 | // 16 | // Example: 17 | // 18 | // a, err := databricks.NewAccountClient() 19 | // if err != nil { 20 | // panic(err) 21 | // } 22 | // ctx := context.Background() 23 | // workspaces, err := a.Workspaces.List(ctx) 24 | // if err != nil { 25 | // panic(err) 26 | // } 27 | // w, err := a.GetWorkspaceClient(workspaces[0]) 28 | // if err != nil { 29 | // panic(err) 30 | // } 31 | // me, err := w.CurrentUser.Me(ctx) 32 | func (c *AccountClient) GetWorkspaceClient(ws provisioning.Workspace) (*WorkspaceClient, error) { 33 | host := c.Config.Environment().DeploymentURL(ws.DeploymentName) 34 | cfg, err := c.Config.NewWithWorkspaceHost(host) 35 | if err != nil { 36 | return nil, err 37 | } 38 | cfg.AzureResourceID = ws.AzureResourceId() 39 | cfg.WorkspaceId = fmt.Sprintf("%d", ws.WorkspaceId) 40 | w, err := NewWorkspaceClient((*Config)(cfg)) 41 | if err != nil { 42 | return nil, err 43 | } 44 | return w, nil 45 | } 46 | -------------------------------------------------------------------------------- /config/azure_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "net/http" 7 | "testing" 8 | 9 | "github.com/stretchr/testify/assert" 10 | ) 11 | 12 | func TestLoadAzureTenantId(t *testing.T) { 13 | c := &Config{ 14 | Host: "https://adb-xyz.c.azuredatabricks.net/", 15 | azureTenantIdFetchClient: makeClient(&http.Response{ 16 | StatusCode: 302, 17 | Header: http.Header{"Location": []string{"https://login.microsoftonline.com/123-abc/oauth2/token"}}, 18 | }), 19 | } 20 | err := c.loadAzureTenantId(context.Background()) 21 | assert.NoError(t, err) 22 | assert.Equal(t, c.AzureTenantID, "123-abc") 23 | } 24 | 25 | func TestLoadAzureTenantId_Failure(t *testing.T) { 26 | testErr := errors.New("Failed to fetch login page") 27 | c := &Config{ 28 | Host: "https://adb-xyz.c.azuredatabricks.net/", 29 | azureTenantIdFetchClient: makeFailingClient(testErr), 30 | } 31 | err := c.loadAzureTenantId(context.Background()) 32 | assert.ErrorIs(t, err, testErr) 33 | } 34 | 35 | func TestLoadAzureTenantId_SkipNotInAzure(t *testing.T) { 36 | testErr := errors.New("Failed to fetch login page") 37 | c := &Config{ 38 | Host: "https://test.cloud.databricks.com/", 39 | azureTenantIdFetchClient: makeFailingClient(testErr), 40 | } 41 | err := c.loadAzureTenantId(context.Background()) 42 | assert.NoError(t, err) 43 | assert.Empty(t, c.AzureTenantID) 44 | } 45 | -------------------------------------------------------------------------------- /httpclient/errors_test.go: -------------------------------------------------------------------------------- 1 | package httpclient 2 | 3 | import ( 4 | "context" 5 | "io" 6 | "net/http" 7 | "strings" 8 | "testing" 9 | 10 | "github.com/stretchr/testify/assert" 11 | "github.com/stretchr/testify/require" 12 | ) 13 | 14 | func TestSimpleRequestAPIError(t *testing.T) { 15 | c := NewApiClient(ClientConfig{ 16 | Transport: hc(func(r *http.Request) (*http.Response, error) { 17 | return &http.Response{ 18 | StatusCode: 400, 19 | Body: io.NopCloser(strings.NewReader(`{ 20 | "error_code": "NOT_FOUND", 21 | "message": "Something was not found" 22 | }`)), 23 | Request: r, 24 | }, nil 25 | }), 26 | }) 27 | err := c.Do(context.Background(), "PATCH", "/a", WithRequestData(map[string]any{})) 28 | var httpErr *HttpError 29 | if assert.ErrorAs(t, err, &httpErr) { 30 | require.Equal(t, 400, httpErr.StatusCode) 31 | } 32 | } 33 | 34 | func TestSimpleRequestErrReaderBody(t *testing.T) { 35 | c := NewApiClient(ClientConfig{ 36 | Transport: hc(func(r *http.Request) (*http.Response, error) { 37 | return &http.Response{ 38 | StatusCode: 200, 39 | Body: errReader(false), 40 | Request: r, 41 | }, nil 42 | }), 43 | }) 44 | headers := map[string]string{"Accept": "application/json"} 45 | err := c.Do(context.Background(), "PATCH", "/a", WithRequestHeaders(headers), WithRequestData(map[string]any{})) 46 | require.EqualError(t, err, "response body: test error") 47 | } 48 | -------------------------------------------------------------------------------- /marshal/types_childfs_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestChildFSDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | ChildFS: structFS{}, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestChildFSValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | ChildFS: structFS{ 24 | Int: 3, 25 | }, 26 | }, 27 | jsonString: `{"childfs":{"childint":3},"childnofs":{}}`, 28 | matchClassic: true, 29 | }, 30 | ) 31 | } 32 | 33 | func TestChildFSForce(t *testing.T) { 34 | executeBasicMarshalTest(t, 35 | basicMarshalTest{ 36 | st: customStruct{ 37 | ChildFS: structFS{}, 38 | ForceSendFields: []string{"ChildFS"}, 39 | }, 40 | jsonString: `{"childfs":{},"childnofs":{}}`, 41 | matchClassic: true, 42 | }, 43 | ) 44 | } 45 | 46 | func TestChildFSChildForce(t *testing.T) { 47 | executeBasicMarshalTest(t, 48 | basicMarshalTest{ 49 | st: customStruct{ 50 | ChildFS: structFS{ 51 | ForceSendFields: []string{"Int"}, 52 | }, 53 | ForceSendFields: []string{"Int"}, 54 | }, 55 | jsonString: `{"childfs":{"childint":0},"childnofs":{}, "int":0}`, 56 | matchClassic: false, 57 | unmarshalForceSendField: []string{"Int"}, 58 | }, 59 | ) 60 | } 61 | -------------------------------------------------------------------------------- /httpclient/fixtures/slice_transport.go: -------------------------------------------------------------------------------- 1 | package fixtures 2 | 3 | import ( 4 | "fmt" 5 | "net/http" 6 | ) 7 | 8 | var Failures = SliceTransport{ 9 | { 10 | MatchAny: true, 11 | ReuseRequest: true, 12 | Status: 418, 13 | Response: `{"error_code": "NONSENSE", "message": "I'm a teapot"}`, 14 | }, 15 | } 16 | 17 | type SliceTransport []HTTPFixture 18 | 19 | func (fixtures SliceTransport) SkipRetryOnIO() bool { 20 | return true 21 | } 22 | 23 | func (fixtures SliceTransport) RoundTrip(req *http.Request) (*http.Response, error) { 24 | resource := resourceFromRequest(req) 25 | for i, f := range fixtures { 26 | if !f.Match(req.Method, resource) { 27 | continue 28 | } 29 | if f.Status == 0 { 30 | f.Status = 200 31 | } 32 | err := f.AssertRequest(req) 33 | if err != nil { 34 | return nil, fmt.Errorf("expected: %w", err) 35 | } 36 | // Reset the request if it is already used 37 | if !f.ReuseRequest { 38 | fixtures[i] = HTTPFixture{} 39 | } 40 | return f.Reply(req) 41 | } 42 | expectedRequest, err := bodyStub(req) 43 | if err != nil { 44 | return nil, fmt.Errorf("marshal: %w", err) 45 | } 46 | // whitespace in this string is very important for unit tests 47 | stub := fmt.Sprintf(`{ 48 | Method: "%s", 49 | Resource: "%s", 50 | %sResponse: XXX { 51 | // fill in specific fields... 52 | }, 53 | },`, req.Method, resource, expectedRequest) 54 | return nil, fmt.Errorf("missing stub, please add: %s", stub) 55 | } 56 | -------------------------------------------------------------------------------- /service/tags/interface.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | package tags 4 | 5 | import ( 6 | "context" 7 | ) 8 | 9 | // The Tag Policy API allows you to manage policies for governed tags in 10 | // Databricks. Permissions for tag policies can be managed using the [Account 11 | // Access Control Proxy API]. 12 | // 13 | // Deprecated: Do not use this interface, it will be removed in a future version of the SDK. 14 | // 15 | // [Account Access Control Proxy API]: https://docs.databricks.com/api/workspace/accountaccesscontrolproxy 16 | type TagPoliciesService interface { 17 | 18 | // Creates a new tag policy, making the associated tag key governed. 19 | CreateTagPolicy(ctx context.Context, request CreateTagPolicyRequest) (*TagPolicy, error) 20 | 21 | // Deletes a tag policy by its associated governed tag's key, leaving that 22 | // tag key ungoverned. 23 | DeleteTagPolicy(ctx context.Context, request DeleteTagPolicyRequest) error 24 | 25 | // Gets a single tag policy by its associated governed tag's key. 26 | GetTagPolicy(ctx context.Context, request GetTagPolicyRequest) (*TagPolicy, error) 27 | 28 | // Lists the tag policies for all governed tags in the account. 29 | ListTagPolicies(ctx context.Context, request ListTagPoliciesRequest) (*ListTagPoliciesResponse, error) 30 | 31 | // Updates an existing tag policy for a single governed tag. 32 | UpdateTagPolicy(ctx context.Context, request UpdateTagPolicyRequest) (*TagPolicy, error) 33 | } 34 | -------------------------------------------------------------------------------- /httpclient/traceparent/traceparent_test.go: -------------------------------------------------------------------------------- 1 | package traceparent 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestNew(t *testing.T) { 10 | tp := NewTraceparent() 11 | assert.Equal(t, byte(0), tp.version) 12 | assert.Equal(t, byte(1), byte(tp.flags)) 13 | } 14 | 15 | func TestEqual(t *testing.T) { 16 | tp1 := NewTraceparent() 17 | tp2 := &Traceparent{ 18 | version: tp1.version, 19 | traceId: tp1.traceId, 20 | parentId: tp1.parentId, 21 | flags: tp1.flags, 22 | } 23 | assert.True(t, tp1.Equals(tp2)) 24 | } 25 | 26 | func TestTwoNewTraceparentsAreNotEqual(t *testing.T) { 27 | tp1 := NewTraceparent() 28 | tp2 := NewTraceparent() 29 | assert.False(t, tp1.Equals(tp2)) 30 | } 31 | 32 | var testTraceId = [16]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15} 33 | var testParentId = [8]byte{0, 1, 2, 3, 4, 5, 6, 7} 34 | 35 | func TestString(t *testing.T) { 36 | tp := &Traceparent{ 37 | version: 0, 38 | traceId: testTraceId, 39 | parentId: testParentId, 40 | flags: 1, 41 | } 42 | res := tp.String() 43 | assert.Equal(t, "00-000102030405060708090a0b0c0d0e0f-0001020304050607-01", res) 44 | } 45 | 46 | func TestFromString(t *testing.T) { 47 | tp, err := FromString("00-000102030405060708090a0b0c0d0e0f-0001020304050607-01") 48 | assert.NoError(t, err) 49 | assert.Equal(t, byte(0), tp.version) 50 | assert.Equal(t, testTraceId, tp.traceId) 51 | assert.Equal(t, testParentId, tp.parentId) 52 | assert.Equal(t, byte(1), byte(tp.flags)) 53 | } 54 | -------------------------------------------------------------------------------- /service/agentbricks/api.go: -------------------------------------------------------------------------------- 1 | // Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. 2 | 3 | // The Custom LLMs service manages state and powers the UI for the Custom LLM 4 | // product. 5 | package agentbricks 6 | 7 | import ( 8 | "context" 9 | 10 | "github.com/databricks/databricks-sdk-go/client" 11 | ) 12 | 13 | type AgentBricksInterface interface { 14 | 15 | // Cancel a Custom LLM Optimization Run. 16 | CancelOptimize(ctx context.Context, request CancelCustomLlmOptimizationRunRequest) error 17 | 18 | // Create a Custom LLM. 19 | CreateCustomLlm(ctx context.Context, request CreateCustomLlmRequest) (*CustomLlm, error) 20 | 21 | // Delete a Custom LLM. 22 | DeleteCustomLlm(ctx context.Context, request DeleteCustomLlmRequest) error 23 | 24 | // Get a Custom LLM. 25 | GetCustomLlm(ctx context.Context, request GetCustomLlmRequest) (*CustomLlm, error) 26 | 27 | // Start a Custom LLM Optimization Run. 28 | StartOptimize(ctx context.Context, request StartCustomLlmOptimizationRunRequest) (*CustomLlm, error) 29 | 30 | // Update a Custom LLM. 31 | UpdateCustomLlm(ctx context.Context, request UpdateCustomLlmRequest) (*CustomLlm, error) 32 | } 33 | 34 | func NewAgentBricks(client *client.DatabricksClient) *AgentBricksAPI { 35 | return &AgentBricksAPI{ 36 | agentBricksImpl: agentBricksImpl{ 37 | client: client, 38 | }, 39 | } 40 | } 41 | 42 | // The Custom LLMs service manages state and powers the UI for the Custom LLM 43 | // product. 44 | type AgentBricksAPI struct { 45 | agentBricksImpl 46 | } 47 | -------------------------------------------------------------------------------- /examples/zerolog/go.mod: -------------------------------------------------------------------------------- 1 | module main 2 | 3 | go 1.18 4 | 5 | replace github.com/databricks/databricks-sdk-go v0.0.0 => ../.. 6 | 7 | require ( 8 | github.com/databricks/databricks-sdk-go v0.24.0 9 | github.com/rs/zerolog v1.29.0 10 | ) 11 | 12 | require ( 13 | cloud.google.com/go/compute v1.23.0 // indirect 14 | cloud.google.com/go/compute/metadata v0.2.3 // indirect 15 | github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect 16 | github.com/golang/protobuf v1.5.3 // indirect 17 | github.com/google/go-querystring v1.1.0 // indirect 18 | github.com/google/s2a-go v0.1.7 // indirect 19 | github.com/googleapis/enterprise-certificate-proxy v0.3.1 // indirect 20 | github.com/mattn/go-colorable v0.1.13 // indirect 21 | github.com/mattn/go-isatty v0.0.17 // indirect 22 | go.opencensus.io v0.24.0 // indirect 23 | golang.org/x/crypto v0.21.0 // indirect 24 | golang.org/x/exp v0.0.0-20230307190834-24139beb5833 // indirect 25 | golang.org/x/mod v0.13.0 // indirect 26 | golang.org/x/net v0.23.0 // indirect 27 | golang.org/x/oauth2 v0.13.0 // indirect 28 | golang.org/x/sys v0.18.0 // indirect 29 | golang.org/x/text v0.14.0 // indirect 30 | golang.org/x/time v0.3.0 // indirect 31 | google.golang.org/api v0.146.0 // indirect 32 | google.golang.org/appengine v1.6.7 // indirect 33 | google.golang.org/genproto/googleapis/rpc v0.0.0-20230920204549-e6e6cdab5c13 // indirect 34 | google.golang.org/grpc v1.58.3 // indirect 35 | google.golang.org/protobuf v1.33.0 // indirect 36 | gopkg.in/ini.v1 v1.67.0 // indirect 37 | ) 38 | -------------------------------------------------------------------------------- /openapi/extensions.go: -------------------------------------------------------------------------------- 1 | package openapi 2 | 3 | // Pagination is the Databricks OpenAPI Extension for retrieving 4 | // lists of entities through multiple API calls 5 | type Pagination struct { 6 | Offset string `json:"offset,omitempty"` 7 | Limit string `json:"limit,omitempty"` 8 | Results string `json:"results,omitempty"` 9 | Increment int `json:"increment,omitempty"` 10 | Inline bool `json:"inline,omitempty"` 11 | Token *Binding `json:"token,omitempty"` 12 | } 13 | 14 | // Wait is the Databricks OpenAPI Extension for long-running result polling 15 | type Wait struct { 16 | Poll string `json:"poll"` 17 | Bind string `json:"bind"` 18 | BindResponse string `json:"bindResponse,omitempty"` 19 | Binding map[string]Binding `json:"binding,omitempty"` 20 | Field []string `json:"field"` 21 | Message []string `json:"message"` 22 | Success []string `json:"success"` 23 | Failure []string `json:"failure"` 24 | Timeout int `json:"timeout,omitempty"` 25 | } 26 | 27 | // Binding is a relationship between request and/or response 28 | type Binding struct { 29 | Request string `json:"request,omitempty"` 30 | Response string `json:"response,omitempty"` 31 | } 32 | 33 | // DataPlane is the Databricks OpenAPI Extension for direct access to DataPlane APIs 34 | type DataPlane struct { 35 | ConfigMethod string `json:"configMethod"` 36 | Fields []string `json:"field"` 37 | } 38 | -------------------------------------------------------------------------------- /examples/zerolog/zerolog.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/databricks/databricks-sdk-go/logger" 7 | "github.com/rs/zerolog" 8 | ) 9 | 10 | // zerologAdapter makes an zerolog.Logger usable with the Databricks SDK. 11 | type zerologAdapter struct { 12 | zerolog.Logger 13 | } 14 | 15 | func (z *zerologAdapter) Enabled(_ context.Context, level logger.Level) bool { 16 | switch level { 17 | case logger.LevelTrace: 18 | return z.Logger.GetLevel() <= zerolog.TraceLevel 19 | case logger.LevelDebug: 20 | return z.Logger.GetLevel() <= zerolog.DebugLevel 21 | case logger.LevelInfo: 22 | return z.Logger.GetLevel() <= zerolog.InfoLevel 23 | case logger.LevelWarn: 24 | return z.Logger.GetLevel() <= zerolog.WarnLevel 25 | case logger.LevelError: 26 | return z.Logger.GetLevel() <= zerolog.ErrorLevel 27 | default: 28 | return true 29 | } 30 | } 31 | 32 | func (s *zerologAdapter) Tracef(_ context.Context, format string, v ...any) { 33 | s.Logger.Trace().Msgf(format, v...) 34 | } 35 | 36 | func (s *zerologAdapter) Debugf(_ context.Context, format string, v ...any) { 37 | s.Logger.Debug().Msgf(format, v...) 38 | } 39 | 40 | func (s *zerologAdapter) Infof(_ context.Context, format string, v ...any) { 41 | s.Logger.Info().Msgf(format, v...) 42 | } 43 | 44 | func (s *zerologAdapter) Warnf(_ context.Context, format string, v ...any) { 45 | s.Logger.Warn().Msgf(format, v...) 46 | } 47 | 48 | func (s *zerologAdapter) Errorf(_ context.Context, format string, v ...any) { 49 | s.Logger.Error().Msgf(format, v...) 50 | } 51 | -------------------------------------------------------------------------------- /internal/globalinitscripts_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "encoding/base64" 5 | "testing" 6 | 7 | "github.com/databricks/databricks-sdk-go/service/compute" 8 | "github.com/stretchr/testify/assert" 9 | "github.com/stretchr/testify/require" 10 | ) 11 | 12 | func TestAccGlobalInitScripts(t *testing.T) { 13 | ctx, w := workspaceTest(t) 14 | 15 | created, err := w.GlobalInitScripts.Create(ctx, compute.GlobalInitScriptCreateRequest{ 16 | Name: RandomName("go-sdk-"), 17 | Script: base64.StdEncoding.EncodeToString([]byte("echo 1")), 18 | Enabled: true, 19 | Position: 10, 20 | }) 21 | require.NoError(t, err) 22 | 23 | defer w.GlobalInitScripts.DeleteByScriptId(ctx, created.ScriptId) 24 | 25 | err = w.GlobalInitScripts.Update(ctx, compute.GlobalInitScriptUpdateRequest{ 26 | ScriptId: created.ScriptId, 27 | Name: RandomName("go-sdk-updated-"), 28 | Script: base64.StdEncoding.EncodeToString([]byte("echo 2")), 29 | }) 30 | require.NoError(t, err) 31 | 32 | byId, err := w.GlobalInitScripts.GetByScriptId(ctx, created.ScriptId) 33 | require.NoError(t, err) 34 | 35 | all, err := w.GlobalInitScripts.ListAll(ctx) 36 | require.NoError(t, err) 37 | 38 | names, err := w.GlobalInitScripts.GlobalInitScriptDetailsNameToScriptIdMap(ctx) 39 | require.NoError(t, err) 40 | assert.Equal(t, len(all), len(names)) 41 | assert.Equal(t, byId.ScriptId, names[byId.Name]) 42 | 43 | byName, err := w.GlobalInitScripts.GetByName(ctx, byId.Name) 44 | require.NoError(t, err) 45 | assert.Equal(t, byName.ScriptId, byId.ScriptId) 46 | } 47 | -------------------------------------------------------------------------------- /examples/custom-auth/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bufio" 5 | "context" 6 | "fmt" 7 | "net/http" 8 | "os" 9 | "strings" 10 | 11 | "github.com/databricks/databricks-sdk-go" 12 | "github.com/databricks/databricks-sdk-go/config" 13 | "github.com/databricks/databricks-sdk-go/config/credentials" 14 | "github.com/databricks/databricks-sdk-go/service/compute" 15 | ) 16 | 17 | type CustomCredentials struct{} 18 | 19 | func (c *CustomCredentials) Name() string { 20 | return "custom" 21 | } 22 | 23 | func (c *CustomCredentials) Configure( 24 | ctx context.Context, cfg *config.Config, 25 | ) (credentials.CredentialsProvider, error) { 26 | visitor := func(r *http.Request) error { 27 | token := askFor("Token:") 28 | r.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token)) 29 | return nil 30 | } 31 | return credentials.CredentialsProviderFn(visitor), nil 32 | } 33 | 34 | func main() { 35 | w := databricks.Must(databricks.NewWorkspaceClient(&databricks.Config{ 36 | Host: askFor("Host:"), 37 | Credentials: &CustomCredentials{}, 38 | })) 39 | all, err := w.Clusters.ListAll(context.Background(), compute.ListClustersRequest{}) 40 | if err != nil { 41 | panic(err) 42 | } 43 | for _, c := range all { 44 | println(c.ClusterName) 45 | } 46 | } 47 | 48 | func askFor(prompt string) string { 49 | var s string 50 | r := bufio.NewReader(os.Stdin) 51 | for { 52 | fmt.Fprint(os.Stdout, prompt+" ") 53 | s, _ = r.ReadString('\n') 54 | s = strings.TrimSpace(s) 55 | if s != "" { 56 | break 57 | } 58 | } 59 | return s 60 | } 61 | -------------------------------------------------------------------------------- /internal/tokenmanagement_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/iam" 7 | "github.com/databricks/databricks-sdk-go/service/settings" 8 | "github.com/stretchr/testify/assert" 9 | "github.com/stretchr/testify/require" 10 | ) 11 | 12 | func TestAccCreateOboTokenOnAws(t *testing.T) { 13 | ctx, w := workspaceTest(t) 14 | if !w.Config.IsAws() { 15 | t.Skip("works only on aws") 16 | } 17 | groups, err := w.Groups.GroupDisplayNameToIdMap(ctx, iam.ListGroupsRequest{}) 18 | require.NoError(t, err) 19 | 20 | spn, err := w.ServicePrincipals.Create(ctx, iam.ServicePrincipal{ 21 | DisplayName: RandomName(t.Name()), 22 | Groups: []iam.ComplexValue{ 23 | { 24 | Value: groups["admins"], 25 | }, 26 | }, 27 | }) 28 | require.NoError(t, err) 29 | t.Cleanup(func() { 30 | err = w.ServicePrincipals.DeleteById(ctx, spn.Id) 31 | require.NoError(t, err) 32 | }) 33 | 34 | obo, err := w.TokenManagement.CreateOboToken(ctx, settings.CreateOboTokenRequest{ 35 | ApplicationId: spn.ApplicationId, 36 | LifetimeSeconds: 60, 37 | }) 38 | require.NoError(t, err) 39 | t.Cleanup(func() { 40 | err = w.TokenManagement.DeleteByTokenId(ctx, obo.TokenInfo.TokenId) 41 | require.NoError(t, err) 42 | }) 43 | 44 | byId, err := w.TokenManagement.GetByTokenId(ctx, obo.TokenInfo.TokenId) 45 | require.NoError(t, err) 46 | t.Log(byId) 47 | 48 | all, err := w.TokenManagement.ListAll(ctx, settings.ListTokenManagementRequest{}) 49 | require.NoError(t, err) 50 | assert.True(t, len(all) >= 1) 51 | } 52 | -------------------------------------------------------------------------------- /config/testdata/.databrickscfg: -------------------------------------------------------------------------------- 1 | [DEFAULT] 2 | host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/ 3 | token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ 4 | rate_limit=34 5 | debug_truncate=829374 6 | 7 | [nohost] 8 | token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ 9 | 10 | [nohost2] 11 | token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ 12 | rate_limit=34 13 | 14 | [azure-justhost] 15 | host = https://adb-123.4.azuredatabricks.net/ 16 | 17 | [azure-pat] 18 | host = https://adb-123.4.azuredatabricks.net/ 19 | token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ 20 | 21 | [notoken] 22 | host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/ 23 | 24 | [basic] 25 | host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/ 26 | username = abc 27 | password = bcd 28 | 29 | [azure-spn] 30 | host = https://adb-123.4.azuredatabricks.net/ 31 | azure_tenant_id = 12345678 32 | azure_client_id = 87654321 33 | azure_client_secret = 10000002 34 | 35 | [gcp-example] 36 | host = https://123.4.gcp.databricks.com/ 37 | google_credentials = paw48590aw8e09t8apu 38 | 39 | [pat.with.dot] 40 | host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/ 41 | token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ 42 | 43 | [password-with-double-quotes] 44 | host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/ 45 | username = abc 46 | password = "%Y#X$Z" # comment 47 | 48 | [password-with-single-quotes] 49 | host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/ 50 | username = abc 51 | password = '%Y#X$Z' # comment 52 | 53 | [password-without-quotes] 54 | host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/ 55 | username = abc 56 | password = %Y#X$Z # comment 57 | -------------------------------------------------------------------------------- /config/experimental/auth/credentials/pat_test.go: -------------------------------------------------------------------------------- 1 | package credentials 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "testing" 7 | 8 | "github.com/google/go-cmp/cmp" 9 | ) 10 | 11 | func TestNewPATCredentials(t *testing.T) { 12 | testCases := []struct { 13 | desc string 14 | token string 15 | wantErr error 16 | wantHeaders map[string]string 17 | }{ 18 | { 19 | desc: "valid token", 20 | token: "dapi1234567890abcdef", 21 | wantHeaders: map[string]string{ 22 | "Authorization": "Bearer dapi1234567890abcdef", 23 | }, 24 | }, 25 | { 26 | desc: "empty token", 27 | token: "", 28 | wantErr: errTokenRequired, 29 | }, 30 | } 31 | 32 | for _, tc := range testCases { 33 | t.Run(tc.desc, func(t *testing.T) { 34 | gotCreds, gotErr := NewPATCredentials(tc.token) 35 | 36 | if tc.wantErr != nil { 37 | if !errors.Is(gotErr, tc.wantErr) { 38 | t.Fatalf("NewPATCredentials() want error %v, got error: %v", tc.wantErr, gotErr) 39 | } 40 | return // no need to test the rest 41 | } 42 | if gotErr != nil { 43 | t.Fatalf("NewPATCredentials() want no error, got error: %v", gotErr) 44 | } 45 | if gotCreds == nil { 46 | t.Fatalf("NewPATCredentials() returned nil credentials") 47 | } 48 | 49 | gotHeaders, err := gotCreds.AuthHeaders(context.Background()) 50 | if err != nil { 51 | t.Fatalf("AuthHeaders() error = %v", err) 52 | } 53 | 54 | if !cmp.Equal(gotHeaders, tc.wantHeaders) { 55 | t.Errorf("AuthHeaders() = %v, want %v", gotHeaders, tc.wantHeaders) 56 | } 57 | }) 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /internal/ipaccesslists_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/settings" 7 | "github.com/stretchr/testify/assert" 8 | "github.com/stretchr/testify/require" 9 | ) 10 | 11 | func TestAccIpAccessLists(t *testing.T) { 12 | ctx, w := workspaceTest(t) 13 | 14 | created, err := w.IpAccessLists.Create(ctx, settings.CreateIpAccessList{ 15 | Label: RandomName("go-sdk-"), 16 | IpAddresses: []string{"1.0.0.0/16"}, 17 | ListType: settings.ListTypeBlock, 18 | }) 19 | require.NoError(t, err) 20 | 21 | defer w.IpAccessLists.DeleteByIpAccessListId(ctx, created.IpAccessList.ListId) 22 | 23 | err = w.IpAccessLists.Replace(ctx, settings.ReplaceIpAccessList{ 24 | IpAccessListId: created.IpAccessList.ListId, 25 | Label: RandomName("go-sdk-updated-"), 26 | IpAddresses: []string{"1.0.0.0/24"}, 27 | ListType: settings.ListTypeBlock, 28 | Enabled: false, 29 | }) 30 | require.NoError(t, err) 31 | 32 | byId, err := w.IpAccessLists.GetByIpAccessListId(ctx, created.IpAccessList.ListId) 33 | require.NoError(t, err) 34 | 35 | all, err := w.IpAccessLists.ListAll(ctx) 36 | require.NoError(t, err) 37 | 38 | names, err := w.IpAccessLists.IpAccessListInfoLabelToListIdMap(ctx) 39 | require.NoError(t, err) 40 | assert.Equal(t, len(all), len(names)) 41 | assert.Equal(t, byId.IpAccessList.ListId, names[byId.IpAccessList.Label]) 42 | 43 | byName, err := w.IpAccessLists.GetByLabel(ctx, byId.IpAccessList.Label) 44 | require.NoError(t, err) 45 | assert.Equal(t, byId.IpAccessList.ListId, byName.ListId) 46 | } 47 | -------------------------------------------------------------------------------- /config/environments.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "strings" 5 | 6 | "github.com/databricks/databricks-sdk-go/common/environment" 7 | ) 8 | 9 | func (c *Config) Environment() environment.DatabricksEnvironment { 10 | // Use the provided environment if specified. Tests may configure the client with different hostnames, 11 | // like localhost, which are not resolvable to a known environment, while needing to mock a specific environment. 12 | if c.DatabricksEnvironment != nil { 13 | return *c.DatabricksEnvironment 14 | } 15 | envs := environment.AllEnvironments() 16 | if c.Host == "" && c.AzureResourceID != "" { 17 | // azure resource ID can also be used in lieu of host by some 18 | // of the clients, like Terraform. However, in this case, the workspace 19 | // is assumed to be a production workspace. 20 | azureEnv := strings.ToUpper(c.AzureEnvironment) 21 | if azureEnv == "" { 22 | azureEnv = "PUBLIC" 23 | } 24 | for _, v := range envs { 25 | if v.Cloud != environment.CloudAzure { 26 | continue 27 | } 28 | if v.AzureEnvironment.Name != azureEnv { 29 | continue 30 | } 31 | if strings.HasPrefix(v.DnsZone, ".dev") || strings.HasPrefix(v.DnsZone, ".staging") { 32 | // we can't support host-less Azure CLI auth for dev & staging environments, as users will get errors like 33 | // ... `The user or administrator has not consented to use the application with ID '...' named 34 | // 'Microsoft Azure CLI'.`. 35 | continue 36 | } 37 | return v 38 | } 39 | } 40 | hostname := c.CanonicalHostName() 41 | return environment.GetEnvironmentForHostname(hostname) 42 | } 43 | -------------------------------------------------------------------------------- /service/compute/example_test.go: -------------------------------------------------------------------------------- 1 | package compute_test 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | 7 | "github.com/databricks/databricks-sdk-go" 8 | "github.com/databricks/databricks-sdk-go/service/compute" 9 | ) 10 | 11 | func ExampleClustersAPI() { 12 | w := databricks.Must(databricks.NewWorkspaceClient()) 13 | ctx := context.Background() 14 | 15 | // Fetch list of spark runtime versions 16 | sparkVersions, err := w.Clusters.SparkVersions(ctx) 17 | if err != nil { 18 | panic(err) 19 | } 20 | 21 | // Select the latest LTS version 22 | latestLTS, err := sparkVersions.Select(compute.SparkVersionRequest{ 23 | Latest: true, 24 | LongTermSupport: true, 25 | }) 26 | if err != nil { 27 | panic(err) 28 | } 29 | 30 | // Fetch list of available node types 31 | nodeTypes, err := w.Clusters.ListNodeTypes(ctx) 32 | if err != nil { 33 | panic(err) 34 | } 35 | 36 | // Select the smallest node type id 37 | smallestWithDisk, err := nodeTypes.Smallest(compute.NodeTypeRequest{ 38 | LocalDisk: true, 39 | }) 40 | if err != nil { 41 | panic(err) 42 | } 43 | 44 | // Create cluster and wait until it's ready to use 45 | runningCluster, err := w.Clusters.CreateAndWait(ctx, compute.CreateCluster{ 46 | ClusterName: "Test cluster from SDK", 47 | SparkVersion: latestLTS, 48 | NodeTypeId: smallestWithDisk, 49 | AutoterminationMinutes: 15, 50 | NumWorkers: 1, 51 | }) 52 | if err != nil { 53 | panic(err) 54 | } 55 | 56 | fmt.Printf("Cluster is ready: %s#setting/clusters/%s/configuration\n", 57 | w.Config.Host, runningCluster.ClusterId) 58 | } 59 | -------------------------------------------------------------------------------- /internal/instancepools_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/compute" 7 | "github.com/stretchr/testify/assert" 8 | "github.com/stretchr/testify/require" 9 | ) 10 | 11 | func TestAccInstancePools(t *testing.T) { 12 | ctx, w := workspaceTest(t) 13 | 14 | smallest, err := w.Clusters.SelectNodeType(ctx, compute.NodeTypeRequest{ 15 | LocalDisk: true, 16 | }) 17 | require.NoError(t, err) 18 | 19 | created, err := w.InstancePools.Create(ctx, compute.CreateInstancePool{ 20 | InstancePoolName: RandomName("go-sdk-"), 21 | NodeTypeId: smallest, 22 | }) 23 | require.NoError(t, err) 24 | 25 | defer w.InstancePools.DeleteByInstancePoolId(ctx, created.InstancePoolId) 26 | 27 | err = w.InstancePools.Edit(ctx, compute.EditInstancePool{ 28 | InstancePoolId: created.InstancePoolId, 29 | InstancePoolName: RandomName("go-sdk-updated"), 30 | NodeTypeId: smallest, 31 | }) 32 | require.NoError(t, err) 33 | 34 | byId, err := w.InstancePools.GetByInstancePoolId(ctx, created.InstancePoolId) 35 | require.NoError(t, err) 36 | 37 | all, err := w.InstancePools.ListAll(ctx) 38 | require.NoError(t, err) 39 | 40 | names, err := w.InstancePools.InstancePoolAndStatsInstancePoolNameToInstancePoolIdMap(ctx) 41 | require.NoError(t, err) 42 | assert.Equal(t, len(all), len(names)) 43 | assert.Equal(t, byId.InstancePoolId, names[byId.InstancePoolName]) 44 | 45 | byName, err := w.InstancePools.GetByInstancePoolName(ctx, byId.InstancePoolName) 46 | require.NoError(t, err) 47 | assert.Equal(t, byName.InstancePoolId, byId.InstancePoolId) 48 | } 49 | -------------------------------------------------------------------------------- /qa/lock/internal/lockable_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestGenerateBlobName(t *testing.T) { 8 | type testStruct struct { 9 | A string 10 | B int 11 | } 12 | 13 | cases := []struct { 14 | name string 15 | in interface{} 16 | want string 17 | }{ 18 | { 19 | name: "string", 20 | in: "hello", 21 | want: "hello", 22 | }, 23 | { 24 | name: "struct", 25 | in: testStruct{A: "hello", B: 1}, 26 | want: "testStruct:A=hello;B=1", 27 | }, 28 | { 29 | name: "struct pointer", 30 | in: &testStruct{A: "hello", B: 1}, 31 | want: "testStruct:A=hello;B=1", 32 | }, 33 | } 34 | for _, c := range cases { 35 | t.Run(c.name, func(t *testing.T) { 36 | got := GenerateBlobName(c.in) 37 | if got != c.want { 38 | t.Errorf("generateBlobName(%v) == %s, want %s", c.in, got, c.want) 39 | } 40 | }) 41 | } 42 | } 43 | 44 | func TestGenerateBlobName_Panic(t *testing.T) { 45 | type testStruct struct { 46 | C chan int 47 | } 48 | cases := []struct { 49 | name string 50 | in interface{} 51 | }{ 52 | { 53 | name: "unsupported type", 54 | in: 1, 55 | }, 56 | { 57 | name: "unsupported type pointer", 58 | in: &[]int{1}, 59 | }, 60 | { 61 | name: "unsupported type struct", 62 | in: testStruct{C: make(chan int)}, 63 | }, 64 | } 65 | for _, c := range cases { 66 | t.Run(c.name, func(t *testing.T) { 67 | defer func() { 68 | if r := recover(); r == nil { 69 | t.Errorf("generateBlobName(%v) did not panic", c.in) 70 | } 71 | }() 72 | GenerateBlobName(c.in) 73 | }) 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /marshal/types_maptocustom_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestMapToCustomTypeFSDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | MapToCustomTypeFS: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestMapToCustomTypeFSValue(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | MapToCustomTypeFS: map[string]structFS{ 24 | "key": { 25 | Int: 2, 26 | ForceSendFields: []string{"Int"}, 27 | }, 28 | }, 29 | }, 30 | jsonString: `{"maptocustomtypefs":{"key":{"childint":2}},"childfs":{},"childnofs":{}}`, 31 | matchClassic: true, 32 | }, 33 | ) 34 | } 35 | 36 | func TestMapToCustomTypeFSValueForce(t *testing.T) { 37 | executeBasicMarshalTest(t, 38 | basicMarshalTest{ 39 | st: customStruct{ 40 | MapToCustomTypeFS: map[string]structFS{ 41 | "key": { 42 | Int: 0, 43 | ForceSendFields: []string{"Int"}, 44 | }, 45 | }, 46 | }, 47 | jsonString: `{"maptocustomtypefs":{"key":{"childint":0}},"childfs":{},"childnofs":{}}`, 48 | matchClassic: true, 49 | }, 50 | ) 51 | } 52 | 53 | func TestMapToCustomTypeFSForce(t *testing.T) { 54 | executeBasicMarshalTest(t, 55 | basicMarshalTest{ 56 | st: customStruct{ 57 | MapToCustomTypeFS: nil, 58 | ForceSendFields: []string{"MapToCustomTypeFS"}, 59 | }, 60 | jsonString: `{"childfs":{},"childnofs":{}}`, 61 | matchClassic: true, 62 | }, 63 | ) 64 | } 65 | -------------------------------------------------------------------------------- /.github/workflows/tagging.yml: -------------------------------------------------------------------------------- 1 | # Generated file. DO NOT EDIT. 2 | name: tagging 3 | 4 | on: 5 | workflow_dispatch: 6 | # Enable for automatic tagging 7 | #schedule: 8 | # - cron: '0 0 * * TUE' 9 | 10 | # Ensure that only a single instance of the workflow is running at a time. 11 | concurrency: 12 | group: "tagging" 13 | 14 | 15 | jobs: 16 | tag: 17 | environment: "release-is" 18 | runs-on: 19 | group: databricks-deco-testing-runner-group 20 | labels: ubuntu-latest-deco 21 | steps: 22 | - name: Generate GitHub App Token 23 | id: generate-token 24 | uses: actions/create-github-app-token@v2 25 | with: 26 | app-id: ${{ secrets.DECO_SDK_TAGGING_APP_ID }} 27 | private-key: ${{ secrets.DECO_SDK_TAGGING_PRIVATE_KEY }} 28 | 29 | - name: Checkout repository 30 | uses: actions/checkout@v4 31 | with: 32 | fetch-depth: 0 33 | token: ${{ steps.generate-token.outputs.token }} 34 | 35 | #NOTE: email must be the GitHub App email or the commit will not be verified. 36 | - name: Set up Git configuration 37 | run: | 38 | git config user.name "Databricks SDK Release Bot" 39 | git config user.email "DECO-SDK-Tagging[bot]@users.noreply.github.com" 40 | 41 | - name: Install dependencies 42 | run: | 43 | python -m pip install --upgrade pip 44 | pip install PyGithub 45 | 46 | - name: Run script 47 | env: 48 | GITHUB_TOKEN: ${{ steps.generate-token.outputs.token }} 49 | GITHUB_REPOSITORY: ${{ github.repository }} 50 | run: | 51 | python tagging.py 52 | -------------------------------------------------------------------------------- /internal/warehouses_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/sql" 7 | "github.com/stretchr/testify/assert" 8 | "github.com/stretchr/testify/require" 9 | ) 10 | 11 | func TestAccSqlWarehouses(t *testing.T) { 12 | ctx, w := workspaceTest(t) 13 | 14 | created, err := w.Warehouses.CreateAndWait(ctx, sql.CreateWarehouseRequest{ 15 | Name: RandomName("go-sdk-"), 16 | ClusterSize: "2X-Small", 17 | MaxNumClusters: 1, 18 | AutoStopMins: 10, 19 | Tags: &sql.EndpointTags{ 20 | CustomTags: []sql.EndpointTagPair{ 21 | { 22 | Key: "Owner", 23 | Value: "eng-dev-ecosystem-team_at_databricks.com", 24 | }, 25 | }, 26 | }, 27 | }) 28 | require.NoError(t, err) 29 | 30 | t.Cleanup(func() { 31 | err = w.Warehouses.DeleteById(ctx, created.Id) 32 | require.NoError(t, err) 33 | }) 34 | 35 | _, err = w.Warehouses.Edit(ctx, sql.EditWarehouseRequest{ 36 | Id: created.Id, 37 | Name: RandomName("go-sdk-updated-"), 38 | ClusterSize: "2X-Small", 39 | MaxNumClusters: 1, 40 | AutoStopMins: 10, 41 | }) 42 | require.NoError(t, err) 43 | 44 | wh, err := w.Warehouses.GetById(ctx, created.Id) 45 | require.NoError(t, err) 46 | 47 | all, err := w.Warehouses.ListAll(ctx, sql.ListWarehousesRequest{}) 48 | require.NoError(t, err) 49 | 50 | names, err := w.Warehouses.EndpointInfoNameToIdMap(ctx, sql.ListWarehousesRequest{}) 51 | require.NoError(t, err) 52 | assert.Equal(t, len(all), len(names)) 53 | assert.Equal(t, wh.Id, names[wh.Name]) 54 | 55 | byName, err := w.Warehouses.GetByName(ctx, wh.Name) 56 | require.NoError(t, err) 57 | assert.Equal(t, wh.Id, byName.Id) 58 | } 59 | -------------------------------------------------------------------------------- /common/response_test.go: -------------------------------------------------------------------------------- 1 | package common 2 | 3 | import ( 4 | "context" 5 | "io" 6 | "net/http" 7 | "strings" 8 | "testing" 9 | 10 | "github.com/stretchr/testify/require" 11 | ) 12 | 13 | type oneTimeCloser struct { 14 | r io.Reader 15 | isClosed bool 16 | } 17 | 18 | func newOneTimeCloser(r io.Reader) *oneTimeCloser { 19 | return &oneTimeCloser{r: r} 20 | } 21 | 22 | func (c *oneTimeCloser) Read(p []byte) (n int, err error) { 23 | if c.isClosed { 24 | return 0, io.ErrClosedPipe 25 | } 26 | return c.r.Read(p) 27 | } 28 | 29 | func (c *oneTimeCloser) Close() error { 30 | if c.isClosed { 31 | return io.ErrClosedPipe 32 | } 33 | c.isClosed = true 34 | return nil 35 | } 36 | 37 | func TestNewResponseWrapperBodyCanBeConsumed(t *testing.T) { 38 | ctx := context.Background() 39 | inner, err := http.NewRequestWithContext(ctx, "GET", "abc", nil) 40 | inner.Header.Set("Content-Type", "application/json") 41 | require.NoError(t, err) 42 | req, err := NewRequestBody("abc123") 43 | require.NoError(t, err) 44 | resp := &http.Response{ 45 | Body: newOneTimeCloser(strings.NewReader("Response Body")), 46 | Request: inner, 47 | Header: make(http.Header), 48 | } 49 | resp.Header.Set("Content-Type", "application/json") 50 | 51 | wrapper, err := NewResponseWrapper(resp, req) 52 | require.NoError(t, err) 53 | 54 | bs, err := io.ReadAll(wrapper.Response.Body) 55 | require.NoError(t, err) 56 | // The response body should be available both in the Body field of the 57 | // *http.Response as well as the DebugBytes field of the ResponseWrapper, 58 | // because this is a non-streaming response. 59 | require.Equal(t, "Response Body", string(bs)) 60 | require.Equal(t, "Response Body", string(wrapper.DebugBytes)) 61 | } 62 | -------------------------------------------------------------------------------- /examples/slog/slog.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | 7 | "github.com/databricks/databricks-sdk-go/logger" 8 | "golang.org/x/exp/slog" 9 | ) 10 | 11 | // slogAdapter makes an slog.Logger usable with the Databricks SDK. 12 | type slogAdapter struct { 13 | *slog.Logger 14 | } 15 | 16 | func (s *slogAdapter) Enabled(ctx context.Context, level logger.Level) bool { 17 | switch level { 18 | case logger.LevelTrace: 19 | // Note: slog doesn't have a default trace level. 20 | // An application can define their own fine grained levels 21 | // and use those here, if needed. 22 | return s.Logger.Enabled(ctx, slog.LevelDebug) 23 | case logger.LevelDebug: 24 | return s.Logger.Enabled(ctx, slog.LevelDebug) 25 | case logger.LevelInfo: 26 | return s.Logger.Enabled(ctx, slog.LevelInfo) 27 | case logger.LevelWarn: 28 | return s.Logger.Enabled(ctx, slog.LevelWarn) 29 | case logger.LevelError: 30 | return s.Logger.Enabled(ctx, slog.LevelError) 31 | default: 32 | return true 33 | } 34 | } 35 | 36 | func (s *slogAdapter) Tracef(ctx context.Context, format string, v ...any) { 37 | s.DebugContext(ctx, fmt.Sprintf(format, v...)) 38 | } 39 | 40 | func (s *slogAdapter) Debugf(ctx context.Context, format string, v ...any) { 41 | s.DebugContext(ctx, fmt.Sprintf(format, v...)) 42 | } 43 | 44 | func (s *slogAdapter) Infof(ctx context.Context, format string, v ...any) { 45 | s.InfoContext(ctx, fmt.Sprintf(format, v...)) 46 | } 47 | 48 | func (s *slogAdapter) Warnf(ctx context.Context, format string, v ...any) { 49 | s.WarnContext(ctx, fmt.Sprintf(format, v...)) 50 | } 51 | 52 | func (s *slogAdapter) Errorf(ctx context.Context, format string, v ...any) { 53 | s.ErrorContext(ctx, fmt.Sprintf(format, v...), nil) 54 | } 55 | -------------------------------------------------------------------------------- /service/files/doc.go: -------------------------------------------------------------------------------- 1 | // Databricks File System (DBFS) API 2 | // 3 | // We recommend using a client created via [databricks.NewWorkspaceClient] 4 | // to simplify the configuration experience. 5 | // 6 | // # Reading and writing files 7 | // 8 | // You can open a file on DBFS for reading or writing with [DbfsAPI.Open]. 9 | // This function returns a [Handle] that is compatible with a subset of [io] 10 | // interfaces for reading, writing, and closing. 11 | // 12 | // Uploading a file from an [io.Reader]: 13 | // 14 | // upload, _ := os.Open("/path/to/local/file.ext") 15 | // remote, _ := w.Dbfs.Open(ctx, "/path/to/remote/file", files.FileModeWrite|files.FileModeOverwrite) 16 | // io.Copy(remote, upload) 17 | // remote.Close() 18 | // 19 | // Downloading a file to an [io.Writer]: 20 | // 21 | // download, _ := os.Create("/path/to/local") 22 | // remote, _ := w.Dbfs.Open(ctx, "/path/to/remote/file", files.FileModeRead) 23 | // _ = io.Copy(download, remote) 24 | // 25 | // # Reading and writing files from buffers 26 | // 27 | // You can read from or write to a DBFS file directly from a byte slice through 28 | // the convenience functions [DbfsAPI.ReadFile] and [DbfsAPI.WriteFile]. 29 | // 30 | // Uploading a file from a byte slice: 31 | // 32 | // buf := []byte("Hello world!") 33 | // _ = w.Dbfs.WriteFile(ctx, "/path/to/remote/file", buf) 34 | // 35 | // Downloading a file into a byte slice: 36 | // 37 | // buf, err := w.Dbfs.ReadFile(ctx, "/path/to/remote/file") 38 | // 39 | // # Moving files 40 | // 41 | // err := w.Dbfs.Move(ctx, files.Move{ 42 | // SourcePath: "/remote/src/path", 43 | // DestinationPath: "/remote/dst/path", 44 | // }) 45 | // 46 | // # Creating directories 47 | // 48 | // w.Dbfs.MkdirsByPath(ctx, "/remote/dir/path") 49 | package files 50 | -------------------------------------------------------------------------------- /useragent/patterns_test.go: -------------------------------------------------------------------------------- 1 | package useragent 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestMatchSemVer(t *testing.T) { 10 | assert.NoError(t, matchSemVer("1.2.3")) 11 | assert.NoError(t, matchSemVer("0.0.0-dev+2e014739024a")) 12 | assert.Error(t, matchSemVer("1.2.3.4")) 13 | assert.Error(t, matchSemVer("1.2")) 14 | } 15 | 16 | func TestMatchAlphanum(t *testing.T) { 17 | for _, v := range []string{ 18 | "foo", 19 | "FOO", 20 | "FOO123", 21 | "foo_bar", 22 | "foo-bar", 23 | "foo.bar", 24 | } { 25 | assert.NoError(t, matchAlphanum(v)) 26 | } 27 | 28 | for _, v := range []string{ 29 | "foo bar", 30 | "foo/bar", 31 | } { 32 | assert.Error(t, matchAlphanum(v)) 33 | } 34 | } 35 | 36 | func TestMatchAlphanumOrSemVer(t *testing.T) { 37 | for _, v := range []string{ 38 | "foo", 39 | "1.2.3", 40 | "0.0.0-dev+2e014739024a", 41 | "client.0", 42 | } { 43 | assert.NoError(t, matchAlphanumOrSemVer(v)) 44 | } 45 | for _, v := range []string{ 46 | "foo/bar", 47 | "1/2/3", 48 | } { 49 | assert.Error(t, matchAlphanumOrSemVer(v)) 50 | } 51 | } 52 | 53 | func TestSanitize(t *testing.T) { 54 | for _, v := range []string{ 55 | "foo", 56 | "FOO", 57 | "FOO123", 58 | "foo_bar", 59 | "foo-bar", 60 | "foo+bar", 61 | "foo.bar", 62 | "1.2.3", 63 | "client.0", 64 | } { 65 | assert.Equal(t, v, Sanitize(v)) 66 | } 67 | 68 | sanitizeMap := map[string]string{ 69 | "1@2#3?4,5/6!7 8 ": "1-2-3-4-5-6-7-8-", 70 | "foo bar": "foo-bar", 71 | "foo/bar": "foo-bar", 72 | "foo:)bar": "foo--bar", 73 | "foo😊bar": "foo-bar", 74 | } 75 | for k, v := range sanitizeMap { 76 | assert.Equal(t, v, Sanitize(k)) 77 | } 78 | } 79 | -------------------------------------------------------------------------------- /config/auth_azure_cli_federated_token_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/databricks/databricks-sdk-go/internal/env" 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | func TestAzureCliCredentials_FederatedTokenServicePrincipal(t *testing.T) { 12 | // This test verifies the fix where service principals authenticated via 13 | // federated token (like in AKS with workload identity) use a fallback mechanism: 14 | // try with tenant ID first, then retry without tenant ID if it fails. 15 | 16 | env.CleanupEnvironment(t) 17 | t.Setenv("PATH", testdataPath()) 18 | 19 | // Simulate a service principal with a client ID (not systemAssignedIdentity/userAssignedIdentity) 20 | // This represents the federated token scenario that occurs in AKS workload identity 21 | t.Setenv("AZ_USER_NAME", "5817e630-86b3-4f67-a38e-a63e6a1a401c") 22 | t.Setenv("AZ_USER_TYPE", "servicePrincipal") 23 | 24 | // This makes the mock az command fail when --tenant is passed, simulating the federated 25 | // token scenario where tenant ID causes authentication failure 26 | t.Setenv("FAIL_IF_TENANT_ID_SET", "true") 27 | 28 | aa := AzureCliCredentials{} 29 | cfg := &Config{ 30 | Host: "https://adb-1891644720860465.5.azuredatabricks.net/", 31 | AzureTenantID: "e6a2f6d5-ece9-4c0d-9464-9c493497cb8f", 32 | } 33 | 34 | // With the fallback fix, this should work: first attempt with --tenant fails, 35 | // then fallback without --tenant succeeds 36 | visitor, err := aa.Configure(context.Background(), cfg) 37 | 38 | assert.NoError(t, err, "Authentication should work with federated token service principals via fallback mechanism") 39 | assert.NotNil(t, visitor, "Should return a valid credentials provider") 40 | } 41 | -------------------------------------------------------------------------------- /examples/slog/go.mod: -------------------------------------------------------------------------------- 1 | module main 2 | 3 | go 1.18 4 | 5 | replace github.com/databricks/databricks-sdk-go v0.0.0 => ../.. 6 | 7 | require ( 8 | github.com/databricks/databricks-sdk-go v0.0.0 9 | golang.org/x/exp v0.0.0-20240222234643-814bf88cf225 10 | ) 11 | 12 | require ( 13 | cloud.google.com/go/compute v1.23.4 // indirect 14 | cloud.google.com/go/compute/metadata v0.2.3 // indirect 15 | github.com/felixge/httpsnoop v1.0.4 // indirect 16 | github.com/go-logr/logr v1.4.1 // indirect 17 | github.com/go-logr/stdr v1.2.2 // indirect 18 | github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect 19 | github.com/golang/protobuf v1.5.3 // indirect 20 | github.com/google/go-querystring v1.1.0 // indirect 21 | github.com/google/s2a-go v0.1.7 // indirect 22 | github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect 23 | go.opencensus.io v0.24.0 // indirect 24 | go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect 25 | go.opentelemetry.io/otel v1.24.0 // indirect 26 | go.opentelemetry.io/otel/metric v1.24.0 // indirect 27 | go.opentelemetry.io/otel/trace v1.24.0 // indirect 28 | golang.org/x/crypto v0.21.0 // indirect 29 | golang.org/x/mod v0.16.0 // indirect 30 | golang.org/x/net v0.23.0 // indirect 31 | golang.org/x/oauth2 v0.18.0 // indirect 32 | golang.org/x/sys v0.18.0 // indirect 33 | golang.org/x/text v0.14.0 // indirect 34 | golang.org/x/time v0.5.0 // indirect 35 | google.golang.org/api v0.169.0 // indirect 36 | google.golang.org/appengine v1.6.8 // indirect 37 | google.golang.org/genproto/googleapis/rpc v0.0.0-20240304161311-37d4d3c04a78 // indirect 38 | google.golang.org/grpc v1.62.0 // indirect 39 | google.golang.org/protobuf v1.33.0 // indirect 40 | gopkg.in/ini.v1 v1.67.0 // indirect 41 | ) 42 | -------------------------------------------------------------------------------- /marshal/reflect_cache.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "reflect" 5 | "strings" 6 | "sync" 7 | ) 8 | 9 | var mutexType sync.Mutex 10 | 11 | var typeCache = map[reflect.Type][]cachedType{} 12 | 13 | var mutexName sync.Mutex 14 | 15 | var nameCache = map[reflect.Type]string{} 16 | 17 | type cachedType struct { 18 | reflect.StructField 19 | JsonTag jsonTag 20 | IndexInStruct int 21 | } 22 | 23 | func getTypeFields(structType reflect.Type) []cachedType { 24 | mutexType.Lock() 25 | defer mutexType.Unlock() 26 | if res, ok := typeCache[structType]; ok { 27 | return res 28 | } 29 | res := []cachedType{} 30 | for i := 0; i < structType.NumField(); i++ { 31 | field := structType.Field(i) 32 | jsonTag := parseJSONTag(field) 33 | res = append(res, cachedType{ 34 | StructField: field, 35 | JsonTag: jsonTag, 36 | IndexInStruct: i, 37 | }) 38 | } 39 | typeCache[structType] = res 40 | return res 41 | } 42 | 43 | func getTypeName(structType reflect.Type) string { 44 | mutexName.Lock() 45 | defer mutexName.Unlock() 46 | if res, ok := nameCache[structType]; ok { 47 | return res 48 | } 49 | name := structType.Name() 50 | nameCache[structType] = name 51 | return name 52 | } 53 | 54 | func parseJSONTag(field reflect.StructField) jsonTag { 55 | raw := field.Tag.Get("json") 56 | name := field.Name 57 | 58 | if raw == "-" { 59 | return jsonTag{ignore: true} 60 | } 61 | 62 | parts := strings.Split(raw, ",") 63 | 64 | if parts[0] != "" { 65 | name = parts[0] 66 | } 67 | 68 | jsonTag := jsonTag{ 69 | name: name, 70 | } 71 | 72 | for _, v := range parts { 73 | switch v { 74 | case "omitempty": 75 | jsonTag.omitempty = true 76 | case "string": 77 | jsonTag.asString = true 78 | } 79 | } 80 | return jsonTag 81 | } 82 | -------------------------------------------------------------------------------- /service/compute/instance_profiles_usage_test.go: -------------------------------------------------------------------------------- 1 | // Code generated from Databricks SDK for Go integration tests by openapi.roll.TestRegenerateExamples. DO NOT EDIT. 2 | 3 | package compute_test 4 | 5 | import ( 6 | "context" 7 | 8 | "github.com/databricks/databricks-sdk-go" 9 | "github.com/databricks/databricks-sdk-go/logger" 10 | 11 | "github.com/databricks/databricks-sdk-go/service/compute" 12 | ) 13 | 14 | func ExampleInstanceProfilesAPI_Add_awsInstanceProfiles() { 15 | ctx := context.Background() 16 | w, err := databricks.NewWorkspaceClient() 17 | if err != nil { 18 | panic(err) 19 | } 20 | 21 | arn := "arn:aws:iam::000000000000:instance-profile/abc" 22 | 23 | err = w.InstanceProfiles.Add(ctx, compute.AddInstanceProfile{ 24 | InstanceProfileArn: arn, 25 | SkipValidation: true, 26 | IamRoleArn: "arn:aws:iam::000000000000:role/bcd", 27 | }) 28 | if err != nil { 29 | panic(err) 30 | } 31 | 32 | } 33 | 34 | func ExampleInstanceProfilesAPI_Edit_awsInstanceProfiles() { 35 | ctx := context.Background() 36 | w, err := databricks.NewWorkspaceClient() 37 | if err != nil { 38 | panic(err) 39 | } 40 | 41 | arn := "arn:aws:iam::000000000000:instance-profile/abc" 42 | 43 | err = w.InstanceProfiles.Edit(ctx, compute.InstanceProfile{ 44 | InstanceProfileArn: arn, 45 | IamRoleArn: "arn:aws:iam::000000000000:role/bcdf", 46 | }) 47 | if err != nil { 48 | panic(err) 49 | } 50 | 51 | } 52 | 53 | func ExampleInstanceProfilesAPI_List_awsInstanceProfiles() { 54 | ctx := context.Background() 55 | w, err := databricks.NewWorkspaceClient() 56 | if err != nil { 57 | panic(err) 58 | } 59 | 60 | all, err := w.InstanceProfiles.ListAll(ctx) 61 | if err != nil { 62 | panic(err) 63 | } 64 | logger.Infof(ctx, "found %v", all) 65 | 66 | } 67 | -------------------------------------------------------------------------------- /examples/long-running/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | 7 | "github.com/databricks/databricks-sdk-go" 8 | "github.com/databricks/databricks-sdk-go/service/compute" 9 | "github.com/databricks/databricks-sdk-go/service/jobs" 10 | ) 11 | 12 | func main() { 13 | w := databricks.Must(databricks.NewWorkspaceClient()) 14 | ctx := context.Background() 15 | 16 | // Fetch list of spark runtime versions 17 | sparkVersions, err := w.Clusters.SparkVersions(ctx) 18 | if err != nil { 19 | panic(err) 20 | } 21 | 22 | // Select the latest LTS version 23 | latestLTS, err := sparkVersions.Select(compute.SparkVersionRequest{ 24 | Latest: true, 25 | LongTermSupport: true, 26 | }) 27 | if err != nil { 28 | panic(err) 29 | } 30 | 31 | // Fetch list of available node types 32 | nodeTypes, err := w.Clusters.ListNodeTypes(ctx) 33 | if err != nil { 34 | panic(err) 35 | } 36 | 37 | // Select the smallest node type id 38 | smallestWithDisk, err := nodeTypes.Smallest(compute.NodeTypeRequest{ 39 | LocalDisk: true, 40 | }) 41 | if err != nil { 42 | panic(err) 43 | } 44 | 45 | allRuns, err := w.Jobs.ListRunsAll(ctx, jobs.ListRunsRequest{}) 46 | if err != nil { 47 | panic(err) 48 | } 49 | for _, run := range allRuns { 50 | println(run.RunId) 51 | } 52 | 53 | runningCluster, err := w.Clusters.CreateAndWait(ctx, compute.CreateCluster{ 54 | ClusterName: "Test cluster from SDK", 55 | SparkVersion: latestLTS, 56 | NodeTypeId: smallestWithDisk, 57 | AutoterminationMinutes: 15, 58 | NumWorkers: 1, 59 | }) 60 | if err != nil { 61 | panic(err) 62 | } 63 | 64 | fmt.Printf("Cluster is ready: %s#setting/clusters/%s/configuration\n", 65 | w.Config.Host, runningCluster.ClusterId) 66 | } 67 | -------------------------------------------------------------------------------- /marshal/types_pchildfs_test.go: -------------------------------------------------------------------------------- 1 | package marshal 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestPChildFSDefault(t *testing.T) { 8 | executeBasicMarshalTest(t, 9 | basicMarshalTest{ 10 | st: customStruct{ 11 | PChildFS: nil, 12 | }, 13 | jsonString: `{"childfs":{},"childnofs":{}}`, 14 | matchClassic: true, 15 | }, 16 | ) 17 | } 18 | 19 | func TestPChildFSValueEmpty(t *testing.T) { 20 | executeBasicMarshalTest(t, 21 | basicMarshalTest{ 22 | st: customStruct{ 23 | PChildFS: &structFS{}, 24 | }, 25 | jsonString: `{"pchildfs":{},"childfs":{},"childnofs":{}}`, 26 | matchClassic: true, 27 | }, 28 | ) 29 | } 30 | 31 | func TestPChildFSValue(t *testing.T) { 32 | executeBasicMarshalTest(t, 33 | basicMarshalTest{ 34 | st: customStruct{ 35 | PChildFS: &structFS{ 36 | Int: 3, 37 | }, 38 | }, 39 | jsonString: `{"pchildfs":{"childint":3},"childfs":{},"childnofs":{}}`, 40 | matchClassic: true, 41 | }, 42 | ) 43 | } 44 | 45 | func TestPChildFSForce(t *testing.T) { 46 | executeBasicMarshalTest(t, 47 | basicMarshalTest{ 48 | st: customStruct{ 49 | PChildFS: nil, 50 | ForceSendFields: []string{"PChildFS"}, 51 | }, 52 | jsonString: `{"childfs":{},"childnofs":{}}`, 53 | matchClassic: true, 54 | }, 55 | ) 56 | } 57 | 58 | func TestPChildFSChildForce(t *testing.T) { 59 | executeBasicMarshalTest(t, 60 | basicMarshalTest{ 61 | st: customStruct{ 62 | PChildFS: &structFS{ 63 | ForceSendFields: []string{"Int"}, 64 | }, 65 | ForceSendFields: []string{"Int"}, 66 | }, 67 | jsonString: `{"childfs":{},"pchildfs":{"childint":0},"childnofs":{}, "int":0}`, 68 | matchClassic: false, 69 | unmarshalForceSendField: []string{"Int"}, 70 | }, 71 | ) 72 | } 73 | -------------------------------------------------------------------------------- /credentials/u2m/unified_oauth_argument.go: -------------------------------------------------------------------------------- 1 | package u2m 2 | 3 | import ( 4 | "fmt" 5 | ) 6 | 7 | // UnifiedOAuthArgument is an interface that provides the necessary information 8 | // to authenticate using OAuth to a host that supports both account and workspace APIs. 9 | type UnifiedOAuthArgument interface { 10 | OAuthArgument 11 | 12 | // GetHost returns the host to authenticate to. 13 | GetHost() string 14 | 15 | // GetAccountId returns the account ID of the account to authenticate to. 16 | GetAccountId() string 17 | } 18 | 19 | // BasicUnifiedOAuthArgument is a basic implementation of the UnifiedOAuthArgument 20 | // interface that links each account with exactly one OAuth token. 21 | type BasicUnifiedOAuthArgument struct { 22 | host string 23 | accountID string 24 | } 25 | 26 | var _ UnifiedOAuthArgument = BasicUnifiedOAuthArgument{} 27 | 28 | // NewBasicUnifiedOAuthArgument creates a new BasicUnifiedOAuthArgument. 29 | func NewBasicUnifiedOAuthArgument(accountsHost, accountID string) (BasicUnifiedOAuthArgument, error) { 30 | if err := validateHost(accountsHost); err != nil { 31 | return BasicUnifiedOAuthArgument{}, err 32 | } 33 | return BasicUnifiedOAuthArgument{host: accountsHost, accountID: accountID}, nil 34 | } 35 | 36 | // GetAccountHost returns the host of the account to authenticate to. 37 | func (a BasicUnifiedOAuthArgument) GetHost() string { 38 | return a.host 39 | } 40 | 41 | // GetAccountId returns the account ID of the account to authenticate to. 42 | func (a BasicUnifiedOAuthArgument) GetAccountId() string { 43 | return a.accountID 44 | } 45 | 46 | // GetCacheKey returns a unique key for caching the OAuth token for the account. 47 | // The key is in the format "/oidc/accounts/". 48 | func (a BasicUnifiedOAuthArgument) GetCacheKey() string { 49 | return fmt.Sprintf("%s/oidc/accounts/%s", a.host, a.accountID) 50 | } 51 | -------------------------------------------------------------------------------- /common/types/fieldmask/fieldmask.go: -------------------------------------------------------------------------------- 1 | package fieldmask 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "net/url" 7 | "strings" 8 | ) 9 | 10 | // FieldMask represents a field mask as defined in Google's Well Known Types. 11 | // It is used to specify which fields of a resource should be included or excluded 12 | // in a request or response. 13 | type FieldMask struct { 14 | Paths []string 15 | } 16 | 17 | // New creates a FieldMask from a slice of field paths. 18 | func New(paths []string) *FieldMask { 19 | return &FieldMask{Paths: paths} 20 | } 21 | 22 | // MarshalJSON implements the [json.Marshaler] interface by formatting the 23 | // field mask as a string according to Google Well Known Type 24 | func (f FieldMask) MarshalJSON() ([]byte, error) { 25 | return json.Marshal(strings.Join(f.Paths, ",")) 26 | } 27 | 28 | // UnmarshalJSON implements the [json.Unmarshaler] interface by parsing the 29 | // field mask from a string according to Google Well Known Type 30 | func (f *FieldMask) UnmarshalJSON(data []byte) error { 31 | if f == nil { 32 | return fmt.Errorf("FieldMask.UnmarshalJSON on nil pointer") 33 | } 34 | 35 | var s string 36 | if err := json.Unmarshal(data, &s); err != nil { 37 | return err 38 | } 39 | 40 | if s != "" { 41 | f.Paths = strings.Split(s, ",") 42 | } else { 43 | f.Paths = []string{} 44 | } 45 | 46 | return nil 47 | } 48 | 49 | // EncodeValues implements the [query.Encoder] interface by encoding the 50 | // field mask as a string, like "a,b,c". 51 | // If the FieldMask is nil or empty, it returns nil. 52 | // If the url.Values is nil, it returns an error. 53 | func (f *FieldMask) EncodeValues(key string, v *url.Values) error { 54 | if f == nil || len(f.Paths) == 0 { 55 | return nil 56 | } 57 | if v == nil { 58 | return fmt.Errorf("url.Values is nil") 59 | } 60 | 61 | v.Set(key, strings.Join(f.Paths, ",")) 62 | return nil 63 | } 64 | -------------------------------------------------------------------------------- /httpclient/errors.go: -------------------------------------------------------------------------------- 1 | package httpclient 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "io" 8 | "net/http" 9 | "net/url" 10 | "strings" 11 | 12 | "github.com/databricks/databricks-sdk-go/common" 13 | ) 14 | 15 | type HttpError struct { 16 | *http.Response 17 | Message string 18 | err error 19 | } 20 | 21 | func (r *HttpError) Unwrap() error { 22 | return r.err 23 | } 24 | 25 | func (r *HttpError) Error() string { 26 | return fmt.Sprintf("http %d: %s", r.StatusCode, r.Message) 27 | } 28 | 29 | // DefaultErrorMapper returns *HttpError 30 | func DefaultErrorMapper(ctx context.Context, resp common.ResponseWrapper) error { 31 | if resp.Response.StatusCode < 400 { 32 | return nil 33 | } 34 | raw, err := io.ReadAll(resp.ReadCloser) 35 | if err != nil { 36 | return &HttpError{ 37 | Response: resp.Response, 38 | Message: "failed to read response", 39 | err: err, 40 | } 41 | } 42 | return &HttpError{ 43 | Response: resp.Response, 44 | Message: string(raw), 45 | } 46 | } 47 | 48 | func DefaultErrorRetriable(ctx context.Context, err error) bool { 49 | var httpError *HttpError 50 | if errors.As(err, &httpError) { 51 | if httpError.StatusCode == http.StatusTooManyRequests { 52 | return true 53 | } 54 | if httpError.StatusCode == http.StatusGatewayTimeout { 55 | return true 56 | } 57 | } 58 | return false 59 | } 60 | 61 | var urlErrorTransientErrorMessages = []string{ 62 | "connection reset by peer", 63 | "TLS handshake timeout", 64 | "connection refused", 65 | "Unexpected error", 66 | "i/o timeout", 67 | } 68 | 69 | func isRetriableUrlError(err error) bool { 70 | var urlError *url.Error 71 | if !errors.As(err, &urlError) { 72 | return false 73 | } 74 | for _, msg := range urlErrorTransientErrorMessages { 75 | if strings.Contains(err.Error(), msg) { 76 | return true 77 | } 78 | } 79 | return false 80 | } 81 | -------------------------------------------------------------------------------- /openapi/gen/main.go: -------------------------------------------------------------------------------- 1 | // Usage: openapi-codegen 2 | package main 3 | 4 | import ( 5 | "context" 6 | "flag" 7 | "fmt" 8 | "os" 9 | "path" 10 | 11 | "github.com/databricks/databricks-sdk-go/openapi/code" 12 | "github.com/databricks/databricks-sdk-go/openapi/generator" 13 | "github.com/databricks/databricks-sdk-go/openapi/render" 14 | "github.com/databricks/databricks-sdk-go/openapi/roll" 15 | ) 16 | 17 | var c Context 18 | 19 | func main() { 20 | ctx := context.Background() 21 | cfg, err := render.Config() 22 | if err != nil { 23 | fmt.Printf("WARN: %s\n\n", err) 24 | } 25 | workDir, _ := os.Getwd() 26 | flag.StringVar(&c.Spec, "spec", cfg.Spec, "location of the spec file") 27 | flag.StringVar(&c.GoSDK, "gosdk", cfg.GoSDK, "location of the Go SDK") 28 | flag.StringVar(&c.Target, "target", workDir, "path to directory with .codegen.json") 29 | flag.BoolVar(&c.DryRun, "dry-run", false, "print to stdout instead of real files") 30 | flag.Parse() 31 | if c.Spec == "" { 32 | println("USAGE: go run openapi/gen/main.go -spec /path/to/spec.json") 33 | flag.PrintDefaults() 34 | os.Exit(1) 35 | } 36 | err = c.Run(ctx) 37 | if err != nil { 38 | fmt.Printf("ERROR: %s\n\n", err) 39 | os.Exit(1) 40 | } 41 | } 42 | 43 | type Context struct { 44 | Spec string 45 | GoSDK string 46 | Target string 47 | DryRun bool 48 | } 49 | 50 | func (c *Context) Run(ctx context.Context) error { 51 | spec, err := code.NewFromFile(ctx, c.Spec) 52 | if err != nil { 53 | return fmt.Errorf("spec: %w", err) 54 | } 55 | var suite *roll.Suite 56 | if c.GoSDK != "" { 57 | suite, err = roll.NewSuite(path.Join(c.GoSDK, "internal")) 58 | if err != nil { 59 | return fmt.Errorf("examples: %w", err) 60 | } 61 | } 62 | gen, err := generator.NewGenerator(c.Target) 63 | if err != nil { 64 | return fmt.Errorf("config: %w", err) 65 | } 66 | return gen.Apply(ctx, spec, suite) 67 | } 68 | -------------------------------------------------------------------------------- /credentials/u2m/account_oauth_argument.go: -------------------------------------------------------------------------------- 1 | package u2m 2 | 3 | import ( 4 | "fmt" 5 | ) 6 | 7 | // AccountOAuthArgument is an interface that provides the necessary information 8 | // to authenticate using OAuth to a specific account. 9 | type AccountOAuthArgument interface { 10 | OAuthArgument 11 | 12 | // GetAccountHost returns the host of the account to authenticate to. 13 | GetAccountHost() string 14 | 15 | // GetAccountId returns the account ID of the account to authenticate to. 16 | GetAccountId() string 17 | } 18 | 19 | // BasicAccountOAuthArgument is a basic implementation of the AccountOAuthArgument 20 | // interface that links each account with exactly one OAuth token. 21 | type BasicAccountOAuthArgument struct { 22 | accountHost string 23 | accountID string 24 | } 25 | 26 | var _ AccountOAuthArgument = BasicAccountOAuthArgument{} 27 | 28 | // NewBasicAccountOAuthArgument creates a new BasicAccountOAuthArgument. 29 | func NewBasicAccountOAuthArgument(accountsHost, accountID string) (BasicAccountOAuthArgument, error) { 30 | if err := validateHost(accountsHost); err != nil { 31 | return BasicAccountOAuthArgument{}, err 32 | } 33 | return BasicAccountOAuthArgument{accountHost: accountsHost, accountID: accountID}, nil 34 | } 35 | 36 | // GetAccountHost returns the host of the account to authenticate to. 37 | func (a BasicAccountOAuthArgument) GetAccountHost() string { 38 | return a.accountHost 39 | } 40 | 41 | // GetAccountId returns the account ID of the account to authenticate to. 42 | func (a BasicAccountOAuthArgument) GetAccountId() string { 43 | return a.accountID 44 | } 45 | 46 | // GetCacheKey returns a unique key for caching the OAuth token for the account. 47 | // The key is in the format "/oidc/accounts/". 48 | func (a BasicAccountOAuthArgument) GetCacheKey() string { 49 | return fmt.Sprintf("%s/oidc/accounts/%s", a.accountHost, a.accountID) 50 | } 51 | -------------------------------------------------------------------------------- /internal/repos_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | 7 | "github.com/databricks/databricks-sdk-go/service/settings" 8 | "github.com/databricks/databricks-sdk-go/service/workspace" 9 | "github.com/stretchr/testify/assert" 10 | "github.com/stretchr/testify/require" 11 | ) 12 | 13 | func TestAccRepos(t *testing.T) { 14 | ctx, w := workspaceTest(t) 15 | 16 | // Skip this test if "Files in Repos" is not enabled. 17 | conf, err := w.WorkspaceConf.GetStatus(ctx, settings.GetStatusRequest{ 18 | Keys: "enableWorkspaceFilesystem", 19 | }) 20 | require.NoError(t, err) 21 | if (*conf)["enableWorkspaceFilesystem"] == "false" { 22 | t.Skipf("Files in Repos not enabled in this workspace") 23 | } 24 | 25 | // Synthesize unique path for this checkout in this user's home. 26 | root := RandomName(fmt.Sprintf("/Repos/%s/tf-", me(t, w).UserName)) 27 | ri, err := w.Repos.Create(ctx, workspace.CreateRepoRequest{ 28 | Path: root, 29 | Url: "https://github.com/shreyas-goenka/empty-repo.git", 30 | Provider: "github", 31 | }) 32 | require.NoError(t, err) 33 | t.Cleanup(func() { 34 | err = w.Repos.DeleteByRepoId(ctx, ri.Id) 35 | require.NoError(t, err) 36 | }) 37 | 38 | assert.Equal(t, "main", ri.Branch) 39 | err = w.Repos.Update(ctx, workspace.UpdateRepoRequest{ 40 | RepoId: ri.Id, 41 | Branch: "foo", 42 | }) 43 | require.NoError(t, err) 44 | 45 | byId, err := w.Repos.GetByRepoId(ctx, ri.Id) 46 | require.NoError(t, err) 47 | 48 | byName, err := w.Repos.GetByPath(ctx, byId.Path) 49 | require.NoError(t, err) 50 | assert.Equal(t, byId.Id, byName.Id) 51 | 52 | all, err := w.Repos.ListAll(ctx, workspace.ListReposRequest{}) 53 | require.NoError(t, err) 54 | assert.True(t, len(all) >= 1) 55 | 56 | paths, err := w.Repos.RepoInfoPathToIdMap(ctx, workspace.ListReposRequest{ 57 | PathPrefix: "/", 58 | }) 59 | require.NoError(t, err) 60 | assert.Equal(t, len(paths), len(all)) 61 | } 62 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/databricks/databricks-sdk-go 2 | 3 | go 1.18 4 | 5 | require ( 6 | github.com/google/go-cmp v0.6.0 7 | github.com/google/go-querystring v1.1.0 8 | github.com/google/uuid v1.6.0 9 | github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c 10 | github.com/stretchr/testify v1.9.0 11 | golang.org/x/exp v0.0.0-20240222234643-814bf88cf225 12 | golang.org/x/mod v0.17.0 13 | golang.org/x/net v0.33.0 14 | golang.org/x/oauth2 v0.20.0 15 | golang.org/x/text v0.21.0 16 | golang.org/x/time v0.5.0 17 | google.golang.org/api v0.182.0 18 | google.golang.org/protobuf v1.34.1 19 | gopkg.in/ini.v1 v1.67.0 20 | ) 21 | 22 | require ( 23 | cloud.google.com/go/auth v0.4.2 // indirect 24 | cloud.google.com/go/auth/oauth2adapt v0.2.2 // indirect 25 | cloud.google.com/go/compute/metadata v0.3.0 // indirect 26 | github.com/davecgh/go-spew v1.1.1 // indirect 27 | github.com/felixge/httpsnoop v1.0.4 // indirect 28 | github.com/go-logr/logr v1.4.1 // indirect 29 | github.com/go-logr/stdr v1.2.2 // indirect 30 | github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect 31 | github.com/golang/protobuf v1.5.4 // indirect 32 | github.com/google/s2a-go v0.1.7 // indirect 33 | github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect 34 | github.com/pmezard/go-difflib v1.0.0 // indirect 35 | github.com/stretchr/objx v0.5.2 // indirect 36 | go.opencensus.io v0.24.0 // indirect 37 | go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect 38 | go.opentelemetry.io/otel v1.24.0 // indirect 39 | go.opentelemetry.io/otel/metric v1.24.0 // indirect 40 | go.opentelemetry.io/otel/trace v1.24.0 // indirect 41 | golang.org/x/crypto v0.31.0 // indirect 42 | golang.org/x/sys v0.28.0 // indirect 43 | google.golang.org/genproto/googleapis/rpc v0.0.0-20240521202816-d264139d666e // indirect 44 | google.golang.org/grpc v1.64.1 // indirect 45 | gopkg.in/yaml.v3 v3.0.1 // indirect 46 | ) 47 | -------------------------------------------------------------------------------- /config/experimental/auth/auth.go: -------------------------------------------------------------------------------- 1 | // Package auth is an internal package that provides authentication utilities. 2 | // 3 | // IMPORTANT: This package is not meant to be used directly by consumers of the 4 | // SDK and is subject to change without notice. 5 | package auth 6 | 7 | import ( 8 | "context" 9 | 10 | "golang.org/x/oauth2" 11 | ) 12 | 13 | // Credentials is anything that can return authentication headers. 14 | type Credentials interface { 15 | AuthHeaders(context.Context) (map[string]string, error) 16 | } 17 | 18 | // CredentialsFn is an adapter to allow the use of ordinary functions as 19 | // Credentials. 20 | // 21 | // Example: 22 | // 23 | // creds := CredentialsFn(func(ctx context.Context) (map[string]string, error) { 24 | // return map[string]string{"Authorization": "Bearer " + token}, nil 25 | // }) 26 | type CredentialsFn func(context.Context) (map[string]string, error) 27 | 28 | func (fn CredentialsFn) AuthHeaders(ctx context.Context) (map[string]string, error) { 29 | return fn(ctx) 30 | } 31 | 32 | // A TokenSource is anything that can return a token. 33 | type TokenSource interface { 34 | // Token returns a token or an error. Token must be safe for concurrent use 35 | // by multiple goroutines. The returned Token must not be modified. 36 | Token(context.Context) (*oauth2.Token, error) 37 | } 38 | 39 | // TokenSourceFn is an adapter to allow the use of ordinary functions as 40 | // TokenSource. 41 | // 42 | // Example: 43 | // 44 | // ts := TokenSourceFn(func(ctx context.Context) (*oauth2.Token, error) { 45 | // return &oauth2.Token{}, nil 46 | // }) 47 | type TokenSourceFn func(context.Context) (*oauth2.Token, error) 48 | 49 | func (fn TokenSourceFn) Token(ctx context.Context) (*oauth2.Token, error) { 50 | return fn(ctx) 51 | } 52 | 53 | // OAuthCredentials is a Credentials and TokenSource that can be used to 54 | // authenticate with OAuth. 55 | type OAuthCredentials interface { 56 | Credentials 57 | TokenSource 58 | } 59 | -------------------------------------------------------------------------------- /credentials/u2m/doc.go: -------------------------------------------------------------------------------- 1 | /* 2 | Package u2m supports the user-to-machine (U2M) OAuth flow for authenticating with Databricks. 3 | 4 | Databricks uses the authorization code flow from OAuth 2.0 to authenticate users. This flow 5 | consists of four steps: 6 | 1. Retrieve an authorization code for a user by opening a browser and directing them to the 7 | Databricks authorization URL. 8 | 2. Exchange the authorization code for an access token. 9 | 3. Use the access token to authenticate with Databricks. 10 | 4. When the access token expires, use the refresh token to get a new access token. 11 | 12 | The token and authorization endpoints for Databricks vary depending on whether the host is 13 | an account- or workspace-level host. Account-level endpoints are fixed based on the account 14 | ID and host, while workspace-level endpoints are discovered using the OIDC discovery endpoint 15 | at /oidc/.well-known/oauth-authorization-server. 16 | 17 | To trigger the authorization flow, construct a PersistentAuth object and call the 18 | Challenge() method: 19 | 20 | auth, err := oauth.NewPersistentAuth(ctx) 21 | if err != nil { 22 | log.Fatalf("failed to create persistent auth: %v", err) 23 | } 24 | token, err := auth.Challenge(ctx, oauth.BasicAccountOAuthArgument{ 25 | AccountHost: "https://accounts.cloud.databricks.com", 26 | AccountID: "xyz", 27 | }) 28 | 29 | Because the U2M flow requires user interaction, the resulting access token and refresh token 30 | can be stored in a persistent cache to avoid prompting the user for credentials on every 31 | authentication attempt. By default, the cache is stored in ~/.databricks/token-cache.json. 32 | Retrieve the cached token by calling the Load() method: 33 | 34 | token, err := auth.Load(ctx, oauth.BasicAccountOAuthArgument{ 35 | AccountHost: "https://accounts.cloud.databricks.com", 36 | AccountID: "xyz", 37 | }) 38 | 39 | See the cache package for more information on customizing the cache. 40 | */ 41 | package u2m 42 | -------------------------------------------------------------------------------- /common/types/time/time.go: -------------------------------------------------------------------------------- 1 | package time 2 | 3 | import ( 4 | "net/url" 5 | "strings" 6 | "time" 7 | 8 | "google.golang.org/protobuf/encoding/protojson" 9 | "google.golang.org/protobuf/types/known/timestamppb" 10 | ) 11 | 12 | // Time is a wrapper for timestamppb.Timestamp to provide custom marshaling 13 | // for JSON and URL query strings. 14 | // 15 | // It embeds timestamppb.Timestamp and exposes the .AsTime() method to 16 | // easily convert to time.Time. 17 | // 18 | // Example: 19 | // 20 | // customTime := time.New(stdtime.Now()) 21 | // goTime := customTime.AsTime() 22 | type Time struct { 23 | internal *timestamppb.Timestamp 24 | } 25 | 26 | // New creates a custom Time from a standard time.Time. 27 | func New(t time.Time) *Time { 28 | return &Time{internal: timestamppb.New(t)} 29 | } 30 | 31 | // AsTime returns the underlying time.Time value. 32 | func (x *Time) AsTime() time.Time { 33 | if x == nil { 34 | return time.Time{} 35 | } 36 | return x.internal.AsTime() 37 | } 38 | 39 | // MarshalJSON implements the [json.Marshaler] interface 40 | // by marshalling the time as a protobuf Timestamp. 41 | func (t Time) MarshalJSON() ([]byte, error) { 42 | return protojson.Marshal(t.internal) 43 | } 44 | 45 | // UnmarshalJSON implements the [json.Unmarshaler] interface 46 | // by unmarshalling the time from a protobuf Timestamp. 47 | func (t *Time) UnmarshalJSON(b []byte) error { 48 | var pb timestamppb.Timestamp 49 | if err := protojson.Unmarshal(b, &pb); err != nil { 50 | return err 51 | } 52 | *t = *New(pb.AsTime()) 53 | return nil 54 | } 55 | 56 | // EncodeValues implements the [query.Encoder] interface by encoding the 57 | // time as a protobuf Timestamp. 58 | func (t Time) EncodeValues(key string, v *url.Values) error { 59 | res, err := protojson.Marshal(t.internal) 60 | if err != nil { 61 | return err 62 | } 63 | // remove the quotes from the string. 64 | queryValue := strings.Trim(string(res), "\"") 65 | v.Set(key, queryValue) 66 | return nil 67 | } 68 | -------------------------------------------------------------------------------- /credentials/u2m/cache/file_test.go: -------------------------------------------------------------------------------- 1 | package cache 2 | 3 | import ( 4 | "os" 5 | "path/filepath" 6 | "testing" 7 | 8 | "github.com/stretchr/testify/assert" 9 | "github.com/stretchr/testify/require" 10 | "golang.org/x/oauth2" 11 | ) 12 | 13 | func setup(t *testing.T) string { 14 | tempHomeDir := t.TempDir() 15 | return filepath.Join(tempHomeDir, "token-cache.json") 16 | } 17 | 18 | func TestStoreAndLookup(t *testing.T) { 19 | c, err := NewFileTokenCache(WithFileLocation(setup(t))) 20 | require.NoError(t, err) 21 | err = c.Store("x", &oauth2.Token{ 22 | AccessToken: "abc", 23 | }) 24 | require.NoError(t, err) 25 | 26 | err = c.Store("y", &oauth2.Token{ 27 | AccessToken: "bcd", 28 | }) 29 | require.NoError(t, err) 30 | 31 | tok, err := c.Lookup("x") 32 | require.NoError(t, err) 33 | assert.Equal(t, "abc", tok.AccessToken) 34 | 35 | _, err = c.Lookup("z") 36 | assert.Equal(t, ErrNotFound, err) 37 | } 38 | 39 | func TestNoCacheFileReturnsErrNotConfigured(t *testing.T) { 40 | l, err := NewFileTokenCache(WithFileLocation(setup(t))) 41 | require.NoError(t, err) 42 | _, err = l.Lookup("x") 43 | assert.Equal(t, ErrNotFound, err) 44 | } 45 | 46 | func TestLoadCorruptFile(t *testing.T) { 47 | f := setup(t) 48 | err := os.MkdirAll(filepath.Dir(f), ownerExecReadWrite) 49 | require.NoError(t, err) 50 | err = os.WriteFile(f, []byte("abc"), ownerExecReadWrite) 51 | require.NoError(t, err) 52 | 53 | _, err = NewFileTokenCache(WithFileLocation(f)) 54 | assert.EqualError(t, err, "load: parse: invalid character 'a' looking for beginning of value") 55 | } 56 | 57 | func TestLoadWrongVersion(t *testing.T) { 58 | f := setup(t) 59 | err := os.MkdirAll(filepath.Dir(f), ownerExecReadWrite) 60 | require.NoError(t, err) 61 | err = os.WriteFile(f, []byte(`{"version": 823, "things": []}`), ownerExecReadWrite) 62 | require.NoError(t, err) 63 | 64 | _, err = NewFileTokenCache(WithFileLocation(f)) 65 | assert.EqualError(t, err, "load: needs version 1, got version 823") 66 | } 67 | -------------------------------------------------------------------------------- /useragent/patterns.go: -------------------------------------------------------------------------------- 1 | package useragent 2 | 3 | import ( 4 | "fmt" 5 | "regexp" 6 | ) 7 | 8 | // Regular expression copied from https://semver.org/. 9 | const ( 10 | semVerCore = `(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)` 11 | semVerPrerelease = `(?:-((?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?` 12 | semVerBuildmetadata = `(?:\+([0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?` 13 | ) 14 | 15 | var regexpSemVer = regexp.MustCompile(`^` + semVerCore + semVerPrerelease + semVerBuildmetadata + `$`) 16 | 17 | // Sanitize replaces all non-alphanumeric characters with a hyphen. Use this to 18 | // ensure that the user agent value is valid. This is useful when the value is not 19 | // ensured to be valid at compile time. 20 | // 21 | // Example: You want to avoid having '/' and ' ' in the value because it will 22 | // make downstream applications fail. 23 | // 24 | // Note: Semver strings are comprised of alphanumeric characters, hyphens, periods 25 | // and plus signs. This function will not remove these characters. 26 | // see: 27 | // 1. https://semver.org/#spec-item-9 28 | // 2. https://semver.org/#spec-item-10 29 | var regexpAlphanum = regexp.MustCompile(`^[0-9A-Za-z_\.\+-]+$`) 30 | var regexpAlphanumInverse = regexp.MustCompile(`[^0-9A-Za-z_\.\+-]`) 31 | 32 | func Sanitize(s string) string { 33 | return regexpAlphanumInverse.ReplaceAllString(s, "-") 34 | } 35 | 36 | func matchSemVer(s string) error { 37 | if regexpSemVer.MatchString(s) { 38 | return nil 39 | } 40 | return fmt.Errorf("invalid semver string: %s", s) 41 | } 42 | 43 | func matchAlphanum(s string) error { 44 | if regexpAlphanum.MatchString(s) { 45 | return nil 46 | } 47 | return fmt.Errorf("invalid alphanumeric string: %s", s) 48 | } 49 | 50 | func matchAlphanumOrSemVer(s string) error { 51 | if regexpAlphanum.MatchString(s) || regexpSemVer.MatchString(s) { 52 | return nil 53 | } 54 | return fmt.Errorf("invalid alphanumeric or semver string: %s", s) 55 | } 56 | -------------------------------------------------------------------------------- /internal/connections_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | 7 | "github.com/databricks/databricks-sdk-go/apierr" 8 | "github.com/databricks/databricks-sdk-go/service/catalog" 9 | "github.com/stretchr/testify/assert" 10 | "github.com/stretchr/testify/require" 11 | ) 12 | 13 | func TestAccConnections(t *testing.T) { 14 | ctx, w := ucwsTest(t) 15 | 16 | connCreate, err := w.Connections.Create(ctx, catalog.CreateConnection{ 17 | Comment: "Go SDK Acceptance Test Connection", 18 | ConnectionType: catalog.ConnectionTypeDatabricks, 19 | Name: RandomName("go-sdk-connection-"), 20 | Options: map[string]string{ 21 | "host": fmt.Sprintf("%s-fake-workspace.cloud.databricks.com", RandomName("go-sdk-connection-")), 22 | "httpPath": fmt.Sprintf("/sql/1.0/warehouses/%s", RandomName("go-sdk-connection-")), 23 | "personalAccessToken": RandomName("go-sdk-connection-"), 24 | }, 25 | }) 26 | require.NoError(t, err) 27 | t.Cleanup(func() { 28 | err := w.Connections.Delete(ctx, catalog.DeleteConnectionRequest{Name: connCreate.Name}) 29 | require.True(t, err == nil || apierr.IsMissing(err)) 30 | }) 31 | 32 | connUpdate, err := w.Connections.Update(ctx, catalog.UpdateConnection{ 33 | Name: connCreate.Name, 34 | Options: map[string]string{ 35 | "host": fmt.Sprintf("%s-fake-workspace.cloud.databricks.com", RandomName("go-sdk-connection-")), 36 | "httpPath": fmt.Sprintf("/sql/1.0/warehouses/%s", RandomName("go-sdk-connection-")), 37 | "personalAccessToken": RandomName("go-sdk-connection-"), 38 | }, 39 | }) 40 | require.NoError(t, err) 41 | 42 | conn, err := w.Connections.Get(ctx, catalog.GetConnectionRequest{Name: connUpdate.Name}) 43 | require.NoError(t, err) 44 | assert.Equal(t, conn.Options, connUpdate.Options) 45 | 46 | connList, err := w.Connections.ListAll(ctx, catalog.ListConnectionsRequest{}) 47 | require.NoError(t, err) 48 | assert.True(t, len(connList) >= 1) 49 | } 50 | -------------------------------------------------------------------------------- /experimental/api/options.go: -------------------------------------------------------------------------------- 1 | package api 2 | 3 | import ( 4 | "time" 5 | ) 6 | 7 | // Option is an option used by Execute to control the behavior of an API call. 8 | // An Option essentially act as a convenient way to configure Options. 9 | type Option interface { 10 | Apply(*Options) error 11 | } 12 | 13 | type optionFn func(*Options) 14 | 15 | func (fn optionFn) Apply(opts *Options) error { 16 | fn(opts) 17 | return nil 18 | } 19 | 20 | // WithRetrier configures to use the given Retrier provider. If no retrier is 21 | // provided, the API call is not retried. 22 | // 23 | // The provider function must be thread-safe. 24 | func WithRetrier(provider func() Retrier) Option { 25 | return optionFn(func(o *Options) { 26 | o.retrier = provider 27 | }) 28 | } 29 | 30 | // WithTimeout is a convenience option to set the timeout duration in the Call 31 | // context. If the context already has a deadline, that deadline is updated to 32 | // the minimum of the context's deadline and the timeout. 33 | // 34 | // The timeout covers the whole Call execution; it is not a timeout of each 35 | // intermediary API call. 36 | func WithTimeout(t time.Duration) Option { 37 | return optionFn(func(o *Options) { 38 | o.timeout = t 39 | }) 40 | } 41 | 42 | // WithLimiter configures to use the given Limiter provider. 43 | // 44 | // The limiter is used to potential rate limit the API call. If no limiter is 45 | // provided, the API call is not rate limited. 46 | func WithLimiter(l Limiter) Option { 47 | return optionFn(func(o *Options) { 48 | o.rateLimiter = l 49 | }) 50 | } 51 | 52 | // Options to control the behavior of an API call. 53 | type Options struct { 54 | // Provides a new Retrier to be used to execute a Call. The function 55 | // is called for each Call and must be thread-safe. The retrier must be 56 | // fresh within the context of an Execute call (e.g. no need to reset a 57 | // BackoffStrategy). 58 | retrier func() Retrier 59 | 60 | rateLimiter Limiter 61 | 62 | timeout time.Duration 63 | } 64 | -------------------------------------------------------------------------------- /.vscode/testing.code-snippets: -------------------------------------------------------------------------------- 1 | { 2 | // https://snippet-generator.app/ 3 | "explain error": { 4 | "prefix": "ifef", 5 | "body": [ 6 | "if err != nil {", 7 | " return nil, fmt.Errorf(\"failed to XXX: %w\", err)", 8 | "}" 9 | ], 10 | "description": "explain error" 11 | }, 12 | "if not nil": { 13 | "prefix": "inn", 14 | "body": [ 15 | "if $1 != nil {", 16 | " $2", 17 | "}" 18 | ], 19 | "description": "if not nil" 20 | }, 21 | "Test case": { 22 | "prefix": "te", 23 | "body": [ 24 | "func Test$1(t *testing.T) {", 25 | " $2", 26 | "}" 27 | ], 28 | "description": "Test case" 29 | }, 30 | "Assert len": { 31 | "prefix": "al", 32 | "body": [ 33 | "assert.Len(t, $1, $2)" 34 | ], 35 | "description": "Assert len" 36 | }, 37 | "Method with single error result": { 38 | "prefix": "me", 39 | "body": [ 40 | "func (a $1) $2() error {", 41 | " return nil", 42 | "}" 43 | ], 44 | "description": "Method with single error result" 45 | }, 46 | "if err != nil return err": { 47 | "prefix": "ife", 48 | "body": [ 49 | "if err != nil {", 50 | " return fmt.Errorf(\"$1: %w\", err)", 51 | "}" 52 | ], 53 | "description": "if err != nil return err" 54 | }, 55 | "assert.EqualError": { 56 | "prefix": "aee", 57 | "body": [ 58 | "assert.EqualError(t, err, \"..\")" 59 | ], 60 | "description": "assert.EqualError" 61 | }, 62 | "assert.Equal": { 63 | "prefix": "ae", 64 | "body": [ 65 | "assert.Equal(t, \"..\", $1)" 66 | ], 67 | "description": "assert.Equal" 68 | }, 69 | "assert.NoError": { 70 | "prefix": "anoe", 71 | "body": [ 72 | "assert.NoError(t, err)" 73 | ], 74 | "description": "assert.NoError" 75 | }, 76 | "require.NoError": { 77 | "prefix": "rnoe", 78 | "body": [ 79 | "require.NoError(t, err)" 80 | ], 81 | "description": "require.NoError" 82 | }, 83 | "err :=": { 84 | "prefix": "e", 85 | "body": [ 86 | "err := " 87 | ], 88 | "description": "err :=" 89 | } 90 | } -------------------------------------------------------------------------------- /common/request.go: -------------------------------------------------------------------------------- 1 | package common 2 | 3 | import ( 4 | "bytes" 5 | "encoding/json" 6 | "fmt" 7 | "io" 8 | "strings" 9 | ) 10 | 11 | // Represents a request body. 12 | // 13 | // If the provided request data is an io.Reader, DebugBytes is set to 14 | // "". Otherwise, DebugBytes is set to the marshaled JSON 15 | // representation of the request data, and ReadCloser is set to a new 16 | // io.ReadCloser that reads from DebugBytes. 17 | // 18 | // Request bodies are never closed by the client, hence only accepting 19 | // io.Reader. 20 | type RequestBody struct { 21 | Reader io.Reader 22 | ContentType string 23 | DebugBytes []byte 24 | } 25 | 26 | func NewRequestBody(data any) (RequestBody, error) { 27 | switch v := data.(type) { 28 | case io.Reader: 29 | return RequestBody{ 30 | Reader: v, 31 | DebugBytes: []byte(""), 32 | }, nil 33 | case string: 34 | return RequestBody{ 35 | Reader: strings.NewReader(v), 36 | DebugBytes: []byte(v), 37 | }, nil 38 | case []byte: 39 | return RequestBody{ 40 | Reader: bytes.NewReader(v), 41 | DebugBytes: v, 42 | }, nil 43 | default: 44 | bs, err := json.Marshal(data) 45 | if err != nil { 46 | return RequestBody{}, fmt.Errorf("request marshal failure: %w", err) 47 | } 48 | return RequestBody{ 49 | Reader: bytes.NewReader(bs), 50 | ContentType: "application/json", 51 | DebugBytes: bs, 52 | }, nil 53 | } 54 | } 55 | 56 | // Reset a request body to its initial state. 57 | // 58 | // This is used to retry requests with a body that has already been read. 59 | // If the request body is not resettable (i.e. not nil and of type other than 60 | // strings.Reader or bytes.Reader), this will return an error. 61 | func (r RequestBody) Reset() error { 62 | if r.Reader == nil { 63 | return nil 64 | } 65 | if v, ok := r.Reader.(io.Seeker); ok { 66 | _, err := v.Seek(0, io.SeekStart) 67 | return err 68 | } else { 69 | return fmt.Errorf("cannot reset reader of type %T", r.Reader) 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /logger/logger.go: -------------------------------------------------------------------------------- 1 | package logger 2 | 3 | import ( 4 | "context" 5 | ) 6 | 7 | var DefaultLogger Logger = &SimpleLogger{} 8 | 9 | // Level maps to the logging levels in [Logger]. 10 | type Level int 11 | 12 | const ( 13 | LevelTrace = -8 14 | LevelDebug = -4 15 | LevelInfo = 0 16 | LevelWarn = 4 17 | LevelError = 8 18 | ) 19 | 20 | type Logger interface { 21 | // Enabled returns if logging at the specified level is enabled. 22 | // This can be used to avoid computating an expensive value if it 23 | // won't be logged anyway. 24 | Enabled(ctx context.Context, level Level) bool 25 | 26 | // Tracef logs a formatted string at [LevelTrace]. 27 | Tracef(ctx context.Context, format string, v ...any) 28 | 29 | // Debugf logs a formatted string at [LevelDebug]. 30 | Debugf(ctx context.Context, format string, v ...any) 31 | 32 | // Infof logs a formatted string at [LevelInfo]. 33 | Infof(ctx context.Context, format string, v ...any) 34 | 35 | // Warnf logs a formatted string at [LevelWarn]. 36 | Warnf(ctx context.Context, format string, v ...any) 37 | 38 | // Errorf logs a formatted string at [LevelError]. 39 | Errorf(ctx context.Context, format string, v ...any) 40 | } 41 | 42 | // Get returns either the logger configured on the context, 43 | // or the global logger if one isn't defined. 44 | func Get(ctx context.Context) Logger { 45 | logger, ok := FromContext(ctx) 46 | if !ok { 47 | logger = DefaultLogger 48 | } 49 | return logger 50 | } 51 | 52 | func Tracef(ctx context.Context, format string, v ...any) { 53 | Get(ctx).Tracef(ctx, format, v...) 54 | } 55 | 56 | func Debugf(ctx context.Context, format string, v ...any) { 57 | Get(ctx).Debugf(ctx, format, v...) 58 | } 59 | 60 | func Infof(ctx context.Context, format string, v ...any) { 61 | Get(ctx).Infof(ctx, format, v...) 62 | } 63 | 64 | func Warnf(ctx context.Context, format string, v ...any) { 65 | Get(ctx).Warnf(ctx, format, v...) 66 | } 67 | 68 | func Errorf(ctx context.Context, format string, v ...any) { 69 | Get(ctx).Errorf(ctx, format, v...) 70 | } 71 | -------------------------------------------------------------------------------- /config/token_source_strategy_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "net/http" 7 | "testing" 8 | 9 | "github.com/databricks/databricks-sdk-go/config/credentials" 10 | "github.com/databricks/databricks-sdk-go/config/experimental/auth" 11 | "golang.org/x/oauth2" 12 | ) 13 | 14 | // authHeader returns the Authorization header set by the given credentials 15 | // provider. It returns an empty string if the provider is nil. 16 | func authHeader(cp credentials.CredentialsProvider) string { 17 | if cp == nil { 18 | return "" 19 | } 20 | req := &http.Request{Header: http.Header{}} 21 | cp.SetHeaders(req) 22 | return req.Header.Get("Authorization") 23 | } 24 | 25 | func TestTokenSourceStrategy_Configure(t *testing.T) { 26 | testCases := []struct { 27 | desc string 28 | ts auth.TokenSource 29 | wantHeader string 30 | wantError bool 31 | }{ 32 | { 33 | desc: "token source return error", 34 | ts: auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { 35 | return nil, errors.New("test error") 36 | }), 37 | wantError: true, 38 | }, 39 | { 40 | desc: "token source return token", 41 | ts: auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { 42 | return &oauth2.Token{ 43 | AccessToken: "token-123", 44 | }, nil 45 | }), 46 | wantHeader: "Bearer token-123", 47 | }, 48 | } 49 | 50 | for _, tc := range testCases { 51 | t.Run(tc.desc, func(t *testing.T) { 52 | start := &tokenSourceStrategy{ 53 | name: "test-strategy", 54 | ts: tc.ts, 55 | } 56 | 57 | cp, err := start.Configure(context.Background(), &Config{}) 58 | gotHeader := authHeader(cp) 59 | 60 | if tc.wantError && err == nil { 61 | t.Errorf("Expected error, but got none") 62 | } 63 | if !tc.wantError && err != nil { 64 | t.Errorf("Expected no error, but got %q", err) 65 | } 66 | if gotHeader != tc.wantHeader { 67 | t.Errorf("Expected header %q, but got %q", tc.wantHeader, gotHeader) 68 | } 69 | }) 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /common/types/duration/duration.go: -------------------------------------------------------------------------------- 1 | package duration 2 | 3 | import ( 4 | "net/url" 5 | "strings" 6 | "time" 7 | 8 | "google.golang.org/protobuf/encoding/protojson" 9 | "google.golang.org/protobuf/types/known/durationpb" 10 | ) 11 | 12 | // Duration is a wrapper for durationpb.Duration to provide custom marshaling 13 | // for JSON and URL query strings. 14 | // 15 | // It embeds durationpb.Duration and exposes the .AsDuration() method to 16 | // easily convert to time.Duration. 17 | // 18 | // Example: 19 | // 20 | // customDur := duration.New(30 * time.Second) 21 | // goDur := customDur.AsDuration() 22 | type Duration struct { 23 | internal *durationpb.Duration 24 | } 25 | 26 | // New creates a custom Duration from a standard time.Duration. 27 | func New(d time.Duration) *Duration { 28 | return &Duration{internal: durationpb.New(d)} 29 | } 30 | 31 | // AsDuration returns the underlying time.Duration value. 32 | func (x *Duration) AsDuration() time.Duration { 33 | if x == nil { 34 | return 0 35 | } 36 | return x.internal.AsDuration() 37 | } 38 | 39 | // MarshalJSON implements the [json.Marshaler] interface 40 | // by marshalling the duration as a protobuf Duration. 41 | func (d Duration) MarshalJSON() ([]byte, error) { 42 | return protojson.Marshal(d.internal) 43 | } 44 | 45 | // EncodeValues implements the [query.Encoder] interface by encoding the 46 | // duration as a string, like "3.3s". 47 | func (d Duration) EncodeValues(key string, v *url.Values) error { 48 | res, err := protojson.Marshal(d.internal) 49 | if err != nil { 50 | return err 51 | } 52 | // remove the quotes from the string 53 | queryValue := strings.Trim(string(res), "\"") 54 | v.Set(key, queryValue) 55 | return nil 56 | } 57 | 58 | // UnmarshalJSON implements the [json.Unmarshaler] interface. It can parse a 59 | // duration from the protobuf Duration. 60 | func (d *Duration) UnmarshalJSON(b []byte) error { 61 | var pb durationpb.Duration 62 | if err := protojson.Unmarshal(b, &pb); err != nil { 63 | return err 64 | } 65 | *d = *New(pb.AsDuration()) 66 | return nil 67 | } 68 | -------------------------------------------------------------------------------- /experimental/mocks/mocks_test.go: -------------------------------------------------------------------------------- 1 | package mocks 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/databricks/databricks-sdk-go" 8 | "github.com/databricks/databricks-sdk-go/service/compute" 9 | "github.com/databricks/databricks-sdk-go/service/iam" 10 | "github.com/stretchr/testify/assert" 11 | "github.com/stretchr/testify/mock" 12 | ) 13 | 14 | func TestMockWorkspaceClient(t *testing.T) { 15 | ctx := context.Background() 16 | w := NewMockWorkspaceClient(t) 17 | mockClusters := w.GetMockClustersAPI() 18 | mockClusters.EXPECT().ListAll( 19 | ctx, 20 | mock.AnythingOfType("compute.ListClustersRequest"), 21 | ).Return( 22 | []compute.ClusterDetails{ 23 | {ClusterName: "test-cluster-1"}, 24 | {ClusterName: "test-cluster-2"}, 25 | }, nil) 26 | 27 | clusters, err := listClusters(ctx, w.WorkspaceClient) 28 | 29 | assert.NoError(t, err) 30 | assert.NotEmpty(t, clusters) 31 | assert.Equal(t, clusters[0].ClusterName, "test-cluster-1") 32 | assert.Equal(t, clusters[1].ClusterName, "test-cluster-2") 33 | } 34 | 35 | func TestAccountWorkspaceClient(t *testing.T) { 36 | ctx := context.Background() 37 | w := NewMockAccountClient(t) 38 | mockUsers := w.GetMockAccountUsersV2API() 39 | mockUsers.EXPECT().ListAll( 40 | ctx, 41 | mock.AnythingOfType("iam.ListAccountUsersRequest"), 42 | ).Return( 43 | []iam.AccountUser{ 44 | {DisplayName: "test-user-1"}, 45 | {DisplayName: "test-user-2"}, 46 | }, nil) 47 | 48 | users, err := listAccountUsers(ctx, w.AccountClient) 49 | 50 | assert.NoError(t, err) 51 | assert.NotEmpty(t, users) 52 | assert.Equal(t, users[0].DisplayName, "test-user-1") 53 | assert.Equal(t, users[1].DisplayName, "test-user-2") 54 | } 55 | 56 | func listClusters(ctx context.Context, w *databricks.WorkspaceClient) ([]compute.ClusterDetails, error) { 57 | return w.Clusters.ListAll(ctx, compute.ListClustersRequest{}) 58 | } 59 | 60 | func listAccountUsers(ctx context.Context, a *databricks.AccountClient) ([]iam.AccountUser, error) { 61 | return a.UsersV2.ListAll(ctx, iam.ListAccountUsersRequest{}) 62 | } 63 | -------------------------------------------------------------------------------- /internal/clusterpolicies_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/databricks/databricks-sdk-go/service/compute" 7 | "github.com/stretchr/testify/assert" 8 | "github.com/stretchr/testify/require" 9 | ) 10 | 11 | func TestAccClusterPolicyFamilies(t *testing.T) { 12 | ctx, w := workspaceTest(t) 13 | 14 | all, err := w.PolicyFamilies.ListAll(ctx, compute.ListPolicyFamiliesRequest{}) 15 | require.NoError(t, err) 16 | 17 | firstFamily, err := w.PolicyFamilies.Get(ctx, compute.GetPolicyFamilyRequest{ 18 | PolicyFamilyId: all[0].PolicyFamilyId, 19 | }) 20 | require.NoError(t, err) 21 | assert.NotEqual(t, "", firstFamily.Name) 22 | } 23 | 24 | func TestAccClusterPolicies(t *testing.T) { 25 | ctx, w := workspaceTest(t) 26 | 27 | created, err := w.ClusterPolicies.Create(ctx, compute.CreatePolicy{ 28 | Name: RandomName("go-sdk-"), 29 | Definition: `{ 30 | "spark_conf.spark.databricks.delta.preview.enabled": { 31 | "type": "fixed", 32 | "value": true 33 | } 34 | }`, 35 | }) 36 | require.NoError(t, err) 37 | defer w.ClusterPolicies.DeleteByPolicyId(ctx, created.PolicyId) 38 | 39 | policy, err := w.ClusterPolicies.GetByPolicyId(ctx, created.PolicyId) 40 | require.NoError(t, err) 41 | 42 | byName, err := w.ClusterPolicies.GetByName(ctx, policy.Name) 43 | require.NoError(t, err) 44 | assert.Equal(t, policy.PolicyId, byName.PolicyId) 45 | 46 | all, err := w.ClusterPolicies.ListAll(ctx, compute.ListClusterPoliciesRequest{}) 47 | require.NoError(t, err) 48 | 49 | names, err := w.ClusterPolicies.PolicyNameToPolicyIdMap(ctx, compute.ListClusterPoliciesRequest{}) 50 | require.NoError(t, err) 51 | assert.Equal(t, len(names), len(all)) 52 | assert.Equal(t, policy.PolicyId, names[policy.Name]) 53 | 54 | err = w.ClusterPolicies.Edit(ctx, compute.EditPolicy{ 55 | PolicyId: policy.PolicyId, 56 | Name: policy.Name, 57 | Definition: `{ 58 | "spark_conf.spark.databricks.delta.preview.enabled": { 59 | "type": "fixed", 60 | "value": false 61 | } 62 | }`, 63 | }) 64 | require.NoError(t, err) 65 | } 66 | -------------------------------------------------------------------------------- /internal/account_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | 7 | "github.com/databricks/databricks-sdk-go/apierr" 8 | "github.com/databricks/databricks-sdk-go/service/iam" 9 | "github.com/stretchr/testify/assert" 10 | "github.com/stretchr/testify/require" 11 | ) 12 | 13 | func TestAccAccountServicePrincipal(t *testing.T) { 14 | ctx, a := accountTest(t) 15 | 16 | spCreate, err := a.ServicePrincipals.Create(ctx, iam.ServicePrincipal{ 17 | Active: true, 18 | DisplayName: RandomName("go-sdk-sp-"), 19 | }) 20 | require.NoError(t, err) 21 | t.Cleanup(func() { 22 | err := a.ServicePrincipals.Delete(ctx, iam.DeleteAccountServicePrincipalRequest{Id: spCreate.Id}) 23 | require.True(t, err == nil || apierr.IsMissing(err)) 24 | }) 25 | 26 | sp, err := a.ServicePrincipals.GetById(ctx, spCreate.Id) 27 | require.NoError(t, err) 28 | assert.Equal(t, spCreate.Id, sp.Id) 29 | 30 | spGet, err := a.ServicePrincipals.Get(ctx, iam.GetAccountServicePrincipalRequest{Id: sp.Id}) 31 | require.NoError(t, err) 32 | assert.Equal(t, spGet.DisplayName, sp.DisplayName) 33 | 34 | spGetByDisplayName, err := a.ServicePrincipals.GetByDisplayName(ctx, sp.DisplayName) 35 | require.NoError(t, err) 36 | assert.Equal(t, spGetByDisplayName.Id, sp.Id) 37 | 38 | spList, err := a.ServicePrincipals.ListAll(ctx, iam.ListAccountServicePrincipalsRequest{ 39 | Filter: fmt.Sprintf("displayName eq %v", sp.DisplayName), 40 | }) 41 | require.NoError(t, err) 42 | assert.Equal(t, spList[0].Id, sp.Id) 43 | 44 | err = a.ServicePrincipals.Patch(ctx, iam.PartialUpdate{ 45 | Id: sp.Id, 46 | Operations: []iam.Patch{ 47 | { 48 | Op: iam.PatchOpReplace, 49 | Path: "active", 50 | Value: "false", 51 | }, 52 | }, 53 | Schemas: []iam.PatchSchema{ 54 | iam.PatchSchemaUrnIetfParamsScimApiMessages20PatchOp, 55 | }, 56 | }) 57 | require.NoError(t, err) 58 | 59 | err = a.ServicePrincipals.Update(ctx, iam.ServicePrincipal{ 60 | Active: true, 61 | DisplayName: sp.DisplayName, 62 | Id: sp.Id, 63 | }) 64 | require.NoError(t, err) 65 | } 66 | --------------------------------------------------------------------------------