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

NIFI-5121: Added DBCPService API method for passing in flow file attributes when available 

This closes #2658

Signed-off-by: Mike Thomsen <mikerthomsen@gmail.com>
This commit is contained in:
Matthew Burgess 2018-04-25 14:36:16 -04:00 committed by Mike Thomsen
parent 500a254e3f
commit 099bfcdf3a
15 changed files with 41 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nifi-nar-bundles/nifi-extension-utils/nifi-processor-utils/src/main/java/org/apache/nifi/processor/util/pattern/PartialFunctions.java
View File

@ -32,7 +32,7 @@ public class PartialFunctions {
@FunctionalInterface
public interface InitConnection<FC, C> {
C apply(ProcessContext context, ProcessSession session, FC functionContext) throws ProcessException;
C apply(ProcessContext context, ProcessSession session, FC functionContext, FlowFile flowFile) throws ProcessException;
}
@FunctionalInterface

3
nifi-nar-bundles/nifi-extension-utils/nifi-processor-utils/src/main/java/org/apache/nifi/processor/util/pattern/Put.java
View File

@ -93,7 +93,8 @@ public class Put<FC, C extends AutoCloseable> {
return;
}
try (C connection = initConnection.apply(context, session, functionContext)) {
// Only pass in a flow file if there is a single one present
try (C connection = initConnection.apply(context, session, functionContext, flowFiles.size() == 1 ? flowFiles.get(0) : null)) {
try {
// Execute the core function.

2
nifi-nar-bundles/nifi-groovyx-bundle/nifi-groovyx-processors/src/main/java/org/apache/nifi/processors/groovyx/ExecuteGroovyScript.java
View File

@ -336,7 +336,7 @@ public class ExecuteGroovyScript extends AbstractProcessor {
private void onInitSQL(HashMap SQL) throws SQLException {
for (Map.Entry e : (Set<Map.Entry>) SQL.entrySet()) {
DBCPService s = (DBCPService) e.getValue();
OSql sql = new OSql(s.getConnection());
OSql sql = new OSql(s.getConnection(Collections.emptyMap()));
//try to set autocommit to false
try {
if (sql.getConnection().getAutoCommit()) {

2
nifi-nar-bundles/nifi-groovyx-bundle/nifi-groovyx-processors/src/test/resources/groovy/test_ctl_01_access.groovy
View File

@ -17,7 +17,7 @@
//just check that it's possible to access controller services
def ff=session.create()
def con=CTL.mydbcp.getConnection()
def con=CTL.mydbcp.getConnection([:])
assert con instanceof java.sql.Connection
con.close();
ff.write('UTF-8', 'OK')

4
nifi-nar-bundles/nifi-hive-bundle/nifi-hive-processors/src/main/java/org/apache/nifi/processors/hive/PutHiveQL.java
View File

@ -201,9 +201,9 @@ public class PutHiveQL extends AbstractHiveQLProcessor {
}
}
private InitConnection<FunctionContext, Connection> initConnection = (context, session, fc) -> {
private InitConnection<FunctionContext, Connection> initConnection = (context, session, fc, ff) -> {
final HiveDBCPService dbcpService = context.getProperty(HIVE_DBCP_SERVICE).asControllerService(HiveDBCPService.class);
final Connection connection = dbcpService.getConnection();
final Connection connection = dbcpService.getConnection(ff == null ? Collections.emptyMap() : ff.getAttributes());
fc.connectionUrl = dbcpService.getConnectionURL();
return connection;
};

2
nifi-nar-bundles/nifi-hive-bundle/nifi-hive-processors/src/main/java/org/apache/nifi/processors/hive/SelectHiveQL.java
View File

@ -308,7 +308,7 @@ public class SelectHiveQL extends AbstractHiveQLProcessor {
final boolean escape = context.getProperty(HIVEQL_CSV_HEADER).asBoolean();
final String fragmentIdentifier = UUID.randomUUID().toString();
try (final Connection con = dbcpService.getConnection();
try (final Connection con = dbcpService.getConnection(fileToProcess == null ? Collections.emptyMap() : fileToProcess.getAttributes());
final Statement st = (flowbased ? con.prepareStatement(selectQuery) : con.createStatement())
) {

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

@ -46,6 +46,7 @@ import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
@ -260,7 +261,7 @@ public abstract class AbstractDatabaseFetchProcessor extends AbstractSessionFact
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();
try (final Connection con = dbcpService.getConnection(flowFile == null ? Collections.emptyMap() : flowFile.getAttributes());
final Statement st = con.createStatement()) {
// Try a query that returns no rows, for the purposes of getting metadata about the columns. It is possible

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

@ -28,6 +28,7 @@ import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Types;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
@ -315,7 +316,7 @@ public class ConvertJSONToSQL extends AbstractProcessor {
if (schema == null) {
// No schema exists for this table yet. Query the database to determine the schema and put it into the cache.
final DBCPService dbcpService = context.getProperty(CONNECTION_POOL).asControllerService(DBCPService.class);
try (final Connection conn = dbcpService.getConnection()) {
try (final Connection conn = dbcpService.getConnection(flowFile == null ? Collections.emptyMap() : flowFile.getAttributes())) {
schema = TableSchema.from(conn, catalog, schemaName, tableName, translateFieldNames, includePrimaryKeys);
schemaCache.put(schemaKey, schema);
} catch (final SQLException e) {

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

@ -216,7 +216,7 @@ public class ExecuteSQL extends AbstractProcessor {
}
int resultCount=0;
try (final Connection con = dbcpService.getConnection();
try (final Connection con = dbcpService.getConnection(fileToProcess == null ? Collections.emptyMap() : fileToProcess.getAttributes());
final PreparedStatement st = con.prepareStatement(selectQuery)) {
st.setQueryTimeout(queryTimeout); // timeout in seconds

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

@ -277,7 +277,7 @@ public class GenerateTableFetch extends AbstractDatabaseFetchProcessor {
final String selectQuery = dbAdapter.getSelectStatement(tableName, columnsClause, whereClause, null, null, null);
long rowCount = 0;
try (final Connection con = dbcpService.getConnection();
try (final Connection con = dbcpService.getConnection(finalFileToProcess == null ? Collections.emptyMap() : finalFileToProcess.getAttributes());
final Statement st = con.createStatement()) {
final Integer queryTimeout = context.getProperty(QUERY_TIMEOUT).evaluateAttributeExpressions(fileToProcess).asTimePeriod(TimeUnit.SECONDS).intValue();

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

@ -225,7 +225,7 @@ public class ListDatabaseTables extends AbstractProcessor {
throw new ProcessException(ioe);
}
try (final Connection con = dbcpService.getConnection()) {
try (final Connection con = dbcpService.getConnection(Collections.emptyMap())) {
DatabaseMetaData dbMetaData = con.getMetaData();
ResultSet rs = dbMetaData.getTables(catalog, schemaPattern, tableNamePattern, tableTypes);

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

@ -329,8 +329,9 @@ public class PutDatabaseRecord extends AbstractSessionFactoryProcessor {
.build();
}
private final PartialFunctions.InitConnection<FunctionContext, Connection> initConnection = (c, s, fc) -> {
final Connection connection = c.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class).getConnection();
private final PartialFunctions.InitConnection<FunctionContext, Connection> initConnection = (c, s, fc, ff) -> {
final Connection connection = c.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class)
.getConnection(ff == null ? Collections.emptyMap() : ff.getAttributes());
try {
fc.originalAutoCommit = connection.getAutoCommit();
connection.setAutoCommit(false);

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

@ -64,6 +64,7 @@ import java.sql.SQLNonTransientException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
@ -233,8 +234,9 @@ public class PutSQL extends AbstractSessionFactoryProcessor {
return poll.getFlowFiles();
};
private final PartialFunctions.InitConnection<FunctionContext, Connection> initConnection = (c, s, fc) -> {
final Connection connection = c.getProperty(CONNECTION_POOL).asControllerService(DBCPService.class).getConnection();
private final PartialFunctions.InitConnection<FunctionContext, Connection> initConnection = (c, s, fc, ff) -> {
final Connection connection = c.getProperty(CONNECTION_POOL).asControllerService(DBCPService.class)
.getConnection(ff == null ? Collections.emptyMap() : ff.getAttributes());
try {
fc.originalAutoCommit = connection.getAutoCommit();
connection.setAutoCommit(false);

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

@ -286,7 +286,7 @@ public class QueryDatabaseTable extends AbstractDatabaseFetchProcessor {
final StopWatch stopWatch = new StopWatch(true);
final String fragmentIdentifier = UUID.randomUUID().toString();
try (final Connection con = dbcpService.getConnection();
try (final Connection con = dbcpService.getConnection(Collections.emptyMap());
final Statement st = con.createStatement()) {
if (fetchSize != null && fetchSize > 0) {

19
nifi-nar-bundles/nifi-standard-services/nifi-dbcp-service-api/src/main/java/org/apache/nifi/dbcp/DBCPService.java
View File

@ -17,6 +17,7 @@
package org.apache.nifi.dbcp;
import java.sql.Connection;
import java.util.Map;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.Tags;
@ -30,5 +31,21 @@ import org.apache.nifi.processor.exception.ProcessException;
@Tags({"dbcp", "jdbc", "database", "connection", "pooling", "store"})
@CapabilityDescription("Provides Database Connection Pooling Service. Connections can be asked from pool and returned after usage.")
public interface DBCPService extends ControllerService {
public Connection getConnection() throws ProcessException;
Connection getConnection() throws ProcessException;
/**
* Allows a Map of attributes to be passed to the DBCPService for use in configuration, etc.
* An implementation will want to override getConnection() to return getConnection(Collections.emptyMap()),
* and override this method (possibly with its existing getConnection() implementation).
* @param attributes a Map of attributes to be passed to the DBCPService. The use of these
* attributes is implementation-specific, and the source of the attributes
* is processor-specific
* @return a Connection from the specifed/configured connection pool(s)
* @throws ProcessException if an error occurs while getting a connection
*/
default Connection getConnection(Map<String,String> attributes) throws ProcessException {
// default implementation (for backwards compatibility) is to call getConnection()
// without attributes
return getConnection();
}
}