parthchandra commented on code in PR #6023: URL: https://github.com/apache/datafusion-comet/pull/6023#discussion_r4109532531
########## spark/src/main/spark-4.x/org/apache/comet/cloud/s3/HadoopS3ACredentialProviderAdapter.java: ########## @@ -0,0 +1,117 @@ +/* + * 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.comet.cloud.s3; + +import java.net.URI; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.s3a.S3AUtils; +import org.apache.hadoop.fs.s3a.auth.CredentialProviderListFactory; + +import org.apache.comet.util.ClassLoaders; + +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; + +/** + * Delegates credential resolution to Hadoop S3A's own provider construction, so it accepts + * everything the {@code fs.s3a.aws.credentials.provider} chain accepts. This is the spark-4.x (AWS + * SDK v2) body; it calls {@link CredentialProviderListFactory} and returns v2 credentials. + * + * <p>Enable it (leaving {@code fs.s3a.aws.credentials.provider} untouched) with: + * + * <pre> + * spark.hadoop.fs.s3a.comet.credential.provider.class=org.apache.comet.cloud.s3.HadoopS3ACredentialProviderAdapter + * </pre> + */ +public class HadoopS3ACredentialProviderAdapter implements CometS3CredentialProvider { + + private Map<String, String> properties; + // Captured on the thread that runs initialize() (the dispatcher calls it during planning, which + // has Spark's user-jar loader); native worker threads have a null context loader. Handed to the + // Configuration so S3A's factory loads the named provider classes from the right loader. + private volatile ClassLoader classLoader; + // One delegate per bucket: on the Iceberg path the dispatch key is the catalog, so a single + // instance can serve multiple buckets; the Parquet path is per-bucket and uses a single entry. + private final ConcurrentHashMap<String, AwsCredentialsProvider> delegates = + new ConcurrentHashMap<>(); + + @Override + public void initialize(Map<String, String> catalogProperties) { + this.properties = catalogProperties; + this.classLoader = ClassLoaders.contextOrDefault(getClass().getClassLoader()); + } + + @Override + public CometS3Credentials getCredentialsForPath(CometS3CredentialContext context) + throws Exception { + AwsCredentialsProvider provider = ensureDelegate(context.getBucket()); + return SdkCredentialExtraction.toCometCredentials(provider.resolveCredentials()); + } + + private AwsCredentialsProvider ensureDelegate(String bucket) throws Exception { + AwsCredentialsProvider existing = delegates.get(bucket); + if (existing != null) { + return existing; + } + synchronized (this) { + AwsCredentialsProvider delegate = delegates.get(bucket); + if (delegate == null) { + delegate = buildDelegate(bucket); + delegates.put(bucket, delegate); + } + return delegate; + } + } + + private AwsCredentialsProvider buildDelegate(String bucket) throws Exception { + Configuration conf = + S3AUtils.propagateBucketOptions(AdapterSupport.toConfiguration(properties), bucket); + if (classLoader != null) { + // So S3AUtils' factory loads the named provider classes from Spark's user-jar loader, not the + // (possibly null) context loader of the native worker thread. + conf.setClassLoader(classLoader); Review Comment: You're right, it's a no-op on 3.4 and the comment was misleading. Dropped the setClassLoader call and the classLoader field from the spark-4.x body, and rewrote the javadoc to say a named provider must be visible to hadoop-aws's own loader on 3.4 (same as plain Spark). Kept it in the spark-3.x body, where 3.3.4 loads through conf.getClasses and honors the conf loader, with a comment saying so. ########## spark/src/main/spark-3.x/org/apache/comet/cloud/s3/AwsSdkCredentialProviderAdapter.java: ########## @@ -0,0 +1,112 @@ +/* + * 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.comet.cloud.s3; + +import java.net.URI; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.s3a.S3AUtils; + +import com.amazonaws.auth.AWSCredentialsProvider; + +import org.apache.comet.util.ClassLoaders; + +/** + * Wraps a raw AWS SDK v1 {@link AWSCredentialsProvider} named via {@code + * fs.s3a.comet.credential.adapter.class}, for a provider not registered through S3A. This is the + * spark-3.x (SDK v1) body. Prefer {@link HadoopS3ACredentialProviderAdapter} unless the provider is + * a plain SDK class not wired through Hadoop. + * + * <pre> + * spark.hadoop.fs.s3a.comet.credential.provider.class=org.apache.comet.cloud.s3.AwsSdkCredentialProviderAdapter + * spark.hadoop.fs.s3a.comet.credential.adapter.class=<FQCN of an AWSCredentialsProvider> + * </pre> + */ +public class AwsSdkCredentialProviderAdapter implements CometS3CredentialProvider { + + static final String DELEGATE_CLASS_PROPERTY = "comet.credential.adapter.class"; + + private Map<String, String> properties; + // Captured on the thread that runs initialize() (the dispatcher calls it during planning, which + // has Spark's user-jar loader). getCredentialsForPath runs on native worker threads whose context + // ClassLoader is null, so the delegate must be loaded with this captured loader, not the TCCL. + private volatile ClassLoader classLoader; + private final ConcurrentHashMap<String, AWSCredentialsProvider> delegates = + new ConcurrentHashMap<>(); + + @Override + public void initialize(Map<String, String> catalogProperties) { + this.properties = catalogProperties; + this.classLoader = ClassLoaders.contextOrDefault(getClass().getClassLoader()); Review Comment: Added tests for all three bodies that capture the loader (spark-4.x SDK, spark-3.x SDK, spark-3.x Hadoop). Each captures a recording ClassLoader in initialize() on one thread, then makes the first getCredentialsForPath call on a thread with a null context loader, and asserts the captured loader loaded the delegate. Setting classLoader = null makes them fail, so they pin the capture. I used a recording loader instead of a child URLClassLoader so it needs no runtime compilation. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
