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

Example 11 with DBCPService

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

the class QueryDatabaseTableTest method testOutputBatchSize.

@Test
public void testOutputBatchSize() throws ClassNotFoundException, SQLException, InitializationException, IOException {
    // load test data to database
    final Connection con = ((DBCPService) runner.getControllerService("dbcp")).getConnection();
    Statement stmt = con.createStatement();
    InputStream in;
    MockFlowFile mff;
    try {
        stmt.execute("drop table TEST_QUERY_DB_TABLE");
    } catch (final SQLException sqle) {
    // Ignore this error, probably a "table does not exist" since Derby doesn't yet support DROP IF EXISTS [DERBY-4842]
    }
    stmt.execute("create table TEST_QUERY_DB_TABLE (id integer not null, name varchar(100), scale float, created_on timestamp, bignum bigint default 0)");
    int rowCount = 0;
    // Create larger row set
    for (int batch = 0; batch < 100; batch++) {
        stmt.execute("insert into TEST_QUERY_DB_TABLE (id, name, scale, created_on) VALUES (" + rowCount + ", 'Joe Smith', 1.0, '1962-09-23 03:23:34.234')");
        rowCount++;
    }
    runner.setProperty(QueryDatabaseTable.TABLE_NAME, "TEST_QUERY_DB_TABLE");
    runner.setIncomingConnection(false);
    runner.setProperty(QueryDatabaseTable.MAX_VALUE_COLUMN_NAMES, "ID");
    runner.setProperty(QueryDatabaseTable.MAX_ROWS_PER_FLOW_FILE, "${" + MAX_ROWS_KEY + "}");
    runner.setVariable(MAX_ROWS_KEY, "7");
    runner.setProperty(QueryDatabaseTable.OUTPUT_BATCH_SIZE, "${outputBatchSize}");
    runner.setVariable("outputBatchSize", "4");
    runner.run();
    runner.assertAllFlowFilesTransferred(QueryDatabaseTable.REL_SUCCESS, 15);
    // Ensure all but the last file have 7 records each
    for (int ff = 0; ff < 14; ff++) {
        mff = runner.getFlowFilesForRelationship(QueryDatabaseTable.REL_SUCCESS).get(ff);
        in = new ByteArrayInputStream(mff.toByteArray());
        assertEquals(7, getNumberOfRecordsFromStream(in));
        mff.assertAttributeExists("fragment.identifier");
        assertEquals(Integer.toString(ff), mff.getAttribute("fragment.index"));
        // No fragment.count set for flow files sent when Output Batch Size is set
        assertNull(mff.getAttribute("fragment.count"));
    }
    // Last file should have 2 records
    mff = runner.getFlowFilesForRelationship(QueryDatabaseTable.REL_SUCCESS).get(14);
    in = new ByteArrayInputStream(mff.toByteArray());
    assertEquals(2, getNumberOfRecordsFromStream(in));
    mff.assertAttributeExists("fragment.identifier");
    assertEquals(Integer.toString(14), mff.getAttribute("fragment.index"));
    // No fragment.count set for flow files sent when Output Batch Size is set
    assertNull(mff.getAttribute("fragment.count"));
}

Example 12 with DBCPService

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

the class TestExecuteSQL method testWithNullIntColumn.

@Test
public void testWithNullIntColumn() throws SQLException {
    // remove previous test database, if any
    final File dbLocation = new File(DB_LOCATION);
    dbLocation.delete();
    // load test data to database
    final Connection con = ((DBCPService) runner.getControllerService("dbcp")).getConnection();
    Statement stmt = con.createStatement();
    try {
        stmt.execute("drop table TEST_NULL_INT");
    } catch (final SQLException sqle) {
    }
    stmt.execute("create table TEST_NULL_INT (id integer not null, val1 integer, val2 integer, constraint my_pk primary key (id))");
    stmt.execute("insert into TEST_NULL_INT (id, val1, val2) VALUES (0, NULL, 1)");
    stmt.execute("insert into TEST_NULL_INT (id, val1, val2) VALUES (1, 1, 1)");
    runner.setIncomingConnection(false);
    runner.setProperty(ExecuteSQL.SQL_SELECT_QUERY, "SELECT * FROM TEST_NULL_INT");
    runner.run();
    runner.assertAllFlowFilesTransferred(ExecuteSQL.REL_SUCCESS, 1);
    runner.getFlowFilesForRelationship(ExecuteSQL.REL_SUCCESS).get(0).assertAttributeEquals(ExecuteSQL.RESULT_ROW_COUNT, "2");
}

Example 13 with DBCPService

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

the class TestExecuteSQL method setup.

@Before
public void setup() throws InitializationException {
    final DBCPService dbcp = new DBCPServiceSimpleImpl();
    final Map<String, String> dbcpProperties = new HashMap<>();
    runner = TestRunners.newTestRunner(ExecuteSQL.class);
    runner.addControllerService("dbcp", dbcp, dbcpProperties);
    runner.enableControllerService(dbcp);
    runner.setProperty(ExecuteSQL.DBCP_SERVICE, "dbcp");
}

Example 14 with DBCPService

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

the class TestGenerateTableFetch method testMultiplePartitionsIncomingFlowFiles.

