Source code for src.module.tenant.infrastructure.db_router

"""Tenant-aware database router.

Routes models marked as 'shared' to the public schema and 'tenant_specific'
to the active tenant's schema. Prevents cross-migration between shared and
tenant schemas.

For the **schema** backend this is the critical piece: ``allow_migrate()``
checks whether we are currently running inside a tenant-schema migration
(flagged by ``_migrating_tenant_schema`` in the schema backend) and only
allows models whose ``TenantMeta.tenant_type`` matches the current context.

Without this filtering, ``manage.py migrate`` would create *every* table in
both the public schema and every tenant schema, causing duplicates and the
symptom the user reported (approval tables only in public, catalog tables
missing from tenant schemas).
"""

from __future__ import annotations

import threading
from typing import Any

# Thread-local flag set by SchemaIsolationBackend.migrate_tenant()
# to indicate that we are currently migrating inside a tenant schema.
#
# When True  → only tenant_specific models are allowed to migrate
# When False → only shared models (and Django built-ins) are allowed
_migration_context = threading.local()


def _is_migrating_tenant_schema() -> bool:
    """Return True if we are inside a tenant-schema migration."""
    return getattr(_migration_context, "tenant_schema", False)


def _set_migrating_tenant_schema(value: bool) -> None:
    """Set the tenant-schema migration flag (used by the schema backend)."""
    _migration_context.tenant_schema = value


# Django built-in app labels whose tables must live in the PUBLIC schema only.
# They do not carry a TenantMeta and should never be duplicated into tenant
# schemas.
_SHARED_ONLY_APP_LABELS: frozenset[str] = frozenset(
    {
        "admin",
        "auth",
        "contenttypes",
        "sessions",
        "sites",
        # simplejwt blacklist
        "token_blacklist",
    }
)




class TenantAwareDatabaseRouter:
    """Django DATABASE_ROUTER that routes based on TenantMeta.tenant_type.

    Models with TenantMeta.tenant_type = 'shared' are routed to the default
    database (public schema). Models with tenant_type = 'tenant_specific' are
    routed to the tenant's schema (controlled by search_path set in middleware).

    For schema backend: the search_path controls which schema is active, so
    both shared and tenant models use the same 'default' DB connection but
    with different search_path settings.

    For RLS/shared_fk backends: all models use the same DB, and filtering is
    applied at the query level.
    """

    @staticmethod
    def _get_tenant_type(model: type) -> str | None:
        """Extract the tenant_type from a model's TenantMeta class."""
        tenant_meta = getattr(model, "TenantMeta", None)
        if tenant_meta is None:
            return None
        return getattr(tenant_meta, "tenant_type", None)



    def db_for_read(self, model: type, **hints: Any) -> str | None:
        """Route reads to the default database.

        The search_path (set by middleware) determines which schema is queried.
        """
        return "default"




    def db_for_write(self, model: type, **hints: Any) -> str | None:
        """Route writes to the default database."""
        return "default"




    def allow_relation(self, obj1: Any, obj2: Any, **hints: Any) -> bool | None:
        """Allow relations between models in the same database."""
        return True




    def allow_migrate(
        self,
        db: str,
        app_label: str,
        model_name: str | None = None,
        **hints: Any,
    ) -> bool | None:
        """Control which models are migrated in which context.

        When running a normal ``manage.py migrate`` (public schema):
            ➜ Only shared models and Django built-ins are migrated.

        When ``SchemaIsolationBackend.migrate_tenant()`` runs:
            ➜ Only tenant_specific models are migrated.

        This prevents table duplication across schemas.
        """
        from src.share_kernel.settings import get_setting

        # Only enforce for the schema backend — RLS/shared_fk use a single
        # schema so there is no conflict.
        if str(get_setting("ISOLATION_BACKEND")) != "schema":
            return None

        migrating_tenant = _is_migrating_tenant_schema()

        # -----------------------------------------------------------------
        # 1. Model provided via hints (most reliable)
        # -----------------------------------------------------------------
        model = hints.get("model")
        if model is not None:
            tenant_type = self._get_tenant_type(model)

            if tenant_type == "tenant_specific":
                return migrating_tenant
            if tenant_type == "shared":
                return not migrating_tenant

            # No TenantMeta → Django built-in (auth, admin …)
            return not migrating_tenant

        # -----------------------------------------------------------------
        # 2. No model hint — try to resolve from the Django app registry.
        #    Django passes model_name for CreateModel / AlterField etc.
        # -----------------------------------------------------------------
        if model_name is not None:
            tenant_type = self._resolve_tenant_type(app_label, model_name)
            if tenant_type == "tenant_specific":
                return migrating_tenant
            if tenant_type == "shared":
                return not migrating_tenant
            # Unknown → fall through to app_label heuristic

        # -----------------------------------------------------------------
        # 3. Fallback: no model info at all (RunSQL, RunPython, etc.)
        #    Use app_label heuristic.
        # -----------------------------------------------------------------
        if app_label in _SHARED_ONLY_APP_LABELS:
            return not migrating_tenant

        # For library apps that have a mix of shared + tenant_specific models
        # (e.g. authorization_context), we cannot safely determine the type.
        # Let the operation through — the model-level checks above will have
        # handled all CreateModel/AlterField operations already.
        return None


    @staticmethod
    def _resolve_tenant_type(app_label: str, model_name: str) -> str | None:
        """Try to resolve tenant_type from the Django app registry."""
        from django.apps import apps

        try:
            model = apps.get_model(app_label, model_name)
        except LookupError:
            return None
        tenant_meta = getattr(model, "TenantMeta", None)
        if tenant_meta is None:
            return None
        return getattr(tenant_meta, "tenant_type", None)