Skip to content

Commit e2b1d40

Browse files
author
Jaeho Yoo
committed
Fix schema parsing fallthrough and Optional string formatting
1 parent 9e8c11f commit e2b1d40

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

0 commit comments

Comments
 (0)