Add the ability to cleanup files after CqlBulkRecordWriter succeeds

Patch by Paul Pak, reviewed by Piotr Kołaczkowski for CASSANDRA-7777
This commit is contained in:
Brandon Williams 2014-10-14 15:42:43 -05:00
parent 781018cb2b
commit 6c0ee30ea7
2 changed files with 23 additions and 1 deletions

View File

@ -53,6 +53,7 @@ public class CqlBulkOutputFormat extends AbstractBulkOutputFormat<Object, List<B
private static final String OUTPUT_CQL_SCHEMA_PREFIX = "cassandra.columnfamily.schema.";
private static final String OUTPUT_CQL_INSERT_PREFIX = "cassandra.columnfamily.insert.";
private static final String DELETE_SOURCE = "cassandra.output.delete.source";
/** Fills the deprecated OutputFormat interface for streaming. */
@Deprecated
@ -103,4 +104,14 @@ public class CqlBulkOutputFormat extends AbstractBulkOutputFormat<Object, List<B
}
return insert;
}
public static void setDeleteSourceOnSuccess(Configuration conf, boolean deleteSrc)
{
conf.setBoolean(DELETE_SOURCE, deleteSrc);
}
public static boolean getDeleteSourceOnSuccess(Configuration conf)
{
return conf.getBoolean(DELETE_SOURCE, false);
}
}

View File

@ -32,6 +32,8 @@ import org.apache.cassandra.hadoop.ConfigHelper;
import org.apache.cassandra.hadoop.HadoopCompat;
import org.apache.cassandra.io.sstable.CQLSSTableWriter;
import org.apache.cassandra.io.sstable.SSTableLoader;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.streaming.StreamState;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.util.Progressable;
@ -57,6 +59,7 @@ public class CqlBulkRecordWriter extends AbstractBulkRecordWriter<Object, List<B
private String schema;
private String insertStatement;
private File outputDir;
private boolean deleteSrc;
CqlBulkRecordWriter(TaskAttemptContext context) throws IOException
{
@ -84,6 +87,7 @@ public class CqlBulkRecordWriter extends AbstractBulkRecordWriter<Object, List<B
schema = CqlBulkOutputFormat.getColumnFamilySchema(conf, columnFamily);
insertStatement = CqlBulkOutputFormat.getColumnFamilyInsertStatement(conf, columnFamily);
outputDir = getColumnFamilyDirectory();
deleteSrc = CqlBulkOutputFormat.getDeleteSourceOnSuccess(conf);
}
@ -107,7 +111,14 @@ public class CqlBulkRecordWriter extends AbstractBulkRecordWriter<Object, List<B
externalClient.addKnownCfs(keyspace, schema);
this.loader = new SSTableLoader(outputDir, externalClient, new BulkRecordWriter.NullOutputHandler());
this.loader = new SSTableLoader(outputDir, externalClient, new BulkRecordWriter.NullOutputHandler()) {
@Override
public void onSuccess(StreamState finalState)
{
if (deleteSrc)
FileUtils.deleteRecursive(outputDir);
}
};
}
}
catch (Exception e)