HHH-17817 - Add option to Enable Byte Code Enhancement for specific classes

tooling/hibernate-enhance-maven-plugin/src/main/java/org/hibernate/orm/tooling/maven/MavenEnhancePlugin.java

@@ -15,10 +15,8 @@ import java.io.OutputStream;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
-import java.util.ArrayList;
+import java.nio.file.Path;
-import java.util.Collections;
+import java.util.*;
-import java.util.List;
-import java.util.Set;
 
 import org.apache.maven.artifact.Artifact;
 import org.apache.maven.plugin.AbstractMojo;
@@ -67,6 +65,9 @@ public class MavenEnhancePlugin extends AbstractMojo {
 	@Parameter(property = "dir", defaultValue = "${project.build.outputDirectory}")
 	private String dir;
 
+	@Parameter(property = "classNames", defaultValue = "")
+	private String classNames;
+
 	@Parameter(property = "failOnError", defaultValue = "true")
 	private boolean failOnError = true;
 
@@ -110,6 +111,11 @@ public class MavenEnhancePlugin extends AbstractMojo {
 			return;
 		}
 
+		List<String> classesToEnhance = new ArrayList<>();
+		if(classNames.length() >= 1) {
+			classesToEnhance = Arrays.asList(classNames.split(","));
+		}
+
 		log.info( "Starting Hibernate enhancement for classes on " + dir );
 		final ClassLoader classLoader = toClassLoader( Collections.singletonList( new File( base ) ) );
 
@@ -168,6 +174,12 @@ public class MavenEnhancePlugin extends AbstractMojo {
 			}
 			for ( File file : sourceSet ) {
 
+				String className = determineClassName(root, file);
+
+				if(! (classesToEnhance.size()==0 || classesToEnhance.contains(className))) {
+					continue;
+				}
+
 				final byte[] enhancedBytecode = doEnhancement( file, enhancer );
 
 				if ( enhancedBytecode == null ) {
@@ -355,4 +367,14 @@ public class MavenEnhancePlugin extends AbstractMojo {
 			}
 		}
 	}
+
+	private String determineClassName(File root, File javaClassFile) {
+		final Path relativeClassPath = root.toPath().relativize( javaClassFile.toPath() );
+		final String relativeClassPathString = relativeClassPath.toString();
+		final String classNameBase = relativeClassPathString.substring(
+				0,
+				relativeClassPathString.length() - ".class".length()
+		);
+		return classNameBase.replace( File.separatorChar, '.' );
+	}
 }

tooling/hibernate-enhance-maven-plugin/src/main/resources/META-INF/maven/org.hibernate.orm.tooling/hibernate-enhance-maven-plugin/plugin-help.xml

@@ -45,6 +45,13 @@
           <editable>true</editable>
           <description>Base directory where to search for .class files</description>
         </parameter>
+        <parameter>
+          <name>classNames</name>
+          <type>java.lang.String</type>
+          <required>false</required>
+          <editable>true</editable>
+          <description>Comma separated string of class names for which enhancement needs to be done</description>
+        </parameter>
         <parameter>
           <name>failOnError</name>
           <type>java.lang.Boolean</type>

tooling/hibernate-enhance-maven-plugin/src/main/resources/META-INF/maven/plugin.xml

@@ -46,6 +46,13 @@
           <editable>true</editable>
           <description>Base directory where to search for .class files</description>
         </parameter>
+        <parameter>
+          <name>classNames</name>
+          <type>java.lang.String</type>
+          <required>false</required>
+          <editable>true</editable>
+          <description>Comma separated string of class names for which enhancement needs to be done</description>
+        </parameter>
         <parameter>
           <name>failOnError</name>
           <type>java.lang.Boolean</type>

tooling/hibernate-enhance-maven-plugin/src/test/java/org/hibernate/orm/tooling/maven/MavenEnhancePluginTest.java

