Honor aws.endpoint for STS and S3 clients

etc/conf/server.properties

@@ -48,6 +48,8 @@ aws.accessKey=
 aws.secretKey=
 # The region of AWS. If running on-premises, pick the nearest region.
 aws.region=
+# Optional STS/S3 endpoint override for self-hosted S3-compatible backends.
+aws.endpoint=
 
 #### The following "s3.*" configs are legacy per-bucket hardcoded credential config
 
@@ -60,6 +62,8 @@ s3.accessKey.0=
 s3.secretKey.0=
 # Test Only (If you provide a session token, it will just use those session creds, no downscoping)
 s3.sessionToken.0=
+# Optional per-bucket endpoint override.
+s3.endpoint.0=
 
 ## ADLS Storage Config (Multiple configs can be added by incrementing the index)
 adls.storageAccountName.0=

server/src/main/java/io/unitycatalog/server/service/credential/aws/AwsCredentialGenerator.java

@@ -3,6 +3,7 @@
 import io.unitycatalog.server.model.AwsIamRoleResponse;
 import io.unitycatalog.server.persist.dao.CredentialDAO;
 import io.unitycatalog.server.service.credential.CredentialContext;
+import java.net.URI;
 import java.time.Duration;
 import java.util.Optional;
 import java.util.UUID;
@@ -85,7 +86,11 @@ public StsAwsCredentialGenerator(StsClientBuilder builder, S3StorageConfig confi
         credentialsProvider = DefaultCredentialsProvider.create();
       }
 
-      this.stsClient = builder.region(region).credentialsProvider(credentialsProvider).build();
+      StsClientBuilder configured = builder.region(region).credentialsProvider(credentialsProvider);
+      if (config.getEndpoint() != null && !config.getEndpoint().isEmpty()) {
+        configured = configured.endpointOverride(URI.create(config.getEndpoint()));
+      }
+      this.stsClient = configured.build();
       this.staticAwsRoleArn = config.getAwsRoleArn();
     }
 

server/src/main/java/io/unitycatalog/server/service/credential/aws/S3StorageConfig.java

@@ -15,4 +15,5 @@ public class S3StorageConfig {
   private final String secretKey;
   private final String sessionToken;
   private final String credentialGenerator;
+  private final String endpoint;
 }

server/src/main/java/io/unitycatalog/server/service/iceberg/FileIOFactory.java

@@ -25,6 +25,7 @@
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.s3.S3Client;
 
+import java.net.URI;
 import java.util.Map;
 import java.util.Set;
 
@@ -86,20 +87,25 @@ protected S3FileIO getS3FileIO(NormalizedURL location) {
     S3StorageConfig s3StorageConfig = s3Configurations.get(location.getStorageBase());
 
     S3FileIO s3FileIO =
-        new S3FileIO(() -> getS3Client(getAwsCredentialsProvider(location),
-            s3StorageConfig.getRegion()));
+        new S3FileIO(() -> getS3Client(getAwsCredentialsProvider(location), s3StorageConfig));
 
     s3FileIO.initialize(Map.of());
 
     return s3FileIO;
   }
 
