[OPENJPA-2961] Include subclasses of the treated class in the TREAT discriminator filter

openjpa-jdbc/src/main/java/org/apache/openjpa/jdbc/kernel/exps/BindVariableAndExpression.java

@@ -58,7 +58,7 @@ public void appendTo(Select sel, ExpContext ctx, ExpState state,
         BinaryOpExpState bstate = (BinaryOpExpState) state;
         // Append discriminator condition from TREAT (if any)
         if (_bind.hasTreatDiscriminator()) {
-            _bind.appendTreatDiscriminator(sel, bstate.state1, buf);
+            _bind.appendTreatDiscriminator(sel, ctx, bstate.state1, buf);
             buf.append(" AND ");
         }
         boolean or = _exp instanceof OrExpression;

openjpa-jdbc/src/main/java/org/apache/openjpa/jdbc/kernel/exps/BindVariableExpression.java

@@ -80,9 +80,10 @@ boolean hasTreatDiscriminator() {
                 Discriminator disc = sup.getDiscriminator();
                 if (disc != null) {
                     Column[] cols = disc.getColumns();
-                    Object discVal = varMapping.getDiscriminator() != null
-                        ? varMapping.getDiscriminator().getValue() : null;
-                    return cols != null && cols.length > 0 && discVal != null;
+                    return cols != null && cols.length > 0
+                        && varMapping.getDiscriminator() != null
+                        && varMapping.getDiscriminator()
+                            .hasClassConditions(varMapping, true);
                 }
             }
         }
