jumpserver/apps/orgs/mixins.py

# -*- coding: utf-8 -*-
#
from django.db import models
from django.shortcuts import redirect
from django.contrib.auth import get_user_model

from common.utils import get_logger
from .utils import get_current_org, get_model_by_db_table

logger = get_logger(__file__)


__all__ = ['OrgManager', 'OrgViewGenericMixin', 'OrgModelMixin']


class OrgManager(models.Manager):
    def get_queryset(self):
        current_org = get_current_org()
        user_model = get_user_model()
        kwargs = {}

        print("Get queryset ")
        print(self.model)
        print(current_org)

        if not current_org:
            kwargs['id'] = None
        elif issubclass(self.model, user_model):
            kwargs['orgs'] = current_org
        elif current_org.is_real():
            kwargs['org'] = current_org
        elif current_org.is_default():
            kwargs['org'] = None
        print(kwargs)
        return super().get_queryset().filter(**kwargs)


class OrgModelMixin(models.Model):
    org = models.ForeignKey('orgs.Organization', on_delete=models.PROTECT, null=True)

    objects = OrgManager()

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def _do_update(self, base_qs, using, pk_val, values, update_fields, forced_update):
        current_org = get_current_org()

        if current_org and current_org.is_real():
            kwargs = {'org': current_org}
            base_qs = base_qs.filter(**kwargs)
        else:
            logger.warn(
                'Attempting to update %s instance "%s" without a current tenant '
                'set. This may cause issues in a partitioned environment. '
                'Recommend calling set_current_org() before performing this '
                'operation.', self._meta.model.__name__, self
            )
        return super()._do_update(base_qs, using, pk_val, values, update_fields, forced_update)

    def save(self, force_insert=False, force_update=False, using=None,
             update_fields=None):
        user_model = get_user_model()
        current_org = get_current_org()
        if current_org and not current_org.is_real():
            self.org = current_org
        instance = super().save(
            force_insert=force_insert, force_update=force_update,
            using=using, update_fields=update_fields
        )
        if isinstance(instance, user_model):
            instance.orgs.add(current_org)
        return instance

    class Meta:
        abstract = True


class OrgViewGenericMixin:
    def dispatch(self, request, *args, **kwargs):
        current_org = get_current_org()
        if not current_org:
            return redirect('orgs:switch-a-org')
        return super().dispatch(request, *args, **kwargs)