From 245cf2ed31f61748fd99c79a0e49f4366331d73e Mon Sep 17 00:00:00 2001
From: Stepan Burlakov <stepansergeevitch@gmail.com>
Date: Mon, 2 Sep 2024 17:09:55 +0300
Subject: [PATCH] remove account resolve call

---
 .../account/FireboltAccountRetriever.java     | 10 +++---
 .../FireboltConnectionServiceSecret.java      | 15 +++-----
 .../gateway/FireboltGatewayUrlClientTest.java | 35 +++++--------------
 3 files changed, 17 insertions(+), 43 deletions(-)

diff --git a/src/main/java/com/firebolt/jdbc/client/account/FireboltAccountRetriever.java b/src/main/java/com/firebolt/jdbc/client/account/FireboltAccountRetriever.java
index 8b4876af4..3e578829f 100644
--- a/src/main/java/com/firebolt/jdbc/client/account/FireboltAccountRetriever.java
+++ b/src/main/java/com/firebolt/jdbc/client/account/FireboltAccountRetriever.java
@@ -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) {
@@ -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);
         }
     }
 
diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java
index ab08eb342..9870d4288 100644
--- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java
+++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java
@@ -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,
@@ -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();
     }
@@ -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
@@ -102,11 +99,9 @@ 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);
         }
@@ -114,7 +109,7 @@ private FireboltProperties getSessionPropertiesForSystemEngine(String accessToke
                 .toBuilder()
                 .systemEngine(true)
                 .compress(false)
-                .accountId(accountId)
+                .accountId(null)
                 .host(systemEngienUrl.getHost())
                 .build();
     }
diff --git a/src/test/java/com/firebolt/jdbc/client/gateway/FireboltGatewayUrlClientTest.java b/src/test/java/com/firebolt/jdbc/client/gateway/FireboltGatewayUrlClientTest.java
index 360eac32b..c0aa7d13d 100644
--- a/src/test/java/com/firebolt/jdbc/client/gateway/FireboltGatewayUrlClientTest.java
+++ b/src/test/java/com/firebolt/jdbc/client/gateway/FireboltGatewayUrlClientTest.java
@@ -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
@@ -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"));
@@ -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);
@@ -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);
         }
     }
 
@@ -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"));
     }