lang/java/tools/src/main/java/org/apache/avro/tool/DataFileRepairTool.java - avro - 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.avro.tool;

 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.PrintStream;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;

 import org.apache.avro.Schema;
 import org.apache.avro.file.CodecFactory;
 import org.apache.avro.file.DataFileConstants;
 import org.apache.avro.file.DataFileReader;
 import org.apache.avro.file.DataFileWriter;
 import org.apache.avro.generic.GenericDatumReader;
 import org.apache.avro.generic.GenericDatumWriter;

 /** Recovers data from a corrupt Avro Data file */
 public class DataFileRepairTool implements Tool {

   @Override
   public String getName() {
     return "repair";
   }

   @Override
   public String getShortDescription() {
     return "Recovers data from a corrupt Avro Data file";
   }

   private void printInfo(PrintStream output) {
     output.println("Insufficient arguments.  Arguments:  [-o option] "
         + "input_file output_file \n"
         + "   Where option is one of the following: \n"
         + "      " + ALL
         + " (default) recover as many records as possible.\n"
         + "      " + PRIOR
         + "         recover only records prior to the first instance"
         + " of corruption \n"
         + "      " + AFTER
         + "         recover only records after the first instance of"
         + " corruption.\n"
         + "      " + REPORT
         + "        print the corruption report only, reporting the\n"
         + "                    number of valid and corrupted blocks and records\n"
         + "   input_file is the file to read from.  output_file is the file to\n"
         + "   create and write recovered data to.  output_file is ignored if\n"
         + "   using the report option.");
   }

   private static final Set<String> OPTIONS = new HashSet<String>();
   private static final String ALL = "all";
   private static final String PRIOR = "prior";
   private static final String AFTER = "after";
   private static final String REPORT = "report";
   static {
     OPTIONS.add(ALL);
     OPTIONS.add(PRIOR);
     OPTIONS.add(AFTER);
     OPTIONS.add(REPORT);
   }

   @Override
   public int run(InputStream stdin, PrintStream out, PrintStream err,
       List<String> args) throws Exception {
     if (args.size() < 2) {
       printInfo(err);
       return 1;
     }
     int index = 0;
     String input = args.get(index);
     String option = "all";
     if ("-o".equals(input)) {
       option = args.get(1);
       index += 2;
     }
     if (!OPTIONS.contains(option) || (args.size() - index < 1)) {
       printInfo(err);
       return 1;
     }
     input = args.get(index++);
     if (!REPORT.equals(option)) {
       if (args.size() - index < 1) {
         printInfo(err);
         return 1;
       }
     }
     if (ALL.equals(option)) {
       return recoverAll(input, args.get(index), out, err);
     } else if (PRIOR.equals(option)) {
       return recoverPrior(input, args.get(index), out, err);
     } else if (AFTER.equals(option)) {
       return recoverAfter(input, args.get(index), out, err);
     } else if (REPORT.equals(option)) {
       return reportOnly(input, out, err);
     } else {
       return 1;
     }
   }

   private int recover(String input, String output, PrintStream out,
       PrintStream err, boolean recoverPrior, boolean recoverAfter)
       throws IOException {
     File infile = new File(input);
     if (!infile.canRead()) {
       err.println("cannot read file: " + input);
       return 1;
     }
     out.println("Recovering file: " + input);
     GenericDatumReader<Object> reader = new GenericDatumReader<Object>();
     DataFileReader<Object> fileReader = new DataFileReader<Object>(infile,
         reader);
     try {
       Schema schema = fileReader.getSchema();
       String codecStr = fileReader.getMetaString(DataFileConstants.CODEC);
       CodecFactory codecFactory = CodecFactory.fromString("" + codecStr);
       List<String> metas = fileReader.getMetaKeys();
       if (recoverPrior || recoverAfter) {
         GenericDatumWriter<Object> writer = new GenericDatumWriter<Object>();
         DataFileWriter<Object> fileWriter = new DataFileWriter<Object>(writer);
         try {
           File outfile = new File(output);
           for (String key : metas) {
             if (!key.startsWith("avro.")) {
               byte[] val = fileReader.getMeta(key);
               fileWriter.setMeta(key, val);
             }
           }
           fileWriter.setCodec(codecFactory);
           int result = innerRecover(fileReader, fileWriter, out, err, recoverPrior,
               recoverAfter, schema, outfile);
           return result;
         } catch (Exception e) {
           e.printStackTrace(err);
           return 1;
         }
       } else {
         return innerRecover(fileReader, null, out, err, recoverPrior,
             recoverAfter, null, null);
       }

     } finally {
       fileReader.close();
     }
   }

   private int innerRecover(DataFileReader<Object> fileReader,
       DataFileWriter<Object> fileWriter, PrintStream out, PrintStream err,
       boolean recoverPrior, boolean recoverAfter, Schema schema, File outfile) {
     int numBlocks = 0;
     int numCorruptBlocks = 0;
     int numRecords = 0;
     int numCorruptRecords = 0;
     int recordsWritten = 0;
     long position = fileReader.previousSync();
     long blockSize = 0;
     long blockCount = 0;
     boolean fileWritten = false;
     try {
       while (true) {
         try {
           if (!fileReader.hasNext()) {
             out.println("File Summary: ");
             out.println("  Number of blocks: " + numBlocks
                 + " Number of corrupt blocks: " + numCorruptBlocks);
             out.println("  Number of records: " + numRecords
                 + " Number of corrupt records: " + numCorruptRecords);
             if (recoverAfter || recoverPrior) {
               out.println("  Number of records written " + recordsWritten);
             }
             out.println();
             return 0;
           }
           position = fileReader.previousSync();
           blockCount = fileReader.getBlockCount();
           blockSize = fileReader.getBlockSize();
           numRecords += blockCount;
           long blockRemaining = blockCount;
           numBlocks++;
           boolean lastRecordWasBad = false;
           long badRecordsInBlock = 0;
           while (blockRemaining > 0) {
             try {
               Object datum = fileReader.next();
               if ((recoverPrior && numCorruptBlocks == 0)
                   || (recoverAfter && numCorruptBlocks > 0)) {
                 if (!fileWritten) {
                   try {
                     fileWriter.create(schema, outfile);
                     fileWritten = true;
                   } catch (Exception e) {
                     e.printStackTrace(err);
                     return 1;
                   }
                 }
                 try {
                   fileWriter.append(datum);
                   recordsWritten++;
                 } catch (Exception e) {
                   e.printStackTrace(err);
                   throw e;
                 }
               }
               blockRemaining--;
               lastRecordWasBad = false;
             } catch (Exception e) {
               long pos = blockCount - blockRemaining;
               if (badRecordsInBlock == 0) {
                 // first corrupt record
                 numCorruptBlocks++;
                 err.println("Corrupt block: " + numBlocks + " Records in block: "
                     + blockCount + " uncompressed block size: " + blockSize);
                 err.println("Corrupt record at position: "
                     + (pos));
               } else {
                 // second bad record in block, if consecutive skip block.
                 err.println("Corrupt record at position: "
                     + (pos));
                 if (lastRecordWasBad) {
                   // consecutive bad record
                   err.println("Second consecutive bad record in block: " + numBlocks
                       + ". Skipping remainder of block. ");
                   numCorruptRecords += blockRemaining;
                   badRecordsInBlock += blockRemaining;
                   try {
                     fileReader.sync(position);
                   } catch (Exception e2) {
                     err.println("failed to sync to sync marker, aborting");
                     e2.printStackTrace(err);
                     return 1;
                   }
                   break;
                 }
               }
               blockRemaining --;
               lastRecordWasBad = true;
               numCorruptRecords++;
               badRecordsInBlock++;
             }
           }
           if (badRecordsInBlock != 0) {
             err.println("** Number of unrecoverable records in block: "
                 + (badRecordsInBlock));
           }
           position = fileReader.previousSync();
         } catch (Exception e) {
           err.println("Failed to read block " + numBlocks + ". Unknown record "
               + "count in block.  Skipping. Reason: " + e.getMessage());
           numCorruptBlocks++;
           try {
             fileReader.sync(position);
           } catch (Exception e2) {
             err.println("failed to sync to sync marker, aborting");
             e2.printStackTrace(err);
             return 1;
           }
         }
       }
     } finally {
       if (fileWritten) {
         try {
           fileWriter.close();
         } catch (Exception e) {
           e.printStackTrace(err);
           return 1;
         }
       }
     }
   }

   private int reportOnly(String input, PrintStream out, PrintStream err)
       throws IOException {
     return recover(input, null, out, err, false, false);
   }

   private int recoverAfter(String input, String output, PrintStream out,
       PrintStream err) throws IOException {
     return recover(input, output, out, err, false, true);
   }

   private int recoverPrior(String input, String output, PrintStream out,
       PrintStream err) throws IOException {
     return recover(input, output, out, err, true, false);
   }

   private int recoverAll(String input, String output, PrintStream out,
       PrintStream err) throws IOException {
     return recover(input, output, out, err, true, true);
   }
 }
	/**
	* 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.avro.tool;

	import java.io.File;
	import java.io.IOException;
	import java.io.InputStream;
	import java.io.PrintStream;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Set;

	import org.apache.avro.Schema;
	import org.apache.avro.file.CodecFactory;
	import org.apache.avro.file.DataFileConstants;
	import org.apache.avro.file.DataFileReader;
	import org.apache.avro.file.DataFileWriter;
	import org.apache.avro.generic.GenericDatumReader;
	import org.apache.avro.generic.GenericDatumWriter;

	/** Recovers data from a corrupt Avro Data file */
	public class DataFileRepairTool implements Tool {

	@Override
	public String getName() {
	return "repair";
	}

	@Override
	public String getShortDescription() {
	return "Recovers data from a corrupt Avro Data file";
	}

	private void printInfo(PrintStream output) {
	output.println("Insufficient arguments. Arguments: [-o option] "
	+ "input_file output_file \n"
	+ " Where option is one of the following: \n"
	+ " " + ALL
	+ " (default) recover as many records as possible.\n"
	+ " " + PRIOR
	+ " recover only records prior to the first instance"
	+ " of corruption \n"
	+ " " + AFTER
	+ " recover only records after the first instance of"
	+ " corruption.\n"
	+ " " + REPORT
	+ " print the corruption report only, reporting the\n"
	+ " number of valid and corrupted blocks and records\n"
	+ " input_file is the file to read from. output_file is the file to\n"
	+ " create and write recovered data to. output_file is ignored if\n"
	+ " using the report option.");
	}

	private static final Set<String> OPTIONS = new HashSet<String>();
	private static final String ALL = "all";
	private static final String PRIOR = "prior";
	private static final String AFTER = "after";
	private static final String REPORT = "report";
	static {
	OPTIONS.add(ALL);
	OPTIONS.add(PRIOR);
	OPTIONS.add(AFTER);
	OPTIONS.add(REPORT);
	}

	@Override
	public int run(InputStream stdin, PrintStream out, PrintStream err,
	List<String> args) throws Exception {
	if (args.size() < 2) {
	printInfo(err);
	return 1;
	}
	int index = 0;
	String input = args.get(index);
	String option = "all";
	if ("-o".equals(input)) {
	option = args.get(1);
	index += 2;
	}
	if (!OPTIONS.contains(option) \|\| (args.size() - index < 1)) {
	printInfo(err);
	return 1;
	}
	input = args.get(index++);
	if (!REPORT.equals(option)) {
	if (args.size() - index < 1) {
	printInfo(err);
	return 1;
	}
	}
	if (ALL.equals(option)) {
	return recoverAll(input, args.get(index), out, err);
	} else if (PRIOR.equals(option)) {
	return recoverPrior(input, args.get(index), out, err);
	} else if (AFTER.equals(option)) {
	return recoverAfter(input, args.get(index), out, err);
	} else if (REPORT.equals(option)) {
	return reportOnly(input, out, err);
	} else {
	return 1;
	}
	}

	private int recover(String input, String output, PrintStream out,
	PrintStream err, boolean recoverPrior, boolean recoverAfter)
	throws IOException {
	File infile = new File(input);
	if (!infile.canRead()) {
	err.println("cannot read file: " + input);
	return 1;
	}
	out.println("Recovering file: " + input);
	GenericDatumReader<Object> reader = new GenericDatumReader<Object>();
	DataFileReader<Object> fileReader = new DataFileReader<Object>(infile,
	reader);
	try {
	Schema schema = fileReader.getSchema();
	String codecStr = fileReader.getMetaString(DataFileConstants.CODEC);
	CodecFactory codecFactory = CodecFactory.fromString("" + codecStr);
	List<String> metas = fileReader.getMetaKeys();
	if (recoverPrior \|\| recoverAfter) {
	GenericDatumWriter<Object> writer = new GenericDatumWriter<Object>();
	DataFileWriter<Object> fileWriter = new DataFileWriter<Object>(writer);
	try {
	File outfile = new File(output);
	for (String key : metas) {
	if (!key.startsWith("avro.")) {
	byte[] val = fileReader.getMeta(key);
	fileWriter.setMeta(key, val);
	}
	}
	fileWriter.setCodec(codecFactory);
	int result = innerRecover(fileReader, fileWriter, out, err, recoverPrior,
	recoverAfter, schema, outfile);
	return result;
	} catch (Exception e) {
	e.printStackTrace(err);
	return 1;
	}
	} else {
	return innerRecover(fileReader, null, out, err, recoverPrior,
	recoverAfter, null, null);
	}

	} finally {
	fileReader.close();
	}
	}

	private int innerRecover(DataFileReader<Object> fileReader,
	DataFileWriter<Object> fileWriter, PrintStream out, PrintStream err,
	boolean recoverPrior, boolean recoverAfter, Schema schema, File outfile) {
	int numBlocks = 0;
	int numCorruptBlocks = 0;
	int numRecords = 0;
	int numCorruptRecords = 0;
	int recordsWritten = 0;
	long position = fileReader.previousSync();
	long blockSize = 0;
	long blockCount = 0;
	boolean fileWritten = false;
	try {
	while (true) {
	try {
	if (!fileReader.hasNext()) {
	out.println("File Summary: ");
	out.println(" Number of blocks: " + numBlocks
	+ " Number of corrupt blocks: " + numCorruptBlocks);
	out.println(" Number of records: " + numRecords
	+ " Number of corrupt records: " + numCorruptRecords);
	if (recoverAfter \|\| recoverPrior) {
	out.println(" Number of records written " + recordsWritten);
	}
	out.println();
	return 0;
	}
	position = fileReader.previousSync();
	blockCount = fileReader.getBlockCount();
	blockSize = fileReader.getBlockSize();
	numRecords += blockCount;
	long blockRemaining = blockCount;
	numBlocks++;
	boolean lastRecordWasBad = false;
	long badRecordsInBlock = 0;
	while (blockRemaining > 0) {
	try {
	Object datum = fileReader.next();
	if ((recoverPrior && numCorruptBlocks == 0)
	\|\| (recoverAfter && numCorruptBlocks > 0)) {
	if (!fileWritten) {
	try {
	fileWriter.create(schema, outfile);
	fileWritten = true;
	} catch (Exception e) {
	e.printStackTrace(err);
	return 1;
	}
	}
	try {
	fileWriter.append(datum);
	recordsWritten++;
	} catch (Exception e) {
	e.printStackTrace(err);
	throw e;
	}
	}
	blockRemaining--;
	lastRecordWasBad = false;
	} catch (Exception e) {
	long pos = blockCount - blockRemaining;
	if (badRecordsInBlock == 0) {
	// first corrupt record
	numCorruptBlocks++;
	err.println("Corrupt block: " + numBlocks + " Records in block: "
	+ blockCount + " uncompressed block size: " + blockSize);
	err.println("Corrupt record at position: "
	+ (pos));
	} else {
	// second bad record in block, if consecutive skip block.
	err.println("Corrupt record at position: "
	+ (pos));
	if (lastRecordWasBad) {
	// consecutive bad record
	err.println("Second consecutive bad record in block: " + numBlocks
	+ ". Skipping remainder of block. ");
	numCorruptRecords += blockRemaining;
	badRecordsInBlock += blockRemaining;
	try {
	fileReader.sync(position);
	} catch (Exception e2) {
	err.println("failed to sync to sync marker, aborting");
	e2.printStackTrace(err);
	return 1;
	}
	break;
	}
	}
	blockRemaining --;
	lastRecordWasBad = true;
	numCorruptRecords++;
	badRecordsInBlock++;
	}
	}
	if (badRecordsInBlock != 0) {
	err.println("** Number of unrecoverable records in block: "
	+ (badRecordsInBlock));
	}
	position = fileReader.previousSync();
	} catch (Exception e) {
	err.println("Failed to read block " + numBlocks + ". Unknown record "
	+ "count in block. Skipping. Reason: " + e.getMessage());
	numCorruptBlocks++;
	try {
	fileReader.sync(position);
	} catch (Exception e2) {
	err.println("failed to sync to sync marker, aborting");
	e2.printStackTrace(err);
	return 1;
	}
	}
	}
	} finally {
	if (fileWritten) {
	try {
	fileWriter.close();
	} catch (Exception e) {
	e.printStackTrace(err);
	return 1;
	}
	}
	}
	}

	private int reportOnly(String input, PrintStream out, PrintStream err)
	throws IOException {
	return recover(input, null, out, err, false, false);
	}

	private int recoverAfter(String input, String output, PrintStream out,
	PrintStream err) throws IOException {
	return recover(input, output, out, err, false, true);
	}

	private int recoverPrior(String input, String output, PrintStream out,
	PrintStream err) throws IOException {
	return recover(input, output, out, err, true, false);
	}

	private int recoverAll(String input, String output, PrintStream out,
	PrintStream err) throws IOException {
	return recover(input, output, out, err, true, true);
	}
	}