[DROOLS-2906] Implement scenario runner (backend) (apache#2036)

DROOLS-2906: Implement scenario runner (backend)
kiegroup · Sep 7, 2018 · 8265d27 · 8265d27
1 parent 29b7a08
commit 8265d27
Show file tree

Hide file tree

Showing 8 changed files with 324 additions and 189 deletions.
diff --git a/drools-compiler/src/test/java/org/drools/compiler/simulation/BatchRunFluentTest.java b/drools-compiler/src/test/java/org/drools/compiler/simulation/BatchRunFluentTest.java
diff --git a/drools-core/src/main/java/org/drools/core/command/AbstractNewKieContainerCommand.java b/drools-core/src/main/java/org/drools/core/command/AbstractNewKieContainerCommand.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2010 Red Hat, Inc. and/or its affiliates.
+ *
+ * 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.core.command;
+
+import org.drools.core.command.impl.RegistryContext;
+import org.kie.api.KieServices;
+import org.kie.api.builder.ReleaseId;
+import org.kie.api.runtime.KieContainer;
+
+public class AbstractNewKieContainerCommand {
+
+    protected KieContainer getKieContainer(RegistryContext context, ReleaseId releaseId) {
+        KieContainer kieContainer;
+        if (releaseId != null) {
+            // use the new API to retrieve the session by ID
+            KieServices kieServices = KieServices.Factory.get();
+            kieContainer = kieServices.newKieContainer(releaseId);
+        } else {
+            kieContainer = context.lookup(KieContainer.class);
+            if (kieContainer == null) {
+                throw new RuntimeException("ReleaseId was not specified, nor was an existing KieContainer assigned to the Registry");
+            }
+        }
+        return kieContainer;
+    }
+}
diff --git a/drools-core/src/main/java/org/drools/core/command/GetKieContainerCommand.java b/drools-core/src/main/java/org/drools/core/command/GetKieContainerCommand.java
@@ -24,8 +24,8 @@
 import org.kie.api.runtime.KieContainer;
 
 public class GetKieContainerCommand
-    implements
-    ExecutableCommand<KieContainer> {
+        implements
+        ExecutableCommand<KieContainer> {
 
     private static final long serialVersionUID = 8748826714594402049L;
     private ReleaseId releaseId;
@@ -36,10 +36,10 @@ public GetKieContainerCommand(ReleaseId releaseId) {
 
     public KieContainer execute(Context context) {
         // use the new API to retrieve the session by ID
-        KieServices  kieServices  = KieServices.Factory.get();
+        KieServices kieServices = KieServices.Factory.get();
         KieContainer kieContainer = kieServices.newKieContainer(releaseId);
 
-        ((RegistryContext)context).register( KieContainer.class, kieContainer );
+        ((RegistryContext) context).register(KieContainer.class, kieContainer);
         return kieContainer;
     }
 
@@ -50,7 +50,7 @@ public ReleaseId getReleaseId() {
     @Override
     public String toString() {
         return "GetKieContainerCommand{" +
-               "releaseId=" + releaseId +
-               '}';
+                "releaseId=" + releaseId +
+                '}';
     }
 }
diff --git a/drools-core/src/main/java/org/drools/core/command/NewKieSessionCommand.java b/drools-core/src/main/java/org/drools/core/command/NewKieSessionCommand.java
@@ -16,21 +16,23 @@
 
 package org.drools.core.command;
 
+import java.util.function.BiFunction;
+
 import org.drools.core.command.impl.RegistryContext;
-import org.kie.api.KieServices;
 import org.kie.api.builder.ReleaseId;
 import org.kie.api.command.ExecutableCommand;
 import org.kie.api.runtime.Context;
 import org.kie.api.runtime.KieContainer;
 import org.kie.api.runtime.KieSession;
 
-public class NewKieSessionCommand
-    implements
-    ExecutableCommand<KieSession> {
+public class NewKieSessionCommand extends AbstractNewKieContainerCommand
+        implements
+        ExecutableCommand<KieSession> {
 
     private static final long serialVersionUID = 8748826714594402049L;
     private String sessionId;
     private ReleaseId releaseId;
+    private BiFunction<String, KieContainer, KieContainer> beforeSessionCreation = (a, b) -> b;
 
     public NewKieSessionCommand(String sessionId) {
         this.sessionId = sessionId;
@@ -43,31 +45,28 @@ public NewKieSessionCommand(ReleaseId releaseId, String sessionId) {
 
     @Override
     public KieSession execute(Context context) {
-        KieContainer kieContainer;
 
-        if ( releaseId != null ) {
-            // use the new API to retrieve the session by ID
-            KieServices  kieServices  = KieServices.Factory.get();
-            kieContainer = kieServices.newKieContainer(releaseId);
-        } else {
-            kieContainer = ((RegistryContext)context).lookup( KieContainer.class );
-            if ( kieContainer == null ) {
-                throw new RuntimeException("ReleaseId was not specfied, nor was an existing KieContainer assigned to the Registry");
-            }
-        }
+        KieContainer kieContainer = getKieContainer((RegistryContext) context, releaseId);
+
+        kieContainer = beforeSessionCreation.apply(sessionId, kieContainer);
 
-        KieSession ksession  = sessionId != null ? kieContainer.newKieSession(sessionId) : kieContainer.newKieSession();
+        KieSession ksession = sessionId != null ? kieContainer.newKieSession(sessionId) : kieContainer.newKieSession();
 
-        ((RegistryContext)context).register( KieSession.class, ksession );
+        ((RegistryContext) context).register(KieSession.class, ksession);
 
         return ksession;
     }
 
+    public NewKieSessionCommand setBeforeSessionCreation(BiFunction<String, KieContainer, KieContainer> beforeSessionCreation) {
+        this.beforeSessionCreation = beforeSessionCreation;
+        return this;
+    }
+
     @Override
     public String toString() {
         return "NewKieSessionCommand{" +
-               "sessionId='" + sessionId + '\'' +
-               ", releaseId=" + releaseId +
-               '}';
+                "sessionId='" + sessionId + '\'' +
+                ", releaseId=" + releaseId +
+                '}';
     }
 }
diff --git a/drools-core/src/main/java/org/drools/core/command/NewRuleUnitExecutorCommand.java b/drools-core/src/main/java/org/drools/core/command/NewRuleUnitExecutorCommand.java
@@ -17,15 +17,14 @@
 package org.drools.core.command;
 
 import org.drools.core.command.impl.RegistryContext;
-import org.kie.api.KieServices;
 import org.kie.api.builder.ReleaseId;
 import org.kie.api.command.ExecutableCommand;
 import org.kie.api.runtime.Context;
 import org.kie.api.runtime.KieContainer;
 import org.kie.api.runtime.KieSession;
 import org.kie.api.runtime.rule.RuleUnitExecutor;
 
-public class NewRuleUnitExecutorCommand
+public class NewRuleUnitExecutorCommand extends AbstractNewKieContainerCommand
         implements
         ExecutableCommand<RuleUnitExecutor> {
 
@@ -46,16 +45,7 @@ public NewRuleUnitExecutorCommand(ReleaseId releaseId, String sessionId) {
     public RuleUnitExecutor execute(Context context) {
         KieContainer kieContainer;
 
-        if (releaseId != null) {
-            // use the new API to retrieve the session by ID
-            KieServices kieServices = KieServices.Factory.get();
-            kieContainer = kieServices.newKieContainer(releaseId);
-        } else {
-            kieContainer = ((RegistryContext) context).lookup(KieContainer.class);
-            if (kieContainer == null) {
-                throw new RuntimeException("ReleaseId was not specfied, nor was an existing KieContainer assigned to the Registry");
-            }
-        }
+        kieContainer = getKieContainer((RegistryContext) context, releaseId);
 
         RuleUnitExecutor ruleUnitExecutor = sessionId != null ? kieContainer.newRuleUnitExecutor(sessionId) : kieContainer.newRuleUnitExecutor();
 

diff --git a/drools-core/src/main/java/org/drools/core/command/SetKieContainerCommand.java b/drools-core/src/main/java/org/drools/core/command/SetKieContainerCommand.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2010 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.core.command;
+
+import org.drools.core.command.impl.RegistryContext;
+import org.kie.api.command.ExecutableCommand;
+import org.kie.api.runtime.Context;
+import org.kie.api.runtime.KieContainer;
+
+public class SetKieContainerCommand
+        implements
+        ExecutableCommand<KieContainer> {
+
+    private static final long serialVersionUID = 2985535777825271597L;
+    private final KieContainer kieContainer;
+
+    public SetKieContainerCommand(KieContainer kieContainer) {
+        this.kieContainer = kieContainer;
+    }
+
+    public KieContainer execute(Context context) {
+        ((RegistryContext) context).register(KieContainer.class, kieContainer);
+        return kieContainer;
+    }
+
+    @Override
+    public String toString() {
+        return "SetKieContainerCommand{" +
+                "kieContainer=" + kieContainer +
+                '}';
+    }
+}
diff --git a/drools-core/src/main/java/org/drools/core/fluent/impl/ExecutableBuilderImpl.java b/drools-core/src/main/java/org/drools/core/fluent/impl/ExecutableBuilderImpl.java
@@ -17,23 +17,30 @@
 package org.drools.core.fluent.impl;
 
 import org.drools.core.command.GetKieContainerCommand;
+import org.drools.core.command.SetKieContainerCommand;
 import org.kie.api.builder.ReleaseId;
 import org.kie.api.runtime.Executable;
+import org.kie.api.runtime.KieContainer;
 import org.kie.api.runtime.builder.ExecutableBuilder;
 import org.kie.api.runtime.builder.KieContainerFluent;
 
 public class ExecutableBuilderImpl extends BaseBatchFluent<ExecutableBuilder, ExecutableBuilder> implements ExecutableBuilder {
 
-
     public ExecutableBuilderImpl() {
         super(new ExecutableImpl());
-        getFluentContext().setExecutableBuilder( this );
-
+        getFluentContext().setExecutableBuilder(this);
     }
 
     @Override
     public KieContainerFluent getKieContainer(ReleaseId releaseId) {
-        addCommand( new GetKieContainerCommand(releaseId) );
+        addCommand(new GetKieContainerCommand(releaseId));
+        KieContainerFluentImpl fluent = new KieContainerFluentImpl(fluentCtx);
+        return fluent;
+    }
+
+    @Override
+    public KieContainerFluent setKieContainer(KieContainer kieContainer) {
+        addCommand(new SetKieContainerCommand(kieContainer));
         KieContainerFluentImpl fluent = new KieContainerFluentImpl(fluentCtx);
         return fluent;
     }

diff --git a/drools-core/src/main/java/org/drools/core/fluent/impl/KieContainerFluentImpl.java b/drools-core/src/main/java/org/drools/core/fluent/impl/KieContainerFluentImpl.java
@@ -16,8 +16,11 @@
 
 package org.drools.core.fluent.impl;
 
+import java.util.function.BiFunction;
+
 import org.drools.core.command.NewKieSessionCommand;
 import org.drools.core.command.NewRuleUnitExecutorCommand;
+import org.kie.api.runtime.KieContainer;
 import org.kie.api.runtime.builder.ExecutableBuilder;
 import org.kie.api.runtime.builder.KieContainerFluent;
 import org.kie.api.runtime.builder.KieSessionFluent;
@@ -44,6 +47,14 @@ public KieSessionFluent newSession(String sessionId) {
         return new KieSessionFluentImpl(ctx);
     }
 
+    @Override
+    public KieSessionFluent newSessionCustomized(String sessionId, BiFunction<String, KieContainer, KieContainer> customizer) {
+        NewKieSessionCommand cmd = new NewKieSessionCommand(sessionId);
+        cmd.setBeforeSessionCreation(customizer);
+        ctx.addCommand(cmd);
+        return new KieSessionFluentImpl(ctx);
+    }
+
     @Override
     public RuleUnitExecutorFluent newRuleUnitExecutor() {
         return newRuleUnitExecutor(null);