Support non ddtool auth

Author: ashangit

src/spark_history_mcp/common/vault.py

@@ -1,14 +1,26 @@
-import logging
+import os
 
-from dd_internal_authentication.libs.py.dd_internal_authentication.dd_internal_authentication.client import (
-    JWTDDToolAuthClientTokenManager,
-)
 
-logger = logging.getLogger(__name__)
+class JWT:
+    def __init__(self, audience: str, datacenter: str):
+        self.audience = audience
 
+        from dd_internal_authentication.libs.py.dd_internal_authentication.dd_internal_authentication.client import (
+            JWTDDToolAuthClientTokenManager,
+            JWTInternalServiceAuthClientTokenManager,
+        )
 
-def get_token(datacenter: str, audience: str) -> str:
-    token = JWTDDToolAuthClientTokenManager.instance(
-        name=audience, datacenter=datacenter
-    ).get_token(audience)
-    return str(token)
+        if os.getenv("POD_NAME"):
+            token_manager_class = JWTInternalServiceAuthClientTokenManager
+        else:
+            token_manager_class = JWTDDToolAuthClientTokenManager
+
+        self.token_manager = token_manager_class.instance(
+            name=self.audience, datacenter=datacenter
+        )
+
+    def get_token(self) -> str:
+        try:
+            return str(self.token_manager.get_token(self.audience))
+        except Exception as e:
+            raise RuntimeError(f"Failed to get authentication token: {str(e)}")

src/spark_history_mcp/common/yoshi.py

@@ -8,6 +8,7 @@
 )
 
 from spark_history_mcp.common import vault
+from spark_history_mcp.common.vault import JWT
 
 
 class Yoshi:
@@ -19,7 +20,7 @@ def __init__(self, datacenter: str):
             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),
+            access_token=JWT(audience=self.AUDIENCE,datacenter=datacenter).get_token(),
         )
 
     def get_job_definition(self, job_id: str) -> Job: