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

Author: benkermani

src/spark_history_mcp/common/s3_client.py

@@ -0,0 +1,42 @@
+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
+
+
+    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 copy_spark_events_logs(self, spark_app_id: str):
+        src_prefix = f"spark_logs/{spark_app_id}"
+        dst_prefix = "indexed_spark_logs/"
+
+        # don't copy if events are already indexed
+        if self.list_contents_by_prefix(dst_prefix + str(spark_app_id), self.bucket_name):
+            return {"success": True}
+
+        # get spark events logs file to copy/index
+        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 = 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}