try to incorporate changes from #4335 and add IcebergRESTVendedS3Provider.

Author: mbutrovich

common/src/main/java/org/apache/comet/cloud/s3/CometS3CredentialDispatcher.java

@@ -20,6 +20,9 @@
 package org.apache.comet.cloud.s3;
 
 import java.lang.reflect.InvocationTargetException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 
 import org.slf4j.Logger;
@@ -28,51 +31,87 @@
 /**
  * JNI entry point invoked from native code to resolve a {@link CometS3CredentialProvider}.
  *
- * <p>Native code passes the FQCN named in {@code fs.s3a.comet.credential.provider.class} (or its
- * per-bucket / Iceberg-namespaced variants). Each named class is instantiated once via reflection
- * and cached, so a single executor JVM can serve multiple providers (e.g. one per bucket).
+ * <p>Native code names a vendor class via the activation knob ({@code
+ * fs.s3a.comet.credential.provider.class} for the Parquet path, {@code
+ * s3.comet.credential.provider.class} on a Spark catalog property for the Iceberg path) and a
+ * {@code dispatchKey} that scopes the instance: catalog name on the Iceberg path, bucket name on
+ * the Parquet path. Each {@code (FQCN, dispatchKey)} key gets its own instance, so two catalogs
+ * sharing one provider class get isolated state.
  */
 public final class CometS3CredentialDispatcher {
 
   private static final Logger LOG = LoggerFactory.getLogger(CometS3CredentialDispatcher.class);
 
-  private static final ConcurrentHashMap<String, CometS3CredentialProvider> INSTANCES =
+  private static final ConcurrentHashMap<InstanceKey, CometS3CredentialProvider> INSTANCES =
       new ConcurrentHashMap<>();
   private static final CometS3AccessMode[] MODES = CometS3AccessMode.values();
 
   private CometS3CredentialDispatcher() {}
 
   /**
-   * Invoked by native code. {@code mode} is the {@link CometS3AccessMode} ordinal.
-   *
-   * @param providerClassName FQCN configured in {@code fs.s3a.comet.credential.provider.class}
+   * Reflects and initializes the named provider for {@code (FQCN, dispatchKey)} exactly once per
+   * JVM. Subsequent calls with the same key are no-ops. Native code invokes this synchronously when
+   * {@code CometS3CredentialBridge} is constructed at plan time, before any per-request {@link
+   * #getCredentialsForPath} call. {@code catalogProperties} carries the unfiltered FileIO property
+   * bag on the Iceberg path and is empty on the Parquet path.
+   */
+  public static void ensureInitialized(
+      String providerClassName, String dispatchKey, Map<String, String> catalogProperties) {
+    if (providerClassName == null || providerClassName.isEmpty()) {
+      throw new IllegalArgumentException(
+          "providerClassName is empty; native side should not call without a configured class");
+    }
+    InstanceKey key = new InstanceKey(providerClassName, dispatchKey == null ? "" : dispatchKey);
+    Map<String, String> props =
+        catalogProperties == null ? Collections.emptyMap() : catalogProperties;
+    INSTANCES.computeIfAbsent(
+        key,
+        k -> {
+          CometS3CredentialProvider provider = instantiate(k.providerClassName);
+          provider.initialize(props);
+          return provider;
+        });
+  }
+
+  /**
+   * Invoked by native code on every per-request credential fetch. The instance must have been
+   * created by a prior {@link #ensureInitialized} call; otherwise this throws. {@code mode} is the
+   * {@link CometS3AccessMode} ordinal.
    */
   public static CometS3Credentials getCredentialsForPath(
-      String providerClassName, String bucket, String path, int mode) throws Exception {
+      String providerClassName, String dispatchKey, String bucket, String path, int mode)
+      throws Exception {
     if (providerClassName == null || providerClassName.isEmpty()) {
       throw new IllegalArgumentException(
           "providerClassName is empty; native side should not call without a configured class");
     }
     if (mode < 0 || mode >= MODES.length) {
       throw new IllegalArgumentException("Invalid CometS3AccessMode ordinal: " + mode);
     }
-    CometS3CredentialProvider provider = resolve(providerClassName);
+    InstanceKey key = new InstanceKey(providerClassName, dispatchKey == null ? "" : dispatchKey);
+    CometS3CredentialProvider provider = INSTANCES.get(key);
+    if (provider == null) {
+      throw new IllegalStateException(
+          "CometS3CredentialProvider "
+              + providerClassName
+              + " (dispatchKey="
+              + key.dispatchKey
+              + ") was not initialized; ensureInitialized must be called before"
+              + " getCredentialsForPath");
+    }
     CometS3AccessMode accessMode = MODES[mode];
     if (LOG.isDebugEnabled()) {
       LOG.debug(
-          "Fetching credentials via {} for bucket={} path={} mode={}",
+          "Fetching credentials via {} (dispatchKey={}) for bucket={} path={} mode={}",
           providerClassName,
+          key.dispatchKey,
           bucket,
           path,
           accessMode);
     }
     return provider.getCredentialsForPath(bucket, path, accessMode);
   }
 
-  private static CometS3CredentialProvider resolve(String providerClassName) {
-    return INSTANCES.computeIfAbsent(providerClassName, CometS3CredentialDispatcher::instantiate);
-  }
-
   private static CometS3CredentialProvider instantiate(String providerClassName) {
     Class<?> clazz;
     try {
@@ -100,4 +139,28 @@ private static CometS3CredentialProvider instantiate(String providerClassName) {
           "Failed to instantiate CometS3CredentialProvider " + providerClassName, e);
     }
   }
+
+  private static final class InstanceKey {
+    final String providerClassName;
+    final String dispatchKey;
+
+    InstanceKey(String providerClassName, String dispatchKey) {
+      this.providerClassName = providerClassName;
+      this.dispatchKey = dispatchKey;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (!(o instanceof InstanceKey)) return false;
+      InstanceKey other = (InstanceKey) o;
+      return providerClassName.equals(other.providerClassName)
+          && dispatchKey.equals(other.dispatchKey);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(providerClassName, dispatchKey);
+    }
+  }
 }

