class PutHiveStreaming extends AbstractProcessor {

    @Override
    public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
        FlowFile flowFile = session.get();
        if (flowFile == null) {
            return;
        }

        final ComponentLog log = getLogger();
        final Integer recordsPerTxn = context.getProperty(RECORDS_PER_TXN).asInteger();

        // Store the original class loader, then explicitly set it to this class's classloader (for use by the Hive Metastore)
        ClassLoader originalClassloader = Thread.currentThread().getContextClassLoader();
        Thread.currentThread().setContextClassLoader(this.getClass().getClassLoader());

        final List<String> partitionColumnList;
        final String partitionColumns = context.getProperty(PARTITION_COLUMNS).getValue();
        if (partitionColumns == null || partitionColumns.isEmpty()) {
            partitionColumnList = Collections.emptyList();
        } else {
            String[] partitionCols = partitionColumns.split(",");
            partitionColumnList = new ArrayList<>(partitionCols.length);
            for (String col : partitionCols) {
                partitionColumnList.add(col.trim());
            }
        }

        final AtomicInteger recordCount = new AtomicInteger(0);
        final AtomicInteger successfulRecordCount = new AtomicInteger(0);
        List<HiveStreamingRecord> successfulRecords = new LinkedList<>();
        final FlowFile inputFlowFile = flowFile;
        final AtomicBoolean processingFailure = new AtomicBoolean(false);

        // Create output flow files and their Avro writers
        AtomicReference<FlowFile> successFlowFile = new AtomicReference<>(session.create(inputFlowFile));
        final DataFileWriter<GenericRecord> successAvroWriter = new DataFileWriter<>(new GenericDatumWriter<GenericRecord>());
        AtomicReference<FlowFile> failureFlowFile = new AtomicReference<>(session.create(inputFlowFile));
        final DataFileWriter<GenericRecord> failureAvroWriter = new DataFileWriter<>(new GenericDatumWriter<GenericRecord>());

