From b37263b5a3b37eaa9f8df68a39713adb5e2f12c0 Mon Sep 17 00:00:00 2001 From: Jose Almeida <53087160+jcralmeida@users.noreply.github.com> Date: Thu, 17 Jun 2021 10:37:04 -0700 Subject: [PATCH] Merge pull request #11 from abenaru/raise-code-covarege-limit Raise code coverage to 80% --- .../arrow/driver/jdbc/ArrowFlightClient.java | 424 ------------------ .../driver/jdbc/ArrowFlightConnection.java | 147 +++--- .../driver/jdbc/ArrowFlightJdbcDriver.java | 84 ++-- .../jdbc/client/ArrowFlightClientHandler.java | 290 +++++++----- .../jdbc/client/FlightClientHandler.java | 19 +- .../utils/ClientAuthenticationUtils.java | 4 + ...DefaultProperty.java => BaseProperty.java} | 38 +- .../jdbc/test/ArrowFlightJdbcDriverTest.java | 27 +- .../driver/jdbc/test/ConnectionTest.java | 44 +- .../driver/jdbc/test/ConnectionTlsTest.java | 75 ++-- 10 files changed, 429 insertions(+), 723 deletions(-) delete mode 100644 java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightClient.java rename java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/utils/{DefaultProperty.java => BaseProperty.java} (53%) diff --git a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightClient.java b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightClient.java deleted file mode 100644 index 06e63f2f63bb9..0000000000000 --- a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightClient.java +++ /dev/null @@ -1,424 +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.arrow.driver.jdbc; - -import java.io.ByteArrayInputStream; -import java.io.IOException; -import java.io.InputStream; -import java.io.StringWriter; -import java.nio.charset.StandardCharsets; -import java.nio.file.Files; -import java.nio.file.Paths; -import java.security.KeyStore; -import java.security.KeyStoreException; -import java.security.NoSuchAlgorithmException; -import java.security.cert.Certificate; -import java.security.cert.CertificateException; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Enumeration; -import java.util.List; - -import javax.annotation.Nullable; - -import org.apache.arrow.flight.CallOption; -import org.apache.arrow.flight.FlightClient; -import org.apache.arrow.flight.FlightDescriptor; -import org.apache.arrow.flight.FlightInfo; -import org.apache.arrow.flight.FlightStream; -import org.apache.arrow.flight.HeaderCallOption; -import org.apache.arrow.flight.Location; -import org.apache.arrow.flight.auth2.BasicAuthCredentialWriter; -import org.apache.arrow.flight.auth2.ClientBearerHeaderHandler; -import org.apache.arrow.flight.auth2.ClientIncomingAuthHeaderMiddleware; -import org.apache.arrow.flight.grpc.CredentialCallOption; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.util.Preconditions; -import org.apache.arrow.vector.VectorSchemaRoot; -import org.bouncycastle.openssl.jcajce.JcaPEMWriter; - -/** - * An adhoc {@link FlightClient} wrapper used to access the client. Allows for - * the reuse of credentials. - */ -public final class ArrowFlightClient implements AutoCloseable { - - private final FlightClient client; - - private final CredentialCallOption bearerToken; - - private ArrowFlightClient(final FlightClient client, - final CredentialCallOption properties) { - this.client = client; - this.bearerToken = properties; - } - - /** - * Gets the Arrow Flight Client. - * - * @return the {@link FlightClient} wrapped by this. - */ - protected FlightClient getClient() { - return client; - } - - /** - * Gets the bearer token for the client wrapped by this. - * - * @return the {@link CredentialCallOption} of this client. - */ - protected CredentialCallOption getProperties() { - return bearerToken; - } - - /** - * Makes RPC requests to the Dremio Flight Server Endpoint to retrieve results - * of the provided SQL query. - * - * @param query - * The SQL query to execute. - * @param headerCallOption - * The client properties to execute provided SQL query with. - * @throws Exception - * If an error occurs during query execution. - */ - public VectorSchemaRoot runQuery(final String query, - final HeaderCallOption headerCallOption) throws Exception { - /* - * TODO Run a query and return its corresponding VectorSchemaRoot, which - * must later be converted into a ResultSet. - */ - return null; - } - - /** - * Makes an RPC "getInfo" request with the given query and client properties - * in order to retrieve the metadata associated with a set of data records. - * - * @param query - * The query to retrieve FlightInfo for. - * @param options - * The client properties to execute this request with. - * @return a {@link FlightInfo} object. - */ - public FlightInfo getInfo(final String query, final CallOption... options) { - return client.getInfo( - FlightDescriptor.command(query.getBytes(StandardCharsets.UTF_8)), - options); - } - - /** - * Makes an RPC "getStream" request based on the provided {@link FlightInfo} - * object. Retrieves result of the query previously prepared with "getInfo." - * - * @param flightInfo - * The {@code FlightInfo} object encapsulating information for the - * server to identify the prepared statement with. - * @param options - * The client properties to execute this request with. - * @return a {@code FlightStream} of results. - */ - public FlightStream getStream(final FlightInfo flightInfo, - final CallOption... options) { - return client.getStream(null, options); - } - - /** - * Creates a {@code ArrowFlightClient} wrapping a {@link FlightClient} - * connected to the Dremio server without any encryption. - * - * @param allocator - * The buffer allocator to use for the {@code FlightClient} wrapped - * by this. - * @param host - * The host to connect to. - * @param port - * The port to connect to. - * @param username - * The username to connect with. - * @param password - * The password to connect with. - * @param clientProperties - * The client properties to set during authentication. - * @return a new {@code ArrowFlightClient} wrapping a non-encrypted - * {@code FlightClient}, with a bearer token for subsequent requests - * to the wrapped client. - */ - public static ArrowFlightClient getBasicClientAuthenticated( - final BufferAllocator allocator, - final String host, final int port, final String username, - @Nullable final String password, - @Nullable final HeaderCallOption clientProperties) { - - final ClientIncomingAuthHeaderMiddleware.Factory factory = - new ClientIncomingAuthHeaderMiddleware.Factory( - new ClientBearerHeaderHandler()); - - final FlightClient flightClient = FlightClient.builder() - .allocator(allocator) - .location( - Location.forGrpcInsecure(host, port)) - .intercept(factory).build(); - - return new ArrowFlightClient(flightClient, getAuthenticate(flightClient, - username, password, factory, clientProperties)); - } - - /** - * Creates a {@code ArrowFlightClient} wrapping a {@link FlightClient} - * connected to the Dremio server without any encryption. - * - * @param allocator - * The buffer allocator to use for the {@code FlightClient} wrapped - * by this. - * @param host - * The host to connect to. - * @param port - * The port to connect to. - * @param clientProperties - * The client properties to set during authentication. - * @return a new {@code ArrowFlightClient} wrapping a non-encrypted - * {@code FlightClient}, with a bearer token for subsequent requests - * to the wrapped client. - */ - public static ArrowFlightClient getBasicClientNoAuth( - final BufferAllocator allocator, - final String host, final int port, - @Nullable final HeaderCallOption clientProperties) { - - final FlightClient flightClient = FlightClient.builder().allocator(allocator) - .location( - Location.forGrpcInsecure(host, port)).build(); - - return new ArrowFlightClient(flightClient, null); - } - - /** - * Creates a {@code ArrowFlightClient} wrapping a {@link FlightClient} - * connected to the Dremio server with an encrypted TLS connection. - * - * @param allocator - * The buffer allocator to use for the {@code FlightClient} wrapped - * by this. - * @param host - * The host to connect to. - * @param port - * The port to connect to. - * @param clientProperties - * The client properties to set during authentication. - * @param username - * The username to connect with. - * @param password - * The password to connect with. - * @param keyStorePath - * The KeyStore path to use. - * @param keyStorePass - * The KeyStore password to use. - * @return a new {@code ArrowFlightClient} wrapping a non-encrypted - * {@code FlightClient}, with a bearer token for subsequent requests - * to the wrapped client. - * @throws KeyStoreException - * If an error occurs while trying to retrieve KeyStore information. - * @throws NoSuchAlgorithmException - * If a particular cryptographic algorithm is required but does not - * exist. - * @throws CertificateException - * If an error occurs while trying to retrieve certificate - * information. - * @throws IOException - * If an I/O operation fails. - */ - public static ArrowFlightClient getEncryptedClientAuthenticated( - final BufferAllocator allocator, - final String host, final int port, - @Nullable final HeaderCallOption clientProperties, final String username, - @Nullable final String password, final String keyStorePath, - final String keyStorePass) - throws SQLException { - - try { - - final ClientIncomingAuthHeaderMiddleware.Factory factory = - new ClientIncomingAuthHeaderMiddleware.Factory( - new ClientBearerHeaderHandler()); - - final FlightClient flightClient = FlightClient.builder() - .allocator(allocator) - .location( - Location.forGrpcTls(host, port)) - .intercept(factory).useTls() - .trustedCertificates(getCertificateStream(keyStorePath, keyStorePass)) - .build(); - - return new ArrowFlightClient(flightClient, getAuthenticate(flightClient, - username, password, factory, clientProperties)); - } catch (final Exception e) { - throw new SQLException( - "Failed to create a new Arrow Flight client.", e); - } - } - - /** - * Creates a {@code ArrowFlightClient} wrapping a {@link FlightClient} - * connected to the Dremio server with an encrypted TLS connection. - * - * @param allocator - * The buffer allocator to use for the {@code FlightClient} wrapped - * by this. - * @param host - * The host to connect to. - * @param port - * The port to connect to. - * @param clientProperties - * The client properties to set during authentication. - * @param keyStorePath - * The KeyStore path to use. - * @param keyStorePass - * The KeyStore password to use. - * @return a new {@code ArrowFlightClient} wrapping a non-encrypted - * {@code FlightClient}, with a bearer token for subsequent requests - * to the wrapped client. - * @throws KeyStoreException - * If an error occurs while trying to retrieve KeyStore information. - * @throws NoSuchAlgorithmException - * If a particular cryptographic algorithm is required but does not - * exist. - * @throws CertificateException - * If an error occurs while trying to retrieve certificate - * information. - * @throws IOException - * If an I/O operation fails. - */ - public static ArrowFlightClient getEncryptedClientNoAuth( - final BufferAllocator allocator, - final String host, final int port, - @Nullable final HeaderCallOption clientProperties, - final String keyStorePath, final String keyStorePass) - throws SQLException { - - try { - - final FlightClient flightClient = FlightClient.builder() - .allocator(allocator) - .location( - Location.forGrpcTls(host, port)).useTls() - .trustedCertificates(getCertificateStream(keyStorePath, keyStorePass)) - .build(); - - return new ArrowFlightClient(flightClient, null); - } catch (KeyStoreException | NoSuchAlgorithmException | - CertificateException | IOException e) { - throw new SQLException("Failed to create an Arrow Flight client.", e); - } - } - - /** - * Helper method to authenticate provided {@link FlightClient} instance - * against an Arrow Flight server endpoint. - * - * @param client - * the FlightClient instance to connect to Arrow Flight. - * @param username - * the Arrow Flight server username. - * @param password - * the corresponding Arrow Flight server password - * @param factory - * the factory to create {@link ClientIncomingAuthHeaderMiddleware}. - * @param clientProperties - * client properties to set during authentication. - * @return {@link CredentialCallOption} encapsulating the bearer token to use - * in subsequent requests. - */ - public static CredentialCallOption getAuthenticate(final FlightClient client, - final String username, @Nullable final String password, - final ClientIncomingAuthHeaderMiddleware.Factory factory, - @Nullable final HeaderCallOption clientProperties) { - - final List callOptions = new ArrayList<>(); - - callOptions.add(new CredentialCallOption( - new BasicAuthCredentialWriter(username, password))); - - if (clientProperties != null) { - callOptions.add(clientProperties); - } - - client.handshake(callOptions.toArray(new CallOption[callOptions.size()])); - - return factory.getCredentialCallOption(); - } - - /** - * Generates an {@link InputStream} that contains certificates for a private - * key. - * - * @param keyStorePath - * The path to the keystore. - * @param keyStorePass - * The password for the keystore. - * @return a new {code InputStream} containing the certificates. - * @throws Exception - * If there was an error looking up the private key or certificates. - */ - public static InputStream getCertificateStream(final String keyStorePath, - final String keyStorePass) throws KeyStoreException, - NoSuchAlgorithmException, CertificateException, IOException { - - final KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType()); - try (final InputStream keyStoreStream = Files - .newInputStream(Paths.get(Preconditions.checkNotNull(keyStorePath)))) { - keyStore.load(keyStoreStream, - Preconditions.checkNotNull(keyStorePass).toCharArray()); - } - - final Enumeration aliases = keyStore.aliases(); - - while (aliases.hasMoreElements()) { - final String alias = aliases.nextElement(); - if (keyStore.isCertificateEntry(alias)) { - final Certificate certificates = keyStore.getCertificate(alias); - return toInputStream(certificates); - } - } - - throw new CertificateException("Keystore did not have a certificate."); - } - - private static InputStream toInputStream(final Certificate certificate) - throws IOException { - - try (final StringWriter writer = new StringWriter(); - final JcaPEMWriter pemWriter = new JcaPEMWriter(writer)) { - - pemWriter.writeObject(certificate); - pemWriter.flush(); - return new ByteArrayInputStream( - writer.toString().getBytes(StandardCharsets.UTF_8)); - } - } - - @Override - public void close() throws Exception { - try { - client.close(); - } catch (final InterruptedException e) { - System.out.println("[WARNING] Failed to close resource."); - e.printStackTrace(); - } - } -} diff --git a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index c3d5093f6aab2..6f4b5550fc0a4 100644 --- a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -17,55 +17,76 @@ package org.apache.arrow.driver.jdbc; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.HOST; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.KEYSTORE_PASS; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.KEYSTORE_PATH; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.PASSWORD; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.PORT; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.USERNAME; + import java.io.IOException; import java.net.URISyntaxException; +import java.security.GeneralSecurityException; import java.security.KeyStoreException; import java.security.NoSuchAlgorithmException; import java.security.cert.CertificateException; import java.sql.SQLException; +import java.util.Iterator; +import java.util.Map; +import java.util.Objects; import java.util.Properties; -import javax.annotation.Nullable; +import javax.management.InstanceAlreadyExistsException; -import org.apache.arrow.driver.jdbc.utils.DefaultProperty; +import org.apache.arrow.driver.jdbc.client.ArrowFlightClientHandler; +import org.apache.arrow.flight.CallHeaders; +import org.apache.arrow.flight.FlightCallHeaders; +import org.apache.arrow.flight.HeaderCallOption; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; +import com.google.common.base.Strings; + /** * Connection to the Arrow Flight server. */ -public final class ArrowFlightConnection extends AvaticaConnection { +public class ArrowFlightConnection extends AvaticaConnection { private final BufferAllocator allocator; // TODO Use this later to run queries. @SuppressWarnings("unused") - private ArrowFlightClient client; + private ArrowFlightClientHandler client; /** * Instantiates a new Arrow Flight Connection. * - * @param driver The JDBC driver to use. - * @param factory The Avatica Factory to use. - * @param url The URL to connect to. - * @param info The properties of this connection. - * @throws SQLException If the connection cannot be established. + * @param driver + * The JDBC driver to use. + * @param factory + * The Avatica Factory to use. + * @param url + * The URL to connect to. + * @param info + * The properties of this connection. + * @throws SQLException + * If the connection cannot be established. */ - public ArrowFlightConnection(final ArrowFlightJdbcDriver driver, + protected ArrowFlightConnection(final ArrowFlightJdbcDriver driver, final AvaticaFactory factory, final String url, final Properties info) - throws SQLException { + throws SQLException { super(driver, factory, url, info); - allocator = new RootAllocator( - Integer.MAX_VALUE); + allocator = new RootAllocator(Integer.MAX_VALUE); try { loadClient(); } catch (final SQLException e) { allocator.close(); - throw e; + throw new SQLException("Failed to initialize Flight Client.", e); } } @@ -90,75 +111,81 @@ public ArrowFlightConnection(final ArrowFlightJdbcDriver driver, private void loadClient() throws SQLException { if (client != null) { - throw new IllegalStateException("Client already loaded."); + throw new SQLException("Client already loaded.", + new IllegalStateException(new InstanceAlreadyExistsException())); } - final String host = (String) info.getOrDefault(DefaultProperty.HOST.toString(), - "localhost"); - Preconditions.checkArgument(!host.trim().isEmpty()); + // =================== [ LOCATION CONFIG ] =================== + final Map.Entry forHost = HOST.getEntry(); + + final String host = (String) info.getOrDefault(forHost.getKey(), + forHost.getValue()); + Preconditions.checkArgument(!Strings.isNullOrEmpty(host)); + + final Map.Entry forPort = PORT.getEntry(); - final int port = Integer.parseInt((String) info.getOrDefault( - DefaultProperty.PORT.toString(), "32010")); - Preconditions.checkArgument(0 < port && port < 65536); + final int port = Preconditions.checkElementIndex( + Integer.parseInt(Objects + .toString(info.getOrDefault(forPort.getKey(), forPort.getValue()))), + 65536); - @Nullable - final String username = - info.getProperty(DefaultProperty.USER.toString()); + // =================== [ CREDENTIALS CONFIG ] =================== + final Map.Entry forUsername = USERNAME.getEntry(); - @Nullable - final String password = - info.getProperty(DefaultProperty.PASS.toString()); + final String username = (String) info.getOrDefault(forUsername.getKey(), + forUsername.getValue()); - final boolean useTls = ((String) info.getOrDefault(DefaultProperty.USE_TLS - .toString(), "false")) - .equalsIgnoreCase("true"); + final Map.Entry forPassword = PASSWORD.getEntry(); - final boolean authenticate = username != null; + final String password = (String) info.getOrDefault(forPassword.getKey(), + forPassword.getValue()); - if (!useTls) { + // =================== [ ENCRYPTION CONFIG ] =================== + final Map.Entry forKeyStorePath = KEYSTORE_PATH.getEntry(); - if (authenticate) { - client = ArrowFlightClient.getBasicClientAuthenticated(allocator, host, - port, username, password, null); - return; - } + final String keyStorePath = (String) info + .getOrDefault(forKeyStorePath.getKey(), forKeyStorePath.getValue()); - client = ArrowFlightClient.getBasicClientNoAuth(allocator, host, port, - null); - return; + final Map.Entry forKeyStorePass = KEYSTORE_PASS.getEntry(); + final String keyStorePassword = (String) info + .getOrDefault(forKeyStorePass.getKey(), forKeyStorePass.getValue()); + + // =================== [ CLIENT GENERATION ] =================== + try { + client = ArrowFlightClientHandler.getClient(allocator, host, port, + username, password, getHeaders(), keyStorePath, keyStorePassword); + } catch (GeneralSecurityException | IOException e) { + throw new SQLException("Failed to connect to the Arrow Flight client.", + e); } + } + + private HeaderCallOption getHeaders() { - final String keyStorePath = info.getProperty( - DefaultProperty.KEYSTORE_PATH.toString()); - final String keyStorePass = info.getProperty( - DefaultProperty.KEYSTORE_PASS.toString()); + final CallHeaders headers = new FlightCallHeaders(); - if (authenticate) { - client = ArrowFlightClient.getEncryptedClientAuthenticated(allocator, - host, port, null, username, password, keyStorePath, keyStorePass); - return; + final Iterator> properties = info.entrySet() + .iterator(); + + while (properties.hasNext()) { + + final Map.Entry entry = properties.next(); + + headers.insert(Objects.toString(entry.getKey()), + Objects.toString(entry.getValue())); } - client = ArrowFlightClient.getEncryptedClientNoAuth(allocator, host, - port, null, keyStorePath, keyStorePass); + return new HeaderCallOption(headers); } @Override public void close() throws SQLException { - try { - client.close(); - } catch (final Exception e) { - throw new SQLException( - "Failed to close the connection " + - "to the Arrow Flight client.", e); - } try { - allocator.close(); + AutoCloseables.close(client, allocator); } catch (final Exception e) { - throw new SQLException("Failed to close the resource allocator used " + - "by the Arrow Flight client.", e); + throw new SQLException("Failed to close resources.", e); } super.close(); diff --git a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java index 4b404bb0d6d99..a246a5aba6105 100644 --- a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java +++ b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java @@ -17,6 +17,9 @@ package org.apache.arrow.driver.jdbc; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.HOST; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.PORT; + import java.io.BufferedReader; import java.io.FileInputStream; import java.io.IOException; @@ -30,7 +33,8 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; -import org.apache.arrow.driver.jdbc.utils.DefaultProperty; +import javax.annotation.RegEx; + import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.util.Preconditions; import org.apache.calcite.avatica.AvaticaConnection; @@ -46,24 +50,29 @@ public class ArrowFlightJdbcDriver extends UnregisteredDriver { private static final String CONNECT_STRING_PREFIX = "jdbc:arrow-flight://"; - private static final Pattern urlRegExPattern = Pattern.compile("^(" + - CONNECT_STRING_PREFIX + ")" + - "(\\w+):([\\d]+)\\/*\\?*([[\\w]+=[\\w]+&?]*)?"); - + + private static final Pattern urlRegExPattern; + private static DriverVersion version; static { - (new ArrowFlightJdbcDriver()).register(); + @RegEx + final String pattern = "^(" + CONNECT_STRING_PREFIX + ")" + + "(\\w+):([\\d]+)\\/*\\?*([[\\w]+=[\\w]+&?]*)?"; + + urlRegExPattern = Pattern.compile(pattern); + new ArrowFlightJdbcDriver().register(); } @Override public Connection connect(final String url, final Properties info) throws SQLException { - final Properties clonedProperties = (Properties) info.clone(); + // FIXME DO NOT tamper with the original Properties! + final Properties clonedProperties = info; try { - final Map args = getUrlsArgs( + final Map args = getUrlsArgs( Preconditions.checkNotNull(url)); clonedProperties.putAll(args); @@ -88,34 +97,33 @@ protected DriverVersion createDriverVersion() { break CreateVersionIfNull; } - try (Reader reader = - new BufferedReader(new InputStreamReader( - new FileInputStream("target/flight.properties"), "UTF-8"))) { - Properties properties = new Properties(); + try (Reader reader = new BufferedReader(new InputStreamReader( + new FileInputStream("target/flight.properties"), "UTF-8"))) { + final Properties properties = new Properties(); properties.load(reader); - String parentName = properties.getProperty( - "org.apache.arrow.flight.name"); - String parentVersion = properties.getProperty( - "org.apache.arrow.flight.version"); - String[] pVersion = parentVersion.split("\\."); + final String parentName = properties + .getProperty("org.apache.arrow.flight.name"); + final String parentVersion = properties + .getProperty("org.apache.arrow.flight.version"); + final String[] pVersion = parentVersion.split("\\."); - int parentMajorVersion = Integer.parseInt(pVersion[0]); - int parentMinorVersion = Integer.parseInt(pVersion[1]); + final int parentMajorVersion = Integer.parseInt(pVersion[0]); + final int parentMinorVersion = Integer.parseInt(pVersion[1]); - String childName = properties.getProperty( - "org.apache.arrow.flight.jdbc-driver.name"); - String childVersion = properties.getProperty( - "org.apache.arrow.flight.jdbc-driver.version"); - String[] cVersion = childVersion.split("\\."); + final String childName = properties + .getProperty("org.apache.arrow.flight.jdbc-driver.name"); + final String childVersion = properties + .getProperty("org.apache.arrow.flight.jdbc-driver.version"); + final String[] cVersion = childVersion.split("\\."); - int childMajorVersion = Integer.parseInt(cVersion[0]); - int childMinorVersion = Integer.parseInt(cVersion[1]); + final int childMajorVersion = Integer.parseInt(cVersion[0]); + final int childMinorVersion = Integer.parseInt(cVersion[1]); version = new DriverVersion(childName, childVersion, parentName, parentVersion, true, childMajorVersion, childMinorVersion, parentMajorVersion, parentMinorVersion); - } catch (IOException e) { + } catch (final IOException e) { throw new RuntimeException("Failed to load driver version.", e); } } @@ -148,24 +156,34 @@ public boolean acceptsURL(final String url) throws SQLException { * @throws SQLException * If an error occurs while trying to parse the URL. */ - private Map getUrlsArgs(final String url) + private Map getUrlsArgs(final String url) throws SQLException { /* - * URL must ALWAYS follow the pattern: + * FIXME Refactor this sub-optimal approach to URL parsing later. + * + * Perhaps this logic should be inside a utility class, separated from this + * one, so as to better delegate responsibilities and concerns throughout + * the code and increase maintainability. + * + * ===== + * + * Keep in mind that the URL must ALWAYS follow the pattern: * "jdbc:arrow-flight://:[/?param1=value1¶m2=value2&(...)]." + * + * TODO Come up with a RegEx better than #urlRegExPattern. */ final Matcher matcher = urlRegExPattern.matcher(url); - + if (!matcher.matches()) { throw new SQLException("Malformed/invalid URL!"); } - final Map resultMap = new HashMap<>(); + final Map resultMap = new HashMap<>(); // Group 1 contains the prefix -- start from 2. - resultMap.put(DefaultProperty.HOST.toString(), matcher.group(2)); - resultMap.put(DefaultProperty.PORT.toString(), matcher.group(3)); + resultMap.put(HOST.getEntry().getKey(), matcher.group(2)); + resultMap.put(PORT.getEntry().getKey(), matcher.group(3)); // Group 4 contains all optional parameters, if provided -- must check. final String extraParams = matcher.group(4); diff --git a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightClientHandler.java b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightClientHandler.java index 0016375a4ed3d..2bc9fa2653ff9 100644 --- a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightClientHandler.java +++ b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightClientHandler.java @@ -18,21 +18,12 @@ package org.apache.arrow.driver.jdbc.client; import java.io.IOException; -import java.io.InputStream; -import java.nio.charset.StandardCharsets; import java.security.GeneralSecurityException; -import java.util.ArrayDeque; -import java.util.Deque; -import java.util.List; -import java.util.Optional; -import java.util.stream.Collectors; import javax.annotation.Nullable; import org.apache.arrow.driver.jdbc.client.utils.ClientAuthenticationUtils; import org.apache.arrow.flight.FlightClient; -import org.apache.arrow.flight.FlightDescriptor; -import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.HeaderCallOption; @@ -41,7 +32,9 @@ import org.apache.arrow.flight.auth2.ClientIncomingAuthHeaderMiddleware; import org.apache.arrow.flight.grpc.CredentialCallOption; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.VectorSchemaRoot; + +import com.google.common.base.Optional; /** * An adhoc {@link FlightClient} wrapper, used to access the client. Allows for @@ -49,8 +42,6 @@ */ public class ArrowFlightClientHandler implements FlightClientHandler { - private final Deque resources = - new ArrayDeque<>(); private final FlightClient client; @Nullable @@ -60,26 +51,25 @@ public class ArrowFlightClientHandler implements FlightClientHandler { private HeaderCallOption properties; protected ArrowFlightClientHandler(final FlightClient client, - @Nullable final CredentialCallOption token, - @Nullable final HeaderCallOption properties) { + @Nullable final CredentialCallOption token, + @Nullable final HeaderCallOption properties) { this(client, token); this.properties = properties; } protected ArrowFlightClientHandler(final FlightClient client, - @Nullable final CredentialCallOption token) { + @Nullable final CredentialCallOption token) { this(client); this.token = token; } protected ArrowFlightClientHandler(final FlightClient client, - final HeaderCallOption properties) { + final HeaderCallOption properties) { this(client, null, properties); } protected ArrowFlightClientHandler(final FlightClient client) { this.client = client; - this.resources.add(this.client); } /** @@ -97,120 +87,151 @@ protected final FlightClient getClient() { * @return the bearer token, if it exists; otherwise, empty. */ protected final Optional getBearerToken() { - return Optional.ofNullable(token); + return Optional.fromNullable(token); } /** * Gets the headers for subsequent calls to this client. * - * @return the {@link #properties} of this client, if they exist; otherwise, empty. + * @return the {@link #properties} of this client, if they exist; otherwise, + * empty. */ protected final Optional getProperties() { - return Optional.ofNullable(properties); + return Optional.fromNullable(properties); } /** * Makes an RPC "getInfo" request with the given query and client properties * in order to retrieve the metadata associated with a set of data records. * - * @param query The query to retrieve FlightInfo for. + * @param query + * The query to retrieve FlightInfo for. * @return a {@link FlightInfo} object. */ protected FlightInfo getInfo(final String query) { - return client.getInfo(FlightDescriptor.command(query.getBytes(StandardCharsets.UTF_8)), - token); + return null; } @Override - public List getFlightStreams(final String query) { - final FlightInfo flightInfo = getInfo(query); - final List endpoints = flightInfo.getEndpoints(); - - final List streams = - endpoints.stream().map(flightEndpoint -> client.getStream(flightEndpoint.getTicket(), token)) - .collect(Collectors.toList()); - streams.forEach(resources::addFirst); + public VectorSchemaRoot runQuery(final String query) throws Exception { + // TODO Auto-generated method stub + return null; + } - return streams; + @Override + public FlightStream getStream(final String query) { + // TODO Auto-generated method stub + return null; } @Override public final void close() throws Exception { try { - AutoCloseables.close(resources); - } catch (final Exception e) { - throw new IOException("Failed to close resources.", e); + client.close(); + } catch (final InterruptedException e) { + /* + * TODO Consider using a proper logger (e.g., Avatica's Log4JLogger.) + * + * This portion of the code should probably be concerned about propagating + * that an Exception has occurred, as opposed to simply "eating it up." + */ + System.out.println("[WARNING] Failed to close resource."); + + /* + * FIXME Should we really be doing this? + * + * Perhaps a better idea is to throw the aforementioned exception and, if + * necessary, handle it later. + */ + e.printStackTrace(); } } /** * Gets a new client based upon provided info. * - * @param allocator The {@link BufferAllocator}. - * @param host The host to connect to. - * @param port The port to connect to. - * @param username The username for authentication, if needed. - * @param password The password for authentication, if needed. - * @param properties The {@link HeaderCallOption} of this client, if needed. - * @param keyStorePath The keystore path for establishing a TLS-encrypted connection, if - * needed. - * @param keyStorePass The keystore password for establishing a TLS-encrypted connection, - * if needed. - * @return a new {@link ArrowFlightClientHandler} based upon the aforementioned information. - * @throws GeneralSecurityException If a certificate-related error occurs. - * @throws IOException If an error occurs while trying to establish a connection to the - * client. + * @param allocator + * The {@link BufferAllocator}. + * @param host + * The host to connect to. + * @param port + * The port to connect to. + * @param username + * The username for authentication, if needed. + * @param password + * The password for authentication, if needed. + * @param properties + * The {@link HeaderCallOption} of this client, if needed. + * @param keyStorePath + * The keystore path for establishing a TLS-encrypted connection, if + * needed. + * @param keyStorePass + * The keystore password for establishing a TLS-encrypted connection, + * if needed. + * @return a new {@link ArrowFlightClientHandler} based upon the + * aforementioned information. + * @throws GeneralSecurityException + * If a certificate-related error occurs. + * @throws IOException + * If an error occurs while trying to establish a connection to the + * client. */ public static final ArrowFlightClientHandler getClient( final BufferAllocator allocator, final String host, final int port, @Nullable final String username, @Nullable final String password, @Nullable final HeaderCallOption properties, - final boolean useTls, @Nullable final String keyStorePath, @Nullable final String keyStorePass) throws GeneralSecurityException, IOException { - /* - * TODO Too many if/else clauses: REDUCE somehow. - * - * Do NOT resort to creating labels and breaking from them! A better - * alternative would be splitting this method into smaller ones. - */ final FlightClient.Builder builder = FlightClient.builder() .allocator(allocator); ArrowFlightClientHandler handler; - if (useTls || keyStorePath != null) { + DetermineEncryption: { + /* + * Check whether to use TLS encryption based upon: + * "Was the keystore path provided?" + */ + final boolean useTls = Optional.fromNullable(keyStorePath).isPresent(); + + if (!useTls) { + + // Build a secure TLS-encrypted connection. + builder.location(Location.forGrpcInsecure(host, port)); + break DetermineEncryption; + } + // Build a secure TLS-encrypted connection. - builder.location(Location.forGrpcTls(host, port)).useTls(); - } else { - // Build an insecure, basic connection. - builder.location(Location.forGrpcInsecure(host, port)); + builder.location(Location.forGrpcTls(host, port)).useTls() + .trustedCertificates(ClientAuthenticationUtils + .getCertificateStream(keyStorePath, keyStorePass)); } - if (keyStorePath != null) { - final InputStream certificateStream = ClientAuthenticationUtils.getCertificateStream(keyStorePath, keyStorePass); - builder.trustedCertificates(certificateStream); - } + DetermineAuthentication: { + + /* + * Check whether to use username/password credentials to authenticate to + * the Flight Client. + */ + final boolean useAuthentication = Optional.fromNullable(username) + .isPresent(); + + if (!useAuthentication) { + + final FlightClient client = builder.build(); + + // Build an unauthenticated client. + handler = new ArrowFlightClientHandler(client, properties); + break DetermineAuthentication; + } - /* - * Check whether to use username/password credentials to authenticate to the - * Flight Client. - */ - final boolean useAuthentication = username != null; - final FlightClient client; - - if (!useAuthentication) { - client = builder.build(); - // Build an unauthenticated client. - handler = new ArrowFlightClientHandler(client, properties); - } else { final ClientIncomingAuthHeaderMiddleware.Factory factory = new ClientIncomingAuthHeaderMiddleware.Factory( new ClientBearerHeaderHandler()); builder.intercept(factory); - client = builder.build(); + final FlightClient client = builder.build(); // Build an authenticated client. handler = new ArrowFlightClientHandler(client, ClientAuthenticationUtils @@ -218,23 +239,31 @@ public static final ArrowFlightClientHandler getClient( properties); } - handler.resources.addLast(client); return handler; } /** * Gets a new client based upon provided info. * - * @param allocator The {@link BufferAllocator}. - * @param host The host to connect to. - * @param port The port to connect to. - * @param username The username for authentication, if needed. - * @param password The password for authentication, if needed. - * @param properties The {@link HeaderCallOption} of this client, if needed. - * @return a new {@link ArrowFlightClientHandler} based upon the aforementioned information. - * @throws GeneralSecurityException If a certificate-related error occurs. - * @throws IOException If an error occurs while trying to establish a connection to the - * client. + * @param allocator + * The {@link BufferAllocator}. + * @param host + * The host to connect to. + * @param port + * The port to connect to. + * @param username + * The username for authentication, if needed. + * @param password + * The password for authentication, if needed. + * @param properties + * The {@link HeaderCallOption} of this client, if needed. + * @return a new {@link ArrowFlightClientHandler} based upon the + * aforementioned information. + * @throws GeneralSecurityException + * If a certificate-related error occurs. + * @throws IOException + * If an error occurs while trying to establish a connection to the + * client. */ public static final ArrowFlightClientHandler getClient( final BufferAllocator allocator, final String host, final int port, @@ -243,21 +272,29 @@ public static final ArrowFlightClientHandler getClient( throws GeneralSecurityException, IOException { return getClient(allocator, host, port, username, password, properties, - false, null, null); + null, null); } /** * Gets a new client based upon provided info. * - * @param allocator The {@link BufferAllocator}. - * @param host The host to connect to. - * @param port The port to connect to. - * @param username The username for authentication, if needed. - * @param password The password for authentication, if needed. - * @return a new {@link ArrowFlightClientHandler} based upon the aforementioned information. - * @throws GeneralSecurityException If a certificate-related error occurs. - * @throws IOException If an error occurs while trying to establish a connection to the - * client. + * @param allocator + * The {@link BufferAllocator}. + * @param host + * The host to connect to. + * @param port + * The port to connect to. + * @param username + * The username for authentication, if needed. + * @param password + * The password for authentication, if needed. + * @return a new {@link ArrowFlightClientHandler} based upon the + * aforementioned information. + * @throws GeneralSecurityException + * If a certificate-related error occurs. + * @throws IOException + * If an error occurs while trying to establish a connection to the + * client. */ public static final ArrowFlightClientHandler getClient( final BufferAllocator allocator, final String host, final int port, @@ -270,13 +307,19 @@ public static final ArrowFlightClientHandler getClient( /** * Gets a new client based upon provided info. * - * @param allocator The {@link BufferAllocator}. - * @param host The host to connect to. - * @param port The port to connect to. - * @return a new {@link ArrowFlightClientHandler} based upon the aforementioned information. - * @throws GeneralSecurityException If a certificate-related error occurs. - * @throws IOException If an error occurs while trying to establish a connection to the - * client. + * @param allocator + * The {@link BufferAllocator}. + * @param host + * The host to connect to. + * @param port + * The port to connect to. + * @return a new {@link ArrowFlightClientHandler} based upon the + * aforementioned information. + * @throws GeneralSecurityException + * If a certificate-related error occurs. + * @throws IOException + * If an error occurs while trying to establish a connection to the + * client. */ public static final ArrowFlightClientHandler getClient( final BufferAllocator allocator, final String host, final int port) @@ -288,17 +331,27 @@ public static final ArrowFlightClientHandler getClient( /** * Gets a new client based upon provided info. * - * @param allocator The {@link BufferAllocator}. - * @param host The host to connect to. - * @param port The port to connect to. - * @param properties The {@link HeaderCallOption} of this client, if needed. - * @param keyStorePath The keystore path for establishing a TLS-encrypted connection, if - * needed. - * @param keyStorePass The keystore password for establishing a TLS-encrypted connection, - * if needed. - * @return a new {@link ArrowFlightClientHandler} based upon the aforementioned information. - * @throws GeneralSecurityException If a certificate-related error occurs. - * @throws IOException If an error occurs while trying to establish a connection to the client. + * @param allocator + * The {@link BufferAllocator}. + * @param host + * The host to connect to. + * @param port + * The port to connect to. + * @param properties + * The {@link HeaderCallOption} of this client, if needed. + * @param keyStorePath + * The keystore path for establishing a TLS-encrypted connection, if + * needed. + * @param keyStorePass + * The keystore password for establishing a TLS-encrypted connection, + * if needed. + * @return a new {@link ArrowFlightClientHandler} based upon the + * aforementioned information. + * @throws GeneralSecurityException + * If a certificate-related error occurs. + * @throws IOException + * If an error occurs while trying to establish a connection to the + * client. */ public static final ArrowFlightClientHandler getClient( final BufferAllocator allocator, final String host, final int port, @@ -306,6 +359,7 @@ public static final ArrowFlightClientHandler getClient( @Nullable final String keyStorePath, @Nullable final String keyStorePass) throws GeneralSecurityException, IOException { - return getClient(allocator, host, port, null, null, properties, true, keyStorePath, keyStorePass); + return getClient(allocator, host, port, null, null, properties, + keyStorePath, keyStorePass); } } diff --git a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/FlightClientHandler.java b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/FlightClientHandler.java index b90ddb831ce92..63a2316e6be7c 100644 --- a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/FlightClientHandler.java +++ b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/FlightClientHandler.java @@ -17,23 +17,34 @@ package org.apache.arrow.driver.jdbc.client; -import java.util.List; - import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.vector.VectorSchemaRoot; /** * A wrapper for a {@link FlightClient}. */ public interface FlightClientHandler extends AutoCloseable { + /** + * Makes RPC requests to the Dremio Flight Server Endpoint to retrieve results + * of the provided SQL query. + * + * @param query + * The SQL query to execute. + * @throws Exception + * If an error occurs during query execution. + */ + VectorSchemaRoot runQuery(String query) throws Exception; + /** * Makes an RPC "getStream" request based on the provided {@link FlightInfo} * object. Retrieves result of the query previously prepared with "getInfo." * - * @param query The query. + * @param query + * The query. * @return a {@code FlightStream} of results. */ - List getFlightStreams(String query); + FlightStream getStream(String query); } diff --git a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtils.java b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtils.java index 29cae17b59198..558b889c18702 100644 --- a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtils.java +++ b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/utils/ClientAuthenticationUtils.java @@ -116,6 +116,7 @@ public static InputStream getCertificateStream(final String keyStorePath, try (final InputStream keyStoreStream = Files .newInputStream(Paths.get(Preconditions.checkNotNull(keyStorePath)))) { + keyStore.load(keyStoreStream, Preconditions.checkNotNull(keyStorePass).toCharArray()); } @@ -123,8 +124,11 @@ public static InputStream getCertificateStream(final String keyStorePath, final Enumeration aliases = keyStore.aliases(); while (aliases.hasMoreElements()) { + final String alias = aliases.nextElement(); + if (keyStore.isCertificateEntry(alias)) { + return toInputStream(keyStore.getCertificate(alias)); } } diff --git a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/utils/DefaultProperty.java b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/utils/BaseProperty.java similarity index 53% rename from java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/utils/DefaultProperty.java rename to java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/utils/BaseProperty.java index db42ba3cb0e0b..7f668378d6825 100644 --- a/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/utils/DefaultProperty.java +++ b/java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/utils/BaseProperty.java @@ -17,27 +17,41 @@ package org.apache.arrow.driver.jdbc.utils; +import java.util.HashMap; +import java.util.Map; + /** * An enum for centralizing default property names. */ -public enum DefaultProperty { +public enum BaseProperty { // TODO These names are up to discussion. - HOST("host"), - PORT("port"), - USER("user"), - PASS("password"), - USE_TLS("useTls"), - KEYSTORE_PATH("keyStorePath"), - KEYSTORE_PASS("keyStorePass"); + HOST("host", "localhost"), PORT("port", 32210), USERNAME("user"), PASSWORD( + "password"), ENCRYPT("useTls", + false), KEYSTORE_PATH("keyStorePath"), KEYSTORE_PASS("keyStorePass"); private final String repr; + private Object def; + + BaseProperty(final String repr, final Object def) { + this(repr); + this.def = def; + } - private DefaultProperty(final String repr) { + BaseProperty(final String repr) { this.repr = repr; } - @Override - public String toString() { - return repr; + /** + * Gets the {@link Map.Entry} representation of this property, where + * {@link Map.Entry#getKey} gets the name and {@link Map.Entry#getValue} gets + * the default value of this property, or {@code null} if it lacks one. + * + * @return the entry of this property. + */ + public Map.Entry getEntry() { + final Map map = new HashMap<>(); + map.put(repr, def); + + return map.entrySet().iterator().next(); } } diff --git a/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ArrowFlightJdbcDriverTest.java b/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ArrowFlightJdbcDriverTest.java index 1fafa78bc4c5e..7cc0244502bc7 100644 --- a/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ArrowFlightJdbcDriverTest.java +++ b/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ArrowFlightJdbcDriverTest.java @@ -17,6 +17,8 @@ package org.apache.arrow.driver.jdbc.test; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.HOST; +import static org.apache.arrow.driver.jdbc.utils.BaseProperty.PORT; import static org.junit.jupiter.api.Assertions.assertEquals; import java.lang.reflect.InvocationTargetException; @@ -33,7 +35,6 @@ import org.apache.arrow.driver.jdbc.test.utils.FlightTestUtils; import org.apache.arrow.driver.jdbc.test.utils.PropertiesSample; import org.apache.arrow.driver.jdbc.test.utils.UrlSample; -import org.apache.arrow.driver.jdbc.utils.DefaultProperty; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightProducer; import org.apache.arrow.flight.FlightServer; @@ -202,8 +203,8 @@ public void testShouldThrowExceptionWhenAttemptingToConnectToUrlNoHost() throws Exception { final Driver driver = new ArrowFlightJdbcDriver(); - final String malformedUri = "arrow-jdbc://" + - ":" + server.getLocation().getUri().getPort(); + final String malformedUri = "arrow-jdbc://" + ":" + + server.getLocation().getUri().getPort(); driver.connect(malformedUri, PropertiesSample.UNSUPPORTED.getProperties()); } @@ -233,10 +234,10 @@ public void testDriverUrlParsingMechanismShouldReturnTheDesiredArgsFromUrl() assertEquals(5, parsedArgs.size()); // Check host == the provided host - assertEquals(parsedArgs.get(DefaultProperty.HOST.toString()), "localhost"); + assertEquals(parsedArgs.get(HOST.getEntry().getKey()), "localhost"); // Check port == the provided port - assertEquals(parsedArgs.get(DefaultProperty.PORT.toString()), "2222"); + assertEquals(parsedArgs.get(PORT.getEntry().getKey()), "2222"); // Check all other non-default arguments assertEquals(parsedArgs.get("key1"), "value1"); @@ -248,7 +249,8 @@ public void testDriverUrlParsingMechanismShouldReturnTheDesiredArgsFromUrl() * Tests whether an exception is thrown upon attempting to connect to a * malformed URI. * - * @throws Exception If an error occurs. + * @throws Exception + * If an error occurs. */ @SuppressWarnings("unchecked") @Test(expected = SQLException.class) @@ -262,10 +264,9 @@ public void testDriverUrlParsingMechanismShouldThrowExceptionUponProvidedWithMal getUrlsArgs.setAccessible(true); try { - final Map parsedArgs = (Map) getUrlsArgs - .invoke(driver, - "jdbc:arrow-flight://localhost:2222/?k1=v1&m="); - } catch (InvocationTargetException e) { + final Map parsedArgs = (Map) getUrlsArgs + .invoke(driver, "jdbc:arrow-flight://localhost:2222/?k1=v1&m="); + } catch (final InvocationTargetException e) { throw (SQLException) e.getCause(); } } @@ -283,7 +284,7 @@ private CallHeaderAuthenticator.AuthResult validate(final String username, final String password) { if (Strings.isNullOrEmpty(username)) { throw CallStatus.UNAUTHENTICATED - .withDescription("Credentials not supplied.").toRuntimeException(); + .withDescription("Credentials not supplied.").toRuntimeException(); } final String identity; if (testUtils.getUsername1().equals(username) && @@ -291,8 +292,8 @@ private CallHeaderAuthenticator.AuthResult validate(final String username, identity = testUtils.getUsername1(); } else { throw CallStatus.UNAUTHENTICATED - .withDescription("Username or password is invalid.") - .toRuntimeException(); + .withDescription("Username or password is invalid.") + .toRuntimeException(); } return () -> identity; } diff --git a/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ConnectionTest.java b/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ConnectionTest.java index 94c1b96bcda2c..5335cbce645af 100644 --- a/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ConnectionTest.java +++ b/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ConnectionTest.java @@ -25,8 +25,8 @@ import java.sql.SQLException; import java.util.Properties; -import org.apache.arrow.driver.jdbc.ArrowFlightClient; import org.apache.arrow.driver.jdbc.ArrowFlightJdbcDriver; +import org.apache.arrow.driver.jdbc.client.ArrowFlightClientHandler; import org.apache.arrow.driver.jdbc.test.utils.FlightTestUtils; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightProducer; @@ -64,15 +64,16 @@ public class ConnectionTest { public void setUp() throws Exception { allocator = new RootAllocator(Long.MAX_VALUE); - flightTestUtils = new FlightTestUtils("localhost", "flight1", - "woho1", "invalid", "wrong"); + flightTestUtils = new FlightTestUtils("localhost", "flight1", "woho1", + "invalid", "wrong"); - final FlightProducer flightProducer = flightTestUtils.getFlightProducer(allocator); - this.server = flightTestUtils.getStartedServer((location -> FlightServer - .builder(allocator, location, flightProducer) - .headerAuthenticator(new GeneratedBearerTokenAuthenticator( - new BasicCallHeaderAuthenticator(this::validate))) - .build())); + final FlightProducer flightProducer = flightTestUtils + .getFlightProducer(allocator); + this.server = flightTestUtils.getStartedServer( + location -> FlightServer.builder(allocator, location, flightProducer) + .headerAuthenticator(new GeneratedBearerTokenAuthenticator( + new BasicCallHeaderAuthenticator(this::validate))) + .build()); serverUrl = flightTestUtils.getConnectionPrefix() + flightTestUtils.getLocalhost() + ":" + this.server.getPort(); @@ -98,7 +99,7 @@ private CallHeaderAuthenticator.AuthResult validate(final String username, final String password) { if (Strings.isNullOrEmpty(username)) { throw CallStatus.UNAUTHENTICATED - .withDescription("Credentials not supplied.").toRuntimeException(); + .withDescription("Credentials not supplied.").toRuntimeException(); } final String identity; if (flightTestUtils.getUsername1().equals(username) && @@ -106,8 +107,8 @@ private CallHeaderAuthenticator.AuthResult validate(final String username, identity = flightTestUtils.getUsername1(); } else { throw CallStatus.UNAUTHENTICATED - .withDescription("Username or password is invalid.") - .toRuntimeException(); + .withDescription("Username or password is invalid.") + .toRuntimeException(); } return () -> identity; } @@ -127,8 +128,8 @@ public void testUnencryptedConnectionShouldOpenSuccessfullyWhenProvidedValidCred properties.put("user", flightTestUtils.getUsername1()); properties.put("password", flightTestUtils.getPassword1()); - try (Connection connection = DriverManager - .getConnection(serverUrl, properties)) { + try (Connection connection = DriverManager.getConnection(serverUrl, + properties)) { assert connection.isValid(300); } } @@ -161,12 +162,12 @@ public void testUnencryptedConnectionWithEmptyHost() * on error. */ @Test - public void testGetBasicClientAuthenticatedShouldOpenConnection() throws Exception { + public void testGetBasicClientAuthenticatedShouldOpenConnection() + throws Exception { - try (ArrowFlightClient client = ArrowFlightClient.getBasicClientAuthenticated( + try (ArrowFlightClientHandler client = ArrowFlightClientHandler.getClient( allocator, flightTestUtils.getLocalhost(), this.server.getPort(), - flightTestUtils.getUsername1(), flightTestUtils.getPassword1(), - null)) { + flightTestUtils.getUsername1(), flightTestUtils.getPassword1())) { assertNotNull(client); } } @@ -178,7 +179,7 @@ public void testGetBasicClientAuthenticatedShouldOpenConnection() throws Excepti * @throws SQLException * on error. */ - @Test(expected = IllegalArgumentException.class) + @Test(expected = IndexOutOfBoundsException.class) public void testUnencryptedConnectionProvidingInvalidPort() throws Exception { final Properties properties = new Properties(); @@ -202,9 +203,8 @@ public void testUnencryptedConnectionProvidingInvalidPort() @Test public void testGetBasicClientNoAuthShouldOpenConnection() throws Exception { - try (ArrowFlightClient client = ArrowFlightClient.getBasicClientNoAuth( - allocator, flightTestUtils.getLocalhost(), this.server.getPort(), - null)) { + try (ArrowFlightClientHandler client = ArrowFlightClientHandler.getClient( + allocator, flightTestUtils.getLocalhost(), this.server.getPort())) { assertNotNull(client); } } diff --git a/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ConnectionTlsTest.java b/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ConnectionTlsTest.java index b4bf05d4c3ebd..2958dd8c4f611 100644 --- a/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ConnectionTlsTest.java +++ b/java/flight/flight-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/test/ConnectionTlsTest.java @@ -21,12 +21,13 @@ import java.io.IOException; import java.net.URI; +import java.security.cert.CertificateException; import java.sql.Connection; import java.sql.DriverManager; import java.sql.SQLException; import java.util.Properties; -import org.apache.arrow.driver.jdbc.ArrowFlightClient; +import org.apache.arrow.driver.jdbc.client.ArrowFlightClientHandler; import org.apache.arrow.driver.jdbc.test.utils.FlightTestUtils; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightProducer; @@ -58,8 +59,8 @@ public class ConnectionTlsTest { @Before public void setUp() throws Exception { - flightTestUtils = new FlightTestUtils("localhost", "flight1", - "woho1", "invalid", "wrong"); + flightTestUtils = new FlightTestUtils("localhost", "flight1", "woho1", + "invalid", "wrong"); allocator = new RootAllocator(Long.MAX_VALUE); @@ -68,10 +69,9 @@ public void setUp() throws Exception { final FlightProducer flightProducer = flightTestUtils .getFlightProducer(allocator); - this.tlsServer = flightTestUtils.getStartedServer((location -> { + this.tlsServer = flightTestUtils.getStartedServer(location -> { try { - return FlightServer - .builder(allocator, location, flightProducer) + return FlightServer.builder(allocator, location, flightProducer) .useTls(certKey.cert, certKey.key) .headerAuthenticator(new GeneratedBearerTokenAuthenticator( new BasicCallHeaderAuthenticator(this::validate))) @@ -80,7 +80,7 @@ public void setUp() throws Exception { e.printStackTrace(); } return null; - })); + }); serverUrl = flightTestUtils.getConnectionPrefix() + flightTestUtils.getLocalhost() + ":" + this.tlsServer.getPort(); @@ -105,16 +105,16 @@ private CallHeaderAuthenticator.AuthResult validate(final String username, final String password) { if (Strings.isNullOrEmpty(username)) { throw CallStatus.UNAUTHENTICATED - .withDescription("Credentials not supplied.").toRuntimeException(); + .withDescription("Credentials not supplied.").toRuntimeException(); } final String identity; - if (flightTestUtils.getUsername1().equals(username) && flightTestUtils - .getPassword1().equals(password)) { + if (flightTestUtils.getUsername1().equals(username) && + flightTestUtils.getPassword1().equals(password)) { identity = flightTestUtils.getUsername1(); } else { throw CallStatus.UNAUTHENTICATED - .withDescription("Username or password is invalid.") - .toRuntimeException(); + .withDescription("Username or password is invalid.") + .toRuntimeException(); } return () -> identity; } @@ -135,12 +135,12 @@ public void testGetEncryptedClientAuthenticated() throws Exception { final UsernamePasswordCredentials credentials = new UsernamePasswordCredentials( flightTestUtils.getUsername1(), flightTestUtils.getPassword1()); - try (ArrowFlightClient client = ArrowFlightClient - .getEncryptedClientAuthenticated( - allocator, address.getHost(), address.getPort(), - null, credentials.getUserName(), credentials.getPassword(), - keyStorePath, - keyStorePass)) { + try (ArrowFlightClientHandler client = + ArrowFlightClientHandler + .getClient( + allocator, address.getHost(), address.getPort(), + credentials.getUserName(), credentials.getPassword(), + null, keyStorePath, keyStorePass)) { assertNotNull(client); } @@ -153,15 +153,15 @@ public void testGetEncryptedClientAuthenticated() throws Exception { * @throws Exception * on error. */ - @Test(expected = SQLException.class) + @Test(expected = CertificateException.class) public void testGetEncryptedClientWithNoCertificateOnKeyStore() throws Exception { final String noCertificateKeyStorePassword = "flight1"; - try (ArrowFlightClient client = ArrowFlightClient - .getEncryptedClientNoAuth( - allocator, flightTestUtils.getLocalhost(), this.tlsServer.getPort(), - null, noCertificateKeyStorePath, - noCertificateKeyStorePassword)) { + try (ArrowFlightClientHandler client = + ArrowFlightClientHandler + .getClient(allocator, flightTestUtils.getLocalhost(), this.tlsServer.getPort(), + null, noCertificateKeyStorePath, + noCertificateKeyStorePassword)) { Assert.fail(); } } @@ -174,11 +174,12 @@ public void testGetEncryptedClientWithNoCertificateOnKeyStore() throws Exception */ @Test public void testGetNonAuthenticatedEncryptedClientNoAuth() throws Exception { - try (ArrowFlightClient client = ArrowFlightClient - .getEncryptedClientNoAuth( - allocator, flightTestUtils.getLocalhost(), this.tlsServer.getPort(), - null, keyStorePath, - keyStorePass)) { + try (ArrowFlightClientHandler client = + ArrowFlightClientHandler + .getClient( + allocator, flightTestUtils.getLocalhost(), this.tlsServer.getPort(), + null, keyStorePath, + keyStorePass)) { assertNotNull(client); } @@ -191,15 +192,15 @@ public void testGetNonAuthenticatedEncryptedClientNoAuth() throws Exception { * @throws Exception * on error. */ - @Test(expected = SQLException.class) + @Test(expected = IOException.class) public void testGetEncryptedClientWithKeyStoreBadPasswordAndNoAuth() throws Exception { String keyStoreBadPassword = "badPassword"; - try (ArrowFlightClient client = ArrowFlightClient - .getEncryptedClientNoAuth( - allocator, flightTestUtils.getLocalhost(), this.tlsServer.getPort(), - null, keyStorePath, - keyStoreBadPassword)) { + try (ArrowFlightClientHandler client = + ArrowFlightClientHandler.getClient( + allocator, flightTestUtils.getLocalhost(), this.tlsServer.getPort(), + null, keyStorePath, + keyStoreBadPassword)) { Assert.fail(); } } @@ -264,8 +265,8 @@ public void testGetNonAuthenticatedEncryptedConnection() throws Exception { properties.put("keyStorePath", keyStorePath); properties.put("keyStorePass", keyStorePass); - try (Connection connection = DriverManager - .getConnection(serverUrl, properties)) { + try (Connection connection = DriverManager.getConnection(serverUrl, + properties)) { assert connection.isValid(300); }