Skip to content

Commit c7a73e1

Browse files
author
Jaeho Yoo
committed
Fix schema parsing fallthrough and Optional string formatting
1 parent 513a41e commit c7a73e1

File tree

2 files changed

+196
-3
lines changed

2 files changed

+196
-3
lines changed

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -392,20 +392,23 @@ private void setCatalogAndSchemaNameFromSchemaQualifiedName(
392392
if (schemaOptional.isEmpty()) {
393393
schemaBuilder.add(defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier));
394394
catalogBuilder.add(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier));
395-
catalogSchemaBuilder.add(format("%s.%s", defaultCatalog, defaultSchema));
395+
catalogSchemaBuilder.add(format("%s.%s", defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier),
396+
defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier)));
396397
}
397398
else {
398399
QualifiedName schema = schemaOptional.orElseThrow();
399400
switch (schema.getParts().size()) {
400401
case 1 -> {
401402
schemaBuilder.add(schema.getParts().getFirst());
402403
catalogBuilder.add(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier));
403-
catalogSchemaBuilder.add(format("%s.%s", defaultCatalog, schema.getParts().getFirst()));
404+
catalogSchemaBuilder.add(format("%s.%s", defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), schema.getParts().getFirst()));
405+
break;
404406
}
405407
case 2 -> {
406408
schemaBuilder.add(schema.getParts().get(1));
407409
catalogBuilder.add(schema.getParts().getFirst());
408410
catalogSchemaBuilder.add(format("%s.%s", schema.getParts().getFirst(), schema.getParts().getLast()));
411+
break;
409412
}
410413
default -> log.error("Schema has >2 parts: %s", schema);
411414
}
@@ -499,7 +502,7 @@ private QualifiedName parseIdentifierStringToQualifiedName(String name)
499502
parts.add(new Identifier(name.substring(start, name.length() - 1)));
500503
}
501504
else {
502-
parts.add(new Identifier(name.substring(start, name.length())));
505+
parts.add(new Identifier(name.substring(start)));
503506
}
504507
return QualifiedName.of(parts);
505508
}

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

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,20 @@
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;
2022

23+
import java.io.BufferedReader;
24+
import java.io.IOException;
25+
import java.io.StringReader;
2126
import java.util.Optional;
2227

28+
import static io.trino.gateway.ha.router.TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME;
29+
import static io.trino.gateway.ha.router.TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME;
2330
import static org.assertj.core.api.Assertions.assertThat;
31+
import static org.mockito.Mockito.mock;
32+
import static org.mockito.Mockito.when;
2433

