<commit_before>from __future__ import absolute_import

import opentracing
from ...log import logger
from ...singletons import tracer

try:
    import celery
    from celery import registry, signals
    from .catalog import task_catalog_get, task_catalog_pop, task_catalog_push, get_task_id

    try:
        from urllib import parse
    except ImportError:
        import urlparse as parse
        import urllib

    def add_broker_tags(span, broker_url):
        try:
            url = parse.urlparse(broker_url)
            span.set_tag("scheme", url.scheme)

            if url.hostname is None:
                span.set_tag("host", 'localhost')
            else:
                span.set_tag("host", url.hostname)

            if url.port is None:
                # Set default port if not specified
                if url.scheme == 'redis':
                    span.set_tag("port", "6379")
                elif 'amqp' in url.scheme:
                    span.set_tag("port", "5672")
                elif 'sqs' in url.scheme:
                    span.set_tag("port", "443")
            else:
                span.set_tag("port", str(url.port))
        except:
            logger.debug("Error parsing broker URL: %s" % broker_url, exc_info=True)

    @signals.task_prerun.connect
    def task_prerun(*args, **kwargs):
        try:
            ctx = None
            task = kwargs.get('sender', None)
            task_id = kwargs.get('task_id', None)
            task = registry.tasks.get(task.name)

            headers = task.request.get('headers', {})
            if headers is not None:
                ctx = tracer.extract(opentracing.Format.HTTP_HEADERS, headers)

            scope = tracer.start_active_span("celery-worker", child_of=ctx)
            scope.span.set_tag("task", task.name)
            scope.span.set_tag("task_id", task_id)<commit_msg>Celery: Add tag parsing safety (#263)<commit_after>from __future__ import absolute_import

import opentracing
from ...log import logger
from ...singletons import tracer

try:
    import celery
    from celery import registry, signals
    from .catalog import task_catalog_get, task_catalog_pop, task_catalog_push, get_task_id

    try:
        from urllib import parse
    except ImportError:
        import urlparse as parse
        import urllib

    def add_broker_tags(span, broker_url):
        try:
            url = parse.urlparse(broker_url)

            # Add safety for edge case where scheme may not be a string
            url_scheme = str(url.scheme)
            span.set_tag("scheme", url_scheme)

            if url.hostname is None:
                span.set_tag("host", 'localhost')
            else:
                span.set_tag("host", url.hostname)

            if url.port is None:
                # Set default port if not specified
                if url_scheme == 'redis':
                    span.set_tag("port", "6379")
                elif 'amqp' in url_scheme:
                    span.set_tag("port", "5672")
                elif 'sqs' in url_scheme:
                    span.set_tag("port", "443")
            else:
                span.set_tag("port", str(url.port))
        except Exception:
            logger.debug("Error parsing broker URL: %s" % broker_url, exc_info=True)

    @signals.task_prerun.connect
    def task_prerun(*args, **kwargs):
        try:
            ctx = None
            task = kwargs.get('sender', None)
            task_id = kwargs.get('task_id', None)
            task = registry.tasks.get(task.name)

            headers = task.request.get('headers', {})
            if headers is not None:
                ctx = tracer.extract(opentracing.Format.HTTP_HEADERS, headers)

            scope = tracer.start_active_span("celery-worker", child_of=ctx)
            scope.span.set_tag("task", task.name)
            scope.span.set_tag("task_id", task_id)
{'fraction_non_alphanumeric': 0.09798345398138573, 'fraction_numerical': 0.006463288521199586, 'mean_word_length': 1.8680504077094144, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '3505466', 'n_tokens_mistral': 1144, 'n_tokens_neox': 1084, 'n_words': 277}