package org.voltdb.sysprocs;
import java.io.File;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.apache.log4j.Logger;
import org.voltdb.DependencySet;
import org.voltdb.ParameterSet;
import org.voltdb.ProcInfo;
import org.voltdb.VoltSystemProcedure;
import org.voltdb.VoltTable;
import org.voltdb.VoltType;
import org.voltdb.catalog.Table;
import au.com.bytecode.opencsv.CSVWriter;
import edu.brown.hstore.PartitionExecutor.SystemProcedureExecutionContext;
@ProcInfo(singlePartition = false)
public class DatabaseDump extends VoltSystemProcedure {
private static final Logger LOG = Logger.getLogger(DatabaseDump.class);
@Override
public void initImpl() {
executor.registerPlanFragment(SysProcFragmentId.PF_dumpDistribute, this);
executor.registerPlanFragment(SysProcFragmentId.PF_dumpAggregate, this);
}
@Override
public DependencySet executePlanFragment(Long txn_id, Map<Integer, List<VoltTable>> dependencies, int fragmentId, ParameterSet params, SystemProcedureExecutionContext context) {
final boolean debug = LOG.isDebugEnabled();
// need to return something ..
VoltTable[] result = new VoltTable[1];
result[0] = new VoltTable(new VoltTable.ColumnInfo("TxnId", VoltType.BIGINT));
result[0].addRow(txn_id);
if (fragmentId == SysProcFragmentId.PF_dumpDistribute) {
// The first parameter should be the directory
File directory = new File(params.toArray()[0].toString());
LOG.debug(String.format("Dumping out database contents for partition %d to '%s'", this.executor.getPartitionId(), directory));
if (directory.exists() == false) {
LOG.debug("Creating dump directory '" + directory + "'");
directory.mkdirs();
}
assert(directory.isDirectory());
int batch_size = 1000;
// Only write column names if it's the first partition (hack)
boolean write_header = (this.executor.getPartitionId() == 0);
for (Table catalog_tbl : catalogContext.database.getTables().values()) {
File csv_file = new File(String.format("%s/%s.%02d.csv", directory, catalog_tbl.getName(), this.executor.getPartitionId()));
CSVWriter writer = null;
try {
writer = new CSVWriter(new FileWriter(csv_file));
} catch (Exception ex) {
LOG.fatal(String.format("Failed to create CSVWriter for '%s'", csv_file), ex);
throw new RuntimeException(ex);
}
if (write_header) {
String cols[] = new String[catalog_tbl.getColumns().size()];
for (int i = 0; i < cols.length; i++) {
cols[i] = catalog_tbl.getColumns().get(i).getName();
} // FOR
writer.writeNext(cols);
}
int total = 0;
while (true) {
LOG.debug(String.format("%s: offset=%d, limit=%d", catalog_tbl.getName(), total, batch_size));
VoltTable vt = this.executor.getExecutionEngine().serializeTable(catalog_tbl.getRelativeIndex());
assert(vt != null) : "Failed to get serialized table for " + catalog_tbl;
if (vt.getRowCount() == 0) break;
total += vt.getRowCount();
LOG.debug(String.format("Writing %d / %d tuples to '%s'", vt.getRowCount(), total, csv_file.getName()));
// Dump table contents
while (vt.advanceRow()) {
String row[] = vt.getRowStringArray();
assert(row != null);
assert(row.length == vt.getColumnCount());
writer.writeNext(row);
} // WHILE
} // WHILE
try {
writer.close();
} catch (Exception ex) {
throw new RuntimeException(ex);
}
} // FOR
return new DependencySet(new int[] { (int)SysProcFragmentId.PF_dumpDistribute }, result);
} else if (fragmentId == SysProcFragmentId.PF_dumpAggregate) {
if (debug) LOG.debug("Aggregating results from loading fragments in txn #" + txn_id);
return new DependencySet(new int[] { (int)SysProcFragmentId.PF_dumpAggregate }, result);
}
assert(false) : "Unexpected FragmentId " + fragmentId;
return null;
}
/**
*
* @param tableName
* @param table
* @return
* @throws VoltAbortException
*/
public VoltTable[] run(String directory) throws VoltAbortException {
final boolean debug = LOG.isDebugEnabled();
final ParameterSet params = new ParameterSet(directory);
// Generate a plan fragment for each site using the sub-tables
final List<SynthesizedPlanFragment> pfs = new ArrayList<SynthesizedPlanFragment>();
for (int i = 0; i < catalogContext.numberOfPartitions; i++) {
int partition = i;
SynthesizedPlanFragment pf = new SynthesizedPlanFragment();
pf.fragmentId = SysProcFragmentId.PF_dumpDistribute;
pf.inputDependencyIds = new int[] { };
pf.outputDependencyIds = new int[] { (int)SysProcFragmentId.PF_dumpDistribute };
pf.multipartition = false;
pf.nonExecSites = false;
pf.destPartitionId = partition;
pf.parameters = params;
pf.last_task = true;
pfs.add(pf);
} // FOR
// a final plan fragment to aggregate the results
// pfs[0] = new SynthesizedPlanFragment();
// pfs[0].destPartitionId = partitionId;
// pfs[0].fragmentId = SysProcFragmentId.PF_dumpAggregate;
// pfs[0].inputDependencyIds = new int[] { (int)SysProcFragmentId.PF_dumpDistribute };
// pfs[0].outputDependencyIds = new int[] { (int)SysProcFragmentId.PF_dumpAggregate };
// pfs[0].multipartition = false;
// pfs[0].nonExecSites = false;
// pfs[0].parameters = new ParameterSet();
// send these forth in to the world .. and wait
if (debug) LOG.debug("Passing " + pfs.size() + " sysproc fragments to executeSysProcPlanFragments()");
VoltTable[] results = executeSysProcPlanFragments(pfs.toArray(new SynthesizedPlanFragment[0]), (int)SysProcFragmentId.PF_dumpDistribute);
return results;
}
}