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

29 statements  

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

1""" 

2Command handler for SQL warehouse selection. 

3 

4This module contains the handler for setting the active SQL warehouse 

5for database operations. 

6""" 

7 

8import logging 

9from typing import Optional 

10 

11from ..clients.databricks import DatabricksAPIClient 

12from ..command_registry import CommandDefinition 

13from ..config import set_warehouse_id 

14from .base import CommandResult 

15 

16 

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

18 """ 

19 Set the active SQL warehouse. 

20 

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.") 

28 

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 ) 

57 

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)) 

72 

73 

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)