Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

GCP: Use catalog endpoint as base when refreshing OAuth2 token #12638

Merged
merged 1 commit into from
Mar 28, 2025
Merged

GCP: Use catalog endpoint as base when refreshing OAuth2 token #12638

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions gcp/src/main/java/org/apache/iceberg/gcp/gcs/OAuth2RefreshCredentialsHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.iceberg.CatalogProperties;
import org.apache.iceberg.gcp.GCPProperties;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
Expand All @@ -41,6 +42,9 @@
public class OAuth2RefreshCredentialsHandler
implements OAuth2CredentialsWithRefresh.OAuth2RefreshHandler, AutoCloseable {
private final Map<String, String> properties;
private final String credentialsEndpoint;
// will be used to refresh the OAuth2 token
private final String catalogEndpoint;
private volatile HTTPClient client;
private AuthManager authManager;
private AuthSession authSession;
Expand All @@ -49,6 +53,11 @@ private OAuth2RefreshCredentialsHandler(Map<String, String> properties) {
Preconditions.checkArgument(
null != properties.get(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT),
"Invalid credentials endpoint: null");
Preconditions.checkArgument(
null != properties.get(CatalogProperties.URI), "Invalid catalog endpoint: null");
this.credentialsEndpoint =
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why aren't we doing the same thing here that we do in ADLS?

RESTUtil.resolveEndpoint(
            properties.get(CatalogProperties.URI),
            properties.get(ADLS_REFRESH_CREDENTIALS_ENDPOINT));

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we're doing it for GCP in GCPProperties:

iceberg/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java

Lines 109 to 112 in 978189c

gcsOauth2RefreshCredentialsEndpoint =
RESTUtil.resolveEndpoint(
properties.get(CatalogProperties.URI),
properties.get(GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT));

For ADLS we're also doing the same in AzureProperties:

iceberg/azure/src/main/java/org/apache/iceberg/azure/AzureProperties.java

Lines 89 to 92 in 68f8053

this.adlsRefreshCredentialsEndpoint =
RESTUtil.resolveEndpoint(
properties.get(CatalogProperties.URI),
properties.get(ADLS_REFRESH_CREDENTIALS_ENDPOINT));

For S3 it happens in AwsClientProperties:

iceberg/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java

Lines 109 to 111 in 7d0395d

this.refreshCredentialsEndpoint =
RESTUtil.resolveEndpoint(
properties.get(CatalogProperties.URI), properties.get(REFRESH_CREDENTIALS_ENDPOINT));

properties.get(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT);
this.catalogEndpoint = properties.get(CatalogProperties.URI);
this.properties = properties;
}

Expand All @@ -58,7 +67,7 @@ public AccessToken refreshAccessToken() {
LoadCredentialsResponse response =
httpClient()
.get(
properties.get(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT),
credentialsEndpoint,
null,
LoadCredentialsResponse.class,
Map.of(),
Expand Down Expand Up @@ -99,10 +108,7 @@ private RESTClient httpClient() {
synchronized (this) {
if (null == client) {
authManager = AuthManagers.loadAuthManager("gcs-credentials-refresh", properties);
HTTPClient httpClient =
HTTPClient.builder(properties)
.uri(properties.get(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT))
.build();
HTTPClient httpClient = HTTPClient.builder(properties).uri(catalogEndpoint).build();
authSession = authManager.catalogSession(httpClient, properties);
client = httpClient.withAuthSession(authSession);
}
Expand Down
3 changes: 3 additions & 0 deletions gcp/src/test/java/org/apache/iceberg/gcp/gcs/GCSFileIOTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
import java.util.Random;
import java.util.stream.StreamSupport;
import org.apache.hadoop.conf.Configuration;
import org.apache.iceberg.CatalogProperties;
import org.apache.iceberg.TestHelpers;
import org.apache.iceberg.common.DynMethods;
import org.apache.iceberg.gcp.GCPProperties;
Expand Down Expand Up @@ -238,6 +239,8 @@ public void refreshCredentialsEndpointSet() {
try (GCSFileIO fileIO = new GCSFileIO()) {
fileIO.initialize(
ImmutableMap.of(
CatalogProperties.URI,
"http://catalog-endpoint",
GCS_OAUTH2_TOKEN,
"gcsToken",
GCS_OAUTH2_TOKEN_EXPIRES_AT,
Expand Down
57 changes: 35 additions & 22 deletions gcp/src/test/java/org/apache/iceberg/gcp/gcs/OAuth2RefreshCredentialsHandlerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
import com.google.auth.oauth2.AccessToken;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Map;
import org.apache.iceberg.CatalogProperties;
import org.apache.iceberg.exceptions.BadRequestException;
import org.apache.iceberg.exceptions.RESTException;
import org.apache.iceberg.gcp.GCPProperties;
Expand All @@ -45,7 +47,15 @@

public class OAuth2RefreshCredentialsHandlerTest {
private static final int PORT = 3333;
private static final String URI = String.format("http://127.0.0.1:%d/v1/credentials", PORT);
private static final String CREDENTIALS_URI =
String.format("http://127.0.0.1:%d/v1/credentials", PORT);
private static final String CATALOG_URI = String.format("http://127.0.0.1:%d/v1/", PORT);
private static final Map<String, String> PROPERTIES =
ImmutableMap.of(
GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT,
CREDENTIALS_URI,
CatalogProperties.URI,
CATALOG_URI);
private static ClientAndServer mockServer;

@BeforeAll
Expand All @@ -65,18 +75,33 @@ public void before() {

@Test
public void invalidOrMissingUri() {
assertThatThrownBy(() -> OAuth2RefreshCredentialsHandler.create(ImmutableMap.of()))
assertThatThrownBy(
() ->
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(CatalogProperties.URI, CATALOG_URI)))
.isInstanceOf(IllegalArgumentException.class)
.hasMessage("Invalid credentials endpoint: null");

assertThatThrownBy(
() ->
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(
GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT, CREDENTIALS_URI)))
.isInstanceOf(IllegalArgumentException.class)
.hasMessage("Invalid catalog endpoint: null");

assertThatThrownBy(
() ->
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(
GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT, "invalid uri"))
GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT,
"invalid uri",
CatalogProperties.URI,
CATALOG_URI))
.refreshAccessToken())
.isInstanceOf(RESTException.class)
.hasMessageStartingWith("Failed to create request URI from base invalid uri");
.hasMessageStartingWith(
"Failed to create request URI from base %sinvalid uri", CATALOG_URI);
}

@Test
Expand All @@ -87,9 +112,7 @@ public void badRequest() {
HttpResponse mockResponse = HttpResponse.response().withStatusCode(400);
mockServer.when(mockRequest).respond(mockResponse);

OAuth2RefreshCredentialsHandler handler =
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT, URI));
OAuth2RefreshCredentialsHandler handler = OAuth2RefreshCredentialsHandler.create(PROPERTIES);

assertThatThrownBy(handler::refreshAccessToken)
.isInstanceOf(BadRequestException.class)
Expand All @@ -108,9 +131,7 @@ public void noGcsCredentialInResponse() {
.withStatusCode(200);
mockServer.when(mockRequest).respond(mockResponse);

OAuth2RefreshCredentialsHandler handler =
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT, URI));
OAuth2RefreshCredentialsHandler handler = OAuth2RefreshCredentialsHandler.create(PROPERTIES);

assertThatThrownBy(handler::refreshAccessToken)
.isInstanceOf(IllegalStateException.class)
Expand All @@ -134,9 +155,7 @@ public void noGcsToken() {
.withStatusCode(200);
mockServer.when(mockRequest).respond(mockResponse);

OAuth2RefreshCredentialsHandler handler =
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT, URI));
OAuth2RefreshCredentialsHandler handler = OAuth2RefreshCredentialsHandler.create(PROPERTIES);

assertThatThrownBy(handler::refreshAccessToken)
.isInstanceOf(IllegalStateException.class)
Expand All @@ -160,9 +179,7 @@ public void tokenWithoutExpiration() {
.withStatusCode(200);
mockServer.when(mockRequest).respond(mockResponse);

OAuth2RefreshCredentialsHandler handler =
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT, URI));
OAuth2RefreshCredentialsHandler handler = OAuth2RefreshCredentialsHandler.create(PROPERTIES);

assertThatThrownBy(handler::refreshAccessToken)
.isInstanceOf(IllegalStateException.class)
Expand Down Expand Up @@ -191,9 +208,7 @@ public void tokenWithExpiration() {
.withStatusCode(200);
mockServer.when(mockRequest).respond(mockResponse);

OAuth2RefreshCredentialsHandler handler =
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT, URI));
OAuth2RefreshCredentialsHandler handler = OAuth2RefreshCredentialsHandler.create(PROPERTIES);

AccessToken accessToken = handler.refreshAccessToken();
assertThat(accessToken.getTokenValue())
Expand Down Expand Up @@ -253,9 +268,7 @@ public void multipleGcsCredentials() {
.withStatusCode(200);
mockServer.when(mockRequest).respond(mockResponse);

OAuth2RefreshCredentialsHandler handler =
OAuth2RefreshCredentialsHandler.create(
ImmutableMap.of(GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT, URI));
OAuth2RefreshCredentialsHandler handler = OAuth2RefreshCredentialsHandler.create(PROPERTIES);

assertThatThrownBy(handler::refreshAccessToken)
.isInstanceOf(IllegalStateException.class)
Expand Down