From 72b1330f02f91143084716174b8265b645d39773 Mon Sep 17 00:00:00 2001 From: Will Morrison Date: Thu, 23 Jan 2025 10:39:08 -0500 Subject: [PATCH] Use RESULTS_ROUTING_GROUP_KEY in routing rules --- .../resources/rules/routing_rules_atomic.yml | 4 ++-- .../rules/routing_rules_if_statements.yml | 4 ++-- .../resources/rules/routing_rules_priorities.yml | 4 ++-- .../test/resources/rules/routing_rules_state.yml | 4 ++-- .../routing_rules_trino_query_properties.yml | 16 ++++++++-------- 5 files changed, 16 insertions(+), 16 deletions(-) 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 91f8ea56c..89a5a09da 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(\"routingGroup\", \"etl\")" + - "result.put(FileBasedRoutingGroupSelector.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(\"routingGroup\", \"etl-special\")" + - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\")" 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 7379a49be..e0cf9bdbb 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(\"routingGroup\", \"etl-special\") + result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"etl-special\") } else { - result.put(\"routingGroup\", \"etl\") + result.put(FileBasedRoutingGroupSelector.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 e075b9381..0ca00f1c1 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(\"routingGroup\", \"etl\")" + - "result.put(FileBasedRoutingGroupSelector.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(\"routingGroup\", \"etl-special\")" + - "result.put(FileBasedRoutingGroupSelector.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 463ebd84f..53650058b 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("routingGroup", "etl") + result.put(FileBasedRoutingGroupSelector.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("routingGroup", "etl-special") + result.put(FileBasedRoutingGroupSelector.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 eadd58bb0..834aa008e 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(\"routingGroup\", \"will-group\")" + - "result.put(FileBasedRoutingGroupSelector.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(\"routingGroup\", \"tbl-group\")" + - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"tbl-group\")" --- name: "catalog-schema" description: "test that catalogSchemas with default catalog and schema" @@ -23,21 +23,21 @@ condition: | && trinoQueryProperties.getCatalogSchemas.contains("caty.default") && !trinoQueryProperties.getCatalogSchemas.contains("catx.default") actions: - - "result.put(\"routingGroup\", \"catalog-schema-group\")" + - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"catalog-schema-group\")" --- name: "query-type" description: "test table type" condition: | trinoQueryProperties.getQueryType().toLowerCase.equals("insert") actions: - - "result.put(\"routingGroup\", \"type-group\")" + - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"type-group\")" --- name: "prepared-statement-header" description: "test execute with multiple prepared statements" condition: | trinoQueryProperties.getQueryType().toLowerCase.equals("query") && trinoQueryProperties.tablesContains("cat.schem.foo") actions: - - "result.put(\"routingGroup\", \"statement-header-group\")" + - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"statement-header-group\")" --- name: "defaults-group" description: "test default schema and catalog" @@ -46,14 +46,14 @@ condition: | && trinoQueryProperties.getDefaultSchema().equals(java.util.Optional.of("other_schema")) actions: - - "result.put(\"routingGroup\", \"defaults-group\")" + - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"defaults-group\")" --- name: "system-group" description: "capture queries to system catalog" condition: | trinoQueryProperties.getCatalogs().contains("system") actions: - - "result.put(\"routingGroup\", \"system\")" + - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"system\")" --- name: "nomatch" @@ -61,4 +61,4 @@ priority: -1 description: "default group to catch if no other rules fired" condition: "true" actions: - - "result.put(\"routingGroup\", \"no-match\")" + - "result.put(FileBasedRoutingGroupSelector.RESULTS_ROUTING_GROUP_KEY, \"no-match\")"