Examples with StateManager org.apache.nifi.components.state.StateManager used on opensource projects

Example 6 with StateManager

use of org.apache.nifi.components.state.StateManager in project nifi by apache.

the class ListDatabaseTables method onTrigger.

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    final DBCPService dbcpService = context.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class);
    final String catalog = context.getProperty(CATALOG).getValue();
    final String schemaPattern = context.getProperty(SCHEMA_PATTERN).getValue();
    final String tableNamePattern = context.getProperty(TABLE_NAME_PATTERN).getValue();
    final String[] tableTypes = context.getProperty(TABLE_TYPES).isSet() ? context.getProperty(TABLE_TYPES).getValue().split("\\s*,\\s*") : null;
    final boolean includeCount = context.getProperty(INCLUDE_COUNT).asBoolean();
    final long refreshInterval = context.getProperty(REFRESH_INTERVAL).asTimePeriod(TimeUnit.MILLISECONDS);
    final StateManager stateManager = context.getStateManager();
    final StateMap stateMap;
    final Map<String, String> stateMapProperties;
    try {
        stateMap = stateManager.getState(Scope.CLUSTER);
        stateMapProperties = new HashMap<>(stateMap.toMap());
    } catch (IOException ioe) {
        throw new ProcessException(ioe);
    }
    try (final Connection con = dbcpService.getConnection()) {
        DatabaseMetaData dbMetaData = con.getMetaData();
        ResultSet rs = dbMetaData.getTables(catalog, schemaPattern, tableNamePattern, tableTypes);
        while (rs.next()) {
            final String tableCatalog = rs.getString(1);
            final String tableSchema = rs.getString(2);
            final String tableName = rs.getString(3);
            final String tableType = rs.getString(4);
            final String tableRemarks = rs.getString(5);
            // Build fully-qualified name
            String fqn = Stream.of(tableCatalog, tableSchema, tableName).filter(segment -> !StringUtils.isEmpty(segment)).collect(Collectors.joining("."));
            String lastTimestampForTable = stateMapProperties.get(fqn);
            boolean refreshTable = true;
            try {
                // Refresh state if the interval has elapsed
                long lastRefreshed = -1;
                final long currentTime = System.currentTimeMillis();
                if (!StringUtils.isEmpty(lastTimestampForTable)) {
                    lastRefreshed = Long.parseLong(lastTimestampForTable);
                }
                if (lastRefreshed == -1 || (refreshInterval > 0 && currentTime >= (lastRefreshed + refreshInterval))) {
                    stateMapProperties.remove(lastTimestampForTable);
                } else {
                    refreshTable = false;
                }
            } catch (final NumberFormatException nfe) {
                getLogger().error("Failed to retrieve observed last table fetches from the State Manager. Will not perform " + "query until this is accomplished.", nfe);
                context.yield();
                return;
            }
            if (refreshTable) {
                FlowFile flowFile = session.create();
                logger.info("Found {}: {}", new Object[] { tableType, fqn });
                if (includeCount) {
                    try (Statement st = con.createStatement()) {
                        final String countQuery = "SELECT COUNT(1) FROM " + fqn;
                        logger.debug("Executing query: {}", new Object[] { countQuery });
                        ResultSet countResult = st.executeQuery(countQuery);
                        if (countResult.next()) {
                            flowFile = session.putAttribute(flowFile, DB_TABLE_COUNT, Long.toString(countResult.getLong(1)));
                        }
                    } catch (SQLException se) {
                        logger.error("Couldn't get row count for {}", new Object[] { fqn });
                        session.remove(flowFile);
                        continue;
                    }
                }
                if (tableCatalog != null) {
                    flowFile = session.putAttribute(flowFile, DB_TABLE_CATALOG, tableCatalog);
                }
                if (tableSchema != null) {
                    flowFile = session.putAttribute(flowFile, DB_TABLE_SCHEMA, tableSchema);
                }
                flowFile = session.putAttribute(flowFile, DB_TABLE_NAME, tableName);
                flowFile = session.putAttribute(flowFile, DB_TABLE_FULLNAME, fqn);
                flowFile = session.putAttribute(flowFile, DB_TABLE_TYPE, tableType);
                if (tableRemarks != null) {
                    flowFile = session.putAttribute(flowFile, DB_TABLE_REMARKS, tableRemarks);
                }
                String transitUri;
                try {
                    transitUri = dbMetaData.getURL();
                } catch (SQLException sqle) {
                    transitUri = "<unknown>";
                }
                session.getProvenanceReporter().receive(flowFile, transitUri);
                session.transfer(flowFile, REL_SUCCESS);
                stateMapProperties.put(fqn, Long.toString(System.currentTimeMillis()));
            }
        }
        // Update the timestamps for listed tables
        if (stateMap.getVersion() == -1) {
            stateManager.setState(stateMapProperties, Scope.CLUSTER);
        } else {
            stateManager.replace(stateMap, stateMapProperties, Scope.CLUSTER);
        }
    } catch (final SQLException | IOException e) {
        throw new ProcessException(e);
    }
}

