Skip to content

Commit

Permalink
refactor: FIR-35469 remove account resolve endpoint from jdbc (#452)
Browse files Browse the repository at this point in the history
  • Loading branch information
@stepansergeevitch
stepansergeevitch authored Sep 9, 2024
1 parent 741ccc5 commit b97c8bb
Show file tree
Hide file tree
Showing 6 changed files with 24 additions and 47 deletions.
5 changes: 3 additions & 2 deletions src/integrationTest/java/integration/tests/SystemEngineTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,8 @@ void shouldFailToSelectFromCustomDbUsingSystemEngine() throws SQLException {
}
FireboltException e = assertThrows(FireboltException.class, () -> systemConnection.createStatement().executeQuery("select count(*) from dummy"));
String actualErrorMessage = e.getErrorMessageFromServer().replaceAll("\r?\n", "");
assertTrue(expectedErrorMessages.contains(actualErrorMessage), "Unexpected error message: " + actualErrorMessage);
// Check that at least one error message from expectedErrorMessages is contained in the actual error message
assertTrue(expectedErrorMessages.stream().anyMatch(actualErrorMessage::contains), "Unexpected error message: " + actualErrorMessage);
} finally {
try {
customConnection.createStatement().executeUpdate("DROP TABLE dummy");
Expand Down Expand Up @@ -292,7 +293,7 @@ private String getClientSecret(Connection connection, String serviceAccountName,
FireboltConnection fbConn = (FireboltConnection)connection;
String accessToken = fbConn.getAccessToken().orElseThrow(() -> new IllegalStateException("access token is not found"));
FireboltProperties fbProps = fbConn.getSessionProperties();
URL url = new URL(format("%s/query?output_format=TabSeparatedWithNamesAndTypes&database=%s&account_id=%s", fbProps.getHttpConnectionUrl(), database, fbProps.getAccountId()));
URL url = new URL(format("%s/query?output_format=TabSeparatedWithNamesAndTypes&database=%s", fbProps.getHttpConnectionUrl(), database));
HttpURLConnection con = (HttpURLConnection)url.openConnection();
con.setRequestMethod("POST");
con.setRequestProperty("Content-Type", "application/x-www-form-urlencoded");
Expand Down
2 changes: 1 addition & 1 deletion src/integrationTest/java/integration/tests/TimeoutTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
@CustomLog
class TimeoutTest extends IntegrationTest {
private static final int MIN_TIME_SECONDS = 350;
private static final Map<Integer, Long> SERIES_SIZE = Map.of(1, 80000000000L, 2, 180000000000L);
private static final Map<Integer, Long> SERIES_SIZE = Map.of(1, 80000000000L, 2, 600000000000L);
private long startTime;

@BeforeEach
Expand Down
10 changes: 4 additions & 6 deletions src/main/java/com/firebolt/jdbc/client/account/FireboltAccountRetriever.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,22 +15,20 @@
import static java.net.HttpURLConnection.HTTP_NOT_FOUND;

public class FireboltAccountRetriever<T> extends FireboltClient implements CacheListener {
private static final String URL = "https://%s/web/v3/account/%s/%s";
private static final String URL = "https://%s/web/v3/account/%s/engineUrl";
private final String host;
private final String path;
private final Class<T> type;
private static final Map<String, Object> valueCache = new ConcurrentHashMap<>();

public FireboltAccountRetriever(OkHttpClient httpClient, FireboltConnection connection, String customDrivers, String customClients, String host, String path, Class<T> type) {
public FireboltAccountRetriever(OkHttpClient httpClient, FireboltConnection connection, String customDrivers, String customClients, String host, Class<T> type) {
super(httpClient, connection, customDrivers, customClients);
this.host = host;
this.path = path;
this.type = type;
}

public T retrieve(String accessToken, String accountName) throws SQLException {
try {
String url = format(URL, host, accountName, path);
String url = format(URL, host, accountName);
@SuppressWarnings("unchecked")
T value = (T)valueCache.get(url);
if (value == null) {
Expand All @@ -39,7 +37,7 @@ public T retrieve(String accessToken, String accountName) throws SQLException {
}
return value;
} catch (IOException e) {
throw new FireboltException(String.format("Failed to get %s url for account %s: %s", path, accountName, e.getMessage()), e);
throw new FireboltException(String.format("Failed to get engine url for account %s: %s", accountName, e.getMessage()), e);
}
}

Expand Down
5 changes: 4 additions & 1 deletion src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -314,9 +314,12 @@ private Map<String, String> getAllParameters(FireboltProperties fireboltProperti
}
} else {
if (connection.getInfraVersion() >= 2) {
String engine = fireboltProperties.getEngine();
if (accountId != null) {
params.put(FireboltQueryParameterKey.ACCOUNT_ID.getKey(), accountId);
params.put(FireboltQueryParameterKey.ENGINE.getKey(), fireboltProperties.getEngine());
}
if (engine != null) {
params.put(FireboltQueryParameterKey.ENGINE.getKey(), engine);
}
params.put(FireboltQueryParameterKey.QUERY_LABEL.getKey(), statementInfoWrapper.getLabel()); //QUERY_LABEL
}
Expand Down
14 changes: 4 additions & 10 deletions src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@
public class FireboltConnectionServiceSecret extends FireboltConnection {
private static final String PROTOCOL_VERSION = "2.1";
private final FireboltGatewayUrlService fireboltGatewayUrlService;
private final FireboltAccountIdService fireboltAccountIdService;
private FireboltEngineService fireboltEngineService; // depends on infra version and is discovered during authentication

FireboltConnectionServiceSecret(@NonNull String url,
Expand All @@ -46,7 +45,6 @@ public class FireboltConnectionServiceSecret extends FireboltConnection {
FireboltAccountIdService fireboltAccountIdService) throws SQLException {
super(url, connectionSettings, fireboltAuthenticationService, fireboltStatementService, PROTOCOL_VERSION);
this.fireboltGatewayUrlService = fireboltGatewayUrlService;
this.fireboltAccountIdService = fireboltAccountIdService;
this.fireboltEngineService = fireboltEngineService;
connect();
}
Expand All @@ -55,14 +53,13 @@ public class FireboltConnectionServiceSecret extends FireboltConnection {
FireboltConnectionServiceSecret(@NonNull String url, Properties connectionSettings) throws SQLException {
super(url, connectionSettings, PROTOCOL_VERSION);
OkHttpClient httpClient = getHttpClient(loginProperties);
this.fireboltGatewayUrlService = new FireboltGatewayUrlService(createFireboltAccountRetriever(httpClient,"engineUrl", GatewayUrlResponse.class));
this.fireboltAccountIdService = new FireboltAccountIdService(createFireboltAccountRetriever(httpClient,"resolve", FireboltAccount.class));
this.fireboltGatewayUrlService = new FireboltGatewayUrlService(createFireboltAccountRetriever(httpClient, GatewayUrlResponse.class));
// initialization of fireboltEngineService depends on the infraVersion (the version of engine)
connect();
}

private <T> FireboltAccountRetriever<T> createFireboltAccountRetriever(OkHttpClient httpClient, String path, Class<T> type) {
return new FireboltAccountRetriever<>(httpClient, this, loginProperties.getUserDrivers(), loginProperties.getUserClients(), loginProperties.getHost(), path, type);
private <T> FireboltAccountRetriever<T> createFireboltAccountRetriever(OkHttpClient httpClient, Class<T> type) {
return new FireboltAccountRetriever<>(httpClient, this, loginProperties.getUserDrivers(), loginProperties.getUserClients(), loginProperties.getHost(), type);
}

@Override
Expand Down Expand Up @@ -102,19 +99,16 @@ protected void assertDatabaseExisting(String database) throws SQLException {

private FireboltProperties getSessionPropertiesForSystemEngine(String accessToken, String accountName) throws SQLException {
String systemEngineEndpoint = fireboltGatewayUrlService.getUrl(accessToken, accountName);
FireboltAccount account = fireboltAccountIdService.getValue(accessToken, accountName);
infraVersion = account.getInfraVersion();
infraVersion = 2;
URL systemEngienUrl = UrlUtil.createUrl(systemEngineEndpoint);
Map<String, String> systemEngineUrlUrlParams = UrlUtil.getQueryParameters(systemEngienUrl);
String accountId = systemEngineUrlUrlParams.getOrDefault(ACCOUNT_ID.getKey(), account.getId());
for (Entry<String, String> e : systemEngineUrlUrlParams.entrySet()) {
loginProperties.addProperty(e);
}
return loginProperties
.toBuilder()
.systemEngine(true)
.compress(false)
.accountId(accountId)
.host(systemEngienUrl.getHost())
.build();
}
Expand Down
35 changes: 8 additions & 27 deletions src/test/java/com/firebolt/jdbc/client/gateway/FireboltGatewayUrlClientTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,13 +65,11 @@ class FireboltAccountRetrieverTest {
private FireboltConnection fireboltConnection;

private FireboltAccountRetriever<GatewayUrlResponse> fireboltGatewayUrlClient;
private FireboltAccountRetriever<FireboltAccount> fireboltAccountIdResolver;


@BeforeEach
void setUp() {
fireboltGatewayUrlClient = new FireboltAccountRetriever<>(httpClient, fireboltConnection, null, null, "test-firebolt.io", "engineUrl", GatewayUrlResponse.class);
fireboltAccountIdResolver = new FireboltAccountRetriever<>(httpClient, fireboltConnection, null, null, "test-firebolt.io", "resolve", FireboltAccount.class);
fireboltGatewayUrlClient = new FireboltAccountRetriever<>(httpClient, fireboltConnection, null, null, "test-firebolt.io", GatewayUrlResponse.class);
}

@Test
Expand All @@ -81,21 +79,6 @@ void shouldGetGatewayUrlWhenResponseIsOk() throws SQLException, IOException {
assertEquals(engineUrl, fireboltGatewayUrlClient.retrieve("access_token", "account").getEngineUrl());
}

@Test
void shouldGetAccountId() throws SQLException, IOException {
fireboltAccountIdResolver.cleanup();
FireboltAccount account = new FireboltAccount("12345", "central", 2);
injectMockedResponse(httpClient, HTTP_OK, "{\"id\": \"12345\", \"region\":\"central\", \"infraVersion\":2}");
assertEquals(account, fireboltAccountIdResolver.retrieve("access_token", "account"));
Mockito.verify(httpClient, times(1)).newCall(any());
// Do this again. The response is cached, so the invocation count will remain 1
assertEquals(account, fireboltAccountIdResolver.retrieve("access_token", "account"));
// now clean the cache and call resolve() again. The invocation counter will be incremented.
fireboltAccountIdResolver.cleanup();
assertEquals(account, fireboltAccountIdResolver.retrieve("access_token", "account"));
Mockito.verify(httpClient, times(2)).newCall(any());
}

@Test
void shouldRuntimeExceptionUponRuntimeException() {
when(httpClient.newCall(any())).thenThrow(new IllegalArgumentException("ex"));
Expand All @@ -107,20 +90,19 @@ void shouldThrowFireboltExceptionUponIOException() throws IOException {
Call call = mock(Call.class);
when(httpClient.newCall(any())).thenReturn(call);
when(call.execute()).thenThrow(new IOException("io error"));
assertEquals("Failed to get engineUrl url for account acc: io error", assertThrows(FireboltException.class, () -> fireboltGatewayUrlClient.retrieve("token", "acc")).getMessage());
assertEquals("Failed to get engine url for account acc: io error", assertThrows(FireboltException.class, () -> fireboltGatewayUrlClient.retrieve("token", "acc")).getMessage());
}

@ParameterizedTest(name = "{0}:{1}")
@CsvSource({
"resolve, com.firebolt.jdbc.client.account.FireboltAccount, {}, JSONObject[\"id\"] not found.",
"engineUrl, com.firebolt.jdbc.client.gateway.GatewayUrlResponse, {}, JSONObject[\"engineUrl\"] not found."
"com.firebolt.jdbc.client.gateway.GatewayUrlResponse, {}, JSONObject[\"engineUrl\"] not found."
})
<T> void shouldThrowFireboltExceptionUponWrongJsonFormat(String path, Class<T> clazz, String json, String expectedErrorMessage) throws IOException {
FireboltAccountRetriever<T> fireboltAccountIdResolver = mockAccountRetriever(path, clazz, json);
assertEquals(format("Failed to get %s url for account acc: %s", path, expectedErrorMessage), assertThrows(FireboltException.class, () -> fireboltAccountIdResolver.retrieve("token", "acc")).getMessage());
<T> void shouldThrowFireboltExceptionUponWrongJsonFormat(Class<T> clazz, String json, String expectedErrorMessage) throws IOException {
FireboltAccountRetriever<T> fireboltAccountIdResolver = mockAccountRetriever(clazz, json);
assertEquals(format("Failed to get engine url for account acc: %s", expectedErrorMessage), assertThrows(FireboltException.class, () -> fireboltAccountIdResolver.retrieve("token", "acc")).getMessage());
}

private <T> FireboltAccountRetriever<T> mockAccountRetriever(String path, Class<T> clazz, String json) throws IOException {
private <T> FireboltAccountRetriever<T> mockAccountRetriever(Class<T> clazz, String json) throws IOException {
try (Response response = mock(Response.class)) {
when(response.code()).thenReturn(200);
ResponseBody responseBody = mock(ResponseBody.class);
Expand All @@ -130,7 +112,7 @@ private <T> FireboltAccountRetriever<T> mockAccountRetriever(String path, Class<
Call call = mock();
when(call.execute()).thenReturn(response);
when(okHttpClient.newCall(any())).thenReturn(call);
return new FireboltAccountRetriever<>(okHttpClient, mock(), null, null, "test-firebolt.io", path, clazz);
return new FireboltAccountRetriever<>(okHttpClient, mock(), null, null, "test-firebolt.io", clazz);
}
}

Expand Down Expand Up @@ -165,7 +147,6 @@ private <T> FireboltAccountRetriever<T> mockAccountRetriever(String path, Class<
})
void testFailedAccountDataRetrieving(int statusCode, String errorMessageTemplate) throws IOException {
injectMockedResponse(httpClient, statusCode, null);
assertErrorMessage(fireboltAccountIdResolver, "one", format(errorMessageTemplate, "one", "resolve"));
assertErrorMessage(fireboltGatewayUrlClient, "two", format(errorMessageTemplate, "two", "engineUrl"));
}

Expand Down

0 comments on commit b97c8bb

Please sign in to comment.