Skip to content

Commit 8b63f16

Browse files
author
Jaeho Yoo
committed
Fix schema parsing fallthrough and Optional string formatting
1 parent a37a345 commit 8b63f16

File tree

2 files changed

+184
-13
lines changed

2 files changed

+184
-13
lines changed

‎gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java‎

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -91,26 +91,27 @@
9191

9292
public class TrinoQueryProperties
9393
{
94+
public static final String TRINO_CATALOG_HEADER_NAME = "X-Trino-Catalog";
95+
public static final String TRINO_SCHEMA_HEADER_NAME = "X-Trino-Schema";
96+
public static final String TRINO_PREPARED_STATEMENT_HEADER_NAME = "X-Trino-Prepared-Statement";
97+
9498
private final Logger log = Logger.get(TrinoQueryProperties.class);
9599
private final boolean isClientsUseV2Format;
96100
private final int maxBodySize;
101+
private final Optional<String> defaultCatalog;
102+
private final Optional<String> defaultSchema;
103+
private final ZstdDecompressor decompressor = ZstdDecompressor.create();
104+
97105
private String body = "";
98106
private String queryType = "";
99107
private String resourceGroupQueryType = "";
100108
private Set<QualifiedName> tables = ImmutableSet.of();
101-
private final Optional<String> defaultCatalog;
102-
private final Optional<String> defaultSchema;
103109
private Set<String> catalogs = ImmutableSet.of();
104110
private Set<String> schemas = ImmutableSet.of();
105111
private Set<String> catalogSchemas = ImmutableSet.of();
106112
private boolean isNewQuerySubmission;
107113
private Optional<String> errorMessage = Optional.empty();
108114
private Optional<String> queryId = Optional.empty();
109-
private final ZstdDecompressor decompressor = ZstdDecompressor.create();
110-
111-
public static final String TRINO_CATALOG_HEADER_NAME = "X-Trino-Catalog";
112-
public static final String TRINO_SCHEMA_HEADER_NAME = "X-Trino-Schema";
113-
public static final String TRINO_PREPARED_STATEMENT_HEADER_NAME = "X-Trino-Prepared-Statement";
114115

115116
@JsonCreator
116117
public TrinoQueryProperties(
@@ -361,7 +362,8 @@ private void visitNode(Node node, ImmutableSet.Builder<QualifiedName> tableBuild
361362
}
362363
}
363364
case TableFunctionInvocation s -> tableBuilder.add(qualifyName(s.getName()));
364-
default -> {}
365+
default -> {
366+
}
365367
}
366368

367369
for (Node child : node.getChildren()) {
@@ -391,20 +393,23 @@ private void setCatalogAndSchemaNameFromSchemaQualifiedName(
391393
if (schemaOptional.isEmpty()) {
392394
schemaBuilder.add(defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier));
393395
catalogBuilder.add(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier));
394-
catalogSchemaBuilder.add(format("%s.%s", defaultCatalog, defaultSchema));
396+
catalogSchemaBuilder.add(format("%s.%s", defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier),
397+
defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier)));
395398
}
396399
else {
397400
QualifiedName schema = schemaOptional.orElseThrow();
398401
switch (schema.getParts().size()) {
399402
case 1 -> {
400403
schemaBuilder.add(schema.getParts().getFirst());
401404
catalogBuilder.add(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier));
402-
catalogSchemaBuilder.add(format("%s.%s", defaultCatalog, schema.getParts().getFirst()));
405+
catalogSchemaBuilder.add(format("%s.%s", defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), schema.getParts().getFirst()));
406+
break;
403407
}
404408
case 2 -> {
405409
schemaBuilder.add(schema.getParts().get(1));
406410
catalogBuilder.add(schema.getParts().getFirst());
407411
catalogSchemaBuilder.add(format("%s.%s", schema.getParts().getFirst(), schema.getParts().getLast()));
412+
break;
408413
}
409414
default -> log.error("Schema has >2 parts: %s", schema);
410415
}
@@ -421,8 +426,7 @@ private QualifiedName qualifyName(QualifiedName name)
421426
{
422427
List<String> nameParts = name.getParts();
423428
return switch (nameParts.size()) {
424-
case 1 ->
425-
QualifiedName.of(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier), nameParts.getFirst());
429+
case 1 -> QualifiedName.of(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier), nameParts.getFirst());
426430
case 2 -> QualifiedName.of(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), nameParts.getFirst(), nameParts.get(1));
427431
case 3 -> QualifiedName.of(nameParts.getFirst(), nameParts.get(1), nameParts.get(2));
428432
default -> throw new RequestParsingException("Unexpected qualified name: " + name.getParts());
@@ -499,7 +503,7 @@ private QualifiedName parseIdentifierStringToQualifiedName(String name)
499503
parts.add(new Identifier(name.substring(start, name.length() - 1)));
500504
}
501505
else {
502-
parts.add(new Identifier(name.substring(start, name.length())));
506+
parts.add(new Identifier(name.substring(start)));
503507
}
504508
return QualifiedName.of(parts);
505509
}

