From 05ee90aab7217203e5ea3787343fa648def675b9 Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Fri, 6 Jun 2025 11:15:47 +0000
Subject: [PATCH 1/3] Tolerate query parameters in login

---
 libs/auth/arguments.go      |  5 ++-
 libs/auth/arguments_test.go | 77 +++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 2 deletions(-)
 create mode 100644 libs/auth/arguments_test.go

diff --git a/libs/auth/arguments.go b/libs/auth/arguments.go
index 17957ec511..899fc19a9e 100644
--- a/libs/auth/arguments.go
+++ b/libs/auth/arguments.go
@@ -18,8 +18,9 @@ func (a AuthArguments) ToOAuthArgument() (u2m.OAuthArgument, error) {
 		Host:      a.Host,
 		AccountID: a.AccountID,
 	}
+	host := cfg.CanonicalHostName()
 	if cfg.IsAccountClient() {
-		return u2m.NewBasicAccountOAuthArgument(cfg.Host, cfg.AccountID)
+		return u2m.NewBasicAccountOAuthArgument(host, cfg.AccountID)
 	}
-	return u2m.NewBasicWorkspaceOAuthArgument(cfg.Host)
+	return u2m.NewBasicWorkspaceOAuthArgument(host)
 }
diff --git a/libs/auth/arguments_test.go b/libs/auth/arguments_test.go
new file mode 100644
index 0000000000..7b6690d338
--- /dev/null
+++ b/libs/auth/arguments_test.go
@@ -0,0 +1,77 @@
+package auth
+
+import (
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/credentials/u2m"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestToOAuthArgument(t *testing.T) {
+	tests := []struct {
+		name      string
+		args      AuthArguments
+		wantHost  string
+		wantError bool
+	}{
+		{
+			name: "workspace with no scheme",
+			args: AuthArguments{
+				Host: "my-workspace.cloud.databricks.com",
+			},
+			wantHost: "https://my-workspace.cloud.databricks.com",
+		},
+		{
+			name: "workspace with https",
+			args: AuthArguments{
+				Host: "https://my-workspace.cloud.databricks.com",
+			},
+			wantHost: "https://my-workspace.cloud.databricks.com",
+		},
+		{
+			name: "account with no scheme",
+			args: AuthArguments{
+				Host:      "accounts.cloud.databricks.com",
+				AccountID: "123456789",
+			},
+			wantHost: "https://accounts.cloud.databricks.com",
+		},
+		{
+			name: "account with https",
+			args: AuthArguments{
+				Host:      "https://accounts.cloud.databricks.com",
+				AccountID: "123456789",
+			},
+			wantHost: "https://accounts.cloud.databricks.com",
+		},
+		{
+			name: "workspace with query parameter",
+			args: AuthArguments{
+				Host: "https://my-workspace.cloud.databricks.com?o=123456789",
+			},
+			wantHost: "https://my-workspace.cloud.databricks.com",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := tt.args.ToOAuthArgument()
+			if tt.wantError {
+				assert.Error(t, err)
+				return
+			}
+			assert.NoError(t, err)
+
+			// Check if we got the right type of argument and verify the hostname
+			if tt.args.AccountID != "" {
+				arg, ok := got.(u2m.AccountOAuthArgument)
+				assert.True(t, ok, "expected AccountOAuthArgument for account ID")
+				assert.Equal(t, tt.wantHost, arg.GetAccountHost())
+			} else {
+				arg, ok := got.(u2m.WorkspaceOAuthArgument)
+				assert.True(t, ok, "expected WorkspaceOAuthArgument for workspace")
+				assert.Equal(t, tt.wantHost, arg.GetWorkspaceHost())
+			}
+		})
+	}
+}

From 08043edfb2efd601ed11b8a2cbdbe499b167d137 Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Fri, 6 Jun 2025 11:20:32 +0000
Subject: [PATCH 2/3] fix

---
 NEXT_CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
index 71fec803b4..a932a581a1 100644
--- a/NEXT_CHANGELOG.md
+++ b/NEXT_CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Notable Changes
 
+* Fix `databricks auth login` to tolerate URLs copied from the browser ([#3001](https://github.com/databricks/cli/pull/3001)).
+
 ### Dependency updates
 
 ### CLI

From c0b68d849c1dfc41ff2bef7e398b94c320f92447 Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Fri, 6 Jun 2025 15:04:26 +0000
Subject: [PATCH 3/3] test

---
 libs/auth/arguments_test.go | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/libs/auth/arguments_test.go b/libs/auth/arguments_test.go
index 7b6690d338..d75827a771 100644
--- a/libs/auth/arguments_test.go
+++ b/libs/auth/arguments_test.go
@@ -51,6 +51,13 @@ func TestToOAuthArgument(t *testing.T) {
 			},
 			wantHost: "https://my-workspace.cloud.databricks.com",
 		},
+		{
+			name: "workspace with query parameter and path",
+			args: AuthArguments{
+				Host: "https://my-workspace.cloud.databricks.com/path?o=123456789",
+			},
+			wantHost: "https://my-workspace.cloud.databricks.com",
+		},
 	}
 
 	for _, tt := range tests {
