Skip to content

Commit

Permalink
Moved sqs header classes
Browse files Browse the repository at this point in the history
  • Loading branch information
obenkenobi committed Jan 21, 2025
1 parent 9dd5005 commit d68e95d
Show file tree
Hide file tree
Showing 10 changed files with 377 additions and 508 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -7,30 +7,27 @@

package com.amazonaws.services.sqs;

import com.amazonaws.services.sqs.model.MessageAttributeValue;
import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
import com.amazonaws.services.sqs.model.ReceiveMessageResult;
import com.amazonaws.services.sqs.model.SendMessageBatchRequest;
import com.amazonaws.services.sqs.model.SendMessageBatchRequestEntry;
import com.amazonaws.services.sqs.model.SendMessageBatchResult;
import com.amazonaws.services.sqs.model.SendMessageRequest;
import com.amazonaws.services.sqs.model.SendMessageResult;
import com.newrelic.api.agent.HeaderType;
import com.newrelic.api.agent.Headers;
import com.newrelic.api.agent.MessageConsumeParameters;
import com.newrelic.api.agent.MessageProduceParameters;
import com.newrelic.api.agent.NewRelic;
import com.newrelic.api.agent.Trace;
import com.newrelic.api.agent.weaver.MatchType;
import com.newrelic.api.agent.weaver.Weave;
import com.newrelic.api.agent.weaver.Weaver;
import com.newrelic.utils.SQSBatchRequestHeaders;
import com.newrelic.utils.SQSRequestHeaders;
import com.newrelic.utils.SqsV1Util;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;

