Fix IndexOutOfBoundsException caused by switch-case statements

willmostly · Sep 25, 2024 · e86bedf · e86bedf
1 parent 9b5c57a
commit e86bedf
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 5 deletions.
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java
@@ -357,16 +357,17 @@ private void setCatalogAndSchemaNameFromSchemaQualifiedName(
         else {
             QualifiedName schema = schemaOptional.orElseThrow();
             switch (schema.getParts().size()) {
-                case 1:
+                case 1 -> {
                     schemaBuilder.add(schema.getParts().getFirst());
                     catalogBuilder.add(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier));
                     catalogSchemaBuilder.add(format("%s.%s", defaultCatalog, schema.getParts().getFirst()));
-                case 2:
+                }
+                case 2 -> {
                     schemaBuilder.add(schema.getParts().get(1));
                     catalogBuilder.add(schema.getParts().getFirst());
                     catalogSchemaBuilder.add(format("%s.%s", schema.getParts().getFirst(), schema.getParts().getLast()));
-                default:
-                    log.error("Schema has >2 parts: " + schema);
+                }
+                default -> log.error("Schema has >2 parts: %s", schema);
             }
         }
     }

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
@@ -50,6 +50,10 @@ public class TestRoutingGroupSelector
 {
     public static final String TRINO_SOURCE_HEADER = "X-Trino-Source";
     public static final String TRINO_CLIENT_TAGS_HEADER = "X-Trino-Client-Tags";
+
+    private static final String DEFAULT_CATALOG = "default_catalog";
+    private static final String DEFAULT_SCHEMA = "default_schema";
+
     RequestAnalyzerConfig requestAnalyzerConfig = new RequestAnalyzerConfig();
 
     @BeforeAll
@@ -301,6 +305,7 @@ private Stream<Arguments> provideTableExtractionQueries()
                         ImmutableSet.of("schem", "mvschem"),
                         ImmutableSet.of(QualifiedName.of("cat", "schem", "tbl"), QualifiedName.of("cat", "mvschem", "mv"))),
                 Arguments.of("CREATE SCHEMA kat.schem", ImmutableSet.of("kat"), ImmutableSet.of("schem"), ImmutableSet.of()),
+                Arguments.of("CREATE SCHEMA schem", ImmutableSet.of(DEFAULT_CATALOG), ImmutableSet.of("schem"), ImmutableSet.of()),
                 Arguments.of("CREATE TABLE cat.schem.tbl(c1 varchar)",
                         ImmutableSet.of("cat"),
                         ImmutableSet.of("schem"),
@@ -315,6 +320,7 @@ private Stream<Arguments> provideTableExtractionQueries()
                         ImmutableSet.of(QualifiedName.of("cat", "schem", "tbl"), QualifiedName.of("cat", "schem2", "tbl2"))),
                 Arguments.of("DROP CATALOG kat", ImmutableSet.of("kat"), ImmutableSet.of(), ImmutableSet.of()),
                 Arguments.of("DROP SCHEMA kat.schem", ImmutableSet.of("kat"), ImmutableSet.of("schem"), ImmutableSet.of()),
+                Arguments.of("DROP SCHEMA schem", ImmutableSet.of(DEFAULT_CATALOG), ImmutableSet.of("schem"), ImmutableSet.of()),
                 Arguments.of("DROP TABLE cat.schem.tbl",
                         ImmutableSet.of("cat"),
                         ImmutableSet.of("schem"),
@@ -344,6 +350,10 @@ private Stream<Arguments> provideTableExtractionQueries()
                         ImmutableSet.of("cat"),
                         ImmutableSet.of("schem"),
                         ImmutableSet.of()),
+                Arguments.of("SHOW CREATE SCHEMA schem",
+                        ImmutableSet.of(DEFAULT_CATALOG),
+                        ImmutableSet.of("schem"),
+                        ImmutableSet.of()),
                 Arguments.of("SHOW CREATE TABLE cat.schem.tbl",
                         ImmutableSet.of("cat"),
                         ImmutableSet.of("schem"),
@@ -353,8 +363,11 @@ private Stream<Arguments> provideTableExtractionQueries()
                         ImmutableSet.of("schem"),
                         ImmutableSet.of(QualifiedName.of("cat", "schem", "vw"))),
                 Arguments.of("SHOW SCHEMAS FROM kat", ImmutableSet.of("kat"), ImmutableSet.of(), ImmutableSet.of()),
+                Arguments.of("SHOW SCHEMAS", ImmutableSet.of(DEFAULT_CATALOG), ImmutableSet.of(), ImmutableSet.of()),
                 Arguments.of("SHOW TABLES FROM kat.schem", ImmutableSet.of("kat"), ImmutableSet.of("schem"), ImmutableSet.of()),
+                Arguments.of("SHOW TABLES", ImmutableSet.of(DEFAULT_CATALOG), ImmutableSet.of(DEFAULT_SCHEMA), ImmutableSet.of()),
                 Arguments.of("ALTER SCHEMA kat.schem SET AUTHORIZATION will", ImmutableSet.of("kat"), ImmutableSet.of("schem"), ImmutableSet.of()),
+                Arguments.of("ALTER SCHEMA schem SET AUTHORIZATION will", ImmutableSet.of(DEFAULT_CATALOG), ImmutableSet.of("schem"), ImmutableSet.of()),
                 Arguments.of("ALTER TABLE cat.schem.tbl SET AUTHORIZATION will",
                         ImmutableSet.of("cat"),
                         ImmutableSet.of("schem"),
@@ -377,6 +390,9 @@ public void testTrinoQueryPropertiesTableExtraction(String query, Set<String> ca
         BufferedReader bufferedReader = new BufferedReader(new StringReader(query));
         HttpServletRequest mockRequest = prepareMockRequest();
         when(mockRequest.getReader()).thenReturn(bufferedReader);
+        when(mockRequest.getHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME)).thenReturn(DEFAULT_CATALOG);
+        when(mockRequest.getHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME)).thenReturn(DEFAULT_SCHEMA);
+
         TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, requestAnalyzerConfig);
 
         assertThat(trinoQueryProperties.getTables()).isEqualTo(tables);

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
@@ -43,7 +43,7 @@ name: "defaults-group"
 description: "test default schema and catalog"
 condition: |
   trinoQueryProperties.getDefaultCatalog().equals(java.util.Optional.of("other_catalog"))
-  && trinoQueryProperties.getDefaultCatalog().equals(java.util.Optional.of("other_catalog"))
+  && trinoQueryProperties.getDefaultSchema().equals(java.util.Optional.of("other_schema"))
 
 actions:
   - "result.put(\"routingGroup\", \"defaults-group\")"