from __future__ import absolute_import

from celery.task import current
from functools import wraps

from sentry.celery import app
from sentry.utils import metrics


def instrumented_task(name, stat_suffix=None, **kwargs):
    def wrapped(func):
        @wraps(func)
        def _wrapped(*args, **kwargs):
            key = 'jobs.duration'
            if stat_suffix:
                instance = '{}.{}'.format(name, stat_suffix(*args, **kwargs))
            else:
                instance = name
            with metrics.timer(key, instance=instance):
                result = func(*args, **kwargs)
            return result
        return app.task(name=name, **kwargs)(_wrapped)
    return wrapped


def retry(func):
    @wraps(func)
    def wrapped(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as exc:
            current.retry(exc=exc)
    return wrapped

from __future__ import absolute_import

from celery.task import current
from raven.contrib.django.models import client as Raven
from functools import wraps

from sentry.celery import app
from sentry.utils import metrics


def instrumented_task(name, stat_suffix=None, **kwargs):
    def wrapped(func):
        @wraps(func)
        def _wrapped(*args, **kwargs):
            key = 'jobs.duration'
            if stat_suffix:
                instance = '{}.{}'.format(name, stat_suffix(*args, **kwargs))
            else:
                instance = name
            with metrics.timer(key, instance=instance):
                try:
                    result = func(*args, **kwargs)
                finally:
                    Raven.context.clear()
            return result
        return app.task(name=name, **kwargs)(_wrapped)
    return wrapped


def retry(func):
    @wraps(func)
    def wrapped(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as exc:
            current.retry(exc=exc)
    return wrapped

Clear context for celery tasks
Clear context for celery tasks

  from __future__ import absolute_import
  
  from celery.task import current
+ from raven.contrib.django.models import client as Raven
  from functools import wraps
  
  from sentry.celery import app
  from sentry.utils import metrics
  
  
  def instrumented_task(name, stat_suffix=None, **kwargs):
      def wrapped(func):
          @wraps(func)
          def _wrapped(*args, **kwargs):
              key = 'jobs.duration'
              if stat_suffix:
                  instance = '{}.{}'.format(name, stat_suffix(*args, **kwargs))
              else:
                  instance = name
              with metrics.timer(key, instance=instance):
+                 try:
-                 result = func(*args, **kwargs)
+                     result = func(*args, **kwargs)
+                 finally:
+                     Raven.context.clear()
              return result
          return app.task(name=name, **kwargs)(_wrapped)
      return wrapped
  
  
  def retry(func):
      @wraps(func)
      def wrapped(*args, **kwargs):
          try:
              return func(*args, **kwargs)
          except Exception as exc:
              current.retry(exc=exc)
      return wrapped
  
Clear context for celery tasks
## Code Before:
from __future__ import absolute_import

from celery.task import current
from functools import wraps

from sentry.celery import app
from sentry.utils import metrics


def instrumented_task(name, stat_suffix=None, **kwargs):
    def wrapped(func):
        @wraps(func)
        def _wrapped(*args, **kwargs):
            key = 'jobs.duration'
            if stat_suffix:
                instance = '{}.{}'.format(name, stat_suffix(*args, **kwargs))
            else:
                instance = name
            with metrics.timer(key, instance=instance):
                result = func(*args, **kwargs)
            return result
        return app.task(name=name, **kwargs)(_wrapped)
    return wrapped


def retry(func):
    @wraps(func)
    def wrapped(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as exc:
            current.retry(exc=exc)
    return wrapped

## Instruction:
Clear context for celery tasks
## Code After:
from __future__ import absolute_import

from celery.task import current
from raven.contrib.django.models import client as Raven
from functools import wraps

from sentry.celery import app
from sentry.utils import metrics


def instrumented_task(name, stat_suffix=None, **kwargs):
    def wrapped(func):
        @wraps(func)
        def _wrapped(*args, **kwargs):
            key = 'jobs.duration'
            if stat_suffix:
                instance = '{}.{}'.format(name, stat_suffix(*args, **kwargs))
            else:
                instance = name
            with metrics.timer(key, instance=instance):
                try:
                    result = func(*args, **kwargs)
                finally:
                    Raven.context.clear()
            return result
        return app.task(name=name, **kwargs)(_wrapped)
    return wrapped


def retry(func):
    @wraps(func)
    def wrapped(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as exc:
            current.retry(exc=exc)
    return wrapped
