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

Author: benkermani

src/spark_history_mcp/common/s3_client.py

@@ -0,0 +1,43 @@
+import os
+
+import boto3
+
+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/"
+
+
+    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
+
+    def copy_spark_events_logs(self, spark_app_id: str):
+        # get spark events logs file to copy/index
+        src_prefix = f"spark_logs/{spark_app_id}"
+        base_logs = self.client.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)
+
+        return {"success": True}

src/spark_history_mcp/tools/tools.py

@@ -22,6 +22,7 @@
     TaskMetricDistributions,
 )
 from ..common.yoshi import Yoshi
+from ..common.s3_client import S3Client
 
 from ..utils.utils import parallel_execute
 
@@ -155,6 +156,11 @@ 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):
+        s3_client.copy_spark_events_logs(app_id)
+
     return client.get_application(app_id)