Coverage for src/commands/warehouse_selection.py: 86%
29 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 handler for SQL warehouse selection.
4This module contains the handler for setting the active SQL warehouse
5for database operations.
6"""
8import logging
9from typing import Optional
11from src.clients.databricks import DatabricksAPIClient
12from src.command_registry import CommandDefinition
13from src.config import set_warehouse_id
14from .base import CommandResult
17def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult:
18 """
19 Set the active SQL warehouse.
21 Args:
22 client: API client instance
23 **kwargs: warehouse_id (str)
24 """
25 warehouse_id: str = kwargs.get("warehouse_id")
26 if not warehouse_id:
27 return CommandResult(False, message="warehouse_id parameter is required.")
29 try:
30 # Verify that the warehouse exists
31 if client:
32 try:
33 warehouse = client.get_warehouse(warehouse_id)
34 if not warehouse:
35 return CommandResult(
36 False,
37 message=f"Warehouse '{warehouse_id}' not found.",
38 )
39 warehouse_name = warehouse.get("name", "Unknown")
40 warehouse_state = warehouse.get("state", "Unknown")
41 except Exception:
42 # Set anyway if verification fails
43 set_warehouse_id(warehouse_id)
44 return CommandResult(
45 True,
46 message=f"Warning: Could not verify warehouse '{warehouse_id}'. Setting anyway.",
47 data={"warehouse_id": warehouse_id},
48 )
49 else:
50 # No client available, set without verification
51 set_warehouse_id(warehouse_id)
52 return CommandResult(
53 True,
54 message=f"Warning: No API client available to verify warehouse '{warehouse_id}'. Setting anyway.",
55 data={"warehouse_id": warehouse_id},
56 )
58 # Set the active warehouse
59 set_warehouse_id(warehouse_id)
60 return CommandResult(
61 True,
62 message=f"Active SQL warehouse is now set to '{warehouse_name}' (ID: {warehouse_id}, State: {warehouse_state}).",
63 data={
64 "warehouse_id": warehouse_id,
65 "warehouse_name": warehouse_name,
66 "state": warehouse_state,
67 },
68 )
69 except Exception as e:
70 logging.error(f"Failed to set warehouse '{warehouse_id}': {e}", exc_info=True)
71 return CommandResult(False, error=e, message=str(e))
74DEFINITION = CommandDefinition(
75 name="select-warehouse",
76 description="Set the active SQL warehouse for database operations",
77 handler=handle_command,
78 parameters={
79 "warehouse_id": {
80 "type": "string",
81 "description": "ID of the SQL warehouse to set as active",
82 }
83 },
84 required_params=["warehouse_id"],
85 tui_aliases=["/select-warehouse"],
86 needs_api_client=True,
87 visible_to_user=True,
88 visible_to_agent=True,
89 usage_hint="Usage: /select-warehouse --warehouse_id <warehouse_id>",
90)