def ready(self):
from waldur_core.quotas import handlers, utils
Quota = self.get_model('Quota')
for index, model in enumerate(utils.get_models_with_quotas()):
signals.post_save.connect(
handlers.increase_global_quota,
sender=model,
dispatch_uid='waldur_core.quotas.handlers.increase_global_quota_%s_%s'
% (model.__name__, index),
)
signals.post_delete.connect(
handlers.decrease_global_quota,
sender=model,
dispatch_uid='waldur_core.quotas.handlers.decrease_global_quota_%s_%s'
% (model.__name__, index),
)
signals.post_migrate.connect(
handlers.create_global_quotas,
dispatch_uid="waldur_core.quotas.handlers.create_global_quotas",
)
# new quotas
from waldur_core.quotas import fields
for model_index, model in enumerate(utils.get_models_with_quotas()):
# quota initialization
signals.post_save.connect(
handlers.init_quotas,
sender=model,
dispatch_uid='waldur_core.quotas.init_quotas_%s_%s'
% (model.__name__, model_index),
)
# Counter quota signals
# How it works:
# Each counter quota field has list of target models. Change of target model should increase or decrease
# counter quota. So we connect generated handler to each of target models.
for counter_field in model.get_quotas_fields(
field_class=fields.CounterQuotaField
):
self.register_counter_field_signals(model, counter_field)
# Aggregator quotas signals
signals.post_save.connect(
handlers.handle_aggregated_quotas,
sender=Quota,
dispatch_uid='waldur_core.quotas.handle_aggregated_quotas_post_save',
)
signals.pre_delete.connect(
handlers.handle_aggregated_quotas,
sender=Quota,
dispatch_uid='waldur_core.quotas.handle_aggregated_quotas_pre_delete',
)
def delete_stale_quotas(self):
self.stdout.write('Deleting stale quotas')
for model in get_models_with_quotas():
for obj in model.objects.all():
quotas_names = model.get_quotas_names()
obj.quotas.exclude(name__in=quotas_names).delete()
self.stdout.write('...done')
def recalculate_counter_quotas(self):
self.stdout.write('Recalculating counter quotas')
for model in get_models_with_quotas():
for counter_field in model.get_quotas_fields(field_class=fields.CounterQuotaField):
for instance in model.objects.all():
counter_field.recalculate(scope=instance)
self.stdout.write('...done')
def handle(self, *args, **options):
for model in get_models_with_quotas():
if hasattr(model, 'GLOBAL_COUNT_QUOTA_NAME'):
quota, _ = models.Quota.objects.get_or_create(
name=model.GLOBAL_COUNT_QUOTA_NAME
)
for index, instance in enumerate(
model.objects.all().order_by('created')
):
revision = Revision.objects.create()
revision.date_created = instance.created
revision.save()
quota.usage = index + 1
serializer = django_serializers.get_serializer('json')()
serialized_data = serializer.serialize([quota])
Version.objects.create(
revision=revision,
object_id=quota.id,
object_id_int=quota.id,
content_type=ContentType.objects.get_for_model(quota),
format='json',
serialized_data=serialized_data,
object_repr=str(quota),
)
def recalculate_aggregator_quotas(self):
# TODO: recalculate child quotas first
self.stdout.write('Recalculating aggregator quotas')
for model in get_models_with_quotas():
for aggregator_field in model.get_quotas_fields(field_class=fields.AggregatorQuotaField):
for instance in model.objects.all():
aggregator_field.recalculate(scope=instance)
self.stdout.write('...done')
def recalculate_global_quotas(self):
self.stdout.write('Recalculating global quotas')
for model in get_models_with_quotas():
if hasattr(model, 'GLOBAL_COUNT_QUOTA_NAME'):
with transaction.atomic():
quota, _ = models.Quota.objects.get_or_create(name=model.GLOBAL_COUNT_QUOTA_NAME)
quota.usage = model.objects.count()
quota.save()
self.stdout.write('...done')
def init_missing_quotas(self):
self.stdout.write('Initializing missing quotas')
for model in get_models_with_quotas():
for obj in model.objects.all():
for field in obj.get_quotas_fields():
try:
field.get_or_create_quota(scope=obj)
except exceptions.CreationConditionFailedQuotaError:
pass
self.stdout.write('...done')
class QuotaSerializer(serializers.HyperlinkedModelSerializer):
scope = GenericRelatedField(related_models=utils.get_models_with_quotas(),
read_only=True)
class Meta:
model = models.Quota
fields = ('url', 'uuid', 'name', 'limit', 'usage', 'scope',
'threshold')
read_only_fields = ('uuid', 'name', 'usage')
extra_kwargs = {
'url': {
'lookup_field': 'uuid'
},
}
def filtered_for_user(self, user, queryset=None):
from waldur_core.quotas import utils
if queryset is None:
queryset = self.get_queryset()
# XXX: This circular dependency will be removed then filter_queryset_for_user
# will be moved to model manager method
from waldur_core.structure.managers import filter_queryset_for_user
if user.is_staff or user.is_support:
return queryset
quota_scope_models = utils.get_models_with_quotas()
query = Q()
for model in quota_scope_models:
user_object_ids = filter_queryset_for_user(
model.objects.all(), user).values_list('id', flat=True)
content_type_id = ct_models.ContentType.objects.get_for_model(
model).id
query |= Q(object_id__in=user_object_ids,
content_type_id=content_type_id)
return queryset.filter(query)
def create_global_quotas(**kwargs):
for model in utils.get_models_with_quotas():
if hasattr(model, 'GLOBAL_COUNT_QUOTA_NAME'):
models.Quota.objects.get_or_create(
name=getattr(model, 'GLOBAL_COUNT_QUOTA_NAME'))
def lookups(self, request, model_admin):
models = utils.get_models_with_quotas()
return [(ct_models.ContentType.objects.get_for_model(m).id, m.__name__)
for m in models]
