Examples with DBCPService org.apache.nifi.dbcp.DBCPService used on opensource projects

Example 41 with DBCPService

use of org.apache.nifi.dbcp.DBCPService in project nifi by apache.

the class GenerateTableFetch method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSessionFactory sessionFactory) throws ProcessException {
    // Fetch the column/table info once (if the table name and max value columns are not dynamic). Otherwise do the setup later
    if (!isDynamicTableName && !isDynamicMaxValues && !setupComplete.get()) {
        super.setup(context);
    }
    ProcessSession session = sessionFactory.createSession();
    FlowFile fileToProcess = null;
    if (context.hasIncomingConnection()) {
        fileToProcess = session.get();
        if (fileToProcess == null) {
            // Incoming connection with no flow file available, do no work (see capability description)
            return;
        }
    }
    final ComponentLog logger = getLogger();
    final DBCPService dbcpService = context.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class);
    final DatabaseAdapter dbAdapter = dbAdapters.get(context.getProperty(DB_TYPE).getValue());
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(fileToProcess).getValue();
    final String columnNames = context.getProperty(COLUMN_NAMES).evaluateAttributeExpressions(fileToProcess).getValue();
    final String maxValueColumnNames = context.getProperty(MAX_VALUE_COLUMN_NAMES).evaluateAttributeExpressions(fileToProcess).getValue();
    final int partitionSize = context.getProperty(PARTITION_SIZE).evaluateAttributeExpressions(fileToProcess).asInteger();
    final String customWhereClause = context.getProperty(WHERE_CLAUSE).evaluateAttributeExpressions(fileToProcess).getValue();
    final StateManager stateManager = context.getStateManager();
    final StateMap stateMap;
    FlowFile finalFileToProcess = fileToProcess;
    try {
        stateMap = stateManager.getState(Scope.CLUSTER);
    } catch (final IOException ioe) {
        logger.error("Failed to retrieve observed maximum values from the State Manager. Will not perform " + "query until this is accomplished.", ioe);
        context.yield();
        return;
    }
    try {
        // Make a mutable copy of the current state property map. This will be updated by the result row callback, and eventually
        // set as the current state map (after the session has been committed)
        final Map<String, String> statePropertyMap = new HashMap<>(stateMap.toMap());
        // If an initial max value for column(s) has been specified using properties, and this column is not in the state manager, sync them to the state property map
        for (final Map.Entry<String, String> maxProp : maxValueProperties.entrySet()) {
            String maxPropKey = maxProp.getKey().toLowerCase();
            String fullyQualifiedMaxPropKey = getStateKey(tableName, maxPropKey);
            if (!statePropertyMap.containsKey(fullyQualifiedMaxPropKey)) {
                String newMaxPropValue;
                // but store the new initial max value under the fully-qualified key.
                if (statePropertyMap.containsKey(maxPropKey)) {
                    newMaxPropValue = statePropertyMap.get(maxPropKey);
                } else {
                    newMaxPropValue = maxProp.getValue();
                }
                statePropertyMap.put(fullyQualifiedMaxPropKey, newMaxPropValue);
            }
        }
        // Build a WHERE clause with maximum-value columns (if they exist), and a list of column names that will contain MAX(<column>) aliases. The
        // executed SQL query will retrieve the count of all records after the filter(s) have been applied, as well as the new maximum values for the
        // specified columns. This allows the processor to generate the correctly partitioned SQL statements as well as to update the state with the
        // latest observed maximum values.
        String whereClause = null;
        List<String> maxValueColumnNameList = StringUtils.isEmpty(maxValueColumnNames) ? new ArrayList<>(0) : Arrays.asList(maxValueColumnNames.split("\\s*,\\s*"));
        List<String> maxValueClauses = new ArrayList<>(maxValueColumnNameList.size());
        String columnsClause = null;
        List<String> maxValueSelectColumns = new ArrayList<>(maxValueColumnNameList.size() + 1);
        maxValueSelectColumns.add("COUNT(*)");
        // For each maximum-value column, get a WHERE filter and a MAX(column) alias
        IntStream.range(0, maxValueColumnNameList.size()).forEach((index) -> {
            String colName = maxValueColumnNameList.get(index);
            maxValueSelectColumns.add("MAX(" + colName + ") " + colName);
            String maxValue = getColumnStateMaxValue(tableName, statePropertyMap, colName);
            if (!StringUtils.isEmpty(maxValue)) {
                if (columnTypeMap.isEmpty() || getColumnType(tableName, colName) == null) {
                    // This means column type cache is clean after instance reboot. We should re-cache column type
                    super.setup(context, false, finalFileToProcess);
                }
                Integer type = getColumnType(tableName, colName);
                // Add a condition for the WHERE clause
                maxValueClauses.add(colName + (index == 0 ? " > " : " >= ") + getLiteralByType(type, maxValue, dbAdapter.getName()));
            }
        });
        if (customWhereClause != null) {
            // adding the custom WHERE clause (if defined) to the list of existing clauses.
            maxValueClauses.add("(" + customWhereClause + ")");
        }
        whereClause = StringUtils.join(maxValueClauses, " AND ");
        columnsClause = StringUtils.join(maxValueSelectColumns, ", ");
        // Build a SELECT query with maximum-value columns (if present)
        final String selectQuery = dbAdapter.getSelectStatement(tableName, columnsClause, whereClause, null, null, null);
        long rowCount = 0;
        try (final Connection con = dbcpService.getConnection();
            final Statement st = con.createStatement()) {
            final Integer queryTimeout = context.getProperty(QUERY_TIMEOUT).evaluateAttributeExpressions(fileToProcess).asTimePeriod(TimeUnit.SECONDS).intValue();
            // timeout in seconds
            st.setQueryTimeout(queryTimeout);
            logger.debug("Executing {}", new Object[] { selectQuery });
            ResultSet resultSet;
            resultSet = st.executeQuery(selectQuery);
            if (resultSet.next()) {
                // Total row count is in the first column
                rowCount = resultSet.getLong(1);
                // Update the state map with the newly-observed maximum values
                ResultSetMetaData rsmd = resultSet.getMetaData();
                for (int i = 2; i <= rsmd.getColumnCount(); i++) {
                    // Some JDBC drivers consider the columns name and label to be very different things.
                    // Since this column has been aliased lets check the label first,
                    // if there is no label we'll use the column name.
                    String resultColumnName = (StringUtils.isNotEmpty(rsmd.getColumnLabel(i)) ? rsmd.getColumnLabel(i) : rsmd.getColumnName(i)).toLowerCase();
                    String fullyQualifiedStateKey = getStateKey(tableName, resultColumnName);
                    String resultColumnCurrentMax = statePropertyMap.get(fullyQualifiedStateKey);
                    if (StringUtils.isEmpty(resultColumnCurrentMax) && !isDynamicTableName) {
                        // If we can't find the value at the fully-qualified key name and the table name is static, it is possible (under a previous scheme)
                        // the value has been stored under a key that is only the column name. Fall back to check the column name; either way, when a new
                        // maximum value is observed, it will be stored under the fully-qualified key from then on.
                        resultColumnCurrentMax = statePropertyMap.get(resultColumnName);
                    }
                    int type = rsmd.getColumnType(i);
                    if (isDynamicTableName) {
                        // We haven't pre-populated the column type map if the table name is dynamic, so do it here
                        columnTypeMap.put(fullyQualifiedStateKey, type);
                    }
                    try {
                        String newMaxValue = getMaxValueFromRow(resultSet, i, type, resultColumnCurrentMax, dbAdapter.getName());
                        if (newMaxValue != null) {
                            statePropertyMap.put(fullyQualifiedStateKey, newMaxValue);
                        }
                    } catch (ParseException | IOException pie) {
                        // Fail the whole thing here before we start creating flow files and such
                        throw new ProcessException(pie);
                    }
                }
            } else {
                // Something is very wrong here, one row (even if count is zero) should be returned
                throw new SQLException("No rows returned from metadata query: " + selectQuery);
            }
            // for each maximum-value column get a right bounding WHERE condition
            IntStream.range(0, maxValueColumnNameList.size()).forEach((index) -> {
                String colName = maxValueColumnNameList.get(index);
                maxValueSelectColumns.add("MAX(" + colName + ") " + colName);
                String maxValue = getColumnStateMaxValue(tableName, statePropertyMap, colName);
                if (!StringUtils.isEmpty(maxValue)) {
                    if (columnTypeMap.isEmpty() || getColumnType(tableName, colName) == null) {
                        // This means column type cache is clean after instance reboot. We should re-cache column type
                        super.setup(context, false, finalFileToProcess);
                    }
                    Integer type = getColumnType(tableName, colName);
                    // Add a condition for the WHERE clause
                    maxValueClauses.add(colName + " <= " + getLiteralByType(type, maxValue, dbAdapter.getName()));
                }
            });
            // Update WHERE list to include new right hand boundaries
            whereClause = StringUtils.join(maxValueClauses, " AND ");
            final long numberOfFetches = (partitionSize == 0) ? 1 : (rowCount / partitionSize) + (rowCount % partitionSize == 0 ? 0 : 1);
            // Generate SQL statements to read "pages" of data
            for (long i = 0; i < numberOfFetches; i++) {
                Long limit = partitionSize == 0 ? null : (long) partitionSize;
                Long offset = partitionSize == 0 ? null : i * partitionSize;
                final String maxColumnNames = StringUtils.join(maxValueColumnNameList, ", ");
                final String query = dbAdapter.getSelectStatement(tableName, columnNames, whereClause, maxColumnNames, limit, offset);
                FlowFile sqlFlowFile = (fileToProcess == null) ? session.create() : session.create(fileToProcess);
                sqlFlowFile = session.write(sqlFlowFile, out -> out.write(query.getBytes()));
                sqlFlowFile = session.putAttribute(sqlFlowFile, "generatetablefetch.tableName", tableName);
                if (columnNames != null) {
                    sqlFlowFile = session.putAttribute(sqlFlowFile, "generatetablefetch.columnNames", columnNames);
                }
                if (StringUtils.isNotBlank(whereClause)) {
                    sqlFlowFile = session.putAttribute(sqlFlowFile, "generatetablefetch.whereClause", whereClause);
                }
                if (StringUtils.isNotBlank(maxColumnNames)) {
                    sqlFlowFile = session.putAttribute(sqlFlowFile, "generatetablefetch.maxColumnNames", maxColumnNames);
                }
                sqlFlowFile = session.putAttribute(sqlFlowFile, "generatetablefetch.limit", String.valueOf(limit));
                if (partitionSize != 0) {
                    sqlFlowFile = session.putAttribute(sqlFlowFile, "generatetablefetch.offset", String.valueOf(offset));
                }
                session.transfer(sqlFlowFile, REL_SUCCESS);
            }
            if (fileToProcess != null) {
                session.remove(fileToProcess);
            }
        } catch (SQLException e) {
            if (fileToProcess != null) {
                logger.error("Unable to execute SQL select query {} due to {}, routing {} to failure", new Object[] { selectQuery, e, fileToProcess });
                fileToProcess = session.putAttribute(fileToProcess, "generatetablefetch.sql.error", e.getMessage());
                session.transfer(fileToProcess, REL_FAILURE);
            } else {
                logger.error("Unable to execute SQL select query {} due to {}", new Object[] { selectQuery, e });
                throw new ProcessException(e);
            }
        }
        session.commit();
        try {
            // Update the state
            stateManager.setState(statePropertyMap, Scope.CLUSTER);
        } catch (IOException ioe) {
            logger.error("{} failed to update State Manager, observed maximum values will not be recorded. " + "Also, any generated SQL statements may be duplicated.", new Object[] { this, ioe });
        }
    } catch (final ProcessException pe) {
        // Log the cause of the ProcessException if it is available
        Throwable t = (pe.getCause() == null ? pe : pe.getCause());
        logger.error("Error during processing: {}", new Object[] { t.getMessage() }, t);
        session.rollback();
        context.yield();
    }
}

