@Test public void onTrigger() throws Exception { runner.setProperty(PutHiveStreaming.METASTORE_URI, "thrift: runner.setProperty(PutHiveStreaming.DB_NAME, "default"); runner.setProperty(PutHiveStreaming.TABLE_NAME, "users"); runner.setProperty(PutHiveStreaming.TXNS_PER_BATCH, "100"); runner.setValidateExpressionUsage(false); Map<String, Object> user1 = new HashMap<String, Object>() { { put("name", "Joe"); put("favorite_number", 146); } }; runner.enqueue(createAvroRecord(Collections.singletonList(user1))); runner.run(); runner.assertTransferCount(PutHiveStreaming.REL_SUCCESS, 1); assertEquals("1", runner.getFlowFilesForRelationship(PutHiveStreaming.REL_SUCCESS).get(0).getAttribute(HIVE_STREAMING_RECORD_COUNT_ATTR)); }
@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(); 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); 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; 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)); 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(); 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) { 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(); if (!partitionColumnList.contains(fieldName)) { Object value = currRecord.get(fieldName); try { obj.put(fieldName, value); } catch (JSONException je) { throw new IOException(je); } } } } catch (IOException ioe) { 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) { log.error("Error connecting to Hive endpoint: table {} at {}", new Object[]{options.getTableName(), options.getMetaStoreURI()}); abortAndCloseWriters(); throw new ProcessException(connectionError); } try { try { hiveWriter.write(record.getRecord().toString().getBytes(StandardCharsets.UTF_8)); successfulRecords.add(record); } catch (InterruptedException | HiveWriter.WriteFailure wf) { log.error("Error writing record to Hive Streaming transaction", wf); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); } if (hiveWriter.getTotalRecords() >= recordsPerTxn) { hiveWriter.flush(true); try { appendRecordsToFlowFile(session, successfulRecords, successFlowFile, successAvroWriter, reader); successfulRecordCount.accumulateAndGet(successfulRecords.size(), (current, incr) -> current + incr); successfulRecords.clear(); } catch (IOException ioe) { 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); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); if (!(writeException instanceof SerializationError)) { try { hiveWriter.abort(); } catch (Exception e) { throw new ProcessException(e); } } } } } try { flushAllWriters(true); closeAllWriters(); 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) { appendRecordsToFlowFile(session, successfulRecords, failureFlowFile, failureAvroWriter, reader); } } catch (IOException ioe) { 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) { 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()) { 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 { session.remove(successFlowFile.get()); session.remove(failureFlowFile.get()); } successFlowFile.set(null); failureFlowFile.set(null); 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); if (successFlowFile.get() != null) { session.remove(successFlowFile.get()); } if (failureFlowFile.get() != null) { session.remove(failureFlowFile.get()); } } else { throw pe; } } else { throw pe; } } finally { Thread.currentThread().setContextClassLoader(originalClassloader); } }
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(); 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); 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; 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)); 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(); 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) { 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(); if (!partitionColumnList.contains(fieldName)) { Object value = currRecord.get(fieldName); try { obj.put(fieldName, value); } catch (JSONException je) { throw new IOException(je); } } } } catch (IOException ioe) { 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) { log.error("Error connecting to Hive endpoint: table {} at {}", new Object[]{options.getTableName(), options.getMetaStoreURI()}); abortAndCloseWriters(); throw new ProcessException(connectionError); } try { try { hiveWriter.write(record.getRecord().toString().getBytes(StandardCharsets.UTF_8)); successfulRecords.add(record); } catch (InterruptedException | HiveWriter.WriteFailure wf) { log.error("Error writing record to Hive Streaming transaction", wf); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); } if (hiveWriter.getTotalRecords() >= recordsPerTxn) { hiveWriter.flush(true); try { appendRecordsToFlowFile(session, successfulRecords, successFlowFile, successAvroWriter, reader); successfulRecordCount.accumulateAndGet(successfulRecords.size(), (current, incr) -> current + incr); successfulRecords.clear(); } catch (IOException ioe) { 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); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); if (!(writeException instanceof SerializationError)) { try { hiveWriter.abort(); } catch (Exception e) { throw new ProcessException(e); } } } } } try { flushAllWriters(true); closeAllWriters(); 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) { appendRecordsToFlowFile(session, successfulRecords, failureFlowFile, failureAvroWriter, reader); } } catch (IOException ioe) { 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) { 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()) { 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 { session.remove(successFlowFile.get()); session.remove(failureFlowFile.get()); } successFlowFile.set(null); failureFlowFile.set(null); 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); if (successFlowFile.get() != null) { session.remove(successFlowFile.get()); } if (failureFlowFile.get() != null) { session.remove(failureFlowFile.get()); } } else { throw pe; } } else { throw pe; } } finally { Thread.currentThread().setContextClassLoader(originalClassloader); } } }
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(); 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); 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; 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)); 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(); 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) { 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(); if (!partitionColumnList.contains(fieldName)) { Object value = currRecord.get(fieldName); try { obj.put(fieldName, value); } catch (JSONException je) { throw new IOException(je); } } } } catch (IOException ioe) { 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) { log.error("Error connecting to Hive endpoint: table {} at {}", new Object[]{options.getTableName(), options.getMetaStoreURI()}); abortAndCloseWriters(); throw new ProcessException(connectionError); } try { try { hiveWriter.write(record.getRecord().toString().getBytes(StandardCharsets.UTF_8)); successfulRecords.add(record); } catch (InterruptedException | HiveWriter.WriteFailure wf) { log.error("Error writing record to Hive Streaming transaction", wf); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); } if (hiveWriter.getTotalRecords() >= recordsPerTxn) { hiveWriter.flush(true); try { appendRecordsToFlowFile(session, successfulRecords, successFlowFile, successAvroWriter, reader); successfulRecordCount.accumulateAndGet(successfulRecords.size(), (current, incr) -> current + incr); successfulRecords.clear(); } catch (IOException ioe) { 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); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); if (!(writeException instanceof SerializationError)) { try { hiveWriter.abort(); } catch (Exception e) { throw new ProcessException(e); } } } } } try { flushAllWriters(true); closeAllWriters(); 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) { appendRecordsToFlowFile(session, successfulRecords, failureFlowFile, failureAvroWriter, reader); } } catch (IOException ioe) { 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) { 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()) { 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 { session.remove(successFlowFile.get()); session.remove(failureFlowFile.get()); } successFlowFile.set(null); failureFlowFile.set(null); 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); if (successFlowFile.get() != null) { session.remove(successFlowFile.get()); } if (failureFlowFile.get() != null) { session.remove(failureFlowFile.get()); } } else { throw pe; } } else { throw pe; } } finally { Thread.currentThread().setContextClassLoader(originalClassloader); } }  }
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(); 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); 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; 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)); 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(); 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) { 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(); if (!partitionColumnList.contains(fieldName)) { Object value = currRecord.get(fieldName); try { obj.put(fieldName, value); } catch (JSONException je) { throw new IOException(je); } } } } catch (IOException ioe) { 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) { log.error("Error connecting to Hive endpoint: table {} at {}", new Object[]{options.getTableName(), options.getMetaStoreURI()}); abortAndCloseWriters(); throw new ProcessException(connectionError); } try { try { hiveWriter.write(record.getRecord().toString().getBytes(StandardCharsets.UTF_8)); successfulRecords.add(record); } catch (InterruptedException | HiveWriter.WriteFailure wf) { log.error("Error writing record to Hive Streaming transaction", wf); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); } if (hiveWriter.getTotalRecords() >= recordsPerTxn) { hiveWriter.flush(true); try { appendRecordsToFlowFile(session, successfulRecords, successFlowFile, successAvroWriter, reader); successfulRecordCount.accumulateAndGet(successfulRecords.size(), (current, incr) -> current + incr); successfulRecords.clear(); } catch (IOException ioe) { 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); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); if (!(writeException instanceof SerializationError)) { try { hiveWriter.abort(); } catch (Exception e) { throw new ProcessException(e); } } } } } try { flushAllWriters(true); closeAllWriters(); 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) { appendRecordsToFlowFile(session, successfulRecords, failureFlowFile, failureAvroWriter, reader); } } catch (IOException ioe) { 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) { 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()) { 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 { session.remove(successFlowFile.get()); session.remove(failureFlowFile.get()); } successFlowFile.set(null); failureFlowFile.set(null); 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); if (successFlowFile.get() != null) { session.remove(successFlowFile.get()); } if (failureFlowFile.get() != null) { session.remove(failureFlowFile.get()); } } else { throw pe; } } else { throw pe; } } finally { Thread.currentThread().setContextClassLoader(originalClassloader); } }  @Override Set<Relationship> getRelationships(); @OnScheduled void setup(final ProcessContext context); @Override void onTrigger(ProcessContext context, ProcessSession session); @OnStopped void cleanup(); }
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(); 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); 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; 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)); 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(); 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) { 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(); if (!partitionColumnList.contains(fieldName)) { Object value = currRecord.get(fieldName); try { obj.put(fieldName, value); } catch (JSONException je) { throw new IOException(je); } } } } catch (IOException ioe) { 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) { log.error("Error connecting to Hive endpoint: table {} at {}", new Object[]{options.getTableName(), options.getMetaStoreURI()}); abortAndCloseWriters(); throw new ProcessException(connectionError); } try { try { hiveWriter.write(record.getRecord().toString().getBytes(StandardCharsets.UTF_8)); successfulRecords.add(record); } catch (InterruptedException | HiveWriter.WriteFailure wf) { log.error("Error writing record to Hive Streaming transaction", wf); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); } if (hiveWriter.getTotalRecords() >= recordsPerTxn) { hiveWriter.flush(true); try { appendRecordsToFlowFile(session, successfulRecords, successFlowFile, successAvroWriter, reader); successfulRecordCount.accumulateAndGet(successfulRecords.size(), (current, incr) -> current + incr); successfulRecords.clear(); } catch (IOException ioe) { 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); appendRecordsToFlowFile(session, Collections.singletonList(record), failureFlowFile, failureAvroWriter, reader); if (!(writeException instanceof SerializationError)) { try { hiveWriter.abort(); } catch (Exception e) { throw new ProcessException(e); } } } } } try { flushAllWriters(true); closeAllWriters(); 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) { appendRecordsToFlowFile(session, successfulRecords, failureFlowFile, failureAvroWriter, reader); } } catch (IOException ioe) { 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) { 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()) { 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 { session.remove(successFlowFile.get()); session.remove(failureFlowFile.get()); } successFlowFile.set(null); failureFlowFile.set(null); 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); if (successFlowFile.get() != null) { session.remove(successFlowFile.get()); } if (failureFlowFile.get() != null) { session.remove(failureFlowFile.get()); } } else { throw pe; } } else { throw pe; } } finally { Thread.currentThread().setContextClassLoader(originalClassloader); } }  @Override Set<Relationship> getRelationships(); @OnScheduled void setup(final ProcessContext context); @Override void onTrigger(ProcessContext context, ProcessSession session); @OnStopped void cleanup(); static final String HIVE_STREAMING_RECORD_COUNT_ATTR; static final PropertyDescriptor METASTORE_URI; static final PropertyDescriptor HIVE_CONFIGURATION_RESOURCES; static final PropertyDescriptor DB_NAME; static final PropertyDescriptor TABLE_NAME; static final PropertyDescriptor PARTITION_COLUMNS; static final PropertyDescriptor AUTOCREATE_PARTITIONS; static final PropertyDescriptor MAX_OPEN_CONNECTIONS; static final PropertyDescriptor HEARTBEAT_INTERVAL; static final PropertyDescriptor TXNS_PER_BATCH; static final PropertyDescriptor RECORDS_PER_TXN; static final Relationship REL_SUCCESS; static final Relationship REL_FAILURE; static final Relationship REL_RETRY; }