Github user ijokarumawak commented on a diff in the pull request:
https://github.com/apache/nifi/pull/2162#discussion_r166280838
--- Diff:
nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/AbstractDatabaseFetchProcessor.java
---
@@ -249,34 +260,56 @@ public void setup(final ProcessContext context,
boolean shouldCleanCache, FlowFi
return;
}
- // Try to fill the columnTypeMap with the types of the desired
max-value columns
- final DBCPService dbcpService =
context.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class);
- final String tableName =
context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();
+ // Try to fill the columnTypeMap with the types of the desired
max-value columns
+ final DBCPService dbcpService =
context.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class);
+ final String tableName =
context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();
+ final String sqlQuery =
context.getProperty(SQL_QUERY).evaluateAttributeExpressions().getValue();
final DatabaseAdapter dbAdapter =
dbAdapters.get(context.getProperty(DB_TYPE).getValue());
try (final Connection con = dbcpService.getConnection();
final Statement st = con.createStatement()) {
- // Try a query that returns no rows, for the purposes of
getting metadata about the columns. It is possible
- // to use DatabaseMetaData.getColumns(), but not all
drivers support this, notably the schema-on-read
- // approach as in Apache Drill
- String query = dbAdapter.getSelectStatement(tableName,
maxValueColumnNames, "1 = 0", null, null, null);
- ResultSet resultSet = st.executeQuery(query);
- ResultSetMetaData resultSetMetaData =
resultSet.getMetaData();
- int numCols = resultSetMetaData.getColumnCount();
- if (numCols > 0) {
- if (shouldCleanCache) {
- columnTypeMap.clear();
- }
- for (int i = 1; i <= numCols; i++) {
- String colName =
resultSetMetaData.getColumnName(i).toLowerCase();
- String colKey = getStateKey(tableName, colName);
- int colType = resultSetMetaData.getColumnType(i);
- columnTypeMap.putIfAbsent(colKey, colType);
+ // Try a query that returns no rows, for the purposes of
getting metadata about the columns. It is possible
+ // to use DatabaseMetaData.getColumns(), but not all drivers
support this, notably the schema-on-read
+ // approach as in Apache Drill
+ String query;
+
+ if(StringUtils.isEmpty(sqlQuery)) {
+ query = dbAdapter.getSelectStatement(tableName,
maxValueColumnNames, "1 = 0", null, null, null);
+ } else {
+ StringBuilder sbQuery = getWrappedQuery(sqlQuery,
tableName);
+ sbQuery.append(" WHERE 1=0");
+
+ query = sbQuery.toString();
+ }
+
+ ResultSet resultSet = st.executeQuery(query);
+ ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+ int numCols = resultSetMetaData.getColumnCount();
+ if (numCols > 0) {
+ if (shouldCleanCache){
+ columnTypeMap.clear();
+ }
+ for (int i = 1; i <= numCols; i++) {
+ String colName =
resultSetMetaData.getColumnName(i).toLowerCase();
+ String colKey = getStateKey(tableName, colName);
+ int colType = resultSetMetaData.getColumnType(i);
+ columnTypeMap.putIfAbsent(colKey, colType);
+ }
+
+ List<String> maxValueColumnNameList =
org.apache.commons.lang3.StringUtils.isEmpty(maxValueColumnNames)
+ ? null
+ :
Arrays.asList(maxValueColumnNames.split("\\s*,\\s*"));
--- End diff --
I guess the aim of the regex `\s*,\s` is to split and trim whitespaces. But
it leaves whitespaces at the head and tail. I'd suggest simply split with `,`
and trim it in the for loop below before `toLowerCase()`.
---
