Coverage for src/commands/create_warehouse.py: 28%
25 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
« 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"""
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
12def handle_command(
13 client: Optional[DatabricksAPIClient], **kwargs: Any
14) -> CommandResult:
15 """
16 Create a new SQL warehouse in the Databricks workspace.
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)
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 )
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")
43 # Prepare warehouse configuration
44 warehouse_config = {
45 "name": name,
46 "size": size,
47 "auto_stop_mins": auto_stop_mins,
48 }
50 # Add optional cluster configuration if provided
51 if min_num_clusters is not None:
52 warehouse_config["min_num_clusters"] = min_num_clusters
54 if max_num_clusters is not None:
55 warehouse_config["max_num_clusters"] = max_num_clusters
57 try:
58 # Create the warehouse
59 result = client.create_warehouse(warehouse_config)
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 )
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)