Skip to content

Commit

Permalink
Separate query and description in QueryAssert
Browse files Browse the repository at this point in the history 
Previously, the `QueryAssert.query` could represent original SQL query,
or a description (in case of column projections). This commit separates
those duties. Original query is not available in case of projections
because it would not produce the current (projected) result.
  • Loading branch information
@findepi
findepi committed Jan 31, 2024
1 parent 4baa745 commit baaef6d
Showing 1 changed file with 41 additions and 21 deletions.
62 changes: 41 additions & 21 deletions core/trino-main/src/test/java/io/trino/sql/query/QueryAssertions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@
import org.assertj.core.api.AbstractAssert;
import org.assertj.core.api.AssertProvider;
import org.assertj.core.api.ListAssert;
import org.assertj.core.description.Description;
import org.assertj.core.description.TextDescription;
import org.assertj.core.presentation.Representation;
import org.assertj.core.presentation.StandardRepresentation;
import org.assertj.core.util.CanIgnoreReturnValue;
Expand All @@ -53,6 +55,7 @@
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.stream.Collectors;
Expand All @@ -67,7 +70,6 @@
import static io.trino.testing.TestingHandles.TEST_CATALOG_NAME;
import static io.trino.testing.TestingSession.testSessionBuilder;
import static io.trino.testing.TransactionBuilder.transaction;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.toList;
import static org.assertj.core.api.Assertions.assertThat;
Expand Down Expand Up @@ -288,21 +290,31 @@ private String formatRowElement(Object value)

private final QueryRunner runner;
private final Session session;
private final String query;
private final Optional<String> query;
private final Description description;
private boolean ordered;
private boolean skipTypesCheck;
private boolean skipResultsCorrectnessCheckForPushdown;

static AssertProvider<QueryAssert> newQueryAssert(String query, QueryRunner runner, Session session)
{
MaterializedResult result = runner.execute(session, query);
return () -> new QueryAssert(runner, session, query, result, false, false, false);
return () -> new QueryAssert(
runner,
session,
Optional.of(query),
new TextDescription("%s", query),
result,
false,
false,
false);
}

