Skip to content

Commit

Permalink
Merge pull request #11 from abenaru/raise-code-covarege-limit
Browse files Browse the repository at this point in the history 
Raise code coverage to 80%
  • Loading branch information
@jcralmeida @rafael-telles
jcralmeida authored and rafael-telles committed Dec 6, 2021
1 parent 62b4197 commit b37263b
Show file tree
Hide file tree
Showing 10 changed files with 429 additions and 723 deletions.
424 changes: 0 additions & 424 deletions ...ight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightClient.java
View file

This file was deleted.

147 changes: 87 additions & 60 deletions .../flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand All @@ -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<Object, Object> forHost = HOST.getEntry();

final String host = (String) info.getOrDefault(forHost.getKey(),
forHost.getValue());
Preconditions.checkArgument(!Strings.isNullOrEmpty(host));

final Map.Entry<Object, Object> 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<Object, Object> 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<Object, Object> 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<Object, Object> 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<Object, Object> 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<Map.Entry<Object, Object>> properties = info.entrySet()
.iterator();

while (properties.hasNext()) {

final Map.Entry<Object, Object> 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();
Expand Down
84 changes: 51 additions & 33 deletions .../flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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<String, String> args = getUrlsArgs(
final Map<Object, Object> args = getUrlsArgs(
Preconditions.checkNotNull(url));

clonedProperties.putAll(args);
Expand All @@ -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);
}
}
Expand Down Expand Up @@ -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<String, String> getUrlsArgs(final String url)
private Map<Object, Object> 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://<host>:<port>[/?param1=value1&param2=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<String, String> resultMap = new HashMap<>();
final Map<Object, Object> 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);
Expand Down
Loading

0 comments on commit b37263b

Please sign in to comment.