Add tests for REST API, update docs and RoutingRule class

prakhar10 · prakhar10 · commit 71d618958036 · 2025-01-28T16:00:37.000+05:30
diff --git a/docs/gateway-api.md b/docs/gateway-api.md
@@ -96,10 +96,11 @@ curl -X POST http://localhost:8080/gateway/backend/activate/trino-2
 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 for the routing rules file. If multiple replicas are used with local storage, then rules will get out of sync when updated.
+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,
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java
@@ -14,11 +14,11 @@
 package io.trino.gateway.ha.domain;
 
 import com.google.common.collect.ImmutableList;
-import jakarta.annotation.Nullable;
 
 import java.util.List;
 
 import static java.util.Objects.requireNonNull;
+import static java.util.Objects.requireNonNullElse;
 
 /**
  * RoutingRules
@@ -31,13 +31,15 @@
  */
 public record RoutingRule(
         String name,
-        @Nullable String description,
-        @Nullable Integer priority,
+        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");
     }
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRulesManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRulesManager.java
@@ -27,7 +27,6 @@
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.nio.file.StandardOpenOption;
-import java.util.ArrayList;
 import java.util.List;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -65,32 +64,23 @@ public synchronized List<RoutingRule> updateRoutingRule(RoutingRule routingRule)
             throws IOException
     {
         ImmutableList.Builder<RoutingRule> updatedRoutingRulesBuilder = ImmutableList.builder();
-        List<RoutingRule> currentRoutingRulesList = new ArrayList<>();
-        ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory());
-        try {
-            String content = Files.readString(Paths.get(rulesConfigPath), UTF_8);
-            YAMLParser parser = new YAMLFactory().createParser(content);
-            while (parser.nextToken() != null) {
-                RoutingRule currentRoutingRule = yamlReader.readValue(parser, RoutingRule.class);
-                currentRoutingRulesList.add(currentRoutingRule);
-            }
-            for (int i = 0; i < currentRoutingRulesList.size(); i++) {
-                if (currentRoutingRulesList.get(i).name().equals(routingRule.name())) {
-                    currentRoutingRulesList.set(i, routingRule);
-                    break;
-                }
-            }
+        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) {
-                yamlContent.append(yamlWriter.writeValueAsString(rule));
-                updatedRoutingRulesBuilder.add(rule);
-            }
-            try (FileChannel fileChannel = FileChannel.open(Paths.get(rulesConfigPath), StandardOpenOption.WRITE, StandardOpenOption.READ);
-                    FileLock lock = fileChannel.lock()) {
-                Files.writeString(Paths.get(rulesConfigPath), yamlContent.toString(), UTF_8);
-                lock.release();
+                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);
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java
@@ -0,0 +1,168 @@
+/*
+ * 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.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.airlift.json.JsonCodec;
+import io.trino.gateway.ha.HaGatewayLauncher;
+import io.trino.gateway.ha.HaGatewayTestUtils;
+import io.trino.gateway.ha.config.UIConfiguration;
+import io.trino.gateway.ha.domain.RoutingRule;
+import okhttp3.MediaType;
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import okhttp3.RequestBody;
+import okhttp3.Response;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.testcontainers.containers.TrinoContainer;
+
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.testcontainers.utility.MountableFile.forClasspathResource;
+
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+final class TestRoutingAPI
+{
+    private final OkHttpClient httpClient = new OkHttpClient();
+    private TrinoContainer trino;
+    int routerPort = 21001 + (int) (Math.random() * 1000);
+    int backendPort;
+
+    @BeforeAll
+    void setup()
+            throws Exception
+    {
+        trino = new TrinoContainer("trinodb/trino");
+        trino.withCopyFileToContainer(forClasspathResource("trino-config.properties"), "/etc/trino/config.properties");
+        trino.start();
+
+        backendPort = trino.getMappedPort(8080);
+
+        // seed database
+        HaGatewayTestUtils.TestConfig testConfig =
+                HaGatewayTestUtils.buildGatewayConfigAndSeedDb(routerPort, "test-config-with-routing-rules-api.yml");
+        // Start Gateway
+        String[] args = {testConfig.configFilePath()};
+        HaGatewayLauncher.main(args);
+        // Now populate the backend
+        HaGatewayTestUtils.setUpBackend(
+                "trino1", "http://localhost:" + backendPort, "externalUrl", true, "adhoc", routerPort);
+    }
+
+    @Test
+    void testGetRoutingRulesAPI()
+            throws Exception
+    {
+        Request request =
+                new Request.Builder()
+                        .url("http://localhost:" + routerPort + "/webapp/getRoutingRules")
+                        .get()
+                        .build();
+        Response response = httpClient.newCall(request).execute();
+
+        String responseBody = response.body().string();
+        ObjectMapper objectMapper = new ObjectMapper();
+        JsonNode rootNode = objectMapper.readTree(responseBody);
+        JsonNode dataNode = rootNode.path("data");
+
+        JsonCodec<RoutingRule[]> responseCodec = JsonCodec.jsonCodec(RoutingRule[].class);
+        RoutingRule[] routingRules = responseCodec.fromJson(dataNode.toString());
+
+        assertThat(response.code()).isEqualTo(200);
+        assertThat(routingRules[0].name()).isEqualTo("airflow");
+        assertThat(routingRules[0].description()).isEqualTo("if query from airflow, route to etl group");
+        assertThat(routingRules[0].priority()).isEqualTo(0);
+        assertThat(routingRules[0].condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"airflow\"");
+        assertThat(routingRules[0].actions()).first().isEqualTo("result.put(\"routingGroup\", \"etl\")");
+    }
+
+    @Test
+    void testUpdateRoutingRulesAPI()
+            throws Exception
+    {
+        //Update routing rules with a new rule
+        RoutingRule updatedRoutingRules = new RoutingRule("airflow", "if query from airflow, route to adhoc group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\"");
+        ObjectMapper objectMapper = new ObjectMapper();
+        RequestBody requestBody = RequestBody.create(objectMapper.writeValueAsString(updatedRoutingRules), MediaType.parse("application/json; charset=utf-8"));
+        Request request = new Request.Builder()
+                        .url("http://localhost:" + routerPort + "/webapp/updateRoutingRules")
+                        .addHeader("Content-Type", "application/json")
+                        .post(requestBody)
+                        .build();
+        Response response = httpClient.newCall(request).execute();
+
+        assertThat(response.code()).isEqualTo(200);
+
+        //Fetch the routing rules to see if the update was successful
+        Request request2 =
+                new Request.Builder()
+                        .url("http://localhost:" + routerPort + "/webapp/getRoutingRules")
+                        .get()
+                        .build();
+        Response response2 = httpClient.newCall(request2).execute();
+
+        String responseBody = response2.body().string();
+        ObjectMapper objectMapper2 = new ObjectMapper();
+        JsonNode rootNode = objectMapper2.readTree(responseBody);
+        JsonNode dataNode = rootNode.path("data");
+
+        JsonCodec<RoutingRule[]> responseCodec = JsonCodec.jsonCodec(RoutingRule[].class);
+        RoutingRule[] routingRules = responseCodec.fromJson(dataNode.toString());
+
+        assertThat(response.code()).isEqualTo(200);
+        assertThat(routingRules[0].name()).isEqualTo("airflow");
+        assertThat(routingRules[0].description()).isEqualTo("if query from airflow, route to adhoc group");
+        assertThat(routingRules[0].priority()).isEqualTo(0);
+        assertThat(routingRules[0].condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"JDBC\"");
+        assertThat(routingRules[0].actions()).first().isEqualTo("result.put(\"routingGroup\", \"adhoc\")");
+
+        //Revert back to old routing rules to avoid any test failures
+        RoutingRule revertRoutingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\"");
+        ObjectMapper objectMapper3 = new ObjectMapper();
+        RequestBody requestBody3 = RequestBody.create(objectMapper3.writeValueAsString(revertRoutingRules), MediaType.parse("application/json; charset=utf-8"));
+        Request request3 = new Request.Builder()
+                .url("http://localhost:" + routerPort + "/webapp/updateRoutingRules")
+                .addHeader("Content-Type", "application/json")
+                .post(requestBody3)
+                .build();
+        httpClient.newCall(request3).execute();
+    }
+
+    @Test
+    void testUIConfigurationAPI()
+            throws Exception
+    {
+        Request request = new Request.Builder()
+                .url("http://localhost:" + routerPort + "/webapp/getUIConfiguration")
+                .get()
+                .build();
+
+        Response response = httpClient.newCall(request).execute();
+        String responseBody = response.body().string();
+
+        ObjectMapper objectMapper2 = new ObjectMapper();
+        JsonNode rootNode = objectMapper2.readTree(responseBody);
+        JsonNode dataNode = rootNode.path("data");
+
+        ObjectMapper objectMapper = new ObjectMapper();
+        UIConfiguration uiConfiguration = objectMapper.readValue(dataNode.toString(), UIConfiguration.class);
+
+        assertThat(response.code()).isEqualTo(200);
+        assertThat(uiConfiguration.getDisablePages()).contains("routing-rules");
+    }
+}
diff --git a/gateway-ha/src/test/resources/rules/routing_rules_update.yml b/gateway-ha/src/test/resources/rules/routing_rules_update.yml
@@ -5,3 +5,11 @@ priority: 0
 actions:
 - "result.put(\"routingGroup\", \"etl\")"
 condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\""
+---
+name: "airflow special"
+description: "if query from airflow with special label, route to etl-special group"
+priority: 1
+actions:
+- "result.put(\"routingGroup\", \"etl-special\")"
+condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"\
+  X-Trino-Client-Tags\") contains \"label=special\""
diff --git a/gateway-ha/src/test/resources/test-config-with-routing-rules-api.yml b/gateway-ha/src/test/resources/test-config-with-routing-rules-api.yml
@@ -0,0 +1,46 @@
+serverConfig:
+  node.environment: test
+  http-server.http.port: REQUEST_ROUTER_PORT
+
+dataStore:
+  jdbcUrl: jdbc:h2:DB_FILE_PATH
+  user: sa
+  password: sa
+  driver: org.h2.Driver
+
+modules:
+  - io.trino.gateway.ha.module.HaGatewayProviderModule
+  - io.trino.gateway.ha.module.ClusterStateListenerModule
+  - io.trino.gateway.ha.module.ClusterStatsMonitorModule
+
+managedApps:
+  - io.trino.gateway.ha.clustermonitor.ActiveClusterMonitor
+
+clusterStatsConfiguration:
+  monitorType: INFO_API
+
+monitor:
+  taskDelaySeconds: 1
+
+extraWhitelistPaths:
+  - '/v1/custom.*'
+  - '/custom/logout.*'
+
+gatewayCookieConfiguration:
+  enabled: true
+  cookieSigningSecret: "kjlhbfrewbyuo452cds3dc1234ancdsjh"
+
+oauth2GatewayCookieConfiguration:
+  deletePaths:
+    - "/custom/logout"
+
+requestAnalyzerConfig:
+  analyzeRequest: true
+
+uiConfiguration:
+  disablePages:
+    - 'routing-rules'
+
+routingRules:
+  rulesEngineEnabled: true
+  rulesConfigPath: "RESOURCES_DIR/rules/routing_rules_update.yml"
