Show and modify routing rules from the UI

Author: prakhar10

docs/gateway-api.md

@@ -91,3 +91,37 @@ Will return a JSON array of active Trino cluster backends:
 curl -X POST http://localhost:8080/gateway/backend/activate/trino-2
 ```
 
+## Update Routing Rules
+
+This API can be used to programmatically update the Routing Rules.
+Rule will be updated based on the rule name.
+
+For this feature to work with multiple replicas of the Trino Gateway, you will need to provide a shared storage that supports file locking for the routing rules file. If multiple replicas are used with local storage, then rules will get out of sync when updated.
+
+```shell
+curl -X POST http://localhost:8080/webapp/updateRoutingRules \
+ -H 'Content-Type: application/json' \
+ -d '{  "name": "trino-rule",
+        "description": "updated rule description",
+        "priority": 0,
+        "actions": ["updated action"],
+        "condition": "updated condition"
+    }'
+```
+### Disable Routing Rules UI
+
+You can set the `disablePages` config to disable pages on the UI.
+
+The following pages are available:
+- `dashboard`
+- `cluster`
+- `resource-group`
+- `selector`
+- `history`
+- `routing-rules`
+
+```yaml
+uiConfiguration:
+  disablePages:
+    - 'routing-rules'
+```

gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java

@@ -34,6 +34,7 @@
 import io.trino.gateway.ha.resource.PublicResource;
 import io.trino.gateway.ha.resource.TrinoResource;
 import io.trino.gateway.ha.router.ForRouter;
+import io.trino.gateway.ha.router.RoutingRulesManager;
 import io.trino.gateway.ha.security.AuthorizedExceptionMapper;
 import io.trino.gateway.proxyserver.ForProxy;
 import io.trino.gateway.proxyserver.ProxyRequestHandler;
@@ -143,6 +144,7 @@ public void configure(Binder binder)
         jaxrsBinder(binder).bind(AuthorizedExceptionMapper.class);
         binder.bind(ProxyHandlerStats.class).in(Scopes.SINGLETON);
         newExporter(binder).export(ProxyHandlerStats.class).withGeneratedName();
+        binder.bind(RoutingRulesManager.class);
     }
 
     private static void addManagedApps(HaGatewayConfiguration configuration, Binder binder)

gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java

@@ -46,6 +46,8 @@ public class HaGatewayConfiguration
 
     private RequestAnalyzerConfig requestAnalyzerConfig = new RequestAnalyzerConfig();
 
+    private UIConfiguration uiConfiguration = new UIConfiguration();
+
     // List of Modules with FQCN (Fully Qualified Class Name)
     private List<String> modules;
 
@@ -214,6 +216,16 @@ public void setRequestAnalyzerConfig(RequestAnalyzerConfig requestAnalyzerConfig
         this.requestAnalyzerConfig = requestAnalyzerConfig;
     }
 
+    public UIConfiguration getUiConfiguration()
+    {
+        return uiConfiguration;
+    }
+
+    public void setUiConfiguration(UIConfiguration uiConfiguration)
+    {
+        this.uiConfiguration = uiConfiguration;
+    }
+
     public List<String> getModules()
     {
         return this.modules;

gateway-ha/src/main/java/io/trino/gateway/ha/config/UIConfiguration.java

@@ -0,0 +1,34 @@
+/*
+ * 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.config;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.List;
+
+public class UIConfiguration
+{
+    private List<String> disablePages;
+
+    @JsonProperty
+    public List<String> getDisablePages()
+    {
+        return disablePages;
+    }
+
+    public void setDisablePages(List<String> disablePages)
+    {
+        this.disablePages = disablePages;
+    }
+}

gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java

@@ -0,0 +1,46 @@
+/*
+ * 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.domain;
+
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
+import static java.util.Objects.requireNonNull;
+import static java.util.Objects.requireNonNullElse;
+
+/**
+ * RoutingRules
+ *
+ * @param name name of the routing rule
+ * @param description description of the routing rule
+ * @param priority priority of the routing rule. Higher number represents higher priority. If two rules have same priority then order of execution is not guaranteed.
+ * @param actions actions of the routing rule
+ * @param condition condition of the routing rule
+ */
+public record RoutingRule(
+        String name,
+        String description,
+        Integer priority,
+        List<String> actions,
+        String condition)
+{
+    public RoutingRule {
+        requireNonNull(name, "name is null");
+        requireNonNullElse(description, "");
+        requireNonNullElse(priority, 0);
+        actions = ImmutableList.copyOf(actions);
+        requireNonNull(condition, "condition is null");
+    }
+}

gateway-ha/src/main/java/io/trino/gateway/ha/resource/GatewayWebAppResource.java

@@ -16,8 +16,11 @@
 import com.google.common.base.Strings;
 import com.google.inject.Inject;
 import io.trino.gateway.ha.clustermonitor.ClusterStats;
+import io.trino.gateway.ha.config.HaGatewayConfiguration;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
+import io.trino.gateway.ha.config.UIConfiguration;
 import io.trino.gateway.ha.domain.Result;
+import io.trino.gateway.ha.domain.RoutingRule;
 import io.trino.gateway.ha.domain.TableData;
 import io.trino.gateway.ha.domain.request.GlobalPropertyRequest;
 import io.trino.gateway.ha.domain.request.QueryDistributionRequest;
@@ -34,8 +37,10 @@
 import io.trino.gateway.ha.router.HaGatewayManager;
 import io.trino.gateway.ha.router.QueryHistoryManager;
 import io.trino.gateway.ha.router.ResourceGroupsManager;
+import io.trino.gateway.ha.router.RoutingRulesManager;
 import jakarta.annotation.security.RolesAllowed;
 import jakarta.ws.rs.Consumes;
+import jakarta.ws.rs.GET;
 import jakarta.ws.rs.POST;
 import jakarta.ws.rs.Path;
 import jakarta.ws.rs.Produces;
@@ -44,6 +49,7 @@
 import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.SecurityContext;
 
+import java.io.IOException;
 import java.time.LocalDateTime;
 import java.time.ZoneId;
 import java.time.ZoneOffset;
@@ -67,18 +73,24 @@ public class GatewayWebAppResource
     private final QueryHistoryManager queryHistoryManager;
     private final BackendStateManager backendStateManager;
     private final ResourceGroupsManager resourceGroupsManager;
+    private final UIConfiguration uiConfiguration;
+    private final RoutingRulesManager routingRulesManager;
 
     @Inject
     public GatewayWebAppResource(
             GatewayBackendManager gatewayBackendManager,
             QueryHistoryManager queryHistoryManager,
             BackendStateManager backendStateManager,
-            ResourceGroupsManager resourceGroupsManager)
+            ResourceGroupsManager resourceGroupsManager,
+            HaGatewayConfiguration configuration,
+            RoutingRulesManager routingRulesManager)
     {
         this.gatewayBackendManager = requireNonNull(gatewayBackendManager, "gatewayBackendManager is null");
         this.queryHistoryManager = requireNonNull(queryHistoryManager, "queryHistoryManager is null");
         this.backendStateManager = requireNonNull(backendStateManager, "backendStateManager is null");
         this.resourceGroupsManager = requireNonNull(resourceGroupsManager, "resourceGroupsManager is null");
+        this.uiConfiguration = configuration.getUiConfiguration();
+        this.routingRulesManager = requireNonNull(routingRulesManager, "routingRulesManager is null");
     }
 
     @POST
@@ -424,4 +436,36 @@ public Response readExactMatchSourceSelector()
         List<ResourceGroupsManager.ExactSelectorsDetail> selectorsDetailList = resourceGroupsManager.readExactMatchSourceSelector();
         return Response.ok(Result.ok(selectorsDetailList)).build();
     }
+
+    @GET
+    @RolesAllowed("USER")
+    @Produces(MediaType.APPLICATION_JSON)
+    @Path("/getRoutingRules")
+    public Response getRoutingRules()
+            throws IOException
+    {
+        List<RoutingRule> routingRulesList = routingRulesManager.getRoutingRules();
+        return Response.ok(Result.ok(routingRulesList)).build();
+    }
+
+    @POST
+    @RolesAllowed("ADMIN")
+    @Consumes(MediaType.APPLICATION_JSON)
+    @Produces(MediaType.APPLICATION_JSON)
+    @Path("/updateRoutingRules")
+    public Response updateRoutingRules(RoutingRule routingRule)
+            throws IOException
+    {
+        List<RoutingRule> routingRulesList = routingRulesManager.updateRoutingRule(routingRule);
+        return Response.ok(Result.ok(routingRulesList)).build();
+    }
+
+    @GET
+    @RolesAllowed("USER")
+    @Produces(MediaType.APPLICATION_JSON)
+    @Path("/getUIConfiguration")
+    public Response getUIConfiguration()
+    {
+        return Response.ok(Result.ok(uiConfiguration)).build();
+    }
 }

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

@@ -0,0 +1,90 @@
+/*
+ * 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.ImmutableList;
+import com.google.inject.Inject;
+import io.trino.gateway.ha.config.HaGatewayConfiguration;
+import io.trino.gateway.ha.domain.RoutingRule;
+
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileLock;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
+import java.util.List;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+public class RoutingRulesManager
+{
+    private final String rulesConfigPath;
+
+    @Inject
+    public RoutingRulesManager(HaGatewayConfiguration configuration)
+    {
+        this.rulesConfigPath = configuration.getRoutingRules().getRulesConfigPath();
+    }
+
+    public List<RoutingRule> getRoutingRules()
+            throws IOException
+    {
+        ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory());
+        ImmutableList.Builder<RoutingRule> routingRulesBuilder = ImmutableList.builder();
+        try {
+            String content = Files.readString(Paths.get(rulesConfigPath), UTF_8);
+            YAMLParser parser = new YAMLFactory().createParser(content);
+            while (parser.nextToken() != null) {
+                RoutingRule routingRule = yamlReader.readValue(parser, RoutingRule.class);
+                routingRulesBuilder.add(routingRule);
+            }
+            return routingRulesBuilder.build();
+        }
+        catch (IOException e) {
+            throw new IOException("Failed to read or parse routing rules configuration from path : " + rulesConfigPath, e);
+        }
+    }
+
+    public synchronized List<RoutingRule> updateRoutingRule(RoutingRule routingRule)
+            throws IOException
+    {
+        ImmutableList.Builder<RoutingRule> updatedRoutingRulesBuilder = ImmutableList.builder();
+        List<RoutingRule> currentRoutingRulesList = getRoutingRules();
+        try (FileChannel fileChannel = FileChannel.open(Paths.get(rulesConfigPath), StandardOpenOption.WRITE, StandardOpenOption.READ);
+                FileLock lock = fileChannel.lock()) {
+            ObjectMapper yamlWriter = new ObjectMapper(new YAMLFactory());
+            StringBuilder yamlContent = new StringBuilder();
+            for (RoutingRule rule : currentRoutingRulesList) {
+                if (rule.name().equals(routingRule.name())) {
+                    yamlContent.append(yamlWriter.writeValueAsString(routingRule));
+                    updatedRoutingRulesBuilder.add(routingRule);
+                }
+                else {
+                    yamlContent.append(yamlWriter.writeValueAsString(rule));
+                    updatedRoutingRulesBuilder.add(rule);
+                }
+            }
+            Files.writeString(Paths.get(rulesConfigPath), yamlContent.toString(), UTF_8);
+            lock.release();
+        }
+        catch (IOException e) {
+            throw new IOException("Failed to parse or update routing rules configuration form path : " + rulesConfigPath, e);
+        }
+        return updatedRoutingRulesBuilder.build();
+    }
+}