feat: Use DH S3Instructions to build Iceberg AWS clients

extensions/iceberg/s3/build.gradle

@@ -24,12 +24,19 @@ dependencies {
     implementation libs.awssdk.s3
     implementation libs.awssdk.crt.client
     runtimeOnly libs.awssdk.sts
-    runtimeOnly libs.awssdk.glue
+    implementation libs.awssdk.glue
+
+    // We don't want to explicitly pull in dependencies for dynamodb (org.apache.iceberg.aws.dynamodb.DynamoDbCatalog),
+    // but we need to be able to compile against it to implement AwsClientFactory
+    compileOnly libs.awssdk.dynamodb
+
+    // We don't want to explicitly pull in dependencies for KMS (there doesn't seem to be anything in Iceberg that
+    // actually calls it?), but we need to be able to compile against it to implement AwsClientFactory
+    compileOnly libs.awssdk.kms
 
     compileOnly libs.autoservice
     annotationProcessor libs.autoservice.compiler
 
-    testImplementation libs.junit4
     testImplementation project(':engine-test-utils')
 
     testImplementation libs.testcontainers
@@ -43,12 +50,27 @@ dependencies {
     testRuntimeOnly project(':test-configs')
     testRuntimeOnly project(':log-to-slf4j')
     testRuntimeOnly libs.slf4j.simple
+
+    testImplementation platform(libs.junit.bom)
+    testImplementation libs.junit.jupiter
+    testRuntimeOnly libs.junit.jupiter.engine
+    testRuntimeOnly libs.junit.platform.launcher
 }
 
-TestTools.addEngineOutOfBandTest(project)
+test {
+    useJUnitPlatform {
+        excludeTags("testcontainers")
+    }
+}
+
+tasks.register('testOutOfBand', Test) {
+    useJUnitPlatform {
+        includeTags("testcontainers")
+    }
 
-testOutOfBand.dependsOn Docker.registryTask(project, 'localstack')
-testOutOfBand.systemProperty 'testcontainers.localstack.image', Docker.localImageName('localstack')
+    dependsOn Docker.registryTask(project, 'localstack')
+    systemProperty 'testcontainers.localstack.image', Docker.localImageName('localstack')
 
-testOutOfBand.dependsOn Docker.registryTask(project, 'minio')
-testOutOfBand.systemProperty 'testcontainers.minio.image', Docker.localImageName('minio')
+    dependsOn Docker.registryTask(project, 'minio')
+    systemProperty 'testcontainers.minio.image', Docker.localImageName('minio')
+}

extensions/iceberg/s3/src/main/java/io/deephaven/extensions/s3/DeephavenAwsClientFactory.java

@@ -0,0 +1,121 @@
+//
+// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
+//
+package io.deephaven.extensions.s3;
+
+import org.apache.iceberg.aws.AwsClientFactory;
+import org.apache.iceberg.aws.AwsProperties;
+import org.apache.iceberg.aws.s3.S3FileIOAwsClientFactory;
+import org.apache.iceberg.aws.s3.S3FileIOProperties;
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import software.amazon.awssdk.services.glue.GlueClient;
+import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.awssdk.services.s3.S3Client;
+
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * An {@link AwsClientFactory} and {@link S3FileIOAwsClientFactory} implementation that assumes ownership of AWS client
+ * creation as configured via {@link S3Instructions}.
+ */
+public final class DeephavenAwsClientFactory implements AwsClientFactory, S3FileIOAwsClientFactory {
+
+    private static final String UUID_KEY = DeephavenAwsClientFactory.class.getName() + ".__uuid";
+
+    /**
+     * Adds {@link DeephavenAwsClientFactory} to {@code propertiesOut} with the keys
+     * {@value AwsProperties#CLIENT_FACTORY} and {@value S3FileIOProperties#CLIENT_FACTORY}; it is an error if either of
+     * these properties is already set. After the corresponding {@link org.apache.iceberg.catalog.Catalog} is no longer
+     * in use, the caller should invoke the returned {@link Runnable} to clean up.
+     *
+     * @param instructions the instructions
+     * @param propertiesOut the properties
+     * @return the runnable to be invoked after initialization
+     */
+    public static Runnable addToProperties(S3Instructions instructions, Map<String, String> propertiesOut) {
+        Objects.requireNonNull(instructions);
+        put(propertiesOut, AwsProperties.CLIENT_FACTORY, DeephavenAwsClientFactory.class.getName());
+        put(propertiesOut, S3FileIOProperties.CLIENT_FACTORY, DeephavenAwsClientFactory.class.getName());
+        final String uuid = UUID.randomUUID().toString();
+        put(propertiesOut, UUID_KEY, uuid);
+        S3_INSTRUCTIONS_MAP.put(uuid, instructions);
+        return () -> S3_INSTRUCTIONS_MAP.remove(uuid);
+    }
+
+    /**
+     * Get the {@link S3Instructions} as set in the corresponding {@link #addToProperties(S3Instructions, Map)} if the
+     * properties were built with that. If the properties were built with {@link #addToProperties(S3Instructions, Map)},
+     * but the {@link Runnable} was already invoked for cleanup, an {@link IllegalStateException} will be thrown.
+     *
+     * @param properties the properties
+     * @return the instructions
+     */
+    public static Optional<S3Instructions> get(Map<String, String> properties) {
+        final String uuid = properties.get(UUID_KEY);
+        if (uuid == null) {
+            return Optional.empty();
+        }
+        final S3Instructions instructions = S3_INSTRUCTIONS_MAP.get(uuid);
+        if (instructions == null) {
+            throw new IllegalStateException(
+                    "This S3Iinstructions were already cleaned up; please ensure that the returned Runnable from addToProperties is not invoked until the Catalog is no longer in use.");
+        }
+        return Optional.of(instructions);
+    }
+
+    private static <K, V> void put(Map<K, V> map, K key, V value) {
+        if (map.putIfAbsent(key, value) != null) {
+            throw new IllegalArgumentException(String.format("Key '%s' already exist in map", key));
+        }
+    }
+
+    private static final Map<String, S3Instructions> S3_INSTRUCTIONS_MAP = new ConcurrentHashMap<>();
+
+    private S3Instructions instructions;
+
+    @Override
+    public void initialize(Map<String, String> properties) {
+        this.instructions = get(properties).orElseThrow(() -> new IllegalArgumentException(
+                "DeephavenAwsClientFactory was setup improperly; it must be configured with DeephavenAwsClientFactory.addToProperties"));
+    }
+
+    private void checkInit() {
+        if (instructions == null) {
+            throw new IllegalStateException("Must initialize before use");
+        }
+    }
+
+    @Override
+    public S3Client s3() {
+        checkInit();
+        return S3ClientFactory.getSyncClient(instructions);
+    }
+
+    @Override
+    public GlueClient glue() {
+        checkInit();
+        return GlueClient.builder()
+                .applyMutation(b -> S3ClientFactory.applyAllSharedSync(b, instructions))
+                .build();
+    }
+
+    @Override
+    public KmsClient kms() {
+        checkInit();
+        return KmsClient.builder()
+                .applyMutation(b -> S3ClientFactory.applyAllSharedSync(b, instructions))
+                .build();
+    }
+
+    @Override
+    public DynamoDbClient dynamo() {
+        checkInit();
+        return DynamoDbClient.builder()
+                .applyMutation(b -> S3ClientFactory.applyAllSharedSync(b, instructions))
+                .build();
+    }
+}

extensions/iceberg/s3/src/main/java/io/deephaven/iceberg/util/IcebergToolsS3.java

@@ -4,24 +4,32 @@
 package io.deephaven.iceberg.util;
 
 import com.google.common.base.Strings;
+import io.deephaven.extensions.s3.DeephavenAwsClientFactory;
+import io.deephaven.extensions.s3.S3Instructions;
 import org.apache.iceberg.CatalogProperties;
 import org.apache.iceberg.aws.AwsClientProperties;
+import org.apache.iceberg.aws.AwsProperties;
+import org.apache.iceberg.aws.HttpClientProperties;
 import org.apache.iceberg.aws.glue.GlueCatalog;
 import org.apache.iceberg.aws.s3.S3FileIOProperties;
 import org.apache.iceberg.rest.RESTCatalog;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.lang.ref.Cleaner;
 import java.util.HashMap;
 import java.util.Map;
 
 /**
  * Tools for accessing tables in the Iceberg table format.
  */
 @SuppressWarnings("unused")
-public class IcebergToolsS3 extends IcebergTools {
+public final class IcebergToolsS3 {
     private static final String S3_FILE_IO_CLASS = "org.apache.iceberg.aws.s3.S3FileIO";
 
+    // Note: this could move to a shared location
+    private static final Cleaner CLEANER = Cleaner.create();
+
     /**
      * Create an Iceberg catalog adapter for a REST catalog backed by S3 storage. If {@code null} is provided for a
      * value, the system defaults will be used.
@@ -102,4 +110,40 @@ public static IcebergCatalogAdapter createGlue(
 
         return new IcebergCatalogAdapter(catalog, properties);
     }
+
+    /**
+     * Create an Iceberg catalog adapter. Instead of the AWS clients being configured via the various Iceberg-specific
+     * properties (found amongst {@link S3FileIOProperties}, {@link HttpClientProperties}, and {@link AwsProperties}),
+     * the clients are created in the same way that Deephaven's AWS clients are configured with respect to
+     * {@code instructions}. This ensures, amongst other things, that Iceberg's AWS configuration and credentials are
+     * in-sync with Deephaven's AWS configuration and credentials for S3 access. The {@code instructions} will
+     * automatically be used as special instructions if {@link IcebergInstructions#dataInstructions()} if not explicitly
+     * set.
+     *
+     * <p>
+     * The caller is still responsible for providing the properties necessary as specified in
+     * {@link IcebergTools#createAdapter(String, Map, Map)}.
+     *
+     * @param name the name of the catalog; if omitted, the catalog URI will be used to generate a name
+     * @param properties a map containing the Iceberg catalog properties to use
+     * @param hadoopConfig a map containing Hadoop configuration properties to use
+     * @param instructions the s3 instructions
+     * @return the Iceberg catalog adapter
+     *
+     * @see DeephavenAwsClientFactory#addToProperties(S3Instructions, Map)
+     */
+    public static IcebergCatalogAdapter createAdapter(
+            @Nullable final String name,
+            @NotNull final Map<String, String> properties,
+            @NotNull final Map<String, String> hadoopConfig,
+            @NotNull final S3Instructions instructions) {
+        final Map<String, String> newProperties = new HashMap<>(properties);
+        final Runnable cleanup = DeephavenAwsClientFactory.addToProperties(instructions, newProperties);
+        final IcebergCatalogAdapter adapter = IcebergTools.createAdapter(name, newProperties, hadoopConfig);
+        // When the Catalog becomes unreachable, we can invoke the DeephavenAwsClientFactory cleanup.
+        // Note: it would be incorrect to register the cleanup against the adapter since the Catalog can outlive the
+        // adapter (and the DeephavenAwsClientFactory properties are needed by the Catalog).
+        CLEANER.register(adapter.catalog(), cleanup);
+        return adapter;
+    }
 }

extensions/iceberg/s3/src/main/java/io/deephaven/iceberg/util/S3InstructionsProviderPlugin.java

@@ -5,6 +5,7 @@
 
 import com.google.auto.service.AutoService;
 import io.deephaven.extensions.s3.Credentials;
+import io.deephaven.extensions.s3.DeephavenAwsClientFactory;
 import io.deephaven.extensions.s3.S3Instructions;
 import io.deephaven.iceberg.internal.DataInstructionsProviderPlugin;
 import org.apache.iceberg.aws.AwsClientProperties;
@@ -15,13 +16,18 @@
 import java.util.Map;
 
 /**
- * {@link io.deephaven.iceberg.internal.DataInstructionsProviderPlugin} implementation used for reading files from S3.
+ * {@link DataInstructionsProviderPlugin} implementation used for reading files from S3.
  */
-@AutoService(io.deephaven.iceberg.internal.DataInstructionsProviderPlugin.class)
+@AutoService(DataInstructionsProviderPlugin.class)
 @SuppressWarnings("unused")
 public final class S3InstructionsProviderPlugin implements DataInstructionsProviderPlugin {
     @Override
     public Object createInstructions(@NotNull final URI uri, @NotNull final Map<String, String> properties) {
+        final S3Instructions s3Instructions = DeephavenAwsClientFactory.get(properties).orElse(null);
+        if (s3Instructions != null) {
+            return s3Instructions;
+        }
+
         // If the URI scheme is "s3","s3a","s3n" or if the properties contain one of these specific keys, we can
         // create a useful S3Instructions object.
         if (uri.getScheme().equals("s3")

extensions/iceberg/s3/src/test/java/io/deephaven/iceberg/util/IcebergLocalStackTest.java

@@ -5,7 +5,8 @@
 
 import io.deephaven.extensions.s3.S3Instructions.Builder;
 import io.deephaven.extensions.s3.testlib.SingletonContainers.LocalStack;
-import org.junit.BeforeClass;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Tag;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 
 import java.util.Map;
@@ -15,9 +16,11 @@
 import static org.apache.iceberg.aws.s3.S3FileIOProperties.ENDPOINT;
 import static org.apache.iceberg.aws.s3.S3FileIOProperties.SECRET_ACCESS_KEY;
 
+@Tag("testcontainers")
+@Deprecated
 public class IcebergLocalStackTest extends IcebergToolsTest {
 
-    @BeforeClass
+    @BeforeAll
     public static void initContainer() {
         // ensure container is started so container startup time isn't associated with a specific test
         LocalStack.init();
@@ -34,7 +37,7 @@ public S3AsyncClient s3AsyncClient() {
     }
 
     @Override
-    public Map<String, String> s3Properties() {
+    public Map<String, String> properties() {
         return Map.of(
                 ENDPOINT, LocalStack.s3Endpoint(),
                 CLIENT_REGION, LocalStack.region(),

extensions/iceberg/s3/src/test/java/io/deephaven/iceberg/util/IcebergMinIOTest.java

@@ -7,7 +7,8 @@
 import io.deephaven.extensions.s3.testlib.SingletonContainers.MinIO;
 import io.deephaven.stats.util.OSUtil;
 import org.junit.jupiter.api.Assumptions;
-import org.junit.BeforeClass;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Tag;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 
 import java.util.Map;
@@ -17,9 +18,11 @@
 import static org.apache.iceberg.aws.s3.S3FileIOProperties.ENDPOINT;
 import static org.apache.iceberg.aws.s3.S3FileIOProperties.SECRET_ACCESS_KEY;
 
+@Tag("testcontainers")
+@Deprecated
 public class IcebergMinIOTest extends IcebergToolsTest {
 
-    @BeforeClass
+    @BeforeAll
     public static void initContainer() {
         // TODO(deephaven-core#5116): MinIO testcontainers does not work on OS X
         Assumptions.assumeFalse(OSUtil.runningMacOS(), "OSUtil.runningMacOS()");
@@ -38,12 +41,11 @@ public S3AsyncClient s3AsyncClient() {
     }
 
     @Override
-    public Map<String, String> s3Properties() {
+    public Map<String, String> properties() {
         return Map.of(
                 ENDPOINT, MinIO.s3Endpoint(),
                 CLIENT_REGION, MinIO.region(),
                 ACCESS_KEY_ID, MinIO.accessKey(),
                 SECRET_ACCESS_KEY, MinIO.secretAccessKey());
     }
-
 }