diff --git a/gateway-ha/pom.xml b/gateway-ha/pom.xml index 93c921fcf..da36f5fff 100644 --- a/gateway-ha/pom.xml +++ b/gateway-ha/pom.xml @@ -492,6 +492,7 @@ org.apache.maven.plugins maven-surefire-plugin + -Djol.skipHotspotSAAttach=true false 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 33302dc9a..5f496b472 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 @@ -35,7 +35,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.router.ForwardingRoutingRulesManager; import io.trino.gateway.ha.security.AuthorizedExceptionMapper; import io.trino.gateway.proxyserver.ForProxy; import io.trino.gateway.proxyserver.ProxyRequestHandler; @@ -145,7 +145,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); + binder.bind(ForwardingRoutingRulesManager.class); } private static void addManagedApps(HaGatewayConfiguration configuration, Binder binder) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/RulesType.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/RulesType.java index dc10ed964..66a79fb28 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/config/RulesType.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/config/RulesType.java @@ -37,4 +37,9 @@ public enum RulesType * The service URL can implement dynamic rule changes. */ EXTERNAL, + + /** + * Routing rules stored in the database. + */ + DB } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java index 0a4a7e2e7..c706c507a 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java @@ -201,10 +201,7 @@ public RoutingGroupSelector getRoutingGroupSelector(@ForRouter HttpClient httpCl if (routingRulesConfig.isRulesEngineEnabled()) { try { return switch (routingRulesConfig.getRulesType()) { - case FILE -> RoutingGroupSelector.byRoutingRulesEngine( - routingRulesConfig.getRulesConfigPath(), - routingRulesConfig.getRulesRefreshPeriod(), - configuration.getRequestAnalyzerConfig()); + case DB, FILE -> RoutingGroupSelector.byRoutingRulesEngine(configuration); case EXTERNAL -> { RulesExternalConfiguration rulesExternalConfiguration = routingRulesConfig.getRulesExternalConfiguration(); yield RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, configuration.getRequestAnalyzerConfig()); diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/module/RouterBaseModule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/module/RouterBaseModule.java index 3903e252b..81242bda5 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/module/RouterBaseModule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/module/RouterBaseModule.java @@ -36,10 +36,11 @@ public class RouterBaseModule public RouterBaseModule(HaGatewayConfiguration configuration) { Jdbi jdbi = Jdbi.create(configuration.getDataStore().getJdbcUrl(), configuration.getDataStore().getUser(), configuration.getDataStore().getPassword()); + boolean isOracleBackend = configuration.getDataStore().getJdbcUrl().startsWith("jdbc:oracle"); connectionManager = new JdbcConnectionManager(jdbi, configuration.getDataStore()); resourceGroupsManager = new HaResourceGroupsManager(connectionManager); - gatewayBackendManager = new HaGatewayManager(jdbi); - queryHistoryManager = new HaQueryHistoryManager(jdbi, configuration.getDataStore().getJdbcUrl().startsWith("jdbc:oracle")); + gatewayBackendManager = new HaGatewayManager(jdbi, isOracleBackend); + queryHistoryManager = new HaQueryHistoryManager(jdbi, isOracleBackend); } @Provides diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackendDao.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackendDao.java index 87009a03b..9f50ad1a3 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackendDao.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/GatewayBackendDao.java @@ -13,40 +13,102 @@ */ package io.trino.gateway.ha.persistence.dao; +import org.jdbi.v3.core.mapper.RowMapper; +import org.jdbi.v3.core.statement.StatementContext; import org.jdbi.v3.sqlobject.statement.SqlQuery; import org.jdbi.v3.sqlobject.statement.SqlUpdate; +import org.jdbi.v3.sqlobject.statement.UseRowMapper; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; import java.util.List; +import static java.util.Locale.ENGLISH; + public interface GatewayBackendDao { + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) @SqlQuery("SELECT * FROM gateway_backend") List findAll(); + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) @SqlQuery(""" SELECT * FROM gateway_backend WHERE active = true """) List findActiveBackend(); + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) + @SqlQuery(""" + SELECT * FROM gateway_backend + WHERE active = 1 + """) + List findActiveBackendNoBoolean(); + + default List findActiveBackend(boolean isSupportsBooleanColumn) + { + if (isSupportsBooleanColumn) { + return findActiveBackend(); + } + + return findActiveBackendNoBoolean(); + } + + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) @SqlQuery(""" SELECT * FROM gateway_backend WHERE active = true AND routing_group = 'adhoc' """) List findActiveAdhocBackend(); + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) + @SqlQuery(""" + SELECT * FROM gateway_backend + WHERE active = 1 AND routing_group = 'adhoc' + """) + List findActiveAdhocBackendNoBoolean(); + + default List findActiveAdhocBackend(boolean isSupportsBooleanColumn) + { + if (isSupportsBooleanColumn) { + return findActiveAdhocBackend(); + } + + return findActiveAdhocBackendNoBoolean(); + } + + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) @SqlQuery(""" SELECT * FROM gateway_backend WHERE active = true AND routing_group = :routingGroup """) List findActiveBackendByRoutingGroup(String routingGroup); + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) + @SqlQuery(""" + SELECT * FROM gateway_backend + WHERE active = 1 AND routing_group = :routingGroup + """) + List findActiveBackendByRoutingGroupNoBoolean(String routingGroup); + + default List findActiveBackendByRoutingGroup(String routingGroup, boolean isSupportsBooleanColumn) + { + if (isSupportsBooleanColumn) { + return findActiveBackendByRoutingGroup(routingGroup); + } + + return findActiveBackendByRoutingGroupNoBoolean(routingGroup); + } + + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) @SqlQuery(""" SELECT * FROM gateway_backend WHERE name = :name """) List findByName(String name); + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) @SqlQuery(""" SELECT * FROM gateway_backend WHERE name = :name @@ -54,12 +116,43 @@ public interface GatewayBackendDao """) GatewayBackend findFirstByName(String name); + @UseRowMapper(GatewayBackendIntToBooleanMapper.class) + @SqlQuery(""" + SELECT * FROM gateway_backend + WHERE name = :name + FETCH FIRST 1 ROWS ONLY + """) + GatewayBackend findFirstByNameWithFetch(String name); + + default GatewayBackend findFirstByName(String name, boolean isLimitUnsupported){ + if (isLimitUnsupported) { + return findFirstByNameWithFetch(name); + } + + return findFirstByName(name); + } + @SqlUpdate(""" INSERT INTO gateway_backend (name, routing_group, backend_url, external_url, active) VALUES (:name, :routingGroup, :backendUrl, :externalUrl, :active) """) void create(String name, String routingGroup, String backendUrl, String externalUrl, boolean active); + @SqlUpdate(""" + INSERT INTO gateway_backend (name, routing_group, backend_url, external_url, active) + VALUES (:name, :routingGroup, :backendUrl, :externalUrl, :active) + """) + void createNoBoolean(String name, String routingGroup, String backendUrl, String externalUrl, int active); + + default void create(String name, String routingGroup, String backendUrl, String externalUrl, boolean active, boolean isSupportsBooleanColumn) + { + if (isSupportsBooleanColumn) { + create(name, routingGroup, backendUrl, externalUrl, active); + return; + } + createNoBoolean(name, routingGroup, backendUrl, externalUrl, active ? 1 : 0); + } + @SqlUpdate(""" UPDATE gateway_backend SET routing_group = :routingGroup, backend_url = :backendUrl, external_url = :externalUrl, active = :active @@ -67,6 +160,22 @@ INSERT INTO gateway_backend (name, routing_group, backend_url, external_url, act """) void update(String name, String routingGroup, String backendUrl, String externalUrl, boolean active); + @SqlUpdate(""" + UPDATE gateway_backend + SET routing_group = :routingGroup, backend_url = :backendUrl, external_url = :externalUrl, active = :active + WHERE name = :name + """) + void updateNoBoolean(String name, String routingGroup, String backendUrl, String externalUrl, int active); + + default void update(String name, String routingGroup, String backendUrl, String externalUrl, boolean active, boolean isSupportsBooleanColumn) + { + if (isSupportsBooleanColumn) { + update(name, routingGroup, backendUrl, externalUrl, active); + return; + } + updateNoBoolean(name, routingGroup, backendUrl, externalUrl, active ? 1 : 0); + } + @SqlUpdate(""" UPDATE gateway_backend SET active = false @@ -86,4 +195,28 @@ INSERT INTO gateway_backend (name, routing_group, backend_url, external_url, act WHERE name = :name """) void deleteByName(String name); + + class GatewayBackendIntToBooleanMapper + implements RowMapper + { + @Override + public GatewayBackend map(ResultSet resultSet, StatementContext ctx) + throws SQLException + { + boolean active; + ResultSetMetaData resultSetMetaData = resultSet.getMetaData(); + if (resultSetMetaData.getColumnClassName(5).toLowerCase(ENGLISH).startsWith("int")) { + active = resultSet.getInt(5) != 0; + } + else { + active = resultSet.getBoolean(5); + } + return new GatewayBackend( + resultSet.getString("name"), + resultSet.getString("routing_group"), + resultSet.getString("backend_url"), + resultSet.getString("external_url"), + active); + } + } } 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/persistence/dao/RoutingRule.java similarity index 54% rename from gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java rename to gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/RoutingRule.java index db3f905ef..fef872ecd 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/domain/RoutingRule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/RoutingRule.java @@ -11,14 +11,14 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.gateway.ha.domain; +package io.trino.gateway.ha.persistence.dao; -import com.google.common.collect.ImmutableList; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.jdbi.v3.core.mapper.reflect.ColumnName; import java.util.List; import static java.util.Objects.requireNonNull; -import static java.util.Objects.requireNonNullElse; /** * RoutingRules @@ -28,19 +28,22 @@ * @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 + * @param routingRuleEngine the engine used for rule evaluation */ + public record RoutingRule( - String name, - String description, - Integer priority, - List actions, - String condition) + @JsonProperty("name") @ColumnName("name") String name, + @JsonProperty("description") @ColumnName("description") String description, + @JsonProperty("priority") @ColumnName("priority") Integer priority, + // "conditionExpression" is used as a column name because "condition" is a reserved word in MySQL + @JsonProperty("condition") @ColumnName("conditionExpression") String condition, + @JsonProperty("actions") @ColumnName("actions") List actions, + @JsonProperty("routingRuleEngine") @ColumnName("routingRuleEngine") RoutingRuleEngine routingRuleEngine) { - public RoutingRule { - requireNonNull(name, "name is null"); - description = requireNonNullElse(description, ""); - priority = requireNonNullElse(priority, 0); - actions = ImmutableList.copyOf(actions); - requireNonNull(condition, "condition is null"); + public RoutingRule + { + requireNonNull(name); + requireNonNull(condition); + requireNonNull(actions); } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/RoutingRuleEngine.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/RoutingRuleEngine.java new file mode 100644 index 000000000..92f0e9669 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/RoutingRuleEngine.java @@ -0,0 +1,19 @@ +/* + * 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.persistence.dao; + +public enum RoutingRuleEngine +{ + MVEL +} diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/RoutingRulesDao.java b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/RoutingRulesDao.java new file mode 100644 index 000000000..0c83f1a79 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/persistence/dao/RoutingRulesDao.java @@ -0,0 +1,97 @@ +/* + * 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.persistence.dao; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.jdbi.v3.core.mapper.RowMapper; +import org.jdbi.v3.core.statement.StatementContext; +import org.jdbi.v3.sqlobject.statement.SqlQuery; +import org.jdbi.v3.sqlobject.statement.SqlUpdate; +import org.jdbi.v3.sqlobject.statement.UseRowMapper; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.List; +import java.util.Optional; + +public interface RoutingRulesDao +{ + @SqlQuery("SELECT * FROM routing_rules") + List getAll(); + + // "conditionExpression" is used as a column name because "condition" is a reserved word in MySQL + @SqlUpdate(""" + INSERT INTO routing_rules (name, description, priority, conditionExpression, actions, routingRuleEngine) + VALUES (:name, :description, :priority, :condition, :actions, :routingRuleEngine) + """) + void create(String name, String description, Integer priority, String condition, List actions, RoutingRuleEngine routingRuleEngine); + + @SqlUpdate(""" + UPDATE routing_rules + SET description = :description, priority = :priority, conditionExpression = :condition, actions = :actions, routingRuleEngine = :routingRuleEngine + WHERE name = :name + """) + void update(String name, String description, Integer priority, String condition, List actions, RoutingRuleEngine routingRuleEngine); + + @SqlQuery("SELECT * FROM routing_rules") + @UseRowMapper(RoutingRulesStringToListMapper.class) + List getAllNoListSupport(); + + @SqlUpdate(""" + INSERT INTO routing_rules (name, description, priority, conditionExpression, actions, routingRuleEngine) + VALUES (:name, :description, :priority, :condition, :actions, :routingRuleEngine) + """) + void createNoListSupport(String name, String description, Integer priority, String condition, String actions, RoutingRuleEngine routingRuleEngine); + + @SqlUpdate(""" + UPDATE routing_rules + SET description = :description, priority = :priority, conditionExpression = :condition, actions = :actions, routingRuleEngine = :routingRuleEngine + WHERE name = :name + """) + void updateNoListSupport(String name, String description, Integer priority, String condition, String actions, RoutingRuleEngine routingRuleEngine); + + @SqlUpdate(""" + DELETE FROM routing_rules + WHERE name = :name + """ + ) + void delete(String name); + + class RoutingRulesStringToListMapper + implements RowMapper + { + ObjectMapper objectMapper = new ObjectMapper(); + TypeReference> actionsTypeReference = new TypeReference>() {}; + + @Override + public RoutingRule map(ResultSet rs, StatementContext ctx) + throws SQLException + { + try { + return new RoutingRule( + rs.getString("name"), + rs.getString("description"), + rs.getInt("priority"), + rs.getString("conditionExpression"), + objectMapper.readValue(rs.getString("actions"), actionsTypeReference), + RoutingRuleEngine.valueOf(Optional.ofNullable(rs.getString("routingRuleEngine")).orElse("MVEL"))); + } + catch (JsonProcessingException e) { + throw new RuntimeException(e); + } + } + } +} 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 34ddb9abd..ba8e91e87 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 @@ -20,7 +20,6 @@ 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; @@ -32,17 +31,21 @@ import io.trino.gateway.ha.domain.request.SelectorsRequest; import io.trino.gateway.ha.domain.response.BackendResponse; import io.trino.gateway.ha.domain.response.DistributionResponse; +import io.trino.gateway.ha.persistence.dao.RoutingRule; import io.trino.gateway.ha.router.BackendStateManager; +import io.trino.gateway.ha.router.ForwardingRoutingRulesManager; import io.trino.gateway.ha.router.GatewayBackendManager; import io.trino.gateway.ha.router.HaGatewayManager; +import io.trino.gateway.ha.router.IRoutingRulesManager; 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.DELETE; import jakarta.ws.rs.GET; import jakarta.ws.rs.POST; import jakarta.ws.rs.Path; +import jakarta.ws.rs.PathParam; import jakarta.ws.rs.Produces; import jakarta.ws.rs.core.Context; import jakarta.ws.rs.core.MediaType; @@ -75,7 +78,7 @@ public class GatewayWebAppResource private final ResourceGroupsManager resourceGroupsManager; // TODO Avoid putting mutable objects in fields private final UIConfiguration uiConfiguration; - private final RoutingRulesManager routingRulesManager; + private final IRoutingRulesManager routingRulesManager; @Inject public GatewayWebAppResource( @@ -83,7 +86,7 @@ public GatewayWebAppResource( QueryHistoryManager queryHistoryManager, BackendStateManager backendStateManager, ResourceGroupsManager resourceGroupsManager, - RoutingRulesManager routingRulesManager, + ForwardingRoutingRulesManager routingRulesManager, HaGatewayConfiguration configuration) { this.gatewayBackendManager = requireNonNull(gatewayBackendManager, "gatewayBackendManager is null"); @@ -450,6 +453,17 @@ public Response getRoutingRules() return Response.ok(Result.ok(routingRulesList)).build(); } + @DELETE + @RolesAllowed("ADMIN") + @Produces(MediaType.APPLICATION_JSON) + @Path("/deleteRoutingRule/{name}") + public Response deleteRoutingRules(@PathParam("name") String name) + throws IOException + { + routingRulesManager.deleteRoutingRule(name); + return Response.ok(routingRulesManager.getRoutingRules()).build(); + } + @POST @RolesAllowed("ADMIN") @Consumes(MediaType.APPLICATION_JSON) @@ -462,6 +476,16 @@ public Response updateRoutingRules(RoutingRule routingRule) return Response.ok(Result.ok(routingRulesList)).build(); } + @POST + @RolesAllowed("ADMIN") + @Consumes(MediaType.APPLICATION_JSON) + @Path("/createRoutingRule") + public Response createRoutingRule(RoutingRule routingRule) + { + routingRulesManager.createRoutingRule(routingRule); + return Response.ok(routingRulesManager.getRoutingRules()).build(); + } + @GET @RolesAllowed("ADMIN") @Produces(MediaType.APPLICATION_JSON) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/DbRoutingRulesManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/DbRoutingRulesManager.java new file mode 100644 index 000000000..d36f4645d --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/DbRoutingRulesManager.java @@ -0,0 +1,107 @@ +/* + * 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.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.trino.gateway.ha.persistence.dao.RoutingRule; +import io.trino.gateway.ha.persistence.dao.RoutingRulesDao; + +import java.util.List; + +public class DbRoutingRulesManager + implements IRoutingRulesManager +{ + private final RoutingRulesDao routingRulesDao; + private final boolean dbSupportsArray; + + ObjectMapper objectMapper = new ObjectMapper(); + + public DbRoutingRulesManager(RoutingRulesDao routingRulesDao, boolean dbSupportsArray) + { + this.routingRulesDao = routingRulesDao; + this.dbSupportsArray = dbSupportsArray; + } + + @Override + public List getRoutingRules() + { + if (dbSupportsArray) { + return routingRulesDao.getAll(); + } + return routingRulesDao.getAllNoListSupport(); + } + + @Override + public List updateRoutingRule(RoutingRule routingRule) + { + if (dbSupportsArray) { + routingRulesDao.update( + routingRule.name(), + routingRule.description(), + routingRule.priority(), + routingRule.condition(), + routingRule.actions(), + routingRule.routingRuleEngine()); + return routingRulesDao.getAll(); + } + else { + routingRulesDao.updateNoListSupport(routingRule.name(), + routingRule.description(), + routingRule.priority(), + routingRule.condition(), + toJson(routingRule.actions()), + routingRule.routingRuleEngine()); + return routingRulesDao.getAllNoListSupport(); + } + } + + @Override + public void deleteRoutingRule(String name) + { + routingRulesDao.delete(name); + } + + @Override + public void createRoutingRule(RoutingRule routingRule) + { + if (dbSupportsArray) { + routingRulesDao.create( + routingRule.name(), + routingRule.description(), + routingRule.priority(), + routingRule.condition(), + routingRule.actions(), + routingRule.routingRuleEngine()); + } + else { + routingRulesDao.createNoListSupport(routingRule.name(), + routingRule.description(), + routingRule.priority(), + routingRule.condition(), + toJson(routingRule.actions()), + routingRule.routingRuleEngine()); + } + } + + private String toJson(Object object) + { + try { + return objectMapper.writeValueAsString(object); + } + catch (JsonProcessingException e) { + throw new RuntimeException("Failed serializing routing rules to JSON", e); + } + } +} 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/FileBasedRoutingRulesManager.java similarity index 66% rename from gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRulesManager.java rename to gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingRulesManager.java index 817a16a86..0b2d6f05f 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/FileBasedRoutingRulesManager.java @@ -17,9 +17,8 @@ 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 io.trino.gateway.ha.persistence.dao.RoutingRule; import java.io.IOException; import java.io.UncheckedIOException; @@ -29,20 +28,22 @@ import java.nio.file.Path; import java.util.List; +import static com.google.common.collect.ImmutableList.toImmutableList; import static java.nio.charset.StandardCharsets.UTF_8; import static java.nio.file.StandardOpenOption.READ; import static java.nio.file.StandardOpenOption.WRITE; -public class RoutingRulesManager +public class FileBasedRoutingRulesManager + implements IRoutingRulesManager { private final String rulesConfigPath; - @Inject - public RoutingRulesManager(HaGatewayConfiguration configuration) + public FileBasedRoutingRulesManager(HaGatewayConfiguration configuration) { this.rulesConfigPath = configuration.getRoutingRules().getRulesConfigPath(); } + @Override public List getRoutingRules() { YAMLFactory yamlFactory = new YAMLFactory(); @@ -62,6 +63,7 @@ public List getRoutingRules() } } + @Override public synchronized List updateRoutingRule(RoutingRule routingRule) { ImmutableList.Builder updatedRoutingRulesBuilder = ImmutableList.builder(); @@ -89,4 +91,39 @@ public synchronized List updateRoutingRule(RoutingRule routingRule) } return updatedRoutingRulesBuilder.build(); } + + @Override + public void deleteRoutingRule(String name) + { + List currentRoutingRulesList = getRoutingRules(); + List updatedRulesList = currentRoutingRulesList.stream().filter(routingRule -> !routingRule.name().equals(name)).collect(toImmutableList()); + Path path = Path.of(rulesConfigPath); + writeRulesFile(path, updatedRulesList); + } + + @Override + public void createRoutingRule(RoutingRule routingRule) + { + List currentRoutingRulesList = getRoutingRules(); + List updatedRulesList = ImmutableList.builder().addAll(currentRoutingRulesList).add(routingRule).build(); + Path path = Path.of(rulesConfigPath); + writeRulesFile(path, updatedRulesList); + } + + private void writeRulesFile(Path path, List routingRules) + { + try (FileChannel fileChannel = FileChannel.open(path, WRITE, READ); + FileLock lock = fileChannel.lock()) { + ObjectMapper yamlWriter = new ObjectMapper(new YAMLFactory()); + StringBuilder yamlContent = new StringBuilder(); + for (RoutingRule rule : routingRules) { + yamlContent.append(yamlWriter.writeValueAsString(rule)); + } + Files.writeString(path, yamlContent.toString(), UTF_8); + lock.release(); + } + catch (IOException e) { + throw new UncheckedIOException("Failed to parse or update routing rules configuration form path : " + rulesConfigPath, e); + } + } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/ForwardingRoutingRulesManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/ForwardingRoutingRulesManager.java new file mode 100644 index 000000000..fe6049bc4 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/ForwardingRoutingRulesManager.java @@ -0,0 +1,76 @@ +/* + * 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.google.inject.Inject; +import io.trino.gateway.ha.config.HaGatewayConfiguration; +import io.trino.gateway.ha.config.RoutingRulesConfiguration; +import io.trino.gateway.ha.persistence.RecordAndAnnotatedConstructorMapper; +import io.trino.gateway.ha.persistence.dao.RoutingRule; +import io.trino.gateway.ha.persistence.dao.RoutingRulesDao; +import org.jdbi.v3.core.Jdbi; +import org.jdbi.v3.sqlobject.SqlObjectPlugin; + +import java.util.List; + +public class ForwardingRoutingRulesManager + implements IRoutingRulesManager +{ + IRoutingRulesManager delegate; + + @Inject + ForwardingRoutingRulesManager(HaGatewayConfiguration haGatewayConfiguration) + { + RoutingRulesConfiguration routingRulesConfig = haGatewayConfiguration.getRoutingRules(); + delegate = switch (routingRulesConfig.getRulesType()) { + case FILE -> new FileBasedRoutingRulesManager(haGatewayConfiguration); + case DB -> { + String jdbcUrl = haGatewayConfiguration.getDataStore().getJdbcUrl(); + Jdbi jdbi = Jdbi.create( + jdbcUrl, + haGatewayConfiguration.getDataStore().getUser(), + haGatewayConfiguration.getDataStore().getPassword()) + .installPlugin(new SqlObjectPlugin()) + .registerRowMapper(new RecordAndAnnotatedConstructorMapper()); + + yield new DbRoutingRulesManager(jdbi.onDemand(RoutingRulesDao.class), jdbcUrl.startsWith("jdbc:postgresql")); + } + default -> throw new RuntimeException("No routing manager for " + routingRulesConfig.getRulesType()); + }; + } + + @Override + public List getRoutingRules() + { + return delegate.getRoutingRules(); + } + + @Override + public List updateRoutingRule(RoutingRule routingRule) + { + return delegate.updateRoutingRule(routingRule); + } + + @Override + public void deleteRoutingRule(String name) + { + delegate.deleteRoutingRule(name); + } + + @Override + public void createRoutingRule(RoutingRule routingRule) + { + delegate.createRoutingRule(routingRule); + } +} diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaGatewayManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaGatewayManager.java index efa5a8a97..9caa05ad3 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaGatewayManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/HaGatewayManager.java @@ -32,10 +32,12 @@ public class HaGatewayManager private static final Logger log = Logger.get(HaGatewayManager.class); private final GatewayBackendDao dao; + private final boolean isOracleBackend; - public HaGatewayManager(Jdbi jdbi) + public HaGatewayManager(Jdbi jdbi, boolean isOracleBackend) { dao = requireNonNull(jdbi, "jdbi is null").onDemand(GatewayBackendDao.class); + this.isOracleBackend = isOracleBackend; } @Override @@ -48,7 +50,7 @@ public List getAllBackends() @Override public List getAllActiveBackends() { - List proxyBackendList = dao.findActiveBackend(); + List proxyBackendList = dao.findActiveBackend(!isOracleBackend); return upcast(proxyBackendList); } @@ -56,7 +58,7 @@ public List getAllActiveBackends() public List getActiveAdhocBackends() { try { - List proxyBackendList = dao.findActiveAdhocBackend(); + List proxyBackendList = dao.findActiveAdhocBackend(!isOracleBackend); return upcast(proxyBackendList); } catch (Exception e) { @@ -68,7 +70,7 @@ public List getActiveAdhocBackends() @Override public List getActiveBackends(String routingGroup) { - List proxyBackendList = dao.findActiveBackendByRoutingGroup(routingGroup); + List proxyBackendList = dao.findActiveBackendByRoutingGroup(routingGroup, !isOracleBackend); return upcast(proxyBackendList); } @@ -94,19 +96,19 @@ public void activateBackend(String backendName) @Override public ProxyBackendConfiguration addBackend(ProxyBackendConfiguration backend) { - dao.create(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive()); + dao.create(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive(), !isOracleBackend); return backend; } @Override public ProxyBackendConfiguration updateBackend(ProxyBackendConfiguration backend) { - GatewayBackend model = dao.findFirstByName(backend.getName()); + GatewayBackend model = dao.findFirstByName(backend.getName(), isOracleBackend); if (model == null) { - dao.create(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive()); + dao.create(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive(), !isOracleBackend); } else { - dao.update(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive()); + dao.update(backend.getName(), backend.getRoutingGroup(), backend.getProxyTo(), backend.getExternalUrl(), backend.isActive(), !isOracleBackend); } return backend; } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/IRoutingRulesManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/IRoutingRulesManager.java new file mode 100644 index 000000000..991bf8b71 --- /dev/null +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/IRoutingRulesManager.java @@ -0,0 +1,29 @@ +/* + * 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.persistence.dao.RoutingRule; + +import java.util.List; + +public interface IRoutingRulesManager +{ + List getRoutingRules(); + + List updateRoutingRule(RoutingRule routingRule); + + void deleteRoutingRule(String name); + + void createRoutingRule(RoutingRule routingRule); +} diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/MVELRoutingRule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/MVELRoutingRule.java index d115dd8a4..809ec5dce 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/MVELRoutingRule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/MVELRoutingRule.java @@ -59,6 +59,17 @@ public MVELRoutingRule( this.actions = actions.stream().map(this::compileExpressionIfNecessary).collect(toImmutableList()); } + public MVELRoutingRule(String name, String description, Integer priority, String condition, List actions) + { + initializeParserContext(parserContext); + + this.name = requireNonNull(name, "name is null"); + this.description = requireNonNullElse(description, ""); + this.priority = requireNonNullElse(priority, 0); + this.condition = requireNonNull(compileExpression(condition, parserContext)); + this.actions = actions.stream().map(a -> compileExpression(a, parserContext)).collect(toImmutableList()); + } + private Serializable compileExpressionIfNecessary(Serializable expression) { if (expression instanceof String stringExpression) { @@ -87,7 +98,7 @@ private void initializeParserContext(ParserContext parserContext) parserContext.addImport(String.class); parserContext.addImport(StringBuffer.class); parserContext.addImport(StringBuilder.class); - parserContext.addImport(FileBasedRoutingGroupSelector.class); + parserContext.addImport(RulesRoutingGroupSelector.class); } @Override diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java index 26a0291b9..6bdc256d7 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java @@ -14,8 +14,9 @@ package io.trino.gateway.ha.router; import io.airlift.http.client.HttpClient; -import io.airlift.units.Duration; +import io.trino.gateway.ha.config.HaGatewayConfiguration; import io.trino.gateway.ha.config.RequestAnalyzerConfig; +import io.trino.gateway.ha.config.RoutingRulesConfiguration; import io.trino.gateway.ha.config.RulesExternalConfiguration; import jakarta.servlet.http.HttpServletRequest; @@ -39,9 +40,11 @@ static RoutingGroupSelector byRoutingGroupHeader() * Routing group selector that uses routing engine rules * to determine the right routing group. */ - static RoutingGroupSelector byRoutingRulesEngine(String rulesConfigPath, Duration rulesRefreshPeriod, RequestAnalyzerConfig requestAnalyzerConfig) + static RoutingGroupSelector byRoutingRulesEngine(HaGatewayConfiguration configuration) { - return new FileBasedRoutingGroupSelector(rulesConfigPath, rulesRefreshPeriod, requestAnalyzerConfig); + IRoutingRulesManager routingRulesManager = new ForwardingRoutingRulesManager(configuration); + RoutingRulesConfiguration routingRulesConfig = configuration.getRoutingRules(); + return new RulesRoutingGroupSelector(routingRulesManager, routingRulesConfig.getRulesRefreshPeriod(), configuration.getRequestAnalyzerConfig()); } /** diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java index d9110aeff..204a3f7ef 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingRule.java @@ -23,4 +23,13 @@ public interface RoutingRule void evaluateAction(Map result, Map data, Map state); Integer getPriority(); + + static RoutingRule fromPersistedRoutingRule(io.trino.gateway.ha.persistence.dao.RoutingRule rule) + { + return switch (rule.routingRuleEngine()) { + case MVEL -> new MVELRoutingRule(rule.name(), rule.description(), rule.priority(), rule.condition(), rule.actions()); + // type will not be defined in legacy file based rules, but they are all MVEL + case null -> new MVELRoutingRule(rule.name(), rule.description(), rule.priority(), rule.condition(), rule.actions()); + }; + } } diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RulesRoutingGroupSelector.java similarity index 63% rename from gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java rename to gateway-ha/src/main/java/io/trino/gateway/ha/router/RulesRoutingGroupSelector.java index baff070d1..e4fc8460f 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/RulesRoutingGroupSelector.java @@ -13,9 +13,6 @@ */ 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.base.Supplier; import com.google.common.collect.ImmutableMap; import io.airlift.log.Logger; @@ -23,40 +20,35 @@ import io.trino.gateway.ha.config.RequestAnalyzerConfig; import jakarta.servlet.http.HttpServletRequest; -import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import static com.google.common.base.Suppliers.memoizeWithExpiration; -import static java.nio.charset.StandardCharsets.UTF_8; -import static java.util.Collections.sort; -public class FileBasedRoutingGroupSelector +public class RulesRoutingGroupSelector implements RoutingGroupSelector { - private static final Logger log = Logger.get(FileBasedRoutingGroupSelector.class); + private static final Logger log = Logger.get(RulesRoutingGroupSelector.class); public static final String RESULTS_ROUTING_GROUP_KEY = "routingGroup"; - private static final ObjectMapper yamlReader = new ObjectMapper(new YAMLFactory()); - private final Supplier> rules; + private final IRoutingRulesManager routingRulesManager; + private final boolean analyzeRequest; private final boolean clientsUseV2Format; private final int maxBodySize; private final TrinoRequestUser.TrinoRequestUserProvider trinoRequestUserProvider; - public FileBasedRoutingGroupSelector(String rulesPath, Duration rulesRefreshPeriod, RequestAnalyzerConfig requestAnalyzerConfig) + public RulesRoutingGroupSelector(IRoutingRulesManager routingRulesManager, Duration rulesRefreshPeriod, RequestAnalyzerConfig requestAnalyzerConfig) { analyzeRequest = requestAnalyzerConfig.isAnalyzeRequest(); clientsUseV2Format = requestAnalyzerConfig.isClientsUseV2Format(); maxBodySize = requestAnalyzerConfig.getMaxBodySize(); trinoRequestUserProvider = new TrinoRequestUser.TrinoRequestUserProvider(requestAnalyzerConfig); - - rules = memoizeWithExpiration(() -> readRulesFromPath(Path.of(rulesPath)), rulesRefreshPeriod.toJavaTime()); + this.routingRulesManager = routingRulesManager; + rules = memoizeWithExpiration(() -> this.routingRulesManager.getRoutingRules().stream().map(RoutingRule::fromPersistedRoutingRule).sorted().toList(), + rulesRefreshPeriod.toJavaTime()); } @Override @@ -86,22 +78,4 @@ public String findRoutingGroup(HttpServletRequest request) }); return result.get(RESULTS_ROUTING_GROUP_KEY); } - - public List readRulesFromPath(Path rulesPath) - { - try { - String content = Files.readString(rulesPath, UTF_8); - YAMLParser parser = new YAMLFactory().createParser(content); - List routingRulesList = new ArrayList<>(); - while (parser.nextToken() != null) { - MVELRoutingRule routingRules = yamlReader.readValue(parser, MVELRoutingRule.class); - routingRulesList.add(routingRules); - } - sort(routingRulesList); - return routingRulesList; - } - catch (IOException e) { - throw new RuntimeException("Failed to read or parse routing rules configuration from path: " + rulesPath, e); - } - } } diff --git a/gateway-ha/src/main/resources/mysql/V2__add_routing_rules.sql b/gateway-ha/src/main/resources/mysql/V2__add_routing_rules.sql new file mode 100644 index 000000000..edb3e5432 --- /dev/null +++ b/gateway-ha/src/main/resources/mysql/V2__add_routing_rules.sql @@ -0,0 +1,10 @@ +CREATE TABLE IF NOT EXISTS routing_rules ( + name varchar(128) primary key, + description varchar(256), + priority INT, + conditionExpression varchar(512), + actions varchar(1024), + routingRuleEngine varchar(128) +) + + diff --git a/gateway-ha/src/main/resources/oracle/V2__add_routing_rules.sql b/gateway-ha/src/main/resources/oracle/V2__add_routing_rules.sql new file mode 100644 index 000000000..faa7d807f --- /dev/null +++ b/gateway-ha/src/main/resources/oracle/V2__add_routing_rules.sql @@ -0,0 +1,9 @@ +CREATE TABLE routing_rules ( + name varchar(128), + description varchar(512), + priority int, + conditionExpression varchar(512), + actions varchar(1024), + routingRuleEngine varchar(128), + PRIMARY KEY (name) +); diff --git a/gateway-ha/src/main/resources/postgresql/V2__add_routing_rules.sql b/gateway-ha/src/main/resources/postgresql/V2__add_routing_rules.sql new file mode 100644 index 000000000..4c503b133 --- /dev/null +++ b/gateway-ha/src/main/resources/postgresql/V2__add_routing_rules.sql @@ -0,0 +1,8 @@ +CREATE TABLE IF NOT EXISTS routing_rules ( + name varchar primary key, + description varchar, + priority int, + conditionExpression varchar, + actions varchar[], + routingRuleEngine varchar +); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/BaseTestDbRoutingRulesMultipleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/BaseTestDbRoutingRulesMultipleBackend.java new file mode 100644 index 000000000..c602077af --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/BaseTestDbRoutingRulesMultipleBackend.java @@ -0,0 +1,236 @@ +/* + * 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; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableList; +import io.trino.gateway.ha.config.ProxyBackendConfiguration; +import io.trino.gateway.ha.persistence.dao.RoutingRule; +import io.trino.gateway.ha.persistence.dao.RoutingRuleEngine; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.testcontainers.containers.JdbcDatabaseContainer; +import org.testcontainers.containers.TrinoContainer; + +import java.io.File; +import java.io.IOException; +import java.net.URI; +import java.util.HashMap; +import java.util.List; + +import static java.util.Objects.requireNonNull; +import static org.assertj.core.api.Assertions.assertThat; +import static org.testcontainers.utility.MountableFile.forClasspathResource; + +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public abstract class BaseTestDbRoutingRulesMultipleBackend +{ + private static final MediaType MEDIA_TYPE = MediaType.parse("application/json; charset=utf-8"); + + private TrinoContainer adhocTrino; + private TrinoContainer systemTrino; + private final JdbcDatabaseContainer database; + private int adhocTrinoMappedPort; + private int systemTrinoMappedPort; + + final int routerPort = 20000 + (int) (Math.random() * 1000); + + private final OkHttpClient httpClient = new OkHttpClient(); + + public BaseTestDbRoutingRulesMultipleBackend(JdbcDatabaseContainer container) + { + this.database = requireNonNull(container, "container is null"); + this.database.start(); + } + + @BeforeAll + void setup() + throws Exception + { + adhocTrino = new TrinoContainer("trinodb/trino"); + adhocTrino.withCopyFileToContainer(forClasspathResource("trino-config.properties"), "/etc/trino/config.properties"); + adhocTrino.start(); + systemTrino = new TrinoContainer("trinodb/trino"); + systemTrino.withCopyFileToContainer(forClasspathResource("trino-config.properties"), "/etc/trino/config.properties"); + systemTrino.start(); + database.start(); + + adhocTrinoMappedPort = adhocTrino.getMappedPort(8080); + systemTrinoMappedPort = systemTrino.getMappedPort(8080); + + File testConfigFile = + HaGatewayTestUtils.buildGatewayConfig(database, routerPort, "test-config-with-db-routing-rules.yaml", database.getDriverClassName()); + + String[] args = {testConfigFile.getAbsolutePath()}; + HaGatewayLauncher.main(args); + + HaGatewayTestUtils.setUpBackend( + "trino1", "http://localhost:" + adhocTrinoMappedPort, "externalUrl", true, "adhoc", routerPort); + HaGatewayTestUtils.setUpBackend( + "trino2", "http://localhost:" + systemTrinoMappedPort, "externalUrl", true, "system", routerPort); + + getRules().forEach(routingRule -> applyRule(routingRule, Operation.CREATE)); + Thread.sleep(2000); //ensure rules are loaded + } + + private List getRules() + { + return ImmutableList.of(new RoutingRule( + "system-group", + "capture queries to system catalog", + 0, + "trinoQueryProperties.getCatalogs().contains(\"system\")", + ImmutableList.of("result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"system\")"), + RoutingRuleEngine.MVEL)); + } + + private void applyRule(RoutingRule routingRule, Operation operation) + { + try { + ObjectMapper objectMapper = new ObjectMapper(); + RequestBody requestBody = RequestBody.create(objectMapper.writeValueAsString(routingRule), MediaType.parse("application/json; charset=utf-8")); + + String path = switch (operation) { + case CREATE -> "/webapp/createRoutingRule"; + case UPDATE -> "/webapp/updateRoutingRules"; + case DELETE -> "/webapp/deleteRoutingRules"; + }; + + Request request = new Request.Builder() + .url("http://localhost:" + routerPort + path) + .addHeader("Content-Type", "application/json") + .post(requestBody) + .build(); + Response response = httpClient.newCall(request).execute(); + if (!response.isSuccessful()) { + throw new RuntimeException("Rule creation failed with response code " + response.code() + " and body " + response.body().string()); + } + } + catch (JsonProcessingException e) { + throw new RuntimeException("Could not serialize rule as JSON", e); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Test + void testQueryDeliveryToMultipleRoutingGroups() + throws Exception + { + // Default request should be routed to adhoc backend + submitQueryAndCheckDelivery("SELECT 1", adhocTrinoMappedPort); + + submitQueryAndCheckDelivery("SELECT * from system.runtime.nodes", systemTrinoMappedPort); + } + + @Test + void testRoutingWithUpdates() + throws Exception + { + submitQueryAndCheckDelivery("SELECT * from system.runtime.nodes", systemTrinoMappedPort); + + //update rule to only capture system.jdbc queries + applyRule(new RoutingRule( + "system-group", + "capture queries to system catalog", + 0, + "trinoQueryProperties.getCatalogSchemas().contains(\"system.jdbc\")", + ImmutableList.of("result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"system\")"), + RoutingRuleEngine.MVEL), Operation.UPDATE); + Thread.sleep(2000); + + submitQueryAndCheckDelivery("SELECT * from system.runtime.nodes", adhocTrinoMappedPort); + submitQueryAndCheckDelivery("SELECT * from system.jdbc.tables", systemTrinoMappedPort); + + Request deleteRequest = new Request.Builder() + .url("http://localhost:" + routerPort + "/webapp/deleteRoutingRule/system-group") + .delete() + .build(); + Response response = httpClient.newCall(deleteRequest).execute(); + assertThat(response.code()).isIn(200, 202, 204); + Thread.sleep(2000); + submitQueryAndCheckDelivery("SELECT * from system.jdbc.tables", adhocTrinoMappedPort); + + getRules().forEach(routingRule -> applyRule(routingRule, Operation.CREATE)); // Reset for other tests + Thread.sleep(2000); + submitQueryAndCheckDelivery("SELECT * from system.runtime.nodes", systemTrinoMappedPort); + } + + @Test + void testBackendConfiguration() + throws Exception + { + Request request = new Request.Builder() + .url("http://localhost:" + routerPort + "/entity/GATEWAY_BACKEND") + .method("GET", null) + .build(); + Response response = httpClient.newCall(request).execute(); + + final ObjectMapper objectMapper = new ObjectMapper(); + ProxyBackendConfiguration[] backendConfiguration = + objectMapper.readValue(response.body().string(), ProxyBackendConfiguration[].class); + + assertThat(backendConfiguration).isNotNull(); + assertThat(backendConfiguration).hasSize(2); + assertThat(backendConfiguration[0].isActive()).isTrue(); + assertThat(backendConfiguration[1].isActive()).isTrue(); + assertThat(backendConfiguration[0].getRoutingGroup()).isEqualTo("adhoc"); + assertThat(backendConfiguration[1].getRoutingGroup()).isEqualTo("system"); + assertThat(backendConfiguration[0].getExternalUrl()).isEqualTo("externalUrl"); + assertThat(backendConfiguration[1].getExternalUrl()).isEqualTo("externalUrl"); + } + + private void submitQueryAndCheckDelivery(String queryText, int expectedBackendPort) + throws Exception + { + // Intended for use with addXForwardedHeaders: false, so that the backend host is used for nextUri + ObjectMapper objectMapper = new ObjectMapper(); + RequestBody requestBody = RequestBody.create(queryText, MEDIA_TYPE); + Request request = + new Request.Builder() + .url("http://localhost:" + routerPort + "/v1/statement") + .addHeader("X-Trino-User", "test") + .post(requestBody) + .build(); + Response response = httpClient.newCall(request).execute(); + String responseBody = response.body().string(); + + HashMap results = objectMapper.readValue(responseBody, HashMap.class); + URI nextUri = URI.create(results.get("nextUri")); + assertThat(nextUri.getPort()).isEqualTo(expectedBackendPort); + } + + @AfterAll + void cleanup() + { + adhocTrino.stop(); + systemTrino.stop(); + } + + enum Operation + { + CREATE, + UPDATE, + DELETE + } +} diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java b/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java index 792021dc2..711589e75 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java @@ -27,6 +27,7 @@ import okhttp3.mockwebserver.MockWebServer; import org.jdbi.v3.core.Handle; import org.jdbi.v3.core.Jdbi; +import org.testcontainers.containers.JdbcDatabaseContainer; import org.testcontainers.containers.OracleContainer; import org.testcontainers.containers.PostgreSQLContainer; @@ -78,16 +79,17 @@ public static void seedRequiredData(String h2DbFilePath) } } - public static File buildGatewayConfig(PostgreSQLContainer postgreSqlContainer, int routerPort, String configFile) + public static File buildGatewayConfig(JdbcDatabaseContainer databaseContainer, int routerPort, String configFile, String driver) throws Exception { URL resource = HaGatewayTestUtils.class.getClassLoader().getResource("auth/localhost.jks"); String configStr = getResourceFileContent(configFile) .replace("REQUEST_ROUTER_PORT", String.valueOf(routerPort)) - .replace("POSTGRESQL_JDBC_URL", postgreSqlContainer.getJdbcUrl()) - .replace("POSTGRESQL_USER", postgreSqlContainer.getUsername()) - .replace("POSTGRESQL_PASSWORD", postgreSqlContainer.getPassword()) + .replace("POSTGRESQL_JDBC_URL", databaseContainer.getJdbcUrl()) + .replace("POSTGRESQL_USER", databaseContainer.getUsername()) + .replace("POSTGRESQL_PASSWORD", databaseContainer.getPassword()) + .replace("DRIVER_CLASS", driver) .replace( "APPLICATION_CONNECTOR_PORT", String.valueOf(30000 + (int) (Math.random() * 1000))) .replace("ADMIN_CONNECTOR_PORT", String.valueOf(31000 + (int) (Math.random() * 1000))) @@ -104,6 +106,12 @@ public static File buildGatewayConfig(PostgreSQLContainer postgreSqlContainer, i return target; } + public static File buildGatewayConfig(PostgreSQLContainer postgreSqlContainer, int routerPort, String configFile) + throws Exception + { + return buildGatewayConfig(postgreSqlContainer, routerPort, configFile, "org.postgresql.Driver"); + } + public static String getResourceFileContent(String fileName) { StringBuilder sb = new StringBuilder(); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestDbRoutingRulesMultipleBackendMySql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestDbRoutingRulesMultipleBackendMySql.java new file mode 100644 index 000000000..e0eecf18c --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestDbRoutingRulesMultipleBackendMySql.java @@ -0,0 +1,25 @@ +/* + * 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; + +import org.testcontainers.containers.MySQLContainer; + +public class TestDbRoutingRulesMultipleBackendMySql + extends BaseTestDbRoutingRulesMultipleBackend +{ + public TestDbRoutingRulesMultipleBackendMySql() + { + super(new MySQLContainer<>("mysql:5.7")); + } +} diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestDbRoutingRulesMultipleBackendPostgreSql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestDbRoutingRulesMultipleBackendPostgreSql.java new file mode 100644 index 000000000..d4af7fe24 --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestDbRoutingRulesMultipleBackendPostgreSql.java @@ -0,0 +1,25 @@ +/* + * 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; + +import org.testcontainers.containers.PostgreSQLContainer; + +public class TestDbRoutingRulesMultipleBackendPostgreSql + extends BaseTestDbRoutingRulesMultipleBackend +{ + public TestDbRoutingRulesMultipleBackendPostgreSql() + { + super(new PostgreSQLContainer<>("postgres:16")); + } +} diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithRoutingRulesSingleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithFileBasedRoutingRulesSingleBackend.java similarity index 98% rename from gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithRoutingRulesSingleBackend.java rename to gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithFileBasedRoutingRulesSingleBackend.java index 3bc3cec55..58475faaa 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithRoutingRulesSingleBackend.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithFileBasedRoutingRulesSingleBackend.java @@ -33,7 +33,7 @@ import static org.testcontainers.utility.MountableFile.forClasspathResource; @TestInstance(TestInstance.Lifecycle.PER_CLASS) -final class TestGatewayHaWithRoutingRulesSingleBackend +final class TestGatewayHaWithFileBasedRoutingRulesSingleBackend { private final OkHttpClient httpClient = new OkHttpClient(); private TrinoContainer trino; diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java index 9cdf80f5d..92dabb756 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java @@ -13,6 +13,7 @@ */ package io.trino.gateway.ha.persistence; +import io.airlift.log.Logger; import io.trino.gateway.ha.config.DataStoreConfiguration; import org.jdbi.v3.core.Handle; import org.jdbi.v3.core.Jdbi; @@ -41,6 +42,7 @@ public abstract class BaseTestDatabaseMigrations private final JdbcDatabaseContainer container; protected final String schema; protected final Jdbi jdbi; + Logger log = Logger.get(BaseTestDatabaseMigrations.class); public BaseTestDatabaseMigrations(JdbcDatabaseContainer container, String schema) { @@ -107,6 +109,7 @@ protected void verifyGatewaySchema(int expectedPropertiesCount) verifyResultSetCount("SELECT name FROM resource_groups", 0); verifyResultSetCount("SELECT user_regex FROM selectors", 0); verifyResultSetCount("SELECT environment FROM exact_match_source_selectors", 0); + verifyResultSetCount("SELECT name FROM routing_rules", 0); } protected void verifyResultSetCount(String sql, int expectedCount) @@ -125,9 +128,10 @@ protected void dropAllTables() String selectorsTable = "DROP TABLE IF EXISTS selectors"; String exactMatchTable = "DROP TABLE IF EXISTS exact_match_source_selectors"; String flywayHistoryTable = "DROP TABLE IF EXISTS flyway_schema_history"; + String routingRulesTable = "DROP TABLE IF EXISTS routing_rules"; Handle jdbiHandle = jdbi.open(); String sql = format("SELECT 1 FROM information_schema.tables WHERE table_schema = '%s'", schema); - verifyResultSetCount(sql, 7); + verifyResultSetCount(sql, 8); jdbiHandle.execute(gatewayBackendTable); jdbiHandle.execute(queryHistoryTable); jdbiHandle.execute(propertiesTable); @@ -135,6 +139,7 @@ protected void dropAllTables() jdbiHandle.execute(resourceGroupsTable); jdbiHandle.execute(exactMatchTable); jdbiHandle.execute(flywayHistoryTable); + jdbiHandle.execute(routingRulesTable); verifyResultSetCount(sql, 0); jdbiHandle.close(); } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsOracle.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsOracle.java index 67d8b8811..fa672c208 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsOracle.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsOracle.java @@ -48,10 +48,17 @@ protected void dropAllTables() * For this reason, if you remove the double quotes on flyway_schema_history, * you will get a table not found error. */ - List tables = ImmutableList.of("gateway_backend", "query_history", "resource_groups_global_properties", "selectors", "resource_groups", "exact_match_source_selectors", "\"flyway_schema_history\""); + List tables = ImmutableList.of( + "gateway_backend", + "query_history", + "resource_groups_global_properties", + "selectors", "resource_groups", + "exact_match_source_selectors", + "routing_rules", + "\"flyway_schema_history\""); Handle jdbiHandle = jdbi.open(); String sql = format("SELECT 1 FROM all_tables WHERE owner = '%s'", schema); - verifyResultSetCount(sql, 7); + verifyResultSetCount(sql, 8); tables.forEach(table -> jdbiHandle.execute("DROP TABLE " + table)); verifyResultSetCount(sql, 0); jdbiHandle.close(); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgreSql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgreSql.java index 7bd51488f..c9c06bee6 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgreSql.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgreSql.java @@ -83,6 +83,17 @@ protected void createGatewaySchema() " PRIMARY KEY (environment, source, query_type),\n" + " UNIQUE (source, environment, query_type, resource_group_id)\n" + ");"; + + String routingRulesTable = """ + CREATE TABLE routing_rules ( + name VARCHAR PRIMARY KEY, + description VARCHAR, + priority INT, + condition VARCHAR, + actions VARCHAR[], + routingRuleEngine VARCHAR) + """; + Handle jdbiHandle = jdbi.open(); jdbiHandle.execute(gatewayBackendTable); jdbiHandle.execute(queryHistoryTable); @@ -90,6 +101,7 @@ protected void createGatewaySchema() jdbiHandle.execute(resourceGroupsTable); jdbiHandle.execute(selectorsTable); jdbiHandle.execute(exactMatchSourceSelectorsTable); + jdbiHandle.execute(routingRulesTable); jdbiHandle.close(); } } 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/TestFileBasedRoutingRulesManager.java similarity index 70% rename from gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingRulesManager.java rename to gateway-ha/src/test/java/io/trino/gateway/ha/router/TestFileBasedRoutingRulesManager.java index 668374c17..faa78c85d 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/TestFileBasedRoutingRulesManager.java @@ -15,7 +15,7 @@ import io.trino.gateway.ha.config.HaGatewayConfiguration; import io.trino.gateway.ha.config.RoutingRulesConfiguration; -import io.trino.gateway.ha.domain.RoutingRule; +import io.trino.gateway.ha.persistence.dao.RoutingRule; import org.junit.jupiter.api.Test; import java.io.IOException; @@ -28,7 +28,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; -final class TestRoutingRulesManager +final class TestFileBasedRoutingRulesManager { @Test void testGetRoutingRules() @@ -39,7 +39,7 @@ void testGetRoutingRules() String rulesConfigPath = "src/test/resources/rules/routing_rules_atomic.yml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); configuration.setRoutingRules(routingRulesConfiguration); - RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); + FileBasedRoutingRulesManager routingRulesManager = new FileBasedRoutingRulesManager(configuration); List result = routingRulesManager.getRoutingRules(); @@ -49,15 +49,17 @@ void testGetRoutingRules() "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())")); + "request.getHeader(\"X-Trino-Source\") == \"airflow\" && (request.getHeader(\"X-Trino-Client-Tags\") == null || request.getHeader(\"X-Trino-Client-Tags\").isEmpty())", + List.of("result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")"), + null)); //Legacy file based rules will not have an engine assertThat(result.get(1)).isEqualTo( new RoutingRule( "airflow special", "if query from airflow with special label, route to etl-special group", null, - List.of("result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")"), - "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\"")); + "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\"", + List.of("result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")"), + null)); } @Test @@ -68,7 +70,7 @@ void testRoutingRulesNoSuchFileException() String rulesConfigPath = "src/test/resources/rules/routing_rules_test.yaml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); configuration.setRoutingRules(routingRulesConfiguration); - RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); + FileBasedRoutingRulesManager routingRulesManager = new FileBasedRoutingRulesManager(configuration); assertThatThrownBy(routingRulesManager::getRoutingRules).hasRootCauseInstanceOf(NoSuchFileException.class); } @@ -82,15 +84,27 @@ void testUpdateRoutingRulesFile() String rulesConfigPath = "src/test/resources/rules/routing_rules_update.yml"; routingRulesConfiguration.setRulesConfigPath(rulesConfigPath); configuration.setRoutingRules(routingRulesConfiguration); - RoutingRulesManager routingRulesManager = new RoutingRulesManager(configuration); + FileBasedRoutingRulesManager routingRulesManager = new FileBasedRoutingRulesManager(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\""); + RoutingRule routingRules = new RoutingRule( + "airflow", + "if query from airflow, route to etl group", + 0, + "request.getHeader(\"X-Trino-Source\") == \"JDBC\"", + List.of("result.put(\"routingGroup\", \"adhoc\")"), + null); 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\""); + RoutingRule originalRoutingRules = new RoutingRule( + "airflow", + "if query from airflow, route to etl group", + 0, + "request.getHeader(\"X-Trino-Source\") == \"airflow\"", + List.of("result.put(\"routingGroup\", \"etl\")"), + null); List updateRoutingRules = routingRulesManager.updateRoutingRule(originalRoutingRules); assertThat(updateRoutingRules).hasSize(2); @@ -106,8 +120,14 @@ void testUpdateRoutingRulesNoSuchFileException() 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\""); + FileBasedRoutingRulesManager routingRulesManager = new FileBasedRoutingRulesManager(configuration); + RoutingRule routingRules = new RoutingRule( + "airflow", + "if query from airflow, route to etl group", + 0, + "request.getHeader(\"X-Trino-Source\") == \"JDBC\"", + List.of("result.put(\"routingGroup\", \"adhoc\")"), + null); assertThatThrownBy(() -> routingRulesManager.updateRoutingRule(routingRules)).hasRootCauseInstanceOf(NoSuchFileException.class); } @@ -121,10 +141,22 @@ void testConcurrentUpdateRoutingRule() 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\""); + FileBasedRoutingRulesManager routingRulesManager = new FileBasedRoutingRulesManager(configuration); + + RoutingRule routingRule1 = new RoutingRule( + "airflow", + "if query from airflow, route to etl group", + 0, + "request.getHeader(\"X-Trino-Source\") == \"airflow\"", + List.of("result.put(\"routingGroup\", \"etl\")"), + null); + RoutingRule routingRule2 = new RoutingRule( + "airflow", + "if query from airflow, route to adhoc group", + 0, + "request.getHeader(\"X-Trino-Source\") == \"datagrip\"", + List.of("result.put(\"routingGroup\", \"adhoc\")"), + null); ExecutorService executorService = Executors.newFixedThreadPool(2); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java index 191ed4a16..b6c3a3fc0 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestHaGatewayManager.java @@ -32,7 +32,7 @@ final class TestHaGatewayManager void setUp() { JdbcConnectionManager connectionManager = createTestingJdbcConnectionManager(); - haGatewayManager = new HaGatewayManager(connectionManager.getJdbi()); + haGatewayManager = new HaGatewayManager(connectionManager.getJdbi(), false); } @Test 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 index 293e77bdf..05ef86332 100644 --- 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 @@ -19,7 +19,8 @@ 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 io.trino.gateway.ha.persistence.dao.RoutingRule; +import io.trino.gateway.ha.persistence.dao.RoutingRuleEngine; import okhttp3.MediaType; import okhttp3.OkHttpClient; import okhttp3.Request; @@ -32,6 +33,7 @@ import org.testcontainers.containers.TrinoContainer; import java.io.File; +import java.io.IOException; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @@ -101,7 +103,12 @@ 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\""); + RoutingRule updatedRoutingRules = new RoutingRule( + "airflow", + "if query from airflow, route to adhoc group", + 0, "request.getHeader(\"X-Trino-Source\") == \"JDBC\"", + List.of("result.put(\"routingGroup\", \"adhoc\")"), + RoutingRuleEngine.MVEL); RequestBody requestBody = RequestBody.create(OBJECT_MAPPER.writeValueAsString(updatedRoutingRules), MediaType.parse("application/json; charset=utf-8")); Request request = new Request.Builder() .url("http://localhost:" + routerPort + "/webapp/updateRoutingRules") @@ -113,19 +120,7 @@ void testUpdateRoutingRulesAPI() 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(); - JsonNode rootNode = OBJECT_MAPPER.readTree(responseBody); - JsonNode dataNode = rootNode.path("data"); - - JsonCodec responseCodec = JsonCodec.jsonCodec(RoutingRule[].class); - RoutingRule[] routingRules = responseCodec.fromJson(dataNode.toString()); + RoutingRule[] routingRules = getRoutingRules(); assertThat(response.code()).isEqualTo(200); assertThat(routingRules[0].name()).isEqualTo("airflow"); @@ -135,7 +130,12 @@ void testUpdateRoutingRulesAPI() 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\""); + RoutingRule revertRoutingRules = new RoutingRule( + "airflow", + "if query from airflow, route to etl group", + 0, "request.getHeader(\"X-Trino-Source\") == \"airflow\"", + List.of("result.put(\"routingGroup\", \"etl\")"), + null); RequestBody requestBody3 = RequestBody.create(OBJECT_MAPPER.writeValueAsString(revertRoutingRules), MediaType.parse("application/json; charset=utf-8")); Request request3 = new Request.Builder() .url("http://localhost:" + routerPort + "/webapp/updateRoutingRules") @@ -145,6 +145,60 @@ void testUpdateRoutingRulesAPI() httpClient.newCall(request3).execute(); } + private RoutingRule[] getRoutingRules() + throws IOException + { + Request request2 = + new Request.Builder() + .url("http://localhost:" + routerPort + "/webapp/getRoutingRules") + .get() + .build(); + return decodeRoutingRulesResponse(httpClient.newCall(request2).execute()); + } + + private RoutingRule[] decodeRoutingRulesResponse(Response response) + throws IOException + { + String responseBody = response.body().string(); + JsonNode rootNode = OBJECT_MAPPER.readTree(responseBody); + JsonNode dataNode = rootNode.path("data"); + + JsonCodec responseCodec = JsonCodec.jsonCodec(RoutingRule[].class); + return responseCodec.fromJson(dataNode.toString()); + } + + @Test + void testRoutingRuleAPICreateDelete() + throws Exception + { + //Update routing rules with a new rule + RoutingRule newRoutingRule = new RoutingRule( + "new-rule", + "newly created rule", + 0, + "false", + List.of("result.put(\"routingGroup\", \"adhoc\")"), + RoutingRuleEngine.MVEL); + RequestBody createRequestBody = RequestBody.create(OBJECT_MAPPER.writeValueAsString(newRoutingRule), MediaType.parse("application/json; charset=utf-8")); + Request createRequest = new Request.Builder() + .url("http://localhost:" + routerPort + "/webapp/createRoutingRule") + .addHeader("Content-Type", "application/json") + .post(createRequestBody) + .build(); + httpClient.newCall(createRequest).execute(); + + RoutingRule[] rules = getRoutingRules(); + assertThat(rules).contains(newRoutingRule); + + Request deleteRequest = new Request.Builder() + .url("http://localhost:" + routerPort + "/webapp/deleteRoutingRule/new-rule") + .delete() + .build(); + httpClient.newCall(deleteRequest).execute(); + rules = getRoutingRules(); + assertThat(rules).doesNotContain(newRoutingRule); + } + @Test void testUIConfigurationAPI() throws Exception diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java index c0091cc7e..c3fb0603a 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java @@ -15,7 +15,9 @@ import com.google.common.collect.ImmutableSet; import io.airlift.units.Duration; +import io.trino.gateway.ha.config.HaGatewayConfiguration; import io.trino.gateway.ha.config.RequestAnalyzerConfig; +import io.trino.gateway.ha.config.RoutingRulesConfiguration; import io.trino.sql.tree.QualifiedName; import jakarta.servlet.http.HttpServletRequest; import jakarta.ws.rs.HttpMethod; @@ -67,6 +69,18 @@ void initialize() requestAnalyzerConfig.setAnalyzeRequest(true); } + private HaGatewayConfiguration getHaGatewayConfiguration(String rulesPath, Duration refreshPeriod) + { + HaGatewayConfiguration haGatewayConfiguration = new HaGatewayConfiguration(); + RoutingRulesConfiguration routingRulesConfiguration = new RoutingRulesConfiguration(); + routingRulesConfiguration.setRulesConfigPath(rulesPath); + routingRulesConfiguration.setRulesRefreshPeriod(refreshPeriod); + haGatewayConfiguration.setRoutingRules(routingRulesConfiguration); + haGatewayConfiguration.setRequestAnalyzerConfig(requestAnalyzerConfig); + + return haGatewayConfiguration; + } + static Stream provideRoutingRuleConfigFiles() { String rulesDir = "src/test/resources/rules/"; @@ -97,7 +111,7 @@ void testByRoutingGroupHeader() void testByRoutingRulesEngine(String rulesConfigPath) { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration(rulesConfigPath, oneHourRefreshPeriod)); HttpServletRequest mockRequest = prepareMockRequest(); @@ -110,10 +124,9 @@ void testByRoutingRulesEngine(String rulesConfigPath) void testGetUserFromBasicAuth() { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration( "src/test/resources/rules/routing_rules_trino_query_properties.yml", - oneHourRefreshPeriod, - requestAnalyzerConfig); + oneHourRefreshPeriod)); String encodedUsernamePassword = Base64.getEncoder().encodeToString("will:supersecret".getBytes(UTF_8)); HttpServletRequest mockRequest = prepareMockRequest(); @@ -128,10 +141,9 @@ void testTrinoQueryPropertiesQueryDetails() throws IOException { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration( "src/test/resources/rules/routing_rules_trino_query_properties.yml", - oneHourRefreshPeriod, - requestAnalyzerConfig); + oneHourRefreshPeriod)); String query = "SELECT x.*, y.*, z.* FROM catx.schemx.tblx x, schemy.tbly y, tblz z"; Reader reader = new StringReader(query); BufferedReader bufferedReader = new BufferedReader(reader); @@ -148,10 +160,9 @@ void testTrinoQueryPropertiesCatalogSchemas() throws IOException { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration( "src/test/resources/rules/routing_rules_trino_query_properties.yml", - oneHourRefreshPeriod, - requestAnalyzerConfig); + oneHourRefreshPeriod)); String query = "SELECT x.*, y.* FROM catx.nondefault.tblx x, caty.default.tbly y"; Reader reader = new StringReader(query); BufferedReader bufferedReader = new BufferedReader(reader); @@ -167,10 +178,9 @@ void testTrinoQueryPropertiesCatalogSchemas() void testTrinoQueryPropertiesSessionDefaults() { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration( "src/test/resources/rules/routing_rules_trino_query_properties.yml", - oneHourRefreshPeriod, - requestAnalyzerConfig); + oneHourRefreshPeriod)); HttpServletRequest mockRequest = prepareMockRequest(); when(mockRequest.getHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME)).thenReturn("other_catalog"); @@ -184,10 +194,9 @@ void testTrinoQueryPropertiesQueryType() throws IOException { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration( "src/test/resources/rules/routing_rules_trino_query_properties.yml", - oneHourRefreshPeriod, - requestAnalyzerConfig); + oneHourRefreshPeriod)); String query = "INSERT INTO foo SELECT 1"; Reader reader = new StringReader(query); BufferedReader bufferedReader = new BufferedReader(reader); @@ -202,10 +211,9 @@ void testTrinoQueryPropertiesResourceGroupQueryType() throws IOException { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration( "src/test/resources/rules/routing_rules_trino_query_properties.yml", - oneHourRefreshPeriod, - requestAnalyzerConfig); + oneHourRefreshPeriod)); HttpServletRequest mockRequest = prepareMockRequest(); when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader("CREATE TABLE cat.schem.foo (c1 int)"))); @@ -218,10 +226,9 @@ void testTrinoQueryPropertiesAlternateStatementFormat() { requestAnalyzerConfig.setClientsUseV2Format(true); RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration( "src/test/resources/rules/routing_rules_trino_query_properties.yml", - oneHourRefreshPeriod, - requestAnalyzerConfig); + oneHourRefreshPeriod)); String body = "{\"preparedStatements\" : {\"statement1\":\"INSERT INTO foo SELECT 1\"}, \"query\": \"EXECUTE statement1\"}"; Reader reader = new StringReader(body); BufferedReader bufferedReader = new BufferedReader(reader); @@ -239,10 +246,9 @@ void testTrinoQueryPropertiesPreparedStatementInHeader() String body = "EXECUTE statement4"; RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine( + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration( "src/test/resources/rules/routing_rules_trino_query_properties.yml", - oneHourRefreshPeriod, - requestAnalyzerConfig); + oneHourRefreshPeriod)); Reader reader = new StringReader(body); BufferedReader bufferedReader = new BufferedReader(reader); HttpServletRequest mockRequest = prepareMockRequest(); @@ -260,7 +266,7 @@ void testTrinoQueryPropertiesPreparedStatementInHeader() void testByRoutingRulesEngineSpecialLabel(String rulesConfigPath) { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration(rulesConfigPath, oneHourRefreshPeriod)); HttpServletRequest mockRequest = prepareMockRequest(); @@ -276,7 +282,7 @@ void testByRoutingRulesEngineSpecialLabel(String rulesConfigPath) void testByRoutingRulesEngineNoMatch(String rulesConfigPath) { RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine(rulesConfigPath, oneHourRefreshPeriod, requestAnalyzerConfig); + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration(rulesConfigPath, oneHourRefreshPeriod)); HttpServletRequest mockRequest = prepareMockRequest(); // even though special label is present, query is not from airflow. @@ -304,7 +310,7 @@ void testByRoutingRulesEngineFileChange() Duration refreshPeriod = new Duration(1, MILLISECONDS); RoutingGroupSelector routingGroupSelector = - RoutingGroupSelector.byRoutingRulesEngine(file.getPath(), refreshPeriod, requestAnalyzerConfig); + RoutingGroupSelector.byRoutingRulesEngine(getHaGatewayConfiguration(file.getPath(), refreshPeriod)); HttpServletRequest mockRequest = prepareMockRequest(); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestStochasticRoutingManager.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestStochasticRoutingManager.java index bf8b2a287..689080e17 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestStochasticRoutingManager.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestStochasticRoutingManager.java @@ -35,7 +35,7 @@ final class TestStochasticRoutingManager void setUp() { JdbcConnectionManager connectionManager = createTestingJdbcConnectionManager(); - backendManager = new HaGatewayManager(connectionManager.getJdbi()); + backendManager = new HaGatewayManager(connectionManager.getJdbi(), false); historyManager = new HaQueryHistoryManager(connectionManager.getJdbi(), false); haRoutingManager = new StochasticRoutingManager(backendManager, historyManager); } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/util/TestDbRoutingRulesMultipleBackendsOracle.java b/gateway-ha/src/test/java/io/trino/gateway/ha/util/TestDbRoutingRulesMultipleBackendsOracle.java new file mode 100644 index 000000000..7a294243b --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/util/TestDbRoutingRulesMultipleBackendsOracle.java @@ -0,0 +1,27 @@ +/* + * 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.util; + +import io.trino.gateway.ha.BaseTestDbRoutingRulesMultipleBackend; + +import static io.trino.gateway.ha.HaGatewayTestUtils.getOracleContainer; + +public class TestDbRoutingRulesMultipleBackendsOracle + extends BaseTestDbRoutingRulesMultipleBackend +{ + public TestDbRoutingRulesMultipleBackendsOracle() + { + super(getOracleContainer()); + } +} diff --git a/gateway-ha/src/test/resources/rules/routing_rules_atomic.yml b/gateway-ha/src/test/resources/rules/routing_rules_atomic.yml index 89a5a09da..da79d5dd8 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_atomic.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_atomic.yml @@ -3,10 +3,10 @@ name: "airflow" description: "if query from airflow, route to etl group" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && (request.getHeader(\"X-Trino-Client-Tags\") == null || request.getHeader(\"X-Trino-Client-Tags\").isEmpty())" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")" --- name: "airflow special" description: "if query from airflow with special label, route to etl-special group" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\"" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml b/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml index f492dfeaa..850a43fca 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_concurrent.yml @@ -2,6 +2,7 @@ name: "airflow" description: "if query from airflow, route to adhoc group" priority: 0 +condition: "request.getHeader(\"X-Trino-Source\") == \"datagrip\"" actions: - "result.put(\"routingGroup\", \"adhoc\")" -condition: "request.getHeader(\"X-Trino-Source\") == \"datagrip\"" +routingRuleEngine: null diff --git a/gateway-ha/src/test/resources/rules/routing_rules_if_statements.yml b/gateway-ha/src/test/resources/rules/routing_rules_if_statements.yml index e0cf9bdbb..ed9082556 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_if_statements.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_if_statements.yml @@ -4,8 +4,8 @@ description: "if query from airflow, route to etl group" condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" actions: - "if (request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\") { - result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\") + result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\") } else { - result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\") + result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\") }" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_priorities.yml b/gateway-ha/src/test/resources/rules/routing_rules_priorities.yml index 0ca00f1c1..974e23cc0 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_priorities.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_priorities.yml @@ -4,11 +4,11 @@ description: "if query from airflow, route to etl group" priority: 0 condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl\")" --- name: "airflow special" description: "if query from airflow with special label, route to etl-special group" priority: 1 condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\" && request.getHeader(\"X-Trino-Client-Tags\") contains \"label=special\"" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_state.yml b/gateway-ha/src/test/resources/rules/routing_rules_state.yml index 53650058b..d01f33c60 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_state.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_state.yml @@ -17,7 +17,7 @@ condition: | request.getHeader("X-Trino-Source") == "airflow" actions: - | - result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, "etl") + result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, "etl") - | state.get("triggeredRules").add("airflow") --- @@ -28,4 +28,4 @@ condition: | state.get("triggeredRules").contains("airflow") && request.getHeader("X-Trino-Client-Tags") contains "label=special" actions: - | - result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, "etl-special") + result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, "etl-special") diff --git a/gateway-ha/src/test/resources/rules/routing_rules_trino_query_properties.yml b/gateway-ha/src/test/resources/rules/routing_rules_trino_query_properties.yml index b76c65769..477c9c906 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_trino_query_properties.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_trino_query_properties.yml @@ -3,7 +3,7 @@ name: "user" description: "if user is will, route to will-group" condition: "trinoRequestUser.userExistsAndEquals(\"will\")" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"will-group\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"will-group\")" --- name: "query" description: "test extraction of tables and schemas in conjunction with default catalog and schema" @@ -14,7 +14,7 @@ condition: | && trinoQueryProperties.getSchemas().contains("schemy") && trinoQueryProperties.getCatalogs().contains("catx") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"tbl-group\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"tbl-group\")" --- name: "catalog-schema" description: "test that catalogSchemas with default catalog and schema" @@ -23,14 +23,14 @@ condition: | && trinoQueryProperties.getCatalogSchemas.contains("caty.default") && !trinoQueryProperties.getCatalogSchemas.contains("catx.default") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"catalog-schema-group\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"catalog-schema-group\")" --- name: "query-type" description: "test table type" condition: | trinoQueryProperties.getQueryType().toLowerCase.equals("insert") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"type-group\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"type-group\")" --- name: "resource-group-query-type" description: "test table type" @@ -44,7 +44,7 @@ description: "test execute with multiple prepared statements" condition: | trinoQueryProperties.getQueryType().toLowerCase.equals("query") && trinoQueryProperties.tablesContains("cat.schem.foo") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"statement-header-group\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"statement-header-group\")" --- name: "defaults-group" description: "test default schema and catalog" @@ -53,14 +53,14 @@ condition: | && trinoQueryProperties.getDefaultSchema().equals(java.util.Optional.of("other_schema")) actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"defaults-group\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"defaults-group\")" --- name: "system-group" description: "capture queries to system catalog" condition: | trinoQueryProperties.getCatalogs().contains("system") actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"system\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"system\")" --- name: "nomatch" @@ -68,4 +68,4 @@ priority: -1 description: "default group to catch if no other rules fired" condition: "true" actions: - - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"no-match\")" + - "result.put(RulesRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"no-match\")" diff --git a/gateway-ha/src/test/resources/rules/routing_rules_update.yml b/gateway-ha/src/test/resources/rules/routing_rules_update.yml index 3478bc45e..27cdbf85e 100644 --- a/gateway-ha/src/test/resources/rules/routing_rules_update.yml +++ b/gateway-ha/src/test/resources/rules/routing_rules_update.yml @@ -2,14 +2,16 @@ name: "airflow" description: "if query from airflow, route to etl group" priority: 0 +condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" actions: - "result.put(\"routingGroup\", \"etl\")" -condition: "request.getHeader(\"X-Trino-Source\") == \"airflow\"" +routingRuleEngine: null --- 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\"" +actions: +- "result.put(\"routingGroup\", \"etl-special\")" +routingRuleEngine: null diff --git a/gateway-ha/src/test/resources/test-config-with-db-routing-rules.yaml b/gateway-ha/src/test/resources/test-config-with-db-routing-rules.yaml new file mode 100644 index 000000000..8debc9177 --- /dev/null +++ b/gateway-ha/src/test/resources/test-config-with-db-routing-rules.yaml @@ -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: DRIVER_CLASS + +clusterStatsConfiguration: + monitorType: INFO_API + +monitor: + taskDelaySeconds: 1 + +extraWhitelistPaths: + - '/v1/custom.*' + - '/custom/logout.*' + +gatewayCookieConfiguration: + enabled: true + cookieSigningSecret: "kjlhbfrewbyuo452cds3dc1234ancdsjh" + +oauth2GatewayCookieConfiguration: + deletePaths: + - "/custom/logout" + +requestAnalyzerConfig: + analyzeRequest: true + +routingRules: + rulesEngineEnabled: true + rulesType: "DB" + rulesRefreshPeriod: "1s" + +routing: + addXForwardedHeaders: false