org.apache.drill.exec.store.jdbc.JdbcStoragePlugin.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.drill.exec.store.jdbc.JdbcStoragePlugin.java

Source

/**
 * 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.drill.exec.store.jdbc;

import java.io.IOException;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.sql.DataSource;

import org.apache.calcite.adapter.jdbc.JdbcConvention;
import org.apache.calcite.adapter.jdbc.JdbcRules;
import org.apache.calcite.adapter.jdbc.JdbcRules.JdbcJoin;
import org.apache.calcite.adapter.jdbc.JdbcSchema;
import org.apache.calcite.linq4j.tree.ConstantUntypedNull;
import org.apache.calcite.plan.RelOptPlanner;
import org.apache.calcite.plan.RelOptRule;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.convert.ConverterRule;
import org.apache.calcite.rel.rules.FilterSetOpTransposeRule;
import org.apache.calcite.rel.rules.ProjectRemoveRule;
import org.apache.calcite.rex.RexCall;
import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.schema.Function;
import org.apache.calcite.schema.Schema;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.schema.Table;
import org.apache.calcite.sql.SqlDialect;
import org.apache.commons.dbcp.BasicDataSource;
import org.apache.drill.common.JSONOptions;
import org.apache.drill.common.expression.SchemaPath;
import org.apache.drill.exec.ops.OptimizerRulesContext;
import org.apache.drill.exec.physical.base.AbstractGroupScan;
import org.apache.drill.exec.planner.logical.DrillRel;
import org.apache.drill.exec.planner.physical.Prel;
import org.apache.drill.exec.server.DrillbitContext;
import org.apache.drill.exec.store.AbstractSchema;
import org.apache.drill.exec.store.AbstractStoragePlugin;
import org.apache.drill.exec.store.SchemaConfig;
import org.apache.drill.exec.store.jdbc.DrillJdbcRuleBase.DrillJdbcFilterRule;
import org.apache.drill.exec.store.jdbc.DrillJdbcRuleBase.DrillJdbcProjectRule;

import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;

public class JdbcStoragePlugin extends AbstractStoragePlugin {
    private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(JdbcStoragePlugin.class);

    // Rules from Calcite's JdbcRules class that we want to avoid using.
    private static String[] RULES_TO_AVOID = { "JdbcToEnumerableConverterRule", "JdbcFilterRule",
            "JdbcProjectRule" };

    private final JdbcStorageConfig config;
    private final DrillbitContext context;
    private final DataSource source;
    private final String name;
    private final SqlDialect dialect;
    private final DrillJdbcConvention convention;

    public JdbcStoragePlugin(JdbcStorageConfig config, DrillbitContext context, String name) {
        this.context = context;
        this.config = config;
        this.name = name;
        BasicDataSource source = new BasicDataSource();
        source.setDriverClassName(config.getDriver());
        source.setUrl(config.getUrl());

        if (config.getUsername() != null) {
            source.setUsername(config.getUsername());
        }

        if (config.getPassword() != null) {
            source.setPassword(config.getPassword());
        }

        this.source = source;
        this.dialect = JdbcSchema.createDialect(source);
        this.convention = new DrillJdbcConvention(dialect, name);
    }

    class DrillJdbcConvention extends JdbcConvention {

        private final ImmutableSet<RelOptRule> rules;

        public DrillJdbcConvention(SqlDialect dialect, String name) {
            super(dialect, ConstantUntypedNull.INSTANCE, name);

            // build rules for this convention.
            ImmutableSet.Builder<RelOptRule> builder = ImmutableSet.builder();

            builder.add(JDBC_PRULE_INSTANCE);
            builder.add(new JdbcDrelConverterRule(this));
            builder.add(new DrillJdbcProjectRule(this));
            builder.add(new DrillJdbcFilterRule(this));

            outside: for (RelOptRule rule : JdbcRules.rules(this)) {
                final String description = rule.toString();

                // we want to black list some rules but the parent Calcite package is all or none.
                // Therefore, we remove rules with names we don't like.
                for (String black : RULES_TO_AVOID) {
                    if (description.equals(black)) {
                        continue outside;
                    }

                }

                builder.add(rule);
            }

            builder.add(FilterSetOpTransposeRule.INSTANCE);
            builder.add(ProjectRemoveRule.INSTANCE);

            rules = builder.build();
        }

        @Override
        public void register(RelOptPlanner planner) {
            for (RelOptRule rule : rules) {
                planner.addRule(rule);
            }
        }

        public Set<RelOptRule> getRules() {
            return rules;
        }

        public JdbcStoragePlugin getPlugin() {
            return JdbcStoragePlugin.this;
        }
    }

    /**
     * Returns whether a condition is supported by {@link JdbcJoin}.
     *
     * <p>Corresponds to the capabilities of
     * {@link JdbcJoin#convertConditionToSqlNode}.
     *
     * @param node Condition
     * @return Whether condition is supported
     */
    private static boolean canJoinOnCondition(RexNode node) {
        final List<RexNode> operands;
        switch (node.getKind()) {
        case AND:
        case OR:
            operands = ((RexCall) node).getOperands();
            for (RexNode operand : operands) {
                if (!canJoinOnCondition(operand)) {
                    return false;
                }
            }
            return true;

        case EQUALS:
        case IS_NOT_DISTINCT_FROM:
        case NOT_EQUALS:
        case GREATER_THAN:
        case GREATER_THAN_OR_EQUAL:
        case LESS_THAN:
        case LESS_THAN_OR_EQUAL:
            operands = ((RexCall) node).getOperands();
            if ((operands.get(0) instanceof RexInputRef) && (operands.get(1) instanceof RexInputRef)) {
                return true;
            }
            // fall through

        default:
            return false;
        }
    }

    private static final JdbcPrule JDBC_PRULE_INSTANCE = new JdbcPrule();

    private static class JdbcPrule extends ConverterRule {

        private JdbcPrule() {
            super(JdbcDrel.class, DrillRel.DRILL_LOGICAL, Prel.DRILL_PHYSICAL, "JDBC_PREL_Converter");
        }

        @Override
        public RelNode convert(RelNode in) {

            return new JdbcIntermediatePrel(in.getCluster(), in.getTraitSet().replace(getOutTrait()),
                    in.getInput(0));
        }

    }

    private class JdbcDrelConverterRule extends ConverterRule {

        public JdbcDrelConverterRule(DrillJdbcConvention in) {
            super(RelNode.class, in, DrillRel.DRILL_LOGICAL, "JDBC_DREL_Converter" + in.getName());
        }

        @Override
        public RelNode convert(RelNode in) {
            return new JdbcDrel(in.getCluster(), in.getTraitSet().replace(DrillRel.DRILL_LOGICAL),
                    convert(in, in.getTraitSet().replace(this.getInTrait())));
        }

    }

    private class CapitalizingJdbcSchema extends AbstractSchema {

        final Map<String, CapitalizingJdbcSchema> schemaMap = Maps.newHashMap();
        private final JdbcSchema inner;

        public CapitalizingJdbcSchema(List<String> parentSchemaPath, String name, DataSource dataSource,
                SqlDialect dialect, JdbcConvention convention, String catalog, String schema) {
            super(parentSchemaPath, name);
            inner = new JdbcSchema(dataSource, dialect, convention, catalog, schema);
        }

        @Override
        public String getTypeName() {
            return JdbcStorageConfig.NAME;
        }

        @Override
        public Collection<Function> getFunctions(String name) {
            return inner.getFunctions(name);
        }

        @Override
        public Set<String> getFunctionNames() {
            return inner.getFunctionNames();
        }

        @Override
        public CapitalizingJdbcSchema getSubSchema(String name) {
            return schemaMap.get(name);
        }

        void setHolder(SchemaPlus plusOfThis) {
            for (String s : getSubSchemaNames()) {
                CapitalizingJdbcSchema inner = getSubSchema(s);
                SchemaPlus holder = plusOfThis.add(s, inner);
                inner.setHolder(holder);
            }
        }

        @Override
        public Set<String> getSubSchemaNames() {
            return schemaMap.keySet();
        }

        @Override
        public Set<String> getTableNames() {
            return inner.getTableNames();
        }

        public String toString() {
            return Joiner.on(".").join(getSchemaPath());
        }

        @Override
        public Table getTable(String name) {
            Table table = inner.getTable(name);
            if (table != null) {
                return table;
            }
            return inner.getTable(name.toUpperCase());

        }

    }

    private class JdbcCatalogSchema extends AbstractSchema {

        private final Map<String, CapitalizingJdbcSchema> schemaMap = Maps.newHashMap();
        private final CapitalizingJdbcSchema defaultSchema;

        public JdbcCatalogSchema(String name) {
            super(ImmutableList.<String>of(), name);

            try (Connection con = source.getConnection(); ResultSet set = con.getMetaData().getCatalogs()) {
                while (set.next()) {
                    final String catalogName = set.getString(1);
                    CapitalizingJdbcSchema schema = new CapitalizingJdbcSchema(getSchemaPath(), catalogName, source,
                            dialect, convention, catalogName, null);
                    schemaMap.put(catalogName, schema);
                }
            } catch (SQLException e) {
                logger.warn("Failure while attempting to load JDBC schema.", e);
            }

            // unable to read catalog list.
            if (schemaMap.isEmpty()) {

                // try to add a list of schemas to the schema map.
                boolean schemasAdded = addSchemas();

                if (!schemasAdded) {
                    // there were no schemas, just create a default one (the jdbc system doesn't support catalogs/schemas).
                    schemaMap.put("default", new CapitalizingJdbcSchema(ImmutableList.<String>of(), name, source,
                            dialect, convention, null, null));
                }
            } else {
                // We already have catalogs. Add schemas in this context of their catalogs.
                addSchemas();
            }

            defaultSchema = schemaMap.values().iterator().next();

        }

        void setHolder(SchemaPlus plusOfThis) {
            for (String s : getSubSchemaNames()) {
                CapitalizingJdbcSchema inner = getSubSchema(s);
                SchemaPlus holder = plusOfThis.add(s, inner);
                inner.setHolder(holder);
            }
        }

        private boolean addSchemas() {
            boolean added = false;
            try (Connection con = source.getConnection(); ResultSet set = con.getMetaData().getSchemas()) {
                while (set.next()) {
                    final String schemaName = set.getString(1);
                    final String catalogName = set.getString(2);

                    CapitalizingJdbcSchema parentSchema = schemaMap.get(catalogName);
                    if (parentSchema == null) {
                        CapitalizingJdbcSchema schema = new CapitalizingJdbcSchema(getSchemaPath(), schemaName,
                                source, dialect, convention, catalogName, schemaName);

                        // if a catalog schema doesn't exist, we'll add this at the top level.
                        schemaMap.put(schemaName, schema);
                    } else {
                        CapitalizingJdbcSchema schema = new CapitalizingJdbcSchema(parentSchema.getSchemaPath(),
                                schemaName, source, dialect, convention, catalogName, schemaName);
                        parentSchema.schemaMap.put(schemaName, schema);

                    }
                    added = true;
                }
            } catch (SQLException e) {
                logger.warn("Failure while attempting to load JDBC schema.", e);
            }

            return added;
        }

        @Override
        public String getTypeName() {
            return JdbcStorageConfig.NAME;
        }

        @Override
        public Schema getDefaultSchema() {
            return defaultSchema;
        }

        @Override
        public CapitalizingJdbcSchema getSubSchema(String name) {
            return schemaMap.get(name);
        }

        @Override
        public Set<String> getSubSchemaNames() {
            return schemaMap.keySet();
        }

        @Override
        public Table getTable(String name) {
            Schema schema = getDefaultSchema();

            if (schema != null) {
                try {
                    Table t = schema.getTable(name);
                    if (t != null) {
                        return t;
                    }
                    return schema.getTable(name.toUpperCase());
                } catch (RuntimeException e) {
                    logger.warn("Failure while attempting to read table '{}' from JDBC source.", name, e);

                }
            }

            // no table was found.
            return null;
        }

        @Override
        public Set<String> getTableNames() {
            return defaultSchema.getTableNames();
        }

    }

    @Override
    public void registerSchemas(SchemaConfig config, SchemaPlus parent) {
        JdbcCatalogSchema schema = new JdbcCatalogSchema(name);
        SchemaPlus holder = parent.add(name, schema);
        schema.setHolder(holder);
    }

    @Override
    public JdbcStorageConfig getConfig() {
        return config;
    }

    public DrillbitContext getContext() {
        return this.context;
    }

    public String getName() {
        return this.name;
    }

    @Override
    public boolean supportsRead() {
        return true;
    }

    public DataSource getSource() {
        return source;
    }

    public SqlDialect getDialect() {
        return dialect;
    }

    @Override
    public AbstractGroupScan getPhysicalScan(String userName, JSONOptions selection, List<SchemaPath> columns)
            throws IOException {
        throw new UnsupportedOperationException();
    }

    @Override
    public Set<RelOptRule> getPhysicalOptimizerRules(OptimizerRulesContext context) {
        return convention.getRules();
    }
}