SOLR-8519: Implement ResultSetMetaDataImpl.getColumnCount()

solr/core/src/java/org/apache/solr/handler/SQLHandler.java

@@ -25,7 +25,6 @@ import java.util.Iterator;
 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 class SQLHandler extends RequestHandlerBase implements SolrCoreAware {
     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 class SQLHandler extends RequestHandlerBase implements SolrCoreAware {
                                     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 class SQLHandler extends RequestHandlerBase implements SolrCoreAware {
         sqlStream = doSelect(sqlVistor);
       }
 
+      if(includeMetadata) {
+        sqlStream = new MetadataStream(sqlStream, sqlVistor);
+      }
+
       return sqlStream;
     }
   }
@@ -1337,6 +1347,53 @@ public class SQLHandler extends RequestHandlerBase implements SolrCoreAware {
     }
   }
 
+  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<TupleStream> 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<Boolean, Tuple> {
 
     private Map<String,String> reverseAliasMap;

solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/DriverImpl.java

@@ -67,6 +67,9 @@ public class DriverImpl implements Driver {
       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);

solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetImpl.java

@@ -17,6 +17,7 @@ package org.apache.solr.client.solrj.io.sql;
  * limitations under the License.
  */
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.Reader;
 import java.math.BigDecimal;
@@ -45,6 +46,8 @@ import org.apache.solr.client.solrj.io.stream.SolrStream;
 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 @@ class ResultSetImpl implements ResultSet {
       } else {
         return true;
       }
-    } catch (Exception e) {
+    } catch (IOException e) {
       throw new SQLException(e);
     }
   }
@@ -279,7 +300,7 @@ class ResultSetImpl implements ResultSet {
 
   @Override
   public ResultSetMetaData getMetaData() throws SQLException {
-    return new ResultSetMetaDataImpl(this);
+    return this.resultSetMetaData;
   }
 
   @Override

solr/solrj/src/java/org/apache/solr/client/solrj/io/sql/ResultSetMetaDataImpl.java

@@ -19,17 +19,26 @@ package org.apache.solr.client.solrj.io.sql;
 
 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<String> fields = metadataTuple.getStrings("fields");
+    if(fields == null) {
+      throw new SQLException("Unable to determine fields for column count");
+    }
+    return fields.size();
   }
 
   @Override

solr/solrj/src/test/org/apache/solr/client/solrj/io/sql/JdbcTest.java

@@ -355,6 +355,8 @@ public class JdbcTest extends AbstractFullDistribZkTestBase {
     ResultSetMetaData resultSetMetaData = rs.getMetaData();
 
     assertNotNull(resultSetMetaData);
+
+    assertEquals(4, resultSetMetaData.getColumnCount());
   }
 
   private void checkResultSet(ResultSet rs) throws Exception {