SQLHandler.java example

Explorer
lucene-solr-master
- lucene
- solr
/*
 * 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.solr.handler;

import java.io.IOException;
import java.lang.invoke.MethodHandles;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Properties;

import org.apache.calcite.config.Lex;
import org.apache.solr.client.solrj.io.Tuple;
import org.apache.solr.client.solrj.io.comp.StreamComparator;
import org.apache.solr.client.solrj.io.stream.ExceptionStream;
import org.apache.solr.client.solrj.io.stream.TupleStream;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.core.CoreContainer;
import org.apache.solr.core.SolrCore;
import org.apache.solr.handler.sql.CalciteSolrDriver;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.security.AuthorizationContext;
import org.apache.solr.security.PermissionNameProvider;
import org.apache.solr.util.plugin.SolrCoreAware;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SQLHandler extends RequestHandlerBase implements SolrCoreAware, PermissionNameProvider {

  private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

  private static String defaultZkhost = null;
  private static String defaultWorkerCollection = null;

  static final String sqlNonCloudErrorMsg = "/sql handler only works in Solr Cloud mode";

  private boolean isCloud = false;

  public void inform(SolrCore core) {
    CoreContainer coreContainer = core.getCoreContainer();

    if(coreContainer.isZooKeeperAware()) {
      defaultZkhost = core.getCoreContainer().getZkController().getZkServerAddress();
      defaultWorkerCollection = core.getCoreDescriptor().getCollectionName();
      isCloud = true;
    }
  }

  @Override
  public PermissionNameProvider.Name getPermissionName(AuthorizationContext request) {
    return PermissionNameProvider.Name.READ_PERM;
  }

  public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throws Exception {
    ModifiableSolrParams params = new ModifiableSolrParams(req.getParams());
    params = adjustParams(params);
    req.setParams(params);

    String sql = params.get("stmt");
    // Set defaults for parameters
    params.set("numWorkers", params.getInt("numWorkers", 1));
    params.set("workerCollection", params.get("workerCollection", defaultWorkerCollection));
    params.set("workerZkhost", params.get("workerZkhost", defaultZkhost));
    params.set("aggregationMode", params.get("aggregationMode", "facet"));

    TupleStream tupleStream = null;
    try {

      if(!isCloud) {
        throw new IllegalStateException(sqlNonCloudErrorMsg);
      }

      if(sql == null) {
        throw new Exception("stmt parameter cannot be null");
      }

      String url = CalciteSolrDriver.CONNECT_STRING_PREFIX;

      Properties properties = new Properties();
      // Add all query parameters
      Iterator<String> parameterNamesIterator = params.getParameterNamesIterator();
      while(parameterNamesIterator.hasNext()) {
        String param = parameterNamesIterator.next();
        properties.setProperty(param, params.get(param));
      }

      // Set these last to ensure that they are set properly
      properties.setProperty("lex", Lex.MYSQL.toString());
      properties.setProperty("zk", defaultZkhost);

      String driverClass = CalciteSolrDriver.class.getCanonicalName();

      // JDBC driver requires metadata from the SQLHandler. Default to false since this adds a new Metadata stream.
      boolean includeMetadata = params.getBool("includeMetadata", false);
      tupleStream = new SqlHandlerStream(url, sql, null, properties, driverClass, includeMetadata);

      tupleStream = new StreamHandler.TimerStream(new ExceptionStream(tupleStream));

      rsp.add("result-set", tupleStream);
    } catch(Exception e) {
      //Catch the SQL parsing and query transformation exceptions.
      if(tupleStream != null) {
        tupleStream.close();
      }
      SolrException.log(logger, e);
      rsp.add("result-set", new StreamHandler.DummyErrorStream(e));
    }
  }

  public String getDescription() {
    return "SQLHandler";
  }

  public String getSource() {
    return null;
  }

  /*
   * Only necessary for SolrJ JDBC driver since metadata has to be passed back
   */
  private static class SqlHandlerStream extends CalciteJDBCStream {
    private final boolean includeMetadata;
    private boolean firstTuple = true;
    List<String> metadataFields = new ArrayList<>();
    Map<String, String> metadataAliases = new HashMap<>();

    SqlHandlerStream(String connectionUrl, String sqlQuery, StreamComparator definedSort,
                     Properties connectionProperties, String driverClassName, boolean includeMetadata)
        throws IOException {
      super(connectionUrl, sqlQuery, definedSort, connectionProperties, driverClassName);

      this.includeMetadata = includeMetadata;
    }

    @Override
    public Tuple read() throws IOException {
      // Return a metadata tuple as the first tuple and then pass through to the JDBCStream.
      if(firstTuple) {
        try {
          Map<String, Object> fields = new HashMap<>();

          firstTuple = false;

          ResultSetMetaData resultSetMetaData = resultSet.getMetaData();

          for(int i = 1; i <= resultSetMetaData.getColumnCount(); i++) {
            String columnName = resultSetMetaData.getColumnName(i);
            String columnLabel = resultSetMetaData.getColumnLabel(i);
            metadataFields.add(columnName);
            metadataAliases.put(columnName, columnLabel);
          }

          if(includeMetadata) {
            fields.put("isMetadata", true);
            fields.put("fields", metadataFields);
            fields.put("aliases", metadataAliases);
            return new Tuple(fields);
          }
        } catch (SQLException e) {
          throw new IOException(e);
        }
      }

      Tuple tuple = super.read();
      if(!tuple.EOF) {
        tuple.fieldNames = metadataFields;
        tuple.fieldLabels = metadataAliases;
      }
      return tuple;
    }
  }

  private ModifiableSolrParams adjustParams(SolrParams params) {
    ModifiableSolrParams adjustedParams = new ModifiableSolrParams();
    adjustedParams.add(params);
    adjustedParams.add(CommonParams.OMIT_HEADER, "true");
    return adjustedParams;
  }
}