Skip to main content

C4 Code Diagram: Multi-Tenant Context Manager

Purpose: Class-level detail of the tenant context management system, showing thread-safe context storage, PostgreSQL session variable integration, and usage patterns for maintaining tenant isolation throughout request lifecycle.

Scope: Tenant context layer - Thread-safe context management

Related Diagrams:

Mermaid Class Diagram​

Component Details​

1. ContextVar (Python Standard Library)​

Module: contextvars (Python 3.7+)

Purpose: Thread-safe and async-safe context variable storage.

Why ContextVar?

  • Thread-safe: Each thread has its own value (prevents cross-thread leakage)
  • Async-safe: Each async task has its own value (prevents cross-task leakage)
  • Request-scoped: Value automatically cleared when execution context ends
  • No global state: Safer than thread-local storage

Implementation:

from contextvars import ContextVar
from typing import Optional

# Module-level context variable (defined once)
_current_tenant: ContextVar[Optional['Tenant']] = ContextVar(
    'current_tenant',  # Name for debugging
    default=None       # Default value if not set
)

Core Methods:

# Get current value (returns default if not set)
tenant = _current_tenant.get()  # Returns None if not set

# Set new value (returns Token for resetting)
token = _current_tenant.set(my_tenant)

# Reset to previous value
_current_tenant.reset(token)

Thread Safety Example:

# Thread 1
_current_tenant.set(tenant1)
print(_current_tenant.get())  # Prints tenant1

# Thread 2 (simultaneously)
_current_tenant.set(tenant2)
print(_current_tenant.get())  # Prints tenant2 (not tenant1!)

# Values are isolated per thread - no interference

2. Tenant Context Module​

Module: apps.tenants.context

Purpose: High-level API for tenant context management with PostgreSQL integration.

Implementation:

# apps/tenants/context.py

from contextvars import ContextVar
from typing import Optional
from django.db import connection

# Type hint for Tenant model (avoid circular import)
if False:  # TYPE_CHECKING
    from .models import Tenant

# Module-level context variable
_current_tenant: ContextVar[Optional['Tenant']] = ContextVar(
    'current_tenant',
    default=None
)

def set_current_tenant(tenant: Optional['Tenant']) -> None:
    """
    Set the current tenant for this execution context.

    Sets both:
    1. Python ContextVar (for Django ORM queries)
    2. PostgreSQL session variable (for Row-Level Security enforcement)

    Args:
        tenant: Tenant object to set as current, or None to clear

    Example:
        from apps.tenants.models import Tenant
        from apps.tenants.context import set_current_tenant

        tenant = Tenant.objects.get(slug='acme-corp')
        set_current_tenant(tenant)

        # Now all queries automatically filter to acme-corp
        users = User.objects.all()  # Only acme-corp users
    """
    # Set Python context variable
    _current_tenant.set(tenant)

    # Set PostgreSQL session variable for RLS enforcement
    if tenant:
        with connection.cursor() as cursor:
            cursor.execute("SELECT set_current_tenant(%s)", [str(tenant.id)])
    else:
        # Clear PostgreSQL variable if tenant is None
        with connection.cursor() as cursor:
            cursor.execute("SELECT clear_current_tenant()")

def get_current_tenant() -> Optional['Tenant']:
    """
    Get the current tenant from execution context.

    Returns:
        Tenant object if set, None otherwise

    Example:
        from apps.tenants.context import get_current_tenant

        tenant = get_current_tenant()
        if tenant:
            print(f"Current tenant: {tenant.name}")
        else:
            print("No tenant context set")
    """
    return _current_tenant.get()