‎gateway-ha/src/test/java/io/trino/gateway/ha/router/TestTrinoQueryProperties.java‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,75 @@
1616
import com.google.common.collect.ImmutableList;
1717
import com.google.common.collect.ImmutableSet;
1818
import io.airlift.json.JsonCodec;
19+
import jakarta.servlet.http.HttpServletRequest;
20+
import jakarta.ws.rs.HttpMethod;
1921
import org.junit.jupiter.api.Test;
22+
import org.junit.jupiter.params.ParameterizedTest;
23+
import org.junit.jupiter.params.provider.Arguments;
24+
import org.junit.jupiter.params.provider.MethodSource;
2025

26+
import java.io.BufferedReader;
27+
import java.io.IOException;
28+
import java.io.StringReader;
2129
import java.util.Optional;
30+
import java.util.stream.Stream;
2231

2332
import static org.assertj.core.api.Assertions.assertThat;
33+
import static org.mockito.Mockito.mock;
34+
import static org.mockito.Mockito.when;
2435

2536
final class TestTrinoQueryProperties
2637
{
38+
/**
39+
* Test data for schema qualified name switch statement cases.
40+
* Tests the break statements added to prevent fallthrough.
41+
*/
42+
static Stream<Arguments> provideSchemaQualifiedNameTestCases()
43+
{
44+
return Stream.of(
45+
// Test case 1: Single part schema (should use default catalog)
46+
Arguments.of("CREATE SCHEMA myschema",
47+
"default_catalog", "default_schema",
48+
ImmutableSet.of("default_catalog"),
49+
ImmutableSet.of("myschema"),
50+
ImmutableSet.of("default_catalog.myschema")),
51+
52+
// Test case 2: Two part schema (catalog.schema)
53+
Arguments.of("CREATE SCHEMA mycatalog.myschema",
54+
"default_catalog", "default_schema",
55+
ImmutableSet.of("mycatalog"),
56+
ImmutableSet.of("myschema"),
57+
ImmutableSet.of("mycatalog.myschema")),
58+
59+
// Test DROP SCHEMA with single part
60+
Arguments.of("DROP SCHEMA testschema",
61+
"test_catalog", "test_schema",
62+
ImmutableSet.of("test_catalog"),
63+
ImmutableSet.of("testschema"),
64+
ImmutableSet.of("test_catalog.testschema")),
65+
66+
// Test DROP SCHEMA with two parts
67+
Arguments.of("DROP SCHEMA testcatalog.testschema",
68+
"test_catalog", "test_schema",
69+
ImmutableSet.of("testcatalog"),
70+
ImmutableSet.of("testschema"),
71+
ImmutableSet.of("testcatalog.testschema")),
72+
73+
// Test SHOW TABLES with single part schema
74+
Arguments.of("SHOW TABLES FROM myschema",
75+
"default_catalog", "default_schema",
76+
ImmutableSet.of("default_catalog"),
77+
ImmutableSet.of("myschema"),
78+
ImmutableSet.of("default_catalog.myschema")),
79+
80+
// Test SHOW TABLES with two part schema
81+
Arguments.of("SHOW TABLES FROM mycatalog.myschema",
82+
"default_catalog", "default_schema",
83+
ImmutableSet.of("mycatalog"),
84+
ImmutableSet.of("myschema"),
85+
ImmutableSet.of("mycatalog.myschema")));
86+
}
87+
2788
@Test
2889
void testJsonCreator()
2990
{
@@ -91,4 +152,110 @@ void testJsonCreatorWithEmptyProperties()
91152
assertThat(deserializedTrinoQueryProperties.isQueryParsingSuccessful()).isEqualTo(trinoQueryProperties.isQueryParsingSuccessful());
92153
assertThat(deserializedTrinoQueryProperties.getErrorMessage()).isEqualTo(trinoQueryProperties.getErrorMessage());
93154
}
155+
156+
/**
157+
* Test the switch statement with break statements to prevent fallthrough.
158+
* This test verifies that each case in setCatalogAndSchemaNameFromSchemaQualifiedName
159+
* properly handles different schema qualified name formats and doesn't fall through
160+
* to subsequent cases.
161+
*/
162+
@ParameterizedTest
163+
@MethodSource("provideSchemaQualifiedNameTestCases")
164+
void testSchemaQualifiedNameSwitchStatement(String query, String defaultCatalog, String defaultSchema,
165+
ImmutableSet<String> expectedCatalogs, ImmutableSet<String> expectedSchemas,
166+
ImmutableSet<String> expectedCatalogSchemas)
167+
throws IOException
168+
{
169+
HttpServletRequest mockRequest = prepareMockRequest();
170+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
171+
when(mockRequest.getHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME)).thenReturn(defaultCatalog);
172+
when(mockRequest.getHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME)).thenReturn(defaultSchema);
173+
174+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(
175+
mockRequest,
176+
false, // isClientsUseV2Format
177+
1024 * 1024); // maxBodySize
178+
179+
// Verify that the switch statement correctly processed the schema qualified
180+
// name without falling through to other cases
181+
assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(expectedCatalogs);
182+
assertThat(trinoQueryProperties.getSchemas()).isEqualTo(expectedSchemas);
183+
assertThat(trinoQueryProperties.getCatalogSchemas()).isEqualTo(expectedCatalogSchemas);
184+
185+
// Ensure query parsing was successful
186+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
187+
}
188+
189+
@Test
190+
void testSchemaQualifiedNameWithoutDefaults()
191+
throws IOException
192+
{
193+
String query = "CREATE SCHEMA myschema";
194+
HttpServletRequest mockRequest = prepareMockRequest();
195+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
196+
// Don't set default catalog or schema headers
197+
198+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
199+
200+
// The exception is caught and stored in errorMessage, not re-thrown
201+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isFalse();
202+
assertThat(trinoQueryProperties.getErrorMessage()).isPresent();
203+
assertThat(trinoQueryProperties.getErrorMessage().get()).contains("Name not fully qualified");
204+
}
205+
206+
@Test
207+
void testSchemaQualifiedNameWithEmptyDefaults()
208+
throws IOException
209+
{
210+
String query = "SHOW TABLES"; // This will use empty schema optional
211+
HttpServletRequest mockRequest = prepareMockRequest();
212+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
213+
when(mockRequest.getHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME)).thenReturn("test_catalog");
214+
when(mockRequest.getHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME)).thenReturn("test_schema");
215+
216+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(
217+
mockRequest,
218+
false,
219+
1024 * 1024);
220+
221+
// When schema optional is empty, it should use defaults
222+
assertThat(trinoQueryProperties.getCatalogs()).contains("test_catalog");
223+
assertThat(trinoQueryProperties.getSchemas()).contains("test_schema");
224+
assertThat(trinoQueryProperties.getCatalogSchemas()).contains("test_catalog.test_schema");
225+
}
226+
227+
@Test
228+
void testSchemaQualifiedNameWithInvalidParts()
229+
throws IOException
230+
{
231+
// Note: This test demonstrates that the default case in the switch statement is
232+
// reached for schemas with >2 parts, which will log an error but continue processing.
233+
// Since we cannot easily mock the SQL parser to create a QualifiedName with >2 parts
234+
// from a valid SQL statement, this test documents the expected behavior.
235+
236+
// Test that a valid two-part schema name works correctly (boundary case)
237+
String query = "CREATE SCHEMA catalog.schema";
238+
HttpServletRequest mockRequest = prepareMockRequest();
239+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
240+
when(mockRequest.getHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME)).thenReturn("default_catalog");
241+
when(mockRequest.getHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME)).thenReturn("default_schema");
242+
243+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(
244+
mockRequest,
245+
false,
246+
1024 * 1024);
247+
248+
// Verify the two-part case works and doesn't fall through to default
249+
assertThat(trinoQueryProperties.getCatalogs()).containsExactly("catalog");
250+
assertThat(trinoQueryProperties.getSchemas()).containsExactly("schema");
251+
assertThat(trinoQueryProperties.getCatalogSchemas()).containsExactly("catalog.schema");
252+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
253+
}
254+
255+
private HttpServletRequest prepareMockRequest()
256+
{
257+
HttpServletRequest mockRequest = mock(HttpServletRequest.class);
258+
when(mockRequest.getMethod()).thenReturn(HttpMethod.POST);
259+
return mockRequest;
260+
}
94261
}

0 commit comments

Comments
 (0)