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

Added regex support for attribute header selection on HandleHTTPResponse

This commit is contained in:
r65535 2020-10-29 15:44:24 +00:00 committed by markap14
parent 8b78277a45
commit d8d9aa9cec
2 changed files with 61 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -80,6 +80,12 @@ public class HandleHttpResponse extends AbstractProcessor {
.required(true)
.identifiesControllerService(HttpContextMap.class)
.build();
public static final PropertyDescriptor ATTRIBUTES_AS_HEADERS_REGEX = new PropertyDescriptor.Builder()
.name("Attributes to add to the HTTP Response (Regex)")
.description("Specifies the Regular Expression that determines the names of FlowFile attributes that should be added to the HTTP response")
.addValidator(StandardValidators.REGULAR_EXPRESSION_VALIDATOR)
.required(false)
.build();
public static final Relationship REL_SUCCESS = new Relationship.Builder()
.name("success")
@ -96,6 +102,7 @@ public class HandleHttpResponse extends AbstractProcessor {
final List<PropertyDescriptor> properties = new ArrayList<>();
properties.add(STATUS_CODE);
properties.add(HTTP_CONTEXT_MAP);
properties.add(ATTRIBUTES_AS_HEADERS_REGEX);
return properties;
}
@ -166,6 +173,21 @@ public class HandleHttpResponse extends AbstractProcessor {
}
}
final String attributeHeaderRegex = context.getProperty(ATTRIBUTES_AS_HEADERS_REGEX).getValue();
if (attributeHeaderRegex != null) {
final Pattern pattern = Pattern.compile(attributeHeaderRegex);
final Map<String, String> attributes = flowFile.getAttributes();
for (final Map.Entry<String, String> entry : attributes.entrySet()) {
final String key = entry.getKey();
if (pattern.matcher(key).matches()) {
if (!entry.getValue().trim().isEmpty()){
response.setHeader(entry.getKey(), entry.getValue());
}
}
}
}
try {
session.exportTo(flowFile, response.getOutputStream());
response.flushBuffer();

39
nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestHandleHttpResponse.java
View File

@ -93,6 +93,45 @@ public class TestHandleHttpResponse {
assertTrue(contextMap.headersWithNoValue.isEmpty());
}
@Test
public void testRegexHeaders() throws InitializationException {
final TestRunner runner = TestRunners.newTestRunner(HandleHttpResponse.class);
final MockHttpContextMap contextMap = new MockHttpContextMap("my-id", "");
runner.addControllerService("http-context-map", contextMap);
runner.enableControllerService(contextMap);
runner.setProperty(HandleHttpResponse.HTTP_CONTEXT_MAP, "http-context-map");
runner.setProperty(HandleHttpResponse.STATUS_CODE, "${status.code}");
runner.setProperty(HandleHttpResponse.ATTRIBUTES_AS_HEADERS_REGEX, "^(my.*)$");
final Map<String, String> attributes = new HashMap<>();
attributes.put(HTTPUtils.HTTP_CONTEXT_ID, "my-id");
attributes.put(HTTPUtils.HTTP_REQUEST_URI, "/test");
attributes.put(HTTPUtils.HTTP_LOCAL_NAME, "server");
attributes.put(HTTPUtils.HTTP_PORT, "8443");
attributes.put(HTTPUtils.HTTP_REMOTE_HOST, "client");
attributes.put(HTTPUtils.HTTP_SSL_CERT, "sslDN");
attributes.put("my-attr", "hello");
attributes.put("my-blank-attr", "");
attributes.put("status.code", "201");
runner.enqueue("hello".getBytes(), attributes);
runner.run();
runner.assertAllFlowFilesTransferred(HandleHttpResponse.REL_SUCCESS, 1);
assertTrue(runner.getProvenanceEvents().size() == 1);
assertEquals(ProvenanceEventType.SEND, runner.getProvenanceEvents().get(0).getEventType());
assertEquals("https://client@server:8443/test", runner.getProvenanceEvents().get(0).getTransitUri());
assertEquals("hello", contextMap.baos.toString());
assertEquals("hello", contextMap.headersSent.get("my-attr"));
assertNull(contextMap.headersSent.get("my-blank-attr"));
assertEquals(201, contextMap.statusCode);
assertEquals(1, contextMap.getCompletionCount());
assertTrue(contextMap.headersWithNoValue.isEmpty());
}
@Test
public void testWithExceptionThrown() throws InitializationException {
final TestRunner runner = TestRunners.newTestRunner(HandleHttpResponse.class);