<commit_before>/*
 * Copyright 2012 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.jobExecutor.ProcessJob;
import azkaban.utils.JSONUtils;
import azkaban.utils.Props;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.log4j.ConsoleAppender;
import org.apache.log4j.Layout;
import org.apache.log4j.Logger;
import org.apache.log4j.PatternLayout;

import azkaban.security.commons.HadoopSecurityManager;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.security.PrivilegedExceptionAction;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Properties;

import static azkaban.security.commons.SecurityUtils.MAPREDUCE_JOB_CREDENTIALS_BINARY;
import static org.apache.hadoop.security.UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION;

public class HadoopJavaJobRunnerMain {

  public static final String JOB_CLASS = "job.class";
  public static final String DEFAULT_RUN_METHOD = "run";
  public static final String DEFAULT_CANCEL_METHOD = "cancel";

  // This is the Job interface method to get the properties generated by the
  // job.
  public static final String GET_GENERATED_PROPERTIES_METHOD =
      "getJobGeneratedProperties";

  public static final String CANCEL_METHOD_PARAM = "method.cancel";
  public static final String RUN_METHOD_PARAM = "method.run";
  public static final String[] PROPS_CLASSES = new String[] {
      "azkaban.utils.Props", "azkaban.common.utils.Props" };

  private static final Layout DEFAULT_LAYOUT = new PatternLayout("%p %m\n");

  public final Logger _logger;

  public String _cancelMethod;
  public String _jobName;
  public Object _javaObject;
  private boolean _isFinished = false;

  private static boolean securityEnabled;

  public static void main(String[] args) throws Exception {
    @SuppressWarnings("unused")
    HadoopJavaJobRunnerMain wrapper = new HadoopJavaJobRunnerMain();
  }

  public HadoopJavaJobRunnerMain() throws Exception {
    Runtime.getRuntime().addShutdownHook(new Thread() {
      @Override
      public void run() {
        cancelJob();
      }
    });

    try {
      _jobName = System.getenv(ProcessJob.JOB_NAME_ENV);
      String propsFile = System.getenv(ProcessJob.JOB_PROP_ENV);

      _logger = Logger.getRootLogger();
      _logger.removeAllAppenders();
      ConsoleAppender appender = new ConsoleAppender(DEFAULT_LAYOUT);
      appender.activateOptions();
      _logger.addAppender(appender);

      Properties prop = new Properties();
      prop.load(new BufferedReader(new FileReader(propsFile)));

      HadoopConfigurationInjector.injectLinks();

      final Configuration conf = new Configuration();

      UserGroupInformation.setConfiguration(conf);
      securityEnabled = UserGroupInformation.isSecurityEnabled();

      _logger.info("Running job " + _jobName);
      String className = prop.getProperty(JOB_CLASS);
      if (className == null) {
        throw new Exception("Class name is not set.");
      }
      _logger.info("Class name " + className);

      UserGroupInformation loginUser = null;
      UserGroupInformation proxyUser = null;

      if (shouldProxy(prop)) {
        String userToProxy = prop.getProperty("user.to.proxy");
        if (securityEnabled) {
          String filelocation = System.getenv(HADOOP_TOKEN_FILE_LOCATION);
          _logger.info("Found token file " + filelocation);
          _logger.info("Security enabled is "
              + UserGroupInformation.isSecurityEnabled());

          _logger.info("Setting mapreduce.job.credentials.binary to "
              + filelocation);
          System.setProperty("mapreduce.job.credentials.binary", filelocation);

          _logger.info("Proxying enabled.");

          loginUser = UserGroupInformation.getLoginUser();

          _logger.info("Current logged in user is " + loginUser.getUserName());

          proxyUser =
              UserGroupInformation.createProxyUser(userToProxy, loginUser);
          for (Token<?> token : loginUser.getTokens()) {
            proxyUser.addToken(token);
          }
        } else {
          proxyUser = UserGroupInformation.createRemoteUser(userToProxy);
        }
        _logger.info("Proxied as user " + userToProxy);
      }

      // Create the object using proxy
      if (shouldProxy(prop)) {
        _javaObject =
            getObjectAsProxyUser(prop, _logger, _jobName, className, proxyUser);
      } else {
        _javaObject = getObject(_jobName, className, prop, _logger);
      }

      if (_javaObject == null) {
        _logger.info("Could not create java object to run job: " + className);
        throw new Exception("Could not create running object");
      }
      _logger.info("Got object " + _javaObject.toString());

      _cancelMethod =
          prop.getProperty(CANCEL_METHOD_PARAM, DEFAULT_CANCEL_METHOD);

      final String runMethod =
          prop.getProperty(RUN_METHOD_PARAM, DEFAULT_RUN_METHOD);
      _logger.info("Invoking method " + runMethod);

      if (shouldProxy(prop)) {
        _logger.info("Proxying enabled.");
        runMethodAsUser(prop, _javaObject, runMethod, proxyUser);
      } else {
        _logger.info("Proxy check failed, not proxying run.");
        runMethod(_javaObject, runMethod);
      }

      _isFinished = true;

      // Get the generated properties and store them to disk, to be read
      // by ProcessJob.
      try {
        final Method generatedPropertiesMethod =
            _javaObject.getClass().getMethod(GET_GENERATED_PROPERTIES_METHOD,
                new Class<?>[] {});
        Object outputGendProps =
            generatedPropertiesMethod.invoke(_javaObject, new Object[] {});
        final Method toPropertiesMethod =
            outputGendProps.getClass().getMethod("toProperties",
                new Class<?>[] {});
        Properties properties =
            (Properties) toPropertiesMethod.invoke(outputGendProps,
                new Object[] {});

        Props outputProps = new Props(null, properties);
        outputGeneratedProperties(outputProps);
      } catch (NoSuchMethodException e) {
        _logger.info(String.format(
            "Apparently there isn't a method[%s] on object[%s], using "
                + "empty Props object instead.",
            GET_GENERATED_PROPERTIES_METHOD, _javaObject));
        outputGeneratedProperties(new Props());
      }
    } catch (Exception e) {
      _isFinished = true;<commit_msg>HadoopJavaJobRunnerMain doesn't properly handle returned Props object from job. Issue #146<commit_after>/*
 * Copyright 2012 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 static azkaban.security.commons.SecurityUtils.MAPREDUCE_JOB_CREDENTIALS_BINARY;
import static org.apache.hadoop.security.UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.security.PrivilegedExceptionAction;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Properties;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.log4j.ConsoleAppender;
import org.apache.log4j.Layout;
import org.apache.log4j.Logger;
import org.apache.log4j.PatternLayout;

import azkaban.jobExecutor.ProcessJob;
import azkaban.security.commons.HadoopSecurityManager;
import azkaban.utils.JSONUtils;
import azkaban.utils.Props;

public class HadoopJavaJobRunnerMain {

  public static final String JOB_CLASS = "job.class";
  public static final String DEFAULT_RUN_METHOD = "run";
  public static final String DEFAULT_CANCEL_METHOD = "cancel";

  // This is the Job interface method to get the properties generated by the
  // job.
  public static final String GET_GENERATED_PROPERTIES_METHOD =
      "getJobGeneratedProperties";

  public static final String CANCEL_METHOD_PARAM = "method.cancel";
  public static final String RUN_METHOD_PARAM = "method.run";
  public static final String[] PROPS_CLASSES = new String[] {
      "azkaban.utils.Props", "azkaban.common.utils.Props" };

  private static final Layout DEFAULT_LAYOUT = new PatternLayout("%p %m\n");

  public final Logger _logger;

  public String _cancelMethod;
  public String _jobName;
  public Object _javaObject;
  private boolean _isFinished = false;

  private static boolean securityEnabled;

  public static void main(String[] args) throws Exception {
    @SuppressWarnings("unused")
    HadoopJavaJobRunnerMain wrapper = new HadoopJavaJobRunnerMain();
  }

  public HadoopJavaJobRunnerMain() throws Exception {
    Runtime.getRuntime().addShutdownHook(new Thread() {
      @Override
      public void run() {
        cancelJob();
      }
    });

    try {
      _jobName = System.getenv(ProcessJob.JOB_NAME_ENV);
      String propsFile = System.getenv(ProcessJob.JOB_PROP_ENV);

      _logger = Logger.getRootLogger();
      _logger.removeAllAppenders();
      ConsoleAppender appender = new ConsoleAppender(DEFAULT_LAYOUT);
      appender.activateOptions();
      _logger.addAppender(appender);

      Properties prop = new Properties();
      prop.load(new BufferedReader(new FileReader(propsFile)));

      HadoopConfigurationInjector.injectLinks();

      final Configuration conf = new Configuration();

      UserGroupInformation.setConfiguration(conf);
      securityEnabled = UserGroupInformation.isSecurityEnabled();

      _logger.info("Running job " + _jobName);
      String className = prop.getProperty(JOB_CLASS);
      if (className == null) {
        throw new Exception("Class name is not set.");
      }
      _logger.info("Class name " + className);

      UserGroupInformation loginUser = null;
      UserGroupInformation proxyUser = null;

      if (shouldProxy(prop)) {
        String userToProxy = prop.getProperty("user.to.proxy");
        if (securityEnabled) {
          String filelocation = System.getenv(HADOOP_TOKEN_FILE_LOCATION);
          _logger.info("Found token file " + filelocation);
          _logger.info("Security enabled is "
              + UserGroupInformation.isSecurityEnabled());

          _logger.info("Setting mapreduce.job.credentials.binary to "
              + filelocation);
          System.setProperty("mapreduce.job.credentials.binary", filelocation);

          _logger.info("Proxying enabled.");

          loginUser = UserGroupInformation.getLoginUser();

          _logger.info("Current logged in user is " + loginUser.getUserName());

          proxyUser =
              UserGroupInformation.createProxyUser(userToProxy, loginUser);
          for (Token<?> token : loginUser.getTokens()) {
            proxyUser.addToken(token);
          }
        } else {
          proxyUser = UserGroupInformation.createRemoteUser(userToProxy);
        }
        _logger.info("Proxied as user " + userToProxy);
      }

      // Create the object using proxy
      if (shouldProxy(prop)) {
        _javaObject =
            getObjectAsProxyUser(prop, _logger, _jobName, className, proxyUser);
      } else {
        _javaObject = getObject(_jobName, className, prop, _logger);
      }

      if (_javaObject == null) {
        _logger.info("Could not create java object to run job: " + className);
        throw new Exception("Could not create running object");
      }
      _logger.info("Got object " + _javaObject.toString());

      _cancelMethod =
          prop.getProperty(CANCEL_METHOD_PARAM, DEFAULT_CANCEL_METHOD);

      final String runMethod =
          prop.getProperty(RUN_METHOD_PARAM, DEFAULT_RUN_METHOD);
      _logger.info("Invoking method " + runMethod);

      if (shouldProxy(prop)) {
        _logger.info("Proxying enabled.");
        runMethodAsUser(prop, _javaObject, runMethod, proxyUser);
      } else {
        _logger.info("Proxy check failed, not proxying run.");
        runMethod(_javaObject, runMethod);
      }

      _isFinished = true;

      // Get the generated properties and store them to disk, to be read
      // by ProcessJob.
      try {
        final Method generatedPropertiesMethod =
            _javaObject.getClass().getMethod(GET_GENERATED_PROPERTIES_METHOD,
                new Class<?>[] {});
        Object outputGendProps =
            generatedPropertiesMethod.invoke(_javaObject, new Object[] {});

        if (outputGendProps != null) {
          final Method toPropertiesMethod =
              outputGendProps.getClass().getMethod("toProperties",
                  new Class<?>[] {});
          Properties properties =
              (Properties) toPropertiesMethod.invoke(outputGendProps,
                  new Object[] {});

          Props outputProps = new Props(null, properties);
          outputGeneratedProperties(outputProps);
        } else {
          _logger.info(GET_GENERATED_PROPERTIES_METHOD
              + " method returned null.  No properties to pass along");
        }
      } catch (NoSuchMethodException e) {
        _logger.info(String.format(
            "Apparently there isn't a method[%s] on object[%s], using "
                + "empty Props object instead.",
            GET_GENERATED_PROPERTIES_METHOD, _javaObject));
        outputGeneratedProperties(new Props());
      }
    } catch (Exception e) {
      _isFinished = true;
{'fraction_non_alphanumeric': 0.09318060021436227, 'fraction_numerical': 0.0018086816720257234, 'mean_word_length': 3.230376877302352, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7059059', 'n_tokens_mistral': 4072, 'n_tokens_neox': 3989, 'n_words': 1022}