-
Notifications
You must be signed in to change notification settings - Fork 41
/
config.go
239 lines (201 loc) · 7.76 KB
/
config.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
package config
import (
"context"
"fmt"
"net/http"
"net/url"
"strings"
"sync"
"github.com/databricks/databricks-sdk-go/logger"
)
// CredentialsProvider responsible for configuring static or refreshable
// authentication credentials for Databricks REST APIs
type CredentialsProvider interface {
// Name returns human-addressable name of this credentials provider name
Name() string
// Configure creates HTTP Request Visitor or returns nil if a given credetials
// are not configured. It returns an error if credentials are misconfigured.
// Takes a context and a pointer to a Config instance, that holds auth mutex.
Configure(context.Context, *Config) (func(*http.Request) error, error)
}
type Loader interface {
// Name is human-addressable representation of this config resolver
Name() string
Configure(*Config) error
}
// Config represents configuration for Databricks Connectivity
type Config struct {
// Credentials holds an instance of Credentials Provider to authenticate with Databricks REST APIs.
// If no credentials provider is specified, `DefaultCredentials` are implicitly used.
Credentials CredentialsProvider
// Databricks host (either of workspace endpoint or Accounts API endpoint)
Host string `name:"host" env:"DATABRICKS_HOST"`
// Databricks Account ID for Accounts API. This field is used in dependencies.
AccountID string `name:"account_id" env:"DATABRICKS_ACCOUNT_ID"`
Token string `name:"token" env:"DATABRICKS_TOKEN" auth:"pat,sensitive"`
Username string `name:"username" env:"DATABRICKS_USERNAME" auth:"basic"`
Password string `name:"password" env:"DATABRICKS_PASSWORD" auth:"basic,sensitive"`
// Connection profile specified within ~/.databrickscfg.
Profile string `name:"profile" env:"DATABRICKS_CONFIG_PROFILE"`
// Location of the Databricks CLI credentials file, that is created
// by `databricks configure --token` command. By default, it is located
// in ~/.databrickscfg.
ConfigFile string `name:"config_file" env:"DATABRICKS_CONFIG_FILE"`
GoogleServiceAccount string `name:"google_service_account" env:"DATABRICKS_GOOGLE_SERVICE_ACCOUNT" auth:"google"`
GoogleCredentials string `name:"google_credentials" env:"GOOGLE_CREDENTIALS" auth:"google,sensitive"`
// Azure Resource Manager ID for Azure Databricks workspace, which is exhanged for a Host
AzureResourceID string `name:"azure_workspace_resource_id" env:"DATABRICKS_AZURE_RESOURCE_ID" auth:"azure"`
AzureUseMSI bool `name:"azure_use_msi" env:"ARM_USE_MSI" auth:"azure"`
AzureClientSecret string `name:"azure_client_secret" env:"ARM_CLIENT_SECRET" auth:"azure,sensitive"`
AzureClientID string `name:"azure_client_id" env:"ARM_CLIENT_ID" auth:"azure"`
AzureTenantID string `name:"azure_tenant_id" env:"ARM_TENANT_ID" auth:"azure"`
// AzureEnvironment (Public, UsGov, China, Germany) has specific set of API endpoints.
AzureEnvironment string `name:"azure_environment" env:"ARM_ENVIRONMENT"`
// Azure Login Application ID. Must be set if authenticating for non-production workspaces.
AzureLoginAppID string `name:"azure_login_app_id" env:"DATABRICKS_AZURE_LOGIN_APP_ID" auth:"azure"`
ClientID string `name:"client_id" env:"DATABRICKS_CLIENT_ID"`
ClientSecret string `name:"client_secret" env:"DATABRICKS_CLIENT_SECRET"`
// When multiple auth attributes are available in the environment, use the auth type
// specified by this argument. This argument also holds currently selected auth.
AuthType string `name:"auth_type" env:"DATABRICKS_AUTH_TYPE" auth:"-"`
// Skip SSL certificate verification for HTTP calls.
// Use at your own risk or for unit testing purposes.
InsecureSkipVerify bool `name:"skip_verify" auth:"-"`
// Number of seconds for HTTP timeout
HTTPTimeoutSeconds int `name:"http_timeout_seconds" auth:"-"`
// Truncate JSON fields in JSON above this limit. Default is 96.
DebugTruncateBytes int `name:"debug_truncate_bytes" env:"DATABRICKS_DEBUG_TRUNCATE_BYTES" auth:"-"`
// Debug HTTP headers of requests made by the provider. Default is false.
DebugHeaders bool `name:"debug_headers" env:"DATABRICKS_DEBUG_HEADERS" auth:"-"`
// Maximum number of requests per second made to Databricks REST API.
RateLimitPerSecond int `name:"rate_limit" env:"DATABRICKS_RATE_LIMIT" auth:"-"`
// Number of seconds to keep retrying HTTP requests. Default is 300 (5 minutes)
RetryTimeoutSeconds int `name:"retry_timeout_seconds" auth:"-"`
Loaders []Loader
// marker for configuration resolving
resolved bool
// Mutex used by Authenticate method to guard `auth`, which
// has to be lazily created on the first request to Databricks API.
// It is done because databricks host and token may often be available
// only in the middle of Terraform DAG execution.
// This mutex is also used for config resolution.
mu sync.Mutex
// HTTP request interceptor, that assigns Authorization header
auth func(r *http.Request) error
}
// Authenticate adds special headers to HTTP request to authorize it to work with Databricks REST API
func (c *Config) Authenticate(r *http.Request) error {
err := c.EnsureResolved()
if err != nil {
return err
}
err = c.authenticateIfNeeded(r.Context())
if err != nil {
return err
}
return c.auth(r)
}
// IsAzure returns true if client is configured for Azure Databricks
func (c *Config) IsAzure() bool {
return strings.Contains(c.Host, ".azuredatabricks.net") || c.AzureResourceID != ""
}
// IsGcp returns true if client is configured for GCP
func (c *Config) IsGcp() bool {
return strings.Contains(c.Host, ".gcp.databricks.com")
}
// IsAws returns true if client is configured for AWS
func (c *Config) IsAws() bool {
return !c.IsAzure() && !c.IsGcp()
}
// IsAccountClient returns true if client is configured for Accounts API
func (c *Config) IsAccountClient() bool {
return strings.HasPrefix(c.Host, "https://accounts.")
}
func (c *Config) EnsureResolved() error {
if c.resolved {
return nil
}
c.mu.Lock()
defer c.mu.Unlock()
if c.resolved {
return nil
}
if len(c.Loaders) == 0 {
c.Loaders = []Loader{
ConfigAttributes,
KnownConfigLoader{},
}
}
for _, loader := range c.Loaders {
logger.Tracef("Loading config via %s", loader.Name())
err := loader.Configure(c)
if err != nil {
return c.wrapDebug(fmt.Errorf("resolve: %w", err))
}
}
err := ConfigAttributes.Validate(c)
if err != nil {
return c.wrapDebug(fmt.Errorf("validate: %w", err))
}
c.resolved = true
return nil
}
func (c *Config) wrapDebug(err error) error {
debug := ConfigAttributes.DebugString(c)
if debug == "" {
return err
}
return fmt.Errorf("%w. %s", err, debug)
}
// authenticateIfNeeded lazily authenticates across authorizers or returns error
func (c *Config) authenticateIfNeeded(ctx context.Context) error {
if c.auth != nil {
return nil
}
c.mu.Lock()
defer c.mu.Unlock()
if c.auth != nil {
return nil
}
if c.Credentials == nil {
c.Credentials = &DefaultCredentials{}
}
c.fixHostIfNeeded()
visitor, err := c.Credentials.Configure(ctx, c)
if err != nil {
return c.wrapDebug(fmt.Errorf("%s auth: %w", c.Credentials.Name(), err))
}
if visitor == nil {
return c.wrapDebug(fmt.Errorf("%s auth: not configured", c.Credentials.Name()))
}
c.auth = visitor
c.AuthType = c.Credentials.Name()
c.fixHostIfNeeded()
// TODO: error customization
return nil
}
func (c *Config) fixHostIfNeeded() error {
// Nothing to fix if the host isn't set.
if c.Host == "" {
return nil
}
parsedHost, err := url.Parse(c.Host)
if err != nil {
return err
}
// If the host is empty, assume the scheme wasn't included.
if parsedHost.Host == "" {
parsedHost, err = url.Parse("https://" + c.Host)
if err != nil {
return err
}
}
// Create new instance to ensure other fields are initialized as empty.
parsedHost = &url.URL{
Scheme: parsedHost.Scheme,
Host: parsedHost.Host,
}
// Store sanitized version of c.Host.
c.Host = parsedHost.String()
return nil
}