Skip to content

Commit ff0597a

Browse files
tejaskochar-dbtejaskochar-db
committed
add SPOG flags when calling auth token
1 parent ea6b181 commit ff0597a

File tree

2 files changed

+113
-35
lines changed

2 files changed

+113
-35
lines changed

‎config/cli_token_source.go‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,11 +39,28 @@ func NewCliTokenSource(cfg *Config) (*CliTokenSource, error) {
3939
if err != nil {
4040
return nil, err
4141
}
42+
cmd := buildCliCommand(cliPath, cfg)
43+
return &CliTokenSource{cmd: cmd}, nil
44+
}
45+
46+
// buildCliCommand constructs the CLI command arguments for fetching an auth token.
47+
// It handles unified hosts (with optional account_id and workspace_id), account hosts,
48+
// and workspace hosts.
49+
func buildCliCommand(cliPath string, cfg *Config) []string {
4250
cmd := []string{cliPath, "auth", "token", "--host", cfg.Host}
43-
if cfg.HostType() == AccountHost {
51+
if cfg.Experimental_IsUnifiedHost {
52+
// For unified hosts, pass account_id, workspace_id, and experimental flag
53+
cmd = append(cmd, "--experimental-is-unified-host")
54+
if cfg.AccountID != "" {
55+
cmd = append(cmd, "--account-id", cfg.AccountID)
56+
}
57+
if cfg.WorkspaceId != "" {
58+
cmd = append(cmd, "--workspace-id", cfg.WorkspaceId)
59+
}
60+
} else if cfg.HostType() == AccountHost {
4461
cmd = append(cmd, "--account-id", cfg.AccountID)
4562
}
46-
return &CliTokenSource{cmd: cmd}, nil
63+
return cmd
4764
}
4865

4966
func (c *CliTokenSource) Token(ctx context.Context) (*oauth2.Token, error) {

‎config/cli_token_source_test.go‎

Lines changed: 94 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"os"
88
"path/filepath"
99
"runtime"
10+
"golang.org/x/exp/slices"
1011
"strings"
1112
"testing"
1213
"time"
@@ -123,62 +124,122 @@ func TestFindDatabricksCli(t *testing.T) {
123124
}
124125
}
125126

126-
func TestNewCliTokenSource(t *testing.T) {
127-
tempDir := t.TempDir()
128-
129-
cliName := "databricks"
130-
if runtime.GOOS == "windows" {
131-
cliName = "databricks.exe"
132-
}
133-
validCliPath := filepath.Join(tempDir, cliName)
134-
if err := os.WriteFile(validCliPath, make([]byte, databricksCliMinSize+1), 0755); err != nil {
135-
t.Fatalf("failed to create mock CLI: %v", err)
136-
}
127+
func TestBuildCliCommand(t *testing.T) {
128+
const (
129+
cliPath = "/path/to/databricks"
130+
host = "https://workspace.cloud.databricks.com"
131+
accountHost = "https://accounts.cloud.databricks.com"
132+
unifiedHost = "https://unified.cloud.databricks.com"
133+
accountID = "abc-123"
134+
workspaceID = "456"
135+
)
137136

138137
testCases := []struct {
139138
name string
140139
cfg *Config
141140
wantCmd []string
142-
wantErr error
143141
}{
144142
{
145143
name: "workspace host",
146-
cfg: &Config{DatabricksCliPath: validCliPath, Host: "https://workspace.cloud.databricks.com"},
147-
wantCmd: []string{validCliPath, "auth", "token", "--host", "https://workspace.cloud.databricks.com"},
144+
cfg: &Config{Host: host},
145+
wantCmd: []string{cliPath, "auth", "token", "--host", host},
148146
},
149147
{
150148
name: "account host",
151-
cfg: &Config{DatabricksCliPath: validCliPath, Host: "https://accounts.cloud.databricks.com", AccountID: "abc-123"},
152-
wantCmd: []string{validCliPath, "auth", "token", "--host", "https://accounts.cloud.databricks.com", "--account-id", "abc-123"},
149+
cfg: &Config{Host: accountHost, AccountID: accountID},
150+
wantCmd: []string{cliPath, "auth", "token", "--host", accountHost, "--account-id", accountID},
151+
},
152+
{
153+
name: "unified host with account ID and workspace ID",
154+
cfg: &Config{
155+
Host: unifiedHost,
156+
Experimental_IsUnifiedHost: true,
157+
AccountID: accountID,
158+
WorkspaceId: workspaceID,
159+
},
160+
wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--experimental-is-unified-host", "--account-id", accountID, "--workspace-id", workspaceID},
153161
},
154162
{
155-
name: "CLI not found",
156-
cfg: &Config{DatabricksCliPath: filepath.Join(tempDir, "nonexistent"), Host: "https://workspace.cloud.databricks.com"},
157-
wantErr: ErrCliNotFound,
163+
name: "unified host with account ID only",
164+
cfg: &Config{
165+
Host: unifiedHost,
166+
Experimental_IsUnifiedHost: true,
167+
AccountID: accountID,
168+
},
169+
wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--experimental-is-unified-host", "--account-id", accountID},
170+
},
171+
{
172+
name: "unified host with workspace ID only",
173+
cfg: &Config{
174+
Host: unifiedHost,
175+
Experimental_IsUnifiedHost: true,
176+
WorkspaceId: workspaceID,
177+
},
178+
wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--experimental-is-unified-host", "--workspace-id", workspaceID},
179+
},
180+
{
181+
name: "unified host with no account ID or workspace ID",
182+
cfg: &Config{
183+
Host: unifiedHost,
184+
Experimental_IsUnifiedHost: true,
185+
},
186+
wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--experimental-is-unified-host"},
187+
},
188+
{
189+
// Explicit false should fall back to the standard host type detection
190+
name: "unified host false with account host",
191+
cfg: &Config{
192+
Host: accountHost,
193+
Experimental_IsUnifiedHost: false,
194+
AccountID: accountID,
195+
},
196+
wantCmd: []string{cliPath, "auth", "token", "--host", accountHost, "--account-id", accountID},
158197
},
159198
}
160199

161200
for _, tc := range testCases {
162201
t.Run(tc.name, func(t *testing.T) {
163-
ts, err := NewCliTokenSource(tc.cfg)
164-
if tc.wantErr != nil {
165-
if !errors.Is(err, tc.wantErr) {
166-
t.Errorf("NewCliTokenSource() error = %v, want %v", err, tc.wantErr)
167-
}
168-
return
169-
}
170-
if err != nil {
171-
t.Fatalf("NewCliTokenSource() unexpected error: %v", err)
172-
}
173-
for i := range ts.cmd {
174-
if ts.cmd[i] != tc.wantCmd[i] {
175-
t.Errorf("cmd[%d] = %q, want %q", i, ts.cmd[i], tc.wantCmd[i])
176-
}
202+
got := buildCliCommand(cliPath, tc.cfg)
203+
if !slices.Equal(got, tc.wantCmd) {
204+
t.Errorf("buildCliCommand() = %v, want %v", got, tc.wantCmd)
177205
}
178206
})
179207
}
180208
}
181209

210+
func TestNewCliTokenSource(t *testing.T) {
211+
tempDir := t.TempDir()
212+
213+
cliName := "databricks"
214+
if runtime.GOOS == "windows" {
215+
cliName = "databricks.exe"
216+
}
217+
validCliPath := filepath.Join(tempDir, cliName)
218+
if err := os.WriteFile(validCliPath, make([]byte, databricksCliMinSize+1), 0755); err != nil {
219+
t.Fatalf("failed to create mock CLI: %v", err)
220+
}
221+
222+
t.Run("success", func(t *testing.T) {
223+
cfg := &Config{DatabricksCliPath: validCliPath, Host: "https://example.databricks.com"}
224+
ts, err := NewCliTokenSource(cfg)
225+
if err != nil {
226+
t.Fatalf("NewCliTokenSource() unexpected error: %v", err)
227+
}
228+
// Verify CLI path was resolved and used
229+
if ts.cmd[0] != validCliPath {
230+
t.Errorf("cmd[0] = %q, want %q", ts.cmd[0], validCliPath)
231+
}
232+
})
233+
234+
t.Run("CLI not found", func(t *testing.T) {
235+
cfg := &Config{DatabricksCliPath: filepath.Join(tempDir, "nonexistent"), Host: "https://example.databricks.com"}
236+
_, err := NewCliTokenSource(cfg)
237+
if !errors.Is(err, ErrCliNotFound) {
238+
t.Errorf("NewCliTokenSource() error = %v, want %v", err, ErrCliNotFound)
239+
}
240+
})
241+
}
242+
182243
func TestCliTokenSource_Token(t *testing.T) {
183244
if runtime.GOOS == "windows" {
184245
t.Skip("Skipping shell script test on Windows")

0 commit comments

Comments
 (0)