[FLINK-39115][s3] Support URI Handling Utility

Author: Samrat002

flink-filesystems/flink-s3-fs-native/src/main/java/org/apache/flink/fs/s3native/NativeS3BulkCopyHelper.java

@@ -78,9 +78,6 @@
  * atomically moved into place after the stream has been fully copied. Cancellation through the
  * provided {@link ICloseableRegistry} aborts active S3 response streams, cancels pending futures,
  * stops the worker pool, and deletes incomplete temporary files.
- *
- * <p><b>TODO:</b> Consider extracting URI parsing logic to a shared S3UriUtils utility class to
- * consolidate S3 URI handling across the codebase.
  */
 @Internal
 class NativeS3BulkCopyHelper {
@@ -185,8 +182,8 @@ public void copyFiles(
             for (int i = 0; i < requests.size(); i++) {
                 PathsCopyingFileSystem.CopyRequest request = requests.get(i);
                 String sourceUri = request.getSource().toUri().toString();
-                if (isSupportedS3Scheme(request.getSource())
-                        && isSupportedLocalScheme(request.getDestination())) {
+                if (S3UriUtils.isSupportedS3Scheme(request.getSource())
+                        && S3UriUtils.isSupportedLocalScheme(request.getDestination())) {
                     copyFutures.add(copyS3ToLocal(request, downloadPool, cancellation));
                 } else {
                     throw new UnsupportedOperationException(
@@ -237,8 +234,8 @@ private CompletableFuture<Void> copyS3ToLocal(
             throws IOException {
 
         String sourceUri = request.getSource().toUri().toString();
-        String bucket = extractBucket(sourceUri);
-        String key = extractKey(sourceUri);
+        String bucket = S3UriUtils.extractBucketName(sourceUri);
+        String key = S3UriUtils.extractKey(sourceUri);
         Path destination = new File(request.getDestination().getPath()).toPath().toAbsolutePath();
 
         Path parent = destination.getParent();
@@ -388,16 +385,6 @@ private void waitForCopies(List<CompletableFuture<Void>> futures) throws IOExcep
         }
     }
 
-    static boolean isSupportedS3Scheme(org.apache.flink.core.fs.Path path) {
-        String scheme = path.toUri().getScheme();
-        return "s3".equalsIgnoreCase(scheme) || "s3a".equalsIgnoreCase(scheme);
-    }
-
-    static boolean isSupportedLocalScheme(org.apache.flink.core.fs.Path path) {
-        String scheme = path.toUri().getScheme();
-        return scheme == null || "file".equalsIgnoreCase(scheme);
-    }
-
     private static void abortAndClose(ResponseInputStream<GetObjectResponse> stream) {
         try {
             stream.abort();
@@ -502,40 +489,4 @@ static boolean isConnectionPoolExhausted(Throwable throwable) {
         }
         return false;
     }
-
-    /**
-     * Extracts the bucket name from an S3 URI.
-     *
-     * <p>Supports both s3:// and s3a:// schemes (s3a is normalized to s3).
-     *
-     * @param s3Uri the S3 URI
-     * @return the bucket name
-     */
-    String extractBucket(String s3Uri) {
-        String uri = s3Uri.replaceFirst("s3a://", "s3://");
-        int bucketStart = uri.indexOf("://") + 3;
-        int bucketEnd = uri.indexOf("/", bucketStart);
-        if (bucketEnd == -1) {
-            return uri.substring(bucketStart);
-        }
-        return uri.substring(bucketStart, bucketEnd);
-    }
-
-    /**
-     * Extracts the object key from an S3 URI.
-     *
-     * <p>Supports both s3:// and s3a:// schemes (s3a is normalized to s3).
-     *
-     * @param s3Uri the S3 URI
-     * @return the object key (empty string if no key in URI)
-     */
-    String extractKey(String s3Uri) {
-        String uri = s3Uri.replaceFirst("s3a://", "s3://");
-        int bucketStart = uri.indexOf("://") + 3;
-        int keyStart = uri.indexOf("/", bucketStart);
-        if (keyStart == -1) {
-            return "";
-        }
-        return uri.substring(keyStart + 1);
-    }
 }

flink-filesystems/flink-s3-fs-native/src/main/java/org/apache/flink/fs/s3native/NativeS3FileSystem.java

@@ -205,7 +205,7 @@ public Path getHomeDirectory() {
     @Override
     public FileStatus getFileStatus(Path path) throws IOException {
         checkNotClosed();
-        final String key = NativeS3ObjectOperations.extractKey(path);
+        final String key = S3UriUtils.extractKey(path);
         final S3Client s3Client = clientProvider.getS3Client();
 
         LOG.debug("Getting file status for s3://{}/{}", bucketName, key);
@@ -304,7 +304,7 @@ public BlockLocation[] getFileBlockLocations(FileStatus file, long start, long l
     @Override
     public FSDataInputStream open(Path path, int bufferSize) throws IOException {
         checkNotClosed();
-        final String key = NativeS3ObjectOperations.extractKey(path);
+        final String key = S3UriUtils.extractKey(path);
         final S3Client s3Client = clientProvider.getS3Client();
         final long fileSize = getFileStatus(path).getLen();
         return new NativeS3InputStream(s3Client, bucketName, key, fileSize, bufferSize);
@@ -313,7 +313,7 @@ public FSDataInputStream open(Path path, int bufferSize) throws IOException {
     @Override
     public FSDataInputStream open(Path path) throws IOException {
         checkNotClosed();
-        final String key = NativeS3ObjectOperations.extractKey(path);
+        final String key = S3UriUtils.extractKey(path);
         final S3Client s3Client = clientProvider.getS3Client();
         final long fileSize = getFileStatus(path).getLen();
         return new NativeS3InputStream(s3Client, bucketName, key, fileSize, readBufferSize);
@@ -337,7 +337,7 @@ public FSDataInputStream open(Path path) throws IOException {
     @Override
     public FileStatus[] listStatus(Path path) throws IOException {
         checkNotClosed();
-        String key = NativeS3ObjectOperations.extractKey(path);
+        String key = S3UriUtils.extractKey(path);
         if (!key.isEmpty() && !key.endsWith("/")) {
             key = key + "/";
         }
@@ -384,7 +384,7 @@ public FileStatus[] listStatus(Path path) throws IOException {
     @Override
     public boolean delete(Path path, boolean recursive) throws IOException {
         checkNotClosed();
-        final String key = NativeS3ObjectOperations.extractKey(path);
+        final String key = S3UriUtils.extractKey(path);
         final S3Client s3Client = clientProvider.getS3Client();
 
         try {
@@ -452,7 +452,7 @@ public FSDataOutputStream create(Path path, WriteMode overwriteMode) throws IOEx
             }
         }
 
-        final String key = NativeS3ObjectOperations.extractKey(path);
+        final String key = S3UriUtils.extractKey(path);
         return new NativeS3OutputStream(
                 clientProvider.getS3Client(),
                 bucketName,
@@ -471,8 +471,8 @@ public FSDataOutputStream create(Path path, WriteMode overwriteMode) throws IOEx
     @Override
     public boolean rename(Path src, Path dst) throws IOException {
         checkNotClosed();
-        final String srcKey = NativeS3ObjectOperations.extractKey(src);
-        final String dstKey = NativeS3ObjectOperations.extractKey(dst);
+        final String srcKey = S3UriUtils.extractKey(src);
+        final String dstKey = S3UriUtils.extractKey(dst);
         final S3Client s3Client = clientProvider.getS3Client();
 
         final FileStatus srcStatus = getFileStatus(src);
@@ -519,8 +519,8 @@ public String generateEntropy() {
     @Override
     public boolean canCopyPaths(Path source, Path destination) {
         return bulkCopyHelper != null
-                && NativeS3BulkCopyHelper.isSupportedS3Scheme(source)
-                && NativeS3BulkCopyHelper.isSupportedLocalScheme(destination);
+                && S3UriUtils.isSupportedS3Scheme(source)
+                && S3UriUtils.isSupportedLocalScheme(destination);
     }
 
     @Override

flink-filesystems/flink-s3-fs-native/src/main/java/org/apache/flink/fs/s3native/S3UriUtils.java

@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.fs.s3native;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.core.fs.Path;
+
+import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/** Shared S3 URI handling utilities. */
+@Internal
+public final class S3UriUtils {
+
+    private S3UriUtils() {}
+
+    private static final String S3 = "s3://";
+    private static final String S3A = "s3a://";
+
+    /** Extracts the S3 object key from a Flink {@link Path}. */
+    public static String extractKey(Path path) {
+        String pathStr = path.toUri().getPath();
+        if (pathStr.startsWith("/")) {
+            pathStr = pathStr.substring(1);
+        }
+        return pathStr;
+    }
+
+    /** Extracts the S3 bucket name from a Flink {@link Path}. */
+    public static String extractBucketName(Path path) {
+        return path.toUri().getHost();
+    }
+
+    /**
+     * Extracts the bucket name from a raw S3 URI string.
+     *
+     * @throws IllegalArgumentException if the scheme is not s3:// or s3a://
+     */
+    public static String extractBucketName(String s3Uri) {
+        String uri = requireSupportedScheme(s3Uri);
+        int bucketEnd = uri.indexOf('/', S3.length());
+        return bucketEnd == -1 ? uri.substring(S3.length()) : uri.substring(S3.length(), bucketEnd);
+    }
+
+    /**
+     * Extracts the S3 object key from a raw S3 URI string, or an empty string if the URI addresses
+     * a bucket only.
+     *
+     * @throws IllegalArgumentException if the scheme is not s3:// or s3a://
+     */
+    public static String extractKey(String s3Uri) {
+        String uri = requireSupportedScheme(s3Uri);
+        int keyStart = uri.indexOf('/', S3.length());
+        return keyStart == -1 ? "" : uri.substring(keyStart + 1);
+    }
+
+    /** Returns whether the path uses a supported S3 scheme ({@code s3} or {@code s3a}). */
+    public static boolean isSupportedS3Scheme(Path path) {
+        String scheme = path.toUri().getScheme();
+        return "s3".equalsIgnoreCase(scheme) || "s3a".equalsIgnoreCase(scheme);
+    }
+
+    /** Returns whether the path addresses the local file system ({@code file} or no scheme). */
+    public static boolean isSupportedLocalScheme(Path path) {
+        String scheme = path.toUri().getScheme();
+        return scheme == null || "file".equalsIgnoreCase(scheme);
+    }
+
+    private static String requireSupportedScheme(String s3Uri) {
+        checkNotNull(s3Uri, "s3Uri must not be null");
+        String uri = startsWithIgnoreCase(s3Uri, S3A) ? S3 + s3Uri.substring(S3A.length()) : s3Uri;
+        checkArgument(
+                startsWithIgnoreCase(uri, S3),
+                "Unsupported S3 URI (expected s3:// or s3a:// scheme): %s",
+                s3Uri);
+        return uri;
+    }
+
+    private static boolean startsWithIgnoreCase(String value, String prefix) {
+        return value.regionMatches(true, 0, prefix, 0, prefix.length());
+    }
+}

flink-filesystems/flink-s3-fs-native/src/main/java/org/apache/flink/fs/s3native/writer/NativeS3ObjectOperations.java

@@ -19,7 +19,6 @@
 package org.apache.flink.fs.s3native.writer;
 
 import org.apache.flink.annotation.Internal;
-import org.apache.flink.core.fs.Path;
 import org.apache.flink.fs.s3native.NativeS3FileIoUtils;
 import org.apache.flink.fs.s3native.S3EncryptionConfig;
 import org.apache.flink.fs.s3native.S3ExceptionUtils;
@@ -89,11 +88,6 @@
  *   <li>SSE-C (customer-provided keys) via a KeyProvider interface
  *   <li>Client-side encryption via an EncryptionHandler interface
  * </ul>
- *
- * <p><b>S3 URI Handling:</b> The {@link #extractKey(Path)} and {@link #extractBucketName(Path)}
- * methods expect URIs in the standard {@code s3://bucket/key} format. Other formats like path-style
- * ({@code https://s3.amazonaws.com/bucket/key}) or virtual-hosted-style ({@code
- * https://bucket.s3.amazonaws.com/key}) are not currently supported.
  */
 @Internal
 public class NativeS3ObjectOperations {
@@ -454,44 +448,6 @@ public String getBucketName() {
         return bucketName;
     }
 
-    /**
-     * Extracts the S3 object key from a Flink Path.
-     *
-     * <p>Expected URI format: {@code s3://bucket-name/path/to/object}
-     *
-     * <p><b>Limitations:</b> This method only supports the standard S3 URI format. Other URI
-     * formats are NOT supported:
-     *
-     * <ul>
-     *   <li>{@code https://bucket.s3.amazonaws.com/path/to/object} (virtual-hosted style)
-     *   <li>{@code https://s3.amazonaws.com/bucket/path/to/object} (path style)
-     *   <li>{@code s3a://} or {@code s3n://} schemes (Hadoop-specific)
-     * </ul>
-     *
-     * @param path the Flink Path with s3:// scheme
-     * @return the object key (path portion without leading slash)
-     */
-    public static String extractKey(Path path) {
-        String pathStr = path.toUri().getPath();
-        if (pathStr.startsWith("/")) {
-            pathStr = pathStr.substring(1);
-        }
-        return pathStr;
-    }
-
-    /**
-     * Extracts the S3 bucket name from a Flink Path.
-     *
-     * <p>Expected URI format: {@code s3://bucket-name/path/to/object}
-     *
-     * @param path the Flink Path with s3:// scheme
-     * @return the bucket name (host portion of the URI)
-     * @see #extractKey(Path) for URI format limitations
-     */
-    public static String extractBucketName(Path path) {
-        return path.toUri().getHost();
-    }
-
     public static class UploadPartResult {
         private final int partNumber;
         private final String eTag;

flink-filesystems/flink-s3-fs-native/src/main/java/org/apache/flink/fs/s3native/writer/NativeS3RecoverableWriter.java

@@ -22,6 +22,7 @@
 import org.apache.flink.core.fs.RecoverableFsDataOutputStream;
 import org.apache.flink.core.fs.RecoverableWriter;
 import org.apache.flink.core.io.SimpleVersionedSerializer;
+import org.apache.flink.fs.s3native.S3UriUtils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -64,7 +65,7 @@ private NativeS3RecoverableWriter(
     @Override
     public RecoverableFsDataOutputStream open(Path path) throws IOException {
         checkNotClosed();
-        String key = NativeS3ObjectOperations.extractKey(path);
+        String key = S3UriUtils.extractKey(path);
         LOG.debug("Opening recoverable stream for key: {}", key);
 
         String uploadId = s3AccessHelper.startMultiPartUpload(key);