mirror of https://github.com/jumpserver/jumpserver
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
85 lines
2.7 KiB
85 lines
2.7 KiB
# -*- 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)
|