Coverage for src/commands/jobs.py: 93%
46 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"""
2Jobs command handlers for Chuck.
4This module contains job submission and status checking commands.
5"""
7import logging
8from typing import Optional
10from src.clients.databricks import DatabricksAPIClient
11from src.commands.base import CommandResult
12from src.command_registry import CommandDefinition
15def handle_launch_job(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult:
16 """Submits a one-time Databricks job run.
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))
54def handle_job_status(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult:
55 """Query Databricks for the status of a one-time run.
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))
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)
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)
125# Combined definition for module
126DEFINITION = [LAUNCH_JOB_DEFINITION, JOB_STATUS_DEFINITION]