Example 7 with StateManager

use of org.apache.nifi.components.state.StateManager in project nifi by apache.

the class ScriptedLookupService method reloadScript.

/**
 * Reloads the script RecordReaderFactory. This must be called within the lock.
 *
 * @param scriptBody An input stream associated with the script content
 * @return Whether the script was successfully reloaded
 */
@Override
protected boolean reloadScript(final String scriptBody) {
    // note we are starting here with a fresh listing of validation
    // results since we are (re)loading a new/updated script. any
    // existing validation results are not relevant
    final Collection<ValidationResult> results = new HashSet<>();
    try {
        // get the engine and ensure its invocable
        if (scriptEngine instanceof Invocable) {
            final Invocable invocable = (Invocable) scriptEngine;
            // Find a custom configurator and invoke their eval() method
            ScriptEngineConfigurator configurator = scriptingComponentHelper.scriptEngineConfiguratorMap.get(scriptingComponentHelper.getScriptEngineName().toLowerCase());
            if (configurator != null) {
                configurator.eval(scriptEngine, scriptBody, scriptingComponentHelper.getModules());
            } else {
                // evaluate the script
                scriptEngine.eval(scriptBody);
            }
            // get configured LookupService from the script (if it exists)
            final Object obj = scriptEngine.get("lookupService");
            if (obj != null) {
                final ComponentLog logger = getLogger();
                try {
                    // set the logger if the processor wants it
                    invocable.invokeMethod(obj, "setLogger", logger);
                } catch (final NoSuchMethodException nsme) {
                    if (logger.isDebugEnabled()) {
                        logger.debug("Scripted LookupService does not contain a setLogger method.");
                    }
                }
                // record the processor for use later
                final LookupService<Object> scriptedLookupService = invocable.getInterface(obj, LookupService.class);
                lookupService.set(scriptedLookupService);
                if (scriptedLookupService != null) {
                    try {
                        scriptedLookupService.initialize(new ControllerServiceInitializationContext() {

                            @Override
                            public String getIdentifier() {
                                return ScriptedLookupService.this.getIdentifier();
                            }

                            @Override
                            public ComponentLog getLogger() {
                                return logger;
                            }

                            @Override
                            public StateManager getStateManager() {
                                return ScriptedLookupService.this.getStateManager();
                            }

                            @Override
                            public ControllerServiceLookup getControllerServiceLookup() {
                                return ScriptedLookupService.super.getControllerServiceLookup();
                            }

                            @Override
                            public String getKerberosServicePrincipal() {
                                return ScriptedLookupService.this.kerberosServicePrincipal;
                            }

                            @Override
                            public File getKerberosServiceKeytab() {
                                return ScriptedLookupService.this.kerberosServiceKeytab;
                            }

                            @Override
                            public File getKerberosConfigurationFile() {
                                return ScriptedLookupService.this.kerberosConfigFile;
                            }
                        });
                    } catch (final Exception e) {
                        logger.error("Unable to initialize scripted LookupService: " + e.getLocalizedMessage(), e);
                        throw new ProcessException(e);
                    }
                }
            } else {
                throw new ScriptException("No LookupService was defined by the script.");
            }
        } else {
            throw new ScriptException("Script engine is not Invocable, cannot be used for ScriptedLookupService");
        }
    } catch (final Exception ex) {
        final ComponentLog logger = getLogger();
        final String message = "Unable to load script: " + ex.getLocalizedMessage();
        logger.error(message, ex);
        results.add(new ValidationResult.Builder().subject("ScriptedLookupServiceValidation").valid(false).explanation("Unable to load script due to " + ex.getLocalizedMessage()).input(scriptingComponentHelper.getScriptPath()).build());
    }
    // store the updated validation results
    validationResults.set(results);
    // return whether there was any issues loading the configured script
    return results.isEmpty();
}

Example 8 with StateManager

use of org.apache.nifi.components.state.StateManager in project nifi by apache.

the class QueryDatabaseTableTest method testGetQuery.

