From 1eec41a136b3405e104b6c9cc9aa200e3f8cd8fe Mon Sep 17 00:00:00 2001
From: Felix Feisst <feisst.felix@gmail.com>
Date: Thu, 16 Mar 2017 21:02:08 +0100
Subject: [PATCH] HHH-11573 Query based on type expressions

---
 .../internal/tools/query/Parameters.java      |   5 +
 .../hibernate/envers/query/AuditEntity.java   |  44 ++++++-
 .../internal/EntityTypeAuditExpression.java   |  47 +++++++
 .../query/EntityTypeQueryTest.java            | 117 ++++++++++++++++++
 4 files changed, 212 insertions(+), 1 deletion(-)
 create mode 100644 hibernate-envers/src/main/java/org/hibernate/envers/query/criteria/internal/EntityTypeAuditExpression.java
 create mode 100644 hibernate-envers/src/test/java/org/hibernate/envers/test/integration/query/EntityTypeQueryTest.java

diff --git a/hibernate-envers/src/main/java/org/hibernate/envers/internal/tools/query/Parameters.java b/hibernate-envers/src/main/java/org/hibernate/envers/internal/tools/query/Parameters.java
index 3212a26bfa..32507b54bd 100644
--- a/hibernate-envers/src/main/java/org/hibernate/envers/internal/tools/query/Parameters.java
+++ b/hibernate-envers/src/main/java/org/hibernate/envers/internal/tools/query/Parameters.java
@@ -482,6 +482,11 @@ public class Parameters {
 		expressions.add( expression.toString() );
 	}
 
