Coverage for src/chuck_data/commands/create_warehouse.py: 0%

25 statements  

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

1""" 

2Command for creating a new SQL warehouse in Databricks. 

3""" 

4 

5from typing import Optional, Any 

6from ..clients.databricks import DatabricksAPIClient 

7from ..commands.base import CommandResult 

8from ..command_registry import CommandDefinition 

9import logging 

10 

11 

12def handle_command( 

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

14) -> CommandResult: 

15 """ 

16 Create a new SQL warehouse in the Databricks workspace. 

17 

18 Args: 

19 client: DatabricksAPIClient instance for API calls 

20 **kwargs: Command parameters 

21 - name: Name for the new warehouse 

22 - size: Size of the warehouse (e.g., "2X-Small", "Small", "Medium", "Large") 

23 - auto_stop_mins: Minutes of inactivity before the warehouse stops automatically 

24 - min_num_clusters: Min number of clusters for serverless warehouses (optional) 

25 - max_num_clusters: Max number of clusters for serverless warehouses (optional) 

26 

27 Returns: 

28 CommandResult with created warehouse details if successful 

29 """ 

30 if not client: 

31 return CommandResult( 

32 False, 

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

34 ) 

35 

36 # Extract parameters 

37 name = kwargs.get("name") 

38 size = kwargs.get("size") 

39 auto_stop_mins = kwargs.get("auto_stop_mins") 

40 min_num_clusters = kwargs.get("min_num_clusters") 

41 max_num_clusters = kwargs.get("max_num_clusters") 

42 

43 # Prepare warehouse configuration 

44 warehouse_config = { 

45 "name": name, 

46 "size": size, 

47 "auto_stop_mins": auto_stop_mins, 

48 } 

49 

50 # Add optional cluster configuration if provided 

51 if min_num_clusters is not None: 

52 warehouse_config["min_num_clusters"] = min_num_clusters 

53 

54 if max_num_clusters is not None: 

55 warehouse_config["max_num_clusters"] = max_num_clusters 

56 

57 try: 

58 # Create the warehouse 

59 result = client.create_warehouse(warehouse_config) 

60 

61 return CommandResult( 

62 True, 

63 data=result, 

64 message=f"Successfully created SQL warehouse '{name}' (ID: {result.get('id')}).", 

65 ) 

66 except Exception as e: 

67 logging.error(f"Error creating warehouse: {str(e)}") 

68 return CommandResult( 

69 False, message=f"Failed to create warehouse: {str(e)}", error=e 

70 ) 

71 

72 

73DEFINITION = CommandDefinition( 

74 name="create-warehouse", 

75 description="Create a new SQL warehouse in the Databricks workspace.", 

76 handler=handle_command, 

77 parameters={ 

78 "name": {"type": "string", "description": "Name for the new warehouse."}, 

79 "size": { 

80 "type": "string", 

81 "description": "Size of the warehouse (e.g., '2X-Small', 'Small', 'Medium', 'Large').", 

82 "default": "Small", 

83 }, 

84 "auto_stop_mins": { 

85 "type": "integer", 

86 "description": "Minutes of inactivity before the warehouse stops automatically.", 

87 "default": 120, 

88 }, 

89 "min_num_clusters": { 

90 "type": "integer", 

91 "description": "Min number of clusters for serverless warehouses. (Optional)", 

92 }, 

93 "max_num_clusters": { 

94 "type": "integer", 

95 "description": "Max number of clusters for serverless warehouses. (Optional)", 

96 }, 

97 }, 

98 required_params=["name"], 

99 tui_aliases=["/create-warehouse"], 

100 needs_api_client=True, 

101 visible_to_user=True, 

102 visible_to_agent=True, 

103 usage_hint="Usage: /create-warehouse --name <name> --size <size> --auto_stop_mins <minutes>", 

104)