Coverage for src/commands/list_warehouses.py: 81%

21 statements  

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

1""" 

2Command for listing all SQL warehouses in the Databricks workspace. 

3""" 

4 

5from typing import Optional, Any 

6from src.clients.databricks import DatabricksAPIClient 

7from src.commands.base import CommandResult 

8from src.command_registry import CommandDefinition 

9import logging 

10 

11 

12def handle_command( 

13 client: Optional[DatabricksAPIClient], **kwargs: Any 

14) -> CommandResult: 

15 """ 

16 Fetch and return a list of all SQL warehouses in the Databricks workspace. 

17 

18 Args: 

19 client: DatabricksAPIClient instance for API calls 

20 **kwargs: No additional parameters required 

21 

22 Returns: 

23 CommandResult with list of warehouses if successful 

24 """ 

25 if not client: 

26 return CommandResult( 

27 False, 

28 message="No Databricks client available. Please set up your workspace first.", 

29 ) 

30 

31 try: 

32 # Fetch the list of warehouses 

33 warehouses = client.list_warehouses() 

34 

35 if not warehouses: 

36 return CommandResult( 

37 True, message="No SQL warehouses found in this workspace." 

38 ) 

39 

40 # Format the warehouse information for display 

41 formatted_warehouses = [] 

42 for warehouse in warehouses: 

43 formatted_warehouse = { 

44 "id": warehouse.get("id"), 

45 "name": warehouse.get("name"), 

46 "size": warehouse.get("size"), 

47 "state": warehouse.get("state"), 

48 "creator_name": warehouse.get("creator_name"), 

49 "auto_stop_mins": warehouse.get("auto_stop_mins", "N/A"), 

50 } 

51 formatted_warehouses.append(formatted_warehouse) 

52 

53 return CommandResult( 

54 True, 

55 data={ 

56 "warehouses": formatted_warehouses, 

57 "total_count": len(formatted_warehouses), 

58 }, 

59 message=f"Found {len(formatted_warehouses)} SQL warehouse(s).", 

60 ) 

61 except Exception as e: 

62 logging.error(f"Error fetching warehouses: {str(e)}") 

63 return CommandResult( 

64 False, message=f"Failed to fetch warehouses: {str(e)}", error=e 

65 ) 

66 

67 

68DEFINITION = CommandDefinition( 

69 name="list-warehouses", 

70 description="Lists all SQL warehouses in the current Databricks workspace.", 

71 handler=handle_command, 

72 parameters={}, # No parameters needed 

73 required_params=[], 

74 tui_aliases=["/list-warehouses", "/warehouses"], 

75 needs_api_client=True, 

76 visible_to_user=True, 

77 visible_to_agent=True, 

78 agent_display="full", # Show full warehouse list in tables 

79 usage_hint="Usage: /list-warehouses", 

80)