add s3 client class to interact with spark-history-server bucket

Author: benkermani

mcp-server.sh

@@ -2,4 +2,4 @@ export VAULT_ADDR=https://vault.us1.staging.dog
 vault login -method=oidc
 export VAULT_TOKEN=$(vault print token)
 
-uv run -m spark_history_mcp.core.main
+aws-vault exec sso-staging-engineering -- uv run -m spark_history_mcp.core.main

src/spark_history_mcp/common/decorators.py

@@ -0,0 +1,19 @@
+import time
+
+
+def backoff_retry(delay=2, retries=3):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            current_retry = 0
+            current_delay = delay
+            while current_retry < retries:
+                try:
+                    return func(*args, **kwargs)
+                except Exception as e:
+                    current_retry += 1
+                    if current_retry >= retries:
+                        raise e
+                    time.sleep(current_delay)
+                    current_delay *= 2
+        return wrapper
+    return decorator

src/spark_history_mcp/common/s3_client.py

@@ -0,0 +1,74 @@
+import os
+
+import boto3
+import requests
+
+from src.spark_history_mcp.common.decorators import backoff_retry
+from src.spark_history_mcp.common.variable import POD_NAME
+
+
+class S3Client:
+    def __init__(self, datacenter: str):
+        self.client = boto3.resource("s3")
+        self.bucket_name = f"dd-spark-history-server-{datacenter.replace(".", "-")}" # e.g dd-spark-history-server-us1-staging-dog
+        self.dst_prefix = "indexed_spark_logs/"
+
+        shs_url_prefix =f"https://spark-history-server.{datacenter}"
+        if POD_NAME:
+            shs_url_prefix = "https://spark-history-server.spark.all-clusters.local-dc.fabric.dog:5554"
+        self.shs_url_prefix = shs_url_prefix
+
+    def list_contents_by_prefix(self, prefix, bucket):
+        b = self.client.Bucket(bucket)
+        keys = [obj.key for obj in b.objects.filter(Prefix=prefix)]
+
+        return keys
+
+    def is_spark_event_logs_already_indexed(self, spark_app_id: str) -> bool:
+        prefix = self.dst_prefix + str(spark_app_id)
+        if self.list_contents_by_prefix(prefix, self.bucket_name):
+            return True
+
+        return False
+
+    @backoff_retry(retries=5, delay=2)
+    def poll_spark_history_server(self, spark_app_id: str) -> Exception | None:
+        print("entered function")
+        full_url = f"{self.shs_url_prefix}/history/{spark_app_id}/jobs/"
+        try:
+           resp = requests.get(full_url, timeout=3)
+        except requests.exceptions.Timeout:
+            raise Exception(f"Spark History Server request timed out: {full_url}", 408)
+        except requests.exceptions.ConnectionError:
+            raise Exception("Spark History Server unavailable, please try again shortly", 503)
+
+        if resp.status_code == 404:
+            raise Exception(f"Spark History Server didn't finish parsing event logs: {full_url}", 404)
+
+        return None
+
+    def copy_spark_events_logs(self, spark_app_id: str) -> Exception | None:
+        # get spark events logs file to copy/index
+        src_prefix = f"spark_logs/{spark_app_id}"
+        base_logs = self.list_contents_by_prefix(src_prefix, self.bucket_name)
+        if not base_logs:
+            raise Exception(f"Logs for {spark_app_id} not found. Is the job older than one month?", 404)
+
+        # copy log file to new prefix
+        src_key = base_logs[0]
+        dst_key = self.dst_prefix + os.path.basename(src_key)
+        copy_source = {
+            'Bucket': self.bucket_name,
+            'Key': src_key
+        }
+
+        bucket = self.client.Bucket(self.bucket_name)
+        bucket.copy(copy_source, dst_key)
+
+        # poll SHS until event logs are parsed
+        try:
+            self.poll_spark_history_server(spark_app_id)
+        except Exception as e:
+            raise Exception(f"Error polling Spark History Server: {e}") from e
+
+        return None

src/spark_history_mcp/tools/tools.py

@@ -23,6 +23,7 @@
 )
 from ..common.datadog import Datadog, LogDD
 from ..common.yoshi import Yoshi
+from ..common.s3_client import S3Client
 
 from ..utils.utils import parallel_execute
 
@@ -156,6 +157,14 @@ def get_application(app_id: str, server: Optional[str] = None) -> ApplicationInf
     ctx = mcp.get_context()
     client = get_client_or_default(ctx, server, app_id)
 
+    # Index spark event logs if missing
+    s3_client = S3Client(datacenter=DATACENTER)
+    if not s3_client.is_spark_event_logs_already_indexed(app_id):
+        try:
+            s3_client.copy_spark_events_logs(app_id)
+        except Exception as e:
+            raise Exception(f"Failed to copy events logs for app_id {app_id}: {e}") from e
+
     return client.get_application(app_id)