Example 42 with DBCPService

use of org.apache.nifi.dbcp.DBCPService in project nifi by apache.

the class TestPutHiveQL method testMultipleStatementsWithinFlowFile.

@Test
public void testMultipleStatementsWithinFlowFile() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutHiveQL.class);
    final File tempDir = folder.getRoot();
    final File dbDir = new File(tempDir, "db");
    final DBCPService service = new MockDBCPService(dbDir.getAbsolutePath());
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);
    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate(createPersons);
        }
    }
    runner.setProperty(PutHiveQL.HIVE_DBCP_SERVICE, "dbcp");
    final String sql = "INSERT INTO PERSONS (ID, NAME, CODE) VALUES (?, ?, ?); " + "UPDATE PERSONS SET NAME='George' WHERE ID=?; ";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("hiveql.args.1.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.1.value", "1");
    attributes.put("hiveql.args.2.type", String.valueOf(Types.VARCHAR));
    attributes.put("hiveql.args.2.value", "Mark");
    attributes.put("hiveql.args.3.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.3.value", "84");
    attributes.put("hiveql.args.4.type", String.valueOf(Types.INTEGER));
    attributes.put("hiveql.args.4.value", "1");
    runner.enqueue(sql.getBytes(), attributes);
    runner.run();
    // should fail because of the semicolon
    runner.assertAllFlowFilesTransferred(PutHiveQL.REL_SUCCESS, 1);
    runner.getFlowFilesForRelationship(PutHiveQL.REL_SUCCESS).forEach(f -> f.assertAttributeEquals(PutHiveQL.ATTR_OUTPUT_TABLES, "PERSONS"));
    // Now we can check that the values were inserted by the multi-statement script.
    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            final ResultSet rs = stmt.executeQuery("SELECT * FROM PERSONS");
            assertTrue(rs.next());
            assertEquals("Record ID mismatch", 1, rs.getInt(1));
            assertEquals("Record NAME mismatch", "George", rs.getString(2));
        }
    }
}

