CM-819: added adjustments to support download with custom filters for…

… group beneficiaires

social_protection/export_mixin.py

@@ -0,0 +1,76 @@
+import datetime
+import json
+import logging
+import types
+import uuid
+from typing import Dict, List, Callable
+
+import graphene
+import pandas as pd
+from django.db import models
+from graphene.types.generic import GenericScalar
+from pandas import DataFrame
+
+from core import fields
+from core.custom_filters import CustomFilterWizardStorage
+from core.models import ExportableQueryModel
+from graphql.utils.ast_to_dict import ast_to_dict
+from core.gql.export_mixin import ExportableQueryMixin
+
+logger = logging.getLogger(__file__)
+
+
+class ExportableSocialProtectionQueryMixin(ExportableQueryMixin):
+
+    @classmethod
+    def create_export_function(cls, field_name):
+        new_function_name = f"resolve_{field_name}_export"
+        default_resolve = getattr(cls, F"resolve_{field_name}", None)
+
+        if not default_resolve:
+            raise AttributeError(
+                f"Query {cls} doesn't provide resolve function for {field_name}. "
+                f"CSV export cannot be created")
+
+        def exporter(cls, self, info, **kwargs):
+            custom_filters = kwargs.pop("customFilters", None)
+            export_fields = [cls._adjust_notation(f) for f in kwargs.pop('fields')]
+            fields_mapping = json.loads(kwargs.pop('fields_columns'))
+
+            source_field = getattr(cls, field_name)
+            filter_kwargs = {k: v for k, v in kwargs.items() if k in source_field.filtering_args}
+
+            qs = default_resolve(None, info, **kwargs)
+            qs = qs.filter(**filter_kwargs)
+            qs = cls.__append_custom_filters(custom_filters, qs, fields_mapping)
+            export_file = ExportableQueryModel\
+                .create_csv_export(qs, export_fields, info.context.user, column_names=fields_mapping,
+                                   patches=cls.get_patches_for_field(field_name))
+
+            return export_file.name
+
+        setattr(cls, new_function_name, types.MethodType(exporter, cls))
+
+    @classmethod
+    def __append_custom_filters(cls, custom_filters, queryset, fields_mapping):
+        if custom_filters:
+            module_name = cls.get_module_name()
+            object_type = cls.get_object_type()
+            related_field = cls.get_related_field()
+            if "group__id" in fields_mapping:
+                queryset = CustomFilterWizardStorage.build_custom_filters_queryset(
+                    "individual",
+                    "GroupIndividual",
+                    custom_filters,
+                    queryset,
+                    relation="group"
+                )
+            else:
+                queryset = CustomFilterWizardStorage.build_custom_filters_queryset(
+                    module_name,
+                    object_type,
+                    custom_filters,
+                    queryset,
+                    relation=related_field
+                )
+        return queryset

social_protection/gql_queries.py

@@ -41,6 +41,7 @@ class Meta:
             "date_valid_to": ["exact", "lt", "lte", "gt", "gte"],
             "max_beneficiaries": ["exact", "lt", "lte", "gt", "gte"],
             "institution": ["exact", "iexact", "startswith", "istartswith", "contains", "icontains"],
+            "type": ["exact", "iexact", "startswith", "istartswith", "contains", "icontains"],
 
             "date_created": ["exact", "lt", "lte", "gt", "gte"],
             "date_updated": ["exact", "lt", "lte", "gt", "gte"],

social_protection/schema.py

@@ -5,8 +5,6 @@
 from django.db.models import Q
 from django.core.exceptions import PermissionDenied
 
-from core.gql.export_mixin import ExportableQueryMixin
-
 from django.utils.translation import gettext as _
 from core.custom_filters import CustomFilterWizardStorage
 from core.gql_queries import ValidationMessageGQLType
@@ -28,6 +26,7 @@
     BenefitPlanDataUploadQGLType, BenefitPlanSchemaFieldsGQLType,
     BenefitPlanHistoryGQLType
 )
+from social_protection.export_mixin import ExportableSocialProtectionQueryMixin
 from social_protection.models import (
     BenefitPlan,
     Beneficiary, GroupBeneficiary, BenefitPlanDataUploadRecords
@@ -50,7 +49,7 @@ class BfTypeEnum(graphene.Enum):
     GROUP = BenefitPlan.BenefitPlanType.GROUP_TYPE
 
 
-class Query(ExportableQueryMixin, graphene.ObjectType):
+class Query(ExportableSocialProtectionQueryMixin, graphene.ObjectType):
     export_patches = {
         'beneficiary': [
             patch_details
@@ -91,7 +90,8 @@ class Query(ExportableQueryMixin, graphene.ObjectType):
         dateValidFrom__Gte=graphene.DateTime(),
         dateValidTo__Lte=graphene.DateTime(),
         applyDefaultValidityFilter=graphene.Boolean(),
-        client_mutation_id=graphene.String()
+        client_mutation_id=graphene.String(),
+        customFilters=graphene.List(of_type=graphene.String),
     )
 
     beneficiary_data_upload_history = OrderedDjangoFilterConnectionField(
@@ -238,6 +238,16 @@ def resolve_group_beneficiary(self, info, **kwargs):
             SocialProtectionConfig.gql_beneficiary_search_perms
         )
         query = GroupBeneficiary.objects.filter(*filters)
+
+        custom_filters = kwargs.get("customFilters", None)
+        if custom_filters:
+            query = CustomFilterWizardStorage.build_custom_filters_queryset(
+                "individual",
+                "GroupIndividual",
+                custom_filters,
+                query,
+                "group",
+            )
         return gql_optimizer.query(query, info)
 
     def resolve_awaiting_beneficiary(self, info, **kwargs):