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

46 statements  

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

1""" 

2Jobs command handlers for Chuck. 

3 

4This module contains job submission and status checking commands. 

5""" 

6 

7import logging 

8from typing import Optional 

9 

10from ..clients.databricks import DatabricksAPIClient 

11from ..commands.base import CommandResult 

12from ..command_registry import CommandDefinition 

13 

14 

15def handle_launch_job(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult: 

16 """Submits a one-time Databricks job run. 

17 

18 Args: 

19 client: API client instance 

20 **kwargs: config_path (str), init_script_path (str), run_name (str, optional) 

21 """ 

22 config_path: str = kwargs.get("config_path") 

23 init_script_path: str = kwargs.get("init_script_path") 

24 run_name: Optional[str] = kwargs.get("run_name") 

25 if not config_path or not init_script_path: 

26 return CommandResult( 

27 False, message="config_path and init_script_path are required." 

28 ) 

29 if not client: 

30 return CommandResult(False, message="Client required to launch job.") 

31 try: 

32 run_data = client.submit_job_run( 

33 config_path=config_path, 

34 init_script_path=init_script_path, 

35 run_name=run_name, 

36 ) 

37 run_id = run_data.get("run_id") 

38 if run_id: 

39 return CommandResult( 

40 True, 

41 data={"run_id": str(run_id)}, 

42 message=f"Job submitted. Run ID: {run_id}", 

43 ) 

44 else: 

45 logging.error(f"Failed to launch job, no run_id: {run_data}") 

46 return CommandResult( 

47 False, message="Failed to submit job (no run_id).", data=run_data 

48 ) 

49 except Exception as e: 

50 logging.error(f"Failed to submit job: {e}", exc_info=True) 

51 return CommandResult(False, error=e, message=str(e)) 

52 

53 

54def handle_job_status(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult: 

55 """Query Databricks for the status of a one-time run. 

56 

57 Args: 

58 client: API client instance 

59 **kwargs: run_id (str) 

60 """ 

61 run_id_str: str = kwargs.get("run_id") 

62 if not run_id_str: 

63 return CommandResult(False, message="run_id parameter is required.") 

64 if not client: 

65 return CommandResult(False, message="Client required to get job status.") 

66 try: 

67 data = client.get_job_run_status(run_id_str) 

68 status = data.get("status", data.get("state", {})) 

69 life_cycle_state = status.get("life_cycle_state", status.get("state")) 

70 result_state = status.get( 

71 "result_state", status.get("termination_details", {}).get("code") 

72 ) 

73 state_message = status.get( 

74 "state_message", status.get("termination_details", {}).get("message") 

75 ) 

76 run_page_url = data.get("run_page_url") 

77 msg = f"Run {run_id_str}: Status: {life_cycle_state or 'N/A'}, Result: {result_state or 'N/A'}, Message: {state_message or ''}" 

78 if run_page_url: 

79 msg += f" URL: {run_page_url}" 

80 return CommandResult(True, data=data, message=msg) 

81 except Exception as e: 

82 logging.error( 

83 f"Failed to get job status for run '{run_id_str}': {e}", exc_info=True 

84 ) 

85 return CommandResult(False, error=e, message=str(e)) 

86 

87 

88# Command definitions 

89LAUNCH_JOB_DEFINITION = CommandDefinition( 

90 name="launch-job", 

91 description="Launch a Databricks job using a config file", 

92 usage_hint="launch_job --config_path=/path/to/config.json --init_script_path=/init/script.sh", 

93 parameters={ 

94 "config_path": { 

95 "type": "string", 

96 "description": "Path to the job configuration file", 

97 }, 

98 "init_script_path": { 

99 "type": "string", 

100 "description": "Path to the init script", 

101 }, 

102 "run_name": {"type": "string", "description": "Optional name for the job run"}, 

103 }, 

104 required_params=["config_path", "init_script_path"], 

105 handler=handle_launch_job, 

106 needs_api_client=True, 

107 visible_to_agent=True, 

108 tui_aliases=["launch-job"], 

109) 

110 

111JOB_STATUS_DEFINITION = CommandDefinition( 

112 name="job-status", 

113 description="Get the status of a Databricks job run", 

114 usage_hint="job_status --run_id=123456", 

115 parameters={ 

116 "run_id": {"type": "string", "description": "ID of the job run to check"}, 

117 }, 

118 required_params=["run_id"], 

119 handler=handle_job_status, 

120 needs_api_client=True, 

121 visible_to_agent=True, 

122 tui_aliases=["job-status"], 

123) 

124 

125# Combined definition for module 

126DEFINITION = [LAUNCH_JOB_DEFINITION, JOB_STATUS_DEFINITION]