+	public void addEntityTypeRestriction(String alias, String entityName) {
+		String expression = String.format( "type(%s) = %s", alias, entityName );
+		expressions.add( expression );
+	}
+
 	private void append(StringBuilder sb, String toAppend, MutableBoolean isFirst) {
 		if ( !isFirst.isSet() ) {
 			sb.append( " " ).append( connective ).append( " " );
diff --git a/hibernate-envers/src/main/java/org/hibernate/envers/query/AuditEntity.java b/hibernate-envers/src/main/java/org/hibernate/envers/query/AuditEntity.java
index 5cf2889b70..9e7264e499 100644
--- a/hibernate-envers/src/main/java/org/hibernate/envers/query/AuditEntity.java
+++ b/hibernate-envers/src/main/java/org/hibernate/envers/query/AuditEntity.java
@@ -11,6 +11,7 @@ import java.util.Collections;
 import java.util.List;
 
 import org.hibernate.envers.RevisionType;
+import org.hibernate.envers.internal.tools.EntityTools;
 import org.hibernate.envers.query.criteria.AuditConjunction;
 import org.hibernate.envers.query.criteria.AuditCriterion;
 import org.hibernate.envers.query.criteria.AuditDisjunction;
@@ -18,6 +19,7 @@ import org.hibernate.envers.query.criteria.AuditFunction;
 import org.hibernate.envers.query.criteria.AuditId;
 import org.hibernate.envers.query.criteria.AuditProperty;
 import org.hibernate.envers.query.criteria.AuditRelatedId;
+import org.hibernate.envers.query.criteria.internal.EntityTypeAuditExpression;
 import org.hibernate.envers.query.criteria.internal.LogicalAuditExpression;
 import org.hibernate.envers.query.criteria.internal.NotAuditExpression;
 import org.hibernate.envers.query.internal.property.EntityPropertyName;
@@ -192,7 +194,7 @@ public class AuditEntity {
 	 * 		<li>AuditEntity.function("concat", AuditEntity.function("upper", AuditEntity.property("prop1")),
 	 * 				AuditEntity.function("substring", AuditEntity.property("prop2"), 1, 2))</li>
 	 * </ul>
-	 * 
+	 *
 	 * @param function the name of the function
 	 * @param arguments the arguments of the function. A function argument can either be
 	 * <ul>
@@ -207,4 +209,44 @@ public class AuditEntity {
 		Collections.addAll( argumentList, arguments );
 		return new AuditFunction( function, argumentList );
 	}
+
+	/**
+	 * Adds a restriction for the type of the current entity.
+	 *
+	 * @param the entity type to restrict the current alias to
+	 */
+	public static AuditCriterion entityType(final Class<?> type) {
+		return entityType( null, type );
+	}
+
+	/**
+	 * Adds a restriction for the type of the current entity.
+	 *
+	 * @param entityName the entity name to restrict the current alias to
+	 */
+	public static AuditCriterion entityType(final String entityName) {
+		return entityType( null, entityName );
+	}
+
+	/**
+	 * Adds a restriction for the type of the entity of the specified alias.
+	 *
+	 * @param alias the alias to restrict. If null is specified, the current alias is used
+	 * @param type the entity type to restrict the alias to
+	 */
+	public static AuditCriterion entityType(final String alias, final Class<?> type) {
+		Class<?> unproxiedType = EntityTools.getTargetClassIfProxied( type );
+		return entityType( alias, unproxiedType.getName() );
+	}
+
+	/**
+	 * Adds a restriction for the type of the entity of the specified alias.
+	 *
+	 * @param alias the alias to restrict. If null is specified, the current alias is used
+	 * @param entityName the entity name to restrict the alias to
+	 * @return
+	 */
+	public static AuditCriterion entityType(final String alias, final String entityName) {
+		return new EntityTypeAuditExpression( alias, entityName );
+	}
 }
diff --git a/hibernate-envers/src/main/java/org/hibernate/envers/query/criteria/internal/EntityTypeAuditExpression.java b/hibernate-envers/src/main/java/org/hibernate/envers/query/criteria/internal/EntityTypeAuditExpression.java
new file mode 100644
index 0000000000..74d627c6cc
--- /dev/null
+++ b/hibernate-envers/src/main/java/org/hibernate/envers/query/criteria/internal/EntityTypeAuditExpression.java
@@ -0,0 +1,47 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.envers.query.criteria.internal;
+
+import java.util.Map;
+
+import org.hibernate.envers.boot.internal.EnversService;
+import org.hibernate.envers.internal.reader.AuditReaderImplementor;
+import org.hibernate.envers.internal.tools.query.Parameters;
+import org.hibernate.envers.internal.tools.query.QueryBuilder;
+import org.hibernate.envers.query.criteria.AuditCriterion;
+
+/**
+ * @author Felix Feisst (feisst dot felix at gmail dot com)
+ */
+public class EntityTypeAuditExpression implements AuditCriterion {
+
+	private String alias;
+	private String entityName;
+
+	public EntityTypeAuditExpression(
+			String alias,
+			String entityName) {
+		this.alias = alias;
+		this.entityName = entityName;
+	}
+
+	@Override
+	public void addToQuery(
+			EnversService enversService,
+			AuditReaderImplementor auditReader,
+			Map<String, String> aliasToEntityNameMap,
+			String baseAlias,
+			QueryBuilder qb,
+			Parameters parameters) {
+		String effectiveAlias = alias == null ? baseAlias : alias;
+		String effectiveEntityName = entityName;
+		if ( enversService.getEntitiesConfigurations().isVersioned( effectiveEntityName ) ) {
+			effectiveEntityName = enversService.getConfig().getAuditEntityName( effectiveEntityName );
+		}
+		parameters.addEntityTypeRestriction( effectiveAlias, effectiveEntityName );
+	}
+}
diff --git a/hibernate-envers/src/test/java/org/hibernate/envers/test/integration/query/EntityTypeQueryTest.java b/hibernate-envers/src/test/java/org/hibernate/envers/test/integration/query/EntityTypeQueryTest.java
new file mode 100644
index 0000000000..e4b56c8bf4
--- /dev/null
+++ b/hibernate-envers/src/test/java/org/hibernate/envers/test/integration/query/EntityTypeQueryTest.java
@@ -0,0 +1,117 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.envers.test.integration.query;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import jakarta.persistence.Entity;
+import jakarta.persistence.EntityManager;
+import jakarta.persistence.GeneratedValue;
+import jakarta.persistence.Id;
+import jakarta.persistence.Inheritance;
+import jakarta.persistence.InheritanceType;
+
+import org.hibernate.envers.Audited;
+import org.hibernate.envers.query.AuditEntity;
+import org.hibernate.orm.test.envers.BaseEnversJPAFunctionalTestCase;
+import org.hibernate.orm.test.envers.Priority;
+import org.junit.Test;
+
+/**
+ * @author Felix Feisst (feisst dot felix at gmail dot com)
+ */
+public class EntityTypeQueryTest extends BaseEnversJPAFunctionalTestCase {
+
+	@Entity
+	@Audited
+	@Inheritance(strategy = InheritanceType.JOINED)
+	public static class EntityA {
+
+		@Id
+		@GeneratedValue
+		private Long id;
+		private String name;
+
+		public Long getId() {
+			return id;
+		}
+
+		public void setId(Long id) {
+			this.id = id;
+		}
+
+		public String getName() {
+			return name;
+		}
+
+		public void setName(String name) {
+			this.name = name;
+		}
+
+	}
+
+	@Entity
+	@Audited
+	public static class EntityB extends EntityA {
+
+	}
+
+	private EntityA a1;
+	private EntityA a2;
+	private EntityB b1;
+	private EntityB b2;
+
+	@Override
+	protected Class<?>[] getAnnotatedClasses() {
+		return new Class<?>[]{ EntityA.class, EntityB.class };
+	}
+
+	@Test
+	@Priority(10)
+	public void initData() {
+		EntityManager em = getEntityManager();
+		em.getTransaction().begin();
+		a1 = new EntityA();
+		a1.setName( "a1" );
+		em.persist( a1 );
+		a2 = new EntityA();
+		a2.setName( "a2" );
+		em.persist( a2 );
+		b1 = new EntityB();
+		b1.setName( "b1" );
+		em.persist( b1 );
+		b2 = new EntityB();
+		b2.setName( "b2" );
+		em.persist( b2 );
+		em.getTransaction().commit();
+	}
+
+	@Test
+	public void testRestrictToSubType() {
+		List<?> list = getAuditReader().createQuery().forEntitiesAtRevision( EntityA.class, 1 )
+				.add( AuditEntity.entityType( EntityB.class ) ).addProjection( AuditEntity.property( "name" ) ).getResultList();
+		assertEquals( "Expected only entities of type EntityB to be selected", list( "b1", "b2" ), list );
+	}
+
+	@Test
+	public void testRestrictToSuperType() {
+		List<?> list = getAuditReader().createQuery().forEntitiesAtRevision( EntityA.class, 1 )
+				.add( AuditEntity.entityType( EntityA.class ) ).addProjection( AuditEntity.property( "name" ) ).getResultList();
+		assertEquals( "Expected only entities of type EntityA to be selected", list( "a1", "a2" ), list );
+	}
+
+	private List<String> list(final String... elements) {
+		final List<String> result = new ArrayList<>();
+		Collections.addAll( result, elements );
+		return result;
+	}
+
+}