Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Exclude with statement names from tables #532

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+79 −30
Open

Exclude with statement names from tables #532

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 49 additions & 30 deletions gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
import io.trino.sql.tree.Statement;
import io.trino.sql.tree.Table;
import io.trino.sql.tree.TableFunctionInvocation;
import io.trino.sql.tree.WithQuery;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.ws.rs.HttpMethod;

Expand All @@ -74,6 +75,7 @@
import java.util.Set;
import java.util.stream.Collectors;

import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.io.BaseEncoding.base64Url;
import static io.airlift.json.JsonCodec.jsonCodec;
import static java.lang.Math.toIntExact;
Expand Down Expand Up @@ -194,13 +196,25 @@ else if (statement instanceof ExecuteImmediate executeImmediate) {

queryType = statement.getClass().getSimpleName();
resourceGroupQueryType = StatementUtils.getQueryType(statement).toString();
ImmutableSet.Builder<QualifiedName> tableBuilder = ImmutableSet.builder();
List<QualifiedName> allQUalifiedNames = new ArrayList<>();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why did you change to ArrayList? #Also, the new name has a typo. U should be a lowercase.

ImmutableSet.Builder<String> catalogBuilder = ImmutableSet.builder();
ImmutableSet.Builder<String> schemaBuilder = ImmutableSet.builder();
ImmutableSet.Builder<String> catalogSchemaBuilder = ImmutableSet.builder();

getNames(statement, tableBuilder, catalogBuilder, schemaBuilder, catalogSchemaBuilder);
tables = tableBuilder.build();
ImmutableSet.Builder<Identifier> withQueryNameBuilder = ImmutableSet.builder();

getNames(statement, allQUalifiedNames, catalogBuilder, schemaBuilder, catalogSchemaBuilder, withQueryNameBuilder);
ImmutableSet<Identifier> withQueryNames = withQueryNameBuilder.build();
tables = allQUalifiedNames.stream()
.filter(n -> (n.getOriginalParts().size() > 1 || !withQueryNames.contains(n.getOriginalParts().getFirst())))
.map(t -> {
try {
return qualifyName(t);
}
catch (RequestParsingException e) {
throw new RuntimeException(e);
}
})
.collect(toImmutableSet());
catalogBuilder.addAll(tables.stream().map(q -> q.getParts().getFirst()).iterator());
catalogs = catalogBuilder.build();
schemaBuilder.addAll(tables.stream().map(q -> q.getParts().get(1)).iterator());
Expand Down Expand Up @@ -260,27 +274,30 @@ private String decodePreparedStatementFromHeader(String headerValue)
return new String(preparedStatement, UTF_8);
}