private QueryAssert(
QueryRunner runner,
Session session,
String query,
Optional<String> query,
Description description,
MaterializedResult actual,
boolean ordered,
boolean skipTypesCheck,
Expand All @@ -312,6 +324,7 @@ private QueryAssert(
this.runner = requireNonNull(runner, "runner is null");
this.session = requireNonNull(session, "session is null");
this.query = requireNonNull(query, "query is null");
this.description = requireNonNull(description, "description is null");
this.ordered = ordered;
this.skipTypesCheck = skipTypesCheck;
this.skipResultsCorrectnessCheckForPushdown = skipResultsCorrectnessCheckForPushdown;
Expand All @@ -322,7 +335,8 @@ public QueryAssert exceptColumns(String... columnNamesToExclude)
return new QueryAssert(
runner,
session,
format("%s except columns %s", query, Arrays.toString(columnNamesToExclude)),
Optional.empty(), // original query would not produce projected result
new TextDescription("%s except columns %s", description, Arrays.toString(columnNamesToExclude)),
actual.exceptColumns(columnNamesToExclude),
ordered,
skipTypesCheck,
Expand All @@ -334,7 +348,8 @@ public QueryAssert projected(String... columnNamesToInclude)
return new QueryAssert(
runner,
session,
format("%s projected with %s", query, Arrays.toString(columnNamesToInclude)),
Optional.empty(), // original query would not produce projected result
new TextDescription("%s projected with %s", description, Arrays.toString(columnNamesToInclude)),
actual.project(columnNamesToInclude),
ordered,
skipTypesCheck,
Expand Down Expand Up @@ -382,11 +397,11 @@ public QueryAssert matches(MaterializedResult expected)
{
return satisfies(actual -> {
if (!skipTypesCheck) {
assertTypes(query, actual, expected.getTypes());
assertTypes(description, actual, expected.getTypes());
}

ListAssert<MaterializedRow> assertion = assertThat(actual.getMaterializedRows())
.as("Rows for query [%s]", query)
.as("Rows for query [%s]", description)
.withRepresentation(ROWS_REPRESENTATION);

if (ordered) {
Expand All @@ -404,7 +419,7 @@ public QueryAssert matches(PlanMatchPattern expectedPlan)
Metadata metadata = runner.getPlannerContext().getMetadata();
transaction(runner.getTransactionManager(), metadata, runner.getAccessControl())
.execute(session, session -> {
Plan plan = runner.createPlan(session, query);
Plan plan = runner.createPlan(session, query());
assertPlan(
session,
metadata,
Expand All @@ -428,11 +443,11 @@ public QueryAssert containsAll(MaterializedResult expected)
{
return satisfies(actual -> {
if (!skipTypesCheck) {
assertTypes(query, actual, expected.getTypes());
assertTypes(description, actual, expected.getTypes());
}

assertThat(actual.getMaterializedRows())
.as("Rows for query [%s]", query)
.as("Rows for query [%s]", description)
.withRepresentation(ROWS_REPRESENTATION)
.containsAll(expected.getMaterializedRows());
});
Expand All @@ -442,7 +457,7 @@ public QueryAssert containsAll(MaterializedResult expected)
public QueryAssert hasOutputTypes(List<Type> expectedTypes)
{
return satisfies(actual -> {
assertTypes(query, actual, expectedTypes);
assertTypes(description, actual, expectedTypes);
});
}

Expand All @@ -451,23 +466,23 @@ public QueryAssert outputHasType(int index, Type expectedType)
{
return satisfies(actual -> {
assertThat(actual.getTypes())
.as("Output types for query [%s]", query)
.as("Output types for query [%s]", description)
.element(index).isEqualTo(expectedType);
});
}

private static void assertTypes(String query, MaterializedResult actual, List<Type> expectedTypes)
private static void assertTypes(Description queryDescription, MaterializedResult actual, List<Type> expectedTypes)
{
assertThat(actual.getTypes())
.as("Output types for query [%s]", query)
.as("Output types for query [%s]", queryDescription)
.isEqualTo(expectedTypes);
}

@CanIgnoreReturnValue
public QueryAssert returnsEmptyResult()
{
return satisfies(actual -> {
assertThat(actual.getMaterializedRows()).as("Rows for query [%s]", query).isEmpty();
assertThat(actual.getMaterializedRows()).as("Rows for query [%s]", description).isEmpty();
});
}

Expand All @@ -482,7 +497,7 @@ public QueryAssert isFullyPushedDown()
Metadata metadata = runner.getPlannerContext().getMetadata();
transaction(runner.getTransactionManager(), metadata, runner.getAccessControl())
.execute(session, session -> {
Plan plan = runner.createPlan(session, query);
Plan plan = runner.createPlan(session, query());
assertPlan(
session,
metadata,
Expand Down Expand Up @@ -511,7 +526,7 @@ public QueryAssert isReplacedWithEmptyValues()
Metadata metadata = runner.getPlannerContext().getMetadata();
transaction(runner.getTransactionManager(), metadata, runner.getAccessControl())
.execute(session, session -> {
Plan plan = runner.createPlan(session, query);
Plan plan = runner.createPlan(session, query());
assertPlan(
session,
metadata,
Expand Down Expand Up @@ -601,7 +616,7 @@ private QueryAssert hasPlan(PlanMatchPattern expectedPlan, Consumer<Plan> additi
Metadata metadata = runner.getPlannerContext().getMetadata();
transaction(runner.getTransactionManager(), metadata, runner.getAccessControl())
.execute(session, session -> {
Plan plan = runner.createPlan(session, query);
Plan plan = runner.createPlan(session, query());
assertPlan(
session,
metadata,
Expand All @@ -623,7 +638,7 @@ private QueryAssert verifyPlan(Consumer<Plan> planVerification)
{
transaction(runner.getTransactionManager(), runner.getPlannerContext().getMetadata(), runner.getAccessControl())
.execute(session, session -> {
Plan plan = runner.createPlan(session, query);
Plan plan = runner.createPlan(session, query());
planVerification.accept(plan);
});

Expand All @@ -640,9 +655,14 @@ public QueryAssert hasCorrectResultsRegardlessOfPushdown()
Session withoutPushdown = Session.builder(session)
.setSystemProperty("allow_pushdown_into_connectors", "false")
.build();
matches(runner.execute(withoutPushdown, query));
matches(runner.execute(withoutPushdown, query()));
return this;
}

private String query()
{
return query.orElseThrow(() -> new IllegalStateException("Original query is not available"));
}
}

public static class ExpressionAssertProvider
Expand Down

0 comments on commit baaef6d

Please sign in to comment.