* Copyright (c) 2015. Qubole Inc
 * Licensed 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,
 * See the License for the specific language governing permissions and
 *    limitations under the License.

package com.qubole.quark.plugins.qubole;

import org.apache.calcite.schema.Schema;

import org.apache.commons.lang.Validate;

import com.google.common.collect.ImmutableMap;

import com.qubole.qds.sdk.java.client.DefaultQdsConfiguration;
import com.qubole.qds.sdk.java.client.QdsClient;
import com.qubole.qds.sdk.java.client.QdsClientFactory;
import com.qubole.qds.sdk.java.client.QdsConfiguration;
import com.qubole.qds.sdk.java.entities.ResultValue;
import com.qubole.qds.sdk.java.entities.SchemaOrdinal;

import com.qubole.quark.QuarkException;
import com.qubole.quark.plugins.Executor;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;

 * Created by dev on 11/13/15.
public abstract class QuboleDB implements Executor {
  private static final Logger LOG = LoggerFactory.getLogger(QuboleDB.class.getName());
  private static final String ROW_DELIMETER = "\r\n";
  private static final String COLUMN_DELIMETER = "\t";
  protected String token;
  protected String endpoint;
  QuboleDB(Map<String, Object> properties) {
    this.endpoint = (String) properties.get("endpoint");
    this.token = (String) properties.get("token");

  protected abstract Map<String, List<SchemaOrdinal>>
                        getSchemaListDescribed() throws ExecutionException, InterruptedException;

  protected abstract ImmutableMap<String, Integer> getDataTypes();

  public abstract boolean isCaseSensitive();

  private void validate(Map<String, Object> properties) {
        "Field \"endpoint\" specifying Qubole's endpoint needs "
            + "to be defined for Qubole Data Source in JSON");
        "Field \"token\" specifying Authentication token needs "
            + "to be defined for Qubole Data Source in JSON");

  protected QdsClient getQdsClient() {
    QdsConfiguration configuration = new DefaultQdsConfiguration(endpoint, token);
    return QdsClientFactory.newClient(configuration);

  public void cleanup() throws Exception {}

  public ImmutableMap<String, Schema> getSchemas() throws QuarkException {
    Map<String, List<SchemaOrdinal>> schemas;
    try {
      schemas = getSchemaListDescribed();
    } catch (ExecutionException | InterruptedException e) {
      LOG.error("Getting Schema metadata for " + this.endpoint
          + " failed. Error: " + e.getMessage(), e);
      throw new QuarkException(e);
    ImmutableMap.Builder<String, Schema> schemaBuilder = new ImmutableMap.Builder<>();

    for (String schemaName: schemas.keySet()) {
      String schemaKey = schemaName;
      if (!this.isCaseSensitive()) {
        schemaKey = schemaName.toUpperCase();
      schemaBuilder.put(schemaKey, new QuboleSchema(schemaKey,
          schemas.get(schemaName), this.isCaseSensitive(), this.getDataTypes()));
    return schemaBuilder.build();

  protected Iterator<Object> convertToIterator(ResultValue resultValue) {
    String[] result = resultValue.getResults().split(ROW_DELIMETER);
    List<Object> resultSet = new ArrayList<>();

    for (int i = 0; i < result.length; i++) {
      //check if Empty Result Set
      // QDS returns row of empty column
      if (result[i].length() == 0) {
        resultSet = new ArrayList<>();
      String[] row = result[i].split(COLUMN_DELIMETER);
    return resultSet.iterator();

   * Merges two schema Maps. 'putAll' method can't be used because if the any of the
   * key is same in the maps, then it will get overwritten.
   * for e.g.
   * if map1 = {'public' -&gt; list1}, map2 = {'public' -&gt; list2}
   * then
   *   map1.putAll(map2) will return map2 ({'public' -&gt; list2})
   *   i.e. map1 gets overwritten
   * @param schemas1 schema map1
   * @param schemas2 schema map2
   * @return returns merged Map
   * */
  protected Map<String, List<SchemaOrdinal>> mergeSchemas(
      Map<String, List<SchemaOrdinal>> schemas1,
      Map<String, List<SchemaOrdinal>> schemas2) {
    for (String schemaName : schemas2.keySet()) {
      if (schemas1.containsKey(schemaName)) {
      } else {
        schemas1.put(schemaName, schemas2.get(schemaName));
    return schemas1;