com.qubole.quark.sql.QueryContext Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of quark-optimizer Show documentation
The newest version!
/*
 * 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,
 * 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 com.qubole.quark.sql;

import org.apache.calcite.DataContext;
import org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.calcite.config.CalciteConnectionConfig;
import org.apache.calcite.config.CalciteConnectionConfigImpl;
import org.apache.calcite.config.CalciteConnectionProperty;
import org.apache.calcite.jdbc.CalcitePrepare;
import org.apache.calcite.jdbc.CalciteSchema;
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
import org.apache.calcite.rel.type.RelDataTypeSystem;
import org.apache.calcite.schema.SchemaPlus;

import com.fasterxml.jackson.databind.ObjectMapper;


import com.google.common.collect.ImmutableList;

import com.qubole.quark.QuarkException;

import com.qubole.quark.planner.DataSourceSchema;
import com.qubole.quark.planner.MetadataSchema;
import com.qubole.quark.planner.QuarkFactory;
import com.qubole.quark.planner.QuarkFactoryResult;
import com.qubole.quark.planner.QuarkSchema;
import com.qubole.quark.planner.TestFactory;
import com.qubole.quark.plugin.DataSource;

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

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;

/**
 * Created by amargoor on 10/26/15.
 */
public class QueryContext {

  protected static final Logger LOG = LoggerFactory.getLogger(QueryContext.class);
  private final SchemaPlus rootSchema = CalciteSchema.createRootSchema(true).plus();
  private DataSourceSchema defaultDataSource;
  private boolean unitTestMode;
  private final List defaultSchema;

  public CalciteConnectionConfig getCfg() {
    return cfg;
  }

  private final CalciteConnectionConfig cfg;

  public JavaTypeFactory getTypeFactory() {
    return typeFactory;
  }

  private final JavaTypeFactory typeFactory;

  private List parseDefaultSchema(Properties info) throws IOException, QuarkException {
    if (info.getProperty("defaultSchema") != null) {
      final ObjectMapper mapper = new ObjectMapper();
      return Arrays.asList(mapper.readValue(info.getProperty("defaultSchema"), String[].class));
    } else if (defaultDataSource != null) {
      DataSource dataSource = defaultDataSource.getDataSource();
      return ImmutableList.of(defaultDataSource.getName(), dataSource.getDefaultSchema());
    } else {
      return ImmutableList.of(MetadataSchema.NAME);
    }
  }

  public QueryContext(Properties info) throws QuarkException {
    try {
      Class schemaFactoryClazz = Class.forName(info.getProperty("schemaFactory"));
      if (!info.contains(CalciteConnectionProperty.FUN.camelName())) {
        info.put(CalciteConnectionProperty.FUN.camelName(), "standard,hive");
      }
      this.cfg = new CalciteConnectionConfigImpl(info);
      final RelDataTypeSystem typeSystem =
          cfg.typeSystem(RelDataTypeSystem.class, RelDataTypeSystem.DEFAULT);
      this.typeFactory = new JavaTypeFactoryImpl(typeSystem);
      this.unitTestMode = Boolean.parseBoolean(info.getProperty("unitTestMode", "false"));

      Object obj = schemaFactoryClazz.newInstance();
      if (obj instanceof QuarkFactory) {
        final QuarkFactory schemaFactory = (QuarkFactory) schemaFactoryClazz.newInstance();
        QuarkFactoryResult factoryResult = schemaFactory.create(info);

        this.defaultDataSource = factoryResult.defaultSchema;
        defaultSchema = parseDefaultSchema(info);

        for (DataSourceSchema schema : factoryResult.dataSourceSchemas) {
          SchemaPlus schemaPlus = rootSchema.add(schema.getName(), schema);
          schema.initialize(this, schemaPlus);
        }

        SchemaPlus metadataPlus = rootSchema.add(factoryResult.metadataSchema.getName(),
            factoryResult.metadataSchema);
        factoryResult.metadataSchema.initialize(this, metadataPlus);
      } else {
        final TestFactory schemaFactory = (TestFactory) schemaFactoryClazz.newInstance();
        List schemas = schemaFactory.create(info);
        if (info.getProperty("defaultSchema") == null) {
          throw new QuarkException(new Throwable("Default schema has to be specified"));
        }
        final ObjectMapper mapper = new ObjectMapper();
        defaultSchema =
              Arrays.asList(mapper.readValue(info.getProperty("defaultSchema"), String[].class));
        for (QuarkSchema schema : schemas) {
          SchemaPlus schemaPlus = rootSchema.add(schema.getName(), schema);
          schema.initialize(this, schemaPlus);
        }
      }
    } catch (ClassNotFoundException | InstantiationException | IllegalAccessException
        | IOException e) {
      throw new QuarkException(e);
    }
  }

  public SchemaPlus getRootSchema() {
    return rootSchema;
  }

  public List getDefaultSchemaPath() {
    return defaultSchema;
  }

  public SchemaPlus getDefaultSchema() {
    SchemaPlus defaultSchemaPlus = this.rootSchema;
    for (String schemaName : this.defaultSchema) {
      defaultSchemaPlus = defaultSchemaPlus.getSubSchema(schemaName);
    }
    return defaultSchemaPlus;
  }

  public CalcitePrepare.Context getPrepareContext() {
    return new CalcitePrepare.Context() {

      @Override
      public JavaTypeFactory getTypeFactory() {
        return typeFactory;
      }

      @Override
      public CalciteSchema getRootSchema() {
        return CalciteSchema.from(rootSchema);
      }

      @Override
      public List getDefaultSchemaPath() {
        return defaultSchema;
      }

      @Override
      public CalciteConnectionConfig config() {
        return cfg;
      }

      @Override
      public CalcitePrepare.SparkHandler spark() {
        return CalcitePrepare.Dummy.getSparkHandler(false);
      }

      @Override
      public DataContext getDataContext() {
        return null;
      }
    };
  }

  public DataSourceSchema getDefaultDataSource() {
    return defaultDataSource;
  }

  public boolean isUnitTestMode() {
    return unitTestMode;
  }
}