|
| 1 | +"""Job commands — status, download plan/profile, and analyze.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import json |
| 6 | +import os |
| 7 | +from typing import Any |
| 8 | + |
| 9 | +import click |
| 10 | + |
| 11 | +from cz_cli import output |
| 12 | +from cz_cli.cli_group import CLIGroup |
| 13 | +from cz_cli.guide_builder import examples_epilog |
| 14 | +from cz_cli.logger import log_operation |
| 15 | +from cz_cli.studio_client import StudioClient, configure_mcp_logging, studio_connection_kwargs |
| 16 | + |
| 17 | + |
| 18 | +def _new_client(ctx: click.Context) -> StudioClient: |
| 19 | + configure_mcp_logging("DEBUG" if ctx.obj.get("debug") else None) |
| 20 | + return StudioClient( |
| 21 | + profile=ctx.obj.get("profile"), |
| 22 | + jdbc_url=ctx.obj.get("jdbc_url"), |
| 23 | + **studio_connection_kwargs(ctx.obj), |
| 24 | + ) |
| 25 | + |
| 26 | + |
| 27 | +_DOWNLOAD_EXAMPLES = [ |
| 28 | + {"cmd": "cz-cli job download 2026012808001805432z9g3fx1sok", "desc": "Download job plan & profile to ./job_<id>/"}, |
| 29 | + {"cmd": "cz-cli job download 2026012808001805432z9g3fx1sok --output-dir /tmp/mydir", "desc": "Download to a custom directory"}, |
| 30 | +] |
| 31 | + |
| 32 | +_ANALYZE_EXAMPLES = [ |
| 33 | + {"cmd": "cz-cli job analyze 2026012808001805432z9g3fx1sok", "desc": "Quick performance analysis"}, |
| 34 | + {"cmd": "cz-cli job analyze 2026012808001805432z9g3fx1sok --mode detailed", "desc": "Detailed analysis"}, |
| 35 | + {"cmd": "cz-cli job analyze --path /tmp/mydir", "desc": "Analyze from local files (no job_id needed)"}, |
| 36 | +] |
| 37 | + |
| 38 | + |
| 39 | +@click.group("job", cls=CLIGroup) |
| 40 | +@click.pass_context |
| 41 | +def job_cmd(ctx: click.Context) -> None: |
| 42 | + """Job performance: download plan/profile and analyze.""" |
| 43 | + pass |
| 44 | + |
| 45 | + |
| 46 | +@job_cmd.command("download", epilog=examples_epilog(_DOWNLOAD_EXAMPLES)) |
| 47 | +@click.argument("job_id") |
| 48 | +@click.option("--workspace", "-w", "workspace_name", default=None, help="Workspace name (uses profile default if omitted).") |
| 49 | +@click.option("--output-dir", "output_dir", default=None, help="Output directory (default: ./job_<job_id>).") |
| 50 | +@click.pass_context |
| 51 | +def job_download(ctx: click.Context, job_id: str, workspace_name: str | None, output_dir: str | None) -> None: |
| 52 | + """Download job plan and profile JSON for a job ID.""" |
| 53 | + fmt: str = ctx.obj.get("format", "json") |
| 54 | + timer = output.Timer() |
| 55 | + |
| 56 | + try: |
| 57 | + client = _new_client(ctx) |
| 58 | + except Exception as exc: |
| 59 | + log_operation("job download", ok=False, error_code="CONNECTION_ERROR") |
| 60 | + output.error("CONNECTION_ERROR", str(exc), fmt=fmt) |
| 61 | + return |
| 62 | + |
| 63 | + args: dict[str, Any] = {"job_id": job_id} |
| 64 | + if workspace_name: |
| 65 | + args["workspace_name"] = workspace_name |
| 66 | + |
| 67 | + try: |
| 68 | + with timer: |
| 69 | + result = client.invoke("fetch_job_plan_profile_data", args) |
| 70 | + except Exception as exc: |
| 71 | + log_operation("job download", ok=False, error_code="DOWNLOAD_ERROR") |
| 72 | + output.error("DOWNLOAD_ERROR", str(exc), fmt=fmt) |
| 73 | + return |
| 74 | + |
| 75 | + payload = result.payload |
| 76 | + dest = output_dir or f"job_{job_id}" |
| 77 | + os.makedirs(dest, exist_ok=True) |
| 78 | + |
| 79 | + # Save job_profile.json |
| 80 | + profile_data = payload.get("job_profile_data") |
| 81 | + if profile_data: |
| 82 | + with open(os.path.join(dest, "job_profile.json"), "w", encoding="utf-8") as f: |
| 83 | + json.dump(profile_data, f, ensure_ascii=False, indent=2) |
| 84 | + |
| 85 | + # Save job_plan.json |
| 86 | + plan_data = payload.get("job_plan_data") |
| 87 | + if plan_data: |
| 88 | + with open(os.path.join(dest, "job_plan.json"), "w", encoding="utf-8") as f: |
| 89 | + json.dump(plan_data, f, ensure_ascii=False, indent=2) |
| 90 | + |
| 91 | + log_operation("job download", ok=True, time_ms=timer.elapsed_ms) |
| 92 | + output.success( |
| 93 | + { |
| 94 | + "job_id": job_id, |
| 95 | + "output_dir": os.path.abspath(dest), |
| 96 | + "job_plan_exists": plan_data is not None, |
| 97 | + "job_profile_exists": profile_data is not None, |
| 98 | + }, |
| 99 | + time_ms=timer.elapsed_ms, |
| 100 | + fmt=fmt, |
| 101 | + ) |
| 102 | + |
| 103 | + |
| 104 | +@job_cmd.command("analyze", epilog=examples_epilog(_ANALYZE_EXAMPLES)) |
| 105 | +@click.argument("job_id", required=False, default=None) |
| 106 | +@click.option("--workspace", "-w", "workspace_name", default=None, help="Workspace name (uses profile default if omitted).") |
| 107 | +@click.option("--mode", "analysis_mode", type=click.Choice(["quick", "detailed", "expert"]), default="quick", help="Analysis mode.") |
| 108 | +@click.option("--path", "local_path", default=None, help="Local folder with job_plan.json & job_profile.json (job_id optional when set).") |
| 109 | +@click.option("--state-table/--no-state-table", "enable_state_table", default=True, help="Enable state table optimization.") |
| 110 | +@click.option("--incremental/--no-incremental", "enable_incremental", default=False, help="Enable incremental algorithm.") |
| 111 | +@click.pass_context |
| 112 | +def job_analyze( |
| 113 | + ctx: click.Context, |
| 114 | + job_id: str | None, |
| 115 | + workspace_name: str | None, |
| 116 | + analysis_mode: str, |
| 117 | + local_path: str | None, |
| 118 | + enable_state_table: bool, |
| 119 | + enable_incremental: bool, |
| 120 | +) -> None: |
| 121 | + """Analyze job performance and provide optimization recommendations.""" |
| 122 | + fmt: str = ctx.obj.get("format", "json") |
| 123 | + |
| 124 | + if not job_id and not local_path: |
| 125 | + output.error("MISSING_ARGUMENT", "JOB_ID is required unless --path is provided.", fmt=fmt) |
| 126 | + return |
| 127 | + |
| 128 | + timer = output.Timer() |
| 129 | + |
| 130 | + try: |
| 131 | + client = _new_client(ctx) |
| 132 | + except Exception as exc: |
| 133 | + log_operation("job analyze", ok=False, error_code="CONNECTION_ERROR") |
| 134 | + output.error("CONNECTION_ERROR", str(exc), fmt=fmt) |
| 135 | + return |
| 136 | + |
| 137 | + args: dict[str, Any] = { |
| 138 | + "analysis_mode": analysis_mode, |
| 139 | + "enable_state_table": enable_state_table, |
| 140 | + "enable_incremental_algorithm": enable_incremental, |
| 141 | + } |
| 142 | + if job_id: |
| 143 | + args["job_id"] = job_id |
| 144 | + if workspace_name: |
| 145 | + args["workspace_name"] = workspace_name |
| 146 | + if local_path: |
| 147 | + args["path"] = local_path |
| 148 | + |
| 149 | + try: |
| 150 | + with timer: |
| 151 | + result = client.invoke("fetch_job_performance_data", args) |
| 152 | + except Exception as exc: |
| 153 | + log_operation("job analyze", ok=False, error_code="ANALYSIS_ERROR") |
| 154 | + output.error("ANALYSIS_ERROR", str(exc), fmt=fmt) |
| 155 | + return |
| 156 | + |
| 157 | + log_operation("job analyze", ok=True, time_ms=timer.elapsed_ms) |
| 158 | + output.success(result.payload, time_ms=timer.elapsed_ms, fmt=fmt) |
| 159 | + |
| 160 | + |
| 161 | +@job_cmd.command("status") |
| 162 | +@click.argument("job_id") |
| 163 | +@click.pass_context |
| 164 | +def job_status(ctx: click.Context, job_id: str) -> None: |
| 165 | + """Check status / summary of a SQL job.""" |
| 166 | + from cz_cli.connection import get_connection |
| 167 | + from cz_cli.connection_ctx import connection_kwargs_from_ctx |
| 168 | + |
| 169 | + fmt: str = ctx.obj.get("format", "json") |
| 170 | + profile: str | None = ctx.obj.get("profile") |
| 171 | + jdbc_url: str | None = ctx.obj.get("jdbc_url") |
| 172 | + |
| 173 | + try: |
| 174 | + conn = get_connection(jdbc_url=jdbc_url, profile=profile, **connection_kwargs_from_ctx(ctx)) |
| 175 | + except Exception as exc: |
| 176 | + log_operation("job status", ok=False, error_code="CONNECTION_ERROR") |
| 177 | + output.error("CONNECTION_ERROR", str(exc), fmt=fmt) |
| 178 | + return |
| 179 | + |
| 180 | + try: |
| 181 | + summary = conn.get_job_summary(job_id) |
| 182 | + log_operation("job status", ok=True) |
| 183 | + output.success(summary, fmt=fmt) |
| 184 | + except Exception as exc: |
| 185 | + log_operation("job status", ok=False, error_code="JOB_STATUS_ERROR") |
| 186 | + output.error("JOB_STATUS_ERROR", str(exc), fmt=fmt) |
| 187 | + finally: |
| 188 | + conn.close() |
0 commit comments