def parse_response(self, url: str, status_code: int, headers, body: bytes) -> dict:
aws_request_id = headers.get("x-amzn-RequestId")
apigw_request_id = headers.get("Apigw-Requestid")
message_id = aws_request_id or apigw_request_id
return recursive_json_join(
{"info": {"messageId": message_id}},
super().parse_response(url, status_code, headers, body),
)
def add_request_event(self, parse_params: HttpRequest):
"""
This function parses an request event and add it to the span.
"""
parser = get_parser(parse_params.host)()
msg = parser.parse_request(parse_params)
self.previous_request = parse_params.headers, parse_params.body
self.http_spans.append(recursive_json_join(msg, self.base_msg))
self.http_span_ids_to_send.add(msg["id"])
def parse_response(self, url: str, status_code: int, headers, body: bytes) -> dict:
return recursive_json_join(
{
"info": {
"messageId": safe_key_from_xml(body, "PublishResponse/PublishResult/MessageId")
}
},
super().parse_response(url, status_code, headers, body),
)
def add_step_end_event(self, ret_val):
message_id = str(uuid.uuid4())
step_function_span = StepFunctionParser().create_span(message_id)
self.http_spans.append(
recursive_json_join(step_function_span, self.base_msg))
self.http_span_ids_to_send.add(step_function_span["id"])
if isinstance(ret_val, dict):
ret_val[LUMIGO_EVENT_KEY] = {STEP_FUNCTION_UID_KEY: message_id}
get_logger().debug(
f"Added key {LUMIGO_EVENT_KEY} to the user's return value")
def parse_request(self, parse_params: HttpRequest) -> dict:
return recursive_json_join(
{
"name": safe_split_get(
str(parse_params.headers.get("path", "")), "/", 3 # type: ignore
),
"invocationType": parse_params.headers.get("x-amz-invocation-type"), # type: ignore
},
super().parse_request(parse_params),
)
def parse_request(self, parse_params: HttpRequest) -> dict:
return recursive_json_join(
{
"info": {
"resourceName": safe_key_from_query(parse_params.body, "TopicArn"),
"targetArn": safe_key_from_query(parse_params.body, "TopicArn"),
}
},
super().parse_request(parse_params),
)
def parse_response(self, url: str, status_code: int, headers, body: bytes) -> dict:
additional_info = {}
message_id = headers.get("x-amzn-RequestId")
if message_id and self.should_add_message_id:
additional_info["info"] = {"messageId": message_id}
span_id = headers.get("x-amzn-requestid") or headers.get("x-amz-requestid")
if span_id:
additional_info["id"] = span_id
return recursive_json_join(
additional_info, super().parse_response(url, status_code, headers, body)
)
def create_span(self, message_id: str) -> dict:
return recursive_json_join(
{
"info": {
"resourceName": "StepFunction",
"httpInfo": {"host": "StepFunction"},
"messageId": message_id,
}
},
super().parse_request(None), # type: ignore
)
def parse_request(self, parse_params: HttpRequest) -> dict:
target: str = str(parse_params.headers.get("x-amz-target", "")) # type: ignore
return recursive_json_join(
{
"info": {
"resourceName": safe_key_from_json(parse_params.body, "TableName"),
"dynamodbMethod": safe_split_get(target, ".", 1),
}
},
super().parse_request(parse_params),
)
def update_event_response(self, host: Optional[str], status_code: int,
headers: http.client.HTTPMessage,
body: bytes) -> None:
"""
:param host: If None, use the host from the last span, otherwise this is the first chuck and we can empty
the aggregated response body
This function assumes synchronous execution - we update the last http event.
"""
if self.http_spans:
last_event = self.http_spans.pop()
if not host:
host = last_event.get("info",
{}).get("httpInfo",
{}).get("host", "unknown")
else:
self.previous_response_body = b""
parser = get_parser(host, headers)() # type: ignore
self.previous_response_body += body
update = parser.parse_response( # type: ignore
host, status_code, headers, self.previous_response_body)
self.http_spans.append(recursive_json_join(update, last_event))
self.http_span_ids_to_send.add(
update.get("id") or last_event["id"])
def __init__(
self,
name: str = None,
started: int = None,
region: str = None,
runtime: str = None,
memory_allocated: str = None,
log_stream_name: str = None,
log_group_name: str = None,
trace_root: str = None,
transaction_id: str = None,
request_id: str = None,
account: str = None,
trace_id_suffix: str = None,
trigger_by: dict = None,
max_finish_time: int = None,
event: str = None,
envs: str = None,
):
version = open(
_VERSION_PATH,
"r").read() if os.path.exists(_VERSION_PATH) else "unknown"
version = version.strip()
self.name = name
self.region = region
self.trace_root = trace_root
self.trace_id_suffix = trace_id_suffix
self.transaction_id = transaction_id
self.max_finish_time = max_finish_time
self.base_msg = {
"started": started,
"transactionId": transaction_id,
"account": account,
"region": region,
"parentId": request_id,
"info": {
"tracer": {
"version": version
},
"traceId": {
"Root": trace_root
}
},
"event": event,
"envs": envs,
}
self.function_span = recursive_json_join(
{
"id": request_id,
"type": "function",
"name": name,
"runtime": runtime,
"memoryAllocated": memory_allocated,
"readiness": "cold" if SpansContainer.is_cold else "warm",
"info": {
"logStreamName": log_stream_name,
"logGroupName": log_group_name,
**(trigger_by or {}),
},
EXECUTION_TAGS_KEY: [],
},
self.base_msg,
)
self.previous_request: Tuple[Optional[http.client.HTTPMessage],
bytes] = (None, b"")
self.previous_response_body: bytes = b""
self.http_span_ids_to_send: Set[str] = set()
self.http_spans: List[Dict] = []
SpansContainer.is_cold = False
def test_recursive_json_join(d1, d2, result):
assert utils.recursive_json_join(d1, d2) == result
def parse_response(self, url: str, status_code: int, headers, body: bytes) -> dict:
return recursive_json_join(
{"info": {"messageId": headers.get("x-amz-request-id")}},
super().parse_response(url, status_code, headers, body),
)
def parse_request(self, parse_params: HttpRequest) -> dict:
return recursive_json_join(
{"info": {"resourceName": safe_split_get(parse_params.host, ".", 0)}},
super().parse_request(parse_params),
)
def parse_response(self, url: str, status_code: int, headers, body: bytes) -> dict:
return recursive_json_join(
{"info": {"messageId": SqsParser._extract_message_id(body)}},
super().parse_response(url, status_code, headers, body),
)
def parse_request(self, parse_params: HttpRequest) -> dict:
return recursive_json_join(
{"info": {"resourceName": safe_key_from_query(parse_params.body, "QueueUrl")}},
super().parse_request(parse_params),
)
