Fix apply_async not working with positional kwargs

Celery's apply_async function accepts 'kwargs' as (get ready to be
confused) either a positional argument, or a keyword argument:

Positional: apply_async(['args'], {'kw': 'args'})

Keyword: apply_async(args=['args'], kwargs={'kw': 'args'})

We rely on the positional form in at least one place [1]. This fixes
the overload of apply_async to cope with both forms, and continue to
pass through any other (confusion time again) keyword args to super(),
such as queue="queue".

Note that we've also decided to stop accepting other positional args,
since this is unnecessarily confusing, and we don't currently rely on
it in our code. This stops it creeping in in future.

[1]: fde927e00e/app/job/rest.py (L186)
This commit is contained in:
Ben Thorner
2021-04-15 12:57:12 +01:00
parent fde927e00e
commit be02573147
2 changed files with 21 additions and 7 deletions

View File

@@ -73,15 +73,15 @@ def make_task(app):
return super().__call__(*args, **kwargs) return super().__call__(*args, **kwargs)
def apply_async(self, *args, **kwargs): def apply_async(self, args=None, kwargs=None, **other_kwargs):
kwargs['kwargs'] = kwargs.get('kwargs', {}) kwargs = kwargs or {}
if has_request_context() and hasattr(request, 'request_id'): if has_request_context() and hasattr(request, 'request_id'):
kwargs['kwargs']['request_id'] = request.request_id kwargs['request_id'] = request.request_id
elif has_app_context() and 'request_id' in g: elif has_app_context() and 'request_id' in g:
kwargs['kwargs']['request_id'] = g.request_id kwargs['request_id'] = g.request_id
return super().apply_async(*args, **kwargs) return super().apply_async(args, kwargs, **other_kwargs)
return NotifyTask return NotifyTask

View File

@@ -89,7 +89,21 @@ def test_apply_async_injects_global_request_id_into_kwargs(mocker, celery_task):
super_apply = mocker.patch('celery.app.task.Task.apply_async') super_apply = mocker.patch('celery.app.task.Task.apply_async')
g.request_id = '1234' g.request_id = '1234'
celery_task.apply_async() celery_task.apply_async()
super_apply.assert_called_with(kwargs={'request_id': '1234'}) super_apply.assert_called_with(None, {'request_id': '1234'})
def test_apply_async_inject_request_id_with_other_kwargs(mocker, celery_task):
super_apply = mocker.patch('celery.app.task.Task.apply_async')
g.request_id = '1234'
celery_task.apply_async(kwargs={'something': 'else'})
super_apply.assert_called_with(None, {'request_id': '1234', 'something': 'else'})
def test_apply_async_inject_request_id_with_positional_args(mocker, celery_task):
super_apply = mocker.patch('celery.app.task.Task.apply_async')
g.request_id = '1234'
celery_task.apply_async(['args'], {'something': 'else'})
super_apply.assert_called_with(['args'], {'request_id': '1234', 'something': 'else'})
def test_apply_async_injects_id_into_kwargs_from_request(mocker, notify_api, celery_task): def test_apply_async_injects_id_into_kwargs_from_request(mocker, notify_api, celery_task):
@@ -100,4 +114,4 @@ def test_apply_async_injects_id_into_kwargs_from_request(mocker, notify_api, cel
with notify_api.test_request_context(headers=request_headers): with notify_api.test_request_context(headers=request_headers):
celery_task.apply_async() celery_task.apply_async()
super_apply.assert_called_with(kwargs={'request_id': '1234'}) super_apply.assert_called_with(None, {'request_id': '1234'})