private void getNames(Node node, ImmutableSet.Builder<QualifiedName> tableBuilder,
private void getNames(
Node node,
List<QualifiedName> allQualifiedNames,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand the difference between tableBuilder and allQualifiedNames. What's the motivation of this renaming?

ImmutableSet.Builder<String> catalogBuilder,
ImmutableSet.Builder<String> schemaBuilder,
ImmutableSet.Builder<String> catalogSchemaBuilder)
ImmutableSet.Builder<String> catalogSchemaBuilder,
ImmutableSet.Builder<Identifier> withQueryNamesBuilder)
throws RequestParsingException
{
switch (node) {
case AddColumn s -> tableBuilder.add(qualifyName(s.getName()));
case Analyze s -> tableBuilder.add(qualifyName(s.getTableName()));
case AddColumn s -> allQualifiedNames.add(s.getName());
case Analyze s -> allQualifiedNames.add(s.getTableName());
case CreateCatalog s -> catalogBuilder.add(s.getCatalogName().getValue());
case CreateMaterializedView s -> tableBuilder.add(qualifyName(s.getName()));
case CreateMaterializedView s -> allQualifiedNames.add(s.getName());
case CreateSchema s -> setCatalogAndSchemaNameFromSchemaQualifiedName(Optional.of(s.getSchemaName()), catalogBuilder, schemaBuilder, catalogSchemaBuilder);
case CreateTable s -> tableBuilder.add(qualifyName(s.getName()));
case CreateView s -> tableBuilder.add(qualifyName(s.getName()));
case CreateTableAsSelect s -> tableBuilder.add(qualifyName(s.getName()));
case CreateTable s -> allQualifiedNames.add(s.getName());
case CreateView s -> allQualifiedNames.add(s.getName());
case CreateTableAsSelect s -> allQualifiedNames.add(s.getName());
case DropCatalog s -> catalogBuilder.add(s.getCatalogName().getValue());
case DropSchema s -> setCatalogAndSchemaNameFromSchemaQualifiedName(Optional.of(s.getSchemaName()), catalogBuilder, schemaBuilder, catalogSchemaBuilder);
case DropTable s -> tableBuilder.add(qualifyName(s.getTableName()));
case DropTable s -> allQualifiedNames.add(s.getTableName());
case RenameMaterializedView s -> {
tableBuilder.add(qualifyName(s.getSource()));
tableBuilder.add(qualifyName(s.getTarget()));
allQualifiedNames.add(s.getSource());
allQualifiedNames.add(s.getTarget());
}
case RenameSchema s -> {
setCatalogAndSchemaNameFromSchemaQualifiedName(Optional.of(s.getSource()), catalogBuilder, schemaBuilder, catalogSchemaBuilder);
Expand All @@ -301,48 +318,49 @@ private void getNames(Node node, ImmutableSet.Builder<QualifiedName> tableBuilde
}
case RenameTable s -> {
QualifiedName qualifiedSource = qualifyName(s.getSource());
tableBuilder.add(qualifiedSource);
allQualifiedNames.add(qualifiedSource);
QualifiedName target = s.getTarget();
if (target.getParts().size() == 1) {
tableBuilder.add(QualifiedName.of(qualifiedSource.getParts().getFirst(), qualifiedSource.getParts().get(1), target.getParts().getFirst()));
allQualifiedNames.add(QualifiedName.of(qualifiedSource.getParts().getFirst(), qualifiedSource.getParts().get(1), target.getParts().getFirst()));
}
else {
tableBuilder.add(QualifiedName.of(qualifiedSource.getParts().getFirst(), target.getParts().getFirst(), target.getParts().get(1)));
allQualifiedNames.add(QualifiedName.of(qualifiedSource.getParts().getFirst(), target.getParts().getFirst(), target.getParts().get(1)));
}
}
case RenameView s -> {
QualifiedName qualifiedSource = qualifyName(s.getSource());
tableBuilder.add(qualifiedSource);
allQualifiedNames.add(qualifiedSource);
QualifiedName target = s.getTarget();
if (target.getParts().size() == 1) {
tableBuilder.add(QualifiedName.of(qualifiedSource.getParts().getFirst(), qualifiedSource.getParts().get(1), target.getParts().getFirst()));
allQualifiedNames.add(QualifiedName.of(qualifiedSource.getParts().getFirst(), qualifiedSource.getParts().get(1), target.getParts().getFirst()));
}
else {
tableBuilder.add(QualifiedName.of(qualifiedSource.getParts().getFirst(), target.getParts().getFirst(), target.getParts().get(1)));
allQualifiedNames.add(QualifiedName.of(qualifiedSource.getParts().getFirst(), target.getParts().getFirst(), target.getParts().get(1)));
}
}
case SetProperties s -> tableBuilder.add(qualifyName(s.getName()));
case ShowColumns s -> tableBuilder.add(qualifyName(s.getTable()));
case SetProperties s -> allQualifiedNames.add(s.getName());
case ShowColumns s -> allQualifiedNames.add(s.getTable());
case ShowCreate s -> {
if (s.getType() == ShowCreate.Type.SCHEMA) {
setCatalogAndSchemaNameFromSchemaQualifiedName(Optional.of(s.getName()), catalogBuilder, schemaBuilder, catalogSchemaBuilder);
}
else {
tableBuilder.add(qualifyName(s.getName()));
allQualifiedNames.add(s.getName());
}
}
case ShowSchemas s -> catalogBuilder.add(s.getCatalog().map(Identifier::getValue).or(() -> defaultCatalog).orElseThrow(this::unsetDefaultExceptionSupplier));
case ShowTables s -> setCatalogAndSchemaNameFromSchemaQualifiedName(s.getSchema(), catalogBuilder, schemaBuilder, catalogSchemaBuilder);
case SetSchemaAuthorization s -> setCatalogAndSchemaNameFromSchemaQualifiedName(Optional.of(s.getSource()), catalogBuilder, schemaBuilder, catalogSchemaBuilder);
case SetTableAuthorization s -> tableBuilder.add(qualifyName(s.getSource()));
case SetViewAuthorization s -> tableBuilder.add(qualifyName(s.getSource()));
case Table s -> tableBuilder.add(qualifyName(s.getName()));
case TableFunctionInvocation s -> tableBuilder.add(qualifyName(s.getName()));
case SetTableAuthorization s -> allQualifiedNames.add(s.getSource());
case SetViewAuthorization s -> allQualifiedNames.add(s.getSource());
case Table s -> allQualifiedNames.add(s.getName());
case TableFunctionInvocation s -> allQualifiedNames.add(s.getName());
case WithQuery withQuery -> withQueryNamesBuilder.add(withQuery.getName());
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

didn't know you could do this directly, i thought withQuery comes with Query.

default -> {}
}

for (Node child : node.getChildren()) {
getNames(child, tableBuilder, catalogBuilder, schemaBuilder, catalogSchemaBuilder);
getNames(child, allQualifiedNames, catalogBuilder, schemaBuilder, catalogSchemaBuilder, withQueryNamesBuilder);
}
}

Expand Down Expand Up @@ -386,7 +404,8 @@ private QualifiedName qualifyName(QualifiedName table)
{
List<String> tableParts = table.getParts();
return switch (tableParts.size()) {
case 1 -> QualifiedName.of(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier), tableParts.getFirst());
case 1 ->
QualifiedName.of(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier), tableParts.getFirst());
case 2 -> QualifiedName.of(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), tableParts.getFirst(), tableParts.get(1));
case 3 -> QualifiedName.of(tableParts.getFirst(), tableParts.get(1), tableParts.get(2));
default -> throw new RequestParsingException("Unexpected table name: " + table.getParts());
Expand Down
30 changes: 30 additions & 0 deletions gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -404,6 +404,36 @@ void testTrinoQueryPropertiesTableExtraction(String query, Set<String> catalogs,
assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(catalogs);
}

