forked from databricks/databricks-sdk-go
/
azure.go
115 lines (102 loc) · 3.67 KB
/
azure.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
package config
import (
"context"
"encoding/json"
"fmt"
"io"
"strings"
"github.com/xuxiaoshuo/databricks-sdk-go/logger"
"golang.org/x/oauth2"
)
type azureEnvironment struct {
Name string `json:"name"`
ServiceManagementEndpoint string `json:"serviceManagementEndpoint"`
ResourceManagerEndpoint string `json:"resourceManagerEndpoint"`
ActiveDirectoryEndpoint string `json:"activeDirectoryEndpoint"`
}
// based on github.com/Azure/go-autorest/autorest/azure/azureEnvironments.go
var (
publicCloud = azureEnvironment{
Name: "AzurePublicCloud",
ServiceManagementEndpoint: "https://management.core.windows.net/",
ResourceManagerEndpoint: "https://management.azure.com/",
ActiveDirectoryEndpoint: "https://login.microsoftonline.com/",
}
usGovernmentCloud = azureEnvironment{
Name: "AzureUSGovernmentCloud",
ServiceManagementEndpoint: "https://management.core.usgovcloudapi.net/",
ResourceManagerEndpoint: "https://management.usgovcloudapi.net/",
ActiveDirectoryEndpoint: "https://login.microsoftonline.us/",
}
chinaCloud = azureEnvironment{
Name: "AzureChinaCloud",
ServiceManagementEndpoint: "https://management.core.chinacloudapi.cn/",
ResourceManagerEndpoint: "https://management.chinacloudapi.cn/",
ActiveDirectoryEndpoint: "https://login.chinacloudapi.cn/",
}
germanCloud = azureEnvironment{
Name: "AzureGermanCloud",
ServiceManagementEndpoint: "https://management.core.cloudapi.de/",
ResourceManagerEndpoint: "https://management.microsoftazure.de/",
ActiveDirectoryEndpoint: "https://login.microsoftonline.de/",
}
azureEnvironments = map[string]azureEnvironment{
"CHINA": chinaCloud,
"GERMAN": germanCloud,
"PUBLIC": publicCloud,
"USGOVERNMENT": usGovernmentCloud,
}
)
func (c *Config) GetAzureEnvironment() (azureEnvironment, error) {
if c.AzureEnvironment == "" {
c.AzureEnvironment = "public"
}
env, ok := azureEnvironments[strings.ToUpper(c.AzureEnvironment)]
if !ok {
return env, fmt.Errorf("azure envionment not found: %s", c.AzureEnvironment)
}
return env, nil
}
type azureHostResolver interface {
tokenSourceFor(ctx context.Context, cfg *Config, env azureEnvironment, resource string) oauth2.TokenSource
}
func (c *Config) azureEnsureWorkspaceUrl(ctx context.Context, ahr azureHostResolver) error {
if c.AzureResourceID == "" || c.Host != "" {
return nil
}
env, err := c.GetAzureEnvironment()
if err != nil {
return err
}
// azure resource ID can also be used in lieu of host by some of the clients, like Terraform
management := ahr.tokenSourceFor(ctx, c, env, env.ResourceManagerEndpoint)
resourceManager := oauth2.NewClient(ctx, management)
resp, err := resourceManager.Get(env.ResourceManagerEndpoint + c.AzureResourceID + "?api-version=2018-04-01")
if err != nil {
return fmt.Errorf("cannot resolve workspace: %w", err)
}
raw, err := io.ReadAll(resp.Body)
if err != nil {
return fmt.Errorf("cannot read: %w", err)
}
var workspaceMetadata struct {
Properties struct {
WorkspaceURL string `json:"workspaceUrl"`
} `json:"properties"`
}
err = json.Unmarshal(raw, &workspaceMetadata)
if err != nil {
return fmt.Errorf("cannot unmarshal: %w", err)
}
c.Host = fmt.Sprintf("https://%s", workspaceMetadata.Properties.WorkspaceURL)
logger.Debugf(ctx, "Discovered workspace url: %s", c.Host)
return nil
}
// Resource ID of the Azure application we need to log in.
const azureDatabricksLoginAppID string = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
func (c *Config) getAzureLoginAppID() string {
if c.AzureLoginAppID != "" {
return c.AzureLoginAppID
}
return azureDatabricksLoginAppID
}