Repository: PrefectHQ/prefect
Pull request: Improve interactive workspace selection
closes https://github.com/PrefectHQ/prefect/issues/12381

This PR updates the interactive `prefect cloud workspace set` (without a handle specified) to prompt for account selection first, only when the user has more than 10 workspaces _and_ more than one account. This makes the experience better for those with lots of workspaces / account

**NOTE** the below is in slow motion 🙂 (its faster IRL)


![better-ws-set](https://github.com/PrefectHQ/prefect/assets/31014960/87a723b8-d8b7-45f4-b849-9038af2a5f2e)


Problem statement:
improve interactive `prefect cloud workspace set`
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar request and didn't find it.
- [X] I searched the Prefect documentation for this feature.

### Prefect Version

2.x

### Describe the current behavior

if you have a lot of workspaces, you have to use the arrow key a lot

### Describe the proposed behavior

select account, then workspace, some intelligent sorting

### Example Use

_No response_

### Additional context

_No response_

Hints from issue discussion:

Commit message: Merge branch 'main' into better-interactive-workspace-set

Patch to review:
diff --git a/src/prefect/cli/cloud/__init__.py b/src/prefect/cli/cloud/__init__.py
index 2eaafafed26d..71e913e687d7 100644
--- a/src/prefect/cli/cloud/__init__.py
+++ b/src/prefect/cli/cloud/__init__.py
@@ -594,23 +594,44 @@ async def set(
                 "Unable to authenticate. Please ensure your credentials are correct."
             )
 
-    if workspace_handle:
-        # Search for the given workspace
-        for workspace in workspaces:
-            if workspace.handle == workspace_handle:
-                break
+        if workspace_handle:
+            # Search for the given workspace
+            for workspace in workspaces:
+                if workspace.handle == workspace_handle:
+                    break
+            else:
+                exit_with_error(f"Workspace {workspace_handle!r} not found.")
         else:
-            exit_with_error(f"Workspace {workspace_handle!r} not found.")
-    else:
-        workspace = prompt_select_from_list(
-            app.console,
-            "Which workspace would you like to use?",
-            [(workspace, workspace.handle) for workspace in workspaces],
-        )
+            # don't force two choices if there's not that many workspaces
+            if len(workspaces) > 10:
+                accounts = await client.read_accounts()
+
+                if len(accounts) == 1:
+                    account = accounts[0]
+                    workspaces = await client.read_workspaces(
+                        account_id=account["account_id"]
+                    )
+                else:
+                    account = prompt_select_from_list(
+                        app.console,
+                        "Which account would you like to use?",
+                        [(account, account["account_handle"]) for account in accounts],
+                    )
+                    workspaces = await client.read_workspaces(
+                        account_id=account["account_id"]
+                    )
+
+            if not workspaces:
+                exit_with_error("No workspaces found in the selected account.")
 
-    profile = update_current_profile({PREFECT_API_URL: workspace.api_url()})
+            workspace = prompt_select_from_list(
+                app.console,
+                "Which workspace would you like to use?",
+                [(workspace, workspace.handle) for workspace in workspaces],
+            )
 
-    exit_with_success(
-        f"Successfully set workspace to {workspace.handle!r} in profile"
-        f" {profile.name!r}."
-    )
+        profile = update_current_profile({PREFECT_API_URL: workspace.api_url()})
+        exit_with_success(
+            f"Successfully set workspace to {workspace.handle!r} in profile"
+            f" {profile.name!r}."
+        )
diff --git a/src/prefect/client/cloud.py b/src/prefect/client/cloud.py
index 479edf5b7c4a..20e874d1fec2 100644
--- a/src/prefect/client/cloud.py
+++ b/src/prefect/client/cloud.py
@@ -1,5 +1,6 @@
 import re
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
 
 import anyio
 import httpx
@@ -84,8 +85,20 @@ async def api_healthcheck(self):
         with anyio.fail_after(10):
             await self.read_workspaces()
 
-    async def read_workspaces(self) -> List[Workspace]:
-        return pydantic.parse_obj_as(List[Workspace], await self.get("/me/workspaces"))
+    async def read_workspaces(
+        self, account_id: Optional[Union[UUID, str]] = None
+    ) -> List[Workspace]:
+        workspaces = pydantic.parse_obj_as(
+            List[Workspace], await self.get("/me/workspaces")
+        )
+        if account_id:
+            if isinstance(account_id, str):
+                account_id = UUID(account_id)
+            return [w for w in workspaces if w.account_id == account_id]
+        return workspaces
+
+    async def read_accounts(self) -> List[Dict[str, Any]]:
+        return await self.get("/me/accounts")
 
     async def read_worker_metadata(self) -> Dict[str, Any]:
         configured_url = prefect.settings.PREFECT_API_URL.value()
diff --git a/tests/cli/test_cloud.py b/tests/cli/test_cloud.py
index 34d85037f39e..52713142121f 100644
--- a/tests/cli/test_cloud.py
+++ b/tests/cli/test_cloud.py
@@ -971,6 +971,121 @@ def test_set_workspace_updates_profile(respx_mock):
     }
 
 
