Apache
/
nifi
mirror of https://github.com/apache/nifi.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

NIFI-4749: This closes #2386. Pass the record reader's schema along to the ResultSetRecordSet so that it is able to resolve the schema for Record fields 

Signed-off-by: joewitt <joewitt@apache.org>
This commit is contained in:
Mark Payne 2018-01-08 16:18:34 -05:00 committed by joewitt
parent a7f1eb89c2
commit 953e922d32
2 changed files with 41 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
nifi-commons/nifi-record/src/main/java/org/apache/nifi/serialization/record/ResultSetRecordSet.java
View File

@ -30,6 +30,7 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
@ -45,10 +46,10 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
private final Set<String> rsColumnNames;
private boolean moreRows;
public ResultSetRecordSet(final ResultSet rs) throws SQLException {
public ResultSetRecordSet(final ResultSet rs, final RecordSchema readerSchema) throws SQLException {
this.rs = rs;
moreRows = rs.next();
this.schema = createSchema(rs);
this.schema = createSchema(rs, readerSchema);
rsColumnNames = new HashSet<>();
final ResultSetMetaData metadata = rs.getMetaData();
@ -118,7 +119,7 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
return value;
}
private static RecordSchema createSchema(final ResultSet rs) throws SQLException {
private static RecordSchema createSchema(final ResultSet rs, final RecordSchema readerSchema) throws SQLException {
final ResultSetMetaData metadata = rs.getMetaData();
final int numCols = metadata.getColumnCount();
final List<RecordField> fields = new ArrayList<>(numCols);
@ -127,7 +128,7 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
final int column = i + 1;
final int sqlType = metadata.getColumnType(column);
final DataType dataType = getDataType(sqlType, rs, column);
final DataType dataType = getDataType(sqlType, rs, column, readerSchema);
final String fieldName = metadata.getColumnLabel(column);
final int nullableFlag = metadata.isNullable(column);
@ -145,7 +146,7 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
return new SimpleRecordSchema(fields);
}
private static DataType getDataType(final int sqlType, final ResultSet rs, final int columnIndex) throws SQLException {
private static DataType getDataType(final int sqlType, final ResultSet rs, final int columnIndex, final RecordSchema readerSchema) throws SQLException {
switch (sqlType) {
case Types.ARRAY:
// The JDBC API does not allow us to know what the base type of an array is through the metadata.
@ -168,12 +169,18 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
case Types.LONGVARBINARY:
case Types.VARBINARY:
return RecordFieldType.ARRAY.getArrayDataType(RecordFieldType.BYTE.getDataType());
case Types.OTHER:
case Types.OTHER: {
// If we have no records to inspect, we can't really know its schema so we simply use the default data type.
if (rs.isAfterLast()) {
return RecordFieldType.RECORD.getDataType();
}
final String columnName = rs.getMetaData().getColumnName(columnIndex);
Optional<DataType> dataType = readerSchema.getDataType(columnName);
if (dataType.isPresent()) {
return dataType.get();
}
final Object obj = rs.getObject(columnIndex);
if (obj == null || !(obj instanceof Record)) {
final List<DataType> dataTypes = Stream.of(RecordFieldType.BIGINT, RecordFieldType.BOOLEAN, RecordFieldType.BYTE, RecordFieldType.CHAR, RecordFieldType.DATE,
@ -188,8 +195,16 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
final Record record = (Record) obj;
final RecordSchema recordSchema = record.getSchema();
return RecordFieldType.RECORD.getRecordDataType(recordSchema);
default:
}
default: {
final String columnName = rs.getMetaData().getColumnName(columnIndex);
Optional<DataType> dataType = readerSchema.getDataType(columnName);
if (dataType.isPresent()) {
return dataType.get();
}
return getFieldType(sqlType).getDataType();
}
}
}
@ -286,7 +301,8 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
return RecordFieldType.TIMESTAMP.getDataType();
}
if (valueToLookAt instanceof Record) {
return RecordFieldType.RECORD.getDataType();
final Record record = (Record) valueToLookAt;
return RecordFieldType.RECORD.getRecordDataType(record.getSchema());
}
}

18
nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/QueryRecord.java
View File

@ -18,6 +18,7 @@ package org.apache.nifi.processors.standard;
import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.sql.Connection;
import java.sql.DriverManager;
@ -74,6 +75,7 @@ import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.io.OutputStreamCallback;
import org.apache.nifi.queryrecord.FlowFileTable;
import org.apache.nifi.schema.access.SchemaNotFoundException;
import org.apache.nifi.serialization.RecordReader;
import org.apache.nifi.serialization.RecordReaderFactory;
import org.apache.nifi.serialization.RecordSetWriter;
import org.apache.nifi.serialization.RecordSetWriterFactory;
@ -256,6 +258,20 @@ public class QueryRecord extends AbstractProcessor {
final Map<FlowFile, Relationship> transformedFlowFiles = new HashMap<>();
final Set<FlowFile> createdFlowFiles = new HashSet<>();
// Determine the Record Reader's schema
final RecordSchema readerSchema;
try (final InputStream rawIn = session.read(original)) {
final Map<String, String> originalAttributes = original.getAttributes();
final RecordReader reader = recordReaderFactory.createRecordReader(originalAttributes, rawIn, getLogger());
final RecordSchema inputSchema = reader.getSchema();
readerSchema = recordSetWriterFactory.getSchema(originalAttributes, inputSchema);
} catch (final Exception e) {
getLogger().error("Failed to determine Record Schema from {}; routing to failure", new Object[] {original, e});
session.transfer(original, REL_FAILURE);
return;
}
// Determine the schema for writing the data
final Map<String, String> originalAttributes = original.getAttributes();
int recordsRead = 0;
@ -294,7 +310,7 @@ public class QueryRecord extends AbstractProcessor {
final RecordSchema writeSchema;
try {
recordSet = new ResultSetRecordSet(rs);
recordSet = new ResultSetRecordSet(rs, readerSchema);
final RecordSchema resultSetSchema = recordSet.getSchema();
writeSchema = recordSetWriterFactory.getSchema(originalAttributes, resultSetSchema);
} catch (final SQLException | SchemaNotFoundException e) {