src/java/org/apache/sqoop/mapreduce/AsyncSqlOutputFormat.java - sqoop - Git at Google

 /**
  * 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;

 import java.io.IOException;
 import java.sql.BatchUpdateException;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
 import java.util.concurrent.SynchronousQueue;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.mapreduce.JobContext;
 import org.apache.hadoop.mapreduce.OutputCommitter;
 import org.apache.hadoop.mapreduce.OutputFormat;
 import org.apache.hadoop.mapreduce.TaskAttemptContext;
 import org.apache.hadoop.util.StringUtils;

 import org.apache.sqoop.lib.SqoopRecord;

 /**
  * Abstract OutputFormat class that allows the RecordWriter to buffer
  * up SQL commands which should be executed in a separate thread after
  * enough commands are created.
  *
  * This supports a configurable "spill threshold" at which
  * point intermediate transactions are committed.
  *
  * Uses DBOutputFormat/DBConfiguration for configuring the output.
  * This is used in conjunction with the abstract AsyncSqlRecordWriter
  * class.
  *
  * Clients of this OutputFormat must implement getRecordWriter(); the
  * returned RecordWriter is intended to subclass AsyncSqlRecordWriter.
  */
 public abstract class AsyncSqlOutputFormat<K extends SqoopRecord, V>
     extends OutputFormat<K, V> {

   /** conf key: number of rows to export per INSERT statement. */
   public static final String RECORDS_PER_STATEMENT_KEY =
       "sqoop.export.records.per.statement";

   /** conf key: number of INSERT statements to bundle per tx.
    * If this is set to -1, then a single transaction will be used
    * per task. Note that each statement may encompass multiple
    * rows, depending on the value of sqoop.export.records.per.statement.
    */
   public static final String STATEMENTS_PER_TRANSACTION_KEY =
       "sqoop.export.statements.per.transaction";

   /**
    * Default number of records to put in an INSERT statement or
    * other batched update statement.
    */
   public static final int DEFAULT_RECORDS_PER_STATEMENT = 100;

   /**
    * Default number of statements to execute before committing the
    * current transaction.
    */
   public static final int DEFAULT_STATEMENTS_PER_TRANSACTION = 100;

   /**
    * Value for STATEMENTS_PER_TRANSACTION_KEY signifying that we should
    * not commit until the RecordWriter is being closed, regardless of
    * the number of statements we execute.
    */
   public static final int UNLIMITED_STATEMENTS_PER_TRANSACTION = -1;

   private static final Log LOG = LogFactory.getLog(AsyncSqlOutputFormat.class);

   @Override
   /** {@inheritDoc} */
   public void checkOutputSpecs(JobContext context)
       throws IOException, InterruptedException {
   }

   @Override
   /** {@inheritDoc} */
   public OutputCommitter getOutputCommitter(TaskAttemptContext context)
       throws IOException, InterruptedException {
     return new NullOutputCommitter();
   }

   /**
    * Represents a database update operation that should be performed
    * by an asynchronous background thread.
    * AsyncDBOperation objects are immutable.
    * They MAY contain a statement which should be executed. The
    * statement may also be null.
    *
    * They may also set 'commitAndClose' to true. If true, then the
    * executor of this operation should commit the current
    * transaction, even if stmt is null, and then stop the executor
    * thread.
    */
   public static class AsyncDBOperation {
     private final PreparedStatement stmt;
     private final boolean isBatch;
     private final boolean commit;
     private final boolean stopThread;

     @Deprecated
     /** Do not use AsyncDBOperation(PreparedStatement s, boolean
      * commitAndClose, boolean batch). Use AsyncDBOperation(PreparedStatement
      *  s, boolean batch, boolean commit, boolean stopThread) instead.
      */
     public AsyncDBOperation(PreparedStatement s, boolean commitAndClose,
         boolean batch) {
         this(s, batch, commitAndClose, commitAndClose);
     }

     /**
      * Create an asynchronous database operation.
      * @param s the statement, if any, to execute.
      * @param batch is true if this is a batch PreparedStatement, or false
      * if it's a normal singleton statement.
      * @param commit is true if this statement should be committed to the
      * database.
      * @param stopThread if true, the executor thread should stop after this
      * operation.
      */
     public AsyncDBOperation(PreparedStatement s, boolean batch,
         boolean commit, boolean stopThread) {
       this.stmt = s;
       this.isBatch = batch;
       this.commit = commit;
       this.stopThread = stopThread;
     }

     /**
      * @return a statement to run as an update.
      */
     public PreparedStatement getStatement() {
       return stmt;
     }

     /**
      * @return true if the executor should commit the current transaction.
      * If getStatement() is non-null, the statement is run first.
      */
     public boolean requiresCommit() {
       return this.commit;
     }

     /**
      * @return true if the executor should stop after this command.
      */
     public boolean stop() {
       return this.stopThread;
     }

     /**
      * @return true if this is a batch SQL statement.
      */
     public boolean execAsBatch() {
       return this.isBatch;
     }
   }

   /**
    * A thread that runs the database interactions asynchronously
    * from the OutputCollector.
    */
   public static class AsyncSqlExecThread extends Thread {

     private final Connection conn; // The connection to the database.
     private SQLException err; // Error from a previously-run statement.

     // How we receive database operations from the RecordWriter.
     private SynchronousQueue<AsyncDBOperation> opsQueue;

     protected int curNumStatements; // statements executed thus far in the tx.
     protected final int stmtsPerTx;  // statements per transaction.

     /**
      * Create a new update thread that interacts with the database.
      * @param conn the connection to use. This must only be used by this
      * thread.
      * @param stmtsPerTx the number of statements to execute before committing
      * the current transaction.
      */
     public AsyncSqlExecThread(Connection conn, int stmtsPerTx) {
       this.conn = conn;
       this.err = null;
       this.opsQueue = new SynchronousQueue<AsyncDBOperation>();
       this.stmtsPerTx = stmtsPerTx;
     }

     public void run() {
       while (true) {
         AsyncDBOperation op = null;
         try {
           op = opsQueue.take();
         } catch (InterruptedException ie) {
           LOG.warn("Interrupted retrieving from operation queue: "
               + StringUtils.stringifyException(ie));
           continue;
         }

         if (null == op) {
           // This shouldn't be allowed to happen.
           LOG.warn("Null operation in queue; illegal state.");
           continue;
         }

         PreparedStatement stmt = op.getStatement();
         // Synchronize on the connection to ensure it does not conflict
         // with the prepareStatement() call in the main thread.
         synchronized (conn) {
           try {
             if (null != stmt) {
               if (op.execAsBatch()) {
                 stmt.executeBatch();
               } else {
                 stmt.execute();
               }
               stmt.close();
               stmt = null;
               this.curNumStatements++;
             }

             if (op.requiresCommit() || (curNumStatements >= stmtsPerTx
                 && stmtsPerTx != UNLIMITED_STATEMENTS_PER_TRANSACTION)) {
               LOG.debug("Committing transaction of " + curNumStatements
                   + " statements");
               this.conn.commit();
               this.curNumStatements = 0;
             }
           } catch (BatchUpdateException batchE) {
             if (batchE.getNextException() != null) {
               // if a statement in a batch causes an SQLException
               // the database can either set it as the cause of
               // the BatchUpdateException, or set it as the 'next'
               // field of the BatchUpdateException (e.g. HSQLDB 1.8
               // does the former and Postgres 8.4 does the latter).
               // We'll check for this SQLException in both places,
               // and use the 'next' one in preference.
               setLastError(batchE.getNextException());
             } else {
               // same as SQLException block
               setLastError(batchE);
             }
           } catch (SQLException sqlE) {
             setLastError(sqlE);
           } finally {
             // Close the statement on our way out if that didn't happen
             // via the normal execution path.
             if (null != stmt) {
               try {
                 stmt.close();
               } catch (SQLException sqlE) {
                 setLastError(sqlE);
               }
             }

             // Always check whether we should end the loop, regardless
             // of the presence of an exception.
             if (op.stop()) {
               return;
             }
           } // try .. catch .. finally.
         } // synchronized (conn)
       }
     }

     /**
      * Allows a user to enqueue the next database operation to run.
      * Since the connection can only execute a single operation at a time,
      * the put() method may block if another operation is already underway.
      * @param op the database operation to perform.
      */
     public void put(AsyncDBOperation op) throws InterruptedException {
       opsQueue.put(op);
     }

     /**
      * If a previously-executed statement resulted in an error, post it here.
      * If the error slot was already filled, then subsequent errors are
      * squashed until the user calls this method (which clears the error
      * slot).
      * @return any SQLException that occurred due to a previously-run
      * statement.
      */
     public synchronized SQLException getLastError() {
       SQLException e = this.err;
       this.err = null;
       return e;
     }

     private synchronized void setLastError(SQLException e) {
       if (this.err == null) {
         // Just set it.
         LOG.error("Got exception in update thread: "
             + StringUtils.stringifyException(e));
         this.err = e;
       } else {
         // Slot is full. Log it and discard.
         LOG.error("SQLException in update thread but error slot full: "
             + StringUtils.stringifyException(e));
       }
     }
   }
 }
	/**
	* 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;

	import java.io.IOException;
	import java.sql.BatchUpdateException;
	import java.sql.Connection;
	import java.sql.PreparedStatement;
	import java.sql.SQLException;
	import java.util.concurrent.SynchronousQueue;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.hadoop.mapreduce.JobContext;
	import org.apache.hadoop.mapreduce.OutputCommitter;
	import org.apache.hadoop.mapreduce.OutputFormat;
	import org.apache.hadoop.mapreduce.TaskAttemptContext;
	import org.apache.hadoop.util.StringUtils;

	import org.apache.sqoop.lib.SqoopRecord;

	/**
	* Abstract OutputFormat class that allows the RecordWriter to buffer
	* up SQL commands which should be executed in a separate thread after
	* enough commands are created.
	*
	* This supports a configurable "spill threshold" at which
	* point intermediate transactions are committed.
	*
	* Uses DBOutputFormat/DBConfiguration for configuring the output.
	* This is used in conjunction with the abstract AsyncSqlRecordWriter
	* class.
	*
	* Clients of this OutputFormat must implement getRecordWriter(); the
	* returned RecordWriter is intended to subclass AsyncSqlRecordWriter.
	*/
	public abstract class AsyncSqlOutputFormat<K extends SqoopRecord, V>
	extends OutputFormat<K, V> {

	/** conf key: number of rows to export per INSERT statement. */
	public static final String RECORDS_PER_STATEMENT_KEY =
	"sqoop.export.records.per.statement";

	/** conf key: number of INSERT statements to bundle per tx.
	* If this is set to -1, then a single transaction will be used
	* per task. Note that each statement may encompass multiple
	* rows, depending on the value of sqoop.export.records.per.statement.
	*/
	public static final String STATEMENTS_PER_TRANSACTION_KEY =
	"sqoop.export.statements.per.transaction";

	/**
	* Default number of records to put in an INSERT statement or
	* other batched update statement.
	*/
	public static final int DEFAULT_RECORDS_PER_STATEMENT = 100;

	/**
	* Default number of statements to execute before committing the
	* current transaction.
	*/
	public static final int DEFAULT_STATEMENTS_PER_TRANSACTION = 100;

	/**
	* Value for STATEMENTS_PER_TRANSACTION_KEY signifying that we should
	* not commit until the RecordWriter is being closed, regardless of
	* the number of statements we execute.
	*/
	public static final int UNLIMITED_STATEMENTS_PER_TRANSACTION = -1;

	private static final Log LOG = LogFactory.getLog(AsyncSqlOutputFormat.class);

	@Override
	/** {@inheritDoc} */
	public void checkOutputSpecs(JobContext context)
	throws IOException, InterruptedException {
	}

	@Override
	/** {@inheritDoc} */
	public OutputCommitter getOutputCommitter(TaskAttemptContext context)
	throws IOException, InterruptedException {
	return new NullOutputCommitter();
	}

	/**
	* Represents a database update operation that should be performed
	* by an asynchronous background thread.
	* AsyncDBOperation objects are immutable.
	* They MAY contain a statement which should be executed. The
	* statement may also be null.
	*
	* They may also set 'commitAndClose' to true. If true, then the
	* executor of this operation should commit the current
	* transaction, even if stmt is null, and then stop the executor
	* thread.
	*/
	public static class AsyncDBOperation {
	private final PreparedStatement stmt;
	private final boolean isBatch;
	private final boolean commit;
	private final boolean stopThread;

	@Deprecated
	/** Do not use AsyncDBOperation(PreparedStatement s, boolean
	* commitAndClose, boolean batch). Use AsyncDBOperation(PreparedStatement
	* s, boolean batch, boolean commit, boolean stopThread) instead.
	*/
	public AsyncDBOperation(PreparedStatement s, boolean commitAndClose,
	boolean batch) {
	this(s, batch, commitAndClose, commitAndClose);
	}

	/**
	* Create an asynchronous database operation.
	* @param s the statement, if any, to execute.
	* @param batch is true if this is a batch PreparedStatement, or false
	* if it's a normal singleton statement.
	* @param commit is true if this statement should be committed to the
	* database.
	* @param stopThread if true, the executor thread should stop after this
	* operation.
	*/
	public AsyncDBOperation(PreparedStatement s, boolean batch,
	boolean commit, boolean stopThread) {
	this.stmt = s;
	this.isBatch = batch;
	this.commit = commit;
	this.stopThread = stopThread;
	}

	/**
	* @return a statement to run as an update.
	*/
	public PreparedStatement getStatement() {
	return stmt;
	}

	/**
	* @return true if the executor should commit the current transaction.
	* If getStatement() is non-null, the statement is run first.
	*/
	public boolean requiresCommit() {
	return this.commit;
	}

	/**
	* @return true if the executor should stop after this command.
	*/
	public boolean stop() {
	return this.stopThread;
	}

	/**
	* @return true if this is a batch SQL statement.
	*/
	public boolean execAsBatch() {
	return this.isBatch;
	}
	}

	/**
	* A thread that runs the database interactions asynchronously
	* from the OutputCollector.
	*/
	public static class AsyncSqlExecThread extends Thread {

	private final Connection conn; // The connection to the database.
	private SQLException err; // Error from a previously-run statement.

	// How we receive database operations from the RecordWriter.
	private SynchronousQueue<AsyncDBOperation> opsQueue;

	protected int curNumStatements; // statements executed thus far in the tx.
	protected final int stmtsPerTx; // statements per transaction.

	/**
	* Create a new update thread that interacts with the database.
	* @param conn the connection to use. This must only be used by this
	* thread.
	* @param stmtsPerTx the number of statements to execute before committing
	* the current transaction.
	*/
	public AsyncSqlExecThread(Connection conn, int stmtsPerTx) {
	this.conn = conn;
	this.err = null;
	this.opsQueue = new SynchronousQueue<AsyncDBOperation>();
	this.stmtsPerTx = stmtsPerTx;
	}

	public void run() {
	while (true) {
	AsyncDBOperation op = null;
	try {
	op = opsQueue.take();
	} catch (InterruptedException ie) {
	LOG.warn("Interrupted retrieving from operation queue: "
	+ StringUtils.stringifyException(ie));
	continue;
	}

	if (null == op) {
	// This shouldn't be allowed to happen.
	LOG.warn("Null operation in queue; illegal state.");
	continue;
	}

	PreparedStatement stmt = op.getStatement();
	// Synchronize on the connection to ensure it does not conflict
	// with the prepareStatement() call in the main thread.
	synchronized (conn) {
	try {
	if (null != stmt) {
	if (op.execAsBatch()) {
	stmt.executeBatch();
	} else {
	stmt.execute();
	}
	stmt.close();
	stmt = null;
	this.curNumStatements++;
	}

	if (op.requiresCommit() \|\| (curNumStatements >= stmtsPerTx
	&& stmtsPerTx != UNLIMITED_STATEMENTS_PER_TRANSACTION)) {
	LOG.debug("Committing transaction of " + curNumStatements
	+ " statements");
	this.conn.commit();
	this.curNumStatements = 0;
	}
	} catch (BatchUpdateException batchE) {
	if (batchE.getNextException() != null) {
	// if a statement in a batch causes an SQLException
	// the database can either set it as the cause of
	// the BatchUpdateException, or set it as the 'next'
	// field of the BatchUpdateException (e.g. HSQLDB 1.8
	// does the former and Postgres 8.4 does the latter).
	// We'll check for this SQLException in both places,
	// and use the 'next' one in preference.
	setLastError(batchE.getNextException());
	} else {
	// same as SQLException block
	setLastError(batchE);
	}
	} catch (SQLException sqlE) {
	setLastError(sqlE);
	} finally {
	// Close the statement on our way out if that didn't happen
	// via the normal execution path.
	if (null != stmt) {
	try {
	stmt.close();
	} catch (SQLException sqlE) {
	setLastError(sqlE);
	}
	}

	// Always check whether we should end the loop, regardless
	// of the presence of an exception.
	if (op.stop()) {
	return;
	}
	} // try .. catch .. finally.
	} // synchronized (conn)
	}
	}

	/**
	* Allows a user to enqueue the next database operation to run.
	* Since the connection can only execute a single operation at a time,
	* the put() method may block if another operation is already underway.
	* @param op the database operation to perform.
	*/
	public void put(AsyncDBOperation op) throws InterruptedException {
	opsQueue.put(op);
	}

	/**
	* If a previously-executed statement resulted in an error, post it here.
	* If the error slot was already filled, then subsequent errors are
	* squashed until the user calls this method (which clears the error
	* slot).
	* @return any SQLException that occurred due to a previously-run
	* statement.
	*/
	public synchronized SQLException getLastError() {
	SQLException e = this.err;
	this.err = null;
	return e;
	}

	private synchronized void setLastError(SQLException e) {
	if (this.err == null) {
	// Just set it.
	LOG.error("Got exception in update thread: "
	+ StringUtils.stringifyException(e));
	this.err = e;
	} else {
	// Slot is full. Log it and discard.
	LOG.error("SQLException in update thread but error slot full: "
	+ StringUtils.stringifyException(e));
	}
	}
	}
	}