Support for S3A Connector #14312 (#14474)

apache · Nov 21, 2024 · c5e00bb · c5e00bb
1 parent 81f16f3
commit c5e00bb
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 71 deletions.
diff --git a/...inot-file-system/pinot-s3/src/main/java/org/apache/pinot/plugin/filesystem/S3PinotFS.java b/...inot-file-system/pinot-s3/src/main/java/org/apache/pinot/plugin/filesystem/S3PinotFS.java
@@ -93,7 +93,10 @@ public class S3PinotFS extends BasePinotFS {
   private static final Logger LOGGER = LoggerFactory.getLogger(S3PinotFS.class);
 
   private static final String DELIMITER = "/";
-  public static final String S3_SCHEME = "s3://";
+  public static final String S3_SCHEME = "s3";
+  public static final String S3A_SCHEME = "s3a";
+  public static final String SCHEME_SEPARATOR = "://";
+
   private S3Client _s3Client;
   private boolean _disableAcl;
   private ServerSideEncryption _serverSideEncryption = null;
@@ -501,12 +504,14 @@ public long length(URI fileUri)
   public String[] listFiles(URI fileUri, boolean recursive)
       throws IOException {
     ImmutableList.Builder<String> builder = ImmutableList.builder();
+    String scheme = fileUri.getScheme();
+    Preconditions.checkArgument(scheme.equals(S3_SCHEME) || scheme.equals(S3A_SCHEME));
     visitFiles(fileUri, recursive, s3Object -> {
       if (!s3Object.key().equals(fileUri.getPath()) && !s3Object.key().endsWith(DELIMITER)) {
-        builder.add(S3_SCHEME + fileUri.getHost() + DELIMITER + getNormalizedFileKey(s3Object));
+        builder.add(scheme + SCHEME_SEPARATOR + fileUri.getHost() + DELIMITER + getNormalizedFileKey(s3Object));
       }
     }, commonPrefix -> {
-      builder.add(S3_SCHEME + fileUri.getHost() + DELIMITER + getNormalizedFileKey(commonPrefix));
+      builder.add(scheme + SCHEME_SEPARATOR + fileUri.getHost() + DELIMITER + getNormalizedFileKey(commonPrefix));
     });
     String[] listedFiles = builder.build().toArray(new String[0]);
     LOGGER.info("Listed {} files from URI: {}, is recursive: {}", listedFiles.length, fileUri, recursive);
@@ -517,17 +522,19 @@ public String[] listFiles(URI fileUri, boolean recursive)
   public List<FileMetadata> listFilesWithMetadata(URI fileUri, boolean recursive)
       throws IOException {
     ImmutableList.Builder<FileMetadata> listBuilder = ImmutableList.builder();
+    String scheme = fileUri.getScheme();
+    Preconditions.checkArgument(scheme.equals(S3_SCHEME) || scheme.equals(S3A_SCHEME));
     visitFiles(fileUri, recursive, s3Object -> {
       if (!s3Object.key().equals(fileUri.getPath())) {
         FileMetadata.Builder fileBuilder = new FileMetadata.Builder().setFilePath(
-                S3_SCHEME + fileUri.getHost() + DELIMITER + getNormalizedFileKey(s3Object))
+                scheme + SCHEME_SEPARATOR + fileUri.getHost() + DELIMITER + getNormalizedFileKey(s3Object))
             .setLastModifiedTime(s3Object.lastModified().toEpochMilli()).setLength(s3Object.size())
             .setIsDirectory(s3Object.key().endsWith(DELIMITER));
         listBuilder.add(fileBuilder.build());
       }
     }, commonPrefix -> {
       FileMetadata.Builder fileBuilder = new FileMetadata.Builder()
-          .setFilePath(S3_SCHEME + fileUri.getHost() + DELIMITER + getNormalizedFileKey(commonPrefix))
+          .setFilePath(scheme + SCHEME_SEPARATOR + fileUri.getHost() + DELIMITER + getNormalizedFileKey(commonPrefix))
           .setIsDirectory(true);
       listBuilder.add(fileBuilder.build());
     });