Add yoshi get_job_definition

Author: ashangit

TESTING.md

@@ -30,7 +30,8 @@ task install
 task start-spark-bg && task start-mcp-bg && task start-inspector-bg
 
 # Opens http://localhost:6274 automatically in your browser
-# When done: task stop-all
+# When done
+task stop-all
 ```
 
 **Alternative** (if you prefer manual control):

pyproject.toml

@@ -27,7 +27,9 @@ dependencies = [
     "boto3~=1.34",
     "pydantic-settings>=2.9.1",
     "requests[socks]>=2.31.0",
-    "pysocks>=1.7.1"
+    "pysocks>=1.7.1",
+    "yoshi-client",
+    "dd-internal-authentication",
 ]
 
 [build-system]
@@ -103,3 +105,7 @@ name = "pypi-test"
 url = "https://test.pypi.org/simple/"
 publish-url = "https://test.pypi.org/legacy/"
 explicit = true
+
+[tool.uv.sources]
+yoshi-client = { url = "https://binaries.ddbuild.io/dd-source/python/yoshi_client-0.0.76422691-py3-none-any.whl" }
+dd-internal-authentication = { url = "https://binaries.ddbuild.io/dd-source/python/dd_internal_authentication-0.0.79707356-py3-none-any.whl" }

src/spark_history_mcp/common/__init__.py

@@ -0,0 +1,14 @@
+import logging
+
+from dd_internal_authentication.libs.py.dd_internal_authentication.dd_internal_authentication.client import (
+    JWTDDToolAuthClientTokenManager,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_token(datacenter: str, audience: str) -> str:
+    token = JWTDDToolAuthClientTokenManager.instance(
+        name=audience, datacenter=datacenter
+    ).get_token(audience)
+    return str(token)

src/spark_history_mcp/common/vault.py

@@ -0,0 +1,28 @@
+import os
+
+from yoshi_client.domains.data_eng_infra.shared.libs.py.yoshi_client import (
+    Configuration,
+    ApiClient,
+    JobApi,
+    Job,
+)
+
+from spark_history_mcp.common import vault
+
+
+class Yoshi:
+    AUDIENCE = "rapid-data-eng-infra"
+
+    def __init__(self, datacenter: str):
+        host =f"https://yoshi.{datacenter}"
+        if os.getenv("POD_NAME"):
+            host = f"https://yoshi.{self.AUDIENCE}.all-clusters.local-dc.fabric.dog:8443"
+        self.configuration = Configuration(
+            host=host,
+            access_token=vault.get_token(datacenter, self.AUDIENCE),
+        )
+
+    def get_job_definition(self, job_id: str) -> Job:
+        with ApiClient(self.configuration) as api_client:
+            job_api = JobApi(api_client)
+            return job_api.v2_get_job(job_id=job_id)

src/spark_history_mcp/common/yoshi.py

@@ -1,8 +1,11 @@
+import os
 from datetime import datetime
 import heapq
 import logging
 from typing import Any, Dict, List, Optional
 
+from yoshi_client.domains.data_eng_infra.shared.libs.py.yoshi_client import Job
+
 from spark_history_mcp.core.app import mcp
 from spark_history_mcp.models.mcp_types import (
     JobSummary,
@@ -18,11 +21,14 @@
     StageStatus,
     TaskMetricDistributions,
 )
+from ..common.yoshi import Yoshi
 
 from ..utils.utils import parallel_execute
 
 logger = logging.getLogger(__name__)
 
+DATACENTER = os.environ.get("DD_DATACENTER", "us1.staging.dog")
+
 
 def get_client_or_default(
     ctx, server_name: Optional[str] = None, app_id: Optional[str] = None
@@ -1296,39 +1302,72 @@ def get_resource_usage_timeline(
         },
     }
 
+
 @mcp.tool()
-def get_yoshi_job(job_id: str):
-    pass
+def get_job_definition(job_id: str) -> Job:
+    """Get job definition
+
+    Args:
+        job_id: Job identifier
+
+    Returns:
+        Job: Job definition
+    """
+    return Yoshi(DATACENTER).get_job_definition(job_id)
 
 
-#@mcp.tool()
-def get_job_logs(job_id: str, status: Optional[str] = None, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None):
+# @mcp.tool()
+def get_job_logs(
+    job_id: str,
+    status: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+):
     pass
 
 
-#@mcp.tool()
-def get_operator_logs(job_id: str, status: Optional[str] = None, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None):
+# @mcp.tool()
+def get_operator_logs(
+    job_id: str,
+    status: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+):
     # can be merged with get_job_logs???
     pass
 
-#@mcp.tool()
-def get_workflow_logs(workflow_id: str, status: Optional[str] = None, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None):
+
+# @mcp.tool()
+def get_workflow_logs(
+    workflow_id: str,
+    status: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+):
     pass
 
 
-#@mcp.tool()
-def get_admission_logs(job_id: str, status: Optional[str] = None, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None):
+# @mcp.tool()
+def get_admission_logs(
+    job_id: str,
+    status: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+):
     pass
 
-#@mcp.tool()
+
+# @mcp.tool()
 def get_oom_metrics_job(job_id: str):
     # oom_kill.oom_process.count
     pass
 
-#@mcp.tool()
+
+# @mcp.tool()
 def get_cpu_metrics_job(job_id: str):
     pass
 
-#@mcp.tool()
+
+# @mcp.tool()
 def get_memory_metrics_job(job_id: str):
     pass