Implement entity graph for Case and Document entities to optimize dat…

backend/nyaysetu-backend/src/main/java/com/nyaysetu/backend/entity/CaseEntity.java

@@ -1,8 +1,11 @@
 package com.nyaysetu.backend.entity;
 
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import jakarta.persistence.*;
 import lombok.*;
 import java.time.LocalDateTime;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.UUID;
 
 @Entity
@@ -62,6 +65,11 @@ public class CaseEntity {
 
     private Long sourceFirId;
 
+    @Builder.Default
+    @OneToMany(mappedBy = "caseEntity", fetch = FetchType.LAZY, cascade = CascadeType.ALL)
+    @JsonIgnoreProperties("caseEntity")
+    private List<DocumentEntity> documents = new ArrayList<>();
+
     private Boolean hasBsaCert;
 
 

backend/nyaysetu-backend/src/main/java/com/nyaysetu/backend/entity/DocumentEntity.java

@@ -1,5 +1,6 @@
 package com.nyaysetu.backend.entity;
 
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import jakarta.persistence.*;
 import lombok.*;
 import java.time.LocalDateTime;
@@ -18,8 +19,14 @@ public class DocumentEntity {
     @GeneratedValue(strategy = GenerationType.UUID)
     private UUID id;
 
+    @Column(name = "case_id")
     private UUID caseId;
 
+    @ManyToOne(fetch = FetchType.LAZY)
+    @JoinColumn(name = "case_id", insertable = false, updatable = false)
+    @JsonIgnoreProperties({"hibernateLazyInitializer", "handler", "documents"})
+    private CaseEntity caseEntity;
+
     private String fileName;
 
     private String fileUrl;

backend/nyaysetu-backend/src/main/java/com/nyaysetu/backend/repository/CaseRepository.java

@@ -3,40 +3,60 @@
 import com.nyaysetu.backend.entity.CaseEntity;
 import com.nyaysetu.backend.entity.CaseStatus;
 import com.nyaysetu.backend.entity.User;
+import org.springframework.data.jpa.repository.EntityGraph;
 import org.springframework.data.jpa.repository.JpaRepository;
 import org.springframework.data.jpa.repository.Query;
 import org.springframework.stereotype.Repository;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.UUID;
 
 @Repository
 public interface CaseRepository extends JpaRepository<CaseEntity, UUID> {
+    @EntityGraph(attributePaths = {"documents"})
     List<CaseEntity> findByJudgeId(Long judgeId);
+
+    @EntityGraph(attributePaths = {"documents"})
     List<CaseEntity> findByClient(User client);
 
     // For auto-assignment - find cases without judge
+    @EntityGraph(attributePaths = {"documents"})
     List<CaseEntity> findByJudgeIdIsNull();
 
     // Find cases by status
+    @EntityGraph(attributePaths = {"documents"})
     List<CaseEntity> findByStatus(CaseStatus status);
 
     // Count cases assigned to a judge (for round-robin)
     long countByJudgeId(Long judgeId);
 
     // Lawyer-specific queries
+    @EntityGraph(attributePaths = {"documents"})
     List<CaseEntity> findByLawyer(User lawyer);
     long countByLawyer(User lawyer);
 
     // Find cases by assigned judge name/email
+    @EntityGraph(attributePaths = {"documents"})
     List<CaseEntity> findByAssignedJudge(String judgeName);
 
     // Find unassigned cases (both null and empty string)
+    @EntityGraph(attributePaths = {"documents"})
     List<CaseEntity> findByAssignedJudgeIsNull();
 
+    @EntityGraph(attributePaths = {"documents"})
     @Query("SELECT c FROM CaseEntity c WHERE c.assignedJudge IS NULL OR c.assignedJudge = ''")
     List<CaseEntity> findUnassignedCases();
 
     // Find cases by respondent email
+    @EntityGraph(attributePaths = {"documents"})
     List<CaseEntity> findByRespondentEmail(String respondentEmail);
+
+    @EntityGraph(attributePaths = {"documents"})
+    @Override
+    Optional<CaseEntity> findById(UUID id);
+
+    @EntityGraph(attributePaths = {"documents"})
+    @Override
+    List<CaseEntity> findAll();
 }

backend/nyaysetu-backend/src/main/java/com/nyaysetu/backend/repository/DocumentRepository.java

@@ -1,14 +1,26 @@
 package com.nyaysetu.backend.repository;
 
 import com.nyaysetu.backend.entity.DocumentEntity;
+import org.springframework.data.jpa.repository.EntityGraph;
 import org.springframework.data.jpa.repository.JpaRepository;
 import org.springframework.stereotype.Repository;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.UUID;
 
 @Repository
 public interface DocumentRepository extends JpaRepository<DocumentEntity, UUID> {
+    @EntityGraph(attributePaths = {"caseEntity"})
+    @Override
+    List<DocumentEntity> findAll();
+
+    @EntityGraph(attributePaths = {"caseEntity"})
     List<DocumentEntity> findByCaseId(UUID caseId);
+
+    @EntityGraph(attributePaths = {"caseEntity"})
+    @Override
+    Optional<DocumentEntity> findById(UUID id);
+
     List<DocumentEntity> findByCategoryAndDescriptionContaining(String category, String description);
 }

backend/nyaysetu-backend/src/main/java/com/nyaysetu/backend/service/DocumentManagementService.java

@@ -198,7 +198,9 @@ private DocumentDto convertToDto(DocumentEntity entity) {
                 .build();
 
         // Get case title if available
-        if (entity.getCaseId() != null) {
+        if (entity.getCaseEntity() != null) {
+            dto.setCaseTitle(entity.getCaseEntity().getTitle());
+        } else if (entity.getCaseId() != null) {
             caseRepository.findById(entity.getCaseId())
                     .ifPresent(caseEntity -> dto.setCaseTitle(caseEntity.getTitle()));
         }

backend/nyaysetu-backend/src/test/java/com/nyaysetu/backend/repository/CaseDocumentEntityGraphTest.java

@@ -0,0 +1,108 @@
+package com.nyaysetu.backend.repository;
+
+import com.nyaysetu.backend.entity.CaseEntity;
+import com.nyaysetu.backend.entity.DocumentEntity;
+import com.nyaysetu.backend.entity.Role;
+import com.nyaysetu.backend.entity.User;
+import jakarta.persistence.EntityManagerFactory;
+import org.hibernate.SessionFactory;
+import org.hibernate.stat.Statistics;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.test.context.ActiveProfiles;
+import org.springframework.transaction.annotation.Transactional;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+@SpringBootTest
+@ActiveProfiles("test")
+@Transactional
+public class CaseDocumentEntityGraphTest {
+
+    @Autowired
+    private CaseRepository caseRepository;
+
+    @Autowired
+    private DocumentRepository documentRepository;
+
+    @Autowired
+    private UserRepository userRepository;
+
+    @Autowired
+    private EntityManagerFactory entityManagerFactory;
+
+    @Test
+    void whenFindingCasesByClient_thenDocumentsAreLoadedWithEntityGraph() {
+        User client = User.builder()
+                .email("case-doc-graph@example.com")
+                .name("Graph Test Client")
+                .password("password")
+                .role(Role.LITIGANT)
+                .build();
+        client = userRepository.save(client);
+        assertNotNull(client.getId());
+
+        CaseEntity firstCase = CaseEntity.builder()
+                .title("First Graph Case")
+                .description("This case should load documents eagerly via EntityGraph")
+                .client(client)
+                .build();
+        firstCase = caseRepository.save(firstCase);
+
+        CaseEntity secondCase = CaseEntity.builder()
+                .title("Second Graph Case")
+                .description("This case should load documents eagerly via EntityGraph")
+                .client(client)
+                .build();
+        secondCase = caseRepository.save(secondCase);
+
+        DocumentEntity firstDoc = DocumentEntity.builder()
+                .caseId(firstCase.getId())
+                .fileName("first-doc.pdf")
+                .fileUrl("/files/first-doc.pdf")
+                .contentType("application/pdf")
+                .size(123L)
+                .uploadedBy(client.getId())
+                .build();
+        documentRepository.save(firstDoc);
+
+        DocumentEntity secondDoc = DocumentEntity.builder()
+                .caseId(secondCase.getId())
+                .fileName("second-doc.pdf")
+                .fileUrl("/files/second-doc.pdf")
+                .contentType("application/pdf")
+                .size(456L)
+                .uploadedBy(client.getId())
+                .build();
+        documentRepository.save(secondDoc);
+
+        DocumentEntity thirdDoc = DocumentEntity.builder()
+                .caseId(secondCase.getId())
+                .fileName("third-doc.pdf")
+                .fileUrl("/files/third-doc.pdf")
+                .contentType("application/pdf")
+                .size(789L)
+                .uploadedBy(client.getId())
+                .build();
+        documentRepository.save(thirdDoc);
+
+        SessionFactory sessionFactory = entityManagerFactory.unwrap(SessionFactory.class);
+        Statistics statistics = sessionFactory.getStatistics();
+        statistics.setStatisticsEnabled(true);
+        statistics.clear();
+
+        List<CaseEntity> cases = caseRepository.findByClient(client);
+
+        assertFalse(cases.isEmpty(), "Expected cases for client to be returned");
+        assertEquals(2, cases.size(), "Expected exactly two cases for the client");
+
+        int totalDocuments = cases.stream().mapToInt(c -> c.getDocuments().size()).sum();
+        assertEquals(3, totalDocuments, "Expected all case documents to be loaded with the cases");
+
+        long queryCount = statistics.getPrepareStatementCount();
+        assertTrue(queryCount <= 2, "Expected entity graph to fetch cases and documents with a low query count, but executed " + queryCount);
+    }
+}