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
« 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"""
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
12def handle_command(
13 client: Optional[DatabricksAPIClient], **kwargs: Any
14) -> CommandResult:
15 """
16 Check status of a Databricks job run.
18 Args:
19 client: DatabricksAPIClient instance for API calls
20 **kwargs: Command parameters
21 - run_id: The job run ID to check status for
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 )
32 # Extract parameters
33 run_id = kwargs.get("run_id")
35 try:
36 # Get the job run status
37 result = client.get_job_run_status(run_id)
39 if not result:
40 return CommandResult(False, message=f"No job run found with ID: {run_id}")
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 }
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)
72 run_info["tasks"] = task_statuses
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']})"
79 message = f"Job run {run_id} is {state_msg}"
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 )
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)