Coverage for src/commands/workspace_selection.py: 100%

24 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 22:56 -0700

1""" 

2Command handler for workspace selection. 

3 

4This module contains the handler for setting the workspace URL 

5to connect to a Databricks workspace. 

6""" 

7 

8import logging 

9from typing import Optional 

10 

11from src.clients.databricks import DatabricksAPIClient 

12from src.command_registry import CommandDefinition 

13from src.config import set_workspace_url 

14from .base import CommandResult 

15 

16 

17def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult: 

18 """ 

19 Set the workspace URL. 

20 

21 Args: 

22 client: API client instance (not used by this handler) 

23 **kwargs: workspace_url (str) 

24 """ 

25 workspace_url: str = kwargs.get("workspace_url") 

26 if not workspace_url: 

27 return CommandResult(False, message="workspace_url parameter is required.") 

28 

29 try: 

30 from src.databricks.url_utils import ( 

31 validate_workspace_url, 

32 normalize_workspace_url, 

33 format_workspace_url_for_display, 

34 detect_cloud_provider, 

35 ) 

36 

37 is_valid, error_message = validate_workspace_url(workspace_url) 

38 if not is_valid: 

39 return CommandResult(False, message=f"Error: {error_message}") 

40 

41 normalized_url = normalize_workspace_url(workspace_url) 

42 cloud_provider = detect_cloud_provider(workspace_url) 

43 display_url = format_workspace_url_for_display(normalized_url, cloud_provider) 

44 set_workspace_url(workspace_url) 

45 

46 return CommandResult( 

47 True, 

48 message=f"Workspace URL is now set to '{display_url}'. Restart may be needed.", 

49 data={ 

50 "workspace_url": workspace_url, 

51 "display_url": display_url, 

52 "cloud_provider": cloud_provider, 

53 "requires_restart": True, 

54 }, 

55 ) 

56 except Exception as e: 

57 logging.error(f"Failed to set workspace URL: {e}", exc_info=True) 

58 return CommandResult(False, error=e, message=str(e)) 

59 

60 

61DEFINITION = CommandDefinition( 

62 name="select-workspace", 

63 description="Set the Databricks workspace URL", 

64 handler=handle_command, 

65 parameters={ 

66 "workspace_url": { 

67 "type": "string", 

68 "description": "URL of the Databricks workspace (e.g., my-workspace.cloud.databricks.com)", 

69 } 

70 }, 

71 required_params=["workspace_url"], 

72 tui_aliases=["/select-workspace"], 

73 visible_to_user=True, 

74 visible_to_agent=False, # Agent doesn't need to select workspace 

75)