From d3345ec0ab5d39534fc51e6af1d50a2abbc94717 Mon Sep 17 00:00:00 2001 From: Kevin Risden Date: Sun, 24 Jan 2016 11:44:30 -0600 Subject: [PATCH] Implement ResultSetMetaDataImpl.getColumnCount() --- .../org/apache/solr/handler/SQLHandler.java | 63 ++++++++++++++++++- .../solr/client/solrj/io/sql/DriverImpl.java | 3 + .../client/solrj/io/sql/ResultSetImpl.java | 25 +++++++- .../solrj/io/sql/ResultSetMetaDataImpl.java | 11 +++- .../solr/client/solrj/io/sql/JdbcTest.java | 2 + 5 files changed, 98 insertions(+), 6 deletions(-) diff --git a/solr/core/src/java/org/apache/solr/handler/SQLHandler.java b/solr/core/src/java/org/apache/solr/handler/SQLHandler.java index a27206f82069..fa1b3921bb8b 100644 --- a/solr/core/src/java/org/apache/solr/handler/SQLHandler.java +++ b/solr/core/src/java/org/apache/solr/handler/SQLHandler.java @@ -25,7 +25,6 @@ import java.util.Locale; import java.util.Set; -import com.facebook.presto.sql.ExpressionFormatter; import com.facebook.presto.sql.tree.*; import com.google.common.base.Strings; import com.google.common.collect.Iterables; @@ -103,16 +102,22 @@ public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throw String workerZkhost = params.get("workerZkhost",defaultZkhost); String mode = params.get("aggregationMode", "map_reduce"); StreamContext context = new StreamContext(); + + // JDBC driver requires metadata from the SQLHandler. Default to false since this adds a new Metadata stream. + boolean includeMetadata = params.getBool("includeMetadata", false); + try { if(sql == null) { throw new Exception("sql parameter cannot be null"); } - TupleStream tupleStream = SQLTupleStreamParser.parse(sql, numWorkers, workerCollection, workerZkhost, AggregationMode.getMode(mode)); + TupleStream tupleStream = SQLTupleStreamParser.parse(sql, numWorkers, workerCollection, workerZkhost, + AggregationMode.getMode(mode), includeMetadata); context.numWorkers = numWorkers; context.setSolrClientCache(StreamHandler.clientCache); tupleStream.setStreamContext(context); + rsp.add("result-set", new StreamHandler.TimerStream(new ExceptionStream(tupleStream))); } catch(Exception e) { //Catch the SQL parsing and query transformation exceptions. @@ -142,7 +147,8 @@ public static TupleStream parse(String sql, int numWorkers, String workerCollection, String workerZkhost, - AggregationMode aggregationMode) throws IOException { + AggregationMode aggregationMode, + boolean includeMetadata) throws IOException { SqlParser parser = new SqlParser(); Statement statement = parser.createStatement(sql); @@ -169,6 +175,10 @@ public static TupleStream parse(String sql, sqlStream = doSelect(sqlVistor); } + if(includeMetadata) { + sqlStream = new MetadataStream(sqlStream, sqlVistor); + } + return sqlStream; } } @@ -1337,6 +1347,53 @@ public Tuple read() throws IOException { } } + private static class MetadataStream extends TupleStream { + + private final TupleStream stream; + private final SQLVisitor sqlVisitor; + private boolean firstTuple = true; + + public MetadataStream(TupleStream stream, SQLVisitor sqlVistor) { + this.stream = stream; + this.sqlVisitor = sqlVistor; + } + + public List children() { + return this.stream.children(); + } + + public void open() throws IOException { + this.stream.open(); + } + + // Return a metadata tuple as the first tuple and then pass through to the underlying stream. + public Tuple read() throws IOException { + if(firstTuple) { + firstTuple = false; + + Map fields = new HashMap<>(); + fields.put("isMetadata", true); + fields.put("fields", sqlVisitor.fields); + fields.put("aliases", sqlVisitor.columnAliases); + return new Tuple(fields); + } + + return this.stream.read(); + } + + public StreamComparator getStreamSort() { + return this.stream.getStreamSort(); + } + + public void close() throws IOException { + this.stream.close(); + } + + public void setStreamContext(StreamContext context) { + this.stream.setStreamContext(context); + } + } + private static class HavingVisitor extends AstVisitor { private Map reverseAliasMap; diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/DriverImpl.java b/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/DriverImpl.java index d850f69bcada..57d60ca01608 100644 --- a/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/DriverImpl.java +++ b/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/DriverImpl.java @@ -67,6 +67,9 @@ public Connection connect(String url, Properties props) throws SQLException { props.setProperty("aggregationMode", "facet"); } + // JDBC requires metadata like field names from the SQLHandler. Force this property to be true. + props.setProperty("includeMetadata", "true"); + String zkHost = uri.getAuthority() + uri.getPath(); return new ConnectionImpl(url, zkHost, collection, props); diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetImpl.java b/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetImpl.java index 1f55eed5bbcd..f4f73324a981 100644 --- a/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetImpl.java +++ b/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetImpl.java @@ -17,6 +17,7 @@ * limitations under the License. */ +import java.io.IOException; import java.io.InputStream; import java.io.Reader; import java.math.BigDecimal; @@ -45,6 +46,8 @@ class ResultSetImpl implements ResultSet { private final StatementImpl statement; private final SolrStream solrStream; + private final ResultSetMetaData resultSetMetaData; + private final Tuple metadataTuple; private Tuple tuple; private boolean done; private boolean closed; @@ -53,6 +56,24 @@ class ResultSetImpl implements ResultSet { ResultSetImpl(StatementImpl statement) { this.statement = statement; this.solrStream = statement.getSolrStream(); + + // Read the first tuple so that metadata can be gathered + try { + this.metadataTuple = this.solrStream.read(); + + Object isMetadata = this.metadataTuple.get("isMetadata"); + if(isMetadata == null || !isMetadata.equals(true)) { + throw new RuntimeException("First tuple is not a metadata tuple"); + } + } catch (IOException e) { + throw new RuntimeException("Couldn't get metadata tuple"); + } + + this.resultSetMetaData = new ResultSetMetaDataImpl(this); + } + + Tuple getMetadataTuple() { + return this.metadataTuple; } @Override @@ -69,7 +90,7 @@ public boolean next() throws SQLException { } else { return true; } - } catch (Exception e) { + } catch (IOException e) { throw new SQLException(e); } } @@ -279,7 +300,7 @@ public String getCursorName() throws SQLException { @Override public ResultSetMetaData getMetaData() throws SQLException { - return new ResultSetMetaDataImpl(this); + return this.resultSetMetaData; } @Override diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetMetaDataImpl.java b/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetMetaDataImpl.java index 6d615d4cf855..6284b1c9154e 100644 --- a/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetMetaDataImpl.java +++ b/solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetMetaDataImpl.java @@ -19,17 +19,26 @@ import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.util.List; + +import org.apache.solr.client.solrj.io.Tuple; class ResultSetMetaDataImpl implements ResultSetMetaData { private final ResultSetImpl resultSet; + private final Tuple metadataTuple; ResultSetMetaDataImpl(ResultSetImpl resultSet) { this.resultSet = resultSet; + this.metadataTuple = this.resultSet.getMetadataTuple(); } @Override public int getColumnCount() throws SQLException { - return 0; + List fields = metadataTuple.getStrings("fields"); + if(fields == null) { + throw new SQLException("Unable to determine fields for column count"); + } + return fields.size(); } @Override diff --git a/solr/solrj/src/test/org/apache/solr/client/solrj/io/sql/JdbcTest.java b/solr/solrj/src/test/org/apache/solr/client/solrj/io/sql/JdbcTest.java index b1ce84f381a2..644e2b384dc3 100644 --- a/solr/solrj/src/test/org/apache/solr/client/solrj/io/sql/JdbcTest.java +++ b/solr/solrj/src/test/org/apache/solr/client/solrj/io/sql/JdbcTest.java @@ -355,6 +355,8 @@ private void checkResultSetMetadata(ResultSet rs) throws Exception { ResultSetMetaData resultSetMetaData = rs.getMetaData(); assertNotNull(resultSetMetaData); + + assertEquals(4, resultSetMetaData.getColumnCount()); } private void checkResultSet(ResultSet rs) throws Exception {