@Test
public void testGetQuery() throws Exception {
    String query = processor.getQuery(dbAdapter, "myTable", null, null, null, null);
    assertEquals("SELECT * FROM myTable", query);
    query = processor.getQuery(dbAdapter, "myTable", "col1,col2", null, null, null);
    assertEquals("SELECT col1,col2 FROM myTable", query);
    query = processor.getQuery(dbAdapter, "myTable", null, Collections.singletonList("id"), null, null);
    assertEquals("SELECT * FROM myTable", query);
    Map<String, String> maxValues = new HashMap<>();
    maxValues.put("id", "509");
    StateManager stateManager = runner.getStateManager();
    stateManager.setState(maxValues, Scope.CLUSTER);
    processor.putColumnType("mytable" + AbstractDatabaseFetchProcessor.NAMESPACE_DELIMITER + "id", Types.INTEGER);
    query = processor.getQuery(dbAdapter, "myTable", null, Collections.singletonList("id"), null, stateManager.getState(Scope.CLUSTER).toMap());
    assertEquals("SELECT * FROM myTable WHERE id > 509", query);
    maxValues.put("date_created", "2016-03-07 12:34:56");
    stateManager.setState(maxValues, Scope.CLUSTER);
    processor.putColumnType("mytable" + AbstractDatabaseFetchProcessor.NAMESPACE_DELIMITER + "date_created", Types.TIMESTAMP);
    query = processor.getQuery(dbAdapter, "myTable", null, Arrays.asList("id", "DATE_CREATED"), null, stateManager.getState(Scope.CLUSTER).toMap());
    assertEquals("SELECT * FROM myTable WHERE id > 509 AND DATE_CREATED >= '2016-03-07 12:34:56'", query);
    // Test Oracle strategy
    dbAdapter = new OracleDatabaseAdapter();
    query = processor.getQuery(dbAdapter, "myTable", null, Arrays.asList("id", "DATE_CREATED"), "type = \"CUSTOMER\"", stateManager.getState(Scope.CLUSTER).toMap());
    assertEquals("SELECT * FROM myTable WHERE id > 509 AND DATE_CREATED >= timestamp '2016-03-07 12:34:56' AND (type = \"CUSTOMER\")", query);
}

Example 9 with StateManager

use of org.apache.nifi.components.state.StateManager in project nifi by apache.

the class TestGenerateTableFetch method testBackwardsCompatibilityStateKeyStaticTableDynamicMaxValues.

@Test
public void testBackwardsCompatibilityStateKeyStaticTableDynamicMaxValues() throws Exception {
    // 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(true);
    runner.setProperty(GenerateTableFetch.MAX_VALUE_COLUMN_NAMES, "${maxValueCol}");
    runner.enqueue("".getBytes(), new HashMap<String, String>() {

        {
            put("maxValueCol", "id");
        }
    });
    // Pre-populate the state with a key for column name (not fully-qualified)
    StateManager stateManager = runner.getStateManager();
    stateManager.setState(new HashMap<String, String>() {

        {
            put("id", "0");
        }
    }, Scope.CLUSTER);
    // Pre-populate the column type map with an entry for id (not fully-qualified)
    processor.columnTypeMap.put("id", 4);
    runner.run();
    runner.assertAllFlowFilesTransferred(REL_SUCCESS, 1);
    MockFlowFile flowFile = runner.getFlowFilesForRelationship(REL_SUCCESS).get(0);
    assertEquals("SELECT * FROM TEST_QUERY_DB_TABLE WHERE id > 0 AND id <= 1 ORDER BY id FETCH NEXT 10000 ROWS ONLY", new String(flowFile.toByteArray()));
}

Example 10 with StateManager

use of org.apache.nifi.components.state.StateManager in project nifi by apache.

the class AttributeRollingWindow method onScheduled.

@OnScheduled
public void onScheduled(final ProcessContext context) throws IOException {
    timeWindow = context.getProperty(TIME_WINDOW).asTimePeriod(TimeUnit.MILLISECONDS);
    microBatchTime = context.getProperty(SUB_WINDOW_LENGTH).asTimePeriod(TimeUnit.MILLISECONDS);
    if (microBatchTime == null || microBatchTime == 0) {
        StateManager stateManager = context.getStateManager();
        StateMap state = stateManager.getState(SCOPE);
        HashMap<String, String> tempMap = new HashMap<>();
        tempMap.putAll(state.toMap());
        if (!tempMap.containsKey(COUNT_KEY)) {
            tempMap.put(COUNT_KEY, "0");
            context.getStateManager().setState(tempMap, SCOPE);
        }
    }
}