From 1dced124b28166ab73458a6714ba0961adfc4b95 Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Thu, 8 Aug 2024 20:48:50 -0500 Subject: [PATCH] Show and modify routing rules from the UI --- docs/gateway-api.md | 34 ++++ .../io/trino/gateway/baseapp/BaseApp.java | 2 + .../ha/config/HaGatewayConfiguration.java | 12 ++ .../gateway/ha/config/UIConfiguration.java | 34 ++++ .../trino/gateway/ha/domain/RoutingRule.java | 46 +++++ .../ha/resource/GatewayWebAppResource.java | 46 ++++- .../ha/router/RoutingRulesManager.java | 90 +++++++++ .../gateway/ha/router/TestRoutingAPI.java | 173 +++++++++++++++++ .../ha/router/TestRoutingRulesManager.java | 149 +++++++++++++++ .../rules/routing_rules_concurrent.yml | 7 + .../resources/rules/routing_rules_update.yml | 15 ++ .../test-config-with-routing-rules-api.yml | 38 ++++ webapp/src/api/webapp/login.ts | 4 + webapp/src/api/webapp/routing-rules.ts | 11 ++ webapp/src/components/layout.tsx | 49 +++-- webapp/src/components/routing-rules.tsx | 175 ++++++++++++++++++ webapp/src/locales/en_US.ts | 1 + webapp/src/router.tsx | 36 ++-- webapp/src/types/routing-rules.d.ts | 7 + 19 files changed, 905 insertions(+), 24 deletions(-) create mode 100644 gateway-ha/src/main/java/io/trino/gateway/ha/config/UIConfiguration.java create mode 100644 gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java create mode 100644 gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRulesManager.java create mode 100644 gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java create mode 100644 gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java create mode 100644 gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml create mode 100644 gateway-ha/src/test/resources/rules/routing_rules_update.yml create mode 100644 gateway-ha/src/test/resources/test-config-with-routing-rules-api.yml create mode 100644 webapp/src/api/webapp/routing-rules.ts create mode 100644 webapp/src/components/routing-rules.tsx create mode 100644 webapp/src/types/routing-rules.d.ts diff --git a/docs/gateway-api.md b/docs/gateway-api.md index 87f30575a..0b6303375 100644 --- a/docs/gateway-api.md +++ b/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' +``` diff --git a/gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java b/gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java index 391e284dd..ef9106d04 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java +++ b/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) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java index 479abcf96..61a859ce7 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java +++ b/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 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 getModules() { return this.modules; diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/UIConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/UIConfiguration.java new file mode 100644 index 000000000..49dc28675 --- /dev/null +++ b/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 disablePages; + + @JsonProperty + public List getDisablePages() + { + return disablePages; + } + + public void setDisablePages(List disablePages) + { + this.disablePages = disablePages; + } +} 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 new file mode 100644 index 000000000..95e43f7b5 --- /dev/null +++ b/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 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/resource/GatewayWebAppResource.java b/gateway-ha/src/main/java/io/trino/gateway/ha/resource/GatewayWebAppResource.java index a117daa55..15e5f9689 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/resource/GatewayWebAppResource.java +++ b/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 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 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 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(); + } } 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 new file mode 100644 index 000000000..d26abdb8f --- /dev/null +++ b/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.Path; +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 getRoutingRules() + throws IOException + { + ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory()); + ImmutableList.Builder routingRulesBuilder = ImmutableList.builder(); + try { + String content = Files.readString(Path.of(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 updateRoutingRule(RoutingRule routingRule) + throws IOException + { + ImmutableList.Builder updatedRoutingRulesBuilder = ImmutableList.builder(); + List currentRoutingRulesList = getRoutingRules(); + try (FileChannel fileChannel = FileChannel.open(Path.of(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(Path.of(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(); + } +} 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 new file mode 100644 index 000000000..503ea5005 --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java @@ -0,0 +1,173 @@ +/* + * 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.PostgreSQLContainer; +import org.testcontainers.containers.TrinoContainer; + +import java.io.File; +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; + private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:16"); + 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); + + postgresql.start(); + + // seed database + File testConfigFile = + HaGatewayTestUtils.buildGatewayConfig(postgresql, routerPort, "test-config-with-routing-rules-api.yml"); + // Start Gateway + String[] args = {testConfigFile.getAbsolutePath()}; + 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 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 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/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java new file mode 100644 index 000000000..00400c11e --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java @@ -0,0 +1,149 @@ +/* + * 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 io.trino.gateway.ha.config.HaGatewayConfiguration; +import io.trino.gateway.ha.config.RoutingRulesConfiguration; +import io.trino.gateway.ha.domain.RoutingRule; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.NoSuchFileException; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +final class TestRoutingRulesManager +{ + @Test + void testGetRoutingRules() + throws IOException + { + HaGatewayConfiguration configuration = new HaGatewayConfiguration(); + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + String rulesConfigPath = "src/test/resources/rules/routing_rules_atomic.yml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + configuration.setRoutingRules(routingRulesConfiguration); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); + + List result = routingRulesManager.getRoutingRules(); + + assertThat(result).hasSize(2); + assertThat(result.getFirst()).isEqualTo( + new RoutingRule( + "airflow", + "if query from airflow, route to etl group", + null, + List.of("result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")"), + "request.getHeader(\"X-Trino-Source\") == \"airflow\" && (request.getHeader(\"X-Trino-Client-Tags\") == null || request.getHeader(\"X-Trino-Client-Tags\").isEmpty())")); + } + + @Test + void testRoutingRulesNoSuchFileException() + { + HaGatewayConfiguration configuration = new HaGatewayConfiguration(); + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + String rulesConfigPath = "src/test/resources/rules/routing_rules_test.yaml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + configuration.setRoutingRules(routingRulesConfiguration); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); + + assertThatThrownBy(routingRulesManager::getRoutingRules).hasRootCauseInstanceOf(NoSuchFileException.class); + } + + @Test + void testUpdateRoutingRulesFile() + throws IOException + { + HaGatewayConfiguration configuration = new HaGatewayConfiguration(); + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + String rulesConfigPath = "src/test/resources/rules/routing_rules_update.yml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + configuration.setRoutingRules(routingRulesConfiguration); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); + + RoutingRule routingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + + List updatedRoutingRules = routingRulesManager.updateRoutingRule(routingRules); + assertThat(updatedRoutingRules.getFirst().actions().getFirst()).isEqualTo("result.put(\"routingGroup\", \"adhoc\")"); + assertThat(updatedRoutingRules.getFirst().condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + + RoutingRule originalRoutingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); + List updateRoutingRules = routingRulesManager.updateRoutingRule(originalRoutingRules); + + assertThat(updateRoutingRules.getFirst().actions().getFirst()).isEqualTo("result.put(\"routingGroup\", \"etl\")"); + assertThat(updateRoutingRules.getFirst().condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"airflow\""); + } + + @Test + void testUpdateRoutingRulesNoSuchFileException() + { + HaGatewayConfiguration configuration = new HaGatewayConfiguration(); + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + String rulesConfigPath = "src/test/resources/rules/routing_rules_updated.yaml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + configuration.setRoutingRules(routingRulesConfiguration); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); + RoutingRule routingRules = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + + assertThatThrownBy(() -> routingRulesManager.updateRoutingRule(routingRules)).hasRootCauseInstanceOf(NoSuchFileException.class); + } + + @Test + void testConcurrentUpdateRoutingRule() + throws IOException + { + HaGatewayConfiguration configuration = new HaGatewayConfiguration(); + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + String rulesConfigPath = "src/test/resources/rules/routing_rules_concurrent.yml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + configuration.setRoutingRules(routingRulesConfiguration); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); + + RoutingRule routingRule1 = new RoutingRule("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); + RoutingRule routingRule2 = new RoutingRule("airflow", "if query from airflow, route to adhoc group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"datagrip\""); + + ExecutorService executorService = Executors.newFixedThreadPool(2); + + executorService.submit(() -> + { + try { + routingRulesManager.updateRoutingRule(routingRule1); + } + catch (IOException e) { + throw new RuntimeException(e); + } + }); + + executorService.submit(() -> + { + try { + routingRulesManager.updateRoutingRule(routingRule2); + } + catch (IOException e) { + throw new RuntimeException(e); + } + }); + + executorService.shutdown(); + List updatedRoutingRules = routingRulesManager.getRoutingRules(); + assertThat(updatedRoutingRules).hasSize(1); + assertThat(updatedRoutingRules.getFirst().condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"datagrip\""); + assertThat(updatedRoutingRules.getFirst().actions().getFirst()).isEqualTo("result.put(\"routingGroup\", \"adhoc\")"); + assertThat(updatedRoutingRules.getFirst().description()).isEqualTo("if query from airflow, route to adhoc group"); + } +} diff --git a/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml b/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml new file mode 100644 index 000000000..f492dfeaa --- /dev/null +++ b/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml @@ -0,0 +1,7 @@ +--- +name: "airflow" +description: "if query from airflow, route to adhoc group" +priority: 0 +actions: +- "result.put(\"routingGroup\", \"adhoc\")" +condition: "request.getHeader(\"X-Trino-Source\") == \"datagrip\"" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_update.yml b/gateway-ha/src/test/resources/rules/routing_rules_update.yml new file mode 100644 index 000000000..3478bc45e --- /dev/null +++ b/gateway-ha/src/test/resources/rules/routing_rules_update.yml @@ -0,0 +1,15 @@ +--- +name: "airflow" +description: "if query from airflow, route to etl group" +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 new file mode 100644 index 000000000..972313ba6 --- /dev/null +++ b/gateway-ha/src/test/resources/test-config-with-routing-rules-api.yml @@ -0,0 +1,38 @@ +serverConfig: + node.environment: test + http-server.http.port: REQUEST_ROUTER_PORT + +dataStore: + jdbcUrl: POSTGRESQL_JDBC_URL + user: POSTGRESQL_USER + password: POSTGRESQL_PASSWORD + driver: org.postgresql.Driver + +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" diff --git a/webapp/src/api/webapp/login.ts b/webapp/src/api/webapp/login.ts index 392114f46..da5396021 100644 --- a/webapp/src/api/webapp/login.ts +++ b/webapp/src/api/webapp/login.ts @@ -19,3 +19,7 @@ export async function getInfoApi(): Promise { export async function loginTypeApi(): Promise { return api.post('/loginType', {}) } + +export async function getUIConfiguration(): Promise { + return api.get('/webapp/getUIConfiguration') +} diff --git a/webapp/src/api/webapp/routing-rules.ts b/webapp/src/api/webapp/routing-rules.ts new file mode 100644 index 000000000..65e364c11 --- /dev/null +++ b/webapp/src/api/webapp/routing-rules.ts @@ -0,0 +1,11 @@ +import {api} from "../base"; +import {RoutingRulesData} from "../../types/routing-rules"; + +export async function routingRulesApi(): Promise { + const response = await api.get('/webapp/getRoutingRules'); + return response; +} + +export async function updateRoutingRulesApi(body: Record): Promise { + return api.post('/webapp/updateRoutingRules', body) +} diff --git a/webapp/src/components/layout.tsx b/webapp/src/components/layout.tsx index c75341b62..3b54e0b4e 100644 --- a/webapp/src/components/layout.tsx +++ b/webapp/src/components/layout.tsx @@ -1,11 +1,11 @@ import { Nav, Avatar, Layout, Dropdown, Button, Toast, Modal, Tag } from '@douyinfe/semi-ui'; -import { IconDoubleChevronRight, IconDoubleChevronLeft, IconMoon, IconSun, IconMark, IconIdCard } from '@douyinfe/semi-icons'; +import { IconDoubleChevronRight, IconDoubleChevronLeft, IconMoon, IconSun, IconMark, IconIdCard, IconUserSetting, IconUser } from '@douyinfe/semi-icons'; import styles from './layout.module.scss'; import { useEffect, useState } from 'react'; import { Link, useLocation } from "react-router-dom"; import { hasPagePermission, routers, routersMapper } from '../router'; import { Theme, useAccessStore, useConfigStore } from '../store'; -import { logoutApi } from '../api/webapp/login'; +import { getUIConfiguration, logoutApi } from '../api/webapp/login'; import Locale from "../locales"; export const RootLayout = (props: { @@ -18,6 +18,28 @@ export const RootLayout = (props: { const [collapsed, setCollapsed] = useState(false); const [selectedKey, setSelectedKey] = useState(location.pathname.substring(location.pathname.lastIndexOf('/') + 1)); const [userProfile, setUserProfile] = useState(false); + const [disabledPages, setDisabledPages] = useState(['']); + const [filteredRouters, setFilteredRouters] = useState(routers); + + useEffect(() => { + getUIConfiguration().then((res) => { + if (Object.keys(res).length == 0) { + setDisabledPages(res) + } else { + setDisabledPages(res.disablePages) + } + }) + }, []); + + useEffect(() => { + const routerFilters = disabledPages.length > 0 ? + routers + .filter(router => router.itemKey && !disabledPages.includes(router.itemKey)) + .filter(router => hasPagePermission(router, access)) + : routers + .filter(router => hasPagePermission(router, access)) + setFilteredRouters(routerFilters); + }, [disabledPages, access]); useEffect(() => { const router = routersMapper[location.pathname]; @@ -83,14 +105,19 @@ export const RootLayout = (props: { } > - - {access.nickName} - + {access.roles.includes('ADMIN') ? ( + + ) : ( + + )} } @@ -125,7 +152,7 @@ export const RootLayout = (props: { return itemElement } }} - items={routers.filter(router => hasPagePermission(router, access))} + items={filteredRouters} > {collapsed ? ( diff --git a/webapp/src/components/routing-rules.tsx b/webapp/src/components/routing-rules.tsx new file mode 100644 index 000000000..fd66bb2f1 --- /dev/null +++ b/webapp/src/components/routing-rules.tsx @@ -0,0 +1,175 @@ +import {useEffect, useState} from "react"; +import {routingRulesApi, updateRoutingRulesApi} from "../api/webapp/routing-rules.ts"; +import {RoutingRulesData} from "../types/routing-rules"; +import {Button, Card, Form, Toast} from "@douyinfe/semi-ui"; +import {FormApi} from "@douyinfe/semi-ui/lib/es/form"; +import {Role, useAccessStore} from "../store"; + +export function RoutingRules() { + const [rules, setRules] = useState([]); + const [editingStates, setEditingStates] = useState([]); + const [formApis, setFormApis] = useState<(FormApi | null)[]>([]); + const access = useAccessStore(); + + useEffect(() => { + fetchRoutingRules(); + }, []); + + const fetchRoutingRules = () => { + routingRulesApi() + .then(data => { + setRules(data); + setEditingStates(new Array(data.length).fill(false)); + setFormApis(new Array(data.length).fill(null)); + }).catch(() => { + Toast.error("Failed to fetch routing rules"); + }); + }; + + const handleEdit = (index: number) => { + setEditingStates(prev => { + const newStates = [...prev]; + newStates[index] = true; + return newStates; + }); + }; + + const handleSave = async (index: number) => { + const formApi = formApis[index]; + if (formApi) { + try { + const values = formApi.getValues(); + const actionsArray = Array.isArray(values.actions) + ? values.actions.map((action: string) => action.trim()) + : [values.actions.trim()]; + + const updatedRule: RoutingRulesData = { + ...rules[index], + ...values, + actions: actionsArray + }; + + await updateRoutingRulesApi(updatedRule); + + setEditingStates(prev => { + const newStates = [...prev]; + newStates[index] = false; + return newStates; + }); + + setRules(prev => { + const newRules = [...prev]; + newRules[index] = updatedRule; + return newRules; + }); + + Toast.success("Routing rule updated successfully"); + } catch (error) { + Toast.error("Failed to update routing rule"); + } + } + }; + + const setFormApiForIndex = (index: number) => (api: FormApi) => { + setFormApis(prev => { + const newApis = [...prev]; + newApis[index] = api; + return newApis; + }); + }; + + return ( +
+ {rules.map((rule, index) => ( +
+ handleEdit(index)}>Edit + )) + } + footerStyle={{ + display: 'flex', + justifyContent: 'flex-end', + ...(editingStates[index] ? {} : { display: 'none' }) + }} + footer={ + (access.hasRole(Role.ADMIN) && ( + + )) + } + > +
+ + + + + + +
+
+ ))} +
+ ); +} diff --git a/webapp/src/locales/en_US.ts b/webapp/src/locales/en_US.ts index c0e1f00d8..6d6aeca50 100644 --- a/webapp/src/locales/en_US.ts +++ b/webapp/src/locales/en_US.ts @@ -32,6 +32,7 @@ const en_US = { History: "History", ResourceGroup: "Resource Group", Selector: "Selector", + RoutingRules: "Routing Rules" } }, Auth: { diff --git a/webapp/src/router.tsx b/webapp/src/router.tsx index ab72bfd04..35784d20c 100644 --- a/webapp/src/router.tsx +++ b/webapp/src/router.tsx @@ -1,14 +1,16 @@ -import { IconHeart, IconIntro, IconPopover, IconScrollList, IconToast } from "@douyinfe/semi-icons-lab"; -import { NavItemProps, NavItemPropsWithItems, SubNavProps } from "@douyinfe/semi-ui/lib/es/navigation"; +import {IconIntro, IconPopover, IconScrollList, IconTree} from "@douyinfe/semi-icons-lab"; +import {NavItemProps, NavItemPropsWithItems, SubNavProps} from "@douyinfe/semi-ui/lib/es/navigation"; import styles from './components/layout.module.scss'; -import { RouteProps } from "react-router-dom"; +import {RouteProps} from "react-router-dom"; import Locale from "./locales"; -import { Dashboard } from './components/dashboard'; -import { Cluster } from './components/cluster'; -import { History } from './components/history'; -import { Selector } from "./components/selector"; -import { ResourceGroup } from "./components/resource-group"; -import { AccessControlStore, Role } from "./store"; +import {Dashboard} from './components/dashboard'; +import {Cluster} from './components/cluster'; +import {History} from './components/history'; +import {Selector} from "./components/selector"; +import {ResourceGroup} from "./components/resource-group"; +import {AccessControlStore, Role} from "./store"; +import {IconHistory, IconList} from "@douyinfe/semi-icons"; +import {RoutingRules} from "./components/routing-rules.tsx"; export interface SubItemItem extends NavItemPropsWithItems { routeProps: RouteProps, @@ -46,7 +48,7 @@ export const routers: RouterItems = [ { itemKey: 'cluster', text: Locale.Menu.Sider.Cluster, - icon: , + icon: , roles: [], routeProps: { path: '/cluster', @@ -76,13 +78,23 @@ export const routers: RouterItems = [ { itemKey: 'history', text: Locale.Menu.Sider.History, - icon: , + icon: , roles: [], routeProps: { path: '/history', element: < History /> }, - } + }, + { + itemKey: 'routing-rules', + text: Locale.Menu.Sider.RoutingRules, + icon: , + roles: [], + routeProps: { + path: '/routing-rules', + element: < RoutingRules /> + }, + } ] export const routersMapper: Record = routers.reduce((mapper, item) => { diff --git a/webapp/src/types/routing-rules.d.ts b/webapp/src/types/routing-rules.d.ts new file mode 100644 index 000000000..a6af8aa21 --- /dev/null +++ b/webapp/src/types/routing-rules.d.ts @@ -0,0 +1,7 @@ +export interface RoutingRulesData { + name: string; + description: string; + priority: number; + actions: string[]; + condition: string; +}