Use temp directory instead of resources directory

Author: Jithendar12

athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/GlueRegistryReader.java

@@ -19,6 +19,7 @@
  */
 package com.amazonaws.athena.connectors.msk;
 
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connectors.msk.dto.MSKField;
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -27,6 +28,8 @@
 import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
 import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
 import software.amazon.awssdk.services.glue.GlueClient;
+import software.amazon.awssdk.services.glue.model.ErrorDetails;
+import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
 import software.amazon.awssdk.services.glue.model.GetSchemaRequest;
 import software.amazon.awssdk.services.glue.model.GetSchemaResponse;
 import software.amazon.awssdk.services.glue.model.GetSchemaVersionRequest;
@@ -35,16 +38,17 @@
 import software.amazon.awssdk.services.glue.model.SchemaVersionNumber;
 
 import java.io.FileInputStream;
+import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.UUID;
 
 public class GlueRegistryReader
 {
     private static final ObjectMapper objectMapper;
-    private static final String PROTO_DIR = "src/main/resources/proto";
     private static final String PROTO_FILE = "schema.proto";
     private static final String DESC_FILE = "schema.desc";
 
@@ -67,7 +71,8 @@ public List<MSKField> getProtobufFields(String glueRegistryName, String glueSche
         GetSchemaVersionResponse schemaVersionResponse = getSchemaVersionResult(glueRegistryName, glueSchemaName);
         String schemaDef = schemaVersionResponse.schemaDefinition();
 
-        Path protoDir = Paths.get(PROTO_DIR);
+        // Create a unique temp directory using UUID
+        Path protoDir = Paths.get("/tmp", "proto_" + UUID.randomUUID());
         Files.createDirectories(protoDir);
         Path protoFile = protoDir.resolve(PROTO_FILE);
         Path descFile = protoDir.resolve(DESC_FILE);
@@ -107,8 +112,20 @@ public List<MSKField> getProtobufFields(String glueRegistryName, String glueSche
             return fields;
         }
         finally {
-            Files.deleteIfExists(protoFile);
-            Files.deleteIfExists(descFile);
+            // Clean up temporary files
+            try {
+                Files.deleteIfExists(protoFile);
+                Files.deleteIfExists(descFile);
+                Files.deleteIfExists(protoDir);
+            }
+            catch (IOException e) {
+                throw new AthenaConnectorException(
+                        "Failed to clean up temporary proto directory: " + protoDir.toAbsolutePath(),
+                        ErrorDetails.builder()
+                                .errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
+                                .build()
+                );
+            }
         }
     }
     /**

athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskRecordHandlerTest.java

@@ -30,6 +30,7 @@
 import com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintEvaluator;
 import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints;
 import com.amazonaws.athena.connector.lambda.domain.spill.S3SpillLocation;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest;
 import com.amazonaws.athena.connector.lambda.security.EncryptionKey;
 import com.amazonaws.athena.connector.lambda.security.EncryptionKeyFactory;
@@ -70,6 +71,8 @@
 import org.mockito.junit.MockitoJUnitRunner;
 import software.amazon.awssdk.services.athena.AthenaClient;
 import software.amazon.awssdk.services.glue.GlueClient;
+import software.amazon.awssdk.services.glue.model.ErrorDetails;
+import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
 import software.amazon.awssdk.services.glue.model.GetSchemaRequest;
 import software.amazon.awssdk.services.glue.model.GetSchemaResponse;
 import software.amazon.awssdk.services.glue.model.GetSchemaVersionRequest;
@@ -78,6 +81,7 @@
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 
 import java.io.FileInputStream;
+import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -515,9 +519,9 @@ private Schema createAvroSchema(AvroTopicSchema avroTopicSchema) throws Exceptio
     private Schema createProtobufSchema(String protobufSchema) throws Exception
     {
         SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
-        
-        // Create a temporary directory and files
-        Path protoDir = Paths.get("src/test/resources/proto");
+
+        // Use a unique temporary directory under /tmp
+        Path protoDir = Paths.get("/tmp", "proto_" + UUID.randomUUID());
         Files.createDirectories(protoDir);
         Path protoFile = protoDir.resolve("test.proto");
         Path descFile = protoDir.resolve("test.desc");
@@ -557,15 +561,26 @@ private Schema createProtobufSchema(String protobufSchema) throws Exception
                         Field arrowField = new Field(field.getName(), arrowFieldType, null);
                         schemaBuilder.addField(arrowField);
                     }
-                } else {
+                }
+                else {
                     throw new RuntimeException("No message types found in compiled schema");
                 }
             }
-        } finally {
-            // Clean up
-            Files.deleteIfExists(protoFile);
-            Files.deleteIfExists(descFile);
-            Files.deleteIfExists(protoDir);
+        }
+        finally {
+            try {
+                Files.deleteIfExists(protoFile);
+                Files.deleteIfExists(descFile);
+                Files.deleteIfExists(protoDir);
+            }
+            catch (IOException e) {
+                throw new AthenaConnectorException(
+                        "Failed to clean up temporary proto directory: " + protoDir.toAbsolutePath(),
+                        ErrorDetails.builder()
+                                .errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
+                                .build()
+                );
+            }
         }
 
         schemaBuilder.addMetadata("dataFormat", PROTOBUF_DATA_FORMAT);