@Test
public void testMultiplePartitionsIncomingFlowFiles() throws ClassNotFoundException, SQLException, InitializationException, IOException {
    // load test data to database
    final Connection con = ((DBCPService) runner.getControllerService("dbcp")).getConnection();
    Statement stmt = con.createStatement();
    try {
        stmt.execute("drop table TEST_QUERY_DB_TABLE1");
    } catch (final SQLException sqle) {
    // Ignore this error, probably a "table does not exist" since Derby doesn't yet support DROP IF EXISTS [DERBY-4842]
    }
    stmt.execute("create table TEST_QUERY_DB_TABLE1 (id integer not null, bucket integer not null)");
    stmt.execute("insert into TEST_QUERY_DB_TABLE1 (id, bucket) VALUES (0, 0)");
    stmt.execute("insert into TEST_QUERY_DB_TABLE1 (id, bucket) VALUES (1, 0)");
    try {
        stmt.execute("drop table TEST_QUERY_DB_TABLE2");
    } catch (final SQLException sqle) {
    // Ignore this error, probably a "table does not exist" since Derby doesn't yet support DROP IF EXISTS [DERBY-4842]
    }
    stmt.execute("create table TEST_QUERY_DB_TABLE2 (id integer not null, bucket integer not null)");
    stmt.execute("insert into TEST_QUERY_DB_TABLE2 (id, bucket) VALUES (0, 0)");
    runner.setProperty(GenerateTableFetch.TABLE_NAME, "${tableName}");
    runner.setIncomingConnection(true);
    runner.setProperty(GenerateTableFetch.PARTITION_SIZE, "${partSize}");
    runner.enqueue("".getBytes(), new HashMap<String, String>() {

        {
            put("tableName", "TEST_QUERY_DB_TABLE1");
            put("partSize", "1");
        }
    });
    runner.enqueue("".getBytes(), new HashMap<String, String>() {

        {
            put("tableName", "TEST_QUERY_DB_TABLE2");
            put("partSize", "2");
        }
    });
    // The table does not exist, expect the original flow file to be routed to failure
    runner.enqueue("".getBytes(), new HashMap<String, String>() {

        {
            put("tableName", "TEST_QUERY_DB_TABLE3");
            put("partSize", "1");
        }
    });
    runner.run(3);
    runner.assertTransferCount(AbstractDatabaseFetchProcessor.REL_SUCCESS, 3);
    // Two records from table 1
    assertEquals(runner.getFlowFilesForRelationship(AbstractDatabaseFetchProcessor.REL_SUCCESS).stream().filter((ff) -> "TEST_QUERY_DB_TABLE1".equals(ff.getAttribute("tableName"))).count(), 2);
    // One record from table 2
    assertEquals(runner.getFlowFilesForRelationship(AbstractDatabaseFetchProcessor.REL_SUCCESS).stream().filter((ff) -> "TEST_QUERY_DB_TABLE2".equals(ff.getAttribute("tableName"))).count(), 1);
    // Table 3 doesn't exist, should be routed to failure
    runner.assertTransferCount(GenerateTableFetch.REL_FAILURE, 1);
    try {
        stmt.execute("drop table TEST_QUERY_DB_TABLE1");
    } catch (final SQLException sqle) {
    // Ignore this error, probably a "table does not exist" since Derby doesn't yet support DROP IF EXISTS [DERBY-4842]
    }
    try {
        stmt.execute("drop table TEST_QUERY_DB_TABLE2");
    } catch (final SQLException sqle) {
    // Ignore this error, probably a "table does not exist" since Derby doesn't yet support DROP IF EXISTS [DERBY-4842]
    }
}

Example 15 with DBCPService

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

the class TestGenerateTableFetch method testMultiplePartitions.

@Test
public void testMultiplePartitions() throws ClassNotFoundException, SQLException, InitializationException, IOException {
    // load test data to database
    final Connection con = ((DBCPService) runner.getControllerService("dbcp")).getConnection();
    Statement stmt = con.createStatement();
    try {
        stmt.execute("drop table TEST_QUERY_DB_TABLE");
    } catch (final SQLException sqle) {
    // Ignore this error, probably a "table does not exist" since Derby doesn't yet support DROP IF EXISTS [DERBY-4842]
    }
    stmt.execute("create table TEST_QUERY_DB_TABLE (id integer not null, bucket integer not null)");
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (0, 0)");
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (1, 0)");
    runner.setProperty(GenerateTableFetch.TABLE_NAME, "TEST_QUERY_DB_TABLE");
    runner.setIncomingConnection(false);
    runner.setProperty(GenerateTableFetch.MAX_VALUE_COLUMN_NAMES, "ID, BUCKET");
    // Set partition size to 1 so we can compare flow files to records
    runner.setProperty(GenerateTableFetch.PARTITION_SIZE, "1");
    runner.run();
    runner.assertAllFlowFilesTransferred(GenerateTableFetch.REL_SUCCESS, 2);
    runner.clearTransferState();
    // Add a new row in the same bucket
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (2, 0)");
    runner.run();
    runner.assertAllFlowFilesTransferred(GenerateTableFetch.REL_SUCCESS, 1);
    runner.clearTransferState();
    // Add a new row in a new bucket
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (3, 1)");
    runner.run();
    runner.assertAllFlowFilesTransferred(GenerateTableFetch.REL_SUCCESS, 1);
    runner.clearTransferState();
    // Add a new row in an old bucket, it should not be transferred
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (4, 0)");
    runner.run();
    runner.assertTransferCount(GenerateTableFetch.REL_SUCCESS, 0);
    // Add a new row in the second bucket, only the new row should be transferred
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (5, 1)");
    runner.run();
    runner.assertAllFlowFilesTransferred(GenerateTableFetch.REL_SUCCESS, 1);
    runner.clearTransferState();
}