def clear_current_tenant() -> None:
    """
    Clear the current tenant context.

    Clears both:
    1. Python ContextVar
    2. PostgreSQL session variable

    CRITICAL: Always call this after request completes to prevent
    context leakage to subsequent requests on the same thread.

    Example:
        from apps.tenants.context import clear_current_tenant

        try:
            # ... process request ...
        finally:
            clear_current_tenant()  # Always clear in finally block
    """
    # Clear Python context variable
    _current_tenant.set(None)

    # Clear PostgreSQL session variable
    with connection.cursor() as cursor:
        cursor.execute("SELECT clear_current_tenant()")

3. tenant_context Context Manager​

Module: apps.tenants.context

Purpose: Temporarily set tenant context with automatic cleanup (RAII pattern).

Implementation:

# apps/tenants/context.py (continued)

class tenant_context:
    """
    Context manager for temporarily setting tenant context.

    Automatically restores previous tenant context on exit,
    even if exception occurs. Safe for nested usage.

    Example:
        from apps.tenants.context import tenant_context
        from apps.tenants.models import Tenant

        tenant = Tenant.objects.get(slug='acme-corp')

        # Tenant context only set within 'with' block
        with tenant_context(tenant):
            # All queries here scoped to acme-corp
            users = User.objects.all()
            projects = Project.objects.all()

        # Context automatically restored to previous value
        # (even if exception was raised inside block)

    Nested usage (for testing):
        with tenant_context(tenant1):
            users1 = User.objects.all()  # tenant1 users

            with tenant_context(tenant2):
                users2 = User.objects.all()  # tenant2 users

            # Back to tenant1
            users1_again = User.objects.all()
    """

    def __init__(self, tenant: Optional['Tenant']):
        """
        Initialize context manager.

        Args:
            tenant: Tenant to set as current (or None to clear)
        """
        self.tenant = tenant
        self.previous_tenant = None

    def __enter__(self):
        """
        Enter context: Save previous tenant and set new tenant.

        Returns:
            The tenant that was set (for convenience)
        """
        # Save current tenant (to restore later)
        self.previous_tenant = get_current_tenant()

        # Set new tenant
        set_current_tenant(self.tenant)

        # Return tenant for optional 'as' binding
        return self.tenant

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Exit context: Restore previous tenant.

        Args:
            exc_type: Exception type (if raised)
            exc_val: Exception value (if raised)
            exc_tb: Exception traceback (if raised)

        Returns:
            False (do not suppress exceptions)
        """
        # Restore previous tenant (even if exception occurred)
        set_current_tenant(self.previous_tenant)

        # Don't suppress exceptions
        return False

Usage Examples:

Example 1: Celery Background Tasks

# apps/licenses/tasks.py

from celery import shared_task
from apps.tenants.context import tenant_context
from apps.licenses.models import LicenseSession

@shared_task
def send_license_expiry_email(license_id):
    """Send email for expiring license session."""

    # Get license (without tenant context, query won't work)
    license = LicenseSession.objects.get(id=license_id)

    # Set tenant context for subsequent queries
    with tenant_context(license.tenant):
        # Now all queries scoped to license's tenant
        user = license.user
        EmailService.send_expiry_notification(user.email)

Example 2: Admin Operations (Cross-Tenant)

# apps/admin/views.py

from apps.tenants.context import tenant_context
from apps.tenants.models import Tenant

def admin_generate_report(request):
    """Generate report across all tenants (admin only)."""

    report_data = []

    for tenant in Tenant.objects.all():
        with tenant_context(tenant):
            # Queries scoped to current tenant
            active_users = User.objects.filter(is_active=True).count()
            active_sessions = LicenseSession.objects.filter(status='active').count()

            report_data.append({
                'tenant': tenant.name,
                'active_users': active_users,
                'active_sessions': active_sessions
            })

    return render(request, 'admin/report.html', {'data': report_data})

Example 3: Unit Testing

# apps/users/tests/test_models.py

from django.test import TestCase
from apps.tenants.context import tenant_context
from apps.tenants.models import Tenant
from apps.users.models import User

class UserModelTestCase(TestCase):
    def setUp(self):
        self.tenant1 = Tenant.objects.create(name="Tenant 1")
        self.tenant2 = Tenant.objects.create(name="Tenant 2")

    def test_user_creation_with_context(self):
        """Test user creation within tenant context."""

        # Create user in tenant1 context
        with tenant_context(self.tenant1):
            user1 = User.objects.create(email="user@tenant1.com")
            self.assertEqual(user1.tenant, self.tenant1)

        # Create user in tenant2 context
        with tenant_context(self.tenant2):
            user2 = User.objects.create(email="user@tenant2.com")
            self.assertEqual(user2.tenant, self.tenant2)

        # Verify isolation
        with tenant_context(self.tenant1):
            users = User.objects.all()
            self.assertEqual(users.count(), 1)
            self.assertEqual(users.first(), user1)

4. TenantMiddleware​

Module: apps.tenants.middleware

Purpose: Automatically set/clear tenant context for every HTTP request.

Implementation:

# apps/tenants/middleware.py

from django.utils.deprecation import MiddlewareMixin
from django.http import JsonResponse
from .context import set_current_tenant, clear_current_tenant
import logging

logger = logging.getLogger(__name__)

class TenantMiddleware(MiddlewareMixin):
    """
    Middleware to set current tenant from authenticated user.

    CRITICAL ORDERING:
    - MUST come AFTER AuthenticationMiddleware
      (requires request.user to be set)

    Settings configuration:
        MIDDLEWARE = [
            # ...
            'django.contrib.auth.middleware.AuthenticationMiddleware',
            'apps.tenants.middleware.TenantMiddleware',  # After auth!
            # ...
        ]

    Request Lifecycle:
        1. process_request() - Set tenant from request.user.tenant
        2. View executes (with tenant context set)
        3. process_response() - Clear tenant context
        4. process_exception() - Clear tenant context (if exception)
    """

    def process_request(self, request):
        """
        Set tenant from authenticated user.

        Called before view execution.

        Args:
            request: Django HttpRequest

        Side Effects:
            - Sets request.tenant attribute
            - Calls set_current_tenant() (Python + PostgreSQL)
            - Logs tenant access for monitoring
        """
        # Always clear context first (prevent leakage from previous request)
        clear_current_tenant()

        # Check if user is authenticated
        if request.user and request.user.is_authenticated:
            # Check if user has tenant FK
            if hasattr(request.user, 'tenant'):
                # Set tenant context (Python + PostgreSQL)
                set_current_tenant(request.user.tenant)

                # Attach tenant to request for convenience
                request.tenant = request.user.tenant

                # Log tenant access for monitoring
                logger.info(
                    "Tenant context set",
                    extra={
                        'tenant_id': str(request.tenant.id),
                        'tenant_name': request.tenant.name,
                        'user_id': str(request.user.id),
                        'user_email': request.user.email,
                        'path': request.path,
                        'method': request.method
                    }
                )
            else:
                # Superuser or staff without tenant (Django admin)
                request.tenant = None
                logger.debug("User without tenant (superuser/staff)")
        else:
            # Anonymous user
            request.tenant = None
            logger.debug("Anonymous request (no tenant)")

    def process_response(self, request, response):
        """
        Clear tenant context after response.

        CRITICAL: Always clear context to prevent leakage to
        next request on same thread (connection pooling).

        Args:
            request: Django HttpRequest
            response: Django HttpResponse

        Returns:
            HttpResponse (unmodified)
        """
        # Clear tenant context (Python + PostgreSQL)
        clear_current_tenant()

        logger.debug("Tenant context cleared after response")

        return response

    def process_exception(self, request, exception):
        """
        Clear tenant context if exception occurs.

        Args:
            request: Django HttpRequest
            exception: Exception that was raised

        Returns:
            None (let exception propagate)
        """
        # Clear tenant context even on exception
        clear_current_tenant()

        logger.warning(
            "Tenant context cleared after exception",
            extra={
                'exception_type': type(exception).__name__,
                'exception_message': str(exception)
            }
        )

        # Don't handle exception, just clear context
        return None

Request Flow:

1. Request arrives → Django processes middleware chain

2. SecurityMiddleware → HTTPS enforcement
3. SessionMiddleware → Load session
4. AuthenticationMiddleware → Set request.user
5. TenantMiddleware → process_request()
   - clear_current_tenant() (prevent leakage)
   - Check request.user.is_authenticated
   - Set tenant from request.user.tenant
   - set_current_tenant(tenant) → Python + PostgreSQL

6. View executes → All queries automatically filtered by tenant

7. TenantMiddleware → process_response()
   - clear_current_tenant() → Python + PostgreSQL
   - Return response

8. Response sent to client

5. PostgreSQL Integration​

Purpose: Enforce tenant isolation at database level via Row-Level Security (RLS).

PostgreSQL Functions:

-- Set current tenant for session
CREATE OR REPLACE FUNCTION set_current_tenant(tenant_uuid UUID)
RETURNS void AS $$
BEGIN
    -- Set PostgreSQL session variable
    PERFORM set_config('app.current_tenant_id', tenant_uuid::TEXT, false);
    -- false = session-level (not transaction-level)
END;
$$ LANGUAGE plpgsql;

-- Get current tenant from session
CREATE OR REPLACE FUNCTION get_current_tenant()
RETURNS UUID AS $$
BEGIN
    RETURN current_setting('app.current_tenant_id', true)::UUID;
    -- true = don't error if not set
EXCEPTION
    WHEN OTHERS THEN
        RETURN NULL;
END;
$$ LANGUAGE plpgsql;

-- Clear current tenant
CREATE OR REPLACE FUNCTION clear_current_tenant()
RETURNS void AS $$
BEGIN
    PERFORM set_config('app.current_tenant_id', '', false);
END;
$$ LANGUAGE plpgsql;

Row-Level Security Policies:

-- Example: users table RLS policy
CREATE POLICY users_tenant_isolation ON users
    USING (tenant_id = current_setting('app.current_tenant_id')::UUID);

-- Result: PostgreSQL automatically adds WHERE clause to all queries
SELECT * FROM users;
-- Becomes:
SELECT * FROM users WHERE tenant_id = current_setting('app.current_tenant_id')::UUID;

Integration Flow:

# Python code
set_current_tenant(my_tenant)

# Executes SQL:
# SELECT set_current_tenant('tenant-uuid')

# PostgreSQL session variable set:
# app.current_tenant_id = 'tenant-uuid'

# All subsequent queries automatically filtered:
User.objects.all()

# SQL generated by Django ORM:
# SELECT * FROM users

# PostgreSQL RLS adds WHERE clause:
# SELECT * FROM users WHERE tenant_id = 'tenant-uuid'

Sequence Diagram: Request Lifecycle​

Error Handling​

Save Without Tenant Context​

from apps.users.models import User
from apps.tenants.context import get_current_tenant

# WRONG: No tenant context set
user = User(email="test@example.com")

try:
    user.save()
except ValueError as e:
    print(e)
    # Output: "Cannot save User without tenant context. Set tenant explicitly or use TenantMiddleware."

# FIX 1: Use context manager
with tenant_context(my_tenant):
    user.save()  # ✓ Works

# FIX 2: Set tenant explicitly
user.tenant = my_tenant
user.save()  # ✓ Works

Delete from Wrong Tenant​

from apps.tenants.context import set_current_tenant

# Set context to tenant1
set_current_tenant(tenant1)

# Create user in tenant1
user1 = User.objects.create(email="user@tenant1.com")

# Switch to tenant2
set_current_tenant(tenant2)

# Try to delete tenant1's user
try:
    user1.delete()
except PermissionError as e:
    print(e)
    # Output: "Cannot delete User from different tenant"

Testing Examples​

Unit Tests​

from django.test import TestCase
from apps.tenants.context import set_current_tenant, get_current_tenant, clear_current_tenant
from apps.tenants.models import Tenant

class TenantContextTestCase(TestCase):
    def test_set_and_get_tenant(self):
        """Test setting and getting tenant context."""
        tenant = Tenant.objects.create(name="Test Tenant")

        # Initially None
        self.assertIsNone(get_current_tenant())

        # Set tenant
        set_current_tenant(tenant)
        self.assertEqual(get_current_tenant(), tenant)

        # Clear tenant
        clear_current_tenant()
        self.assertIsNone(get_current_tenant())

    def test_context_manager(self):
        """Test tenant_context context manager."""
        tenant1 = Tenant.objects.create(name="Tenant 1")
        tenant2 = Tenant.objects.create(name="Tenant 2")

        # Set tenant1
        set_current_tenant(tenant1)

        # Use context manager to temporarily switch to tenant2
        with tenant_context(tenant2):
            self.assertEqual(get_current_tenant(), tenant2)

        # Context restored to tenant1
        self.assertEqual(get_current_tenant(), tenant1)

    def test_context_manager_exception(self):
        """Test context manager clears on exception."""
        tenant1 = Tenant.objects.create(name="Tenant 1")
        tenant2 = Tenant.objects.create(name="Tenant 2")

        set_current_tenant(tenant1)

        try:
            with tenant_context(tenant2):
                # Raise exception inside context
                raise ValueError("Test exception")
        except ValueError:
            pass

        # Context should be restored despite exception
        self.assertEqual(get_current_tenant(), tenant1)

Middleware Tests​

from django.test import TestCase, RequestFactory
from apps.tenants.middleware import TenantMiddleware
from apps.tenants.context import get_current_tenant
from apps.tenants.models import Tenant
from apps.users.models import User

class TenantMiddlewareTestCase(TestCase):
    def setUp(self):
        self.factory = RequestFactory()
        self.middleware = TenantMiddleware()
        self.tenant = Tenant.objects.create(name="Test Tenant")

        with tenant_context(self.tenant):
            self.user = User.objects.create(
                email="user@test.com",
                tenant=self.tenant
            )

    def test_middleware_sets_tenant_for_authenticated_user(self):
        """Test middleware sets tenant from authenticated user."""
        request = self.factory.get('/')
        request.user = self.user

        self.middleware.process_request(request)

        # Tenant should be set
        self.assertEqual(get_current_tenant(), self.tenant)
        self.assertEqual(request.tenant, self.tenant)

        # Clean up
        self.middleware.process_response(request, None)

    def test_middleware_clears_tenant_after_response(self):
        """Test middleware clears tenant after response."""
        request = self.factory.get('/')
        request.user = self.user

        self.middleware.process_request(request)
        self.assertIsNotNone(get_current_tenant())

        self.middleware.process_response(request, None)
        self.assertIsNone(get_current_tenant())

Performance Considerations​

ContextVar Performance​

ContextVar is fast:

  • O(1) get/set operations
  • No locks required (thread-local storage)
  • Minimal memory overhead (~40 bytes per context)

Benchmark:

import timeit

# ContextVar get/set
setup = "from contextvars import ContextVar; cv = ContextVar('test'); cv.set('value')"
result = timeit.timeit("cv.get()", setup=setup, number=1000000)
print(f"1M get() operations: {result:.3f}s")
# Output: ~0.05s (50 nanoseconds per operation)

PostgreSQL Session Variable Overhead​

Setting session variable:

  • ~1-2ms per set_config() call
  • Acceptable overhead (once per request)

Optimization:

  • Middleware sets variable only once per request
  • Not set on every query (would be expensive)

Last Updated: 2025-11-30 Diagram Type: C4 Code (Mermaid Class Diagram) Scope: Tenant Context Management - Thread-safe Context Storage