common/src/main/java/org/apache/comet/cloud/s3/CometS3CredentialProvider.java

@@ -19,6 +19,8 @@
 
 package org.apache.comet.cloud.s3;
 
+import java.util.Map;
+
 /**
  * SPI for supplying AWS credentials to Comet's native S3 readers, which bypass Spark's Hadoop S3A
  * code path and cannot reach signer-based or path-aware credential mechanisms through the standard
@@ -47,16 +49,47 @@
  *
  * <p>Vendors register an implementation by setting {@code
  * spark.hadoop.fs.s3a.comet.credential.provider.class} (or the per-bucket form {@code
- * spark.hadoop.fs.s3a.bucket.<name>.comet.credential.provider.class}) to the implementing FQCN. The
- * class must have a public no-arg constructor. {@link #getCredentialsForPath} may be invoked
- * concurrently from many native tokio tasks, so implementations must be thread-safe.
+ * spark.hadoop.fs.s3a.bucket.<name>.comet.credential.provider.class}) for the Parquet path, or
+ * {@code spark.sql.catalog.<catalog>.s3.comet.credential.provider.class} for the Iceberg path. The
+ * class must have a public no-arg constructor.
+ *
+ * <h2>Lifecycle</h2>
+ *
+ * <p>Comet keys provider instances by {@code (FQCN, dispatchKey)}, where {@code dispatchKey} is the
+ * Spark V2 catalog name on the Iceberg path and the bucket on the Parquet path. The first time a
+ * given key is seen on an executor, Comet reflects the class, calls {@link #initialize(Map)} once,
+ * and caches the instance. Subsequent requests for the same key reuse it. Two catalogs that share
+ * one FQCN therefore get isolated instances with their own {@code initialize} maps.
+ *
+ * <p>{@link #initialize(Map)} should be cheap and non-blocking; defer real credential fetches to
+ * the first {@link #getCredentialsForPath} call. {@link #getCredentialsForPath} may be invoked
+ * concurrently from many native tokio worker threads, so implementations must be thread-safe.
+ *
+ * <h2>Caching, refresh, and distribution are the vendor's job</h2>
+ *
+ * <p>Comet does not maintain a TTL cache, broadcast catalog state, or schedule refresh. Vendors
+ * decide whether to cache (e.g. by wrapping {@code
+ * org.apache.iceberg.aws.s3.VendedCredentialsProvider}'s {@code CachedSupplier}), when to refresh,
+ * and how to distribute driver-side state to executors (typically by reading {@link #initialize}'s
+ * {@code catalogProperties}, which Comet has already serialized through the native plan op).
  *
  * <p>Returns credentials or throws; there is no fall-through return value. A provider that is only
  * authoritative for some paths should resolve the default AWS chain itself for the rest. See the
  * user guide on cloud credential providers.
  */
 public interface CometS3CredentialProvider {
 
+  /**
+   * Called once per {@code (FQCN, dispatchKey)} on each executor before any {@link
+   * #getCredentialsForPath} call. The {@code catalogProperties} map carries the full FileIO
+   * property bag for the Iceberg path (including {@code credentials.uri}, OAuth tokens, vendor keys
+   * like {@code tenant-id}) and is empty on the Parquet path. The default no-op keeps Parquet
+   * vendors source-compatible.
+   *
+   * @param catalogProperties unfiltered FileIO/catalog properties; may contain secrets, do not log
+   */
+  default void initialize(Map<String, String> catalogProperties) {}
+
   /**
    * @param bucket S3 bucket name (no scheme, no path)
    * @param path object key or prefix, leading slash included (matches the URL path component)

common/src/test/java/org/apache/comet/cloud/s3/CometS3CredentialDispatcherTest.java

@@ -19,6 +19,10 @@
 
 package org.apache.comet.cloud.s3;
 
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
 import org.junit.Before;
 import org.junit.Test;
 
@@ -32,6 +36,7 @@
 public class CometS3CredentialDispatcherTest {
 
   private static final String TEST_PROVIDER = TestCometS3CredentialProvider.class.getName();
+  private static final String DK = "test-dispatch-key";
   private static final int READ = CometS3AccessMode.READ.ordinal();
   private static final int WRITE = CometS3AccessMode.WRITE.ordinal();
 
@@ -40,11 +45,16 @@ public void resetTestProvider() {
     TestCometS3CredentialProvider.reset();
   }
 
+  private void init() {
+    CometS3CredentialDispatcher.ensureInitialized(TEST_PROVIDER, DK, Collections.emptyMap());
+  }
+
   @Test
   public void getCredentialsRoundTripsThroughProvider() throws Exception {
+    init();
     CometS3Credentials creds =
         CometS3CredentialDispatcher.getCredentialsForPath(
-            TEST_PROVIDER, "my-bucket", "path/to/object", READ);
+            TEST_PROVIDER, DK, "my-bucket", "path/to/object", READ);
 
     assertNotNull(creds);
     assertEquals("AKIATEST", creds.getAccessKeyId());
@@ -60,25 +70,30 @@ public void getCredentialsRoundTripsThroughProvider() throws Exception {
 
   @Test
   public void writeModeIsForwarded() throws Exception {
-    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "b", "k", WRITE);
+    init();
+    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, DK, "b", "k", WRITE);
     assertEquals(CometS3AccessMode.WRITE, TestCometS3CredentialProvider.lastMode);
   }
 
   @Test
   public void unknownModeRejected() {
+    init();
     assertThrows(
         IllegalArgumentException.class,
-        () -> CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "b", "k", 99));
+        () -> CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, DK, "b", "k", 99));
   }
 
   @Test
   public void emptyClassNameRejected() {
     assertThrows(
         IllegalArgumentException.class,
-        () -> CometS3CredentialDispatcher.getCredentialsForPath("", "b", "k", READ));
+        () -> CometS3CredentialDispatcher.ensureInitialized("", DK, Collections.emptyMap()));
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> CometS3CredentialDispatcher.ensureInitialized(null, DK, Collections.emptyMap()));
     assertThrows(
         IllegalArgumentException.class,
-        () -> CometS3CredentialDispatcher.getCredentialsForPath(null, "b", "k", READ));
+        () -> CometS3CredentialDispatcher.getCredentialsForPath("", DK, "b", "k", READ));
   }
 
   @Test
@@ -87,8 +102,8 @@ public void missingClassReportsActionableError() {
         assertThrows(
             IllegalStateException.class,
             () ->
-                CometS3CredentialDispatcher.getCredentialsForPath(
-                    "com.does.not.Exist", "b", "k", READ));
+                CometS3CredentialDispatcher.ensureInitialized(
+                    "com.does.not.Exist", DK, Collections.emptyMap()));
     assertTrue(thrown.getMessage().contains("not found"));
   }
 
@@ -98,8 +113,8 @@ public void classNotImplementingInterfaceRejected() {
         assertThrows(
             IllegalStateException.class,
             () ->
-                CometS3CredentialDispatcher.getCredentialsForPath(
-                    NotACredentialProvider.class.getName(), "b", "k", READ));
+                CometS3CredentialDispatcher.ensureInitialized(
+                    NotACredentialProvider.class.getName(), DK, Collections.emptyMap()));
     assertTrue(thrown.getMessage().contains("does not implement"));
   }
 
@@ -109,38 +124,85 @@ public void classWithoutNoArgCtorRejected() {
         assertThrows(
             IllegalStateException.class,
             () ->
-                CometS3CredentialDispatcher.getCredentialsForPath(
-                    NoNoArgCtorProvider.class.getName(), "b", "k", READ));
+                CometS3CredentialDispatcher.ensureInitialized(
+                    NoNoArgCtorProvider.class.getName(), DK, Collections.emptyMap()));
     assertTrue(thrown.getMessage().contains("no-arg constructor"));
   }
 
+  @Test
+  public void getWithoutEnsureInitializedThrows() {
+    Exception thrown =
+        assertThrows(
+            IllegalStateException.class,
+            () ->
+                CometS3CredentialDispatcher.getCredentialsForPath(
+                    TEST_PROVIDER, "never-initialized", "b", "k", READ));
+    assertTrue(thrown.getMessage().contains("not initialized"));
+  }
+
+  @Test
+  public void initializeCalledExactlyOncePerKey() {
+    Map<String, String> props = new HashMap<>();
+    props.put("tenant-id", "T1");
+
+    CometS3CredentialDispatcher.ensureInitialized(TEST_PROVIDER, "cat-a", props);
+    CometS3CredentialDispatcher.ensureInitialized(TEST_PROVIDER, "cat-a", props);
+    CometS3CredentialDispatcher.ensureInitialized(TEST_PROVIDER, "cat-a", props);
+
+    assertEquals(1, TestCometS3CredentialProvider.initCount.get());
+  }
+
+  @Test
+  public void distinctDispatchKeysIsolateInstances() throws Exception {
+    Map<String, String> propsA = new HashMap<>();
+    propsA.put("tenant-id", "T-A");
+    Map<String, String> propsB = new HashMap<>();
+    propsB.put("tenant-id", "T-B");
+
+    CometS3CredentialDispatcher.ensureInitialized(TEST_PROVIDER, "iso-cat-a", propsA);
+    CometS3CredentialDispatcher.ensureInitialized(TEST_PROVIDER, "iso-cat-b", propsB);
+
+    assertEquals(2, TestCometS3CredentialProvider.initCount.get());
+
+    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "iso-cat-a", "b", "k", READ);
+    assertEquals("T-A", TestCometS3CredentialProvider.lastTenantSeen);
+
+    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "iso-cat-b", "b", "k", READ);
+    assertEquals("T-B", TestCometS3CredentialProvider.lastTenantSeen);
+  }
+
   @Test
   public void providerExceptionsPropagate() {
+    init();
     IllegalStateException boom = new IllegalStateException("simulated STS failure");
     TestCometS3CredentialProvider.throwOnNext = boom;
 
     Exception thrown =
         assertThrows(
             Exception.class,
-            () -> CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "b", "k", READ));
+            () ->
+                CometS3CredentialDispatcher.getCredentialsForPath(
+                    TEST_PROVIDER, DK, "b", "k", READ));
     assertSame(boom, thrown);
   }
 
   @Test
   public void nullSessionTokenAllowed() throws Exception {
+    init();
     TestCometS3CredentialProvider.nextResult = new CometS3Credentials("AKIA", "sec", null, 0L);
 
     CometS3Credentials creds =
-        CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "b", "k", READ);
+        CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, DK, "b", "k", READ);
 
     assertNull(creds.getSessionToken());
   }
 
   @Test
   public void providerReceivesEachCallSeparately() throws Exception {
-    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "b1", "k1", READ);
-    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "b2", "k2", READ);
-    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, "b3", "k3", READ);
+    init();
+    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, DK, "b1", "k1", READ);
+    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, DK, "b2", "k2", READ);
+    CometS3CredentialDispatcher.getCredentialsForPath(TEST_PROVIDER, DK, "b3", "k3", READ);
 
     assertEquals(3, TestCometS3CredentialProvider.callCount.get());
     assertEquals("b3", TestCometS3CredentialProvider.lastBucket);