pythondjangodjango-rest-frameworkproduction-environmentthread-local-storage

Django - Troubleshooting thread locals & middleware issue with DRF ViewSets in Production


Anyone have a recommendation for approaching this challenging problem? Not sure how to troubleshoot this further.

I am setting a value in thread locals in middleware using the following:

from threading import local

ORG_ATTR_NAME = getattr(settings, "LOCAL_ORG_ATTR_NAME", "_current_org")

_thread_locals = local()

def _do_set_current_variable(variable_fun, attr_name):
    setattr(_thread_locals, attr_name, variable_fun.__get__(variable_fun, local))

def thread_local_middleware(get_response):

    def middleware(request):
        organization = Organization.objects.get_current(request)
        _do_set_current_variable(lambda self: organization, ORG_ATTR_NAME)

        response = get_response(request)

        return response

    return middleware

def get_current_variable(attr_name: str):
    """Given an attr_name, returns the object or callable if it exists in the local thread"""
    current_variable = getattr(_thread_locals, attr_name, None)
    if callable(current_variable):
        return current_variable()
    return current_variable

I have a Manager class that all models which contain an organization FK field make use of:

class OrganizationAwareManager(models.Manager):
    def get_queryset(self):
        logger.debug(f"_thread_locals attributes in manager: {_thread_locals.__dict__}")
        organization = get_current_variable(ORG_ATTR_NAME)
        queryset = super().get_queryset().filter(organization=organization)
        logger.debug(f"OrganizationAwareManager queryset: {queryset.query}")
        return queryset

This all works like a charm in dev, and also works great on nearly everything in production - except in DRF ViewSets that query models with the OrganizationAwareManager. Other ViewSets work as expected, and normal django views that refer to models with the OrganizationAwareManager in view context also work fine. Literally just DRF Viewsets with OrganizationAwareManager in Production are the only issue.

As you can see, I added logging in the manager to check what attributes are set on _thread_locals.

In dev I get something like: _thread_locals attributes in manager: {'_current_org': <bound method thread_local_middleware.<locals>.middleware.<locals>.<lambda> of <function thread_local_middleware.<locals>.middleware.<locals>.<lambda> at 0x7f1820264a60>>}

In prod, nothing seems to be set: _thread_locals attributes in manager: {}

Is there something I might be missing about how DRF processes requests? It should be running through the middleware stack and setting my attribute whether we're on dev or prod, right? I can't seem to find any difference between the two environments that could possibly account for this. Both are very similar, using docker-compose with nearly identical containers. I have tried both gunicorn and gunicorn+uvicorn in the prod environment with no difference in symptoms.

I have been trying to resolve this for 3 days now, and am running out of ideas, so suggestions to fix this would be much appreciated.


Solution

  • How strange. After some time I had an idea that resolved the issue.

    Still not exactly sure of the root cause (would love any insight why it worked just fine in dev, but not in prod until the changes below were made), but it seems that in ViewSets which define the queryset in the class itself, the query is evaluated when the thread begins?

    I noticed in my logs that when I started the server, I got a whole bunch of log entries from the OrganizationAwareManager saying that _thread_locals had no associated attributes. The number of these log entries seemed to be about the same as the quantity of ViewSets in my project that used OrganizationAwareManager. They all got evaluated initially as the system initiated, with orgnization=None, so any further filtering on organization would be discarded.

    ViewSets like the one below did not end up correctly filtering by organization:

    class AssetTypeViewSet(viewsets.ModelViewSet):
        queryset = AssetType.objects.all()
        serializer_class = AssetTypeSerializer
    

    When I modified to define the queryset inside get_queryset() so that it gets evaluated when the ViewSet is executed, things now work as expected:

    class AssetTypeViewSet(viewsets.ModelViewSet):
        queryset = AssetType.objects.none()
        serializer_class = AssetTypeSerializer
    
        def get_queryset(self):
            return AssetType.objects.all()
    

    Weird.