Coverage for src/commands/job_status.py: 23%

30 statements  

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

1""" 

2Command for checking status of Databricks job runs. 

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 Check status of a Databricks job run. 

17 

18 Args: 

19 client: DatabricksAPIClient instance for API calls 

20 **kwargs: Command parameters 

21 - run_id: The job run ID to check status for 

22 

23 Returns: 

24 CommandResult with job run status details if successful 

25 """ 

26 if not client: 

27 return CommandResult( 

28 False, 

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

30 ) 

31 

32 # Extract parameters 

33 run_id = kwargs.get("run_id") 

34 

35 try: 

36 # Get the job run status 

37 result = client.get_job_run_status(run_id) 

38 

39 if not result: 

40 return CommandResult(False, message=f"No job run found with ID: {run_id}") 

41 

42 # Extract key information 

43 run_info = { 

44 "job_id": result.get("job_id"), 

45 "run_id": result.get("run_id"), 

46 "run_name": result.get("run_name"), 

47 "state": result.get("state", {}).get("life_cycle_state"), 

48 "result_state": result.get("state", {}).get("result_state"), 

49 "start_time": result.get("start_time"), 

50 "setup_duration": result.get("setup_duration"), 

51 "execution_duration": result.get("execution_duration"), 

52 "cleanup_duration": result.get("cleanup_duration"), 

53 "creator_user_name": result.get("creator_user_name"), 

54 } 

55 

56 # Add task status information if available 

57 tasks = result.get("tasks", []) 

58 if tasks: 

59 task_statuses = [] 

60 for task in tasks: 

61 task_status = { 

62 "task_key": task.get("task_key"), 

63 "state": task.get("state", {}).get("life_cycle_state"), 

64 "result_state": task.get("state", {}).get("result_state"), 

65 "start_time": task.get("start_time"), 

66 "setup_duration": task.get("setup_duration"), 

67 "execution_duration": task.get("execution_duration"), 

68 "cleanup_duration": task.get("cleanup_duration"), 

69 } 

70 task_statuses.append(task_status) 

71 

72 run_info["tasks"] = task_statuses 

73 

74 # Create a user-friendly message 

75 state_msg = f"{run_info['state']}" 

76 if run_info.get("result_state"): 

77 state_msg += f" ({run_info['result_state']})" 

78 

79 message = f"Job run {run_id} is {state_msg}" 

80 

81 return CommandResult(True, data=run_info, message=message) 

82 except Exception as e: 

83 logging.error(f"Error getting job run status: {str(e)}") 

84 return CommandResult( 

85 False, message=f"Failed to get job run status: {str(e)}", error=e 

86 ) 

87 

88 

89DEFINITION = CommandDefinition( 

90 name="job-status", 

91 description="Check status of a Databricks job run.", 

92 handler=handle_command, 

93 parameters={ 

94 "run_id": { 

95 "type": "string", 

96 "description": "The job run ID to check status for.", 

97 } 

98 }, 

99 required_params=["run_id"], 

100 tui_aliases=["/job-status", "/job"], 

101 needs_api_client=True, 

102 visible_to_user=True, 

103 visible_to_agent=True, 

104 usage_hint="Usage: /job-status --run_id <run_id>", 

105 condensed_action="Checking job status", 

106)