From 673bd02ce859bbd3977c4ac44cf803daf7e6aa14 Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Wed, 16 Apr 2025 18:23:44 -0700 Subject: [PATCH 1/5] Update the StorageConfiguration to invoke singleton client objects, and add a test --- .../service/storage/StorageConfiguration.java | 9 +- .../service/util/MemoizingSupplier.java | 48 ++++++ .../storage/StorageConfigurationTest.java | 143 ++++++++++++++++++ 3 files changed, 196 insertions(+), 4 deletions(-) create mode 100644 service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java create mode 100644 service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java diff --git a/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java b/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java index 711cf032e9..0583d53d61 100644 --- a/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java +++ b/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java @@ -25,6 +25,7 @@ import java.time.Instant; import java.util.*; import java.util.function.Supplier; +import org.apache.polaris.service.util.MemoizingSupplier; import org.slf4j.LoggerFactory; import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; @@ -60,7 +61,7 @@ public interface StorageConfiguration { Optional gcpAccessTokenLifespan(); default Supplier stsClientSupplier() { - return () -> { + return MemoizingSupplier.memoize(() -> { StsClientBuilder stsClientBuilder = StsClient.builder(); if (awsAccessKey().isPresent() && awsSecretKey().isPresent()) { LoggerFactory.getLogger(StorageConfiguration.class) @@ -71,11 +72,11 @@ default Supplier stsClientSupplier() { stsClientBuilder.credentialsProvider(awsCredentialsProvider); } return stsClientBuilder.build(); - }; + }); } default Supplier gcpCredentialsSupplier() { - return () -> { + return MemoizingSupplier.memoize(() -> { if (gcpAccessToken().isEmpty()) { try { return GoogleCredentials.getApplicationDefault(); @@ -92,6 +93,6 @@ default Supplier gcpCredentialsSupplier() { .toEpochMilli())); return GoogleCredentials.create(accessToken); } - }; + }); } } diff --git a/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java b/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java new file mode 100644 index 0000000000..a053d10312 --- /dev/null +++ b/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java @@ -0,0 +1,48 @@ +/* + * 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.polaris.service.util; + +import java.util.function.Supplier; + +public final class MemoizingSupplier implements Supplier { + private final Supplier delegate; + private volatile T value; + + private MemoizingSupplier(Supplier delegate) { + this.delegate = delegate; + } + + public static Supplier memoize(Supplier delegate) { + return new MemoizingSupplier<>(delegate); + } + + @Override + public T get() { + if (value == null) { + synchronized (this) { + if (value == null) { + value = delegate.get(); + } + } + } + return value; + } +} + diff --git a/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java b/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java new file mode 100644 index 0000000000..e8c5f1c9d1 --- /dev/null +++ b/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java @@ -0,0 +1,143 @@ +/* + * 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.polaris.service.storage; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.*; + +import com.google.auth.oauth2.AccessToken; +import com.google.auth.oauth2.GoogleCredentials; +import java.time.Duration; +import java.time.Instant; +import java.util.Optional; +import java.util.function.Supplier; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.services.sts.StsClient; + +public class StorageConfigurationTest { + + private static final String TEST_ACCESS_KEY = "test-access-key"; + private static final String TEST_SECRET_KEY = "test-secret-key"; + private static final String TEST_GCP_TOKEN = "ya29.test-token"; + private static final Duration TEST_TOKEN_LIFESPAN = Duration.ofMinutes(20); + + private StorageConfiguration configWithAwsCredentialsAndGcpToken() { + return new StorageConfiguration() { + @Override + public Optional awsAccessKey() { + return Optional.of(TEST_ACCESS_KEY); + } + + @Override + public Optional awsSecretKey() { + return Optional.of(TEST_SECRET_KEY); + } + + @Override + public Optional gcpAccessToken() { + return Optional.of(TEST_GCP_TOKEN); + } + + @Override + public Optional gcpAccessTokenLifespan() { + return Optional.of(TEST_TOKEN_LIFESPAN); + } + }; + } + + private StorageConfiguration configWithoutGcpToken() { + return new StorageConfiguration() { + @Override + public Optional awsAccessKey() { + return Optional.empty(); + } + + @Override + public Optional awsSecretKey() { + return Optional.empty(); + } + + @Override + public Optional gcpAccessToken() { + return Optional.empty(); + } + + @Override + public Optional gcpAccessTokenLifespan() { + return Optional.empty(); + } + }; + } + + @Test + public void testSingletonStsClientWithStaticCredentials() { + Supplier supplier = configWithAwsCredentialsAndGcpToken().stsClientSupplier(); + + StsClient client1 = supplier.get(); + StsClient client2 = supplier.get(); + + assertThat(client1).isSameAs(client2); + assertThat(client1).isNotNull(); + + assertThat(client1.serviceClientConfiguration().credentialsProvider()) + .isInstanceOf(StaticCredentialsProvider.class); + StaticCredentialsProvider credentialsProvider = + (StaticCredentialsProvider) client1.serviceClientConfiguration().credentialsProvider(); + assertThat(credentialsProvider.resolveCredentials().accessKeyId()) + .isEqualTo(TEST_ACCESS_KEY); + assertThat(credentialsProvider.resolveCredentials().secretAccessKey()) + .isEqualTo(TEST_SECRET_KEY); + } + + @Test + public void testCreateGcpCredentialsFromStaticToken() { + Supplier supplier = + configWithAwsCredentialsAndGcpToken().gcpCredentialsSupplier(); + + GoogleCredentials credentials = supplier.get(); + assertThat(credentials).isNotNull(); + + AccessToken accessToken = credentials.getAccessToken(); + assertThat(accessToken).isNotNull(); + assertThat(accessToken.getTokenValue()).isEqualTo(TEST_GCP_TOKEN); + long expectedExpiry = Instant.now().plus(Duration.ofMinutes(20)).toEpochMilli(); + long actualExpiry = accessToken.getExpirationTime().getTime(); + assertThat(actualExpiry).isBetween(expectedExpiry - 500, expectedExpiry + 500); + } + + @Test + public void testGcpCredentialsFromDefault() { + GoogleCredentials mockDefaultCreds = mock(GoogleCredentials.class); + + try (MockedStatic mockedStatic = + Mockito.mockStatic(GoogleCredentials.class)) { + + mockedStatic.when(GoogleCredentials::getApplicationDefault).thenReturn(mockDefaultCreds); + + Supplier supplier = configWithoutGcpToken().gcpCredentialsSupplier(); + GoogleCredentials result = supplier.get(); + + assertThat(result).isSameAs(mockDefaultCreds); + mockedStatic.verify(GoogleCredentials::getApplicationDefault, times(1)); + } + } +} From a4ef99bc44adf714fd5a25c17495003e96ba36d7 Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Wed, 16 Apr 2025 19:17:58 -0700 Subject: [PATCH 2/5] Fix formatting --- .../service/storage/StorageConfiguration.java | 62 +++--- .../service/util/MemoizingSupplier.java | 34 ++- .../storage/StorageConfigurationTest.java | 207 +++++++++--------- 3 files changed, 151 insertions(+), 152 deletions(-) diff --git a/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java b/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java index 0583d53d61..cdafcbb187 100644 --- a/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java +++ b/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java @@ -61,38 +61,40 @@ public interface StorageConfiguration { Optional gcpAccessTokenLifespan(); default Supplier stsClientSupplier() { - return MemoizingSupplier.memoize(() -> { - StsClientBuilder stsClientBuilder = StsClient.builder(); - if (awsAccessKey().isPresent() && awsSecretKey().isPresent()) { - LoggerFactory.getLogger(StorageConfiguration.class) - .warn("Using hard-coded AWS credentials - this is not recommended for production"); - StaticCredentialsProvider awsCredentialsProvider = - StaticCredentialsProvider.create( - AwsBasicCredentials.create(awsAccessKey().get(), awsSecretKey().get())); - stsClientBuilder.credentialsProvider(awsCredentialsProvider); - } - return stsClientBuilder.build(); - }); + return MemoizingSupplier.memoize( + () -> { + StsClientBuilder stsClientBuilder = StsClient.builder(); + if (awsAccessKey().isPresent() && awsSecretKey().isPresent()) { + LoggerFactory.getLogger(StorageConfiguration.class) + .warn("Using hard-coded AWS credentials - this is not recommended for production"); + StaticCredentialsProvider awsCredentialsProvider = + StaticCredentialsProvider.create( + AwsBasicCredentials.create(awsAccessKey().get(), awsSecretKey().get())); + stsClientBuilder.credentialsProvider(awsCredentialsProvider); + } + return stsClientBuilder.build(); + }); } default Supplier gcpCredentialsSupplier() { - return MemoizingSupplier.memoize(() -> { - if (gcpAccessToken().isEmpty()) { - try { - return GoogleCredentials.getApplicationDefault(); - } catch (IOException e) { - throw new RuntimeException("Failed to get GCP credentials", e); - } - } else { - AccessToken accessToken = - new AccessToken( - gcpAccessToken().get(), - new Date( - Instant.now() - .plus(gcpAccessTokenLifespan().orElse(DEFAULT_TOKEN_LIFESPAN)) - .toEpochMilli())); - return GoogleCredentials.create(accessToken); - } - }); + return MemoizingSupplier.memoize( + () -> { + if (gcpAccessToken().isEmpty()) { + try { + return GoogleCredentials.getApplicationDefault(); + } catch (IOException e) { + throw new RuntimeException("Failed to get GCP credentials", e); + } + } else { + AccessToken accessToken = + new AccessToken( + gcpAccessToken().get(), + new Date( + Instant.now() + .plus(gcpAccessTokenLifespan().orElse(DEFAULT_TOKEN_LIFESPAN)) + .toEpochMilli())); + return GoogleCredentials.create(accessToken); + } + }); } } diff --git a/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java b/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java index a053d10312..148e97efc7 100644 --- a/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java +++ b/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java @@ -16,33 +16,31 @@ * specific language governing permissions and limitations * under the License. */ - package org.apache.polaris.service.util; import java.util.function.Supplier; public final class MemoizingSupplier implements Supplier { - private final Supplier delegate; - private volatile T value; + private final Supplier delegate; + private volatile T value; - private MemoizingSupplier(Supplier delegate) { - this.delegate = delegate; - } + private MemoizingSupplier(Supplier delegate) { + this.delegate = delegate; + } - public static Supplier memoize(Supplier delegate) { - return new MemoizingSupplier<>(delegate); - } + public static Supplier memoize(Supplier delegate) { + return new MemoizingSupplier<>(delegate); + } - @Override - public T get() { + @Override + public T get() { + if (value == null) { + synchronized (this) { if (value == null) { - synchronized (this) { - if (value == null) { - value = delegate.get(); - } - } + value = delegate.get(); } - return value; + } } + return value; + } } - diff --git a/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java b/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java index e8c5f1c9d1..629400ca39 100644 --- a/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java +++ b/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java @@ -35,109 +35,108 @@ public class StorageConfigurationTest { - private static final String TEST_ACCESS_KEY = "test-access-key"; - private static final String TEST_SECRET_KEY = "test-secret-key"; - private static final String TEST_GCP_TOKEN = "ya29.test-token"; - private static final Duration TEST_TOKEN_LIFESPAN = Duration.ofMinutes(20); - - private StorageConfiguration configWithAwsCredentialsAndGcpToken() { - return new StorageConfiguration() { - @Override - public Optional awsAccessKey() { - return Optional.of(TEST_ACCESS_KEY); - } - - @Override - public Optional awsSecretKey() { - return Optional.of(TEST_SECRET_KEY); - } - - @Override - public Optional gcpAccessToken() { - return Optional.of(TEST_GCP_TOKEN); - } - - @Override - public Optional gcpAccessTokenLifespan() { - return Optional.of(TEST_TOKEN_LIFESPAN); - } - }; - } - - private StorageConfiguration configWithoutGcpToken() { - return new StorageConfiguration() { - @Override - public Optional awsAccessKey() { - return Optional.empty(); - } - - @Override - public Optional awsSecretKey() { - return Optional.empty(); - } - - @Override - public Optional gcpAccessToken() { - return Optional.empty(); - } - - @Override - public Optional gcpAccessTokenLifespan() { - return Optional.empty(); - } - }; - } - - @Test - public void testSingletonStsClientWithStaticCredentials() { - Supplier supplier = configWithAwsCredentialsAndGcpToken().stsClientSupplier(); - - StsClient client1 = supplier.get(); - StsClient client2 = supplier.get(); - - assertThat(client1).isSameAs(client2); - assertThat(client1).isNotNull(); - - assertThat(client1.serviceClientConfiguration().credentialsProvider()) - .isInstanceOf(StaticCredentialsProvider.class); - StaticCredentialsProvider credentialsProvider = - (StaticCredentialsProvider) client1.serviceClientConfiguration().credentialsProvider(); - assertThat(credentialsProvider.resolveCredentials().accessKeyId()) - .isEqualTo(TEST_ACCESS_KEY); - assertThat(credentialsProvider.resolveCredentials().secretAccessKey()) - .isEqualTo(TEST_SECRET_KEY); - } - - @Test - public void testCreateGcpCredentialsFromStaticToken() { - Supplier supplier = - configWithAwsCredentialsAndGcpToken().gcpCredentialsSupplier(); - - GoogleCredentials credentials = supplier.get(); - assertThat(credentials).isNotNull(); - - AccessToken accessToken = credentials.getAccessToken(); - assertThat(accessToken).isNotNull(); - assertThat(accessToken.getTokenValue()).isEqualTo(TEST_GCP_TOKEN); - long expectedExpiry = Instant.now().plus(Duration.ofMinutes(20)).toEpochMilli(); - long actualExpiry = accessToken.getExpirationTime().getTime(); - assertThat(actualExpiry).isBetween(expectedExpiry - 500, expectedExpiry + 500); - } - - @Test - public void testGcpCredentialsFromDefault() { - GoogleCredentials mockDefaultCreds = mock(GoogleCredentials.class); - - try (MockedStatic mockedStatic = - Mockito.mockStatic(GoogleCredentials.class)) { - - mockedStatic.when(GoogleCredentials::getApplicationDefault).thenReturn(mockDefaultCreds); - - Supplier supplier = configWithoutGcpToken().gcpCredentialsSupplier(); - GoogleCredentials result = supplier.get(); - - assertThat(result).isSameAs(mockDefaultCreds); - mockedStatic.verify(GoogleCredentials::getApplicationDefault, times(1)); - } + private static final String TEST_ACCESS_KEY = "test-access-key"; + private static final String TEST_SECRET_KEY = "test-secret-key"; + private static final String TEST_GCP_TOKEN = "ya29.test-token"; + private static final Duration TEST_TOKEN_LIFESPAN = Duration.ofMinutes(20); + + private StorageConfiguration configWithAwsCredentialsAndGcpToken() { + return new StorageConfiguration() { + @Override + public Optional awsAccessKey() { + return Optional.of(TEST_ACCESS_KEY); + } + + @Override + public Optional awsSecretKey() { + return Optional.of(TEST_SECRET_KEY); + } + + @Override + public Optional gcpAccessToken() { + return Optional.of(TEST_GCP_TOKEN); + } + + @Override + public Optional gcpAccessTokenLifespan() { + return Optional.of(TEST_TOKEN_LIFESPAN); + } + }; + } + + private StorageConfiguration configWithoutGcpToken() { + return new StorageConfiguration() { + @Override + public Optional awsAccessKey() { + return Optional.empty(); + } + + @Override + public Optional awsSecretKey() { + return Optional.empty(); + } + + @Override + public Optional gcpAccessToken() { + return Optional.empty(); + } + + @Override + public Optional gcpAccessTokenLifespan() { + return Optional.empty(); + } + }; + } + + @Test + public void testSingletonStsClientWithStaticCredentials() { + Supplier supplier = configWithAwsCredentialsAndGcpToken().stsClientSupplier(); + + StsClient client1 = supplier.get(); + StsClient client2 = supplier.get(); + + assertThat(client1).isSameAs(client2); + assertThat(client1).isNotNull(); + + assertThat(client1.serviceClientConfiguration().credentialsProvider()) + .isInstanceOf(StaticCredentialsProvider.class); + StaticCredentialsProvider credentialsProvider = + (StaticCredentialsProvider) client1.serviceClientConfiguration().credentialsProvider(); + assertThat(credentialsProvider.resolveCredentials().accessKeyId()).isEqualTo(TEST_ACCESS_KEY); + assertThat(credentialsProvider.resolveCredentials().secretAccessKey()) + .isEqualTo(TEST_SECRET_KEY); + } + + @Test + public void testCreateGcpCredentialsFromStaticToken() { + Supplier supplier = + configWithAwsCredentialsAndGcpToken().gcpCredentialsSupplier(); + + GoogleCredentials credentials = supplier.get(); + assertThat(credentials).isNotNull(); + + AccessToken accessToken = credentials.getAccessToken(); + assertThat(accessToken).isNotNull(); + assertThat(accessToken.getTokenValue()).isEqualTo(TEST_GCP_TOKEN); + long expectedExpiry = Instant.now().plus(Duration.ofMinutes(20)).toEpochMilli(); + long actualExpiry = accessToken.getExpirationTime().getTime(); + assertThat(actualExpiry).isBetween(expectedExpiry - 500, expectedExpiry + 500); + } + + @Test + public void testGcpCredentialsFromDefault() { + GoogleCredentials mockDefaultCreds = mock(GoogleCredentials.class); + + try (MockedStatic mockedStatic = + Mockito.mockStatic(GoogleCredentials.class)) { + + mockedStatic.when(GoogleCredentials::getApplicationDefault).thenReturn(mockDefaultCreds); + + Supplier supplier = configWithoutGcpToken().gcpCredentialsSupplier(); + GoogleCredentials result = supplier.get(); + + assertThat(result).isSameAs(mockDefaultCreds); + mockedStatic.verify(GoogleCredentials::getApplicationDefault, times(1)); } + } } From d1c643c9c38b7a002506e932718a3090e94be18f Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Thu, 17 Apr 2025 09:48:17 -0700 Subject: [PATCH 3/5] using guava suppliers --- .../service/storage/StorageConfiguration.java | 6 +-- .../service/util/MemoizingSupplier.java | 46 ------------------- 2 files changed, 3 insertions(+), 49 deletions(-) delete mode 100644 service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java diff --git a/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java b/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java index cdafcbb187..188a2b85c1 100644 --- a/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java +++ b/service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java @@ -20,12 +20,12 @@ import com.google.auth.oauth2.AccessToken; import com.google.auth.oauth2.GoogleCredentials; +import com.google.common.base.Suppliers; import java.io.IOException; import java.time.Duration; import java.time.Instant; import java.util.*; import java.util.function.Supplier; -import org.apache.polaris.service.util.MemoizingSupplier; import org.slf4j.LoggerFactory; import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; @@ -61,7 +61,7 @@ public interface StorageConfiguration { Optional gcpAccessTokenLifespan(); default Supplier stsClientSupplier() { - return MemoizingSupplier.memoize( + return Suppliers.memoize( () -> { StsClientBuilder stsClientBuilder = StsClient.builder(); if (awsAccessKey().isPresent() && awsSecretKey().isPresent()) { @@ -77,7 +77,7 @@ default Supplier stsClientSupplier() { } default Supplier gcpCredentialsSupplier() { - return MemoizingSupplier.memoize( + return Suppliers.memoize( () -> { if (gcpAccessToken().isEmpty()) { try { diff --git a/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java b/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java deleted file mode 100644 index 148e97efc7..0000000000 --- a/service/common/src/main/java/org/apache/polaris/service/util/MemoizingSupplier.java +++ /dev/null @@ -1,46 +0,0 @@ -/* - * 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.polaris.service.util; - -import java.util.function.Supplier; - -public final class MemoizingSupplier implements Supplier { - private final Supplier delegate; - private volatile T value; - - private MemoizingSupplier(Supplier delegate) { - this.delegate = delegate; - } - - public static Supplier memoize(Supplier delegate) { - return new MemoizingSupplier<>(delegate); - } - - @Override - public T get() { - if (value == null) { - synchronized (this) { - if (value == null) { - value = delegate.get(); - } - } - } - return value; - } -} From 58e16105ed20a722e17d1eeec2ac2301f1218b7c Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Thu, 17 Apr 2025 14:05:06 -0700 Subject: [PATCH 4/5] Add aws region --- .../storage/StorageConfigurationTest.java | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java b/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java index 629400ca39..b8c99637ae 100644 --- a/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java +++ b/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java @@ -27,6 +27,8 @@ import java.time.Instant; import java.util.Optional; import java.util.function.Supplier; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.MockedStatic; import org.mockito.Mockito; @@ -36,10 +38,21 @@ public class StorageConfigurationTest { private static final String TEST_ACCESS_KEY = "test-access-key"; - private static final String TEST_SECRET_KEY = "test-secret-key"; private static final String TEST_GCP_TOKEN = "ya29.test-token"; + private static final String TEST_REGION = "us-west-2"; + private static final String TEST_SECRET_KEY = "test-secret-key"; private static final Duration TEST_TOKEN_LIFESPAN = Duration.ofMinutes(20); + @BeforeEach + public void setUpAwsRegion() { + System.setProperty("aws.region", TEST_REGION); + } + + @AfterEach + public void tearDownAwsRegion() { + System.clearProperty("aws.region"); + } + private StorageConfiguration configWithAwsCredentialsAndGcpToken() { return new StorageConfiguration() { @Override @@ -100,6 +113,7 @@ public void testSingletonStsClientWithStaticCredentials() { assertThat(client1.serviceClientConfiguration().credentialsProvider()) .isInstanceOf(StaticCredentialsProvider.class); + assertThat(client1.serviceClientConfiguration().region().toString()).isEqualTo(TEST_REGION); StaticCredentialsProvider credentialsProvider = (StaticCredentialsProvider) client1.serviceClientConfiguration().credentialsProvider(); assertThat(credentialsProvider.resolveCredentials().accessKeyId()).isEqualTo(TEST_ACCESS_KEY); From 5c215b6f8460eae5bec6bfdbf34d14b882b9aff2 Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Thu, 17 Apr 2025 14:29:10 -0700 Subject: [PATCH 5/5] Cleanup and mock test --- .../storage/StorageConfigurationTest.java | 56 +++++++++---------- 1 file changed, 27 insertions(+), 29 deletions(-) diff --git a/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java b/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java index b8c99637ae..61eb174250 100644 --- a/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java +++ b/service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java @@ -27,32 +27,21 @@ import java.time.Instant; import java.util.Optional; import java.util.function.Supplier; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.mockito.MockedStatic; import org.mockito.Mockito; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.StsClientBuilder; public class StorageConfigurationTest { private static final String TEST_ACCESS_KEY = "test-access-key"; private static final String TEST_GCP_TOKEN = "ya29.test-token"; - private static final String TEST_REGION = "us-west-2"; private static final String TEST_SECRET_KEY = "test-secret-key"; private static final Duration TEST_TOKEN_LIFESPAN = Duration.ofMinutes(20); - @BeforeEach - public void setUpAwsRegion() { - System.setProperty("aws.region", TEST_REGION); - } - - @AfterEach - public void tearDownAwsRegion() { - System.clearProperty("aws.region"); - } - private StorageConfiguration configWithAwsCredentialsAndGcpToken() { return new StorageConfiguration() { @Override @@ -103,22 +92,31 @@ public Optional gcpAccessTokenLifespan() { @Test public void testSingletonStsClientWithStaticCredentials() { - Supplier supplier = configWithAwsCredentialsAndGcpToken().stsClientSupplier(); - - StsClient client1 = supplier.get(); - StsClient client2 = supplier.get(); - - assertThat(client1).isSameAs(client2); - assertThat(client1).isNotNull(); - - assertThat(client1.serviceClientConfiguration().credentialsProvider()) - .isInstanceOf(StaticCredentialsProvider.class); - assertThat(client1.serviceClientConfiguration().region().toString()).isEqualTo(TEST_REGION); - StaticCredentialsProvider credentialsProvider = - (StaticCredentialsProvider) client1.serviceClientConfiguration().credentialsProvider(); - assertThat(credentialsProvider.resolveCredentials().accessKeyId()).isEqualTo(TEST_ACCESS_KEY); - assertThat(credentialsProvider.resolveCredentials().secretAccessKey()) - .isEqualTo(TEST_SECRET_KEY); + StsClientBuilder mockBuilder = mock(StsClientBuilder.class); + StsClient mockStsClient = mock(StsClient.class); + ArgumentCaptor credsCaptor = + ArgumentCaptor.forClass(StaticCredentialsProvider.class); + + when(mockBuilder.credentialsProvider(credsCaptor.capture())).thenReturn(mockBuilder); + when(mockBuilder.region(any())).thenReturn(mockBuilder); + when(mockBuilder.build()).thenReturn(mockStsClient); + + try (MockedStatic staticMock = Mockito.mockStatic(StsClient.class)) { + staticMock.when(StsClient::builder).thenReturn(mockBuilder); + + StorageConfiguration config = configWithAwsCredentialsAndGcpToken(); + Supplier supplier = config.stsClientSupplier(); + StsClient client1 = supplier.get(); + StsClient client2 = supplier.get(); + + assertThat(client1).isSameAs(client2); + assertThat(client1).isNotNull(); + + StaticCredentialsProvider credentialsProvider = credsCaptor.getValue(); + assertThat(credentialsProvider.resolveCredentials().accessKeyId()).isEqualTo(TEST_ACCESS_KEY); + assertThat(credentialsProvider.resolveCredentials().secretAccessKey()) + .isEqualTo(TEST_SECRET_KEY); + } } @Test