Add location-based row level security to BenefitPlans (#105)

* added support for location field for beneficiary upload and update, for programs of type individuals

* added beneficiary base fields in base service

* added support for location queries for both type of benefits plans and only show beneficiaries based on row level security

* added tests for location cases for update and upload workflows

* added unique_class_validation back

* updated get_queryset of beneficiaries logic to reuse Individual and Group models get_queryset method functionality

* debug logs removed

Author: Shahzaibahmad97

social_protection/apps.py

@@ -48,6 +48,9 @@
         'json_ext.beneficiary_data_source',
         'json_ext.educated_level'
     ],
+    "beneficiary_base_fields": [
+        'first_name', 'last_name', 'dob', 'location_name', 'location_code', 'id'
+    ],
     "social_protection_masking_enabled": True,
     "enable_python_workflows": True,
     "default_beneficiary_status": "POTENTIAL",
@@ -93,6 +96,7 @@ class SocialProtectionConfig(AppConfig):
     enable_maker_checker_logic_enrollment = None
     beneficiary_mask_fields = None
     group_beneficiary_mask_fields = None
+    beneficiary_base_fields = None
     social_protection_masking_enabled = None
 
     default_beneficiary_status = None

social_protection/models.py

@@ -1,3 +1,4 @@
+from django.conf import settings
 from django.contrib.auth.models import AnonymousUser
 from django.db import models
 from django.db.models import Func
@@ -57,6 +58,18 @@ def clean(self):
 
     def __str__(self):
         return f'{self.individual.first_name} {self.individual.last_name}'
+    
+    @classmethod
+    def get_queryset(cls, queryset, user):
+        if queryset is None:
+            queryset = cls.objects.all()
+
+        individuals = Individual.objects.filter(
+            id__in=queryset.values('individual_id')
+        ).distinct()
+    
+        individual_queryset = Individual.get_queryset(individuals, user)
+        return queryset.filter(individual__in=individual_queryset)
 
 
 class BenefitPlanDataUploadRecords(core_models.HistoryModel):
@@ -80,6 +93,18 @@ def clean(self):
             raise ValidationError(_("Group beneficiary must be associated with a benefit plan type = GROUP."))
 
         super().clean()
+    
+    @classmethod
+    def get_queryset(cls, queryset, user):
+        if queryset is None:
+            queryset = cls.objects.all()
+
+        groups = Group.objects.filter(
+            id__in=queryset.values('group_id')
+        ).distinct()
+    
+        group_queryset = Group.get_queryset(groups, user)
+        return queryset.filter(group__in=group_queryset)
 
 
 class JSONUpdate(Func):

social_protection/schema.py

@@ -35,6 +35,7 @@
 )
 from social_protection.validation import validate_bf_unique_code, validate_bf_unique_name
 import graphene_django_optimizer as gql_optimizer
+from location.apps import LocationConfig
 
 
 def patch_details(beneficiary_df: pd.DataFrame):
@@ -82,6 +83,8 @@ class Query(ExportableSocialProtectionQueryMixin, graphene.ObjectType):
         orderBy=graphene.List(of_type=graphene.String),
         dateValidFrom__Gte=graphene.DateTime(),
         dateValidTo__Lte=graphene.DateTime(),
+        parent_location=graphene.String(),
+        parent_location_level=graphene.Int(),
         applyDefaultValidityFilter=graphene.Boolean(),
         client_mutation_id=graphene.String(),
         customFilters=graphene.List(of_type=graphene.String),
@@ -91,6 +94,8 @@ class Query(ExportableSocialProtectionQueryMixin, graphene.ObjectType):
         orderBy=graphene.List(of_type=graphene.String),
         dateValidFrom__Gte=graphene.DateTime(),
         dateValidTo__Lte=graphene.DateTime(),
+        parent_location=graphene.String(),
+        parent_location_level=graphene.Int(),
         applyDefaultValidityFilter=graphene.Boolean(),
         client_mutation_id=graphene.String(),
         customFilters=graphene.List(of_type=graphene.String),
@@ -271,7 +276,14 @@ def _annotate_is_eligible(query, eligible_uuids, eligibility_check_performed):
             )
 
         filters = _build_filters(info, **kwargs)
