def worker_prerun(task: Task, **kwargs) -> None:
"""
Called before a worker starts executing a task.
Here we make sure to set the appropriate correlation ID for all logs logged
during the tasks, and on the thread in general. In that regard, this does
the Celery equivalent to what the django-guid middleware does for a request.
"""
guid = task.request.get(settings.guid_header_name)
if guid:
logger.info('Setting GUID %s', guid)
set_guid(guid)
else:
generated_guid = generate_guid(
uuid_length=settings.integration_settings.celery.uuid_length)
logger.info('Generated GUID %s', generated_guid)
set_guid(generated_guid)
if settings.integration_settings.celery.log_parent:
origin = task.request.get(parent_header)
if origin:
logger.info('Setting parent ID %s', origin)
celery_parent.set(origin)
generated_current_guid = generate_guid(
uuid_length=settings.integration_settings.celery.uuid_length)
logger.info('Generated current ID %s', generated_current_guid)
celery_current.set(generated_current_guid)
def clean_up(task: Task, **kwargs) -> None:
"""
Called after a task is finished.
Here we make sure to clean up the IDs we set in the pre-run method, so that
the next task executed by the same worker doesn't inherit the same IDs.
"""
logger.debug('Cleaning up GUIDs')
clear_guid()
if settings.integration_settings.celery.log_parent:
celery_current.set(None)
celery_parent.set(None)
def test_cleanup(monkeypatch, mocker: MockerFixture):
"""
Test that cleanup works as expected
"""
set_guid('123')
celery_current.set('123')
celery_parent.set('123')
mocked_settings = deepcopy(django_settings.DJANGO_GUID)
mocked_settings['INTEGRATIONS'] = [CeleryIntegration(log_parent=True)]
with override_settings(DJANGO_GUID=mocked_settings):
settings = Settings()
monkeypatch.setattr('django_guid.integrations.celery.signals.settings',
settings)
clean_up(task=mocker.Mock())
assert [get_guid(), celery_current.get(),
celery_parent.get()] == [None, None, None]
def test_task_publish_includes_correct_depth_headers(monkeypatch):
"""
Test log_parent True
"""
mocked_settings = deepcopy(django_settings.DJANGO_GUID)
mocked_settings['INTEGRATIONS'] = [CeleryIntegration(log_parent=True)]
with override_settings(DJANGO_GUID=mocked_settings):
settings = Settings()
monkeypatch.setattr('django_guid.integrations.celery.signals.settings',
settings)
headers = {}
publish_task_from_worker_or_request(headers=headers)
# The parent header should not be in headers, because
# There should be no celery_current context
assert parent_header not in headers
for correlation_id in ['test', 123, -1]:
headers = {}
celery_current.set(correlation_id)
publish_task_from_worker_or_request(headers=headers)
# Here the celery-parent-id header should exist
assert headers[parent_header] == correlation_id