aws_xray_sdk/core/models/entity.py

import logging
import os
import binascii
import time
import string

import json

from ..utils.compat import annotation_value_types
from ..utils.conversion import metadata_to_dict
from .throwable import Throwable
from . import http
from ..exceptions.exceptions import AlreadyEndedException

log = logging.getLogger(__name__)

# Valid characters can be found at http://docs.aws.amazon.com/xray/latest/devguide/xray-api-segmentdocuments.html
_common_invalid_name_characters = '?;*()!$~^<>'
_valid_annotation_key_characters = string.ascii_letters + string.digits + '_'

ORIGIN_TRACE_HEADER_ATTR_KEY = '_origin_trace_header'


class Entity:
    """
    The parent class for segment/subsegment. It holds common properties
    and methods on segment and subsegment.
    """

    def __init__(self, name, entity_id=None):
        if not entity_id:
            self.id = self._generate_random_id()
        else:
            self.id = entity_id

        # required attributes
        self.name = name
        self.name = ''.join([c for c in name if c not in _common_invalid_name_characters])
        self.start_time = time.time()
        self.parent_id = None

        if self.name != name:
            log.warning("Removing Segment/Subsugment Name invalid characters from {}.".format(name))

        # sampling
        self.sampled = True

        # state
        self.in_progress = True

        # meta fields
        self.http = {}
        self.annotations = {}
        self.metadata = {}
        self.aws = {}
        self.cause = {}

        # child subsegments
        # list is thread-safe
        self.subsegments = []

    def close(self, end_time=None):
        """
        Close the trace entity by setting `end_time`
        and flip the in progress flag to False.

        :param float end_time: Epoch in seconds. If not specified
            current time will be used.
        """
        self._check_ended()

        if end_time:
            self.end_time = end_time
        else:
            self.end_time = time.time()
        self.in_progress = False

    def add_subsegment(self, subsegment):
        """
        Add input subsegment as a child subsegment.
        """
        self._check_ended()
        subsegment.parent_id = self.id

        if not self.sampled and subsegment.sampled:
            log.warning("This sampled subsegment is being added to an unsampled parent segment/subsegment and will be orphaned.")

        self.subsegments.append(subsegment)

    def remove_subsegment(self, subsegment):
        """
        Remove input subsegment from child subsegments.
        """
        self.subsegments.remove(subsegment)

    def put_http_meta(self, key, value):
        """
        Add http related metadata.

        :param str key: Currently supported keys are:
            * url
            * method
            * user_agent
            * client_ip
            * status
            * content_length
        :param value: status and content_length are int and for other
            supported keys string should be used.
        """
        self._check_ended()

        if value is None:
            return

        if key == http.STATUS:
            if isinstance(value, str):
                value = int(value)
            self.apply_status_code(value)

        if key in http.request_keys:
            if 'request' not in self.http:
                self.http['request'] = {}
            self.http['request'][key] = value
        elif key in http.response_keys:
            if 'response' not in self.http:
                self.http['response'] = {}
            self.http['response'][key] = value
        else:
            log.warning("ignoring unsupported key %s in http meta.", key)

    def put_annotation(self, key, value):
        """
        Annotate segment or subsegment with a key-value pair.
        Annotations will be indexed for later search query.

        :param str key: annotation key
        :param object value: annotation value. Any type other than
            string/number/bool will be dropped
        """
        self._check_ended()

        if not isinstance(key, str):
            log.warning("ignoring non string type annotation key with type %s.", type(key))
            return

        if not isinstance(value, annotation_value_types):
            log.warning("ignoring unsupported annotation value type %s.", type(value))
            return

        if any(character not in _valid_annotation_key_characters for character in key):
            log.warning("ignoring annnotation with unsupported characters in key: '%s'.", key)
            return

        self.annotations[key] = value

    def put_metadata(self, key, value, namespace='default'):
        """
        Add metadata to segment or subsegment. Metadata is not indexed
        but can be later retrieved by BatchGetTraces API.

        :param str namespace: optional. Default namespace is `default`.
            It must be a string and prefix `AWS.` is reserved.
        :param str key: metadata key under specified namespace
        :param object value: any object that can be serialized into JSON string
        """
        self._check_ended()

        if not isinstance(namespace, str):
            log.warning("ignoring non string type metadata namespace")
            return

        if namespace.startswith('AWS.'):
            log.warning("Prefix 'AWS.' is reserved, drop metadata with namespace %s", namespace)
            return

        if self.metadata.get(namespace, None):
            self.metadata[namespace][key] = value
        else:
            self.metadata[namespace] = {key: value}

    def set_aws(self, aws_meta):
        """
        set aws section of the entity.
        This method is called by global recorder and botocore patcher
        to provide additonal information about AWS runtime.
        It is not recommended to manually set aws section.
        """
        self._check_ended()
        self.aws = aws_meta

    def add_throttle_flag(self):
        self.throttle = True

    def add_fault_flag(self):
        self.fault = True

    def add_error_flag(self):
        self.error = True

    def apply_status_code(self, status_code):
        """
        When a trace entity is generated under the http context,
        the status code will affect this entity's fault/error/throttle flags.
        Flip these flags based on status code.
        """
        self._check_ended()
        if not status_code:
            return

        if status_code >= 500:
            self.add_fault_flag()
        elif status_code == 429:
            self.add_throttle_flag()
            self.add_error_flag()
        elif status_code >= 400:
            self.add_error_flag()

    def add_exception(self, exception, stack, remote=False):
        """
        Add an exception to trace entities.

        :param Exception exception: the caught exception.
        :param list stack: the output from python built-in
            `traceback.extract_stack()`.
        :param bool remote: If False it means it's a client error
            instead of a downstream service.
        """
        self._check_ended()
        self.add_fault_flag()

        if hasattr(exception, '_recorded'):
            setattr(self, 'cause', getattr(exception, '_cause_id'))
            return

        if not isinstance(self.cause, dict):
            log.warning("The current cause object is not a dict but an id: {}. Resetting the cause and recording the "
                        "current exception".format(self.cause))
            self.cause = {}

        if 'exceptions' in self.cause:
            exceptions = self.cause['exceptions']
        else:
            exceptions = []

        exceptions.append(Throwable(exception, stack, remote))

        self.cause['exceptions'] = exceptions
        self.cause['working_directory'] = os.getcwd()

    def save_origin_trace_header(self, trace_header):
        """
        Temporarily store additional data fields in trace header
        to the entity for later propagation. The data will be
        cleaned up upon serialization.
        """
        setattr(self, ORIGIN_TRACE_HEADER_ATTR_KEY, trace_header)

    def get_origin_trace_header(self):
        """
        Retrieve saved trace header data.
        """
        return getattr(self, ORIGIN_TRACE_HEADER_ATTR_KEY, None)

    def serialize(self):
        """
        Serialize to JSON document that can be accepted by the
        X-Ray backend service. It uses json to perform serialization.
        """
        return json.dumps(self.to_dict(), default=str)

    def to_dict(self):
        """
        Convert Entity(Segment/Subsegment) object to dict
        with required properties that have non-empty values.
        """
        entity_dict = {}

        for key, value in vars(self).items():
            if isinstance(value, bool) or value:
                if key == 'subsegments':
                    # child subsegments are stored as List
                    subsegments = []
                    for subsegment in value:
                        subsegments.append(subsegment.to_dict())
                    entity_dict[key] = subsegments
                elif key == 'cause':
                    if isinstance(self.cause, dict):
                        entity_dict[key] = {}
                        entity_dict[key]['working_directory'] = self.cause['working_directory']
                        # exceptions are stored as List
                        throwables = []
                        for throwable in value['exceptions']:
                            throwables.append(throwable.to_dict())
                        entity_dict[key]['exceptions'] = throwables
                    else:
                        entity_dict[key] = self.cause
                elif key == 'metadata':
                    entity_dict[key] = metadata_to_dict(value)
                elif key != 'sampled' and key != ORIGIN_TRACE_HEADER_ATTR_KEY:
                    entity_dict[key] = value

        return entity_dict

    def _check_ended(self):
        if not self.in_progress:
            raise AlreadyEndedException("Already ended segment and subsegment cannot be modified.")

    def _generate_random_id(self):
        """
        Generate a random 16-digit hex str.
        This is used for generating segment/subsegment id.
        """
        return binascii.b2a_hex(os.urandom(8)).decode('utf-8')