Show profile name alongside workspace URL in picker#114
Conversation
When selecting a workspace during first-time setup, display the Databricks CLI profile alias next to the host URL so users can easily identify which workspace to pick. Co-authored-by: Anthony Ivan <anthony.ivan@databricks.com>
|
Can we only do this when there are duplicates ? Else its a bit redundant |
|
@rohita5l the situation that I am facing is that I am using >5 workspaces in my local setup, and when I try to do the setup with ucode I find it really hard to know which workspace I need to select in order to target the workspace that I want to use. The change also brings the information shown closer to what |
Can you add a screenshot of what it looks like now ? |
- get_databricks_profiles now returns one entry per profile instead of deduping by host, so workspaces with multiple profiles each get a row. - prompt_for_workspace renders a `Profile Name` / `Workspace URL` header (matching `databricks auth profiles`, minus the Valid column) with ljust-padded names so columns align. - The picker's value is still (host, profile_name); that tuple is what flows into configure_shared_state and gets persisted to state.json, so the profile the user actually selects (not just the first one for the host) is what every later databricks CLI invocation uses. Co-authored-by: Isaac
|
@rohita5l I have updated the PR based on what we discussed on Slack |
| return out | ||
|
|
||
|
|
||
| def find_profile_name_for_host(workspace: str) -> str | None: |
There was a problem hiding this comment.
Is this being called by anywhere now ? If not can we delete it ?
There was a problem hiding this comment.
Yes it is still being called in multiple places in the code outside of the function I am modifying, I think if we want to refactor this we should do it as part of a different PR
Resolves tests/test_ui.py — keeps both the layout-coverage tests added on this branch and the return-type-defensiveness tests added on main in PR databricks#104, since they exercise complementary aspects of prompt_for_workspace. Co-authored-by: Isaac
rohita5l
left a comment
rohita5l
left a comment
There was a problem hiding this comment.
Could we add a regression test for the picker → configure_shared_state hop? Existing tests in TestConfigureAgentsSelection only exercise profile=None, so if a future refactor of _configure_shared_workspace_states
(cli.py:232) drops the profile from the tuple unpack, nothing fails. Suggested test (drop into TestConfigureAgentsSelection in tests/test_cli.py):
def test_picker_selected_profile_flows_to_configure_shared_state(self, monkeypatch):
"""Picker's (host, profile) tuple must reach configure_shared_state's
profile kwarg, otherwise downstream --profile calls fall back to
host-based resolution and silently pick the wrong profile."""
import ucode.cli as cli_mod
monkeypatch.setattr(
cli_mod,
"_prompt_for_configuration",
lambda tool=None: ("https://shared.cloud.databricks.com", "picked-profile"),
)
captured: dict = {}
def fake_configure_shared_state(workspace, profile=None, tools=None, force_login=False):
captured["workspace"] = workspace
captured["profile"] = profile
return {**MINIMAL_STATE, "workspace": workspace, "profile": profile}
monkeypatch.setattr(cli_mod, "configure_shared_state", fake_configure_shared_state)
monkeypatch.setattr(cli_mod, "save_state", lambda state: None)
monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: True)
monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: ["claude"])
monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *args, **kwargs: True)
monkeypatch.setattr(
cli_mod,
"configure_selected_tools",
lambda state, tools: {**state, "available_tools": tools},
)
monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None)
assert cli_mod.configure_workspace_command() == 0
assert captured["profile"] == "picked-profile"
rohita5l
left a comment
rohita5l
left a comment
There was a problem hiding this comment.
The "Enter a different URL" branch (ui.py:221-224) returns (url, None) unconditionally, which means downstream falls back to find_profile_name_for_host even when the URL the user typed happens to match a profile
they already have. Since prompt_for_workspace already has profiles in scope, we can resolve the match inline without breaking the Databricks-agnostic boundary:
while True:
raw_value = console.input(f" [bold]Workspace URL[/bold] {muted('›')} ").strip()
try:
url = normalize_workspace_url(raw_value)
except ValueError as exc:
print_err(str(exc))
continue
# If the typed URL matches a known profile, attach it so downstream
# --profile calls disambiguate. We only auto-attach when exactly one
# profile matches — multiple matches re-trigger the same ambiguity the
# picker is meant to resolve, so fall back to None and let the existing
# host-based resolution handle it (or, better, re-prompt).
matches = [name for host, name in (profiles or []) if host == url]
matched_profile = matches[0] if len(matches) == 1 else None
return url, matched_profile
This closes the gap for the common case (user types a URL they actually have a profile for). The ambiguous multi-match case stays on the existing find_profile_name_for_host path — fine for now, but worth a
follow-up to re-prompt with the matching rows.
| for host, profile_name in profiles | ||
| name_header = "Profile Name" | ||
| url_header = "Workspace URL" | ||
| name_width = max(len(name_header), *(len(name) for _, name in profiles)) |
There was a problem hiding this comment.
name_width = max(len(name_header), *(len(name) for _, name in profiles))
name_width is unbounded — whatever the longest profile name is becomes the column width. If someone has a ~/.databrickscfg entry like eng-staging-us-west-2-rohit-personal-sandbox (45 chars) plus a 40-char host
URL, that row is ~87 chars. On an 80-column terminal:
- questionary doesn't wrap rows cleanly — the cursor pointer and padding math assume a single line, so a wrapped row visually breaks the picker.
- The header (Profile Name … Workspace URL) and rows can drift out of alignment.
- Worst case the row runs off-screen and the URL is unreadable without resizing.
Concrete fix — clamp + ellipsize:
MAX_NAME_WIDTH = 40
def _truncate(s: str, width: int) -> str:
return s if len(s) <= width else s[: width - 1] + "…"
name_width = min(MAX_NAME_WIDTH, max(len(name_header), *(len(name) for _, name in profiles)))
header_title = f" {name_header.ljust(name_width)} {url_header}"
choices: list[questionary.Choice | questionary.Separator] = [questionary.Separator(header_title)]
for host, profile_name in profiles:
row_title = f"{_truncate(profile_name, name_width).ljust(name_width)} {host}"
choices.append(questionary.Choice(title=row_title, value=(host, profile_name)))
The value=(host, profile_name) tuple still carries the full untruncated name, so the selection passed to configure_shared_state is unchanged — only the display string is shortened.
…to profile - Cap profile-name column at 40 chars with ellipsis truncation so a long name can't push the URL column off-screen on an 80-col terminal. Value tuple still carries the full untruncated name. - In the "Enter a different URL" branch, if the typed URL matches exactly one known profile, return that profile so downstream `--profile` calls resolve unambiguously. Multi-match falls through to host-based lookup. - Add regression test locking in picker → configure_shared_state(profile=…) flow so a future refactor can't silently drop the profile. Co-authored-by: Isaac
|
@rohita5l I have made the changes to resolve the comments |
`configure_shared_state` already calls `find_profile_name_for_host` when profile is None (cli.py:161-162), which does the exact same host-to-profile lookup. The inline match was equivalent code at a different layer — removing it. Co-authored-by: Isaac
@rohita5l |
2 participants
Summary
Reworks the workspace picker shown during
ucode configure/ first launch so it mirrors the layout ofdatabricks auth profiles(without theValidcolumn) and so the exact profile the user selects is what laterdatabricksCLI calls use.Changes:
get_databricks_profiles()previously collapsed multiple profiles pointing at the same workspace down to the first match. It now returns one row per non-PAT profile, so each profile is selectable.prompt_for_workspace()renders aProfile Name/Workspace URLheader viaquestionary.Separatorand ljust-pads the profile names so columns line up.(host, profile_name); that tuple now travels throughconfigure_shared_state(profile=…)intostate.jsonand into every--profile-bearing CLI invocation. Previously, when a host had multiple profiles, downstream code fell back to "first profile for host" viafind_profile_name_for_host, now this is determined solely by the user's selectionBefore:
(only one row per host even when several profiles target it)
After:
Test plan
uv run pytest(546 passed / 28 skipped) — addsTestGetDatabricksProfiles(duplicate-host preservation, PAT skip, trailing-slash strip, non-zero-exit) andTestPromptForWorkspace(header rendering, duplicate-host preservation, URL normalization).uv run ruff check .clean.uv run ucode configure— visually verify column alignment and that all profiles for a duplicated host show up.uv run ucode status— confirmCLI profilematches the selection (and~/.ucode/state.jsonprofilekey matches).uv run ucode revert && uv run ucode claude— same picker fires on the launch path.This pull request and its description were written by Isaac.