@@ -42,6 +42,7 @@ public class MavenEnhancePluginTest {
 
         setVariableValueToObject( plugin, "base", baseDir.getAbsolutePath() );
         setVariableValueToObject( plugin, "dir", baseDir.getAbsolutePath() );
+        setVariableValueToObject( plugin, "classNames", "" );
 
         setVariableValueToObject( plugin, "failOnError", true );
         setVariableValueToObject( plugin, "enableLazyInitialization", true );

tooling/hibernate-gradle-plugin/src/main/java/org/hibernate/orm/tooling/gradle/enhance/EnhancementHelper.java

@@ -10,6 +10,7 @@ import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.nio.file.Files;
+import java.util.List;
 
 import org.gradle.api.GradleException;
 import org.gradle.api.Project;
@@ -40,6 +41,9 @@ public class EnhancementHelper {
 		final File classesDir = classesDirectory.getAsFile();
 
 		final EnhancementSpec enhancementDsl = ormDsl.getEnhancement();
+
+		List<String> classesToEnhance = enhancementDsl.getClassNames().get();
+
 		if ( !enhancementDsl.getEnableLazyInitialization().get() ) {
 			project.getLogger().warn( "The 'enableLazyInitialization' configuration is deprecated and will be removed. Set the value to 'true' to get rid of this warning" );
 		}
@@ -49,7 +53,7 @@ public class EnhancementHelper {
 		final Enhancer enhancer = generateEnhancer( classLoader, ormDsl );
 
 		discoverTypes( classesDir, classesDir, enhancer, project );
-		doEnhancement( classesDir, classesDir, enhancer, project );
+		doEnhancement( classesDir, classesDir, enhancer, project, classesToEnhance );
 	}
 
 	private static void discoverTypes(File classesDir, File dir, Enhancer enhancer, Project project) {
@@ -72,15 +76,19 @@ public class EnhancementHelper {
 		}
 	}
 
-	private static void doEnhancement(File classesDir, File dir, Enhancer enhancer, Project project) {
+	private static void doEnhancement(File classesDir, File dir, Enhancer enhancer, Project project, List<String> classesToEnhance) {
 		for ( File subLocation : dir.listFiles() ) {
 			if ( subLocation.isDirectory() ) {
-				doEnhancement( classesDir, subLocation, enhancer, project );
+				doEnhancement( classesDir, subLocation, enhancer, project, classesToEnhance );
 			}
 			else if ( subLocation.isFile() && subLocation.getName().endsWith( ".class" ) ) {
 				final String className = determineClassName( classesDir, subLocation );
 				final long lastModified = subLocation.lastModified();
 
+				if(! (classesToEnhance.size()==0 || classesToEnhance.contains(className))) {
+					continue;
+				}
+
 				enhance( subLocation, className, enhancer, project );
 
 				final boolean timestampReset = subLocation.setLastModified( lastModified );

tooling/hibernate-gradle-plugin/src/main/java/org/hibernate/orm/tooling/gradle/enhance/EnhancementSpec.java

@@ -9,11 +9,16 @@ package org.hibernate.orm.tooling.gradle.enhance;
 import javax.inject.Inject;
 
 import org.gradle.api.Project;
+import org.gradle.api.provider.ListProperty;
 import org.gradle.api.provider.Property;
+import org.gradle.api.provider.Provider;
 import org.gradle.api.tasks.SourceSetContainer;
 
 import org.hibernate.orm.tooling.gradle.HibernateOrmSpec;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * DSL extension for configuring bytecode enhancement - available as `project.hibernateOrm.enhancement`
  */
@@ -28,6 +33,7 @@ public class EnhancementSpec {
 	private final Property<Boolean> enableDirtyTracking;
 	private final Property<Boolean> enableAssociationManagement;
 	private final Property<Boolean> enableExtendedEnhancement;
+	private final ListProperty<String> classNames;
 
 
 	@Inject
@@ -38,6 +44,7 @@ public class EnhancementSpec {
 		enableDirtyTracking = makeProperty( project ).convention( true );
 		enableAssociationManagement = makeProperty( project ).convention( false );
 		enableExtendedEnhancement = makeProperty( project ).convention( false );
+		classNames = project.getObjects().listProperty(String.class).convention(new ArrayList<>());
 	}
 
 	@SuppressWarnings( "UnstableApiUsage" )
@@ -83,6 +90,13 @@ public class EnhancementSpec {
 		return enableExtendedEnhancement;
 	}
 
+	/**
+	 * Returns the classes on which enhancement needs to be done
+	 */
+	public ListProperty<String> getClassNames() {
+		return classNames;
+	}
+
 
 	/**
 	 * @deprecated See the Gradle property naming <a href="https://docs.gradle.org/current/userguide/lazy_configuration.html#lazy_configuration_faqs">guidelines</a>
@@ -195,4 +209,16 @@ public class EnhancementSpec {
 	public void extendedEnhancement(boolean enable) {
 		setEnableExtendedEnhancement( enable );
 	}
+
+	public void setClassNames(List<String> classNames) {
+		this.classNames.set(classNames);
+	}
+
+	public void setClassNames(Provider<List<String>> classNames) {
+		this.classNames.set(classNames);
+	}
+
+	public void includeClassName(String className) {
+		this.classNames.add(className);
+	}
 }