1
\$\begingroup\$

I have been typically adding the @transaction.atomic decorator to my endpoints and appreciate if a request fails at any point the data completely rolls back. However, some of my endpoints will make an API call to another API service. If my token for that API service is invalid, I refresh the token. This is all well and good until an exception occurs after the API token has been refreshed. The transaction is rolled back including the token information which then will be stale. This includes the new refresh token information so I'm now unable to refresh the auth token.

I attempted to get around this situation by distributing the token refresh to a celery task and waiting for the result as a way of isolating that database change. Even if the requests fails later on, the token information remains while the rest is rolled back.

Here is a simplified example

@transaction.atomic()
def post(request):
  response = api_call()
  db_operation(response['result']) #  Could potentially have an exception

def api_call(allow_refresh=True)
  response = requests.post(url, params)

  if response.status_code == 401:
    if allow_refresh:
     refresh_token_task.delay().get(timeout=10)
     response = api_call(allow_refresh=False)
    else:
      raise Exception('Invalid authorization') # In the case refreshing token doesn't resolve authorization issues, this prevents infinite loop.

  return response

Phew! I thought I was in the clear except... sometimes I distribute api calls. It didn't occur to me but I realized calling celery tasks from celery tasks easily risks locking. So now I've written out a solution that involves chaining.

def post(request):
  chain(api_call_task.s(), refresh_token_task.s(), api_call_task.s(), db_operation_task.s())()

@task
def api_call_task(initial_response=None)
  if initial_response and initial_response.status_code != 401:
    return response

  response = requests.post(url, params)

  return response

@task
def refresh_token_task(response):
  if response.status_code == 401:
    refresh_token()

  return response

@transaction.atomic()
@task
def db_operation_task(response):
  if response.status_code == 401:
    raise Exception('Invalid authorization')

  db_operation(response['result'])

I think this works to a degree but gosh what a confusing mess. By chaining the functions, the response is continuously passed through which determines if the token needs to be refreshed and/or if the call needs to be made again and then finally to handle the db operation based upon the response to the api.

Is there a simpler way? Marginal improvement? Bad practice involved? Would it better to just let go of transaction.atomic? Let me know if my simplified pseudo code can use clarification.

\$\endgroup\$
3
\$\begingroup\$

Let's start off with what's probably the most important note here: if you are running into issues with atomic transactions, pushing the database work off to Celery (or another worker thread) inherently makes your transactions non-atomic. It sounds like this has been working for you currently, but this will not always work and may have unintended consequences.

So I'm going to review this from two different angles:

  1. Handling atomic transactions within APIs when you only need part of it to be atomic
  2. Chaining Celery tasks that need to execute conditionally

If only part of your handling needs to be atomic, don't make everything atomic

In your hypothetical example, you make it clear that db_operation needs to be atomic and that's why you are wrapping the entire handle in transaction.atomic. I'm assuming that db_operation is actually a series of calls, or something more complex than just a single call, which is why you couldn't just decorate that call with transaction.atomic. But you only appear to care about these calls being atomic, which is why the call to refresh_token still needs to be committed to the database even if there are issues within db_operation.

My suggestion would be to use transaction.atomic as a context manager instead of as a decorator. This would result in your hypothetical view looking something like:

def post(request):
    response = api_call()
    with @transaction.atomic():
        db_operation(response['result']) #  Could potentially have an exception

This still allows db_operation to remain atomic while ensuring that api_call is always committed to the database.

Celery has advanced method chaining, use it to your advantage

You correctly identified that calling one Celery task and waiting for a response can easily cause a deadlock or starvation, which is why Celery explicitly recommends that you don't do it.

Getting an error when making an API call because your access token has expired is generally considered to be an exceptional case. Celery provides functionality for handling exceptions, so don't be afraid to use it. You found Celery's chaining functionality, but what is also useful is the linking functionality between tasks that works very similarly.

def post(request):
    api_call_task.apply_async(
        link=[db_operation_task.s()],
        link_error=[refresh_token_task.s(), api_call_task.s(), db_operation_task.s()]
    )

You'll see in this case that if the original call to api_call_task returns an error, we are instructing Celery to call refresh_token_task, then api_call_task again (it should have a valid token) and then eventually db_operation_task like we originally planned. This allows for api_call_task to error out twice, such as if there was an error not relating to the token refresh, without risking calling db_operation_task with potentially invalid data.

Because Celery provides a clear callback format and allows you to define a chain of methods to be called during the callbacks, you can design your tasks to only focus on what they really care about: calling those other methods.

@task
def api_call_task():
    response = requests.post(url, params)

    if response.status_code == 401:
        raise Exception('Access token needs to be refreshed')

    return response

@task
def refresh_token_task():
    refresh_token()

@task
@transaction.atomic()
def db_operation_task(response):
    db_operation(response['result'])

This has the benefit of allowing you to handle errors naturally within the flow, without forcing every step in the chain to perform the error handling for you.

Bonus: Celery tasks can also be run synchronously, use that to your advantage

In your hypothetical example, you defined api_call and api_call_task when talking about the Celery task and the actual API call itself. If you have control over api_call (it's not a third-party call), you can actually just condense this down to a single method. If you call a method wrapped with @task() without calling .delay() or .apply_async(), the method is called synchronously and Celery is never contacted. This allows you to design methods which can be called either asynchronously using Celery, or synchronously in the course of your own application, without having to duplicate a lot of code or add in additional special cases.

Because Celery allows you to execute the code synchronously that would otherwise be executed asynchronously, this also makes testing of these tasks much easier. In tests you can call them synchronously and ensure the functionality works as expected, while knowing that in the real world there may be a time delay that you need to account for.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.