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

Example 11 with StateMap

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

the class TailFile method persistState.

private void persistState(final Map<String, String> state, final ProcessContext context) {
    try {
        StateMap oldState = context.getStateManager().getState(getStateScope(context));
        Map<String, String> updatedState = new HashMap<String, String>();
        for (String key : oldState.toMap().keySet()) {
            // New states have 'file.<index>.' prefix.
            if (TailFileState.StateKeys.CHECKSUM.equals(key) || TailFileState.StateKeys.FILENAME.equals(key) || TailFileState.StateKeys.POSITION.equals(key) || TailFileState.StateKeys.TIMESTAMP.equals(key)) {
                getLogger().info("Removed state {}={} stored by older version of NiFi.", new Object[] { key, oldState.get(key) });
                continue;
            }
            updatedState.put(key, oldState.get(key));
        }
        updatedState.putAll(state);
        context.getStateManager().setState(updatedState, getStateScope(context));
    } catch (final IOException e) {
        getLogger().warn("Failed to store state due to {}; some data may be duplicated on restart of NiFi", new Object[] { e });
    }
}

Example 12 with StateMap

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

the class GetHTTP method updateStateMap.

private void updateStateMap(ProcessContext context, HttpResponse response, StateMap beforeStateMap, String url) {
    try {
        Map<String, String> workingMap = new HashMap<>();
        workingMap.putAll(beforeStateMap.toMap());
        final StateManager stateManager = context.getStateManager();
        StateMap oldValue = beforeStateMap;
        long currentTime = System.currentTimeMillis();
        final Header receivedLastModified = response.getFirstHeader(HEADER_LAST_MODIFIED);
        if (receivedLastModified != null) {
            workingMap.put(LAST_MODIFIED + ":" + url, currentTime + ":" + receivedLastModified.getValue());
        }
        final Header receivedEtag = response.getFirstHeader(HEADER_ETAG);
        if (receivedEtag != null) {
            workingMap.put(ETAG + ":" + url, currentTime + ":" + receivedEtag.getValue());
        }
        boolean replaceSucceeded = stateManager.replace(oldValue, workingMap, Scope.LOCAL);
        boolean changed;
        while (!replaceSucceeded) {
            oldValue = stateManager.getState(Scope.LOCAL);
            workingMap.clear();
            workingMap.putAll(oldValue.toMap());
            changed = false;
            if (receivedLastModified != null) {
                Tuple<String, String> storedLastModifiedTuple = parseStateValue(workingMap.get(LAST_MODIFIED + ":" + url));
                if (Long.parseLong(storedLastModifiedTuple.getKey()) < currentTime) {
                    workingMap.put(LAST_MODIFIED + ":" + url, currentTime + ":" + receivedLastModified.getValue());
                    changed = true;
                }
            }
            if (receivedEtag != null) {
                Tuple<String, String> storedLastModifiedTuple = parseStateValue(workingMap.get(ETAG + ":" + url));
                if (Long.parseLong(storedLastModifiedTuple.getKey()) < currentTime) {
                    workingMap.put(ETAG + ":" + url, currentTime + ":" + receivedEtag.getValue());
                    changed = true;
                }
            }
            if (changed) {
                replaceSucceeded = stateManager.replace(oldValue, workingMap, Scope.LOCAL);
            } else {
                break;
            }
        }
    } catch (final IOException ioe) {
        throw new ProcessException(ioe);
    }
}

Example 13 with StateMap

use of org.apache.nifi.components.state.StateMap 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 14 with StateMap

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

the class RedisStateProvider method setState.

@Override
public void setState(final Map<String, String> state, final String componentId) throws IOException {
    verifyEnabled();
    final StateMap currStateMap = getState(componentId);
    int attempted = 0;
    boolean updated = false;
    while (!updated && attempted < 20) {
        updated = replace(currStateMap, state, componentId, true);
        attempted++;
    }
    if (!updated) {
        throw new IOException("Unable to update state due to concurrent modifications");
    }
}

Example 15 with StateMap

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

the class RedisStateProvider method clear.

@Override
public void clear(final String componentId) throws IOException {
    int attempted = 0;
    boolean updated = false;
    while (!updated && attempted < 20) {
        final StateMap currStateMap = getState(componentId);
        updated = replace(currStateMap, Collections.emptyMap(), componentId, true);
        final String result = updated ? "successful" : "unsuccessful";
        logger.debug("Attempt # {} to clear state for component {} was {}", new Object[] { attempted + 1, componentId, result });
        attempted++;
    }
    if (!updated) {
        throw new IOException("Unable to update state due to concurrent modifications");
    }
}