-        query = _apply_custom_filters(Beneficiary.objects.filter(*filters), **kwargs)
+        
+        parent_location = kwargs.get('parent_location')
+        parent_location_level = kwargs.get('parent_location_level')
+        if parent_location is not None and parent_location_level is not None:
+            filters.append(Query._get_location_filters(parent_location, parent_location_level, prefix='individual__'))
+
+        query = Beneficiary.get_queryset(None, info.context.user)
+        query = _apply_custom_filters(query.filter(*filters), **kwargs)
 
         eligible_uuids, eligibility_check_performed = _get_eligible_uuids(query, info, **kwargs)
         query = _annotate_is_eligible(query, eligible_uuids, eligibility_check_performed)
@@ -342,7 +354,14 @@ def _annotate_is_eligible(query, eligible_group_uuids, eligibility_check_perform
             )
 
         filters = _build_filters(info, **kwargs)
-        query = _apply_custom_filters(GroupBeneficiary.objects.filter(*filters), **kwargs)
+        
+        parent_location = kwargs.get('parent_location')
+        parent_location_level = kwargs.get('parent_location_level')
+        if parent_location is not None and parent_location_level is not None:
+            filters.append(Query._get_location_filters(parent_location, parent_location_level, prefix='group__'))
+
+        query = GroupBeneficiary.get_queryset(None, info.context.user)
+        query = _apply_custom_filters(query.filter(*filters), **kwargs)
 
         eligible_group_uuids, eligibility_check_performed = _get_eligible_group_uuids(query, info, **kwargs)
         query = _annotate_is_eligible(query, eligible_group_uuids, eligibility_check_performed)
@@ -439,6 +458,14 @@ def resolve_benefit_plan_history(self, info, **kwargs):
         if sort_alphabetically:
             query = query.order_by('code')
         return gql_optimizer.query(query, info)
+    
+    @staticmethod
+    def _get_location_filters(parent_location, parent_location_level, prefix=""):
+        query_key = "uuid"
+        for i in range(len(LocationConfig.location_types) - parent_location_level - 1):
+            query_key = "parent__" + query_key
+        query_key = prefix + "location__" + query_key
+        return Q(**{query_key: parent_location})
 
 
 class Mutation(graphene.ObjectType):

social_protection/tests/__init__.py

@@ -3,5 +3,7 @@
 from .group_beneficiary_service_test import GroupBeneficiaryServiceTest
 from .beneficiary_import_service_test import BeneficiaryImportServiceTest
 from .beneficiary_gql_test import BeneficiaryGQLTest
+from .test_workflows_beneficiaries_upload import ProcessImportBeneficiariesWorkflowTest
+from .test_workflows_beneficiaries_update import ProcessUpdateBeneficiariesWorkflowTest
 # TODO: implement group upload workflow
-# from .group_import_gql_test import GroupBeneficiaryImportGQLTest
+# from .group_import_gql_test import GroupBeneficiaryImportGQLTest

social_protection/tests/test_workflows_beneficiaries_update.py

