azkaban.jobtype.ReportalPrestoRunner.java Source code

Introduction

Here is the source code for azkaban.jobtype.ReportalPrestoRunner.java

Source

/*
 * Copyright 2018 LinkedIn Corp.
 *
 * 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 azkaban.jobtype;

import azkaban.crypto.Decryptions;
import azkaban.reportal.util.ReportalUtil;
import com.google.common.base.Preconditions;
import java.io.IOException;
import java.net.URI;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.Statement;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;

public class ReportalPrestoRunner extends ReportalAbstractRunner {

    public static final String JDBC_DRIVER_KEY = "presto.driver";
    public static final String PRESTO_USER = "presto.driver.user";
    public static final String DRIVER_URL = "presto.driver.jdbc.url";
    private static final String PRESTO_DRIVER_PROP_PREFIX = "presto.driver.";
    private static final String IMPERSONATED_USER_KEY = "presto.execute.user";

    public ReportalPrestoRunner(final String jobName, final Properties props) {
        super(props);

        Preconditions.checkArgument(props.containsKey(JDBC_DRIVER_KEY), "missing " + JDBC_DRIVER_KEY);
        Preconditions.checkArgument(props.containsKey(PRESTO_USER), "missing " + PRESTO_USER);
        Preconditions.checkArgument(props.containsKey(DRIVER_URL), "missing " + DRIVER_URL);
    }

    private String decrypt(final String encrypted, final String keyPath) throws IOException {
        final FileSystem fs = FileSystem.get(URI.create("file:///"), new Configuration());
        return new Decryptions().decrypt(encrypted, keyPath, fs);
    }

    private Properties getProperties() throws IOException {
        final Properties connProperties = new Properties();
        final Map<String, String> prestoProps = this.props.getMapByPrefix(PRESTO_DRIVER_PROP_PREFIX);

        for (final Entry<String, String> entry : prestoProps.entrySet()) {
            final String key = entry.getKey();
            final String value = entry.getValue();
            if (key.contains("encrypted.")) {
                // if props contains "encrypted." then decrypted it with key file and put the decrypted
                // value into jdbc connection props
                //"encrypted.password" => "password"
                final String jdbcProp = key.replaceFirst("encrypted.", "");
                connProperties.put(jdbcProp, decrypt(value, prestoProps.get("jdbc.crypto.key.path")));
            } else if (!key.equals("jdbc.url") && !key.equals("jdbc.crypto.key.path")) {
                connProperties.put(key, value);
            }
        }

        return connProperties;
    }

    private Connection getConnection(final String jdbcUrl, final String userToProxy) {
        try {
            Class.forName(this.props.get(JDBC_DRIVER_KEY));
            final Properties connProperties = getProperties();
            connProperties.put(IMPERSONATED_USER_KEY, userToProxy);
            final Connection conn = DriverManager.getConnection(jdbcUrl, connProperties);
            return conn;
        } catch (final Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    protected void runReportal() throws Exception {
        final Connection conn = getConnection(this.props.get(PRESTO_DRIVER_PROP_PREFIX + "jdbc.url"),
                this.proxyUser);
        final Statement statement = conn.createStatement();
        try {
            String processedQuery = injectVariables(this.jobQuery);
            processedQuery = processedQuery.trim();

            // jdbc presto driver doesn't work with query with trailing semicolon,
            // if user query has semicolon, remove it.
            if (processedQuery.endsWith(";")) {
                processedQuery = processedQuery.substring(0, processedQuery.length() - 1);
            }

            statement.execute(processedQuery);
            ReportalUtil.outputQueryResult(statement.getResultSet(), this.outputStream);
        } finally {
            statement.close();
            conn.close();
        }
    }
}