From d3c78559271612c0bda55137e5f7a869fd99eae7 Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Thu, 8 Aug 2024 20:48:50 -0500 Subject: [PATCH 01/14] Show and modify routing rules from the UI Remove console logs and renamed api Add api documentation and changes related to PR comments --- docs/gateway-api.md | 13 ++ .../trino/gateway/ha/domain/RoutingRules.java | 36 ++++ .../ha/resource/GatewayWebAppResource.java | 76 +++++++- webapp/src/api/webapp/routing-rules.ts | 11 ++ webapp/src/components/layout.tsx | 23 ++- 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 + 9 files changed, 356 insertions(+), 22 deletions(-) create mode 100644 gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java 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..6f8d0afd1 100644 --- a/docs/gateway-api.md +++ b/docs/gateway-api.md @@ -91,3 +91,16 @@ 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. +```shell +curl -X POST http://localhost:8080/webapp/updateRoutingRules \ + -d '{ "name": "trino-rule", + "description": "updated rule description", + "priority": 0, + "actions": ["updated action"], + "condition": "updated condition" + }' +``` diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java new file mode 100644 index 000000000..8fef950d2 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java @@ -0,0 +1,36 @@ +/* + * 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.fasterxml.jackson.annotation.JsonProperty; + +import java.util.List; + +/** + * RoutingRules + * + * @param name name of the routing rule + * @param description description of the routing rule + * @param priority priority of the routing rule + * @param actions actions of the routing rule + * @param condition condition of the routing rule + */ +public record RoutingRules( + @JsonProperty("name") String name, + @JsonProperty("description") String description, + @JsonProperty("priority") Integer priority, + @JsonProperty("actions") List actions, + @JsonProperty("condition") String condition) +{ +} 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 31aa75cfe..b8056fafd 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 @@ -13,11 +13,18 @@ */ package io.trino.gateway.ha.resource; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import com.fasterxml.jackson.dataformat.yaml.YAMLParser; import com.google.common.base.Strings; import com.google.inject.Inject; +import com.google.inject.Singleton; 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.domain.Result; +import io.trino.gateway.ha.domain.RoutingRules; import io.trino.gateway.ha.domain.TableData; import io.trino.gateway.ha.domain.request.GlobalPropertyRequest; import io.trino.gateway.ha.domain.request.QueryDistributionRequest; @@ -36,6 +43,7 @@ import io.trino.gateway.ha.router.ResourceGroupsManager; 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,11 +52,15 @@ import jakarta.ws.rs.core.Response; import jakarta.ws.rs.core.SecurityContext; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Paths; import java.time.LocalDateTime; import java.time.ZoneId; import java.time.ZoneOffset; import java.time.ZonedDateTime; import java.time.format.DateTimeFormatter; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; @@ -59,6 +71,7 @@ import static java.util.Objects.requireNonNullElse; @Path("/webapp") +@Singleton public class GatewayWebAppResource { private static final LocalDateTime START_TIME = LocalDateTime.now(); @@ -67,18 +80,21 @@ public class GatewayWebAppResource private final QueryHistoryManager queryHistoryManager; private final BackendStateManager backendStateManager; private final ResourceGroupsManager resourceGroupsManager; + private final HaGatewayConfiguration configuration; @Inject public GatewayWebAppResource( GatewayBackendManager gatewayBackendManager, QueryHistoryManager queryHistoryManager, BackendStateManager backendStateManager, - ResourceGroupsManager resourceGroupsManager) + ResourceGroupsManager resourceGroupsManager, + HaGatewayConfiguration configuration) { 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.configuration = requireNonNull(configuration, "configuration is null"); } @POST @@ -423,4 +439,62 @@ 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() + { + try { + String rulesConfigPath = configuration.getRoutingRules().getRulesConfigPath(); + YAMLFactory yamlFactory = new YAMLFactory(); + ObjectMapper yamlReader = new ObjectMapper(yamlFactory); + YAMLParser yamlParser = yamlFactory.createParser(new String(Files.readAllBytes(Paths.get(rulesConfigPath)))); + List routingRulesList = yamlReader + .readValues(yamlParser, new TypeReference() {}) + .readAll(); + return Response.ok(Result.ok(routingRulesList)).build(); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @POST + @RolesAllowed("ADMIN") + @Consumes(MediaType.APPLICATION_JSON) + @Produces(MediaType.APPLICATION_JSON) + @Path("/updateRoutingRules") + public synchronized Response updateRoutingRules(RoutingRules routingRules) + { + String rulesConfigPath = configuration.getRoutingRules().getRulesConfigPath(); + ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory()); + List routingRulesList = new ArrayList<>(); + YAMLFactory yamlFactory = new YAMLFactory(); + try { + YAMLParser yamlParser = yamlFactory.createParser(new String(Files.readAllBytes(Paths.get(rulesConfigPath)))); + routingRulesList = yamlReader + .readValues(yamlParser, new TypeReference() {}) + .readAll(); + + for (int i = 0; i < routingRulesList.size(); i++) { + if (routingRulesList.get(i).name().equals(routingRules.name())) { + routingRulesList.set(i, routingRules); + break; + } + } + + ObjectMapper yamlWriter = new ObjectMapper(new YAMLFactory()); + StringBuilder yamlContent = new StringBuilder(); + for (RoutingRules rule : routingRulesList) { + yamlContent.append(yamlWriter.writeValueAsString(rule)); + } + Files.write(Paths.get(rulesConfigPath), yamlContent.toString().getBytes()); + } + catch (IOException e) { + throw new RuntimeException(e); + } + return Response.ok(Result.ok(routingRulesList)).build(); + } } 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 57d10bd30..1420198e3 100644 --- a/webapp/src/components/layout.tsx +++ b/webapp/src/components/layout.tsx @@ -1,5 +1,5 @@ import { Nav, Avatar, Layout, Dropdown, Button, Toast, Modal, Tag } from '@douyinfe/semi-ui'; -import { IconGithubLogo, IconDoubleChevronRight, IconDoubleChevronLeft, IconMoon, IconSun, IconMark, IconIdCard } from '@douyinfe/semi-icons'; +import { IconGithubLogo, 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"; @@ -87,14 +87,19 @@ export const RootLayout = (props: { } > - - {access.nickName} - + {access.roles.includes('ADMIN') ? ( + + ) : ( + + )} } 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; +} From ba2e92f67595078f1c2f71f8d3d0e194b2061f51 Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Mon, 23 Sep 2024 22:37:10 -0500 Subject: [PATCH 02/14] Add option to disable pages and update API docs --- docs/gateway-api.md | 20 +++++++++ .../ha/config/HaGatewayConfiguration.java | 12 ++++++ .../gateway/ha/config/UIConfiguration.java | 42 +++++++++++++++++++ .../ha/resource/GatewayWebAppResource.java | 9 ++++ webapp/src/api/webapp/login.ts | 4 ++ webapp/src/components/layout.tsx | 26 +++++++++++- 6 files changed, 111 insertions(+), 2 deletions(-) create mode 100644 gateway-ha/src/main/java/io/trino/gateway/ha/config/UIConfiguration.java diff --git a/docs/gateway-api.md b/docs/gateway-api.md index 6f8d0afd1..2eee746d2 100644 --- a/docs/gateway-api.md +++ b/docs/gateway-api.md @@ -95,6 +95,9 @@ 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 well you will need to provide a shared storage for the routing rules file. + ```shell curl -X POST http://localhost:8080/webapp/updateRoutingRules \ -d '{ "name": "trino-rule", @@ -104,3 +107,20 @@ curl -X POST http://localhost:8080/webapp/updateRoutingRules \ "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/ha/config/HaGatewayConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java index 4a2134c8a..2058a650f 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 @@ -44,6 +44,18 @@ public class HaGatewayConfiguration private RequestAnalyzerConfig requestAnalyzerConfig = new RequestAnalyzerConfig(); + private UIConfiguration uiConfiguration = new UIConfiguration(); + + public UIConfiguration getUiConfiguration() + { + return uiConfiguration; + } + + public void setUiConfiguration(UIConfiguration uiConfiguration) + { + this.uiConfiguration = uiConfiguration; + } + // List of Modules with FQCN (Fully Qualified Class Name) private List 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..10357e8d2 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/config/UIConfiguration.java @@ -0,0 +1,42 @@ +/* + * 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; + } + + @Override + public String toString() + { + return "UIConfiguration{" + + "disablePages=" + disablePages + + '}'; + } +} 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 b8056fafd..515fa3aa2 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 @@ -497,4 +497,13 @@ public synchronized Response updateRoutingRules(RoutingRules routingRules) } return Response.ok(Result.ok(routingRulesList)).build(); } + + @GET + @RolesAllowed("USER") + @Produces(MediaType.APPLICATION_JSON) + @Path("/getUIConfiguration") + public Response getUIConfiguration() + { + return Response.ok(Result.ok(configuration.getUiConfiguration())).build(); + } } 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/components/layout.tsx b/webapp/src/components/layout.tsx index 1420198e3..6b25ca74a 100644 --- a/webapp/src/components/layout.tsx +++ b/webapp/src/components/layout.tsx @@ -5,7 +5,7 @@ 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]; @@ -134,7 +156,7 @@ export const RootLayout = (props: { return itemElement } }} - items={routers.filter(router => hasPagePermission(router, access))} + items={filteredRouters} > {collapsed ? ( From 034963db17612f98ddd1f59d7c57b0fe21b6baf9 Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Mon, 7 Oct 2024 21:43:43 -0500 Subject: [PATCH 03/14] Move functions to service class and add tests and other fixes --- .../ha/resource/GatewayWebAppResource.java | 66 ++++-------- .../ha/router/RoutingRulesManager.java | 82 ++++++++++++++ .../ha/router/TestRoutingRulesManager.java | 102 ++++++++++++++++++ .../resources/rules/routing_rules_update.yml | 7 ++ 4 files changed, 209 insertions(+), 48 deletions(-) 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/TestRoutingRulesManager.java create mode 100644 gateway-ha/src/test/resources/rules/routing_rules_update.yml 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 515fa3aa2..326740201 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 @@ -13,16 +13,16 @@ */ package io.trino.gateway.ha.resource; -import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; -import com.fasterxml.jackson.dataformat.yaml.YAMLParser; import com.google.common.base.Strings; import com.google.inject.Inject; import com.google.inject.Singleton; 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.RoutingRulesConfiguration; +import io.trino.gateway.ha.config.UIConfiguration; import io.trino.gateway.ha.domain.Result; import io.trino.gateway.ha.domain.RoutingRules; import io.trino.gateway.ha.domain.TableData; @@ -41,6 +41,7 @@ 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; @@ -53,14 +54,11 @@ import jakarta.ws.rs.core.SecurityContext; import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Paths; import java.time.LocalDateTime; import java.time.ZoneId; import java.time.ZoneOffset; import java.time.ZonedDateTime; import java.time.format.DateTimeFormatter; -import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; @@ -80,7 +78,10 @@ public class GatewayWebAppResource private final QueryHistoryManager queryHistoryManager; private final BackendStateManager backendStateManager; private final ResourceGroupsManager resourceGroupsManager; - private final HaGatewayConfiguration configuration; + private final RoutingRulesManager routingRulesManager; + private final ObjectMapper yamlReader; + private final RoutingRulesConfiguration routingRulesConfiguration; + private final UIConfiguration uiConfiguration; @Inject public GatewayWebAppResource( @@ -88,13 +89,17 @@ public GatewayWebAppResource( QueryHistoryManager queryHistoryManager, BackendStateManager backendStateManager, ResourceGroupsManager resourceGroupsManager, + RoutingRulesManager routingRulesManager, HaGatewayConfiguration configuration) { 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.configuration = requireNonNull(configuration, "configuration is null"); + this.yamlReader = new ObjectMapper(new YAMLFactory()); + this.routingRulesManager = requireNonNull(routingRulesManager, "resourceGroupsManager is null"); + this.routingRulesConfiguration = configuration.getRoutingRules(); + this.uiConfiguration = configuration.getUiConfiguration(); } @POST @@ -445,20 +450,10 @@ public Response readExactMatchSourceSelector() @Produces(MediaType.APPLICATION_JSON) @Path("/getRoutingRules") public Response getRoutingRules() + throws IOException { - try { - String rulesConfigPath = configuration.getRoutingRules().getRulesConfigPath(); - YAMLFactory yamlFactory = new YAMLFactory(); - ObjectMapper yamlReader = new ObjectMapper(yamlFactory); - YAMLParser yamlParser = yamlFactory.createParser(new String(Files.readAllBytes(Paths.get(rulesConfigPath)))); - List routingRulesList = yamlReader - .readValues(yamlParser, new TypeReference() {}) - .readAll(); - return Response.ok(Result.ok(routingRulesList)).build(); - } - catch (IOException e) { - throw new RuntimeException(e); - } + List routingRulesList = routingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); + return Response.ok(Result.ok(routingRulesList)).build(); } @POST @@ -467,34 +462,9 @@ public Response getRoutingRules() @Produces(MediaType.APPLICATION_JSON) @Path("/updateRoutingRules") public synchronized Response updateRoutingRules(RoutingRules routingRules) + throws IOException { - String rulesConfigPath = configuration.getRoutingRules().getRulesConfigPath(); - ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory()); - List routingRulesList = new ArrayList<>(); - YAMLFactory yamlFactory = new YAMLFactory(); - try { - YAMLParser yamlParser = yamlFactory.createParser(new String(Files.readAllBytes(Paths.get(rulesConfigPath)))); - routingRulesList = yamlReader - .readValues(yamlParser, new TypeReference() {}) - .readAll(); - - for (int i = 0; i < routingRulesList.size(); i++) { - if (routingRulesList.get(i).name().equals(routingRules.name())) { - routingRulesList.set(i, routingRules); - break; - } - } - - ObjectMapper yamlWriter = new ObjectMapper(new YAMLFactory()); - StringBuilder yamlContent = new StringBuilder(); - for (RoutingRules rule : routingRulesList) { - yamlContent.append(yamlWriter.writeValueAsString(rule)); - } - Files.write(Paths.get(rulesConfigPath), yamlContent.toString().getBytes()); - } - catch (IOException e) { - throw new RuntimeException(e); - } + List routingRulesList = routingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); return Response.ok(Result.ok(routingRulesList)).build(); } @@ -504,6 +474,6 @@ public synchronized Response updateRoutingRules(RoutingRules routingRules) @Path("/getUIConfiguration") public Response getUIConfiguration() { - return Response.ok(Result.ok(configuration.getUiConfiguration())).build(); + 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..f1d81763d --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRulesManager.java @@ -0,0 +1,82 @@ +/* + * 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 io.trino.gateway.ha.config.RoutingRulesConfiguration; +import io.trino.gateway.ha.domain.RoutingRules; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.List; + +public class RoutingRulesManager +{ + public List getRoutingRules(RoutingRulesConfiguration configuration, ObjectMapper yamlReader) + throws IOException + { + String content = null; + try { + String rulesConfigPath = configuration.getRulesConfigPath(); + content = new String(Files.readAllBytes(Paths.get(rulesConfigPath))); + YAMLParser parser = new YAMLFactory().createParser(content); + List routingRulesList = new ArrayList<>(); + while (parser.nextToken() != null) { + RoutingRules routingRules = yamlReader.readValue(parser, RoutingRules.class); + routingRulesList.add(routingRules); + } + return routingRulesList; + } + catch (IOException e) { + throw new IOException(e); + } + } + + public List updateRoutingRules(RoutingRules routingRules, RoutingRulesConfiguration configuration, ObjectMapper yamlReader) + throws IOException + { + String rulesConfigPath = configuration.getRulesConfigPath(); + List routingRulesList = new ArrayList<>(); + try { + String content = new String(Files.readAllBytes(Paths.get(rulesConfigPath))); + YAMLParser parser = new YAMLFactory().createParser(content); + while (parser.nextToken() != null) { + RoutingRules routingRule = yamlReader.readValue(parser, RoutingRules.class); + routingRulesList.add(routingRule); + } + + for (int i = 0; i < routingRulesList.size(); i++) { + if (routingRulesList.get(i).name().equals(routingRules.name())) { + routingRulesList.set(i, routingRules); + break; + } + } + + ObjectMapper yamlWriter = new ObjectMapper(new YAMLFactory()); + StringBuilder yamlContent = new StringBuilder(); + for (RoutingRules rule : routingRulesList) { + yamlContent.append(yamlWriter.writeValueAsString(rule)); + } + Files.write(Paths.get(rulesConfigPath), yamlContent.toString().getBytes()); + } + catch (IOException e) { + throw new IOException(e); + } + return routingRulesList; + } +} 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..9af87e473 --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java @@ -0,0 +1,102 @@ +/* + * 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 io.trino.gateway.ha.config.RoutingRulesConfiguration; +import io.trino.gateway.ha.domain.RoutingRules; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.NoSuchFileException; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatException; + +class TestRoutingRulesManager +{ + private RoutingRulesConfiguration routingRulesConfiguration; + + private ObjectMapper yamlReader; + + private RoutingRulesManager routingRulesManager; + + @BeforeEach + void setUp() + { + routingRulesConfiguration = new RoutingRulesConfiguration(); + routingRulesManager = new RoutingRulesManager(); + yamlReader = new ObjectMapper(new YAMLFactory()); + } + + @Test + void testGetRoutingRules() + throws IOException + { + String rulesConfigPath = "src/test/resources/rules/routing_rules_atomic.yml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + + List result = routingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); + + assertThat(2).isEqualTo(result.size()); + assertThat("airflow").isEqualTo(result.get(0).name()); + assertThat("if query from airflow, route to etl group").isEqualTo(result.get(0).description()); + assertThat("request.getHeader(\"X-Trino-Source\") == \"airflow\" && (request.getHeader(\"X-Trino-Client-Tags\") == null || request.getHeader(\"X-Trino-Client-Tags\").isEmpty())").isEqualTo(result.get(0).condition()); + assertThat("result.put(\"routingGroup\", \"etl\")").isEqualTo(result.get(0).actions().get(0)); + assertThat(1).isEqualTo(result.get(0).actions().size()); + assertThat("airflow special").isEqualTo(result.get(1).name()); + } + + @Test + void testRoutingRulesNoSuchFileException() + { + String rulesConfigPath = "src/test/resources/rules/routing_rules_test.yaml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + + assertThatException().isThrownBy(() -> { + routingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); + }).withCauseExactlyInstanceOf(NoSuchFileException.class); + } + + @Test + void testUpdateRoutingRulesFile() + throws IOException + { + String rulesConfigPath = "src/test/resources/rules/routing_rules_update.yml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + RoutingRules routingRules = new RoutingRules("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.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); + assertThat("result.put(\"routingGroup\", \"adhoc\")").isEqualTo(updatedRoutingRules.get(0).actions().get(0)); + assertThat("request.getHeader(\"X-Trino-Source\") == \"JDBC\"").isEqualTo(updatedRoutingRules.get(0).condition()); + + RoutingRules originalRoutingRules = new RoutingRules("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); + routingRulesManager.updateRoutingRules(originalRoutingRules, routingRulesConfiguration, yamlReader); + } + + @Test + void testUpdateRoutingRulesNoSuchFileException() + { + String rulesConfigPath = "src/test/resources/rules/routing_rules_updated.yaml"; + routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + RoutingRules routingRules = new RoutingRules("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + + assertThatException().isThrownBy(() -> { + routingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); + }).withCauseExactlyInstanceOf(NoSuchFileException.class); + } +} 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..6b3bd67c8 --- /dev/null +++ b/gateway-ha/src/test/resources/rules/routing_rules_update.yml @@ -0,0 +1,7 @@ +--- +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\"" From 38023fc01a76103e384a73afaa0e567ed302f59c Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Mon, 7 Oct 2024 21:57:50 -0500 Subject: [PATCH 04/14] Fix build and remove json property --- .../main/java/io/trino/gateway/baseapp/BaseApp.java | 2 ++ .../io/trino/gateway/ha/domain/RoutingRules.java | 12 +++++------- .../trino/gateway/ha/router/RoutingRulesManager.java | 8 +++++--- 3 files changed, 12 insertions(+), 10 deletions(-) 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 7c169a531..27500a9d1 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 @@ -31,6 +31,7 @@ import io.trino.gateway.ha.resource.LoginResource; import io.trino.gateway.ha.resource.PublicResource; import io.trino.gateway.ha.resource.TrinoResource; +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; @@ -123,6 +124,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).in(Scopes.SINGLETON); } private static void addManagedApps(HaGatewayConfiguration configuration, Binder binder) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java index 8fef950d2..fa4fec58d 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java @@ -13,8 +13,6 @@ */ package io.trino.gateway.ha.domain; -import com.fasterxml.jackson.annotation.JsonProperty; - import java.util.List; /** @@ -27,10 +25,10 @@ * @param condition condition of the routing rule */ public record RoutingRules( - @JsonProperty("name") String name, - @JsonProperty("description") String description, - @JsonProperty("priority") Integer priority, - @JsonProperty("actions") List actions, - @JsonProperty("condition") String condition) + String name, + String description, + Integer priority, + List actions, + String condition) { } 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 index f1d81763d..f647fdce5 100644 --- 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 @@ -25,6 +25,8 @@ import java.util.ArrayList; import java.util.List; +import static java.nio.charset.StandardCharsets.UTF_8; + public class RoutingRulesManager { public List getRoutingRules(RoutingRulesConfiguration configuration, ObjectMapper yamlReader) @@ -33,7 +35,7 @@ public List getRoutingRules(RoutingRulesConfiguration configuratio String content = null; try { String rulesConfigPath = configuration.getRulesConfigPath(); - content = new String(Files.readAllBytes(Paths.get(rulesConfigPath))); + content = new String(Files.readAllBytes(Paths.get(rulesConfigPath)), UTF_8); YAMLParser parser = new YAMLFactory().createParser(content); List routingRulesList = new ArrayList<>(); while (parser.nextToken() != null) { @@ -53,7 +55,7 @@ public List updateRoutingRules(RoutingRules routingRules, RoutingR String rulesConfigPath = configuration.getRulesConfigPath(); List routingRulesList = new ArrayList<>(); try { - String content = new String(Files.readAllBytes(Paths.get(rulesConfigPath))); + String content = new String(Files.readAllBytes(Paths.get(rulesConfigPath)), UTF_8); YAMLParser parser = new YAMLFactory().createParser(content); while (parser.nextToken() != null) { RoutingRules routingRule = yamlReader.readValue(parser, RoutingRules.class); @@ -72,7 +74,7 @@ public List updateRoutingRules(RoutingRules routingRules, RoutingR for (RoutingRules rule : routingRulesList) { yamlContent.append(yamlWriter.writeValueAsString(rule)); } - Files.write(Paths.get(rulesConfigPath), yamlContent.toString().getBytes()); + Files.write(Paths.get(rulesConfigPath), yamlContent.toString().getBytes(UTF_8)); } catch (IOException e) { throw new IOException(e); From 546b85f6c0bbc53513a600c96727563255202703 Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Wed, 16 Oct 2024 16:04:41 -0500 Subject: [PATCH 05/14] Fix PR comments --- .../ha/config/HaGatewayConfiguration.java | 20 +++++------ .../trino/gateway/ha/domain/RoutingRules.java | 16 +++++++-- .../ha/resource/GatewayWebAppResource.java | 7 ++-- .../ha/router/RoutingRulesManager.java | 35 +++++++++---------- .../ha/router/TestRoutingRulesManager.java | 17 ++++----- 5 files changed, 48 insertions(+), 47 deletions(-) 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 2058a650f..13e495546 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,16 +46,6 @@ public class HaGatewayConfiguration private UIConfiguration uiConfiguration = new UIConfiguration(); - public UIConfiguration getUiConfiguration() - { - return uiConfiguration; - } - - public void setUiConfiguration(UIConfiguration uiConfiguration) - { - this.uiConfiguration = uiConfiguration; - } - // List of Modules with FQCN (Fully Qualified Class Name) private List modules; @@ -224,6 +214,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/domain/RoutingRules.java b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java index fa4fec58d..bdfea749e 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java @@ -13,22 +13,32 @@ */ 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; + /** * RoutingRules * * @param name name of the routing rule * @param description description of the routing rule - * @param priority priority 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 RoutingRules( String name, - String description, - Integer priority, + @Nullable String description, + @Nullable Integer priority, List actions, String condition) { + public RoutingRules { + requireNonNull(name, "name must not be null"); + requireNonNull(condition, "condition must not be null"); + actions = ImmutableList.copyOf(actions); + } } 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 326740201..4fce3908f 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 @@ -78,7 +78,6 @@ public class GatewayWebAppResource private final QueryHistoryManager queryHistoryManager; private final BackendStateManager backendStateManager; private final ResourceGroupsManager resourceGroupsManager; - private final RoutingRulesManager routingRulesManager; private final ObjectMapper yamlReader; private final RoutingRulesConfiguration routingRulesConfiguration; private final UIConfiguration uiConfiguration; @@ -89,7 +88,6 @@ public GatewayWebAppResource( QueryHistoryManager queryHistoryManager, BackendStateManager backendStateManager, ResourceGroupsManager resourceGroupsManager, - RoutingRulesManager routingRulesManager, HaGatewayConfiguration configuration) { this.gatewayBackendManager = requireNonNull(gatewayBackendManager, "gatewayBackendManager is null"); @@ -97,7 +95,6 @@ public GatewayWebAppResource( this.backendStateManager = requireNonNull(backendStateManager, "backendStateManager is null"); this.resourceGroupsManager = requireNonNull(resourceGroupsManager, "resourceGroupsManager is null"); this.yamlReader = new ObjectMapper(new YAMLFactory()); - this.routingRulesManager = requireNonNull(routingRulesManager, "resourceGroupsManager is null"); this.routingRulesConfiguration = configuration.getRoutingRules(); this.uiConfiguration = configuration.getUiConfiguration(); } @@ -452,7 +449,7 @@ public Response readExactMatchSourceSelector() public Response getRoutingRules() throws IOException { - List routingRulesList = routingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); + List routingRulesList = RoutingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); return Response.ok(Result.ok(routingRulesList)).build(); } @@ -464,7 +461,7 @@ public Response getRoutingRules() public synchronized Response updateRoutingRules(RoutingRules routingRules) throws IOException { - List routingRulesList = routingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); + List routingRulesList = RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); return Response.ok(Result.ok(routingRulesList)).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 index f647fdce5..7529c05f5 100644 --- 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 @@ -16,26 +16,29 @@ 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 io.trino.gateway.ha.config.RoutingRulesConfiguration; import io.trino.gateway.ha.domain.RoutingRules; import java.io.IOException; import java.nio.file.Files; +import java.nio.file.NoSuchFileException; import java.nio.file.Paths; import java.util.ArrayList; import java.util.List; import static java.nio.charset.StandardCharsets.UTF_8; -public class RoutingRulesManager +public final class RoutingRulesManager { - public List getRoutingRules(RoutingRulesConfiguration configuration, ObjectMapper yamlReader) + private RoutingRulesManager() {} + + public static List getRoutingRules(RoutingRulesConfiguration configuration, ObjectMapper yamlReader) throws IOException { - String content = null; + String rulesConfigPath = configuration.getRulesConfigPath(); try { - String rulesConfigPath = configuration.getRulesConfigPath(); - content = new String(Files.readAllBytes(Paths.get(rulesConfigPath)), UTF_8); + String content = Files.readString(Paths.get(rulesConfigPath), UTF_8); YAMLParser parser = new YAMLFactory().createParser(content); List routingRulesList = new ArrayList<>(); while (parser.nextToken() != null) { @@ -45,40 +48,34 @@ public List getRoutingRules(RoutingRulesConfiguration configuratio return routingRulesList; } catch (IOException e) { - throw new IOException(e); + throw new IOException("Failed to read or parse routing rules configuration from path : " + rulesConfigPath, e); } } - public List updateRoutingRules(RoutingRules routingRules, RoutingRulesConfiguration configuration, ObjectMapper yamlReader) + public static List updateRoutingRules(RoutingRules routingRules, RoutingRulesConfiguration configuration, ObjectMapper yamlReader) throws IOException { + ImmutableList.Builder routingRulesBuilder = ImmutableList.builder(); String rulesConfigPath = configuration.getRulesConfigPath(); - List routingRulesList = new ArrayList<>(); try { - String content = new String(Files.readAllBytes(Paths.get(rulesConfigPath)), UTF_8); - YAMLParser parser = new YAMLFactory().createParser(content); - while (parser.nextToken() != null) { - RoutingRules routingRule = yamlReader.readValue(parser, RoutingRules.class); - routingRulesList.add(routingRule); - } - + List routingRulesList = getRoutingRules(configuration, yamlReader); for (int i = 0; i < routingRulesList.size(); i++) { if (routingRulesList.get(i).name().equals(routingRules.name())) { routingRulesList.set(i, routingRules); break; } } - ObjectMapper yamlWriter = new ObjectMapper(new YAMLFactory()); StringBuilder yamlContent = new StringBuilder(); for (RoutingRules rule : routingRulesList) { yamlContent.append(yamlWriter.writeValueAsString(rule)); + routingRulesBuilder.add(rule); } - Files.write(Paths.get(rulesConfigPath), yamlContent.toString().getBytes(UTF_8)); + Files.writeString(Paths.get(rulesConfigPath), yamlContent.toString(), UTF_8); } catch (IOException e) { - throw new IOException(e); + throw new IOException("Failed to parse or update routing rules configuration form path : " + rulesConfigPath, e); } - return routingRulesList; + return routingRulesBuilder.build(); } } 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 index 9af87e473..f14d225cb 100644 --- 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 @@ -33,13 +33,10 @@ class TestRoutingRulesManager private ObjectMapper yamlReader; - private RoutingRulesManager routingRulesManager; - @BeforeEach void setUp() { routingRulesConfiguration = new RoutingRulesConfiguration(); - routingRulesManager = new RoutingRulesManager(); yamlReader = new ObjectMapper(new YAMLFactory()); } @@ -50,7 +47,7 @@ void testGetRoutingRules() String rulesConfigPath = "src/test/resources/rules/routing_rules_atomic.yml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); - List result = routingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); + List result = RoutingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); assertThat(2).isEqualTo(result.size()); assertThat("airflow").isEqualTo(result.get(0).name()); @@ -68,8 +65,8 @@ void testRoutingRulesNoSuchFileException() routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); assertThatException().isThrownBy(() -> { - routingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); - }).withCauseExactlyInstanceOf(NoSuchFileException.class); + RoutingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); + }).withRootCauseInstanceOf(NoSuchFileException.class); } @Test @@ -80,12 +77,12 @@ void testUpdateRoutingRulesFile() routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); RoutingRules routingRules = new RoutingRules("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.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); + List updatedRoutingRules = RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); assertThat("result.put(\"routingGroup\", \"adhoc\")").isEqualTo(updatedRoutingRules.get(0).actions().get(0)); assertThat("request.getHeader(\"X-Trino-Source\") == \"JDBC\"").isEqualTo(updatedRoutingRules.get(0).condition()); RoutingRules originalRoutingRules = new RoutingRules("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); - routingRulesManager.updateRoutingRules(originalRoutingRules, routingRulesConfiguration, yamlReader); + RoutingRulesManager.updateRoutingRules(originalRoutingRules, routingRulesConfiguration, yamlReader); } @Test @@ -96,7 +93,7 @@ void testUpdateRoutingRulesNoSuchFileException() RoutingRules routingRules = new RoutingRules("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); assertThatException().isThrownBy(() -> { - routingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); - }).withCauseExactlyInstanceOf(NoSuchFileException.class); + RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); + }).withRootCauseInstanceOf(NoSuchFileException.class); } } From 7cf593d7f7b6ee88247c08517934488b770cab84 Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Wed, 16 Oct 2024 16:06:05 -0500 Subject: [PATCH 06/14] Remove unused import --- .../java/io/trino/gateway/ha/router/RoutingRulesManager.java | 1 - 1 file changed, 1 deletion(-) 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 index 7529c05f5..d02fab0eb 100644 --- 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 @@ -22,7 +22,6 @@ import java.io.IOException; import java.nio.file.Files; -import java.nio.file.NoSuchFileException; import java.nio.file.Paths; import java.util.ArrayList; import java.util.List; From e3ea3cd33eb18b6dfd4c6980e29415e5809b7119 Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Wed, 16 Oct 2024 16:18:57 -0500 Subject: [PATCH 07/14] Fix CI Build --- gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java | 2 -- 1 file changed, 2 deletions(-) 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 27500a9d1..7c169a531 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 @@ -31,7 +31,6 @@ import io.trino.gateway.ha.resource.LoginResource; import io.trino.gateway.ha.resource.PublicResource; import io.trino.gateway.ha.resource.TrinoResource; -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; @@ -124,7 +123,6 @@ 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).in(Scopes.SINGLETON); } private static void addManagedApps(HaGatewayConfiguration configuration, Binder binder) From 3c980064b7216e80ba959bd826bfcde4806cbef9 Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Mon, 11 Nov 2024 22:42:14 -0600 Subject: [PATCH 08/14] Add some PR comment fixes --- docs/gateway-api.md | 2 +- .../gateway/ha/config/UIConfiguration.java | 8 --- .../{RoutingRules.java => RoutingRule.java} | 8 +-- .../ha/resource/GatewayWebAppResource.java | 14 ++-- .../ha/router/RoutingRulesManager.java | 17 ++--- .../ha/router/TestRoutingRulesManager.java | 65 +++++++++---------- 6 files changed, 47 insertions(+), 67 deletions(-) rename gateway-ha/src/main/java/io/trino/gateway/ha/domain/{RoutingRules.java => RoutingRule.java} (88%) diff --git a/docs/gateway-api.md b/docs/gateway-api.md index 2eee746d2..a6d52c836 100644 --- a/docs/gateway-api.md +++ b/docs/gateway-api.md @@ -96,7 +96,7 @@ 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 well you will need to provide a shared storage for the routing rules file. +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. ```shell curl -X POST http://localhost:8080/webapp/updateRoutingRules \ 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 index 10357e8d2..49dc28675 100644 --- 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 @@ -31,12 +31,4 @@ public void setDisablePages(List disablePages) { this.disablePages = disablePages; } - - @Override - public String toString() - { - return "UIConfiguration{" + - "disablePages=" + disablePages + - '}'; - } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java similarity index 88% rename from gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java rename to gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java index bdfea749e..c9d8935ee 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRules.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java @@ -29,16 +29,16 @@ * @param actions actions of the routing rule * @param condition condition of the routing rule */ -public record RoutingRules( +public record RoutingRule( String name, @Nullable String description, @Nullable Integer priority, List actions, String condition) { - public RoutingRules { - requireNonNull(name, "name must not be null"); - requireNonNull(condition, "condition must not be null"); + public RoutingRule { + requireNonNull(name, "name is null"); 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 4fce3908f..94b9ce6b2 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 @@ -13,18 +13,15 @@ */ package io.trino.gateway.ha.resource; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import com.google.common.base.Strings; import com.google.inject.Inject; -import com.google.inject.Singleton; 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.RoutingRulesConfiguration; import io.trino.gateway.ha.config.UIConfiguration; import io.trino.gateway.ha.domain.Result; -import io.trino.gateway.ha.domain.RoutingRules; +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; @@ -69,7 +66,6 @@ import static java.util.Objects.requireNonNullElse; @Path("/webapp") -@Singleton public class GatewayWebAppResource { private static final LocalDateTime START_TIME = LocalDateTime.now(); @@ -78,7 +74,6 @@ public class GatewayWebAppResource private final QueryHistoryManager queryHistoryManager; private final BackendStateManager backendStateManager; private final ResourceGroupsManager resourceGroupsManager; - private final ObjectMapper yamlReader; private final RoutingRulesConfiguration routingRulesConfiguration; private final UIConfiguration uiConfiguration; @@ -94,7 +89,6 @@ public GatewayWebAppResource( this.queryHistoryManager = requireNonNull(queryHistoryManager, "queryHistoryManager is null"); this.backendStateManager = requireNonNull(backendStateManager, "backendStateManager is null"); this.resourceGroupsManager = requireNonNull(resourceGroupsManager, "resourceGroupsManager is null"); - this.yamlReader = new ObjectMapper(new YAMLFactory()); this.routingRulesConfiguration = configuration.getRoutingRules(); this.uiConfiguration = configuration.getUiConfiguration(); } @@ -449,7 +443,7 @@ public Response readExactMatchSourceSelector() public Response getRoutingRules() throws IOException { - List routingRulesList = RoutingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); + List routingRulesList = RoutingRulesManager.getRoutingRules(routingRulesConfiguration); return Response.ok(Result.ok(routingRulesList)).build(); } @@ -458,10 +452,10 @@ public Response getRoutingRules() @Consumes(MediaType.APPLICATION_JSON) @Produces(MediaType.APPLICATION_JSON) @Path("/updateRoutingRules") - public synchronized Response updateRoutingRules(RoutingRules routingRules) + public synchronized Response updateRoutingRules(RoutingRule routingRules) throws IOException { - List routingRulesList = RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); + List routingRulesList = RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration); return Response.ok(Result.ok(routingRulesList)).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 index d02fab0eb..744b88fbf 100644 --- 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 @@ -18,7 +18,7 @@ import com.fasterxml.jackson.dataformat.yaml.YAMLParser; import com.google.common.collect.ImmutableList; import io.trino.gateway.ha.config.RoutingRulesConfiguration; -import io.trino.gateway.ha.domain.RoutingRules; +import io.trino.gateway.ha.domain.RoutingRule; import java.io.IOException; import java.nio.file.Files; @@ -32,16 +32,17 @@ public final class RoutingRulesManager { private RoutingRulesManager() {} - public static List getRoutingRules(RoutingRulesConfiguration configuration, ObjectMapper yamlReader) + public static List getRoutingRules(RoutingRulesConfiguration configuration) throws IOException { + ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory()); String rulesConfigPath = configuration.getRulesConfigPath(); try { String content = Files.readString(Paths.get(rulesConfigPath), UTF_8); YAMLParser parser = new YAMLFactory().createParser(content); - List routingRulesList = new ArrayList<>(); + List routingRulesList = new ArrayList<>(); while (parser.nextToken() != null) { - RoutingRules routingRules = yamlReader.readValue(parser, RoutingRules.class); + RoutingRule routingRules = yamlReader.readValue(parser, RoutingRule.class); routingRulesList.add(routingRules); } return routingRulesList; @@ -51,13 +52,13 @@ public static List getRoutingRules(RoutingRulesConfiguration confi } } - public static List updateRoutingRules(RoutingRules routingRules, RoutingRulesConfiguration configuration, ObjectMapper yamlReader) + public static List updateRoutingRules(RoutingRule routingRules, RoutingRulesConfiguration configuration) throws IOException { - ImmutableList.Builder routingRulesBuilder = ImmutableList.builder(); + ImmutableList.Builder routingRulesBuilder = ImmutableList.builder(); String rulesConfigPath = configuration.getRulesConfigPath(); try { - List routingRulesList = getRoutingRules(configuration, yamlReader); + List routingRulesList = getRoutingRules(configuration); for (int i = 0; i < routingRulesList.size(); i++) { if (routingRulesList.get(i).name().equals(routingRules.name())) { routingRulesList.set(i, routingRules); @@ -66,7 +67,7 @@ public static List updateRoutingRules(RoutingRules routingRules, R } ObjectMapper yamlWriter = new ObjectMapper(new YAMLFactory()); StringBuilder yamlContent = new StringBuilder(); - for (RoutingRules rule : routingRulesList) { + for (RoutingRule rule : routingRulesList) { yamlContent.append(yamlWriter.writeValueAsString(rule)); routingRulesBuilder.add(rule); } 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 index f14d225cb..eb1df3626 100644 --- 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 @@ -13,11 +13,8 @@ */ package io.trino.gateway.ha.router; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import io.trino.gateway.ha.config.RoutingRulesConfiguration; -import io.trino.gateway.ha.domain.RoutingRules; -import org.junit.jupiter.api.BeforeEach; +import io.trino.gateway.ha.domain.RoutingRule; import org.junit.jupiter.api.Test; import java.io.IOException; @@ -29,71 +26,67 @@ class TestRoutingRulesManager { - private RoutingRulesConfiguration routingRulesConfiguration; - - private ObjectMapper yamlReader; - - @BeforeEach - void setUp() - { - routingRulesConfiguration = new RoutingRulesConfiguration(); - yamlReader = new ObjectMapper(new YAMLFactory()); - } - @Test void testGetRoutingRules() throws IOException { + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); String rulesConfigPath = "src/test/resources/rules/routing_rules_atomic.yml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); - List result = RoutingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); + List result = RoutingRulesManager.getRoutingRules(routingRulesConfiguration); - assertThat(2).isEqualTo(result.size()); - assertThat("airflow").isEqualTo(result.get(0).name()); - assertThat("if query from airflow, route to etl group").isEqualTo(result.get(0).description()); - assertThat("request.getHeader(\"X-Trino-Source\") == \"airflow\" && (request.getHeader(\"X-Trino-Client-Tags\") == null || request.getHeader(\"X-Trino-Client-Tags\").isEmpty())").isEqualTo(result.get(0).condition()); - assertThat("result.put(\"routingGroup\", \"etl\")").isEqualTo(result.get(0).actions().get(0)); - assertThat(1).isEqualTo(result.get(0).actions().size()); - assertThat("airflow special").isEqualTo(result.get(1).name()); + assertThat(result.size()).isEqualTo(2); + assertThat(result.getFirst().name()).isEqualTo("airflow"); + assertThat(result.getFirst().description()).isEqualTo("if query from airflow, route to etl group"); + assertThat(result.getFirst().condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"airflow\" && (request.getHeader(\"X-Trino-Client-Tags\") == null || request.getHeader(\"X-Trino-Client-Tags\").isEmpty())"); + assertThat(result.getFirst().actions().getFirst()).isEqualTo("result.put(\"routingGroup\", \"etl\")"); + assertThat(result.getFirst().actions().size()).isEqualTo(1); + assertThat(result.get(1).name()).isEqualTo("airflow special"); } @Test void testRoutingRulesNoSuchFileException() { + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); String rulesConfigPath = "src/test/resources/rules/routing_rules_test.yaml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); - assertThatException().isThrownBy(() -> { - RoutingRulesManager.getRoutingRules(routingRulesConfiguration, yamlReader); - }).withRootCauseInstanceOf(NoSuchFileException.class); + assertThatException() + .isThrownBy(() -> RoutingRulesManager.getRoutingRules(routingRulesConfiguration)) + .withRootCauseInstanceOf(NoSuchFileException.class); } @Test void testUpdateRoutingRulesFile() throws IOException { + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); String rulesConfigPath = "src/test/resources/rules/routing_rules_update.yml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); - RoutingRules routingRules = new RoutingRules("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + 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.updateRoutingRules(routingRules, routingRulesConfiguration); + assertThat(updatedRoutingRules.getFirst().actions().getFirst()).isEqualTo("result.put(\"routingGroup\", \"adhoc\")"); + assertThat(updatedRoutingRules.getFirst().condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"JDBC\""); - List updatedRoutingRules = RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); - assertThat("result.put(\"routingGroup\", \"adhoc\")").isEqualTo(updatedRoutingRules.get(0).actions().get(0)); - assertThat("request.getHeader(\"X-Trino-Source\") == \"JDBC\"").isEqualTo(updatedRoutingRules.get(0).condition()); + 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.updateRoutingRules(originalRoutingRules, routingRulesConfiguration); - RoutingRules originalRoutingRules = new RoutingRules("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"etl\")"), "request.getHeader(\"X-Trino-Source\") == \"airflow\""); - RoutingRulesManager.updateRoutingRules(originalRoutingRules, routingRulesConfiguration, yamlReader); + assertThat(updateRoutingRules.getFirst().actions().getFirst()).isEqualTo("result.put(\"routingGroup\", \"etl\")"); + assertThat(updateRoutingRules.getFirst().condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"airflow\""); } @Test void testUpdateRoutingRulesNoSuchFileException() { + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); String rulesConfigPath = "src/test/resources/rules/routing_rules_updated.yaml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); - RoutingRules routingRules = new RoutingRules("airflow", "if query from airflow, route to etl group", 0, List.of("result.put(\"routingGroup\", \"adhoc\")"), "request.getHeader(\"X-Trino-Source\") == \"JDBC\""); + 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\""); - assertThatException().isThrownBy(() -> { - RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration, yamlReader); - }).withRootCauseInstanceOf(NoSuchFileException.class); + assertThatException() + .isThrownBy(() -> RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration)) + .withRootCauseInstanceOf(NoSuchFileException.class); } } From b8ba509ef41c2e648c5376459b4ebd89274ed55b Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Mon, 18 Nov 2024 12:41:21 -0600 Subject: [PATCH 09/14] Move code to service class --- .../java/io/trino/gateway/baseapp/BaseApp.java | 2 ++ .../gateway/ha/resource/GatewayWebAppResource.java | 9 ++++++--- .../gateway/ha/router/RoutingRulesManager.java | 8 +++----- .../gateway/ha/router/TestRoutingRulesManager.java | 14 +++++++++----- 4 files changed, 20 insertions(+), 13 deletions(-) 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 7c169a531..3c1a06081 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 @@ -31,6 +31,7 @@ import io.trino.gateway.ha.resource.LoginResource; import io.trino.gateway.ha.resource.PublicResource; import io.trino.gateway.ha.resource.TrinoResource; +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; @@ -123,6 +124,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/resource/GatewayWebAppResource.java b/gateway-ha/src/main/java/io/trino/gateway/ha/resource/GatewayWebAppResource.java index 94b9ce6b2..7594276e8 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 @@ -76,6 +76,7 @@ public class GatewayWebAppResource private final ResourceGroupsManager resourceGroupsManager; private final RoutingRulesConfiguration routingRulesConfiguration; private final UIConfiguration uiConfiguration; + private final RoutingRulesManager routingRulesManager; @Inject public GatewayWebAppResource( @@ -83,7 +84,8 @@ public GatewayWebAppResource( QueryHistoryManager queryHistoryManager, BackendStateManager backendStateManager, ResourceGroupsManager resourceGroupsManager, - HaGatewayConfiguration configuration) + HaGatewayConfiguration configuration, + RoutingRulesManager routingRulesManager) { this.gatewayBackendManager = requireNonNull(gatewayBackendManager, "gatewayBackendManager is null"); this.queryHistoryManager = requireNonNull(queryHistoryManager, "queryHistoryManager is null"); @@ -91,6 +93,7 @@ public GatewayWebAppResource( this.resourceGroupsManager = requireNonNull(resourceGroupsManager, "resourceGroupsManager is null"); this.routingRulesConfiguration = configuration.getRoutingRules(); this.uiConfiguration = configuration.getUiConfiguration(); + this.routingRulesManager = requireNonNull(routingRulesManager, "routingRulesManager is null"); } @POST @@ -443,7 +446,7 @@ public Response readExactMatchSourceSelector() public Response getRoutingRules() throws IOException { - List routingRulesList = RoutingRulesManager.getRoutingRules(routingRulesConfiguration); + List routingRulesList = routingRulesManager.getRoutingRules(routingRulesConfiguration); return Response.ok(Result.ok(routingRulesList)).build(); } @@ -455,7 +458,7 @@ public Response getRoutingRules() public synchronized Response updateRoutingRules(RoutingRule routingRules) throws IOException { - List routingRulesList = RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration); + List routingRulesList = routingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration); return Response.ok(Result.ok(routingRulesList)).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 index 744b88fbf..28fed4227 100644 --- 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 @@ -28,11 +28,9 @@ import static java.nio.charset.StandardCharsets.UTF_8; -public final class RoutingRulesManager +public class RoutingRulesManager { - private RoutingRulesManager() {} - - public static List getRoutingRules(RoutingRulesConfiguration configuration) + public List getRoutingRules(RoutingRulesConfiguration configuration) throws IOException { ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory()); @@ -52,7 +50,7 @@ public static List getRoutingRules(RoutingRulesConfiguration config } } - public static List updateRoutingRules(RoutingRule routingRules, RoutingRulesConfiguration configuration) + public List updateRoutingRules(RoutingRule routingRules, RoutingRulesConfiguration configuration) throws IOException { ImmutableList.Builder routingRulesBuilder = ImmutableList.builder(); 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 index eb1df3626..6a8df43ae 100644 --- 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 @@ -33,8 +33,9 @@ void testGetRoutingRules() RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); String rulesConfigPath = "src/test/resources/rules/routing_rules_atomic.yml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(); - List result = RoutingRulesManager.getRoutingRules(routingRulesConfiguration); + List result = routingRulesManager.getRoutingRules(routingRulesConfiguration); assertThat(result.size()).isEqualTo(2); assertThat(result.getFirst().name()).isEqualTo("airflow"); @@ -49,11 +50,12 @@ void testGetRoutingRules() void testRoutingRulesNoSuchFileException() { RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(); String rulesConfigPath = "src/test/resources/rules/routing_rules_test.yaml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); assertThatException() - .isThrownBy(() -> RoutingRulesManager.getRoutingRules(routingRulesConfiguration)) + .isThrownBy(() -> routingRulesManager.getRoutingRules(routingRulesConfiguration)) .withRootCauseInstanceOf(NoSuchFileException.class); } @@ -62,16 +64,17 @@ void testUpdateRoutingRulesFile() throws IOException { RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(); String rulesConfigPath = "src/test/resources/rules/routing_rules_update.yml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); 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.updateRoutingRules(routingRules, routingRulesConfiguration); + List updatedRoutingRules = routingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration); 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.updateRoutingRules(originalRoutingRules, routingRulesConfiguration); + List updateRoutingRules = routingRulesManager.updateRoutingRules(originalRoutingRules, routingRulesConfiguration); assertThat(updateRoutingRules.getFirst().actions().getFirst()).isEqualTo("result.put(\"routingGroup\", \"etl\")"); assertThat(updateRoutingRules.getFirst().condition()).isEqualTo("request.getHeader(\"X-Trino-Source\") == \"airflow\""); @@ -81,12 +84,13 @@ void testUpdateRoutingRulesFile() void testUpdateRoutingRulesNoSuchFileException() { RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + RoutingRulesManager routingRulesManager = new RoutingRulesManager(); String rulesConfigPath = "src/test/resources/rules/routing_rules_updated.yaml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); 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\""); assertThatException() - .isThrownBy(() -> RoutingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration)) + .isThrownBy(() -> routingRulesManager.updateRoutingRules(routingRules, routingRulesConfiguration)) .withRootCauseInstanceOf(NoSuchFileException.class); } } From 5799ce8b2b3edcb3398dbe68812d98aa920bf55a Mon Sep 17 00:00:00 2001 From: prakharsapre Date: Mon, 18 Nov 2024 12:46:18 -0600 Subject: [PATCH 10/14] Fix conflict --- webapp/src/components/layout.tsx | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/webapp/src/components/layout.tsx b/webapp/src/components/layout.tsx index 6b25ca74a..3b54e0b4e 100644 --- a/webapp/src/components/layout.tsx +++ b/webapp/src/components/layout.tsx @@ -1,5 +1,5 @@ import { Nav, Avatar, Layout, Dropdown, Button, Toast, Modal, Tag } from '@douyinfe/semi-ui'; -import { IconGithubLogo, IconDoubleChevronRight, IconDoubleChevronLeft, IconMoon, IconSun, IconMark, IconIdCard, IconUserSetting, IconUser } 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"; @@ -96,10 +96,6 @@ export const RootLayout = (props: { aria-label="Switch Theme" onClick={nextTheme} /> -