src/java/org/apache/sqoop/mapreduce/postgresql/PGBulkloadExportMapper.java

/**
 * 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.sqoop.mapreduce.postgresql;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.Mapper.Context;
import org.apache.sqoop.lib.SqoopRecord;
import org.apache.sqoop.mapreduce.AutoProgressMapper;
import org.apache.sqoop.mapreduce.db.DBConfiguration;
import org.apache.sqoop.util.LoggingUtils;
import org.apache.sqoop.util.PostgreSQLUtils;
import org.apache.sqoop.util.Executor;
import org.apache.sqoop.util.JdbcUrl;


/**
 * Mapper that starts a 'pg_bulkload' process and uses that to export rows from
 * HDFS to a PostgreSQL database at high speed.
 *
 * map() methods are actually provided by subclasses that read from
 * SequenceFiles (containing existing SqoopRecords) or text files
 * (containing delimited lines) and deliver these results to the stream
 * used to interface with pg_bulkload.
 */
public class PGBulkloadExportMapper
    extends AutoProgressMapper<LongWritable, Writable, LongWritable, Text> {
  private Configuration conf;
  private DBConfiguration dbConf;
  private Process process;
  private OutputStream out;
  protected BufferedWriter writer;
  private Thread thread;
  protected String tmpTableName;
  private String tableName;
  private String passwordFilename;


  public PGBulkloadExportMapper() {
  }


  protected void setup(Context context)
      throws IOException, InterruptedException {
    super.setup(context);
    conf = context.getConfiguration();
    dbConf = new DBConfiguration(conf);
    tableName = dbConf.getOutputTableName();
    tmpTableName = tableName + "_" + context.getTaskAttemptID().toString();

    Connection conn = null;
    try {
      conn = dbConf.getConnection();
      conn.setAutoCommit(false);
      if (conf.getBoolean("pgbulkload.clear.staging.table", false)) {
        StringBuffer query = new StringBuffer();
        query.append("DROP TABLE IF EXISTS ");
        query.append(tmpTableName);
        doExecuteUpdate(query.toString());
      }
      StringBuffer query = new StringBuffer();
      query.append("CREATE TABLE ");
      query.append(tmpTableName);
      query.append("(LIKE ");
      query.append(tableName);
      query.append(" INCLUDING CONSTRAINTS)");
      if (conf.get("pgbulkload.staging.tablespace") != null) {
        query.append("TABLESPACE ");
        query.append(conf.get("pgbulkload.staging.tablespace"));
      }
      doExecuteUpdate(query.toString());
      conn.commit();
    } catch (ClassNotFoundException ex) {
      LOG.error("Unable to load JDBC driver class", ex);
      throw new IOException(ex);
    } catch (SQLException ex) {
      LoggingUtils.logAll(LOG, "Unable to execute statement", ex);
      throw new IOException(ex);
    } finally {
      try {
        conn.close();
      } catch (SQLException ex) {
        LoggingUtils.logAll(LOG, "Unable to close connection", ex);
      }
    }

    try {
      ArrayList<String> args = new ArrayList<String>();
      List<String> envp = Executor.getCurEnvpStrings();
      args.add(conf.get("pgbulkload.bin", "pg_bulkload"));
      args.add("--username="
          + conf.get(DBConfiguration.USERNAME_PROPERTY));
      args.add("--dbname="
          + JdbcUrl.getDatabaseName(conf.get(DBConfiguration.URL_PROPERTY)));
      args.add("--host="
          + JdbcUrl.getHostName(conf.get(DBConfiguration.URL_PROPERTY)));
      int port = JdbcUrl.getPort(conf.get(DBConfiguration.URL_PROPERTY));
      if (port != -1) {
        args.add("--port=" + port);
      }
      args.add("--input=stdin");
      args.add("--output=" + tmpTableName);
      args.add("-o");
      args.add("TYPE=CSV");
      args.add("-o");
      args.add("DELIMITER=" + conf.get("pgbulkload.input.field.delim", ","));
      args.add("-o");
      args.add("QUOTE=" + conf.get("pgbulkload.input.enclosedby", "\""));
      args.add("-o");
      args.add("ESCAPE=" + conf.get("pgbulkload.input.escapedby", "\""));
      args.add("-o");
      args.add("CHECK_CONSTRAINTS=" + conf.get("pgbulkload.check.constraints"));
      args.add("-o");
      args.add("PARSE_ERRORS=" + conf.get("pgbulkload.parse.errors"));
      args.add("-o");
      args.add("DUPLICATE_ERRORS=" + conf.get("pgbulkload.duplicate.errors"));
      if (conf.get("pgbulkload.null.string") != null) {
        args.add("-o");
        args.add("NULL=" + conf.get("pgbulkload.null.string"));
      }
      if (conf.get("pgbulkload.filter") != null) {
        args.add("-o");
        args.add("FILTER=" + conf.get("pgbulkload.filter"));
      }
      LOG.debug("Starting pg_bulkload with arguments:");
      for (String arg : args) {
        LOG.debug("  " + arg);
      }
      if (conf.get(DBConfiguration.PASSWORD_PROPERTY) != null) {
        String tmpDir = System.getProperty("test.build.data", "/tmp/");
        if (!tmpDir.endsWith(File.separator)) {
          tmpDir = tmpDir + File.separator;
        }
        tmpDir = conf.get("job.local.dir", tmpDir);
        passwordFilename = PostgreSQLUtils.writePasswordFile(tmpDir,
            conf.get(DBConfiguration.PASSWORD_PROPERTY));
        envp.add("PGPASSFILE=" + passwordFilename);
      }
      process = Runtime.getRuntime().exec(args.toArray(new String[0]),
                                          envp.toArray(new String[0]));
      out = process.getOutputStream();
      writer = new BufferedWriter(new OutputStreamWriter(out));
      thread = new ReadThread(process.getErrorStream());
      thread.start();
    } catch (Exception e) {
      LOG.error("Can't start up pg_bulkload process", e);
      cleanup(context);
      doExecuteUpdate("DROP TABLE " + tmpTableName);
      throw new IOException(e);
    }
  }


  public void map(LongWritable key, Writable value, Context context)
    throws IOException, InterruptedException {
    try {
      String str = value.toString();
      if (value instanceof Text) {
        writer.write(str, 0, str.length());
        writer.newLine();
      } else if (value instanceof SqoopRecord) {
        writer.write(str, 0, str.length());
      }
    } catch (Exception e) {
      doExecuteUpdate("DROP TABLE " + tmpTableName);
      cleanup(context);
      throw new IOException(e);
    }
  }


  protected void cleanup(Context context)
    throws IOException, InterruptedException {
    LongWritable taskid =
        new LongWritable(context.getTaskAttemptID().getTaskID().getId());
    context.write(taskid, new Text(tmpTableName));

    if (writer != null) {
      writer.close();
    }
    if (out != null) {
      out.close();
    }
    try {
      if (thread != null) {
        thread.join();
      }
    } finally {
      // block until the process is done.
      if (null != process) {
        while (true) {
          try {
            int returnValue = process.waitFor();

            // Check pg_bulkload's process return value
            if (returnValue != 0) {
              throw new RuntimeException(
                "Unexpected return value from pg_bulkload: "+ returnValue);
            }
          } catch (InterruptedException ie) {
            // interrupted; loop around.
            LOG.debug("Caught interrupted exception waiting for process "
                + "pg_bulkload.bin to exit");
            //Clear the interrupted flag.   We have to call Thread.interrupted
            //to clear for interrupted exceptions from process.waitFor
            //See http://bugs.sun.com/view_bug.do?bug_id=6420270 for more info
            Thread.interrupted();
            continue;
          }
          break;
        }
      }
    }
    if (null != passwordFilename) {
      if (!new File(passwordFilename).delete()) {
        LOG.error("Could not remove postgresql password file "
                  + passwordFilename);
        LOG.error("You should remove this file to protect your credentials.");
      }
    }
  }


  protected int doExecuteUpdate(String query) throws IOException {
    Connection conn = null;
    try {
      conn = dbConf.getConnection();
      conn.setAutoCommit(false);
    } catch (ClassNotFoundException ex) {
      LOG.error("Unable to load JDBC driver class", ex);
      throw new IOException(ex);
    } catch (SQLException ex) {
      LoggingUtils.logAll(LOG, "Unable to connect to database", ex);
      throw new IOException(ex);
    }
    Statement stmt = null;
    try {
      stmt = conn.createStatement();
      int ret = stmt.executeUpdate(query);
      conn.commit();
      return ret;
    } catch (SQLException ex) {
      LoggingUtils.logAll(LOG, "Unable to execute query: "  + query, ex);
      throw new IOException(ex);
    } finally {
      if (stmt != null) {
        try {
          stmt.close();
        } catch (SQLException ex) {
          LoggingUtils.logAll(LOG, "Unable to close statement", ex);
        }
      }
      try {
        conn.close();
      } catch (SQLException ex) {
        LoggingUtils.logAll(LOG, "Unable to close connection", ex);
      }
    }
  }


  private class ReadThread extends Thread {
    private InputStream in;

    ReadThread(InputStream in) {
      this.in = in;
    }

    public void run() {
      BufferedReader reader = new BufferedReader(new InputStreamReader(in));
      String line = null;
      try {
        while((line = reader.readLine()) != null) {
          System.out.println(line);
        }
        reader.close();
      } catch (Exception e) {
        e.printStackTrace();
      }
    }
  }
}