This is an automated email from the ASF dual-hosted git repository.
rongr pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 2dc26ac402 [feature] Add jdbc multistage support (#10292)
2dc26ac402 is described below
commit 2dc26ac402db9fb1d83c0e22b63b66025f60461c
Author: ImprovingRichard <[email protected]>
AuthorDate: Tue Feb 28 12:34:42 2023 -0600
[feature] Add jdbc multistage support (#10292)
* Refactor query option handling
* Remove commented out code
* Remove use of var to support java 8
* Added unit tests for setting query options
* Renamed test cases to match current naming. Added test case for preset
option.
* Updates to address PR comments
* Updated option handling to support booleans and numbers
---
.../org/apache/pinot/client/PinotConnection.java | 49 ++++++++-
.../pinot/client/PinotPreparedStatement.java | 4 +-
.../org/apache/pinot/client/PinotStatement.java | 2 +-
.../org/apache/pinot/client/utils/DriverUtils.java | 38 +++++--
.../pinot/client/PinotPreparedStatementTest.java | 31 ++++++
.../apache/pinot/client/PinotStatementTest.java | 114 ++++++++++++++++++++-
6 files changed, 222 insertions(+), 16 deletions(-)
diff --git
a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotConnection.java
b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotConnection.java
index b7186a7977..9e703b5f3e 100644
---
a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotConnection.java
+++
b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotConnection.java
@@ -24,7 +24,9 @@ import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Arrays;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import java.util.Properties;
import org.apache.pinot.client.base.AbstractBaseConnection;
import org.apache.pinot.client.controller.PinotControllerTransport;
@@ -38,11 +40,16 @@ import org.slf4j.LoggerFactory;
public class PinotConnection extends AbstractBaseConnection {
private static final Logger LOGGER =
LoggerFactory.getLogger(Connection.class);
+ protected static final String[] POSSIBLE_QUERY_OPTIONS = {
+ QueryOptionKey.ENABLE_NULL_HANDLING,
+ QueryOptionKey.USE_MULTISTAGE_ENGINE
+ };
private org.apache.pinot.client.Connection _session;
private boolean _closed;
private String _controllerURL;
private PinotControllerTransport _controllerTransport;
- private final boolean _enableNullHandling;
+ private final Map<String, Object> _queryOptions = new HashMap<String,
Object>();
+
public static final String BROKER_LIST = "brokers";
PinotConnection(String controllerURL, PinotClientTransport transport, String
tenant,
@@ -67,15 +74,49 @@ public class PinotConnection extends AbstractBaseConnection
{
}
_session = new org.apache.pinot.client.Connection(properties, brokers,
transport);
- _enableNullHandling =
Boolean.parseBoolean(properties.getProperty(QueryOptionKey.ENABLE_NULL_HANDLING));
+ for (String possibleQueryOption: POSSIBLE_QUERY_OPTIONS) {
+ Object property = properties.getProperty(possibleQueryOption);
+ if (property != null) {
+ _queryOptions.put(possibleQueryOption, parseOptionValue(property));
+ }
+ }
+ }
+
+ private Object parseOptionValue(Object value) {
+ if (value instanceof String) {
+ String str = (String) value;
+
+ try {
+ Long numVal = Long.valueOf(str);
+ if (numVal != null) {
+ return numVal;
+ }
+ } catch (NumberFormatException e) {
+ }
+
+ try {
+ Double numVal = Double.valueOf(str);
+ if (numVal != null) {
+ return numVal;
+ }
+ } catch (NumberFormatException e) {
+ }
+
+ Boolean boolVal = Boolean.valueOf(str.toLowerCase());
+ if (boolVal != null) {
+ return boolVal;
+ }
+ }
+
+ return value;
}
public org.apache.pinot.client.Connection getSession() {
return _session;
}
- public boolean isNullHandlingEnabled() {
- return _enableNullHandling;
+ public Map<String, Object> getQueryOptions() {
+ return _queryOptions;
}
private List<String> getBrokerList(String controllerURL, String tenant) {
diff --git
a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotPreparedStatement.java
b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotPreparedStatement.java
index 9dbcfd5c50..98b777b3ac 100644
---
a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotPreparedStatement.java
+++
b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotPreparedStatement.java
@@ -51,7 +51,7 @@ public class PinotPreparedStatement extends
AbstractBasePreparedStatement {
if (!DriverUtils.queryContainsLimitStatement(_query)) {
_query += " " + LIMIT_STATEMENT + " " + _maxRows;
}
- _query = DriverUtils.enableNullHandling(_connection, _query);
+ _query = DriverUtils.enableQueryOptions(_query,
_connection.getQueryOptions());
_preparedStatement = new PreparedStatement(_session, _query);
}
@@ -177,7 +177,7 @@ public class PinotPreparedStatement extends
AbstractBasePreparedStatement {
throws SQLException {
validateState();
try {
- _resultSetGroup =
_session.execute(DriverUtils.enableNullHandling(_connection, sql));
+ _resultSetGroup = _session.execute(DriverUtils.enableQueryOptions(sql,
_connection.getQueryOptions()));
if (_resultSetGroup.getResultSetCount() == 0) {
_resultSet = PinotResultSet.empty();
return _resultSet;
diff --git
a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotStatement.java
b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotStatement.java
index 9b325b2f1f..fc36c4b34a 100644
---
a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotStatement.java
+++
b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotStatement.java
@@ -63,7 +63,7 @@ public class PinotStatement extends AbstractBaseStatement {
if (!DriverUtils.queryContainsLimitStatement(sql)) {
sql += " " + LIMIT_STATEMENT + " " + _maxRows;
}
- String enabledSql = DriverUtils.enableNullHandling(_connection, sql);
+ String enabledSql = DriverUtils.enableQueryOptions(sql,
_connection.getQueryOptions());
ResultSetGroup resultSetGroup = _session.execute(enabledSql);
if (resultSetGroup.getResultSetCount() == 0) {
_resultSet = PinotResultSet.empty();
diff --git
a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java
b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java
index 73c1a9a72a..c914f959ac 100644
---
a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java
+++
b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java
@@ -36,12 +36,10 @@ import org.apache.commons.configuration.MapConfiguration;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.NameValuePair;
import org.apache.http.client.utils.URLEncodedUtils;
-import org.apache.pinot.client.PinotConnection;
import org.apache.pinot.common.config.TlsConfig;
import org.apache.pinot.common.utils.TlsUtils;
import org.apache.pinot.core.auth.BasicAuthUtils;
import org.apache.pinot.spi.env.PinotConfiguration;
-import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -220,13 +218,37 @@ public class DriverUtils {
return matcher.find();
}
- public static String enableNullHandling(PinotConnection connection, String
query) {
- if (query.contains(QueryOptionKey.ENABLE_NULL_HANDLING)) {
- return query;
+ public static String enableQueryOptions(String sql, Map<String, Object>
options) {
+ StringBuilder optionsBuilder = new StringBuilder();
+ for (Map.Entry<String, Object> optionEntry: options.entrySet()) {
+ if (!sql.contains(optionEntry.getKey())) {
+
optionsBuilder.append(DriverUtils.createSetQueryOptionString(optionEntry.getKey(),
optionEntry.getValue()));
+ }
+ }
+ optionsBuilder.append(sql);
+ return optionsBuilder.toString();
+ }
+
+ public static String createSetQueryOptionString(String optionKey, Object
optionValue) {
+ StringBuilder optionBuilder = new StringBuilder();
+ optionBuilder.append("SET ").append(optionKey);
+
+ if (optionValue != null) {
+ optionBuilder.append('=');
+
+ if (optionValue instanceof Boolean) {
+ optionBuilder.append(((Boolean) optionValue).booleanValue());
+ } else if (optionValue instanceof Integer || optionValue instanceof
Long) {
+ optionBuilder.append(((Number) optionValue).longValue());
+ } else if (optionValue instanceof Float || optionValue instanceof
Double) {
+ optionBuilder.append(((Number) optionValue).doubleValue());
+ } else {
+ throw new IllegalArgumentException(
+ "Option Type " + optionValue.getClass().getSimpleName() + " is not
supported.");
+ }
}
- return connection.isNullHandlingEnabled()
- ? String.format("SET %s = true; %s",
QueryOptionKey.ENABLE_NULL_HANDLING, query)
- : query;
+ optionBuilder.append(";\n");
+ return optionBuilder.toString();
}
}
diff --git
a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotPreparedStatementTest.java
b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotPreparedStatementTest.java
index 0cc203a994..2b77206c55 100644
---
a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotPreparedStatementTest.java
+++
b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotPreparedStatementTest.java
@@ -26,6 +26,8 @@ import java.sql.Timestamp;
import java.util.Properties;
import org.apache.commons.codec.binary.Hex;
import org.apache.pinot.client.utils.DateTimeUtils;
+import org.apache.pinot.client.utils.DriverUtils;
+import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
import org.testng.Assert;
import org.testng.annotations.Test;
@@ -35,6 +37,7 @@ public class PinotPreparedStatementTest {
"SELECT * FROM dummy WHERE name = ? and age = ? and score = ? and ts = ?
and eligible = ? and sub_score = ?";
public static final String DATE_QUERY = "SELECT * FROM dummy WHERE date = ?
and updated_at = ? and created_at = ?";
public static final String SINGLE_STRING_QUERY = "SELECT * FROM dummy WHERE
value = ?";
+ private static final String BASIC_TEST_QUERY = "SELECT * FROM dummy";
private DummyPinotClientTransport _dummyPinotClientTransport = new
DummyPinotClientTransport();
private DummyPinotControllerTransport _dummyPinotControllerTransport =
DummyPinotControllerTransport.create();
@@ -120,4 +123,32 @@ public class PinotPreparedStatementTest {
Assert.assertEquals(lastExecutedQuery.substring(0,
lastExecutedQuery.indexOf("LIMIT")).trim(),
String.format("SELECT * FROM dummy WHERE value = '%s'",
Hex.encodeHexString(value.getBytes())));
}
+
+ @Test
+ public void testSetEnableNullHandling()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.ENABLE_NULL_HANDLING, "true");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ PreparedStatement preparedStatement =
pinotConnection.prepareStatement(BASIC_TEST_QUERY);
+ preparedStatement.executeQuery();
+ String expectedSql =
+
DriverUtils.createSetQueryOptionString(QueryOptionKey.ENABLE_NULL_HANDLING,
true) + BASIC_TEST_QUERY;
+ Assert.assertEquals(_dummyPinotClientTransport.getLastQuery().substring(0,
expectedSql.length()), expectedSql);
+ }
+
+ @Test
+ public void testSetEnableNullHandling2()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.ENABLE_NULL_HANDLING, "true");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ PreparedStatement preparedStatement = pinotConnection.prepareStatement("");
+ preparedStatement.executeQuery(BASIC_TEST_QUERY);
+ String expectedSql =
+
DriverUtils.createSetQueryOptionString(QueryOptionKey.ENABLE_NULL_HANDLING,
true) + BASIC_TEST_QUERY;
+ Assert.assertEquals(_dummyPinotClientTransport.getLastQuery().substring(0,
expectedSql.length()), expectedSql);
+ }
}
diff --git
a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotStatementTest.java
b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotStatementTest.java
index db4e1dfbb0..fd47c6db90 100644
---
a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotStatementTest.java
+++
b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotStatementTest.java
@@ -20,11 +20,15 @@ package org.apache.pinot.client;
import java.sql.ResultSet;
import java.sql.Statement;
+import java.util.Properties;
+import org.apache.pinot.client.utils.DriverUtils;
+import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
import org.testng.Assert;
import org.testng.annotations.Test;
public class PinotStatementTest {
+ private static final String BASIC_TEST_QUERY = "SELECT * FROM dummy";
private DummyPinotClientTransport _dummyPinotClientTransport = new
DummyPinotClientTransport();
private DummyPinotControllerTransport _dummyPinotControllerTransport =
DummyPinotControllerTransport.create();
@@ -34,8 +38,116 @@ public class PinotStatementTest {
PinotConnection connection =
new PinotConnection("dummy", _dummyPinotClientTransport, "dummy",
_dummyPinotControllerTransport);
Statement statement = new PinotStatement(connection);
- ResultSet resultSet = statement.executeQuery("select * from dummy");
+ ResultSet resultSet = statement.executeQuery(BASIC_TEST_QUERY);
Assert.assertNotNull(resultSet);
Assert.assertEquals(statement.getConnection(), connection);
}
+
+ @Test
+ public void testSetEnableNullHandling()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.ENABLE_NULL_HANDLING, "true");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ Statement statement = pinotConnection.createStatement();
+ Assert.assertNotNull(statement);
+ statement.executeQuery(BASIC_TEST_QUERY);
+ String expectedSql =
+
DriverUtils.createSetQueryOptionString(QueryOptionKey.ENABLE_NULL_HANDLING,
true) + BASIC_TEST_QUERY;
+ Assert.assertEquals(_dummyPinotClientTransport.getLastQuery().substring(0,
expectedSql.length()), expectedSql);
+ }
+
+ @Test
+ public void testSetDisableNullHandling()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.ENABLE_NULL_HANDLING, "false");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ Statement statement = pinotConnection.createStatement();
+ Assert.assertNotNull(statement);
+ statement.executeQuery(BASIC_TEST_QUERY);
+ String expectedSql =
+
DriverUtils.createSetQueryOptionString(QueryOptionKey.ENABLE_NULL_HANDLING,
false) + BASIC_TEST_QUERY;
+ Assert.assertEquals(_dummyPinotClientTransport.getLastQuery().substring(0,
expectedSql.length()), expectedSql);
+ }
+
+ @Test
+ public void testPresetEnableNullHandling()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.ENABLE_NULL_HANDLING, "true");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ Statement statement = pinotConnection.createStatement();
+ Assert.assertNotNull(statement);
+ String presetSql =
+
DriverUtils.createSetQueryOptionString(QueryOptionKey.ENABLE_NULL_HANDLING,
true) + BASIC_TEST_QUERY;
+ statement.executeQuery(presetSql);
+ Assert.assertEquals(_dummyPinotClientTransport.getLastQuery().substring(0,
presetSql.length()), presetSql);
+ }
+
+ @Test
+ public void testSetUseMultistageEngine()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.USE_MULTISTAGE_ENGINE, "true");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ Statement statement = pinotConnection.createStatement();
+ Assert.assertNotNull(statement);
+ statement.executeQuery(BASIC_TEST_QUERY);
+ String expectedSql =
+
DriverUtils.createSetQueryOptionString(QueryOptionKey.USE_MULTISTAGE_ENGINE,
true) + BASIC_TEST_QUERY;
+ Assert.assertEquals(_dummyPinotClientTransport.getLastQuery().substring(0,
expectedSql.length()), expectedSql);
+ }
+
+ @Test
+ public void testSetMultipleQueryOptions()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.ENABLE_NULL_HANDLING, "true");
+ props.put(QueryOptionKey.USE_MULTISTAGE_ENGINE, "true");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ Statement statement = pinotConnection.createStatement();
+ Assert.assertNotNull(statement);
+ statement.executeQuery(BASIC_TEST_QUERY);
+ String resultingQuery = _dummyPinotClientTransport.getLastQuery();
+ Assert.assertTrue(
+
resultingQuery.contains(DriverUtils.createSetQueryOptionString(QueryOptionKey.ENABLE_NULL_HANDLING,
true)));
+ Assert.assertTrue(
+
resultingQuery.contains(DriverUtils.createSetQueryOptionString(QueryOptionKey.USE_MULTISTAGE_ENGINE,
true)));
+ }
+
+ @Test
+ public void testSetOptionAsInteger()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.USE_MULTISTAGE_ENGINE, "2");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ Statement statement = pinotConnection.createStatement();
+ Assert.assertNotNull(statement);
+ statement.executeQuery(BASIC_TEST_QUERY);
+ String expectedSql =
+
DriverUtils.createSetQueryOptionString(QueryOptionKey.USE_MULTISTAGE_ENGINE, 2)
+ BASIC_TEST_QUERY;
+ Assert.assertEquals(_dummyPinotClientTransport.getLastQuery().substring(0,
expectedSql.length()), expectedSql);
+ }
+
+ @Test
+ public void testSetOptionAsFloat()
+ throws Exception {
+ Properties props = new Properties();
+ props.put(QueryOptionKey.USE_MULTISTAGE_ENGINE, "2.5");
+ PinotConnection pinotConnection =
+ new PinotConnection(props, "dummy", _dummyPinotClientTransport,
"dummy", _dummyPinotControllerTransport);
+ Statement statement = pinotConnection.createStatement();
+ Assert.assertNotNull(statement);
+ statement.executeQuery(BASIC_TEST_QUERY);
+ String expectedSql =
+
DriverUtils.createSetQueryOptionString(QueryOptionKey.USE_MULTISTAGE_ENGINE,
2.5) + BASIC_TEST_QUERY;
+ Assert.assertEquals(_dummyPinotClientTransport.getLastQuery().substring(0,
expectedSql.length()), expectedSql);
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