@@ -0,0 +1,242 @@
+from django.db import connection
+from django.test import TestCase
+from core.test_helpers import create_test_interactive_user
+from individual.models import (
+    Individual,
+    IndividualDataSource,
+    IndividualDataSourceUpload,
+)
+from social_protection.models import BenefitPlanDataUploadRecords, Beneficiary, BeneficiaryStatus
+from social_protection.services import BeneficiaryService
+from social_protection.workflows.base_beneficiary_update import process_update_beneficiaries_workflow
+from social_protection.tests.test_helpers import add_individual_to_benefit_plan, create_benefit_plan, create_individual
+from location.test_helpers import create_test_village
+from unittest.mock import patch
+import uuid
+from unittest import skipIf
+
+
+@skipIf(
+    connection.vendor != "postgresql",
+    "Skipping tests due to implementation usage of validate_json_schema, which is a postgres specific extension."
+)
+class ProcessUpdateBeneficiariesWorkflowTest(TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        # Patch validate_dataframe_headers as it is already tested separately
+        cls.validate_headers_patcher = patch(
+            "social_protection.workflows.utils.BasePythonWorkflowExecutor.validate_dataframe_headers",
+            lambda self, is_update: None
+        )
+        cls.validate_headers_patcher.start()
+
+        cls.schema_patcher = patch("individual.apps.IndividualConfig.individual_schema", "{}")
+        cls.schema_patcher.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.validate_headers_patcher.stop()
+        super().tearDownClass()
+
+    def setUp(self):
+        self.user = create_test_interactive_user(username="admin")
+        self.user_uuid = self.user.id
+
+        self.plan_schema = {"properties": {}}
+        self.benefit_plan = create_benefit_plan(self.user.username, {
+                "name": "Test Benefit Plan",
+                "description": "A test benefit plan",
+                "code": "TESTPlan",
+                "max_beneficiaries": 1000,
+                "beneficiary_data_schema": {}
+            })
+
+        individual1_dict = {
+            'first_name': 'Foo 1',
+            'last_name': 'Bar',
+            "dob": "1998-01-01",
+            'json_ext': {},
+        }
+        self.individual1 = create_individual(self.user.username, individual1_dict)
+
+        individual2_dict = {
+            'first_name': 'Foo 2',
+            'last_name': 'Baz',
+            "dob": "1980-01-01",
+            'json_ext': {},
+        }
+        self.individual2 = create_individual(self.user.username, individual2_dict)
+        
+        self.service = BeneficiaryService(self.user)
+        
+        # Add individuals to the benefit plan and create beneficiaries
+        self.beneficiary1_uuid = add_individual_to_benefit_plan(
+            self.service, 
+            self.individual1, 
+            self.benefit_plan, 
+            payload_override={'status': 'ACTIVE'}
+        )
+        self.beneficiary2_uuid = add_individual_to_benefit_plan(
+            self.service, 
+            self.individual2, 
+            self.benefit_plan, 
+            payload_override={'status': 'ACTIVE'}
+        )
+
+        self.upload = IndividualDataSourceUpload(
+            source_name='csv',
+            source_type='update',
+            status="PENDING",
+        )
+        self.upload.save(user=self.user)
+        self.upload_uuid = self.upload.id
+
+        upload_record = BenefitPlanDataUploadRecords(
+            data_upload=self.upload,
+            workflow='Python Beneficiaries Update',
+            benefit_plan=self.benefit_plan,
+            json_ext={}
+        )
+        upload_record.save(user=self.user.user)
+
+        self.village = create_test_village({
+            'name': 'McLean',
+            'code': 'VwA',
+        })
+        self.individual1_updated_first_name = "John"
+        self.beneficiary1_updated_status = BeneficiaryStatus.POTENTIAL
+        self.valid_data_source = IndividualDataSource(
+            upload_id=self.upload_uuid,
+            json_ext={
+                "ID": str(self.beneficiary1_uuid),
+                "first_name": self.individual1_updated_first_name,
+                "last_name": "Doe",
+                "dob": "1980-01-01",
+                "status": self.beneficiary1_updated_status,
+                "location_name": self.village.name,
+                "location_code": self.village.code,
+            }
+        )
+        self.valid_data_source.save(user=self.user)
+
+        self.individual2_updated_first_name = "Jane"
+        self.invalid_data_source = IndividualDataSource(
+            upload_id=self.upload_uuid,
+            json_ext={
+                "ID": str(uuid.uuid4()),
+                "first_name": self.individual2_updated_first_name,
+            }
+        )
+        self.invalid_data_source.save(user=self.user)
+
+    @patch('individual.apps.IndividualConfig.enable_maker_checker_for_individual_update', False)
+    @patch('social_protection.apps.SocialProtectionConfig.enable_maker_checker_for_beneficiary_update', False)
+    @skipIf(
+        connection.vendor != "postgresql",
+        "Skipping tests due to implementation usage of validate_json_schema, which is a postgres specific extension."
+    )
+    def test_process_update_beneficiaries_workflow_successful_execution(self):
+        process_update_beneficiaries_workflow(self.user_uuid, self.benefit_plan.uuid, self.upload_uuid)
+
+        upload = IndividualDataSourceUpload.objects.get(id=self.upload_uuid)
+
+        # Check that the status is 'FAIL' due to the entry with invalid ID
+        self.assertEqual(upload.status, "FAIL")
+        self.assertIsNotNone(upload.error)
+        errors = upload.error['errors']
+        self.assertIn("Invalid entries", errors['error'])
+
+        # Check that the correct failing entries are logged in the error field
+        error_key = "failing_entries_invalid_id"
+        self.assertIn(error_key, errors)
+        self.assertIn(str(self.invalid_data_source.id), errors[error_key]['uuids'])
+        self.assertNotIn(str(self.valid_data_source.id), errors[error_key]['uuids'])
+
+        # individual_id should not be assigned for any data sources
+        data_entries = IndividualDataSource.objects.filter(upload_id=self.upload_uuid)
+        for entry in data_entries:
+            self.assertIsNone(entry.individual_id)
+
+        # individual data should not be updated
+        individual1_from_db = Individual.objects.get(id=self.individual1.id)
+        self.assertNotEqual(individual1_from_db.first_name, self.individual1_updated_first_name)
+        
+        beneficiary1_from_db = Beneficiary.objects.get(id=self.beneficiary1_uuid)
+        self.assertNotEqual(beneficiary1_from_db.status, self.beneficiary1_updated_status)
+
+        individual2_from_db = Individual.objects.get(id=self.individual2.id)
+        self.assertNotEqual(individual2_from_db.first_name, self.individual2_updated_first_name)
+
+    @patch('individual.apps.IndividualConfig.enable_maker_checker_for_individual_update', False)
+    @patch('social_protection.apps.SocialProtectionConfig.enable_maker_checker_for_beneficiary_update', False)
+    @skipIf(
+        connection.vendor != "postgresql",
+        "Skipping tests due to implementation usage of validate_json_schema, which is a postgres specific extension."
+    )
+    def test_process_update_beneficiaries_workflow_with_all_valid_entries(self):
+        # Update invalid entry in IndividualDataSource to valid data
+        self.invalid_data_source.json_ext = {
+            "ID": str(self.beneficiary2_uuid),
+            "first_name": self.individual2_updated_first_name,
+            "location_name": None,
+            "location_code": None,
+        }
+        self.invalid_data_source.save(user=self.user)
+
+        process_update_beneficiaries_workflow(self.user_uuid, self.benefit_plan.uuid, self.upload_uuid)
+
+        upload = IndividualDataSourceUpload.objects.get(id=self.upload_uuid)
+
+        self.assertEqual(upload.status, "SUCCESS", upload.error)
+        self.assertEqual(upload.error, {})
+
+        # Verify that individual IDs have been assigned to data entries in IndividualDataSource
+        data_entries = IndividualDataSource.objects.filter(upload_id=self.upload_uuid)
+        for entry in data_entries:
+            self.assertIsNotNone(entry.individual_id)
+
+        # individual data should be updated
+        individual1_from_db = Individual.objects.get(id=self.individual1.id)
+        self.assertEqual(individual1_from_db.first_name, self.individual1_updated_first_name)
+        self.assertEqual(individual1_from_db.location.name, self.village.name)
+        
+        # Beneficiary status update shouldn't make any effect as it is not supported
+        beneficiary1_from_db = Beneficiary.objects.get(id=self.beneficiary1_uuid)
+        self.assertNotEqual(beneficiary1_from_db.status, self.beneficiary1_updated_status)
+
+        individual2_from_db = Individual.objects.get(id=self.individual2.id)
+        self.assertEqual(individual2_from_db.first_name, self.individual2_updated_first_name)
+        self.assertIsNone(individual2_from_db.location)
+
+    @patch('individual.apps.IndividualConfig.enable_maker_checker_for_individual_update', True)
+    @patch('social_protection.apps.SocialProtectionConfig.enable_maker_checker_for_beneficiary_update', True)
+    def test_process_update_beneficiaries_workflow_with_maker_checker_enabled(self):
+        # Update invalid entry in IndividualDataSource to valid data
+        self.invalid_data_source.json_ext = {
+            "ID": str(self.beneficiary2_uuid),
+            "first_name": "Jane",
+            "location_name": None,
+            "location_code": None,
+        }
+        self.invalid_data_source.save(user=self.user)
+
+        process_update_beneficiaries_workflow(self.user_uuid, self.benefit_plan.uuid, self.upload_uuid)
+
+        upload = IndividualDataSourceUpload.objects.get(id=self.upload_uuid)
+
+        self.assertEqual(upload.status, "WAITING_FOR_VERIFICATION")
+        self.assertEqual(upload.error, {})
+
+        # Verify that individual IDs not yet assigned to data entries in IndividualDataSource
+        data_entries = IndividualDataSource.objects.filter(upload_id=self.upload_uuid)
+        for entry in data_entries:
+            self.assertIsNone(entry.individual_id)
+
+        # individual data should not be updated
+        individual1_from_db = Individual.objects.get(id=self.individual1.id)
+        self.assertNotEqual(individual1_from_db.first_name, self.individual1_updated_first_name)
+
+        individual2_from_db = Individual.objects.get(id=self.individual2.id)
+        self.assertNotEqual(individual2_from_db.first_name, self.individual2_updated_first_name)