@Weave(type = MatchType.Interface, originalName = "com.amazonaws.services.sqs.AmazonSQS")
public class AmazonSQS_Instrumentation {
Expand Down Expand Up @@ -68,126 +65,4 @@ public ReceiveMessageResult receiveMessage(ReceiveMessageRequest receiveMessageR
return Weaver.callOriginal();
}

public static class SQSRequestHeaders implements Headers {
private SendMessageRequest request = null;
public SQSRequestHeaders(SendMessageRequest req) {
request = req;
}
@Override
public HeaderType getHeaderType() {
return HeaderType.MESSAGE;
}
@Override
public String getHeader(String name) {
Map<String, MessageAttributeValue> messageAttributes = request.getMessageAttributes();
if(messageAttributes != null) {
MessageAttributeValue value = messageAttributes.get(name);
if (value != null && value.getDataType().equalsIgnoreCase("string")) {
return value.getStringValue();
}
}
Map<String, String> customRequestHeaders = request.getCustomRequestHeaders();
if(customRequestHeaders != null) {
return customRequestHeaders.get(name);
}
return null;
}
@Override
public Collection<String> getHeaders(String name) {
String value = getHeader(name);
if(value != null) {
return Collections.singletonList(value);
}
return Collections.emptyList();
}
@Override
public void setHeader(String name, String value) {
if(request != null) {
Map<String, MessageAttributeValue> existingAttributes = request.getMessageAttributes();
if(!existingAttributes.containsKey(name)) {
request.addMessageAttributesEntry(name, new MessageAttributeValue().withDataType("String").withStringValue(value));
}
}
}
@Override
public void addHeader(String name, String value) {
if(request != null) {
Map<String, MessageAttributeValue> existingAttributes = request.getMessageAttributes();
if(!existingAttributes.containsKey(name)) {
request.addMessageAttributesEntry(name, new MessageAttributeValue().withDataType("String").withStringValue(value));
}
}
}
@Override
public Collection<String> getHeaderNames() {
Map<String, MessageAttributeValue> messageAttributes = request.getMessageAttributes();
return messageAttributes.keySet();
}
@Override
public boolean containsHeader(String name) {
return getHeaderNames().contains(name);
}
}

public static class SQSBatchRequestHeaders implements Headers {
private SendMessageBatchRequestEntry requestEntry = null;
public SQSBatchRequestHeaders(SendMessageBatchRequestEntry re) {
requestEntry = re;
}
@Override
public HeaderType getHeaderType() {
return HeaderType.MESSAGE;
}
@Override
public String getHeader(String name) {
Map<String, MessageAttributeValue> attributes = requestEntry.getMessageAttributes();
if(attributes != null) {
MessageAttributeValue value = attributes.get(name);
if(value != null) {
String dataType = value.getDataType();
if(dataType.equalsIgnoreCase("String")) {
String stringValue = value.getStringValue();
if(stringValue != null) return stringValue;
}
}
}
return null;
}
@Override
public Collection<String> getHeaders(String name) {
List<String> list = new ArrayList<String>();
String value = getHeader(name);
if(value != null && !value.isEmpty()) {
list.add(value);
}
return list;
}
@Override
public void setHeader(String name, String value) {
if(requestEntry != null) {
requestEntry.addMessageAttributesEntry(name, new MessageAttributeValue().withDataType("String").withStringValue(value));
}
}
@Override
public void addHeader(String name, String value) {
if(requestEntry != null) {
requestEntry.addMessageAttributesEntry(name, new MessageAttributeValue().withDataType("String").withStringValue(value));
}
}
@Override
public Collection<String> getHeaderNames() {
if(requestEntry != null) {
Map<String, MessageAttributeValue> attributes = requestEntry.getMessageAttributes();
if(attributes != null) {
return attributes.keySet();
}
}
return Collections.emptyList();
}
@Override
public boolean containsHeader(String name) {
return getHeaderNames().contains(name);
}
}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
package com.newrelic.utils;

import com.amazonaws.services.sqs.model.MessageAttributeValue;
import com.amazonaws.services.sqs.model.SendMessageBatchRequestEntry;
import com.newrelic.api.agent.HeaderType;
import com.newrelic.api.agent.Headers;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;

public class SQSBatchRequestHeaders implements Headers {
private final SendMessageBatchRequestEntry requestEntry;

public SQSBatchRequestHeaders(SendMessageBatchRequestEntry re) {
requestEntry = re;
}

@Override
public HeaderType getHeaderType() {
return HeaderType.MESSAGE;
}

@Override
public String getHeader(String name) {
Map<String, MessageAttributeValue> attributes = requestEntry.getMessageAttributes();
if (attributes != null) {
MessageAttributeValue value = attributes.get(name);
if (value != null) {
String dataType = value.getDataType();
if (dataType.equalsIgnoreCase("String")) {
String stringValue = value.getStringValue();
if (stringValue != null) {
return stringValue;
}
}
}
}
return null;
}

@Override
public Collection<String> getHeaders(String name) {
List<String> list = new ArrayList<String>();
String value = getHeader(name);
if (value != null && !value.isEmpty()) {
list.add(value);
}
return list;
}

@Override
public void setHeader(String name, String value) {
if (requestEntry != null) {
requestEntry.addMessageAttributesEntry(name, new MessageAttributeValue().withDataType("String").withStringValue(value));
}
}

@Override
public void addHeader(String name, String value) {
if (requestEntry != null) {
requestEntry.addMessageAttributesEntry(name, new MessageAttributeValue().withDataType("String").withStringValue(value));
}
}

@Override
public Collection<String> getHeaderNames() {
if (requestEntry != null) {
Map<String, MessageAttributeValue> attributes = requestEntry.getMessageAttributes();
if (attributes != null) {
return attributes.keySet();
}
}
return Collections.emptyList();
}

@Override
public boolean containsHeader(String name) {
return getHeaderNames().contains(name);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
package com.newrelic.utils;

import com.amazonaws.services.sqs.model.MessageAttributeValue;
import com.amazonaws.services.sqs.model.SendMessageRequest;
import com.newrelic.api.agent.HeaderType;
import com.newrelic.api.agent.Headers;

import java.util.Collection;
import java.util.Collections;
import java.util.Map;

public class SQSRequestHeaders implements Headers {
private final SendMessageRequest request;

public SQSRequestHeaders(SendMessageRequest req) {
request = req;
}

@Override
public HeaderType getHeaderType() {
return HeaderType.MESSAGE;
}

@Override
public String getHeader(String name) {
Map<String, MessageAttributeValue> messageAttributes = request.getMessageAttributes();
if (messageAttributes != null) {
MessageAttributeValue value = messageAttributes.get(name);
if (value != null && value.getDataType().equalsIgnoreCase("string")) {
return value.getStringValue();
}
}
Map<String, String> customRequestHeaders = request.getCustomRequestHeaders();
if (customRequestHeaders != null) {
return customRequestHeaders.get(name);
}
return null;
}

@Override
public Collection<String> getHeaders(String name) {
String value = getHeader(name);
if (value != null) {
return Collections.singletonList(value);
}
return Collections.emptyList();
}

@Override
public void setHeader(String name, String value) {
if (request != null) {
Map<String, MessageAttributeValue> existingAttributes = request.getMessageAttributes();
if (!existingAttributes.containsKey(name)) {
request.addMessageAttributesEntry(name, new MessageAttributeValue().withDataType("String").withStringValue(value));
}
}
}

@Override
public void addHeader(String name, String value) {
if (request != null) {
Map<String, MessageAttributeValue> existingAttributes = request.getMessageAttributes();
if (!existingAttributes.containsKey(name)) {
request.addMessageAttributesEntry(name, new MessageAttributeValue().withDataType("String").withStringValue(value));
}
}
}

@Override
public Collection<String> getHeaderNames() {
Map<String, MessageAttributeValue> messageAttributes = request.getMessageAttributes();
return messageAttributes.keySet();
}

@Override
public boolean containsHeader(String name) {
return getHeaderNames().contains(name);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
import static org.junit.Assert.assertTrue;

@RunWith(InstrumentationTestRunner.class)
@InstrumentationTestConfig(includePrefixes = { "com.amazonaws.services.sqs" }, configName = "dt_enabled.yml")
@InstrumentationTestConfig(includePrefixes = { "com.amazonaws.services.sqs", "com.newrelic.utils" }, configName = "dt_enabled.yml")
public class SqsClientTest {

private static AmazonSQSClient sqsClient;
Expand Down
Loading

0 comments on commit d68e95d

Please sign in to comment.