From 12f2f6643efac591eb2ef458691aa6f3ad33e5bf Mon Sep 17 00:00:00 2001 From: alexradzin Date: Tue, 31 Oct 2023 17:37:12 +0200 Subject: [PATCH] separated FireboltEngineServer to 2 implementations --- .../jdbc/connection/FireboltConnection.java | 19 +- ...ConnectionServiceSecretAuthentication.java | 7 +- ...tConnectionUserPasswordAuthentication.java | 11 +- .../settings/FireboltProperties.java | 6 + .../service/FireboltEngineApiService.java | 109 ++++++++++++ ...ireboltEngineInformationSchemaService.java | 91 ++++++++++ .../jdbc/service/FireboltEngineService.java | 167 +----------------- .../connection/FireboltConnectionTest.java | 15 +- .../service/FireboltEngineServiceTest.java | 2 +- 9 files changed, 239 insertions(+), 188 deletions(-) create mode 100644 src/main/java/com/firebolt/jdbc/service/FireboltEngineApiService.java create mode 100644 src/main/java/com/firebolt/jdbc/service/FireboltEngineInformationSchemaService.java diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java index 8ab4da476..bb5338be4 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java @@ -17,7 +17,7 @@ import com.firebolt.jdbc.metadata.FireboltSystemEngineDatabaseMetadata; import com.firebolt.jdbc.service.FireboltAccountIdService; import com.firebolt.jdbc.service.FireboltAuthenticationService; -import com.firebolt.jdbc.service.FireboltEngineService; +import com.firebolt.jdbc.service.FireboltEngineInformationSchemaService; import com.firebolt.jdbc.service.FireboltGatewayUrlService; import com.firebolt.jdbc.service.FireboltStatementService; import com.firebolt.jdbc.statement.FireboltStatement; @@ -80,12 +80,12 @@ protected FireboltConnection(@NonNull String url, Properties connectionSettings, FireboltAuthenticationService fireboltAuthenticationService, FireboltGatewayUrlService fireboltGatewayUrlService, FireboltStatementService fireboltStatementService, - FireboltEngineService fireboltEngineService, + FireboltEngineInformationSchemaService fireboltEngineService, FireboltAccountIdService fireboltAccountIdService) throws SQLException { this.loginProperties = extractFireboltProperties(url, connectionSettings); this.fireboltAuthenticationService = fireboltAuthenticationService; - this.httpConnectionUrl = getHttpConnectionUrl(loginProperties); + this.httpConnectionUrl = loginProperties.getHttpConnectionUrl(); this.fireboltStatementService = fireboltStatementService; this.statements = new ArrayList<>(); @@ -101,7 +101,7 @@ protected FireboltConnection(@NonNull String url, Properties connectionSettings) ObjectMapper objectMapper = FireboltObjectMapper.getInstance(); this.fireboltAuthenticationService = new FireboltAuthenticationService(createFireboltAuthenticationClient(httpClient, objectMapper)); - this.httpConnectionUrl = getHttpConnectionUrl(loginProperties); + this.httpConnectionUrl = loginProperties.getHttpConnectionUrl(); this.fireboltStatementService = new FireboltStatementService(new StatementClientImpl(httpClient, objectMapper, this, loginProperties.getUserDrivers(), loginProperties.getUserClients())); this.statements = new ArrayList<>(); @@ -112,6 +112,10 @@ protected FireboltConnection(@NonNull String url, Properties connectionSettings) protected abstract FireboltAuthenticationClient createFireboltAuthenticationClient(OkHttpClient httpClient, ObjectMapper objectMapper); public static FireboltConnection create(@NonNull String url, Properties connectionSettings) throws SQLException { + return createConnectionInstance(url, connectionSettings); + } + + private static FireboltConnection createConnectionInstance(@NonNull String url, Properties connectionSettings) throws SQLException { switch(getUrlVersion(url)) { case 1: return new FireboltConnectionUserPasswordAuthentication(url, connectionSettings); case 2: return new FireboltConnectionServiceSecretAuthentication(url, connectionSettings); @@ -312,16 +316,11 @@ public void close() { log.debug("Connection closed"); } - protected FireboltProperties extractFireboltProperties(String jdbcUri, Properties connectionProperties) { + protected FireboltProperties extractFireboltProperties(String jdbcUri, Properties connectionProperties) throws SQLException { Properties propertiesFromUrl = UrlUtil.extractProperties(jdbcUri); return FireboltProperties.of(propertiesFromUrl, connectionProperties); } - private String getHttpConnectionUrl(FireboltProperties newSessionProperties) { - String hostAndPort = newSessionProperties.getHost() + ":" + newSessionProperties.getPort(); - return newSessionProperties.isSsl() ? "https://" + hostAndPort : "http://" + hostAndPort; - } - @Override public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException { diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretAuthentication.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretAuthentication.java index ad93f6793..7e0bcb9a4 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretAuthentication.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretAuthentication.java @@ -13,6 +13,7 @@ import com.firebolt.jdbc.exception.FireboltException; import com.firebolt.jdbc.service.FireboltAccountIdService; import com.firebolt.jdbc.service.FireboltAuthenticationService; +import com.firebolt.jdbc.service.FireboltEngineInformationSchemaService; import com.firebolt.jdbc.service.FireboltEngineService; import com.firebolt.jdbc.service.FireboltGatewayUrlService; import com.firebolt.jdbc.service.FireboltStatementService; @@ -31,7 +32,7 @@ public class FireboltConnectionServiceSecretAuthentication extends FireboltConne private final FireboltAccountIdService fireboltAccountIdService; private final FireboltEngineService fireboltEngineService; - FireboltConnectionServiceSecretAuthentication(@NonNull String url, Properties connectionSettings, FireboltAuthenticationService fireboltAuthenticationService, FireboltGatewayUrlService fireboltGatewayUrlService, FireboltStatementService fireboltStatementService, FireboltEngineService fireboltEngineService, FireboltAccountIdService fireboltAccountIdService) throws SQLException { + FireboltConnectionServiceSecretAuthentication(@NonNull String url, Properties connectionSettings, FireboltAuthenticationService fireboltAuthenticationService, FireboltGatewayUrlService fireboltGatewayUrlService, FireboltStatementService fireboltStatementService, FireboltEngineInformationSchemaService fireboltEngineService, FireboltAccountIdService fireboltAccountIdService) throws SQLException { super(url, connectionSettings, fireboltAuthenticationService, fireboltGatewayUrlService, fireboltStatementService, fireboltEngineService, fireboltAccountIdService); this.fireboltGatewayUrlService = fireboltGatewayUrlService; this.fireboltAccountIdService = fireboltAccountIdService; @@ -45,7 +46,7 @@ public class FireboltConnectionServiceSecretAuthentication extends FireboltConne ObjectMapper objectMapper = FireboltObjectMapper.getInstance(); this.fireboltGatewayUrlService = new FireboltGatewayUrlService(createFireboltAccountRetriever(httpClient, objectMapper, "engineUrl", GatewayUrlResponse.class)); this.fireboltAccountIdService = new FireboltAccountIdService(createFireboltAccountRetriever(httpClient, objectMapper, "resolve", FireboltAccount.class)); - this.fireboltEngineService = new FireboltEngineService(this, new FireboltAccountClient(httpClient, objectMapper, this, loginProperties.getUserDrivers(), loginProperties.getUserClients())); + this.fireboltEngineService = new FireboltEngineInformationSchemaService(this); connect(); } @@ -70,7 +71,7 @@ protected void authenticate() throws SQLException { private FireboltProperties getSessionPropertiesForNonSystemEngine() throws SQLException { sessionProperties = sessionProperties.toBuilder().engine(loginProperties.getEngine()).build(); - Engine engine = fireboltEngineService.getEngine(loginProperties.getEngine(), loginProperties.getDatabase()); + Engine engine = fireboltEngineService.getEngine(loginProperties); return loginProperties.toBuilder().host(engine.getEndpoint()).engine(engine.getName()).systemEngine(false).database(engine.getDatabase()).build(); } diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordAuthentication.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordAuthentication.java index 4df8c9b4b..6a418821b 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordAuthentication.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPasswordAuthentication.java @@ -10,6 +10,8 @@ import com.firebolt.jdbc.connection.settings.FireboltProperties; import com.firebolt.jdbc.service.FireboltAccountIdService; import com.firebolt.jdbc.service.FireboltAuthenticationService; +import com.firebolt.jdbc.service.FireboltEngineApiService; +import com.firebolt.jdbc.service.FireboltEngineInformationSchemaService; import com.firebolt.jdbc.service.FireboltEngineService; import com.firebolt.jdbc.service.FireboltGatewayUrlService; import com.firebolt.jdbc.service.FireboltStatementService; @@ -23,7 +25,7 @@ public class FireboltConnectionUserPasswordAuthentication extends FireboltConnection { private final FireboltEngineService fireboltEngineService; - FireboltConnectionUserPasswordAuthentication(@NonNull String url, Properties connectionSettings, FireboltAuthenticationService fireboltAuthenticationService, FireboltGatewayUrlService fireboltGatewayUrlService, FireboltStatementService fireboltStatementService, FireboltEngineService fireboltEngineService, FireboltAccountIdService fireboltAccountIdService) throws SQLException { + FireboltConnectionUserPasswordAuthentication(@NonNull String url, Properties connectionSettings, FireboltAuthenticationService fireboltAuthenticationService, FireboltGatewayUrlService fireboltGatewayUrlService, FireboltStatementService fireboltStatementService, FireboltEngineInformationSchemaService fireboltEngineService, FireboltAccountIdService fireboltAccountIdService) throws SQLException { super(url, connectionSettings, fireboltAuthenticationService, fireboltGatewayUrlService, fireboltStatementService, fireboltEngineService, fireboltAccountIdService); this.fireboltEngineService = fireboltEngineService; connect(); @@ -33,19 +35,20 @@ public class FireboltConnectionUserPasswordAuthentication extends FireboltConnec super(url, connectionSettings); OkHttpClient httpClient = getHttpClient(loginProperties); ObjectMapper objectMapper = FireboltObjectMapper.getInstance(); - this.fireboltEngineService = new FireboltEngineService(this, new FireboltAccountClient(httpClient, objectMapper, this, loginProperties.getUserDrivers(), loginProperties.getUserClients())); + this.fireboltEngineService = new FireboltEngineApiService(new FireboltAccountClient(httpClient, objectMapper, this, loginProperties.getUserDrivers(), loginProperties.getUserClients())); connect(); } @Override protected void authenticate() throws SQLException { String accessToken = getAccessToken(loginProperties).orElse(StringUtils.EMPTY); - String endpoint = fireboltEngineService.getEngine(httpConnectionUrl, loginProperties, accessToken).getEndpoint(); + FireboltProperties propertiesWithAccessToken = loginProperties.toBuilder().accessToken(accessToken).build(); + String endpoint = fireboltEngineService.getEngine(propertiesWithAccessToken).getEndpoint(); this.sessionProperties = loginProperties.toBuilder().host(endpoint).build(); } @Override - protected FireboltProperties extractFireboltProperties(String jdbcUri, Properties connectionProperties) { + protected FireboltProperties extractFireboltProperties(String jdbcUri, Properties connectionProperties) throws SQLException { FireboltProperties properties = super.extractFireboltProperties(jdbcUri, connectionProperties); boolean systemEngine = "system".equals(properties.getEngine()); return properties.toBuilder().systemEngine(systemEngine).build(); diff --git a/src/main/java/com/firebolt/jdbc/connection/settings/FireboltProperties.java b/src/main/java/com/firebolt/jdbc/connection/settings/FireboltProperties.java index 98fa70ab5..ff9de8cf1 100644 --- a/src/main/java/com/firebolt/jdbc/connection/settings/FireboltProperties.java +++ b/src/main/java/com/firebolt/jdbc/connection/settings/FireboltProperties.java @@ -248,4 +248,10 @@ public void addProperty(@NonNull String key, String value) { public void addProperty(Pair property) { this.addProperty(property.getLeft(), property.getRight()); } + + public String getHttpConnectionUrl() { + String hostAndPort = getHost() + ":" + getPort(); + return isSsl() ? "https://" + hostAndPort : "http://" + hostAndPort; + } + } diff --git a/src/main/java/com/firebolt/jdbc/service/FireboltEngineApiService.java b/src/main/java/com/firebolt/jdbc/service/FireboltEngineApiService.java new file mode 100644 index 000000000..e5dc80205 --- /dev/null +++ b/src/main/java/com/firebolt/jdbc/service/FireboltEngineApiService.java @@ -0,0 +1,109 @@ +package com.firebolt.jdbc.service; + +import com.firebolt.jdbc.client.account.FireboltAccountClient; +import com.firebolt.jdbc.client.account.response.FireboltAccountResponse; +import com.firebolt.jdbc.client.account.response.FireboltDefaultDatabaseEngineResponse; +import com.firebolt.jdbc.client.account.response.FireboltEngineIdResponse; +import com.firebolt.jdbc.client.account.response.FireboltEngineResponse; +import com.firebolt.jdbc.connection.Engine; +import com.firebolt.jdbc.connection.settings.FireboltProperties; +import com.firebolt.jdbc.exception.FireboltException; +import lombok.CustomLog; +import lombok.RequiredArgsConstructor; +import org.apache.commons.lang3.StringUtils; + +import java.io.IOException; +import java.sql.SQLException; +import java.util.Optional; +import java.util.Set; + +@RequiredArgsConstructor +@CustomLog +public class FireboltEngineApiService implements FireboltEngineService { + private static final Set ENGINE_NOT_READY_STATUSES = Set.of( + "ENGINE_STATUS_PROVISIONING_STARTED", "ENGINE_STATUS_PROVISIONING_PENDING", + "ENGINE_STATUS_PROVISIONING_FINISHED", "ENGINE_STATUS_RUNNING_REVISION_STARTING"); + private static final String ERROR_NO_ENGINE_ATTACHED = "There is no Firebolt engine running on %s attached to the database %s. To connect first make sure there is a running engine and then try again."; + private static final String ERROR_NO_ENGINE_WITH_NAME = "There is no Firebolt engine running on %s with the name %s. To connect first make sure there is a running engine and then try again."; + + + private final FireboltAccountClient fireboltAccountClient; + + @Override + public Engine getEngine(FireboltProperties properties) throws SQLException { + return getEngine(properties.getHttpConnectionUrl(), properties, properties.getAccessToken()); + } + + /** + * Returns the engine + * + * @param connectionUrl the connection url + * @param loginProperties properties to login + * @param accessToken the access token + * @return the engine + */ + private Engine getEngine(String connectionUrl, FireboltProperties loginProperties, String accessToken) + throws FireboltException { + String accountId = null; + Engine engine; + try { + if (StringUtils.isNotEmpty(loginProperties.getAccount())) { + accountId = getAccountId(connectionUrl, loginProperties.getAccount(), accessToken).orElse(null); + } + if (StringUtils.isEmpty(loginProperties.getEngine())) { + engine = getDefaultEngine(connectionUrl, accountId, loginProperties.getDatabase(), accessToken); + } else { + engine = getEngineWithName(connectionUrl, accountId, loginProperties.getEngine(), accessToken); + } + } catch (FireboltException e) { + throw e; + } catch (Exception e) { + throw new FireboltException("Failed to get engine", e); + } + validateEngineIsNotStarting(engine); + return engine; + } + + private Engine getEngineWithName(String connectionUrl, String accountId, String engineName, String accessToken) + throws FireboltException, IOException { + FireboltEngineIdResponse response = fireboltAccountClient.getEngineId(connectionUrl, accountId, engineName, + accessToken); + String engineID = Optional.ofNullable(response).map(FireboltEngineIdResponse::getEngine) + .map(FireboltEngineIdResponse.Engine::getEngineId).orElseThrow(() -> new FireboltException( + "Failed to extract engine id field from the server response: the response from the server is invalid.")); + FireboltEngineResponse fireboltEngineResponse = fireboltAccountClient.getEngine(connectionUrl, accountId, + engineName, engineID, accessToken); + + return Optional.ofNullable(fireboltEngineResponse).map(FireboltEngineResponse::getEngine) + .filter(e -> StringUtils.isNotEmpty(e.getEndpoint())) + .map(e -> new Engine(e.getEndpoint(), e.getCurrentStatus(), engineName, null, engineID)) + .orElseThrow(() -> new FireboltException( + String.format(ERROR_NO_ENGINE_WITH_NAME, connectionUrl, engineName))); + } + + private Engine getDefaultEngine(String connectionUrl, String accountId, String database, String accessToken) + throws FireboltException, IOException { + FireboltDefaultDatabaseEngineResponse defaultEngine = fireboltAccountClient + .getDefaultEngineByDatabaseName(connectionUrl, accountId, database, accessToken); + + return Optional.ofNullable(defaultEngine).map(FireboltDefaultDatabaseEngineResponse::getEngineUrl) + .map(url -> new Engine(url, "running", null, database, null)).orElseThrow( + () -> new FireboltException(String.format(ERROR_NO_ENGINE_ATTACHED, connectionUrl, database))); + } + + private Optional getAccountId(String connectionUrl, String account, String accessToken) + throws FireboltException, IOException { + FireboltAccountResponse fireboltAccountResponse = fireboltAccountClient.getAccount(connectionUrl, account, + accessToken); + return Optional.ofNullable(fireboltAccountResponse).map(FireboltAccountResponse::getAccountId); + } + + private void validateEngineIsNotStarting(Engine engine) throws FireboltException { + if (StringUtils.isNotEmpty(engine.getId()) && StringUtils.isNotEmpty(engine.getStatus()) + && ENGINE_NOT_READY_STATUSES.contains(engine.getStatus())) { + throw new FireboltException(String.format( + "The engine %s is currently starting. Please wait until the engine is on and then execute the query again.", engine.getName())); + } + } + +} diff --git a/src/main/java/com/firebolt/jdbc/service/FireboltEngineInformationSchemaService.java b/src/main/java/com/firebolt/jdbc/service/FireboltEngineInformationSchemaService.java new file mode 100644 index 000000000..0c9b7808d --- /dev/null +++ b/src/main/java/com/firebolt/jdbc/service/FireboltEngineInformationSchemaService.java @@ -0,0 +1,91 @@ +package com.firebolt.jdbc.service; + +import com.firebolt.jdbc.connection.Engine; +import com.firebolt.jdbc.connection.FireboltConnection; +import com.firebolt.jdbc.connection.settings.FireboltProperties; +import com.firebolt.jdbc.exception.FireboltException; +import lombok.CustomLog; +import lombok.RequiredArgsConstructor; + +import javax.annotation.Nullable; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Optional; + +import static java.lang.String.format; + +@RequiredArgsConstructor +@CustomLog +public class FireboltEngineInformationSchemaService implements FireboltEngineService { + private static final String ENGINE_URL = "url"; + private static final String STATUS_FIELD = "status"; + private static final String ENGINE_NAME_FIELD = "engine_name"; + private static final String RUNNING_STATUS = "running"; + private static final String ENGINE_QUERY = + "SELECT engs.url, engs.attached_to, dbs.database_name, engs.status, engs.engine_name " + + "FROM information_schema.engines as engs " + + "LEFT JOIN information_schema.databases as dbs ON engs.attached_to = dbs.database_name " + + "WHERE engs.engine_name = ?"; + private static final String DATABASE_QUERY = "SELECT database_name FROM information_schema.databases WHERE database_name=?"; + + private final FireboltConnection fireboltConnection; + + /** + * Extracts the engine name from host + * + * @param engineHost engine host + * @return the engine name + */ + public String getEngineNameByHost(String engineHost) throws FireboltException { + return Optional.ofNullable(engineHost).filter(host -> host.contains(".")).map(host -> host.split("\\.")[0]) + .map(host -> host.replace("-", "_")).orElseThrow(() -> new FireboltException( + format("Could not establish the engine from the host: %s", engineHost))); + } + + @Override + public boolean doesDatabaseExist(String database) throws SQLException { + try (PreparedStatement ps = fireboltConnection.prepareStatement(DATABASE_QUERY)) { + ps.setString(1, database); + try (ResultSet rs = ps.executeQuery()) { + return rs.next(); + } + } + } + + @Override + public Engine getEngine(FireboltProperties properties) throws SQLException { + return getEngine(properties.getEngine(), properties.getDatabase()); + } + + + public Engine getEngine(String engine, @Nullable String database) throws SQLException { + if (engine == null) { + throw new IllegalArgumentException("Cannot retrieve engine parameters because its name is null"); + } + try (PreparedStatement ps = fireboltConnection.prepareStatement(ENGINE_QUERY)) { + ps.setString(1, engine); + try (ResultSet rs = ps.executeQuery()) { + if (!rs.next()) { + throw new FireboltException(format("The engine with the name %s could not be found", engine)); + } + String status = rs.getString(STATUS_FIELD); + if (!isEngineRunning(status)) { + throw new FireboltException(format("The engine with the name %s is not running. Status: %s", engine, status)); + } + String attachedDatabase = rs.getString("attached_to"); + if (attachedDatabase == null) { + throw new FireboltException(format("The engine with the name %s is not attached to any database", engine)); + } + if (database != null && !database.equals(attachedDatabase)) { + throw new FireboltException(format("The engine with the name %s is not attached to database %s", engine, database)); + } + return new Engine(rs.getString(ENGINE_URL), status, rs.getString(ENGINE_NAME_FIELD), attachedDatabase, null); + } + } + } + + private boolean isEngineRunning(String status) { + return RUNNING_STATUS.equalsIgnoreCase(status); + } +} diff --git a/src/main/java/com/firebolt/jdbc/service/FireboltEngineService.java b/src/main/java/com/firebolt/jdbc/service/FireboltEngineService.java index 4e6ea6412..b50c74b5c 100644 --- a/src/main/java/com/firebolt/jdbc/service/FireboltEngineService.java +++ b/src/main/java/com/firebolt/jdbc/service/FireboltEngineService.java @@ -1,174 +1,15 @@ package com.firebolt.jdbc.service; -import com.firebolt.jdbc.client.account.FireboltAccountClient; -import com.firebolt.jdbc.client.account.response.FireboltAccountResponse; -import com.firebolt.jdbc.client.account.response.FireboltDefaultDatabaseEngineResponse; -import com.firebolt.jdbc.client.account.response.FireboltEngineIdResponse; -import com.firebolt.jdbc.client.account.response.FireboltEngineResponse; import com.firebolt.jdbc.connection.Engine; -import com.firebolt.jdbc.connection.FireboltConnection; import com.firebolt.jdbc.connection.settings.FireboltProperties; -import com.firebolt.jdbc.exception.FireboltException; -import lombok.CustomLog; -import lombok.RequiredArgsConstructor; -import org.apache.commons.lang3.StringUtils; -import javax.annotation.Nullable; -import java.io.IOException; -import java.sql.PreparedStatement; -import java.sql.ResultSet; import java.sql.SQLException; -import java.util.Optional; -import java.util.Set; -import static java.lang.String.format; +public interface FireboltEngineService { + Engine getEngine(FireboltProperties properties) throws SQLException; -@RequiredArgsConstructor -@CustomLog -public class FireboltEngineService { - private static final String ENGINE_URL = "url"; - private static final String STATUS_FIELD = "status"; - private static final String ENGINE_NAME_FIELD = "engine_name"; - private static final String RUNNING_STATUS = "running"; - private static final String ENGINE_QUERY = - "SELECT engs.url, engs.attached_to, dbs.database_name, engs.status, engs.engine_name " + - "FROM information_schema.engines as engs " + - "LEFT JOIN information_schema.databases as dbs ON engs.attached_to = dbs.database_name " + - "WHERE engs.engine_name = ?"; - private static final String DATABASE_QUERY = "SELECT database_name FROM information_schema.databases WHERE database_name=?"; - - private static final Set ENGINE_NOT_READY_STATUSES = Set.of( - "ENGINE_STATUS_PROVISIONING_STARTED", "ENGINE_STATUS_PROVISIONING_PENDING", - "ENGINE_STATUS_PROVISIONING_FINISHED", "ENGINE_STATUS_RUNNING_REVISION_STARTING"); - private static final String ERROR_NO_ENGINE_ATTACHED = "There is no Firebolt engine running on %s attached to the database %s. To connect first make sure there is a running engine and then try again."; - private static final String ERROR_NO_ENGINE_WITH_NAME = "There is no Firebolt engine running on %s with the name %s. To connect first make sure there is a running engine and then try again."; - - - private final FireboltConnection fireboltConnection; - private final FireboltAccountClient fireboltAccountClient; - - /** - * Extracts the engine name from host - * - * @param engineHost engine host - * @return the engine name - */ - public String getEngineNameByHost(String engineHost) throws FireboltException { - return Optional.ofNullable(engineHost).filter(host -> host.contains(".")).map(host -> host.split("\\.")[0]) - .map(host -> host.replace("-", "_")).orElseThrow(() -> new FireboltException( - format("Could not establish the engine from the host: %s", engineHost))); - } - - public boolean doesDatabaseExist(String database) throws SQLException { - try (PreparedStatement ps = fireboltConnection.prepareStatement(DATABASE_QUERY)) { - ps.setString(1, database); - try (ResultSet rs = ps.executeQuery()) { - return rs.next(); - } - } - } - - public Engine getEngine(String engine, @Nullable String database) throws SQLException { - if (engine == null) { - throw new IllegalArgumentException("Cannot retrieve engine parameters because its name is null"); - } - try (PreparedStatement ps = fireboltConnection.prepareStatement(ENGINE_QUERY)) { - ps.setString(1, engine); - try (ResultSet rs = ps.executeQuery()) { - if (!rs.next()) { - throw new FireboltException(format("The engine with the name %s could not be found", engine)); - } - String status = rs.getString(STATUS_FIELD); - if (!isEngineRunning(status)) { - throw new FireboltException(format("The engine with the name %s is not running. Status: %s", engine, status)); - } - String attachedDatabase = rs.getString("attached_to"); - if (attachedDatabase == null) { - throw new FireboltException(format("The engine with the name %s is not attached to any database", engine)); - } - if (database != null && !database.equals(attachedDatabase)) { - throw new FireboltException(format("The engine with the name %s is not attached to database %s", engine, database)); - } - return new Engine(rs.getString(ENGINE_URL), status, rs.getString(ENGINE_NAME_FIELD), attachedDatabase, null); - } - } - } - - private boolean isEngineRunning(String status) { - return RUNNING_STATUS.equalsIgnoreCase(status); - } - - - /** - * Returns the engine - * - * @param connectionUrl the connection url - * @param loginProperties properties to login - * @param accessToken the access token - * @return the engine - */ - public Engine getEngine(String connectionUrl, FireboltProperties loginProperties, String accessToken) - throws FireboltException { - String accountId = null; - Engine engine; - try { - if (StringUtils.isNotEmpty(loginProperties.getAccount())) { - accountId = getAccountId(connectionUrl, loginProperties.getAccount(), accessToken).orElse(null); - } - if (StringUtils.isEmpty(loginProperties.getEngine())) { - engine = getDefaultEngine(connectionUrl, accountId, loginProperties.getDatabase(), accessToken); - } else { - engine = getEngineWithName(connectionUrl, accountId, loginProperties.getEngine(), accessToken); - } - } catch (FireboltException e) { - throw e; - } catch (Exception e) { - throw new FireboltException("Failed to get engine", e); - } - validateEngineIsNotStarting(engine); - return engine; - } - - private Engine getEngineWithName(String connectionUrl, String accountId, String engineName, String accessToken) - throws FireboltException, IOException { - FireboltEngineIdResponse response = fireboltAccountClient.getEngineId(connectionUrl, accountId, engineName, - accessToken); - String engineID = Optional.ofNullable(response).map(FireboltEngineIdResponse::getEngine) - .map(FireboltEngineIdResponse.Engine::getEngineId).orElseThrow(() -> new FireboltException( - "Failed to extract engine id field from the server response: the response from the server is invalid.")); - FireboltEngineResponse fireboltEngineResponse = fireboltAccountClient.getEngine(connectionUrl, accountId, - engineName, engineID, accessToken); - - return Optional.ofNullable(fireboltEngineResponse).map(FireboltEngineResponse::getEngine) - .filter(e -> StringUtils.isNotEmpty(e.getEndpoint())) - .map(e -> new Engine(e.getEndpoint(), e.getCurrentStatus(), engineName, null, engineID)) - .orElseThrow(() -> new FireboltException( - String.format(ERROR_NO_ENGINE_WITH_NAME, connectionUrl, engineName))); - } - - private Engine getDefaultEngine(String connectionUrl, String accountId, String database, String accessToken) - throws FireboltException, IOException { - FireboltDefaultDatabaseEngineResponse defaultEngine = fireboltAccountClient - .getDefaultEngineByDatabaseName(connectionUrl, accountId, database, accessToken); - - return Optional.ofNullable(defaultEngine).map(FireboltDefaultDatabaseEngineResponse::getEngineUrl) - .map(url -> new Engine(url, "running", null, database, null)).orElseThrow( - () -> new FireboltException(String.format(ERROR_NO_ENGINE_ATTACHED, connectionUrl, database))); - } - - private Optional getAccountId(String connectionUrl, String account, String accessToken) - throws FireboltException, IOException { - FireboltAccountResponse fireboltAccountResponse = fireboltAccountClient.getAccount(connectionUrl, account, - accessToken); - return Optional.ofNullable(fireboltAccountResponse).map(FireboltAccountResponse::getAccountId); - } - - private void validateEngineIsNotStarting(Engine engine) throws FireboltException { - if (StringUtils.isNotEmpty(engine.getId()) && StringUtils.isNotEmpty(engine.getStatus()) - && ENGINE_NOT_READY_STATUSES.contains(engine.getStatus())) { - throw new FireboltException(String.format( - "The engine %s is currently starting. Please wait until the engine is on and then execute the query again.", engine.getName())); - } + default boolean doesDatabaseExist(String database) throws SQLException { + return true; } } diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java index 27b3e836e..3d722272a 100644 --- a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java +++ b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java @@ -6,6 +6,7 @@ import com.firebolt.jdbc.exception.FireboltException; import com.firebolt.jdbc.service.FireboltAccountIdService; import com.firebolt.jdbc.service.FireboltAuthenticationService; +import com.firebolt.jdbc.service.FireboltEngineInformationSchemaService; import com.firebolt.jdbc.service.FireboltEngineService; import com.firebolt.jdbc.service.FireboltGatewayUrlService; import com.firebolt.jdbc.service.FireboltStatementService; @@ -93,7 +94,7 @@ class FireboltConnectionTest { private FireboltGatewayUrlService fireboltGatewayUrlService; @Mock - private FireboltEngineService fireboltEngineService; + private FireboltEngineInformationSchemaService fireboltEngineService; @Mock private FireboltStatementService fireboltStatementService; @Mock @@ -149,7 +150,7 @@ void init() throws SQLException { .thenReturn(fireboltConnectionTokens); lenient().when(fireboltGatewayUrlService.getUrl(any(), any())).thenReturn("http://foo:8080/bar"); engine = new Engine("endpoint", "id123", "OK", "noname", null); - lenient().when(fireboltEngineService.getEngine(any(), any())).thenReturn(engine); + lenient().when(fireboltEngineService.getEngine(any())).thenReturn(engine); lenient().when(fireboltEngineService.doesDatabaseExist(any())).thenReturn(true); } @@ -427,7 +428,7 @@ void shouldRemoveExpiredToken() throws SQLException { when(FireboltProperties.of(any())).thenReturn(fireboltProperties); when(fireboltAuthenticationService.getConnectionTokens("http://host:8080", fireboltProperties)) .thenReturn(FireboltConnectionTokens.builder().build()); - lenient().when(fireboltEngineService.getEngine(any(), any())).thenReturn(new Engine("http://hello", null, null, null, null)); + lenient().when(fireboltEngineService.getEngine(any())).thenReturn(new Engine("http://hello", null, null, null, null)); try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) { fireboltConnection.removeExpiredTokens(); @@ -444,7 +445,7 @@ void shouldReturnConnectionTokenWhenAvailable() throws SQLException { when(FireboltProperties.of(any())).thenReturn(fireboltProperties); FireboltConnectionTokens connectionTokens = FireboltConnectionTokens.builder().accessToken(accessToken).build(); when(fireboltAuthenticationService.getConnectionTokens(eq("http://host:8080"), any())).thenReturn(connectionTokens); - lenient().when(fireboltEngineService.getEngine(any(), any())).thenReturn(new Engine("http://engineHost", null, null, null, null)); + lenient().when(fireboltEngineService.getEngine(any())).thenReturn(new Engine("http://engineHost", null, null, null, null)); try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) { verify(fireboltAuthenticationService).getConnectionTokens("http://host:8080", fireboltProperties); assertEquals(accessToken, fireboltConnection.getAccessToken().get()); @@ -592,9 +593,9 @@ void shouldReturnEmptyResult(String name, Callable function, Object expected) void shouldGetEngineUrlWhenEngineIsProvided() throws SQLException { connectionProperties.put("engine", "engine"); - when(fireboltEngineService.getEngine(any(), any())).thenReturn(new Engine("http://my_endpoint", null, null, null, null)); + when(fireboltEngineService.getEngine(any())).thenReturn(new Engine("http://my_endpoint", null, null, null, null)); try (FireboltConnection connection = createConnection(URL, connectionProperties)) { - verify(fireboltEngineService).getEngine("engine", "db"); + verify(fireboltEngineService).getEngine(FireboltProperties.builder().engine("engine").database("db").build()); assertEquals("http://my_endpoint", connection.getSessionProperties().getHost()); } } @@ -605,7 +606,7 @@ void shouldNotGetEngineUrlOrDefaultEngineUrlWhenUsingSystemEngine() throws SQLEx when(fireboltGatewayUrlService.getUrl(any(), any())).thenReturn("http://my_endpoint"); try (FireboltConnection connection = createConnection(SYSTEM_ENGINE_URL, connectionProperties)) { - verify(fireboltEngineService, times(0)).getEngine(isNull(), eq("my_db")); + verify(fireboltEngineService, times(0)).getEngine(FireboltProperties.builder().database("my_db").build()); assertEquals("my_endpoint", connection.getSessionProperties().getHost()); } } diff --git a/src/test/java/com/firebolt/jdbc/service/FireboltEngineServiceTest.java b/src/test/java/com/firebolt/jdbc/service/FireboltEngineServiceTest.java index b5717397e..0ade96104 100644 --- a/src/test/java/com/firebolt/jdbc/service/FireboltEngineServiceTest.java +++ b/src/test/java/com/firebolt/jdbc/service/FireboltEngineServiceTest.java @@ -29,7 +29,7 @@ class FireboltEngineServiceTest { @InjectMocks - private FireboltEngineService fireboltEngineService; + private FireboltEngineInformationSchemaService fireboltEngineService; @Mock private FireboltConnection fireboltConnection;