        try {
            session.read(inputFlowFile, in -> {

                try (final DataFileStream<GenericRecord> reader = new DataFileStream<>(in, new GenericDatumReader<GenericRecord>())) {
                    GenericRecord currRecord = null;

                    // Copy codec and schema information to all writers
                    final String codec = reader.getMetaString(DataFileConstants.CODEC) == null
                            ? DataFileConstants.NULL_CODEC
                            : reader.getMetaString(DataFileConstants.CODEC);

                    Arrays.asList(successAvroWriter, failureAvroWriter)
                            .forEach((writer) -> {
                                writer.setCodec(CodecFactory.fromString(codec));
                                // Transfer metadata (this is a subset of the incoming file)
                                for (String metaKey : reader.getMetaKeys()) {
                                    if (!RESERVED_METADATA.contains(metaKey)) {
                                        writer.setMeta(metaKey, reader.getMeta(metaKey));
                                    }
                                }
                            });

                    while (reader.hasNext()) {
                        currRecord = reader.next(currRecord);
                        recordCount.incrementAndGet();

                        // Extract the partition values (they must be put separately into the Hive Streaming API)
                        List<String> partitionValues = new ArrayList<>();

                        try {
                            for (String partition : partitionColumnList) {
                                Object partitionValue = currRecord.get(partition);
                                if (partitionValue == null) {
                                    throw new IOException("Partition column '" + partition + "' not found in Avro record");
                                }
                                partitionValues.add(partitionValue.toString());
                            }
                        } catch (IOException ioe) {
                            // Add the failed record to the failure flow file
                            log.error("Error writing record to Hive Streaming transaction", ioe);
                            appendRecordsToFlowFile(session, Collections.singletonList(new HiveStreamingRecord(null, currRecord)),
                                    failureFlowFile, failureAvroWriter, reader);
                            continue;
                        }

                        List<Schema.Field> fields = currRecord.getSchema().getFields();
                        if (fields != null) {
                            JSONObject obj = new JSONObject();
                            try {
                                for (Schema.Field field : fields) {
                                    String fieldName = field.name();
                                    // Skip fields that are partition columns, we extracted those values above to create an EndPoint
                                    if (!partitionColumnList.contains(fieldName)) {
                                        Object value = currRecord.get(fieldName);
                                        try {
                                            obj.put(fieldName, value);
                                        } catch (JSONException je) {
                                            throw new IOException(je);
                                        }
                                    }
                                }
                            } catch (IOException ioe) {
                                // This really shouldn't happen since we are iterating over the schema fields, but just in case,
                                // add the failed record to the failure flow file.
                                log.error("Error writing record to Hive Streaming transaction", ioe);
                                appendRecordsToFlowFile(session, Collections.singletonList(new HiveStreamingRecord(null, currRecord)),
                                        failureFlowFile, failureAvroWriter, reader);
                                continue;
                            }
                            final HiveStreamingRecord record = new HiveStreamingRecord(partitionValues, currRecord);
                            HiveEndPoint endPoint = null;
                            HiveWriter hiveWriter = null;
                            try {
                                endPoint = makeHiveEndPoint(record.getPartitionValues(), options);
                                hiveWriter = getOrCreateWriter(endPoint);
                            } catch (ConnectionError
                                    | HiveWriter.ConnectFailure
                                    | InterruptedException connectionError) {
                                // Can't connect to Hive endpoint.
                                log.error("Error connecting to Hive endpoint: table {} at {}",
                                        new Object[]{options.getTableName(), options.getMetaStoreURI()});
                                // If we can't connect to the endpoint, exit the loop and let the outer exception handler route the original flow file to retry
                                abortAndCloseWriters();
                                throw new ProcessException(connectionError);
                            }
                            try {
                                try {
                                    hiveWriter.write(record.getRecord().toString().getBytes(StandardCharsets.UTF_8));
                                    successfulRecords.add(record);
                                } catch (InterruptedException | HiveWriter.WriteFailure wf) {
                                    // Add the failed record to the failure flow file
                                    log.error("Error writing record to Hive Streaming transaction", wf);
                                    appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader);
                                }

                                // If we've reached the records-per-transaction limit, flush the Hive Writer and update the Avro Writer for successful records
                                if (hiveWriter.getTotalRecords() >= recordsPerTxn) {
                                    hiveWriter.flush(true);
                                    // Now send the records to the success relationship and update the success count
                                    try {
                                        appendRecordsToFlowFile(session, successfulRecords, successFlowFile, successAvroWriter, reader);
                                        successfulRecordCount.accumulateAndGet(successfulRecords.size(), (current, incr) -> current + incr);

                                        // Clear the list of successful records, we'll use it at the end when we flush whatever records are left
                                        successfulRecords.clear();

                                    } catch (IOException ioe) {
                                        // The records were put to Hive Streaming successfully, but there was an error while writing the
                                        // Avro records to the flow file. Log as an error and move on.
                                        getLogger().error("Error writing Avro records (which were sent successfully to Hive Streaming) to the flow file", ioe);
                                    }
                                }

                            } catch (InterruptedException
                                    | HiveWriter.CommitFailure
                                    | HiveWriter.TxnBatchFailure
                                    | HiveWriter.TxnFailure
                                    | SerializationError writeException) {

                                log.error("Error writing record to Hive Streaming transaction", writeException);
                                // Add the failed record to the failure flow file
                                appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader);

                                if (!(writeException instanceof SerializationError)) {
                                    try {
                                        hiveWriter.abort();
                                    } catch (Exception e) {
                                        // Can't even abort properly, throw a process exception
                                        throw new ProcessException(e);
                                    }
                                }
                            }
                        }
                    }
                    try {
                        // Finish any transactions
                        flushAllWriters(true);
                        closeAllWriters();

                        // Now send any remaining records to the success relationship and update the count
                        appendRecordsToFlowFile(session, successfulRecords, successFlowFile, successAvroWriter, reader);
                        successfulRecordCount.accumulateAndGet(successfulRecords.size(), (current, incr) -> current + incr);
                        successfulRecords.clear();

                    } catch (HiveWriter.CommitFailure
                            | HiveWriter.TxnBatchFailure
                            | HiveWriter.TxnFailure
                            | InterruptedException e) {

                        // If any records are in the successfulRecords list but ended up here, then they actually weren't transferred successfully, so
                        // route them to failure instead
                        appendRecordsToFlowFile(session, successfulRecords, failureFlowFile, failureAvroWriter, reader);
                    }
                } catch (IOException ioe) {
                    // The Avro file is invalid (or may not be an Avro file at all), send it to failure
                    log.error("The incoming flow file can not be read as an Avro file, routing to failure", ioe);
                    processingFailure.set(true);
                }
            });

            if (recordCount.get() > 0) {
                if (successfulRecordCount.get() > 0) {
                    // Transfer the flow file with successful records
                    successFlowFile.set(
                            session.putAttribute(successFlowFile.get(), HIVE_STREAMING_RECORD_COUNT_ATTR, Integer.toString(recordCount.get())));
                    session.getProvenanceReporter().send(successFlowFile.get(), options.getMetaStoreURI());
                    session.transfer(successFlowFile.get(), REL_SUCCESS);
                } else {
                    session.remove(successFlowFile.get());
                }

                if (recordCount.get() != successfulRecordCount.get()) {
                    // There were some failed records, so transfer that flow file to failure
                    failureFlowFile.set(
                            session.putAttribute(failureFlowFile.get(), HIVE_STREAMING_RECORD_COUNT_ATTR,
                                    Integer.toString(recordCount.get() - successfulRecordCount.get())));
                    session.transfer(failureFlowFile.get(), REL_FAILURE);
                } else {
                    session.remove(failureFlowFile.get());
                }
            } else {
                // No records were processed, so remove the output flow files
                session.remove(successFlowFile.get());
                session.remove(failureFlowFile.get());
            }
            successFlowFile.set(null);
            failureFlowFile.set(null);

            // If we got here, we've processed the outgoing flow files correctly, so remove the incoming one if necessary
            if (processingFailure.get()) {
                session.transfer(inputFlowFile, REL_FAILURE);
            } else {
                session.remove(flowFile);
            }

        } catch (ProcessException pe) {
            abortAndCloseWriters();
            Throwable t = pe.getCause();
            if (t != null) {
                if (t instanceof ConnectionError
                        || t instanceof HiveWriter.ConnectFailure
                        || t instanceof HiveWriter.CommitFailure
                        || t instanceof HiveWriter.TxnBatchFailure
                        || t instanceof HiveWriter.TxnFailure
                        || t instanceof InterruptedException) {
                    log.error("Hive Streaming connect/write error, flow file will be penalized and routed to retry", t);
                    flowFile = session.penalize(flowFile);
                    session.transfer(flowFile, REL_RETRY);
                    // Remove the ones we created
                    if (successFlowFile.get() != null) {
                        session.remove(successFlowFile.get());
                    }
                    if (failureFlowFile.get() != null) {
                        session.remove(failureFlowFile.get());
                    }
                } else {
                    throw pe;
                }
            } else {
                throw pe;
            }
        } finally {
            // Restore original class loader, might not be necessary but is good practice since the processor task changed it
            Thread.currentThread().setContextClassLoader(originalClassloader);
        }
    }

}