fix: support stage-scoped data element headers in both short/full formats (TE) (#22918)

Author: luciano-fiandesio

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/common/processing/DimensionIdentifierConverter.java

@@ -126,15 +126,18 @@ private DimensionIdentifier<StringUid> handleTwoPartFormat(
       return handleProgramScoped(allowedPrograms, programOpt.get(), firstElement, dimension);
     }
 
-    // If first element is a stage UID with event-level dimension, handle as stage-scoped
+    // Event-level static dimensions always follow stage-scoped handling.
     if (isEventLevelStaticDimension(dimension.getUid())) {
       return handleStageScoped(allowedPrograms, firstElement, dimension);
     }
 
-    // If first element is a stage UID with non-event-level static dimension, throw specific error
-    if (isProgramStageUid(allowedPrograms, firstUid)
-        && isNonEventLevelStaticDimension(dimension.getUid())) {
-      throw unsupportedStageDimensionError(dimension.getUid(), firstUid);
+    // If first element is a stage UID, infer the program from the stage.
+    // Non-event-level static dimensions remain unsupported for stage-specific scoping.
+    if (isProgramStageUid(allowedPrograms, firstUid)) {
+      if (isNonEventLevelStaticDimension(dimension.getUid())) {
+        throw unsupportedStageDimensionError(dimension.getUid(), firstUid);
+      }
+      return handleStageScoped(allowedPrograms, firstElement, dimension);
     }
 
     // Otherwise, program doesn't exist

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/common/processing/HeaderParamsHandler.java

@@ -35,6 +35,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import org.hisp.dhis.analytics.common.CommonRequestParams;
 import org.hisp.dhis.analytics.common.ContextParams;
@@ -44,6 +45,7 @@
 import org.hisp.dhis.common.Grid;
 import org.hisp.dhis.common.GridHeader;
 import org.hisp.dhis.common.IllegalQueryException;
+import org.hisp.dhis.common.RepeatableStageParams;
 import org.hisp.dhis.feedback.ErrorMessage;
 import org.springframework.stereotype.Component;
 
@@ -80,14 +82,10 @@ public void handle(
 
       // Adds only the headers present in params, in the same order.
       paramHeaders.forEach(
-          header -> {
-            GridHeader gridHeader = new GridHeader(header);
-
-            if (gridHeaders.contains(gridHeader)) {
-              int element = gridHeaders.indexOf(gridHeader);
-              grid.addHeader(gridHeaders.get(element));
-            }
-          });
+          header ->
+              findMatchingHeader(gridHeaders, header)
+                  .map(matched -> withRequestedNameIfNeeded(matched, header))
+                  .ifPresent(grid::addHeader));
     }
 
     checkHeaders(headers, paramHeaders);
@@ -102,13 +100,64 @@ public void handle(
    *     "gridHeaders".
    */
   private void checkHeaders(Set<GridHeader> gridHeaders, Set<String> paramHeaders) {
+    List<GridHeader> requestedGridHeaders = new ArrayList<>(gridHeaders);
     paramHeaders.forEach(
         header -> {
-          GridHeader gridHeader = new GridHeader(header);
-
-          if (!gridHeaders.contains(gridHeader)) {
+          if (findMatchingHeader(requestedGridHeaders, header).isEmpty()) {
             throw new IllegalQueryException(new ErrorMessage(E7230, header));
           }
         });
   }
+
+  /**
+   * Finds a matching header by exact name or supported short/full stage-scoped alias:
+   * programUid.stageUid.dimension <-> stageUid.dimension.
+   */
+  private Optional<GridHeader> findMatchingHeader(List<GridHeader> gridHeaders, String header) {
+    GridHeader requested = new GridHeader(header);
+
+    if (gridHeaders.contains(requested)) {
+      return Optional.of(gridHeaders.get(gridHeaders.indexOf(requested)));
+    }
+
+    return gridHeaders.stream().filter(h -> isStageScopedAlias(h.getName(), header)).findFirst();
+  }
+
+  private boolean isStageScopedAlias(String existingHeaderName, String requestedHeaderName) {
+    long existingDots = existingHeaderName.chars().filter(c -> c == '.').count();
+    long requestedDots = requestedHeaderName.chars().filter(c -> c == '.').count();
+
+    if (existingDots == 2 && requestedDots == 1) {
+      return existingHeaderName.endsWith("." + requestedHeaderName);
+    }
+
+    if (existingDots == 1 && requestedDots == 2) {
+      return requestedHeaderName.endsWith("." + existingHeaderName);
+    }
+
+    return false;
+  }
+
+  private GridHeader withRequestedNameIfNeeded(GridHeader gridHeader, String requestedHeaderName) {
+    if (requestedHeaderName.equals(gridHeader.getName())) {
+      return gridHeader;
+    }
+
+    GridHeader renamed =
+        new GridHeader(
+            requestedHeaderName,
+            gridHeader.getColumn(),
+            gridHeader.getValueType(),
+            gridHeader.isHidden(),
+            gridHeader.isMeta(),
+            gridHeader.getOptionSetObject(),
+            gridHeader.getLegendSetObject());
+
+    if (gridHeader.getStageOffset() != null) {
+      renamed =
+          renamed.withRepeatableStageParams(RepeatableStageParams.of(gridHeader.getStageOffset()));
+    }
+
+    return renamed;
+  }
 }

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/common/query/jsonextractor/SqlRowSetJsonExtractorDelegator.java

@@ -130,6 +130,14 @@ public SqlRowSetJsonExtractorDelegator(
                   + DIMENSION_IDENTIFIER_SEP
                   + "ou";
           dimIdByKey.put(shortFormatKey, dimensionIdentifier);
+        } else if (isEventLevelDataElementDimension(dimensionIdentifier)) {
+          // For stage-scoped data elements, add short format alias
+          // to support headers like programStageUid.dataElementUid.
+          String shortFormatKey =
+              dimensionIdentifier.getProgramStage().getElement().getUid()
+                  + DIMENSION_IDENTIFIER_SEP
+                  + dimensionIdentifier.getDimension().getUid();
+          dimIdByKey.put(shortFormatKey, dimensionIdentifier);
         }
       }
     }
