org.apache.calcite.adapter.elasticsearch.ElasticsearchSchema Maven / Gradle / Ivy

Go to download
/*
 * 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.calcite.adapter.elasticsearch;

import org.apache.calcite.schema.Table;
import org.apache.calcite.schema.impl.AbstractSchema;

import com.carrotsearch.hppc.cursors.ObjectObjectCursor;

import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;

import org.elasticsearch.action.admin.indices.get.GetIndexRequest;
import org.elasticsearch.action.admin.indices.mapping.get.GetMappingsRequest;
import org.elasticsearch.action.admin.indices.mapping.get.GetMappingsResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.cluster.metadata.MappingMetaData;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.common.collect.ImmutableOpenMap;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.InetSocketTransportAddress;
import org.elasticsearch.common.transport.TransportAddress;

import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * Schema mapped onto an index of ELASTICSEARCH types.
 *
 * Each table in the schema is an ELASTICSEARCH type in that index.
 */
public class ElasticsearchSchema extends AbstractSchema {
  final String index;

  private transient Client client;

  /**
   * Creates an Elasticsearch schema.
   *
   * @param coordinates Map of Elasticsearch node locations (host, port)
   * @param userConfig Map of user-specified configurations
   * @param indexName Elasticsearch database name, e.g. "usa".
   */
  ElasticsearchSchema(Map coordinates,
      Map userConfig, String indexName) {
    super();

    final List transportAddresses = new ArrayList<>();
    for (Map.Entry coordinate: coordinates.entrySet()) {
      transportAddresses.add(new InetSocketAddress(coordinate.getKey(), coordinate.getValue()));
    }

    open(transportAddresses, userConfig);

    if (client != null) {
      final String[] indices = client.admin().indices()
          .getIndex(new GetIndexRequest().indices(indexName))
          .actionGet().getIndices();
      if (indices.length == 1) {
        index = indices[0];
      } else {
        index = null;
      }
    } else {
      index = null;
    }
  }

  @Override protected Map getTableMap() {
    final ImmutableMap.Builder builder = ImmutableMap.builder();

    try {
      GetMappingsResponse response = client.admin().indices().getMappings(
        new GetMappingsRequest().indices(index)).get();
      ImmutableOpenMap mapping = response.getMappings().get(index);
      for (ObjectObjectCursor c: mapping) {
        builder.put(c.key, new ElasticsearchTable(client, index, c.key));
      }
    } catch (Exception e) {
      throw Throwables.propagate(e);
    }
    return builder.build();
  }

  private void open(List transportAddresses, Map userConfig) {
    final List transportNodes = new ArrayList<>(transportAddresses.size());
    for (InetSocketAddress address : transportAddresses) {
      transportNodes.add(new InetSocketTransportAddress(address));
    }

    Settings settings = Settings.settingsBuilder().put(userConfig).build();

    final TransportClient transportClient = TransportClient.builder().settings(settings).build();
    for (TransportAddress transport : transportNodes) {
      transportClient.addTransportAddress(transport);
    }

    final List nodes = ImmutableList.copyOf(transportClient.connectedNodes());
    if (nodes.isEmpty()) {
      throw new RuntimeException("Cannot connect to any elasticsearch nodes");
    }

    client = transportClient;
  }
}

// End ElasticsearchSchema.java