2534
final class TestTrinoQueryProperties
2635
{
@@ -91,4 +100,185 @@ void testJsonCreatorWithEmptyProperties()
91100
assertThat(deserializedTrinoQueryProperties.isQueryParsingSuccessful()).isEqualTo(trinoQueryProperties.isQueryParsingSuccessful());
92101
assertThat(deserializedTrinoQueryProperties.getErrorMessage()).isEqualTo(trinoQueryProperties.getErrorMessage());
93102
}
103+
104+
@Test
105+
void testCreateSchemaSinglePart()
106+
throws IOException
107+
{
108+
String query = "CREATE SCHEMA myschema";
109+
HttpServletRequest mockRequest = prepareMockRequest();
110+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
111+
when(mockRequest.getHeader(TRINO_CATALOG_HEADER_NAME)).thenReturn("default_catalog");
112+
when(mockRequest.getHeader(TRINO_SCHEMA_HEADER_NAME)).thenReturn("default_schema");
113+
114+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
115+
116+
// Single part schema should use the default catalog
117+
assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(ImmutableSet.of("default_catalog"));
118+
assertThat(trinoQueryProperties.getSchemas()).isEqualTo(ImmutableSet.of("myschema"));
119+
assertThat(trinoQueryProperties.getCatalogSchemas()).isEqualTo(ImmutableSet.of("default_catalog.myschema"));
120+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
121+
}
122+
123+
@Test
124+
void testCreateSchemaTwoPart()
125+
throws IOException
126+
{
127+
String query = "CREATE SCHEMA mycatalog.myschema";
128+
HttpServletRequest mockRequest = prepareMockRequest();
129+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
130+
when(mockRequest.getHeader(TRINO_CATALOG_HEADER_NAME)).thenReturn("default_catalog");
131+
when(mockRequest.getHeader(TRINO_SCHEMA_HEADER_NAME)).thenReturn("default_schema");
132+
133+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
134+
135+
// Two-part schema should use a specified catalog
136+
assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(ImmutableSet.of("mycatalog"));
137+
assertThat(trinoQueryProperties.getSchemas()).isEqualTo(ImmutableSet.of("myschema"));
138+
assertThat(trinoQueryProperties.getCatalogSchemas()).isEqualTo(ImmutableSet.of("mycatalog.myschema"));
139+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
140+
}
141+
142+
@Test
143+
void testDropSchemaSinglePart()
144+
throws IOException
145+
{
146+
String query = "DROP SCHEMA testschema";
147+
HttpServletRequest mockRequest = prepareMockRequest();
148+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
149+
when(mockRequest.getHeader(TRINO_CATALOG_HEADER_NAME)).thenReturn("test_catalog");
150+
when(mockRequest.getHeader(TRINO_SCHEMA_HEADER_NAME)).thenReturn("test_schema");
151+
152+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
153+
154+
// Single part schema should use the default catalog
155+
assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(ImmutableSet.of("test_catalog"));
156+
assertThat(trinoQueryProperties.getSchemas()).isEqualTo(ImmutableSet.of("testschema"));
157+
assertThat(trinoQueryProperties.getCatalogSchemas()).isEqualTo(ImmutableSet.of("test_catalog.testschema"));
158+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
159+
}
160+
161+
@Test
162+
void testDropSchemaTwoPart()
163+
throws IOException
164+
{
165+
String query = "DROP SCHEMA testcatalog.testschema";
166+
HttpServletRequest mockRequest = prepareMockRequest();
167+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
168+
when(mockRequest.getHeader(TRINO_CATALOG_HEADER_NAME)).thenReturn("test_catalog");
169+
when(mockRequest.getHeader(TRINO_SCHEMA_HEADER_NAME)).thenReturn("test_schema");
170+
171+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
172+
173+
// Two-part schema should use a specified catalog
174+
assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(ImmutableSet.of("testcatalog"));
175+
assertThat(trinoQueryProperties.getSchemas()).isEqualTo(ImmutableSet.of("testschema"));
176+
assertThat(trinoQueryProperties.getCatalogSchemas()).isEqualTo(ImmutableSet.of("testcatalog.testschema"));
177+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
178+
}
179+
180+
@Test
181+
void testShowTablesSinglePart()
182+
throws IOException
183+
{
184+
String query = "SHOW TABLES FROM myschema";
185+
HttpServletRequest mockRequest = prepareMockRequest();
186+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
187+
when(mockRequest.getHeader(TRINO_CATALOG_HEADER_NAME)).thenReturn("default_catalog");
188+
when(mockRequest.getHeader(TRINO_SCHEMA_HEADER_NAME)).thenReturn("default_schema");
189+
190+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
191+
192+
// Single part schema should use the default catalog
193+
assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(ImmutableSet.of("default_catalog"));
194+
assertThat(trinoQueryProperties.getSchemas()).isEqualTo(ImmutableSet.of("myschema"));
195+
assertThat(trinoQueryProperties.getCatalogSchemas()).isEqualTo(ImmutableSet.of("default_catalog.myschema"));
196+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
197+
}
198+
199+
@Test
200+
void testShowTablesTwoPart()
201+
throws IOException
202+
{
203+
String query = "SHOW TABLES FROM mycatalog.myschema";
204+
HttpServletRequest mockRequest = prepareMockRequest();
205+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
206+
when(mockRequest.getHeader(TRINO_CATALOG_HEADER_NAME)).thenReturn("default_catalog");
207+
when(mockRequest.getHeader(TRINO_SCHEMA_HEADER_NAME)).thenReturn("default_schema");
208+
209+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
210+
211+
// Two-part schema should use a specified catalog
212+
assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(ImmutableSet.of("mycatalog"));
213+
assertThat(trinoQueryProperties.getSchemas()).isEqualTo(ImmutableSet.of("myschema"));
214+
assertThat(trinoQueryProperties.getCatalogSchemas()).isEqualTo(ImmutableSet.of("mycatalog.myschema"));
215+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
216+
}
217+
218+
@Test
219+
void testSchemaQualifiedNameWithoutDefaults()
220+
throws IOException
221+
{
222+
String query = "CREATE SCHEMA myschema";
223+
HttpServletRequest mockRequest = prepareMockRequest();
224+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
225+
// Don't set default catalog or schema headers
226+
227+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
228+
229+
// The exception is caught and stored in errorMessage, not re-thrown
230+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isFalse();
231+
assertThat(trinoQueryProperties.getErrorMessage()).isPresent();
232+
assertThat(trinoQueryProperties.getErrorMessage().get()).contains("Name not fully qualified");
233+
}
234+
235+
@Test
236+
void testSchemaQualifiedNameWithEmptyDefaults()
237+
throws IOException
238+
{
239+
String query = "SHOW TABLES"; // This will use empty schema optional
240+
HttpServletRequest mockRequest = prepareMockRequest();
241+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
242+
when(mockRequest.getHeader(TRINO_CATALOG_HEADER_NAME)).thenReturn("test_catalog");
243+
when(mockRequest.getHeader(TRINO_SCHEMA_HEADER_NAME)).thenReturn("test_schema");
244+
245+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
246+
247+
// When schema optional is empty, it should use defaults
248+
assertThat(trinoQueryProperties.getCatalogs()).contains("test_catalog");
249+
assertThat(trinoQueryProperties.getSchemas()).contains("test_schema");
250+
assertThat(trinoQueryProperties.getCatalogSchemas()).contains("test_catalog.test_schema");
251+
}
252+
253+
@Test
254+
void testSchemaQualifiedNameWithInvalidParts()
255+
throws IOException
256+
{
257+
// Note: This test demonstrates that the default case in the switch statement is
258+
// reached for schemas with >2 parts, which will log an error but continue processing.
259+
// Since we cannot easily mock the SQL parser to create a QualifiedName with >2 parts
260+
// from a valid SQL statement, this test documents the expected behavior.
261+
262+
// Test that a valid two-part schema name works correctly (boundary case)
263+
String query = "CREATE SCHEMA catalog.schema";
264+
HttpServletRequest mockRequest = prepareMockRequest();
265+
when(mockRequest.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
266+
when(mockRequest.getHeader(TRINO_CATALOG_HEADER_NAME)).thenReturn("default_catalog");
267+
when(mockRequest.getHeader(TRINO_SCHEMA_HEADER_NAME)).thenReturn("default_schema");
268+
269+
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(mockRequest, false, 1024 * 1024);
270+
271+
// Verify the two-part case works and doesn't fall through to default
272+
assertThat(trinoQueryProperties.getCatalogs()).containsExactly("catalog");
273+
assertThat(trinoQueryProperties.getSchemas()).containsExactly("schema");
274+
assertThat(trinoQueryProperties.getCatalogSchemas()).containsExactly("catalog.schema");
275+
assertThat(trinoQueryProperties.isQueryParsingSuccessful()).isTrue();
276+
}
277+
278+
private HttpServletRequest prepareMockRequest()
279+
{
280+
HttpServletRequest mockRequest = mock(HttpServletRequest.class);
281+
when(mockRequest.getMethod()).thenReturn(HttpMethod.POST);
282+
return mockRequest;
283+
}
94284
}

0 commit comments

Comments
 (0)