Example 43 with DBCPService

use of org.apache.nifi.dbcp.DBCPService in project nifi by apache.

the class TestPutHiveQL method testFailInMiddleWithBadParameterValue.

@Test
public void testFailInMiddleWithBadParameterValue() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutHiveQL.class);
    final File tempDir = folder.getRoot();
    final File dbDir = new File(tempDir, "db");
    final DBCPService service = new MockDBCPService(dbDir.getAbsolutePath());
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);
    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate(createPersonsAutoId);
        }
    }
    runner.setProperty(PutHiveQL.HIVE_DBCP_SERVICE, "dbcp");
    final Map<String, String> goodAttributes = new HashMap<>();
    goodAttributes.put("hiveql.args.1.type", String.valueOf(Types.INTEGER));
    goodAttributes.put("hiveql.args.1.value", "84");
    final Map<String, String> badAttributes = new HashMap<>();
    badAttributes.put("hiveql.args.1.type", String.valueOf(Types.INTEGER));
    // Constraint violation, up to 100
    badAttributes.put("hiveql.args.1.value", "101");
    final byte[] data = "INSERT INTO PERSONS (NAME, CODE) VALUES ('Mark', ?)".getBytes();
    runner.enqueue(data, goodAttributes);
    runner.enqueue(data, badAttributes);
    runner.enqueue(data, goodAttributes);
    runner.enqueue(data, goodAttributes);
    runner.run();
    runner.assertTransferCount(PutHiveQL.REL_SUCCESS, 3);
    runner.assertTransferCount(PutHiveQL.REL_FAILURE, 1);
    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            final ResultSet rs = stmt.executeQuery("SELECT * FROM PERSONS");
            assertTrue(rs.next());
            assertEquals(1, rs.getInt(1));
            assertEquals("Mark", rs.getString(2));
            assertEquals(84, rs.getInt(3));
            assertTrue(rs.next());
            assertTrue(rs.next());
            assertFalse(rs.next());
        }
    }
}

