Use MVEL for rule evaluation

Co-Authored-By: Prakhar Sapre <[email protected]>

Author: willmostly

docs/routing-rules.md

@@ -92,11 +92,10 @@ return a result with the following criteria:
 
 ### Configure routing rules with a file
 
-To express and fire routing rules, we use the
-[easy-rules](https://github.com/j-easy/easy-rules) engine. These rules must be
-stored in a YAML file. Rules consist of a name, description, condition, and list
+Rules consist of a name, description, condition, and list
 of actions. If the condition of a particular rule evaluates to `true`, its
-actions are fired.
+actions are fired. Rules are stored as a 
+[multi-document](https://www.yaml.info/learn/document.html) YAML file. 
 
 ```yaml
 ---
@@ -113,20 +112,37 @@ actions:
   - 'result.put("routingGroup", "etl-special")'
 ```
 
-In the condition, you can access the methods of a
-[HttpServletRequest](https://docs.oracle.com/javaee/6/api/javax/servlet/http/HttpServletRequest.html)
-object called `request`. Rules may also utilize
+Three objects are available by default. They are
+* `request`, the incoming request as an [HttpServletRequest](https://docs.oracle.com/javaee/6/api/javax/servlet/http/HttpServletRequest.html)
+* `state`, a `HashMap<String, Object>` that allows passing arbitrary state from one rule evaluation to the next
+* `result`, a `HashMap<String, String>` that is used to return the result of rule evaluation to the engine
+
+In addition to the default objects, rules may optionally utilize
 [trinoRequestUser](#trinorequestuser) and
 [trinoQueryProperties](#trinoqueryproperties)
-objects, which provide information about the user and query respectively.
+, which provide information about the user and query respectively.
 You must include an action of the form `result.put(\"routingGroup\", \"foo\")`
 to trigger routing of a request  that satisfies the condition to the specific 
 routing group. Without this action, the default adhoc group is used and the 
 whole routing rule is redundant.
 
 The condition and actions are written in [MVEL](http://mvel.documentnode.com/),
-an expression language with Java-like syntax. In most cases, you can write
-conditions and actions in Java syntax and expect it to work. There are some
+an expression language with Java-like syntax. Classes from `java.util`, data-type 
+classes from `java.lang` such as `Integer` and `String`, as well as `java.lang.Math`
+and `java.lang.StrictMath` are available in rules. Rules may not use `java.lang.System`
+and other classes that allow access the Java runtime and operating system.
+In most cases, you can write
+conditions and actions in Java syntax and expect it to work. One exception is 
+parametrized types. Exclude type parameters, for example to add a `HashSet` to the
+`state` variable, use an action such as:
+```java
+actions:
+  - |
+    state.put("triggeredRules",new HashSet())
+```
+This is equivalent to `new HashSet<Object>()`. 
+
+There are some
 MVEL-specific operators. For example, instead of doing a null-check before
 accessing the `String.contains` method like this:
 
@@ -296,8 +312,8 @@ actions:
 ```
 
 This can difficult to maintain with more rules. To have better control over the
-execution of rules, we can use rule priorities and composite rules. Overall,
-priorities, composite rules, and other constructs that MVEL support allows
+execution of rules, we can use rule priorities. Overall,
+priorities and other constructs that MVEL support allows
 you to express your routing logic.
 
 #### Rule priority
@@ -328,99 +344,52 @@ that the first rule (priority 0) is fired before the second rule (priority 1).
 Thus `routingGroup` is set to `etl` and then to `etl-special`, so the
 `routingGroup` is always `etl-special` in the end.
 
-More specific rules must be set to a lesser priority so they are evaluated last
-to set a `routingGroup`. To further control the execution of rules, for example
-to have only one rule fire, you can use composite rules.
+More specific rules must be set to a higher priority so they are evaluated last
+to set a `routingGroup`.
 
-##### Composite rules
+##### Passing State
 
-First, please refer to the [easy-rule composite rules documentation](https://github.com/j-easy/easy-rules/wiki/defining-rules#composite-rules).
-
-The preceding section covers how to control the order of rule execution using
-priorities. In addition, you can configure evaluation so that only the first
-rule matched fires (the highest priority one) and the rest is ignored. You can
-use `ActivationRuleGroup` to achieve this:
+The `state` object may be used to pass information from one rule evaluation to
+the next. This allows an author to avoid duplicating logic in multiple rules.
+Priority should be used to ensure that `state` is updated before being used 
+in downstream rules. For example, the atomic rules may be re-implemented as
 
 ```yaml
 ---
-name: "airflow rule group"
-description: "routing rules for query from airflow"
-compositeRuleType: "ActivationRuleGroup"
-composingRules:
-  - name: "airflow special"
-    description: "if query from airflow with special label, route to etl-special group"
-    priority: 0
-    condition: 'request.getHeader("X-Trino-Source") == "airflow" && request.getHeader("X-Trino-Client-Tags") contains "label=special"'
-    actions:
-      - 'result.put("routingGroup", "etl-special")'
-  - name: "airflow"
-    description: "if query from airflow, route to etl group"
-    priority: 1
-    condition: 'request.getHeader("X-Trino-Source") == "airflow"'
-    actions:
-      - 'result.put("routingGroup", "etl")'
-```
-
-Note that the priorities have switched. The more specific rule has a higher
-priority, since it should fire first. A query coming from airflow with special
-label is matched to the "airflow special" rule first, since it's higher
-priority, and the second rule is ignored. A query coming from airflow with no
-labels does not match the first rule, and is then tested and matched to the
-second rule.
-
-You can also use `ConditionalRuleGroup` and `ActivationRuleGroup` to implement
-an if/else workflow. The following logic in pseudocode:
-
-```text
-if source == "airflow":
-  if clientTags["label"] == "foo":
-    return "etl-foo"
-  else if clientTags["label"] = "bar":
-    return "etl-bar"
-  else
-    return "etl"
-```
-
-This logic can be implemented with the following rules:
+name: "initialize state"
+description: "Add a set to the state map to track rules that have evaluated to true"
+priority: 0
+condition: "true"
+actions:
+  - |
+    state.put("triggeredRules",new HashSet())
+  # MVEL does not support type parameters! HashSet<String>() would result in an error.
+---
+name: "airflow"
+description: "if query from airflow, route to etl group"
+priority: 1
+condition: |
+  request.getHeader("X-Trino-Source") == "airflow"
+actions:
+  - |
+    result.put("routingGroup", "etl")
+  - |
+    state.get("triggeredRules").add("airflow")
+---
+name: "airflow special"
+description: "if query from airflow with special label, route to etl-special group"
+priority: 2
+condition: |
+  state.get("triggeredRules").contains("airflow") && request.getHeader("X-Trino-Client-Tags") contains "label=special"
+actions:
+  - |
+    result.put("routingGroup", "etl-special")
 
-```yaml
-name: "airflow rule group"
-description: "routing rules for query from airflow"
-compositeRuleType: "ConditionalRuleGroup"
-composingRules:
-  - name: "main condition"
-    description: "source is airflow"
-    priority: 0 # rule with the highest priority acts as main condition
-    condition: 'request.getHeader("X-Trino-Source") == "airflow"'
-    actions:
-      - ""
-  - name: "airflow subrules"
-    compositeRuleType: "ActivationRuleGroup" # use ActivationRuleGroup to simulate if/else
-    composingRules:
-      - name: "label foo"
-        description: "label client tag is foo"
-        priority: 0
-        condition: 'request.getHeader("X-Trino-Client-Tags") contains "label=foo"'
-        actions:
-          - 'result.put("routingGroup", "etl-foo")'
-      - name: "label bar"
-        description: "label client tag is bar"
-        priority: 0
-        condition: 'request.getHeader("X-Trino-Client-Tags") contains "label=bar"'
-        actions:
-          - 'result.put("routingGroup", "etl-bar")'
-      - name: "airflow default"
-        description: "airflow queries default to etl"
-        condition: "true"
-        actions:
-          - 'result.put("routingGroup", "etl")'
 ```
 
 ##### If statements (MVEL Flow Control)
 
-In the preceding section you see how `ConditionalRuleGroup` and
-`ActivationRuleGroup` are used to implement an `if/else` workflow. You can
-use MVEL support for `if` statements and other flow control. The following logic
+You can use MVEL support for `if` statements and other flow control. The following logic
 in pseudocode:
 
 ```text

gateway-ha/pom.xml

@@ -21,7 +21,6 @@
         <frontend.pnpmRegistryURL>https://registry.npmmirror.com</frontend.pnpmRegistryURL>
 
         <!-- dependency versions -->
-        <dep.jeasy.version>4.1.0</dep.jeasy.version>
         <dep.mockito.version>5.14.2</dep.mockito.version>
         <dep.okhttp3.version>4.12.0</dep.okhttp3.version>
         <dep.trino.version>464</dep.trino.version>
@@ -253,21 +252,9 @@
         </dependency>
 
         <dependency>
-            <groupId>org.jeasy</groupId>
-            <artifactId>easy-rules-core</artifactId>
-            <version>${dep.jeasy.version}</version>
-        </dependency>
-
-        <dependency>
-            <groupId>org.jeasy</groupId>
-            <artifactId>easy-rules-mvel</artifactId>
-            <version>${dep.jeasy.version}</version>
-        </dependency>
-
-        <dependency>
-            <groupId>org.jeasy</groupId>
-            <artifactId>easy-rules-support</artifactId>
-            <version>${dep.jeasy.version}</version>
+            <groupId>org.mvel</groupId>
+            <artifactId>mvel2</artifactId>
+            <version>2.5.2.Final</version>
         </dependency>
 
         <dependency>
@@ -290,13 +277,6 @@
             <scope>runtime</scope>
         </dependency>
 
-        <dependency>
-            <groupId>org.mvel</groupId>
-            <artifactId>mvel2</artifactId>
-            <version>2.5.2.Final</version>
-            <scope>runtime</scope>
-        </dependency>
-
         <dependency>
             <groupId>org.postgresql</groupId>
             <artifactId>postgresql</artifactId>

gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java

@@ -0,0 +1,132 @@
+/*
+ * 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 io.trino.gateway.ha.router;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
+import com.fasterxml.jackson.dataformat.yaml.YAMLParser;
+import com.google.common.collect.ImmutableMap;
+import io.trino.gateway.ha.config.RequestAnalyzerConfig;
+import jakarta.servlet.http.HttpServletRequest;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Collections.sort;
+
+public class FileBasedRoutingGroupSelector
+        implements RoutingGroupSelector
+{
+    public static final String RESULTS_ROUTING_GROUP_KEY = "routingGroup";
+
+    private List<RoutingRule> rules;
+    final boolean analyzeRequest;
+    final boolean clientsUseV2Format;
+    final int maxBodySize;
+    final TrinoRequestUser.TrinoRequestUserProvider trinoRequestUserProvider;
+    private volatile long lastUpdatedTimeMillis;
+    Path rulesPath;
+
+    public FileBasedRoutingGroupSelector(String rulesPath, RequestAnalyzerConfig requestAnalyzerConfig)
+    {
+        analyzeRequest = requestAnalyzerConfig.isAnalyzeRequest();
+        clientsUseV2Format = requestAnalyzerConfig.isClientsUseV2Format();
+        maxBodySize = requestAnalyzerConfig.getMaxBodySize();
+        trinoRequestUserProvider = new TrinoRequestUser.TrinoRequestUserProvider(requestAnalyzerConfig);
+        this.rulesPath = Paths.get(rulesPath);
+
+        setRules(readRulesFromPath(this.rulesPath));
+    }
+
+    void setRules(List<RoutingRule> rules)
+    {
+        this.rules = new ArrayList<>(rules);
+        lastUpdatedTimeMillis = System.currentTimeMillis();
+        sort(this.rules);
+    }
+
+    @Override
+    public String findRoutingGroup(HttpServletRequest request)
+    {
+        reloadRules(lastUpdatedTimeMillis);
+        Map<String, String> result = new HashMap<>();
+        Map<String, Object> state = new HashMap<>();
+
+        Map<String, Object> data;
+        if (analyzeRequest) {
+            TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(
+                    request,
+                    clientsUseV2Format,
+                    maxBodySize);
+            TrinoRequestUser trinoRequestUser = trinoRequestUserProvider.getInstance(request);
+            data = ImmutableMap.of("request", request, "trinoQueryProperties", trinoQueryProperties, "trinoRequestUser", trinoRequestUser);
+        }
+        else {
+            data = ImmutableMap.of("request", request);
+        }
+
+        rules.forEach(rule -> {
+            if (rule.evaluateCondition(data, state)) {
+                rule.evaluateAction(result, data, state);
+            }
+        });
+        return result.get(RESULTS_ROUTING_GROUP_KEY);
+    }
+
+    void reloadRules(long lastUpdatedTimeMillis)
+    {
+        try {
+            BasicFileAttributes attr = Files.readAttributes(this.rulesPath, BasicFileAttributes.class);
+            if (attr.lastModifiedTime().toMillis() <= lastUpdatedTimeMillis) {
+                return;
+            }
+            synchronized (this) {
+                // Prevent re-entry in case another thread passes the first check while rules are being updated
+                if (attr.lastModifiedTime().toMillis() > lastUpdatedTimeMillis) {
+                    List<RoutingRule> ruleList = readRulesFromPath(this.rulesPath);
+                    setRules(ruleList);
+                }
+            }
+        }
+        catch (IOException e) {
+            throw new RuntimeException("Could not access rules file", e);
+        }
+    }
+
+    public List<RoutingRule> readRulesFromPath(Path rulesPath)
+    {
+        ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory());
+        try {
+            String content = Files.readString(rulesPath, UTF_8);
+            YAMLParser parser = new YAMLFactory().createParser(content);
+            List<RoutingRule> routingRulesList = new ArrayList<>();
+            while (parser.nextToken() != null) {
+                MVELRoutingRule routingRules = yamlReader.readValue(parser, MVELRoutingRule.class);
+                routingRulesList.add(routingRules);
+            }
+            return routingRulesList;
+        }
+        catch (IOException e) {
+            throw new RuntimeException("Failed to read or parse routing rules configuration from path: " + rulesPath, e);
+        }
+    }
+}