Skip to content

Commit

Permalink
add get_current_workspace cloud client method (#15542)
Browse files Browse the repository at this point in the history
  • Loading branch information
zzstoatzz authored Sep 30, 2024
1 parent 7f8e852 commit ef71b83
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 4 deletions.
12 changes: 8 additions & 4 deletions src/prefect/cli/cloud/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -557,7 +557,12 @@ async def logout():
exit_with_success("Logged out from Prefect Cloud.")


@cloud_app.command()
@cloud_app.command(
deprecated=True,
deprecated_name="prefect cloud open",
deprecated_start_date="Oct 2024",
deprecated_help="Use `prefect dashboard open` to open the Prefect UI.",
)
async def open():
"""
Open the Prefect Cloud UI in the browser.
Expand All @@ -570,10 +575,9 @@ async def open():
"There is no current profile set - set one with `prefect profile create"
" <name>` and `prefect profile use <name>`."
)
async with get_cloud_client() as client:
current_workspace = await client.read_current_workspace()

current_workspace = get_current_workspace(
await prefect.client.cloud.get_cloud_client().read_workspaces()
)
if current_workspace is None:
exit_with_error(
"There is no current workspace set - set one with `prefect cloud workspace"
Expand Down
9 changes: 9 additions & 0 deletions src/prefect/client/cloud.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from prefect.exceptions import ObjectNotFound, PrefectException
from prefect.settings import (
PREFECT_API_KEY,
PREFECT_API_URL,
PREFECT_CLOUD_API_URL,
PREFECT_UNIT_TEST_MODE,
)
Expand Down Expand Up @@ -110,6 +111,14 @@ async def read_workspaces(self) -> List[Workspace]:
)
return workspaces

async def read_current_workspace(self) -> Workspace:
workspaces = await self.read_workspaces()
current_api_url = PREFECT_API_URL.value()
for workspace in workspaces:
if workspace.api_url() == current_api_url.rstrip("/"):
return workspace
raise ValueError("Current workspace not found")

async def read_worker_metadata(self) -> Dict[str, Any]:
response = await self.get(
f"{self.workspace_base_url}/collections/work_pool_types"
Expand Down
32 changes: 32 additions & 0 deletions tests/client/test_cloud_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,3 +98,35 @@ async def test_get_cloud_work_pool_types():
async with get_cloud_client() as client:
response = await client.read_worker_metadata()
assert response == mock_work_pool_types_response


async def test_read_current_workspace():
account_id = uuid.uuid4()
workspace_id = uuid.uuid4()
api_url = f"https://api.prefect.cloud/api/accounts/{account_id}/workspaces/{workspace_id}/"

with temporary_settings(updates={PREFECT_API_URL: api_url}):
with respx.mock(
assert_all_mocked=False, base_url=PREFECT_API_URL.value()
) as respx_mock:
respx_mock.get("https://api.prefect.cloud/api/me/workspaces").mock(
return_value=httpx.Response(
200,
json=[
{
"account_id": str(account_id),
"account_name": "Test Account",
"account_handle": "test-account",
"workspace_id": str(workspace_id),
"workspace_name": "Test Workspace",
"workspace_description": "Test workspace description",
"workspace_handle": "test-workspace",
}
],
)
)

async with get_cloud_client() as client:
workspace = await client.read_current_workspace()
assert workspace.workspace_id == workspace_id
assert workspace.account_id == account_id

0 comments on commit ef71b83

Please sign in to comment.