Make Gradle's ProjectImportAction more flexible.

Add model provider interface. It will allow a plugin to decide on-the-fly, what models to collect inside Gradle runtime.

plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/AbstractProjectResolverExtension.java

@@ -31,6 +31,7 @@ import org.gradle.tooling.model.idea.IdeaProject;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.plugins.gradle.model.ProjectImportExtraModelProvider;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -124,6 +125,12 @@ public abstract class AbstractProjectResolverExtension implements GradleProjectR
     return Collections.emptySet();
   }
 
+  @NotNull
+  @Override
+  public ProjectImportExtraModelProvider getExtraModelProvider() {
+    return new ClassSetProjectImportExtraModelProvider(getExtraProjectModelClasses());
+  }
+
   @NotNull
   @Override
   public Set<Class> getToolingExtensionsClasses() {

plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/BaseGradleProjectResolverExtension.java

@@ -693,6 +693,12 @@ public class BaseGradleProjectResolverExtension implements GradleProjectResolver
     return result;
   }
 
+  @NotNull
+  @Override
+  public ProjectImportExtraModelProvider getExtraModelProvider() {
+    return new ClassSetProjectImportExtraModelProvider(getExtraProjectModelClasses());
+  }
+
   @NotNull
   @Override
   public Set<Class> getToolingExtensionsClasses() {

plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/ClassSetProjectImportExtraModelProvider.java

+// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package org.jetbrains.plugins.gradle.service.project;
+
+import org.gradle.tooling.BuildController;
+import org.gradle.tooling.model.idea.IdeaModule;
+import org.gradle.tooling.model.idea.IdeaProject;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.plugins.gradle.model.ProjectImportExtraModelProvider;
+
+import java.util.Collection;
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+public class ClassSetProjectImportExtraModelProvider implements ProjectImportExtraModelProvider {
+  @NotNull private Set<Class> classSet = new LinkedHashSet<>();
+
+  public ClassSetProjectImportExtraModelProvider(@NotNull Collection<Class> classes) {
+    classSet.addAll(classes);
+  }
+
+  @Override
+  public void populateBuildModels(@NotNull BuildController controller, @NotNull IdeaProject project, @NotNull BuildModelConsumer consumer) {
+    // Do nothing, this provider only works on the project model level
+  }
+
+  @Override
+  public void populateProjectModels(@NotNull BuildController controller,
+                                    @Nullable IdeaModule module,
+                                    @NotNull ProjectModelConsumer modelConsumer) {
+    for (Class<?> aClass : classSet) {
+      Object instance = controller.findModel(module, aClass);
+      if (instance != null) {
+        modelConsumer.consume(instance, aClass);
+      }
+    }
+  }
+}

plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/GradleProjectResolver.java

@@ -66,7 +66,8 @@ import java.io.IOException;
 import java.util.*;
 
 import static com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil.*;
-import static org.jetbrains.plugins.gradle.service.project.GradleProjectResolverUtil.*;
+import static org.jetbrains.plugins.gradle.service.project.GradleProjectResolverUtil.getDefaultModuleTypeId;
+import static org.jetbrains.plugins.gradle.service.project.GradleProjectResolverUtil.getModuleId;
 
 /**
  * @author Denis Zhdanov, Vladislav Soroka
@@ -201,7 +202,7 @@ public class GradleProjectResolver implements ExternalSystemProjectResolver<Grad
     if(resolverCtx.isPreviewMode()){
       executionSettings.withArgument("-Didea.isPreviewMode=true");
       final Set<Class> previewLightWeightToolingModels = ContainerUtil.set(ExternalProjectPreview.class, GradleBuild.class);
-      projectImportAction.addExtraProjectModelClasses(previewLightWeightToolingModels);
+      projectImportAction.addProjectImportExtraModelProvider(new ClassSetProjectImportExtraModelProvider(previewLightWeightToolingModels));
     }
     if(resolverCtx.isResolveModulePerSourceSet()) {
       executionSettings.withArgument("-Didea.resolveSourceSetDependencies=true");
@@ -228,7 +229,8 @@ public class GradleProjectResolver implements ExternalSystemProjectResolver<Grad
       if(!resolverCtx.isPreviewMode()){
         // register classes of extra gradle project models required for extensions (e.g. com.android.builder.model.AndroidProject)
         try {
-          projectImportAction.addExtraProjectModelClasses(resolverExtension.getExtraProjectModelClasses());
+          projectImportAction.addProjectImportExtraModelProvider(
+            new ClassSetProjectImportExtraModelProvider(resolverExtension.getExtraProjectModelClasses()));
         }
         catch (Throwable t) {
           LOG.warn(t);

plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/GradleProjectResolverExtension.java

@@ -30,6 +30,7 @@ import org.gradle.tooling.model.idea.IdeaProject;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.gradle.GradleManager;
+import org.jetbrains.plugins.gradle.model.ProjectImportExtraModelProvider;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -95,6 +96,9 @@ public interface GradleProjectResolverExtension extends ParametersEnhancer {
   @NotNull
   Set<Class> getExtraProjectModelClasses();
 
+  @NotNull
+  ProjectImportExtraModelProvider getExtraModelProvider();
+
   /**
    * add paths containing these classes to classpath of gradle tooling extension
    *

plugins/gradle/tooling-extension-api/src/org/jetbrains/plugins/gradle/model/ProjectImportAction.java

@@ -28,6 +28,8 @@ import org.gradle.tooling.model.idea.IdeaModule;
 import org.gradle.tooling.model.idea.IdeaProject;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.plugins.gradle.model.ProjectImportExtraModelProvider.BuildModelConsumer;
+import org.jetbrains.plugins.gradle.model.ProjectImportExtraModelProvider.ProjectModelConsumer;
 
 import java.io.Serializable;
 import java.lang.reflect.Field;
@@ -38,7 +40,7 @@ import java.util.*;
  */
 public class ProjectImportAction implements BuildAction<ProjectImportAction.AllModels>, Serializable {
 
-  private final Set<Class> myExtraProjectModelClasses = new LinkedHashSet<Class>();
+  private final Set<ProjectImportExtraModelProvider> myExtraProjectModelProviders = new LinkedHashSet<ProjectImportExtraModelProvider>();
   private final Set<Class> myTargetTypes = new LinkedHashSet<Class>();
   private final boolean myIsPreviewMode;
   private final boolean myIsGradleProjectDirSupported;
@@ -50,12 +52,12 @@ public class ProjectImportAction implements BuildAction<ProjectImportAction.AllM
 
   public ProjectImportAction(boolean isPreviewMode, boolean isGradleProjectDirSupported, boolean isCompositeBuildsSupported) {
     myIsPreviewMode = isPreviewMode;
-    myIsGradleProjectDirSupported= isGradleProjectDirSupported;
+    myIsGradleProjectDirSupported = isGradleProjectDirSupported;
     myIsCompositeBuildsSupported = isCompositeBuildsSupported;
   }
 
-  public void addExtraProjectModelClasses(@NotNull Set<Class> projectModelClasses) {
-    myExtraProjectModelClasses.addAll(projectModelClasses);
+  public void addProjectImportExtraModelProvider(@NotNull ProjectImportExtraModelProvider provider) {
+    myExtraProjectModelProviders.add(provider);
   }
 
   public void addTargetTypes(@NotNull Set<Class> targetTypes) {
@@ -83,10 +85,11 @@ public class ProjectImportAction implements BuildAction<ProjectImportAction.AllM
     allModels.setBuildEnvironment(buildEnvironment);
     allModels.logPerformance("Get model BuildEnvironment", System.currentTimeMillis() - startTimeBuildEnv);
 
-    addExtraProject(controller, allModels, null);
+    addExtraProjectModels(controller, allModels, null);
     for (IdeaModule module : ideaProject.getModules()) {
-      addExtraProject(controller, allModels, module);
+      addExtraProjectModels(controller, allModels, module);
     }
+    addExtraBuildModels(controller, allModels, ideaProject);
 
     if (myIsCompositeBuildsSupported) {
       long startTimeGradleBuild = System.currentTimeMillis();
@@ -96,8 +99,9 @@ public class ProjectImportAction implements BuildAction<ProjectImportAction.AllM
         IdeaProject ideaIncludedProject = controller.findModel(build, IdeaProject.class);
         allModels.getIncludedBuilds().add(ideaIncludedProject);
         for (IdeaModule module : ideaIncludedProject.getModules()) {
-          addExtraProject(controller, allModels, module);
+          addExtraProjectModels(controller, allModels, module);
         }
+        addExtraBuildModels(controller, allModels, ideaIncludedProject);
       }
     }
 
@@ -139,26 +143,81 @@ public class ProjectImportAction implements BuildAction<ProjectImportAction.AllM
     }
   }
 
-  private void addExtraProject(@NotNull BuildController controller, @NotNull AllModels allModels, @Nullable IdeaModule model) {
-    for (Class aClass : myExtraProjectModelClasses) {
-      try {
+  private void addExtraProjectModels(@NotNull BuildController controller,
+                                     @NotNull final AllModels allModels,
+                                     @Nullable final IdeaModule model) {
+    try {
+      for (ProjectImportExtraModelProvider extension : myExtraProjectModelProviders) {
+        final Set<String> obtainedModels = new HashSet<String>();
         long startTime = System.currentTimeMillis();
-        Object extraProject = controller.findModel(model, aClass);
-        if (extraProject == null) continue;
-        allModels.addExtraProject(extraProject, aClass, model != null ? model.getGradleProject() : null);
-        allModels.logPerformance("Get model " + aClass.getName()
-                                 + (model == null ? " without target" : " for module " + model.getName()),
-                                 System.currentTimeMillis() - startTime);
+        ProjectModelConsumer modelConsumer = new ProjectModelConsumer() {
+          @Override
+          public void consume(@NotNull Object object, @NotNull Class clazz) {
+            allModels.addExtraProject(object, clazz, model != null ? model.getGradleProject() : null);
+            obtainedModels.add(clazz.getName());
+          }
+        };
+        extension.populateProjectModels(controller, model, modelConsumer);
+        allModels.logPerformance(
+          "Ran extension " +
+          extension.getClass().getName() +
+          (model == null ? " without target" : " for module " + model.getName()) +
+          " obtained " + obtainedModels.size() + " model(s): " + joinClassNamesToString(obtainedModels),
+          System.currentTimeMillis() - startTime);
       }
-      catch (Exception e) {
-        // do not fail project import in a preview mode
-        if (!myIsPreviewMode) {
-          throw new ExternalSystemException(e);
-        }
+    }
+    catch (Exception e) {
+      // do not fail project import in a preview mode
+      if (!myIsPreviewMode) {
+        throw new ExternalSystemException(e);
       }
     }
   }
 
+  private void addExtraBuildModels(@NotNull BuildController controller,
+                                   @NotNull final AllModels allModels,
+                                   @NotNull final IdeaProject project) {
+    try {
+      for (ProjectImportExtraModelProvider extension : myExtraProjectModelProviders) {
+        final Set<String> obtainedModels = new HashSet<String>();
+        long startTime = System.currentTimeMillis();
+        BuildModelConsumer modelConsumer = new BuildModelConsumer() {
+          @Override
+          public void consume(@Nullable IdeaModule module, @NotNull Object object, @NotNull Class clazz) {
+            allModels.addExtraProject(object, clazz, module != null ? module.getGradleProject() : null);
+            obtainedModels.add(clazz.getName());
+          }
+        };
+        extension.populateBuildModels(controller, project, modelConsumer);
+        allModels.logPerformance(
+          "Ran extension " +
+          extension.getClass().getName() +
+          " for project " + project.getName() +
+          " obtained " + obtainedModels.size() + " model(s): " + joinClassNamesToString(obtainedModels),
+          System.currentTimeMillis() - startTime);
+      }
+    }
+    catch (Exception e) {
+      // do not fail project import in a preview mode
+      if (!myIsPreviewMode) {
+        throw new ExternalSystemException(e);
+      }
+    }
+  }
+
+  @NotNull
+  private static String joinClassNamesToString(@NotNull Set<String> names) {
+    StringBuilder sb = new StringBuilder();
+    for (Iterator<String> it = names.iterator(); it.hasNext();) {
+      sb.append(it.next());
+      if (it.hasNext()) {
+        sb.append(", ");
+      }
+    }
+
+    return sb.toString();
+  }
+
   public static class AllModels extends ModelsHolder<IdeaProject, GradleProject> {
 
     private final List<IdeaProject> includedBuilds = new ArrayList<IdeaProject>();

plugins/gradle/tooling-extension-api/src/org/jetbrains/plugins/gradle/model/ProjectImportExtraModelProvider.java

+// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package org.jetbrains.plugins.gradle.model;
+
+
+import org.gradle.tooling.BuildController;
+import org.gradle.tooling.model.idea.IdeaModule;
+import org.gradle.tooling.model.idea.IdeaProject;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.io.Serializable;
+
+/**
+ * Allows the {@link ProjectImportAction} to be extended to allow extra flexibility to extensions when requesting the models.
+ *
+ * {@link #populateProjectModels(BuildController, IdeaModule, ProjectModelConsumer)} is called once for each {@link IdeaModule} obtained
+ * from the Gradle Tooling API (this includes modules from included builds).
+ *
+ * {@link #populateBuildModels(BuildController, IdeaProject, BuildModelConsumer)} is called once for each {@link IdeaProject} that is
+ * obtained from the Gradle Tooling API, for none-composite builds this will be called exactly once, for composite builds this will be
+ * called once for each included build and once for the name build. This will always be called after
+ * {@link #populateProjectModels(BuildController, IdeaModule, ProjectModelConsumer)}.
+ */
+public interface ProjectImportExtraModelProvider extends Serializable {
+  interface ProjectModelConsumer {
+    void consume(@NotNull Object object, @NotNull Class clazz);
+  }
+
+  interface BuildModelConsumer {
+    void consume(@Nullable IdeaModule module, @NotNull Object object, @NotNull Class  clazz);
+  }
+
+  void populateBuildModels(@NotNull BuildController controller, @NotNull IdeaProject project, @NotNull BuildModelConsumer consumer);
+
+  void populateProjectModels(@NotNull BuildController controller, @Nullable IdeaModule module, @NotNull ProjectModelConsumer modelConsumer);
+}

plugins/gradle/tooling-extension-impl/testSources/org/jetbrains/plugins/gradle/tooling/builder/AbstractModelBuilderTest.java

@@ -37,6 +37,7 @@ import org.jetbrains.plugins.gradle.model.ClasspathEntryModel;
 import org.jetbrains.plugins.gradle.model.ExternalProject;
 import org.jetbrains.plugins.gradle.model.ProjectImportAction;
 import org.jetbrains.plugins.gradle.service.execution.GradleExecutionHelper;
+import org.jetbrains.plugins.gradle.service.project.ClassSetProjectImportExtraModelProvider;
 import org.jetbrains.plugins.gradle.tooling.VersionMatcherRule;
 import org.jetbrains.plugins.gradle.util.GradleConstants;
 import org.junit.After;
@@ -155,7 +156,7 @@ public abstract class AbstractModelBuilderTest {
       boolean isCompositeBuildsSupported = isGradleProjectDirSupported && _gradleVersion.compareTo(GradleVersion.version("3.1")) >= 0;
       final ProjectImportAction projectImportAction = new ProjectImportAction(false, isGradleProjectDirSupported,
                                                                               isCompositeBuildsSupported);
-      projectImportAction.addExtraProjectModelClasses(getModels());
+      projectImportAction.addProjectImportExtraModelProvider(new ClassSetProjectImportExtraModelProvider(getModels()));
       BuildActionExecuter<ProjectImportAction.AllModels> buildActionExecutor = connection.action(projectImportAction);
       File initScript = GradleExecutionHelper.generateInitScript(false, getToolingExtensionClasses());
       assertNotNull(initScript);