support multi KieBases build

drools-model/drools-model-compiler/src/main/java/org/drools/modelcompiler/CanonicalKieModule.java

@@ -18,6 +18,7 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Comparator;
@@ -59,6 +60,8 @@ public class CanonicalKieModule extends ZipKieModule {
 
     private final Map<String, CanonicalKiePackages> pkgsInKbase = new HashMap<>();
 
+    private final Map<String, Model> models = new HashMap<>();
+
     private ProjectClassLoader moduleClassLoader;
 
     public CanonicalKieModule( ReleaseId releaseId, KieModuleModel kieProject, File file ) {
@@ -85,16 +88,16 @@ public ResultsImpl build() {
     public InternalKnowledgeBase createKieBase( KieBaseModelImpl kBaseModel, KieProject kieProject, ResultsImpl messages, KieBaseConfiguration conf ) {
         this.moduleClassLoader = (( ProjectClassLoader ) kieProject.getClassLoader());
         KieBaseConfiguration kBaseConf = getKieBaseConfiguration( kBaseModel, moduleClassLoader, conf );
-        CanonicalKiePackages kpkgs = pkgsInKbase.computeIfAbsent( kBaseModel.getName(), k -> createKiePackages(kBaseConf) );
+        CanonicalKiePackages kpkgs = pkgsInKbase.computeIfAbsent( kBaseModel.getName(), k -> createKiePackages(kBaseModel, kBaseConf) );
         return new KieBaseBuilder( kBaseModel, kBaseConf ).createKieBase(kpkgs);
     }
 
-    private CanonicalKiePackages createKiePackages( KieBaseConfiguration conf ) {
-        return new KiePackagesBuilder(conf, getModels().values(), this.moduleClassLoader).build();
+    private CanonicalKiePackages createKiePackages( KieBaseModelImpl kBaseModel, KieBaseConfiguration conf ) {
+        return new KiePackagesBuilder(conf, getModelForKBase(kBaseModel), this.moduleClassLoader).build();
     }
 
     public CanonicalKiePackages getKiePackages( KieBaseModelImpl kBaseModel ) {
-        return pkgsInKbase.computeIfAbsent( kBaseModel.getName(), k -> createKiePackages(getKnowledgeBaseConfiguration(kBaseModel, getModuleClassLoader())) );
+        return pkgsInKbase.computeIfAbsent( kBaseModel.getName(), k -> createKiePackages(kBaseModel, getKnowledgeBaseConfiguration(kBaseModel, getModuleClassLoader())) );
     }
 
     public ProjectClassLoader getModuleClassLoader() {
@@ -107,17 +110,34 @@ public ProjectClassLoader getModuleClassLoader() {
 
     public void setModuleClassLoader(ProjectClassLoader moduleClassLoader) {
         pkgsInKbase.clear();
+        models.clear();
         this.moduleClassLoader = moduleClassLoader;
     }
 
     private Map<String, Model> getModels() {
-        if (ruleClassesNames == null) {
-            ruleClassesNames = findRuleClassesNames( getModuleClassLoader() );
+        if ( models.isEmpty() ) {
+            if ( ruleClassesNames == null ) {
+                ruleClassesNames = findRuleClassesNames( getModuleClassLoader() );
+            }
+            for (String rulesFile : ruleClassesNames) {
+                Model model = createInstance( getModuleClassLoader(), rulesFile );
+                models.put( model.getName(), model );
+            }
+        }
+        return models;
+    }
+
+    private Collection<Model> getModelForKBase(KieBaseModelImpl kBaseModel) {
+        Map<String, Model> modelsMap = getModels();
+        if (kBaseModel.getPackages().isEmpty()) {
+            return modelsMap.values();
         }
-        Map<String, Model> models = new HashMap<>();
-        for (String rulesFile : ruleClassesNames) {
-            Model model = createInstance( getModuleClassLoader(), rulesFile );
-            models.put( model.getName(), model );
+        Collection<Model> models = new ArrayList<>();
+        for (String pkg : kBaseModel.getPackages()) {
+            Model model = modelsMap.get(pkg);
+            if ( model != null ) {
+                models.add( model );
+            }
         }
         return models;
     }

drools-model/drools-model-compiler/src/main/java/org/drools/modelcompiler/builder/CanonicalModelKieProject.java

@@ -16,7 +16,9 @@
 
 package org.drools.modelcompiler.builder;
 
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 
 import org.drools.compiler.commons.jci.compilers.CompilationResult;
 import org.drools.compiler.compiler.io.memory.MemoryFileSystem;
@@ -30,30 +32,35 @@
 
 public class CanonicalModelKieProject extends KieModuleKieProject {
 
-    private ModelBuilderImpl modelBuilder;
+    protected List<ModelBuilderImpl> modelBuilders = new ArrayList<>();
 
-    public CanonicalModelKieProject( InternalKieModule kieModule, ClassLoader classLoader ) {
-        super( kieModule, classLoader );
+    public CanonicalModelKieProject(InternalKieModule kieModule, ClassLoader classLoader) {
+        super(kieModule, classLoader);
     }
 
     @Override
-    protected KnowledgeBuilder createKnowledgeBuilder( KieBaseModelImpl kBaseModel, AbstractKieModule kModule ) {
-        modelBuilder = new ModelBuilderImpl( getBuilderConfiguration( kBaseModel, kModule ) );
+    protected KnowledgeBuilder createKnowledgeBuilder(KieBaseModelImpl kBaseModel, AbstractKieModule kModule) {
+        ModelBuilderImpl modelBuilder = new ModelBuilderImpl();
+        modelBuilders.add(modelBuilder);
         return modelBuilder;
     }
 
     @Override
     public void writeProjectOutput( MemoryFileSystem trgMfs ) {
         MemoryFileSystem srcMfs = new MemoryFileSystem();
+        ModelWriter modelWriter = new ModelWriter();
+        List<String> modelFiles = new ArrayList<>();
 
-        final ModelWriter modelWriter = new ModelWriter();
-        final ModelWriter.Result result = modelWriter.writeModel(srcMfs, modelBuilder.getPackageModels());
-        CompilationResult res = getCompiler().compile(result.getSources(), srcMfs, trgMfs, getClassLoader() );
+        for (ModelBuilderImpl modelBuilder : modelBuilders) {
+            final ModelWriter.Result result = modelWriter.writeModel( srcMfs, modelBuilder.getPackageModels() );
+            modelFiles.addAll(result.modelFiles);
+            CompilationResult res = getCompiler().compile( result.getSources(), srcMfs, trgMfs, getClassLoader() );
 
-        if ( res.getErrors().length != 0 ) {
-            throw new RuntimeException( "Compilation errors: " + Arrays.toString( res.getErrors() ) );
+            if ( res.getErrors().length != 0 ) {
+                throw new RuntimeException( "Compilation errors: " + Arrays.toString( res.getErrors() ) );
+            }
         }
 
-        modelWriter.writeModelFile(result.modelFiles, trgMfs);
+        modelWriter.writeModelFile(modelFiles, trgMfs);
     }
 }

drools-model/drools-model-compiler/src/main/java/org/drools/modelcompiler/builder/CanonicalModelMavenPluginKieProject.java

@@ -21,33 +21,20 @@
 
 import org.drools.compiler.compiler.io.Folder;
 import org.drools.compiler.compiler.io.memory.MemoryFileSystem;
-import org.drools.compiler.kie.builder.impl.AbstractKieModule;
 import org.drools.compiler.kie.builder.impl.InternalKieModule;
-import org.drools.compiler.kie.builder.impl.KieModuleKieProject;
-import org.drools.compiler.kproject.models.KieBaseModelImpl;
-import org.kie.internal.builder.KnowledgeBuilder;
 
-public class CanonicalModelMavenPluginKieProject extends KieModuleKieProject {
-
-    private List<ModelBuilderImpl> modelBuilders = new ArrayList<>();
+public class CanonicalModelMavenPluginKieProject extends CanonicalModelKieProject {
 
     public CanonicalModelMavenPluginKieProject(InternalKieModule kieModule, ClassLoader classLoader) {
         super(kieModule, classLoader);
     }
 
-    @Override
-    protected KnowledgeBuilder createKnowledgeBuilder(KieBaseModelImpl kBaseModel, AbstractKieModule kModule) {
-        ModelBuilderImpl modelBuilder = new ModelBuilderImpl();
-        modelBuilders.add(modelBuilder);
-        return modelBuilder;
-    }
-
     @Override
     public void writeProjectOutput(MemoryFileSystem trgMfs) {
         MemoryFileSystem srcMfs = new MemoryFileSystem();
-        final List<String> modelFiles = new ArrayList<>();
-        final ModelWriter modelWriter = new ModelWriter();
-        for(ModelBuilderImpl modelBuilder: modelBuilders) {
+        List<String> modelFiles = new ArrayList<>();
+        ModelWriter modelWriter = new ModelWriter();
+        for (ModelBuilderImpl modelBuilder : modelBuilders) {
             ModelWriter.Result result = modelWriter.writeModel(srcMfs, modelBuilder.getPackageModels());
             modelFiles.addAll(result.modelFiles);
             final Folder sourceFolder = srcMfs.getFolder("src/main/java");

drools-model/drools-model-compiler/src/main/java/org/drools/modelcompiler/builder/ModelWriter.java

@@ -1,6 +1,7 @@
 package org.drools.modelcompiler.builder;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -16,9 +17,9 @@
 
 public class ModelWriter {
 
-    public static final String RULES_FILE_NAME = "Rules";
+    private static final String RULES_FILE_NAME = "Rules";
 
-    public Result writeModel(MemoryFileSystem srcMfs, List<PackageModel> packageModels) {
+    public Result writeModel(MemoryFileSystem srcMfs, Collection<PackageModel> packageModels) {
         List<String> sourceFiles = new ArrayList<>();
         List<String> modelFiles = new ArrayList<>();
 

drools-model/drools-model-compiler/src/main/java/org/drools/modelcompiler/builder/generator/ModelGenerator.java

@@ -165,7 +165,6 @@ public static void generateModel(InternalKnowledgePackage pkg, PackageDescr pack
         new WindowReferenceGenerator(packageModel, pkg).addWindowReferences(packageDescr.getWindowDeclarations());
         packageModel.addAllFunctions(packageDescr.getFunctions().stream().map(FunctionGenerator::toFunction).collect(toList()));
 
-
         for(RuleDescr descr : packageDescr.getRules()) {
             if (descr instanceof QueryDescr) {
                 QueryGenerator.processQueryDef(pkg, packageModel, (QueryDescr) descr);

drools-model/drools-model-compiler/src/test/java/org/drools/modelcompiler/BaseModelTest.java

@@ -34,9 +34,7 @@
 import org.kie.api.builder.KieModule;
 import org.kie.api.builder.Message;
 import org.kie.api.builder.ReleaseId;
-import org.kie.api.builder.model.KieBaseModel;
 import org.kie.api.builder.model.KieModuleModel;
-import org.kie.api.builder.model.KieSessionModel;
 import org.kie.api.runtime.ClassObjectFilter;
 import org.kie.api.runtime.KieContainer;
 import org.kie.api.runtime.KieSession;
@@ -69,12 +67,19 @@ protected KieSession getKieSession( String... rules ) {
     }
 
     protected KieSession getKieSession(KieModuleModel model, String... stringRules) {
+        return getKieContainer( model, stringRules ).newKieSession();
+    }
+
+    protected KieContainer getKieContainer( KieModuleModel model, String... stringRules ) {
+        return getKieContainer( model, toKieFiles( stringRules ) );
+    }
+
+    protected KieContainer getKieContainer( KieModuleModel model, KieFile... stringRules ) {
         KieServices ks = KieServices.get();
         ReleaseId releaseId = ks.newReleaseId( "org.kie", "kjar-test-" + UUID.randomUUID(), "1.0" );
 
         KieBuilder kieBuilder = createKieBuilder( ks, model, releaseId, stringRules );
-
-        return getKieContainer( ks, model, releaseId, kieBuilder ).newKieSession();
+        return getKieContainer( ks, model, releaseId, kieBuilder );
     }
 
     protected KieContainer getKieContainer( KieServices ks, KieModuleModel model, ReleaseId releaseId, KieBuilder kieBuilder ) {
@@ -91,6 +96,10 @@ protected void addKieModuleFromCanonicalModel( KieServices ks, KieModuleModel mo
     }
 
     protected KieBuilder createKieBuilder( KieServices ks, KieModuleModel model, ReleaseId releaseId, String... stringRules ) {
+        return createKieBuilder( ks, model, releaseId, toKieFiles( stringRules ) );
+    }
+
+    protected KieBuilder createKieBuilder( KieServices ks, KieModuleModel model, ReleaseId releaseId, KieFile... stringRules ) {
         ks.getRepository().removeKieModule( releaseId );
 
         KieFileSystem kfs = ks.newKieFileSystem();
@@ -99,7 +108,7 @@ protected KieBuilder createKieBuilder( KieServices ks, KieModuleModel model, Rel
         }
         kfs.writePomXML( KJARUtils.getPom( releaseId ) );
         for (int i = 0; i < stringRules.length; i++) {
-            kfs.write(String.format("src/main/resources/r%d.drl", i), stringRules[i] );
+            kfs.write( stringRules[i].path, stringRules[i].content );
         }
 
         KieBuilder kieBuilder = ( testRunType == RUN_TYPE.USE_CANONICAL_MODEL ) ?
@@ -115,12 +124,34 @@ protected KieBuilder createKieBuilder( KieServices ks, KieModuleModel model, Rel
 
     protected KieModuleModel getDefaultKieModuleModel( KieServices ks ) {
         KieModuleModel kproj = ks.newKieModuleModel();
-        KieBaseModel kieBaseModel1 = kproj.newKieBaseModel( "kbase" ).setDefault( true );
-        KieSessionModel ksession1 = kieBaseModel1.newKieSessionModel( "ksession" ).setDefault( true );
+        kproj.newKieBaseModel( "kbase" ).setDefault( true ).newKieSessionModel( "ksession" ).setDefault( true );
         return kproj;
     }
 
     public static <T> List<T> getObjectsIntoList(KieSession ksession, Class<T> clazz) {
         return (List<T>) ksession.getObjects(new ClassObjectFilter(clazz)).stream().collect(Collectors.toList());
     }
+
+    public static class KieFile {
+
+        public final String path;
+        public final String content;
+
+        public KieFile( int index, String content ) {
+            this( String.format("src/main/resources/r%d.drl", index), content );
+        }
+
+        public KieFile( String path, String content ) {
+            this.path = path;
+            this.content = content;
+        }
+    }
+
+    public KieFile[] toKieFiles(String[] stringRules) {
+        KieFile[] kieFiles = new KieFile[stringRules.length];
+        for (int i = 0; i < stringRules.length; i++) {
+            kieFiles[i] = new KieFile( i, stringRules[i] );
+        }
+        return kieFiles;
+    }
 }

drools-model/drools-model-compiler/src/test/java/org/drools/modelcompiler/MultiKieBaseTest.java

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2005 JBoss Inc
+ *
+ * Licensed 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.drools.modelcompiler;
+
+import org.junit.Test;
+import org.kie.api.KieServices;
+import org.kie.api.builder.model.KieModuleModel;
+import org.kie.api.runtime.KieContainer;
+import org.kie.api.runtime.KieSession;
+
+import static org.junit.Assert.assertEquals;
+
+public class MultiKieBaseTest extends BaseModelTest {
+
+    public MultiKieBaseTest( RUN_TYPE testRunType ) {
+        super( testRunType );
+    }
+
+    @Test
+    public void testHelloWorldWithPackagesAnd2KieBases() throws Exception {
+        String drl1 = "package org.pkg1\n" +
+                "rule R11 when\n" +
+                "   $m : String( this == \"Hello World\" )\n" +
+                "then\n" +
+                "end\n" +
+                "rule R12 when\n" +
+                "   $m : String( this == \"Hi Universe\" )\n" +
+                "then\n" +
+                "end\n";
+
+        String drl2 = "package org.pkg2\n" +
+                "rule R21 when\n" +
+                "   $m : String( this == \"Hello World\" )\n" +
+                "then\n" +
+                "end\n" +
+                "rule R22 when\n" +
+                "   $m : String( this == \"Aloha Earth\" )\n" +
+                "then\n" +
+                "end\n";
+
+        KieContainer kieContainer = getKieContainer( createKieProjectWithPackagesAnd2KieBases(),
+                                                     new KieFile( "src/main/resources/org/pkg1/r1.drl", drl1 ),
+                                                     new KieFile( "src/main/resources/org/pkg2/r2.drl", drl2 ) );
+
+        KieSession ksession = kieContainer.newKieSession("KSession1");
+        ksession.insert("Hello World");
+        assertEquals( 1, ksession.fireAllRules() );
+
+        ksession = kieContainer.newKieSession("KSession1");
+        ksession.insert("Hi Universe");
+        assertEquals( 1, ksession.fireAllRules() );
+
+        ksession = kieContainer.newKieSession("KSession1");
+        ksession.insert("Aloha Earth");
+        assertEquals( 0, ksession.fireAllRules() );
+
+        ksession = kieContainer.newKieSession("KSession2");
+        ksession.insert("Hello World");
+        assertEquals( 1, ksession.fireAllRules() );
+
+        ksession = kieContainer.newKieSession("KSession2");
+        ksession.insert("Hi Universe");
+        assertEquals( 0, ksession.fireAllRules() );
+
+        ksession = kieContainer.newKieSession("KSession2");
+        ksession.insert("Aloha Earth");
+        assertEquals(1, ksession.fireAllRules());
+    }
+
+    private KieModuleModel createKieProjectWithPackagesAnd2KieBases() {
+        KieModuleModel kproj = KieServices.get().newKieModuleModel();
+
+        kproj.newKieBaseModel()
+                .addPackage("org.pkg1")
+                .newKieSessionModel("KSession1");
+
+        kproj.newKieBaseModel()
+                .addPackage("org.pkg2")
+                .newKieSessionModel("KSession2");
+
+        return kproj;
+    }
+}