package com.linkedin.uif.writer;

import java.io.IOException;
import java.net.URI;
import java.util.concurrent.atomic.AtomicLong;

import com.google.common.base.Preconditions;
import org.apache.avro.Schema;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;

import com.linkedin.uif.converter.DataConversionException;
import com.linkedin.uif.writer.converter.DataConverter;

/**
 * An implementation of {@link DataWriter} that writes directly
 * to HDFS in Avro format.
 *
 * @param <S> type of source data record representation
 *
 * @author ynli
 */
class AvroHdfsDataWriter<S> implements DataWriter<S, GenericRecord> {

    private static final Logger LOG = LoggerFactory.getLogger(AvroHdfsDataWriter.class);

    private final FileSystem fs;
    private final Path stagingFile;
    private final Path outputFile;
    private final DataConverter<S, GenericRecord> dataConverter;
    private final DataFileWriter<GenericRecord> writer;

    // Number of records successfully written
    private final AtomicLong count = new AtomicLong(0);

    public AvroHdfsDataWriter(URI uri, String stagingDir, String outputDir,
            String fileName, int bufferSize, DataConverter<S,
            GenericRecord> dataConverter, Schema schema) throws IOException {

        Configuration conf = new Configuration();
        this.fs = FileSystem.get(uri, conf);
        this.stagingFile = new Path(stagingDir, fileName);
        this.outputFile = new Path(outputDir, fileName);
        // Create the parent directory of the output file if it does not exist
        if (!this.fs.exists(this.outputFile.getParent())) {
            this.fs.mkdirs(this.outputFile.getParent());
        }
        this.dataConverter = dataConverter;
        this.writer = createDatumWriter(schema, this.stagingFile, bufferSize);
    }

    @Override
    public void write(S sourceRecord) throws IOException {
        Preconditions.checkNotNull(sourceRecord);

        try {
            this.writer.append(this.dataConverter.convert(sourceRecord));
        } catch (DataConversionException e) {
            LOG.error("Failed to convert and write source data record: " + sourceRecord);
            throw new IOException(e);
        }

        // Only increment when write is successful
        this.count.incrementAndGet();
    }

    @Override
    public void close() throws IOException {
        this.writer.flush();
        this.writer.close();
    }

    @Override
    public void commit() throws IOException {
        if (!this.fs.exists(this.stagingFile)) {
            throw new IOException(String.format("File %s does not exist", this.stagingFile));
        }

        LOG.info(String.format("Moving data from %s to %s", this.stagingFile, this.outputFile));
        if (this.fs.exists(this.outputFile)) {
            throw new IOException(String.format("File %s already exists", this.outputFile));
        }
        this.fs.rename(this.stagingFile, this.outputFile);
    }

    @Override
    public void cleanup() throws IOException {
        // Delete the staging file
        if (this.fs.exists(this.stagingFile)) {
            this.fs.delete(this.stagingFile, false);
        }
    }

    @Override
    public long recordsWritten() {
        return this.count.get();
    }

    @Override
    public long bytesWritten() throws IOException {
        if (!this.fs.exists(this.outputFile)) {
            return 0;
        }

        return this.fs.getFileStatus(this.outputFile).getLen();
    }

    /**
     * Create a new {@link DataFileWriter} for writing Avro records.
     *
     * @param schema Avro schema
     * @param avroFile Avro file to write to
     * @throws IOException
     */
    private DataFileWriter<GenericRecord> createDatumWriter(Schema schema,
            Path avroFile, int bufferSize) throws IOException {

        if (this.fs.exists(avroFile)) {
            throw new IOException(String.format("File %s already exists", avroFile));
        }

        FSDataOutputStream outputStream = this.fs.create(avroFile, true, bufferSize);
        return new DataFileWriter<GenericRecord>(new GenericDatumWriter<GenericRecord>())
                .create(schema, outputStream);
    }
}

@@ -128,53 +128,8 @@
 g;%0A%0A
-import com.google.common.base.Preconditions;%0A
 impo
@@ -300,65 +300,8 @@
 rd;%0A
-import org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0A
 impo
@@ -462,16 +462,119 @@
 fs.Path;
+%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0A%0Aimport com.google.common.base.Preconditions;
 %0A%0Aimport
@@ -1662,16 +1662,17 @@
  conf);%0A
+%0A
         
@@ -1714,32 +1714,539 @@
 Dir, fileName);%0A
+        // Deleting the staging file if it already exists, which can happen if the%0A        // task failed and the staging file didn't get cleaned up for some reason.%0A        // Deleting the staging file prevents the task retry from being blocked.%0A        if (this.fs.exists(this.stagingFile)) %7B%0A            LOG.warn(String.format(%0A                    %22Task staging file %25s already exists, deleting it%22,%0A                    this.stagingFile));%0A            this.fs.delete(this.stagingFile, false);%0A        %7D%0A%0A
         this.out
@@ -2484,32 +2484,33 @@
 t());%0A        %7D%0A
+%0A
         this.dat
@@ -3543,32 +3543,211 @@
 s.outputFile));%0A
+        // For the same reason as deleting the staging file if it already exists, deleting%0A        // the output file if it already exists prevents task retry from being blocked.%0A
         if (this
@@ -3789,36 +3789,23 @@
         
-throw new IOExceptio
+LOG.war
 n(String
@@ -3805,33 +3805,45 @@
 (String.format(%22
-F
+Task output f
 ile %25s already e
@@ -3861,32 +3861,84 @@
 s.outputFile));%0A
+            this.fs.delete(this.outputFile, false);%0A
         %7D%0A      

Delete task staging/output file that already exists in the writer