From 59c87483e64bcc165d4a309bf338c48638613256 Mon Sep 17 00:00:00 2001 From: ibuler Date: Thu, 19 Sep 2024 11:00:36 +0800 Subject: [PATCH] perf: filter gateway with new params --- apps/accounts/tasks/verify_account.py | 5 ++-- apps/assets/api/asset/asset.py | 34 ++++++++++++++++----------- apps/assets/models/asset/common.py | 7 ++++++ apps/assets/models/domain.py | 2 +- apps/assets/models/gateway.py | 4 ---- apps/assets/serializers/domain.py | 2 +- 6 files changed, 31 insertions(+), 23 deletions(-) diff --git a/apps/accounts/tasks/verify_account.py b/apps/accounts/tasks/verify_account.py index be3e085b5..d6526ec12 100644 --- a/apps/accounts/tasks/verify_account.py +++ b/apps/accounts/tasks/verify_account.py @@ -4,7 +4,6 @@ from django.utils.translation import gettext_noop from accounts.const import AutomationTypes from accounts.tasks.common import quickstart_automation_by_snapshot -from assets.const import GATEWAY_NAME from common.utils import get_logger from orgs.utils import org_aware_func @@ -32,13 +31,13 @@ def verify_accounts_connectivity_util(accounts, task_name): asset_ids = [a.asset_id for a in accounts] assets = Asset.objects.filter(id__in=asset_ids) - gateways = assets.filter(platform__name=GATEWAY_NAME) + gateways = assets.gateways() verify_connectivity_util( gateways, AutomationTypes.verify_gateway_account, accounts, task_name ) - common_assets = assets.exclude(platform__name=GATEWAY_NAME) + common_assets = assets.gateways(0) verify_connectivity_util( common_assets, AutomationTypes.verify_account, accounts, task_name diff --git a/apps/assets/api/asset/asset.py b/apps/assets/api/asset/asset.py index d3bcdb785..849b5f3ad 100644 --- a/apps/assets/api/asset/asset.py +++ b/apps/assets/api/asset/asset.py @@ -2,10 +2,10 @@ # from collections import defaultdict -import django_filters from django.conf import settings from django.shortcuts import get_object_or_404 from django.utils.translation import gettext as _ +from django_filters import rest_framework as drf_filters from rest_framework import status from rest_framework.decorators import action from rest_framework.response import Response @@ -33,31 +33,32 @@ __all__ = [ class AssetFilterSet(BaseFilterSet): - platform = django_filters.CharFilter(method='filter_platform') - exclude_platform = django_filters.CharFilter(field_name="platform__name", lookup_expr='exact', exclude=True) - domain = django_filters.CharFilter(method='filter_domain') - type = django_filters.CharFilter(field_name="platform__type", lookup_expr="exact") - category = django_filters.CharFilter(field_name="platform__category", lookup_expr="exact") - protocols = django_filters.CharFilter(method='filter_protocols') - domain_enabled = django_filters.BooleanFilter( + platform = drf_filters.CharFilter(method='filter_platform') + is_gateway = drf_filters.BooleanFilter(method='filter_is_gateway') + exclude_platform = drf_filters.CharFilter(field_name="platform__name", lookup_expr='exact', exclude=True) + domain = drf_filters.CharFilter(method='filter_domain') + type = drf_filters.CharFilter(field_name="platform__type", lookup_expr="exact") + category = drf_filters.CharFilter(field_name="platform__category", lookup_expr="exact") + protocols = drf_filters.CharFilter(method='filter_protocols') + domain_enabled = drf_filters.BooleanFilter( field_name="platform__domain_enabled", lookup_expr="exact" ) - ping_enabled = django_filters.BooleanFilter( + ping_enabled = drf_filters.BooleanFilter( field_name="platform__automation__ping_enabled", lookup_expr="exact" ) - gather_facts_enabled = django_filters.BooleanFilter( + gather_facts_enabled = drf_filters.BooleanFilter( field_name="platform__automation__gather_facts_enabled", lookup_expr="exact" ) - change_secret_enabled = django_filters.BooleanFilter( + change_secret_enabled = drf_filters.BooleanFilter( field_name="platform__automation__change_secret_enabled", lookup_expr="exact" ) - push_account_enabled = django_filters.BooleanFilter( + push_account_enabled = drf_filters.BooleanFilter( field_name="platform__automation__push_account_enabled", lookup_expr="exact" ) - verify_account_enabled = django_filters.BooleanFilter( + verify_account_enabled = drf_filters.BooleanFilter( field_name="platform__automation__verify_account_enabled", lookup_expr="exact" ) - gather_accounts_enabled = django_filters.BooleanFilter( + gather_accounts_enabled = drf_filters.BooleanFilter( field_name="platform__automation__gather_accounts_enabled", lookup_expr="exact" ) @@ -77,6 +78,11 @@ class AssetFilterSet(BaseFilterSet): else: return queryset.filter(platform__name=value) + @staticmethod + def filter_is_gateway(queryset, name, value): + queryset = queryset.gateways(value) + return queryset + @staticmethod def filter_domain(queryset, name, value): if is_uuid(value): diff --git a/apps/assets/models/asset/common.py b/apps/assets/models/asset/common.py index 7388f27a8..833d8f1be 100644 --- a/apps/assets/models/asset/common.py +++ b/apps/assets/models/asset/common.py @@ -38,6 +38,13 @@ class AssetQuerySet(models.QuerySet): def valid(self): return self.active() + def gateways(self, is_gateway=1): + kwargs = {'platform__name__startswith': 'Gateway'} + if is_gateway: + return self.filter(**kwargs) + else: + return self.exclude(**kwargs) + def has_protocol(self, name): return self.filter(protocols__contains=name) diff --git a/apps/assets/models/domain.py b/apps/assets/models/domain.py index 0c75e22b7..4a9311aed 100644 --- a/apps/assets/models/domain.py +++ b/apps/assets/models/domain.py @@ -31,7 +31,7 @@ class Domain(LabeledMixin, JMSOrgBaseModel): @lazyproperty def assets_amount(self): - return self.assets.exclude(platform__name='Gateway').count() + return self.assets.gateways(0).count() def random_gateway(self): gateways = [gw for gw in self.active_gateways if gw.is_connective] diff --git a/apps/assets/models/gateway.py b/apps/assets/models/gateway.py index 9f23cb5fc..7e9f10da7 100644 --- a/apps/assets/models/gateway.py +++ b/apps/assets/models/gateway.py @@ -33,10 +33,6 @@ class Gateway(Host): proxy = True verbose_name = _("Gateway") - def save(self, *args, **kwargs): - self.platform = self.default_platform() - return super().save(*args, **kwargs) - @classmethod def default_platform(cls): return Platform.objects.get(name=GATEWAY_NAME, internal=True) diff --git a/apps/assets/serializers/domain.py b/apps/assets/serializers/domain.py index bba7e7310..e538a3d0c 100644 --- a/apps/assets/serializers/domain.py +++ b/apps/assets/serializers/domain.py @@ -68,7 +68,7 @@ class DomainListSerializer(DomainSerializer): @classmethod def setup_eager_loading(cls, queryset): queryset = queryset.annotate( - assets_amount=Count('assets', filter=~Q(assets__platform__name='Gateway'), distinct=True), + assets_amount=Count('assets', filter=~Q(assets__platform__name__startswith='Gateway'), distinct=True), ) return queryset