+@pytest.mark.usefixtures("interactive_console")
+def test_set_workspace_with_account_selection(respx_mock):
+    foo_workspace = gen_test_workspace(account_handle="test1", workspace_handle="foo")
+    bar_workspace = gen_test_workspace(account_handle="test2", workspace_handle="bar")
+
+    respx_mock.get(PREFECT_CLOUD_API_URL.value() + "/me/workspaces").mock(
+        return_value=httpx.Response(
+            status.HTTP_200_OK,
+            json=[
+                foo_workspace.dict(json_compatible=True),
+                bar_workspace.dict(json_compatible=True),
+            ],
+        )
+    )
+
+    respx_mock.get(PREFECT_CLOUD_API_URL.value() + "/me/accounts").mock(
+        return_value=httpx.Response(
+            status.HTTP_200_OK,
+            json=[
+                {"account_handle": "test1", "account_id": "account1"},
+                {"account_handle": "test2", "account_id": "account2"},
+            ],
+        )
+    )
+
+    respx_mock.get(
+        PREFECT_CLOUD_API_URL.value() + "/me/workspaces?account_id=account2"
+    ).mock(
+        return_value=httpx.Response(
+            status.HTTP_200_OK,
+            json=[bar_workspace.dict(json_compatible=True)],
+        )
+    )
+
+    cloud_profile = "cloud-foo"
+    save_profiles(
+        ProfilesCollection(
+            [
+                Profile(
+                    name=cloud_profile,
+                    settings={
+                        PREFECT_API_URL: foo_workspace.api_url(),
+                        PREFECT_API_KEY: "fake-key",
+                    },
+                )
+            ],
+            active=None,
+        )
+    )
+
+    with use_profile(cloud_profile):
+        invoke_and_assert(
+            ["cloud", "workspace", "set"],
+            expected_code=0,
+            user_input=readchar.key.DOWN + readchar.key.ENTER + readchar.key.ENTER,
+            expected_output_contains=[
+                f"Successfully set workspace to {bar_workspace.handle!r} in profile {cloud_profile!r}.",
+            ],
+        )
+
+    profiles = load_profiles()
+    assert profiles[cloud_profile].settings == {
+        PREFECT_API_URL: bar_workspace.api_url(),
+        PREFECT_API_KEY: "fake-key",
+    }
+
+
+@pytest.mark.usefixtures("interactive_console")
+def test_set_workspace_with_less_than_10_workspaces(respx_mock):
+    foo_workspace = gen_test_workspace(account_handle="test1", workspace_handle="foo")
+    bar_workspace = gen_test_workspace(account_handle="test2", workspace_handle="bar")
+
+    respx_mock.get(PREFECT_CLOUD_API_URL.value() + "/me/workspaces").mock(
+        return_value=httpx.Response(
+            status.HTTP_200_OK,
+            json=[
+                foo_workspace.dict(json_compatible=True),
+                bar_workspace.dict(json_compatible=True),
+            ],
+        )
+    )
+
+    cloud_profile = "cloud-foo"
+    save_profiles(
+        ProfilesCollection(
+            [
+                Profile(
+                    name=cloud_profile,
+                    settings={
+                        PREFECT_API_URL: foo_workspace.api_url(),
+                        PREFECT_API_KEY: "fake-key",
+                    },
+                )
+            ],
+            active=None,
+        )
+    )
+
+    with use_profile(cloud_profile):
+        invoke_and_assert(
+            ["cloud", "workspace", "set"],
+            expected_code=0,
+            user_input=readchar.key.DOWN + readchar.key.ENTER,
+            expected_output_contains=[
+                f"Successfully set workspace to {bar_workspace.handle!r} in profile {cloud_profile!r}.",
+            ],
+        )
+
+    profiles = load_profiles()
+    assert profiles[cloud_profile].settings == {
+        PREFECT_API_URL: bar_workspace.api_url(),
+        PREFECT_API_KEY: "fake-key",
+    }
+
+
 def test_cannot_get_webhook_if_you_are_not_logged_in():
     cloud_profile = "cloud-foo"
     save_profiles(


Code review comments:
diff_hunk: @@ -594,23 +594,44 @@ async def set(
                 "Unable to authenticate. Please ensure your credentials are correct."
             )
 
-    if workspace_handle:
-        # Search for the given workspace
-        for workspace in workspaces:
-            if workspace.handle == workspace_handle:
-                break
+        if workspace_handle:
+            # Search for the given workspace
+            for workspace in workspaces:
+                if workspace.handle == workspace_handle:
+                    break
+            else:
+                exit_with_error(f"Workspace {workspace_handle!r} not found.")
         else:
-            exit_with_error(f"Workspace {workspace_handle!r} not found.")
-    else:
-        workspace = prompt_select_from_list(
-            app.console,
-            "Which workspace would you like to use?",
-            [(workspace, workspace.handle) for workspace in workspaces],
-        )
+            # don't force two choices if there's not that many workspaces
+            if len(workspaces) > 10:
+                accounts = await client.read_accounts()
path: src/prefect/cli/cloud/__init__.py
text: @user1:
You might be able to get away with not making the extra 2 API calls. `workspaces = await client.read_workspaces()` should have all the information needed with a little bit of bucketing.

little code golfy for sure, but only mentioning in case the api round trip adds latency