@Test
void testWithQueryNameExcluded()
throws IOException
{
String query = """
WITH dos (c1) as (select c1 from uno),
uno(c1) as (select c1 from cat.schem.tbl1)
SELECT c1 from uno, dos
Comment on lines +412 to +414
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Uppercase as, select, from.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i don't think this is valid SQL. You cannot use uno as it is defined after dos.

Suggested change
WITH dos (c1) as (select c1 from uno),
uno(c1) as (select c1 from cat.schem.tbl1)
SELECT c1 from uno, dos
WITH dos AS (SELECT c1 from cat.schem.tbl1),
uno as (SELECT c1 FROM dos)
SELECT c1 FROM uno, dos

you can test this in trino via simple sql

// works fine
WITH dos AS (SELECT node_id FROM system.runtime.nodes),
uno AS (SELECT node_id as node_id2 FROM dos)
SELECT node_id, node_id2
FROM uno, dos

// doesn't work as uno is not defined before dos.
// Query 20241023_021715_04215_f8xaz failed: line 1:34: Table 'system.runtime.uno' does not exist
WITH dos AS (SELECT node_id FROM uno),
uno AS (SELECT node_id as node_id2 FROM system.runtime.nodes)
SELECT node_id, node_id2
FROM uno, dos

""";
BufferedReader bufferedReader = new BufferedReader(new StringReader(query));
HttpServletRequest mockRequestWithDefaults = prepareMockRequest();
when(mockRequestWithDefaults.getReader()).thenReturn(bufferedReader);
when(mockRequestWithDefaults.getHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME)).thenReturn(DEFAULT_CATALOG);
when(mockRequestWithDefaults.getHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME)).thenReturn(DEFAULT_SCHEMA);

TrinoQueryProperties trinoQueryPropertiesWithDefaults = new TrinoQueryProperties(mockRequestWithDefaults, requestAnalyzerConfig);

assertThat(trinoQueryPropertiesWithDefaults.tablesContains(String.format("%s.%s.%s", DEFAULT_CATALOG, DEFAULT_SCHEMA, "uno"))).isFalse();
assertThat(trinoQueryPropertiesWithDefaults.tablesContains(String.format("%s.%s.%s", DEFAULT_CATALOG, DEFAULT_SCHEMA, "does"))).isFalse();
Comment on lines +424 to +425
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Static import format method. https://github.com/trinodb/trino/blob/master/.github/DEVELOPMENT.md#prefer-string-formatting

assertThat(trinoQueryPropertiesWithDefaults.tablesContains("cat.schem.tbl1")).isTrue();

HttpServletRequest mockRequestNoDefaults = prepareMockRequest();
bufferedReader = new BufferedReader(new StringReader(query));
when(mockRequestNoDefaults.getReader()).thenReturn(bufferedReader);

TrinoQueryProperties trinoQueryPropertiesNoDefaults = new TrinoQueryProperties(mockRequestNoDefaults, requestAnalyzerConfig);
assertThat(trinoQueryPropertiesNoDefaults.getTables().size()).isEqualTo(1);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AssertJ provides hasSize method.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do you consider temporary tables (with clause) as non-tables?

assertThat(trinoQueryPropertiesNoDefaults.tablesContains("cat.schem.tbl1")).isTrue();
}

private HttpServletRequest prepareMockRequest()
{
HttpServletRequest mockRequest = mock(HttpServletRequest.class);
Expand Down