<commit_before># --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------
from __future__ import unicode_literals

import datetime
import calendar
import json
import logging
import six

from uamqp import BatchMessage, Message, types, constants  # type: ignore
from uamqp.message import MessageHeader  # type: ignore

from azure.core.settings import settings # type: ignore

from azure.eventhub.error import EventDataError

log = logging.getLogger(__name__)

# event_data.encoded_size < 255, batch encode overhead is 5, >=256, overhead is 8 each
_BATCH_MESSAGE_OVERHEAD_COST = [5, 8]


def parse_sas_token(sas_token):
    """Parse a SAS token into its components.

    :param sas_token: The SAS token.
    :type sas_token: str
    :rtype: dict[str, str]
    """
    sas_data = {}
    token = sas_token.partition(' ')[2]
    fields = token.split('&')
    for field in fields:
        key, value = field.split('=', 1)
        sas_data[key.lower()] = value
    return sas_data


class EventData(object):
    """
    The EventData class is a holder of event content.

    Example:
        .. literalinclude:: ../examples/test_examples_eventhub.py
            :start-after: [START create_event_data]
            :end-before: [END create_event_data]
            :language: python
            :dedent: 4
            :caption: Create instances of EventData

    """

    PROP_SEQ_NUMBER = b"x-opt-sequence-number"
    PROP_OFFSET = b"x-opt-offset"
    PROP_PARTITION_KEY = b"x-opt-partition-key"
    PROP_PARTITION_KEY_AMQP_SYMBOL = types.AMQPSymbol(PROP_PARTITION_KEY)
    PROP_TIMESTAMP = b"x-opt-enqueued-time"
    PROP_LAST_ENQUEUED_SEQUENCE_NUMBER = b"last_enqueued_sequence_number"
    PROP_LAST_ENQUEUED_OFFSET = b"last_enqueued_offset"
    PROP_LAST_ENQUEUED_TIME_UTC = b"last_enqueued_time_utc"
    PROP_RUNTIME_INFO_RETRIEVAL_TIME_UTC = b"runtime_info_retrieval_time_utc"

    def __init__(self, body=None):
        """
        Initialize EventData.

        :param body: The data to send in a single message.
        :type body: str, bytes or list
        """

        self._last_enqueued_event_properties = {}
        if body and isinstance(body, list):
            self.message = Message(body[0])
            for more in body[1:]:
                self.message._body.append(more)  # pylint: disable=protected-access
        elif body is None:
            raise ValueError("EventData cannot be None.")
        else:
            self.message = Message(body)
        self.message.annotations = {}

    def __str__(self):
        dic = {
            'body': self.body_as_str(),
            'application_properties': str(self.application_properties)
        }

        if self.sequence_number:
            dic['sequence_number'] = str(self.sequence_number)
        if self.offset:
            dic['offset'] = str(self.offset)
        if self.enqueued_time:
            dic['enqueued_time'] = str(self.enqueued_time)
        if self.partition_key:
            dic['partition_key'] = str(self.partition_key)
        return str(dic)

    def _set_partition_key(self, value):
        """
        Set the partition key of the event data object.

        :param value: The partition key to set.
        :type value: str or bytes
        """
        annotations = dict(self.message.annotations)
        annotations[EventData.PROP_PARTITION_KEY_AMQP_SYMBOL] = value
        header = MessageHeader()
        header.durable = True
        self.message.annotations = annotations
        self.message.header = header

    def _trace_message(self, parent_span=None):
        """Add tracing information to this message.

        Will open and close a "Azure.EventHubs.message" span, and
        add the "DiagnosticId" as app properties of the message.
        """
        span_impl_type = settings.tracing_implementation()  # type: Type[AbstractSpan]
        if span_impl_type is not None:
            current_span = parent_span or span_impl_type(span_impl_type.get_current_span())
            message_span = current_span.span(name="Azure.EventHubs.message")
            message_span.start()
            app_prop = dict(self.application_properties)
            app_prop.setdefault(b"Diagnostic-Id", message_span.get_trace_parent().encode('ascii'))
            self.application_properties = app_prop
            message_span.finish()

    def _trace_link_message(self, parent_span=None):
        """Link the current message to current span.

        Will extract DiagnosticId if available.
        """
        span_impl_type = settings.tracing_implementation()  # type: Type[AbstractSpan]
        if span_impl_type is not None:
            current_span = parent_span or span_impl_type(span_impl_type.get_current_span())
            if current_span and self.application_properties:
                traceparent = self.application_properties.get(b"Diagnostic-Id", "").decode('ascii')
                if traceparent:
                    current_span.link(traceparent)

    def _get_last_enqueued_event_properties(self):
        if self._last_enqueued_event_properties:
            return self._last_enqueued_event_properties

        if self.message.delivery_annotations:
            self._last_enqueued_event_properties = {
                "sequence_number":
                    self.message.delivery_annotations.get(EventData.PROP_LAST_ENQUEUED_SEQUENCE_NUMBER, None),
                "offset":
                    self.message.delivery_annotations.get(EventData.PROP_LAST_ENQUEUED_OFFSET, None),
                "enqueued_time":
                    self.message.delivery_annotations.get(EventData.PROP_LAST_ENQUEUED_TIME_UTC, None),
                "retrieval_time":
                    self.message.delivery_annotations.get(EventData.PROP_RUNTIME_INFO_RETRIEVAL_TIME_UTC, None)
            }
            return self._last_enqueued_event_properties

        return None

    @classmethod
    def _from_message(cls, message):
        # pylint:disable=protected-access
        event_data = cls(body='')
        event_data.message = message
        return event_data

    @property
    def sequence_number(self):
        """
        The sequence number of the event data object.

        :rtype: int or long
        """
        return self.message.annotations.get(EventData.PROP_SEQ_NUMBER, None)

    @property
    def offset(self):
        """
        The offset of the event data object.

        :rtype: str
        """<commit_msg>[EventHubs] returning datetime instead of raw timestamp (#7699)<commit_after># --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------
from __future__ import unicode_literals

import datetime
import calendar
import json
import logging
import six

from uamqp import BatchMessage, Message, types, constants  # type: ignore
from uamqp.message import MessageHeader  # type: ignore

from azure.core.settings import settings # type: ignore

from azure.eventhub.error import EventDataError

log = logging.getLogger(__name__)

# event_data.encoded_size < 255, batch encode overhead is 5, >=256, overhead is 8 each
_BATCH_MESSAGE_OVERHEAD_COST = [5, 8]


def parse_sas_token(sas_token):
    """Parse a SAS token into its components.

    :param sas_token: The SAS token.
    :type sas_token: str
    :rtype: dict[str, str]
    """
    sas_data = {}
    token = sas_token.partition(' ')[2]
    fields = token.split('&')
    for field in fields:
        key, value = field.split('=', 1)
        sas_data[key.lower()] = value
    return sas_data


class EventData(object):
    """
    The EventData class is a holder of event content.

    Example:
        .. literalinclude:: ../examples/test_examples_eventhub.py
            :start-after: [START create_event_data]
            :end-before: [END create_event_data]
            :language: python
            :dedent: 4
            :caption: Create instances of EventData

    """

    PROP_SEQ_NUMBER = b"x-opt-sequence-number"
    PROP_OFFSET = b"x-opt-offset"
    PROP_PARTITION_KEY = b"x-opt-partition-key"
    PROP_PARTITION_KEY_AMQP_SYMBOL = types.AMQPSymbol(PROP_PARTITION_KEY)
    PROP_TIMESTAMP = b"x-opt-enqueued-time"
    PROP_LAST_ENQUEUED_SEQUENCE_NUMBER = b"last_enqueued_sequence_number"
    PROP_LAST_ENQUEUED_OFFSET = b"last_enqueued_offset"
    PROP_LAST_ENQUEUED_TIME_UTC = b"last_enqueued_time_utc"
    PROP_RUNTIME_INFO_RETRIEVAL_TIME_UTC = b"runtime_info_retrieval_time_utc"

    def __init__(self, body=None):
        """
        Initialize EventData.

        :param body: The data to send in a single message.
        :type body: str, bytes or list
        """

        self._last_enqueued_event_properties = {}
        if body and isinstance(body, list):
            self.message = Message(body[0])
            for more in body[1:]:
                self.message._body.append(more)  # pylint: disable=protected-access
        elif body is None:
            raise ValueError("EventData cannot be None.")
        else:
            self.message = Message(body)
        self.message.annotations = {}

    def __str__(self):
        dic = {
            'body': self.body_as_str(),
            'application_properties': str(self.application_properties)
        }

        if self.sequence_number:
            dic['sequence_number'] = str(self.sequence_number)
        if self.offset:
            dic['offset'] = str(self.offset)
        if self.enqueued_time:
            dic['enqueued_time'] = str(self.enqueued_time)
        if self.partition_key:
            dic['partition_key'] = str(self.partition_key)
        return str(dic)

    def _set_partition_key(self, value):
        """
        Set the partition key of the event data object.

        :param value: The partition key to set.
        :type value: str or bytes
        """
        annotations = dict(self.message.annotations)
        annotations[EventData.PROP_PARTITION_KEY_AMQP_SYMBOL] = value
        header = MessageHeader()
        header.durable = True
        self.message.annotations = annotations
        self.message.header = header

    def _trace_message(self, parent_span=None):
        """Add tracing information to this message.

        Will open and close a "Azure.EventHubs.message" span, and
        add the "DiagnosticId" as app properties of the message.
        """
        span_impl_type = settings.tracing_implementation()  # type: Type[AbstractSpan]
        if span_impl_type is not None:
            current_span = parent_span or span_impl_type(span_impl_type.get_current_span())
            message_span = current_span.span(name="Azure.EventHubs.message")
            message_span.start()
            app_prop = dict(self.application_properties)
            app_prop.setdefault(b"Diagnostic-Id", message_span.get_trace_parent().encode('ascii'))
            self.application_properties = app_prop
            message_span.finish()

    def _trace_link_message(self, parent_span=None):
        """Link the current message to current span.

        Will extract DiagnosticId if available.
        """
        span_impl_type = settings.tracing_implementation()  # type: Type[AbstractSpan]
        if span_impl_type is not None:
            current_span = parent_span or span_impl_type(span_impl_type.get_current_span())
            if current_span and self.application_properties:
                traceparent = self.application_properties.get(b"Diagnostic-Id", "").decode('ascii')
                if traceparent:
                    current_span.link(traceparent)

    def _get_last_enqueued_event_properties(self):
        if self._last_enqueued_event_properties:
            return self._last_enqueued_event_properties

        if self.message.delivery_annotations:
            enqueued_time_stamp = \
                self.message.delivery_annotations.get(EventData.PROP_LAST_ENQUEUED_TIME_UTC, None)
            retrieval_time_stamp = \
                self.message.delivery_annotations.get(EventData.PROP_RUNTIME_INFO_RETRIEVAL_TIME_UTC, None)

            self._last_enqueued_event_properties = {
                "sequence_number":
                    self.message.delivery_annotations.get(EventData.PROP_LAST_ENQUEUED_SEQUENCE_NUMBER, None),
                "offset":
                    self.message.delivery_annotations.get(EventData.PROP_LAST_ENQUEUED_OFFSET, None),
                "enqueued_time":
                    datetime.datetime.utcfromtimestamp(
                        float(enqueued_time_stamp)/1000) if enqueued_time_stamp else None,
                "retrieval_time":
                    datetime.datetime.utcfromtimestamp(
                        float(retrieval_time_stamp)/1000) if retrieval_time_stamp else None
            }
            return self._last_enqueued_event_properties

        return None

    @classmethod
    def _from_message(cls, message):
        # pylint:disable=protected-access
        event_data = cls(body='')
        event_data.message = message
        return event_data

    @property
    def sequence_number(self):
        """
        The sequence number of the event data object.

        :rtype: int or long
        """
        return self.message.annotations.get(EventData.PROP_SEQ_NUMBER, None)

    @property
    def offset(self):
        """
        The offset of the event data object.

        :rtype: str
        """
{'fraction_non_alphanumeric': 0.10926864590876177, 'fraction_numerical': 0.0030412744388124547, 'mean_word_length': 2.328753916606411, 'pattern_counts': {'":': 8, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3397178', 'n_tokens_mistral': 4000, 'n_tokens_neox': 3865, 'n_words': 875}