@@ -151,6 +159,11 @@ private static boolean isEventLevelOuDimensionalObject(
             == DimensionParamObjectType.ORGANISATION_UNIT;
   }
 
+  private static boolean isEventLevelDataElementDimension(
+      DimensionIdentifier<DimensionParam> dimIdentifier) {
+    return isDataElement(dimIdentifier);
+  }
+
   @Override
   @SneakyThrows
   public Object getObject(String columnLabel) throws InvalidResultSetAccessException {
@@ -162,6 +175,9 @@ public Object getObject(String columnLabel) throws InvalidResultSetAccessExcepti
     List<JsonEnrollment> enrollments = parseEnrollmentsFromJson(super.getString("enrollments"));
 
     DimensionIdentifier<DimensionParam> dimensionIdentifier = dimIdByKey.get(columnLabel);
+    if (dimensionIdentifier == null) {
+      throw new IllegalQueryException(E7250, columnLabel);
+    }
 
     if (dimensionIdentifier.isEnrollmentDimension()) {
       return getObjectForEnrollments(enrollments, dimensionIdentifier);

dhis-2/dhis-services/dhis-service-analytics/src/main/java/org/hisp/dhis/analytics/trackedentity/query/TrackedEntityFields.java

@@ -34,6 +34,7 @@
 import static org.hisp.dhis.analytics.common.params.dimension.DimensionIdentifierHelper.SUPPORTED_EVENT_STATIC_DIMENSIONS;
 import static org.hisp.dhis.analytics.common.params.dimension.DimensionIdentifierHelper.getCustomLabelOrFullName;
 import static org.hisp.dhis.analytics.common.params.dimension.DimensionIdentifierHelper.getCustomLabelOrHeaderColumnName;
+import static org.hisp.dhis.analytics.common.params.dimension.DimensionIdentifierHelper.isDataElement;
 import static org.hisp.dhis.analytics.common.params.dimension.DimensionIdentifierHelper.isEventLevelStaticDimension;
 import static org.hisp.dhis.analytics.common.params.dimension.DimensionIdentifierHelper.joinedWithPrefixesIfNeeded;
 import static org.hisp.dhis.analytics.trackedentity.query.context.QueryContextConstants.TRACKED_ENTITY_ALIAS;
@@ -251,6 +252,26 @@ private static void addHeaderToMap(
 
       headersMap.put(shortFormatName, shortFormatHeader);
       headersMap.put(offsetHeader.getName(), shortFormatHeader);
+    } else if (isEventLevelDataElementDimension(dimIdentifier)) {
+      // Stage-scoped data elements should also be addressable using short format
+      // (programStageUid.dataElementUid) for consistency with stage-specific static headers.
+      String shortFormatName =
+          dimIdentifier.getProgramStage().getElement().getUid()
+              + DIMENSION_IDENTIFIER_SEP
+              + dimIdentifier.getDimension().getUid();
+
+      GridHeader shortFormatHeader =
+          new GridHeader(
+              shortFormatName,
+              offsetHeader.getColumn(),
+              offsetHeader.getValueType(),
+              offsetHeader.isHidden(),
+              offsetHeader.isMeta(),
+              offsetHeader.getOptionSetObject(),
+              offsetHeader.getLegendSetObject());
+
+      headersMap.put(shortFormatName, shortFormatHeader);
+      headersMap.put(offsetHeader.getName(), offsetHeader);
     } else {
       headersMap.put(offsetHeader.getName(), offsetHeader);
     }
@@ -269,6 +290,11 @@ private static boolean isEventLevelOuDimensionalObject(
             == DimensionParamObjectType.ORGANISATION_UNIT;
   }
 
+  private static boolean isEventLevelDataElementDimension(
+      DimensionIdentifier<DimensionParam> dimIdentifier) {
+    return isDataElement(dimIdentifier);
+  }
+
   /**
    * Adds the stage offset to the given {@link GridHeader} if necessary.
    *

dhis-2/dhis-services/dhis-service-analytics/src/test/java/org/hisp/dhis/analytics/common/processing/DimensionIdentifierConverterTest.java

@@ -671,7 +671,6 @@ void fromStringStageScopedWithOffset() {
   @Test
   void fromStringWithStageUidAndNonStaticDimension() {
     // Given - stage UID (not a program UID) with a non-static dimension (e.g., data element)
-    // This should fail because the first element is not a valid program
     Program program = new Program("prg-1");
     program.setUid("IpHINAT79UW");
     ProgramStage programStage = new ProgramStage("ps-1", program);
@@ -681,16 +680,23 @@ void fromStringWithStageUidAndNonStaticDimension() {
     List<Program> programs = List.of(program);
     String fullDimensionId = "ZkbAXlQUYJG.someDataElement";
 
-    // When - stage UID with non-static, non-event-level dimension
-    IllegalArgumentException thrown =
-        assertThrows(
-            IllegalArgumentException.class, () -> converter.fromString(programs, fullDimensionId));
+    // When
+    DimensionIdentifier<StringUid> dimensionIdentifier =
+        converter.fromString(programs, fullDimensionId);
 
-    // Then - should fail with "program does not exist" since it's not a recognized pattern
+    // Then - should resolve program from stage and keep stage scope
     assertEquals(
-        "Specified program ZkbAXlQUYJG does not exist",
-        thrown.getMessage(),
-        "Should fail as program not found for non-static dimensions");
+        "someDataElement",
+        dimensionIdentifier.getDimension().getUid(),
+        "Dimension uid should be someDataElement");
+    assertEquals(
+        "IpHINAT79UW",
+        dimensionIdentifier.getProgram().getElement().getUid(),
+        "Program uid should be resolved from program stage");
+    assertEquals(
+        "ZkbAXlQUYJG",
+        dimensionIdentifier.getProgramStage().getElement().getUid(),
+        "Stage uid should be ZkbAXlQUYJG");
   }
 
   @Test

dhis-2/dhis-services/dhis-service-analytics/src/test/java/org/hisp/dhis/analytics/common/processing/HeaderParamsHandlerTest.java

@@ -41,11 +41,21 @@
 import org.hisp.dhis.analytics.common.CommonRequestParams;
 import org.hisp.dhis.analytics.common.ContextParams;
 import org.hisp.dhis.analytics.common.params.CommonParsedParams;
+import org.hisp.dhis.analytics.common.params.dimension.DimensionIdentifier;
+import org.hisp.dhis.analytics.common.params.dimension.DimensionParam;
+import org.hisp.dhis.analytics.common.params.dimension.DimensionParamType;
+import org.hisp.dhis.analytics.common.params.dimension.ElementWithOffset;
 import org.hisp.dhis.analytics.common.query.Field;
 import org.hisp.dhis.analytics.trackedentity.TrackedEntityQueryParams;
 import org.hisp.dhis.analytics.trackedentity.TrackedEntityRequestParams;
 import org.hisp.dhis.common.Grid;
+import org.hisp.dhis.common.IdScheme;
 import org.hisp.dhis.common.IllegalQueryException;
+import org.hisp.dhis.common.QueryItem;
+import org.hisp.dhis.common.ValueType;
+import org.hisp.dhis.dataelement.DataElement;
+import org.hisp.dhis.program.Program;
+import org.hisp.dhis.program.ProgramStage;
 import org.hisp.dhis.system.grid.ListGrid;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -156,6 +166,91 @@ void testHandleWithNonExistingParamHeader() {
     assertTrue(ex.getMessage().contains("Header param `non-existing` does not exist"));
   }
 
+  @Test
+  void testHandleWithStageScopedDataElementShortHeader() {
+    // Given
+    String stageUid = "A03MvHHogjR";
+    String dataElementUid = "bx6fsa0t90x";
+    String shortHeader = stageUid + "." + dataElementUid;
+    Program program = new Program("program");
+    program.setUid("IpHINAT79UW");
+    ProgramStage stage = new ProgramStage("stage", program);
+    stage.setUid(stageUid);
+    DataElement dataElement = new DataElement("data element");
+    dataElement.setUid(dataElementUid);
+    dataElement.setValueType(ValueType.TEXT);
+    QueryItem queryItem = new QueryItem(dataElement, null, ValueType.TEXT, null, null);
+    DimensionIdentifier<DimensionParam> dataElementDimension =
+        DimensionIdentifier.of(
+            ElementWithOffset.of(program),
+            ElementWithOffset.of(stage),
+            DimensionParam.ofObject(
+                queryItem, DimensionParamType.DIMENSIONS, IdScheme.UID, List.of()));
+
+    CommonRequestParams requestParams = new CommonRequestParams();
+    requestParams.setHeaders(new LinkedHashSet<>(List.of(shortHeader)));
+    ContextParams<TrackedEntityRequestParams, TrackedEntityQueryParams> contextParams =
+        ContextParams.<TrackedEntityRequestParams, TrackedEntityQueryParams>builder()
+            .commonRaw(requestParams)
+            .commonParsed(
+                CommonParsedParams.builder()
+                    .dimensionIdentifiers(List.of(dataElementDimension))
+                    .build())
+            .build();
+    List<Field> fields = List.of(ofUnquoted("ev", of("anyName"), dataElementDimension.getKey()));
+    Grid grid = new ListGrid();
+
+    // When
+    headerParamsHandler.handle(grid, contextParams, fields);
+
+    // Then
+    assertEquals(1, grid.getHeaders().size());
+    assertEquals(shortHeader, grid.getHeaders().get(0).getName());
+  }
+
+  @Test
+  void testHandleWithStageScopedDataElementFullHeader() {
+    // Given
+    String stageUid = "A03MvHHogjR";
+    String programUid = "IpHINAT79UW";
+    String dataElementUid = "bx6fsa0t90x";
+    String fullHeader = programUid + "." + stageUid + "." + dataElementUid;
+    Program program = new Program("program");
+    program.setUid(programUid);
+    ProgramStage stage = new ProgramStage("stage", program);
+    stage.setUid(stageUid);
+    DataElement dataElement = new DataElement("data element");
+    dataElement.setUid(dataElementUid);
+    dataElement.setValueType(ValueType.TEXT);
+    QueryItem queryItem = new QueryItem(dataElement, null, ValueType.TEXT, null, null);
+    DimensionIdentifier<DimensionParam> dataElementDimension =
+        DimensionIdentifier.of(
+            ElementWithOffset.of(program),
+            ElementWithOffset.of(stage),
+            DimensionParam.ofObject(
+                queryItem, DimensionParamType.DIMENSIONS, IdScheme.UID, List.of()));
+
+    CommonRequestParams requestParams = new CommonRequestParams();
+    requestParams.setHeaders(new LinkedHashSet<>(List.of(fullHeader)));
+    ContextParams<TrackedEntityRequestParams, TrackedEntityQueryParams> contextParams =
+        ContextParams.<TrackedEntityRequestParams, TrackedEntityQueryParams>builder()
+            .commonRaw(requestParams)
+            .commonParsed(
+                CommonParsedParams.builder()
+                    .dimensionIdentifiers(List.of(dataElementDimension))
+                    .build())
+            .build();
+    List<Field> fields = List.of(ofUnquoted("ev", of("anyName"), dataElementDimension.getKey()));
+    Grid grid = new ListGrid();
+
+    // When
+    headerParamsHandler.handle(grid, contextParams, fields);
+
+    // Then
+    assertEquals(1, grid.getHeaders().size());
+    assertEquals(fullHeader, grid.getHeaders().get(0).getName());
+  }
+
   private CommonRequestParams stubCommonParamsWithHeaders() {
     Set<String> headers = new LinkedHashSet<>(List.of("oucode", "ouname", "lastupdated"));