-  protected S3Client getS3Client(AwsCredentialsProvider awsCredentialsProvider, String region) {
-    return S3Client.builder()
-        .region(Region.of(region))
+  protected S3Client getS3Client(
+      AwsCredentialsProvider awsCredentialsProvider, S3StorageConfig s3StorageConfig) {
+    String endpoint = s3StorageConfig.getEndpoint();
+    boolean hasEndpoint = endpoint != null && !endpoint.isEmpty();
+    var builder = S3Client.builder()
+        .region(Region.of(s3StorageConfig.getRegion()))
         .credentialsProvider(awsCredentialsProvider)
-        .forcePathStyle(false)
-        .build();
+        .forcePathStyle(hasEndpoint);
+    if (hasEndpoint) {
+      builder = builder.endpointOverride(URI.create(endpoint));
+    }
+    return builder.build();
   }
 
   private AwsCredentialsProvider getAwsCredentialsProvider(NormalizedURL location) {

server/src/main/java/io/unitycatalog/server/utils/ServerProperties.java

@@ -205,6 +205,7 @@ public enum Property {
     AWS_SECRET_KEY("aws.secretKey"),
     AWS_SESSION_TOKEN("aws.sessionToken"),
     AWS_REGION("aws.region"),
+    AWS_ENDPOINT("aws.endpoint"),
     INCLUDE_STACK_TRACE_IN_ERROR("server.include-stacktrace-in-error", "false", BOOLEAN_VALIDATOR);
     // The is not an exhaustive list. Some property keys like s3.bucketPath.0 with a numbering
     // suffix is not included. They are only accessed internally from functions like
@@ -300,6 +301,7 @@ public S3StorageConfig getS3MasterRoleConfiguration() {
         .accessKey(get(Property.AWS_ACCESS_KEY))
         .secretKey(get(Property.AWS_SECRET_KEY))
         // Does not take AWS_SESSION_TOKEN as it's only part of a temporary credential.
+        .endpoint(get(Property.AWS_ENDPOINT))
         .build();
   }
 
@@ -314,6 +316,7 @@ public Map<NormalizedURL, S3StorageConfig> getS3Configurations() {
       String secretKey = getProperty("s3.secretKey." + i);
       String sessionToken = getProperty("s3.sessionToken." + i);
       String credentialGenerator = getProperty("s3.credentialGenerator." + i);
+      String endpoint = getProperty("s3.endpoint." + i);
       if ((bucketPath == null || region == null || awsRoleArn == null)
           && (accessKey == null || secretKey == null || sessionToken == null)) {
         break;
@@ -327,6 +330,7 @@ public Map<NormalizedURL, S3StorageConfig> getS3Configurations() {
               .secretKey(secretKey)
               .sessionToken(sessionToken)
               .credentialGenerator(credentialGenerator)
+              .endpoint(endpoint)
               .build();
       s3BucketConfigMap.put(NormalizedURL.from(bucketPath), s3StorageConfig);
       i++;

server/src/test/java/io/unitycatalog/server/service/credential/CloudCredentialVendorTest.java

@@ -26,6 +26,7 @@
 import io.unitycatalog.server.service.credential.gcp.TestingCredentialGenerator;
 import io.unitycatalog.server.utils.NormalizedURL;
 import io.unitycatalog.server.utils.ServerProperties;
+import java.net.URI;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -287,4 +288,95 @@ public void testVendCredentialWithExternalLocationCredential() {
       assertThat(capturedRequest.externalId()).isEqualTo(expectedExternalId);
     }
   }
+
+  @Test
+  public void testStsClientHonorsAwsEndpointOverride() {
+    final String S3_PATH = "s3://my-bucket/path/to/data";
+    final String CUSTOM_ENDPOINT = "http://host.docker.internal:8333";
+
+    when(serverProperties.getS3Configurations())
+        .thenReturn(
+            Map.of(
+                NormalizedURL.from("s3://my-bucket"),
+                S3StorageConfig.builder()
+                    .bucketPath("s3://my-bucket")
+                    .region("us-east-1")
+                    .awsRoleArn("arn:aws:iam::000000000000:role/UCVendedRole")
+                    .accessKey("admin")
+                    .secretKey("admin")
+                    .endpoint(CUSTOM_ENDPOINT)
+                    .build()));
+
+    StsClient mockStsClient = Mockito.mock(StsClient.class);
+    when(mockStsClient.assumeRole(any(AssumeRoleRequest.class)))
+        .thenReturn(
+            AssumeRoleResponse.builder()
+                .credentials(
+                    Credentials.builder()
+                        .accessKeyId("ak")
+                        .secretAccessKey("sk")
+                        .sessionToken("st")
+                        .build())
+                .build());
+
+    StsClientBuilder mockBuilder = Mockito.mock(StsClientBuilder.class);
+    when(mockBuilder.region(any())).thenReturn(mockBuilder);
+    when(mockBuilder.credentialsProvider(any())).thenReturn(mockBuilder);
+    when(mockBuilder.endpointOverride(any())).thenReturn(mockBuilder);
+    when(mockBuilder.build()).thenReturn(mockStsClient);
+    ArgumentCaptor<URI> endpointCaptor = ArgumentCaptor.forClass(URI.class);
+
+    try (MockedStatic<StsClient> mockedStsClient = Mockito.mockStatic(StsClient.class)) {
+      mockedStsClient.when(StsClient::builder).thenReturn(mockBuilder);
+      AwsCredentialVendor awsCredentialVendor = new AwsCredentialVendor(serverProperties);
+      credentialsOperations = new CloudCredentialVendor(awsCredentialVendor, null, null);
+      vendCredential(S3_PATH, Set.of(CredentialContext.Privilege.SELECT));
+    }
+
+    verify(mockBuilder).endpointOverride(endpointCaptor.capture());
+    assertThat(endpointCaptor.getValue()).isEqualTo(URI.create(CUSTOM_ENDPOINT));
+  }
+
+  @Test
+  public void testStsClientSkipsEndpointOverrideWhenUnset() {
+    final String S3_PATH = "s3://my-bucket/path/to/data";
+
+    when(serverProperties.getS3Configurations())
+        .thenReturn(
+            Map.of(
+                NormalizedURL.from("s3://my-bucket"),
+                S3StorageConfig.builder()
+                    .bucketPath("s3://my-bucket")
+                    .region("us-east-1")
+                    .awsRoleArn("arn:aws:iam::123:role/r")
+                    .accessKey("ak")
+                    .secretKey("sk")
+                    .build()));
+
+    StsClient mockStsClient = Mockito.mock(StsClient.class);
+    when(mockStsClient.assumeRole(any(AssumeRoleRequest.class)))
+        .thenReturn(
+            AssumeRoleResponse.builder()
+                .credentials(
+                    Credentials.builder()
+                        .accessKeyId("ak")
+                        .secretAccessKey("sk")
+                        .sessionToken("st")
+                        .build())
+                .build());
+
+    StsClientBuilder mockBuilder = Mockito.mock(StsClientBuilder.class);
+    when(mockBuilder.region(any())).thenReturn(mockBuilder);
+    when(mockBuilder.credentialsProvider(any())).thenReturn(mockBuilder);
+    when(mockBuilder.build()).thenReturn(mockStsClient);
+
+    try (MockedStatic<StsClient> mockedStsClient = Mockito.mockStatic(StsClient.class)) {
+      mockedStsClient.when(StsClient::builder).thenReturn(mockBuilder);
+      AwsCredentialVendor awsCredentialVendor = new AwsCredentialVendor(serverProperties);
+      credentialsOperations = new CloudCredentialVendor(awsCredentialVendor, null, null);
+      vendCredential(S3_PATH, Set.of(CredentialContext.Privilege.SELECT));
+    }
+
+    verify(mockBuilder, Mockito.never()).endpointOverride(any());
+  }
 }