Example 44 with DBCPService

use of org.apache.nifi.dbcp.DBCPService in project nifi by apache.

the class TestPutHiveQL method testFailInMiddleWithBadStatementRollbackOnFailure.

@Test
public void testFailInMiddleWithBadStatementRollbackOnFailure() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutHiveQL.class);
    final File tempDir = folder.getRoot();
    final File dbDir = new File(tempDir, "db");
    final DBCPService service = new MockDBCPService(dbDir.getAbsolutePath());
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);
    runner.setProperty(RollbackOnFailure.ROLLBACK_ON_FAILURE, "true");
    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate(createPersonsAutoId);
        }
    }
    runner.setProperty(PutHiveQL.HIVE_DBCP_SERVICE, "dbcp");
    runner.enqueue("INSERT INTO PERSONS (NAME, CODE) VALUES ('Mark', 84)".getBytes());
    // intentionally wrong syntax
    runner.enqueue("INSERT INTO PERSONS".getBytes());
    runner.enqueue("INSERT INTO PERSONS (NAME, CODE) VALUES ('Tom', 3)".getBytes());
    runner.enqueue("INSERT INTO PERSONS (NAME, CODE) VALUES ('Harry', 44)".getBytes());
    runner.run();
    // The 1st one should be routed to success, others should stay in queue.
    assertEquals(3, runner.getQueueSize().getObjectCount());
    runner.assertTransferCount(PutHiveQL.REL_FAILURE, 0);
    runner.assertTransferCount(PutHiveQL.REL_SUCCESS, 1);
}

Example 45 with DBCPService

use of org.apache.nifi.dbcp.DBCPService in project nifi by apache.

the class TestPutHiveQL method testDirectStatements.

@Test
public void testDirectStatements() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutHiveQL.class);
    final File tempDir = folder.getRoot();
    final File dbDir = new File(tempDir, "db");
    final DBCPService service = new MockDBCPService(dbDir.getAbsolutePath());
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);
    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate(createPersons);
        }
    }
    runner.setProperty(PutHiveQL.HIVE_DBCP_SERVICE, "dbcp");
    runner.enqueue("INSERT INTO PERSONS (ID, NAME, CODE) VALUES (1, 'Mark', 84)".getBytes());
    runner.run();
    runner.assertAllFlowFilesTransferred(PutHiveQL.REL_SUCCESS, 1);
    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            final ResultSet rs = stmt.executeQuery("SELECT * FROM PERSONS");
            assertTrue(rs.next());
            assertEquals(1, rs.getInt(1));
            assertEquals("Mark", rs.getString(2));
            assertEquals(84, rs.getInt(3));
            assertFalse(rs.next());
        }
    }
    runner.enqueue("UPDATE PERSONS SET NAME='George' WHERE ID=1".getBytes());
    runner.run();
    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            final ResultSet rs = stmt.executeQuery("SELECT * FROM PERSONS");
            assertTrue(rs.next());
            assertEquals(1, rs.getInt(1));
            assertEquals("George", rs.getString(2));
            assertEquals(84, rs.getInt(3));
            assertFalse(rs.next());
        }
    }
}