@@ -91,26 +92,22 @@ boolean hasTreatDiscriminator() {
 
     /**
      * Appends the discriminator condition for TREAT-narrowed variables.
+     * The condition matches the treated class and all of its subclasses.
      */
-    void appendTreatDiscriminator(Select sel, ExpState state, SQLBuffer buf) {
+    void appendTreatDiscriminator(Select sel, ExpContext ctx, ExpState state,
+        SQLBuffer buf) {
         ClassMapping varMapping = (ClassMapping) _var.getMetaData();
-        ClassMapping sup = varMapping;
-        while (sup.getMappedPCSuperclassMapping() != null) {
-            sup = sup.getMappedPCSuperclassMapping();
-        }
-        Discriminator disc = sup.getDiscriminator();
-        Column[] cols = disc.getColumns();
-        Object discVal = varMapping.getDiscriminator().getValue();
-        buf.append(sel.getColumnAlias(cols[0], state.joins));
-        buf.append(" = ");
-        buf.appendValue(discVal, cols[0]);
+        Discriminator disc = varMapping.getDiscriminator();
+        ctx.store.loadSubclasses(varMapping);
+        buf.append(disc.getClassConditions(sel, state.joins, varMapping,
+            true));
     }
 
     @Override
     public void appendTo(Select sel, ExpContext ctx, ExpState state,
         SQLBuffer buf) {
         if (hasTreatDiscriminator()) {
-            appendTreatDiscriminator(sel, state, buf);
+            appendTreatDiscriminator(sel, ctx, state, buf);
         } else {
             buf.append("1 = 1");
         }

openjpa-kernel/src/main/java/org/apache/openjpa/kernel/jpql/JPQLExpressionBuilder.java

@@ -98,6 +98,9 @@ public class JPQLExpressionBuilder
     private int aliasCount = 0;
     private boolean inAssignSubselectProjection = false;
     private boolean hasParameterizedInExpression = false;
+    // type restrictions of the TREAT paths of the predicate being evaluated
+    private List<Expression> treatRestrictions;
+    private Context treatRestrictionsContext;
 
     /**
      * Constructor.
@@ -2236,7 +2239,27 @@ private Value getTreatPath(JPQLNode node) {
         String schemaName = assemble(schemaNode);
         ClassMetaData treatMeta = getClassMetaData(schemaName, true);
 
-        // Resolve the base path using the variable
+        Path path = getTreatBasePath(node, name);
+        if (treatRestrictions != null && treatRestrictionsContext == ctx()) {
+            treatRestrictions.add(factory.isInstance(getTreatBasePath(node, name),
+                treatMeta.getDescribedType()));
+        }
+
+        // Override the metadata to the treat target type
+        path.setMetaData(treatMeta);
+
+        // Walk through the remaining children (path components after the dot)
+        for (int i = 2; i < node.children.length; i++) {
+            path = (Path) traversePath(path, node.children[i].text, false, true);
+        }
+
+        return path;
+    }
+
+    /**
+     * Resolves the identifier of a TREAT(identifier AS Type) path to a new path.
+     */
+    private Path getTreatBasePath(JPQLNode node, String name) {
         Path path = null;
         final Value val = getVariable(name, false);
 
@@ -2260,15 +2283,6 @@ private Value getTreatPath(JPQLNode node) {
         }
 
         path.setSchemaAlias(name);
-
-        // Override the metadata to the treat target type
-        path.setMetaData(treatMeta);
-
-        // Walk through the remaining children (path components after the dot)
-        for (int i = 2; i < node.children.length; i++) {
-            path = (Path) traversePath(path, node.children[i].text, false, true);
-        }
-
         return path;
     }
 
@@ -2345,12 +2359,26 @@ protected Class<?> getDeclaredVariableType(String name) {
      * Returns an Expression for the given node by eval'ing it.
      */
     private Expression getExpression(JPQLNode node) {
-        Object exp = eval(node);
+        List<Expression> outerRestrictions = treatRestrictions;
+        Context outerRestrictionsContext = treatRestrictionsContext;
+        treatRestrictions = new ArrayList<>();
+        treatRestrictionsContext = ctx();
+        try {
+            Object exp = eval(node);
+
+            // check for boolean values used as expressions
+            Expression result = exp instanceof Expression
+                ? (Expression) exp : factory.asExpression((Value) exp);
 
-        // check for boolean values used as expressions
-        if (!(exp instanceof Expression))
-            return factory.asExpression((Value) exp);
-        return (Expression) exp;
+            // a predicate over TREAT(x AS Type) is false if x is not a Type
+            for (Expression restriction : treatRestrictions) {
+                result = and(restriction, result);
+            }
+            return result;
+        } finally {
+            treatRestrictions = outerRestrictions;
+            treatRestrictionsContext = outerRestrictionsContext;
+        }
     }
 
     /**

openjpa-persistence-jdbc/src/test/java/org/apache/openjpa/persistence/jpql/treatjoinon/TCustomSoftwareProduct.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.openjpa.persistence.jpql.treatjoinon;
+
+import jakarta.persistence.DiscriminatorValue;
+import jakarta.persistence.Entity;
+import jakarta.persistence.Table;
+
+@Entity
+@Table(name = "TPRODUCT")
+@DiscriminatorValue("CS")
+public class TCustomSoftwareProduct extends TSoftwareProduct {
+
+	private int customizationHours;
+
+	public int getCustomizationHours() {
+		return customizationHours;
+	}
+
+	public void setCustomizationHours(int customizationHours) {
+		this.customizationHours = customizationHours;
+	}
+
+}

openjpa-persistence-jdbc/src/test/java/org/apache/openjpa/persistence/jpql/treatjoinon/TGameProduct.java

@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.openjpa.persistence.jpql.treatjoinon;
+
+import jakarta.persistence.DiscriminatorValue;
+import jakarta.persistence.Entity;
+import jakarta.persistence.Table;
+
+/**
+ * Third level of the TProduct single table hierarchy.
+ */
+@Entity
+@Table(name = "TPRODUCT")
+@DiscriminatorValue("GAME")
+public class TGameProduct extends TSoftwareProduct {
+    private String genre;
+
+    public String getGenre() { return genre; }
+    public void setGenre(String genre) { this.genre = genre; }
+}

openjpa-persistence-jdbc/src/test/java/org/apache/openjpa/persistence/jpql/treatjoinon/TestTreatJoinOnEmbeddable.java

@@ -39,10 +39,10 @@ public class TestTreatJoinOnEmbeddable extends SingleEMFTestCase {
 
     @Override
     public void setUp() {
-        setUp(TProduct.class, TSoftwareProduct.class,
-            TLineItem.class, TOrder.class,
-            TCustomer.class, TCountry.class,
-            DROP_TABLES);
+        setUp(TProduct.class, TSoftwareProduct.class, 
+                TCustomSoftwareProduct.class, TLineItem.class, 
+                TOrder.class, TCustomer.class, TCountry.class,
+                DROP_TABLES);
         createTestData();
     }
 
@@ -65,6 +65,12 @@ private void createTestData() {
         sp3.setName("Software C");
         sp3.setRevisionNumber(1.0);
         em.persist(sp3);
+
+        TCustomSoftwareProduct cs1 = new TCustomSoftwareProduct();
+        cs1.setName("Custom Software A");
+        cs1.setRevisionNumber(1.0);
+        cs1.setCustomizationHours(3);
+        em.persist(cs1);
 
         // Create a hardware product (base type)
         TProduct hw = new TProduct();
@@ -95,6 +101,15 @@ private void createTestData() {
         li3.setProduct(hw);
         li3.setOrder(order2);
         em.persist(li3);
+
+        TOrder order3 = new TOrder();
+        em.persist(order3);
+
+        TLineItem li4 = new TLineItem();
+        li4.setQuantity(1);
+        li4.setProduct(cs1);
+        li4.setOrder(order3);
+        em.persist(li4);
 
         // Create customer with embedded country
         TCustomer cust = new TCustomer();
@@ -121,10 +136,25 @@ public void testTreatInWhereClause() {
             "SELECT p.name FROM TProduct p WHERE TREAT(p AS TSoftwareProduct).revisionNumber = 1.0",
             String.class).getResultList();
 
-        Collections.sort(results);
-        assertEquals(2, results.size());
+        assertEquals(3, results.size());
         assertTrue(results.contains("Software A"));
         assertTrue(results.contains("Software C"));
+        assertTrue(results.contains("Custom Software A"));
+
+        em.getTransaction().commit();
+        em.close();
+    }
+
+    public void testTreatLeafInWhereClause() {
+        EntityManager em = emf.createEntityManager();
+        em.getTransaction().begin();
+
+        List<String> results = em.createQuery(
+            "SELECT p.name FROM TProduct p WHERE TREAT(p AS TCustomSoftwareProduct).revisionNumber = 1.0",
+            String.class).getResultList();
+
+        assertEquals(1, results.size());
+        assertTrue(results.contains("Custom Software A"));
 
         em.getTransaction().commit();
         em.close();
@@ -142,9 +172,10 @@ public void testTreatWithDoubleSuffix() {
             "SELECT p.name FROM TProduct p WHERE TREAT(p AS TSoftwareProduct).revisionNumber = 1.0D",
             String.class).getResultList();
 
-        assertEquals(2, results.size());
+        assertEquals(3, results.size());
         assertTrue(results.contains("Software A"));
         assertTrue(results.contains("Software C"));
+        assertTrue(results.contains("Custom Software A"));
 
         em.getTransaction().commit();
         em.close();
@@ -162,9 +193,10 @@ public void testTreatWithScientificNotation() {
             "SELECT p.name FROM TProduct p WHERE TREAT(p AS TSoftwareProduct).revisionNumber = 1E0",
             String.class).getResultList();
 
-        assertEquals(2, results.size());
+        assertEquals(3, results.size());
         assertTrue(results.contains("Software A"));
         assertTrue(results.contains("Software C"));
+        assertTrue(results.contains("Custom Software A"));
 
         em.getTransaction().commit();
         em.close();
@@ -183,12 +215,12 @@ public void testTreatJoinClass() {
             String.class).getResultList();
 
         // TREAT join should only return software products (not hardware)
-        // We have 2 line items with software products (sp1, sp2) and 1 with hardware
+        // We have 3 line items with software products (sp1, sp2, cp1) and 1 with hardware
         assertNotNull(results);
-        Collections.sort(results);
-        assertEquals("TREAT join should return only software products", 2, results.size());
+        assertEquals("TREAT join should return software products, including the especialized ones", 3, results.size());
         assertTrue(results.contains("Software A"));
         assertTrue(results.contains("Software B"));
+        assertTrue(results.contains("Custom Software A"));
         // Hardware X should NOT be in the results due to TREAT filtering
         assertFalse(results.contains("Hardware X"));