Skip to content

Commit de4be2d

Browse files
bhzaeribhzaeri
committed
add unit test for the added updatedRows to QueryStatistics.java
The test (for UPDATE/DELETE queries) was added to BaseConnectorTest.java and it needs the classes like EventsAwaitingQueries and QueryEvents that exist in trino-tests module. So, I moved the needed classes from trino-tests to trino-testing where they are exposed to both trino-tests and trino-testing.
1 parent aa336ca commit de4be2d

File tree

20 files changed

+532
-287
lines changed

20 files changed

+532
-287
lines changed

Diff for: ‎core/trino-main/src/test/java/io/trino/operator/output/TestPartitionedOutputOperator.java

+1
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ public void testOperatorContextStats()
6969

7070
OperatorContext operatorContext = partitionedOutputOperator.getOperatorContext();
7171
assertThat(operatorContext.getOutputDataSize().getTotalCount()).isEqualTo(0);
72+
assertThat(operatorContext.getUpdatedPositions().getTotalCount()).isEqualTo(0);
7273
assertThat(operatorContext.getOutputPositions().getTotalCount()).isEqualTo(page.getPositionCount());
7374

7475
partitionedOutputOperator.finish();

Diff for: ‎core/trino-spi/pom.xml

+2-2
Original file line numberDiff line numberDiff line change
@@ -193,9 +193,9 @@
193193
<classSimpleName>QueryStatistics</classSimpleName>
194194
<methodName>&lt;init&gt;</methodName>
195195
<elementKind>constructor</elementKind>
196-
<oldArchive>io.trino:trino-spi:jar:468</oldArchive>
196+
<oldArchive>io.trino:trino-spi:jar:469</oldArchive>
197197
<oldArchiveRole>primary</oldArchiveRole>
198-
<newArchive>io.trino:trino-spi:jar:469-SNAPSHOT</newArchive>
198+
<newArchive>io.trino:trino-spi:jar:470-SNAPSHOT</newArchive>
199199
<newArchiveRole>primary</newArchiveRole>
200200
</item>
201201
<!-- Allow adding new annotations (such as @Nullable) -->

Diff for: ‎plugin/trino-cassandra/src/test/java/io/trino/plugin/cassandra/TestCassandraConnectorTest.java

+36
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import io.trino.testing.MaterializedRow;
2424
import io.trino.testing.QueryRunner;
2525
import io.trino.testing.TestingConnectorBehavior;
26+
import io.trino.testing.eventlistener.NamedClosable;
2627
import io.trino.testing.sql.TestTable;
2728
import org.intellij.lang.annotations.Language;
2829
import org.junit.jupiter.api.AfterAll;
@@ -39,6 +40,7 @@
3940
import java.util.List;
4041
import java.util.Optional;
4142
import java.util.OptionalInt;
43+
import java.util.function.Supplier;
4244

4345
import static com.datastax.oss.driver.api.core.data.ByteUtils.toHexString;
4446
import static com.google.common.io.BaseEncoding.base16;
@@ -1573,6 +1575,40 @@ public void testDelete()
15731575
}
15741576
}
15751577

1578+
@Test
1579+
@Override
1580+
public void testDeleteStatsWithRaisedEvents()
1581+
{
1582+
Supplier<NamedClosable> supplier = () -> {
1583+
TestCassandraTable testCassandraTable = testTable(
1584+
"table_delete_data",
1585+
ImmutableList.of(
1586+
partitionColumn("partition_one", "bigint"),
1587+
partitionColumn("partition_two", "int"),
1588+
clusterColumn("clust_one", "text"),
1589+
generalColumn("data", "text")),
1590+
ImmutableList.of(
1591+
"1, 1, 'clust_one_1', null",
1592+
"2, 2, 'clust_one_2', null",
1593+
"3, 3, 'clust_one_3', null",
1594+
"4, 4, 'clust_one_4', null",
1595+
"5, 5, 'clust_one_5', null",
1596+
"6, 6, 'clust_one_6', null",
1597+
"7, 7, 'clust_one_7', null",
1598+
"8, 8, 'clust_one_8', null",
1599+
"9, 9, 'clust_one_9', null",
1600+
"1, 1, 'clust_one_2', null",
1601+
"1, 1, 'clust_one_3', null",
1602+
"1, 2, 'clust_one_1', null",
1603+
"1, 2, 'clust_one_2', null",
1604+
"1, 2, 'clust_one_3', null",
1605+
"2, 2, 'clust_one_1', null"));
1606+
return new NamedClosable(testCassandraTable.getTableName(), testCassandraTable);
1607+
};
1608+
String wherePrimaryKey = " WHERE partition_one=3 AND partition_two=3 AND clust_one='clust_one_3'";
1609+
runUpdateDeleteStatsWithRaisedEvents(supplier, name -> "DELETE FROM " + name + wherePrimaryKey);
1610+
}
1611+
15761612
@Test
15771613
@Override
15781614
public void testDeleteWithLike()

Diff for: ‎plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java

+16
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,14 @@ public void testDelete()
304304
.hasStackTraceContaining(MODIFYING_NON_TRANSACTIONAL_TABLE_MESSAGE);
305305
}
306306

307+
@Test
308+
@Override
309+
public void testDeleteStatsWithRaisedEvents()
310+
{
311+
assertThatThrownBy(super::testDeleteStatsWithRaisedEvents)
312+
.hasStackTraceContaining(MODIFYING_NON_TRANSACTIONAL_TABLE_MESSAGE);
313+
}
314+
307315
@Test
308316
@Override
309317
public void testDeleteWithLike()
@@ -352,6 +360,14 @@ public void testUpdateMultipleCondition()
352360
.hasMessage(MODIFYING_NON_TRANSACTIONAL_TABLE_MESSAGE);
353361
}
354362

363+
@Test
364+
@Override
365+
public void testUpdateStatsWithRaisedEvents()
366+
{
367+
assertThatThrownBy(super::testUpdateStatsWithRaisedEvents)
368+
.hasMessage(MODIFYING_NON_TRANSACTIONAL_TABLE_MESSAGE);
369+
}
370+
355371
@Test
356372
@Override
357373
public void testUpdateWithNullValues()

Diff for: ‎plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestHiveOrcWithShortZoneId.java

+1
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ public void testSelectWithoutShortZoneId()
7878
"test_select_without_short_zone_id_",
7979
"(id INT, firstName VARCHAR, lastName VARCHAR)",
8080
ImmutableList.of("2, 'Alice', 'Doe'"))) {
81+
getQueryRunner().execute("update " + testTable.getName() + " set firstName = 'Ali' where id = 2");
8182
assertQuery("SELECT * FROM " + testTable.getName(), "VALUES (2, 'Alice', 'Doe')");
8283
}
8384
}

Diff for: ‎plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java

+40
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,15 @@
1919
import io.trino.testing.MaterializedResult;
2020
import io.trino.testing.QueryRunner;
2121
import io.trino.testing.TestingConnectorBehavior;
22+
import io.trino.testing.eventlistener.NamedClosable;
2223
import io.trino.testing.sql.TestTable;
2324
import org.junit.jupiter.api.Disabled;
2425
import org.junit.jupiter.api.Test;
2526

2627
import java.util.Optional;
2728
import java.util.OptionalInt;
2829
import java.util.function.Consumer;
30+
import java.util.function.Supplier;
2931
import java.util.regex.Matcher;
3032
import java.util.regex.Pattern;
3133

@@ -706,6 +708,25 @@ public void testDelete()
706708
});
707709
}
708710

711+
@Test
712+
@Override
713+
public void testDeleteStatsWithRaisedEvents()
714+
{
715+
Supplier<NamedClosable> supplier = () -> {
716+
String tableName = "%s_%s".formatted("test_delete", randomNameSuffix());
717+
createTableForWrites("CREATE TABLE %s " + ORDER_COLUMNS, tableName, Optional.empty());
718+
assertUpdate("INSERT INTO " + tableName + " SELECT * FROM orders", 15000);
719+
return new NamedClosable(tableName, null) {
720+
@Override
721+
public void close()
722+
{
723+
assertUpdate("DROP TABLE " + tableName);
724+
}
725+
};
726+
};
727+
runUpdateDeleteStatsWithRaisedEvents(supplier, table -> "DELETE FROM " + table + " WHERE custkey <= 100");
728+
}
729+
709730
@Test
710731
@Override
711732
public void testDeleteWithLike()
@@ -970,6 +991,25 @@ public void testUpdate()
970991
});
971992
}
972993

994+
@Test
995+
@Override
996+
public void testUpdateStatsWithRaisedEvents()
997+
{
998+
Supplier<NamedClosable> supplier = () -> {
999+
String tableName = "%s_%s".formatted("test_update", randomNameSuffix());
1000+
createTableForWrites("CREATE TABLE %s " + NATION_COLUMNS, tableName, Optional.empty());
1001+
assertUpdate("INSERT INTO " + tableName + " SELECT * FROM nation", 25);
1002+
return new NamedClosable(tableName, null) {
1003+
@Override
1004+
public void close()
1005+
{
1006+
assertUpdate("DROP TABLE " + tableName);
1007+
}
1008+
};
1009+
};
1010+
runUpdateDeleteStatsWithRaisedEvents(supplier, table -> "UPDATE " + table + " SET regionkey = 100 WHERE regionkey = 2");
1011+
}
1012+
9731013
/**
9741014
* This test fails intermittently because Kudu doesn't have strong enough
9751015
* semantics to support writing from multiple threads.

Diff for: ‎plugin/trino-redshift/src/test/java/io/trino/plugin/redshift/TestRedshiftConnectorTest.java

+13
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import io.trino.sql.planner.plan.TableScanNode;
3131
import io.trino.testing.QueryRunner;
3232
import io.trino.testing.TestingConnectorBehavior;
33+
import io.trino.testing.eventlistener.NamedClosable;
3334
import io.trino.testing.sql.SqlExecutor;
3435
import io.trino.testing.sql.TestTable;
3536
import org.jdbi.v3.core.Handle;
@@ -43,6 +44,7 @@
4344
import java.util.OptionalInt;
4445
import java.util.Set;
4546
import java.util.concurrent.ScheduledThreadPoolExecutor;
47+
import java.util.function.Supplier;
4648
import java.util.stream.Collectors;
4749
import java.util.stream.IntStream;
4850
import java.util.stream.Stream;
@@ -339,6 +341,17 @@ public void testDelete()
339341
}
340342
}
341343

344+
@Test
345+
@Override
346+
public void testDeleteStatsWithRaisedEvents() {
347+
Supplier<NamedClosable> supplier = () -> {
348+
TestTable table = newTrinoTable("test_delete_", "AS SELECT * FROM nation");
349+
return new NamedClosable(table.getName(), table);
350+
};
351+
runUpdateDeleteStatsWithRaisedEvents(supplier,
352+
table -> "DELETE FROM " + table + " WHERE nationkey <= 5");
353+
}
354+
342355
@Test
343356
public void testCaseColumnNames()
344357
{

Diff for: ‎testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java

+62-1
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,15 @@
2727
import io.trino.dispatcher.DispatchManager;
2828
import io.trino.execution.QueryInfo;
2929
import io.trino.execution.QueryManager;
30+
import io.trino.execution.QueryStats;
3031
import io.trino.metadata.FunctionManager;
3132
import io.trino.metadata.Metadata;
3233
import io.trino.metadata.QualifiedObjectName;
3334
import io.trino.server.BasicQueryInfo;
35+
import io.trino.spi.QueryId;
3436
import io.trino.spi.connector.ConnectorSession;
3537
import io.trino.spi.connector.MaterializedViewFreshness;
38+
import io.trino.spi.eventlistener.QueryCompletedEvent;
3639
import io.trino.spi.security.Identity;
3740
import io.trino.sql.planner.OptimizerConfig.JoinDistributionType;
3841
import io.trino.sql.planner.Plan;
@@ -45,6 +48,11 @@
4548
import io.trino.sql.query.QueryAssertions.QueryAssert;
4649
import io.trino.testing.QueryRunner.MaterializedResultWithPlan;
4750
import io.trino.testing.assertions.TrinoExceptionAssert;
51+
import io.trino.testing.eventlistener.EventsAwaitingQueries;
52+
import io.trino.testing.eventlistener.EventsCollector;
53+
import io.trino.testing.eventlistener.NamedClosable;
54+
import io.trino.testing.eventlistener.QueryEvents;
55+
import io.trino.testing.eventlistener.TestingEventListenerPlugin;
4856
import io.trino.testing.sql.TestTable;
4957
import io.trino.testing.sql.TestView;
5058
import io.trino.tpch.TpchTable;
@@ -68,6 +76,7 @@
6876
import java.util.concurrent.CyclicBarrier;
6977
import java.util.concurrent.ExecutorService;
7078
import java.util.concurrent.Future;
79+
import java.util.concurrent.TimeoutException;
7180
import java.util.concurrent.atomic.AtomicReference;
7281
import java.util.function.Consumer;
7382
import java.util.function.Function;
@@ -186,7 +195,7 @@ public abstract class BaseConnectorTest
186195
{
187196
private static final Logger log = Logger.get(BaseConnectorTest.class);
188197

189-
protected static final List<TpchTable<?>> REQUIRED_TPCH_TABLES = ImmutableSet.<TpchTable<?>>builder()
198+
public static final List<TpchTable<?>> REQUIRED_TPCH_TABLES = ImmutableSet.<TpchTable<?>>builder()
190199
.addAll(AbstractTestQueries.REQUIRED_TPCH_TABLES)
191200
.add(CUSTOMER)
192201
.build().asList();
@@ -5030,6 +5039,58 @@ public void testUpdateMultipleCondition()
50305039
}
50315040
}
50325041

5042+
@Test
5043+
public void testUpdateStatsWithRaisedEvents()
5044+
{
5045+
if (!hasBehavior(SUPPORTS_UPDATE)) {
5046+
// Note this change is a no-op, if actually run
5047+
assertQueryFails("UPDATE nation SET nationkey = nationkey + regionkey WHERE regionkey < 1", MODIFYING_ROWS_MESSAGE);
5048+
return;
5049+
}
5050+
Supplier<NamedClosable> supplier = () -> {
5051+
TestTable table = newTrinoTable("test_row_update", "AS SELECT * FROM nation");
5052+
return new NamedClosable(table.getName(), table);
5053+
};
5054+
runUpdateDeleteStatsWithRaisedEvents(supplier,
5055+
table -> "UPDATE " + table + " SET nationkey = 100 WHERE regionkey = 2");
5056+
}
5057+
5058+
@Test
5059+
public void testDeleteStatsWithRaisedEvents()
5060+
{
5061+
skipTestUnless(hasBehavior(SUPPORTS_DELETE));
5062+
// delete successive parts of the table
5063+
Supplier<NamedClosable> supplier = () -> {
5064+
TestTable table = newTrinoTable("test_delete_", "AS SELECT * FROM orders");
5065+
return new NamedClosable(table.getName(), table);
5066+
};
5067+
runUpdateDeleteStatsWithRaisedEvents(supplier,
5068+
table -> "DELETE FROM " + table + " WHERE custkey <= 100");
5069+
}
5070+
5071+
protected void runUpdateDeleteStatsWithRaisedEvents(Supplier<NamedClosable> namedClosable,
5072+
Function<String, String> querySupplier) {
5073+
EventsCollector generatedEvents = new EventsCollector();
5074+
EventsAwaitingQueries queries = new EventsAwaitingQueries(generatedEvents, getQueryRunner());
5075+
getQueryRunner().installPlugin(new TestingEventListenerPlugin(generatedEvents));
5076+
try(NamedClosable table = namedClosable.get()) {
5077+
EventsAwaitingQueries.MaterializedResultWithEvents result = queries.runQueryAndWaitForEvents(querySupplier.apply(table.getName()), getSession());
5078+
QueryEvents queryEvents = result.getQueryEvents();
5079+
try {
5080+
queryEvents.waitForQueryCompletion(new Duration(300, SECONDS));
5081+
SECONDS.sleep(1);
5082+
QueryCompletedEvent c = queryEvents.getQueryCompletedEvent();
5083+
QueryStats queryStats = getDistributedQueryRunner().getCoordinator().getQueryManager().getFullQueryInfo(new QueryId(c.getMetadata().getQueryId())).getQueryStats();
5084+
assertThat(result.getMaterializedResult().getUpdateCount().orElse(0)).isEqualTo(queryStats.getUpdatedPositions());
5085+
} catch (InterruptedException | TimeoutException e) {
5086+
throw new RuntimeException(e);
5087+
}
5088+
} catch (Exception e) {
5089+
throw new RuntimeException(e);
5090+
}
5091+
}
5092+
5093+
50335094
@Test
50345095
public void testUpdateWithNullValues()
50355096
{

Diff for: ‎testing/trino-tests/src/test/java/io/trino/execution/EventsAwaitingQueries.java renamed to ‎testing/trino-testing/src/main/java/io/trino/testing/eventlistener/EventsAwaitingQueries.java

+6-7
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,10 @@
1111
* See the License for the specific language governing permissions and
1212
* limitations under the License.
1313
*/
14-
package io.trino.execution;
14+
package io.trino.testing.eventlistener;
1515

1616
import io.airlift.units.Duration;
1717
import io.trino.Session;
18-
import io.trino.execution.EventsCollector.QueryEvents;
1918
import io.trino.spi.QueryId;
2019
import io.trino.testing.MaterializedResult;
2120
import io.trino.testing.QueryFailedException;
@@ -31,32 +30,32 @@
3130
import static java.util.concurrent.TimeUnit.SECONDS;
3231
import static org.assertj.core.api.Fail.fail;
3332

34-
class EventsAwaitingQueries
33+
public class EventsAwaitingQueries
3534
{
3635
private final EventsCollector eventsCollector;
3736

3837
private final QueryRunner queryRunner;
3938

40-
EventsAwaitingQueries(EventsCollector eventsCollector, QueryRunner queryRunner)
39+
public EventsAwaitingQueries(EventsCollector eventsCollector, QueryRunner queryRunner)
4140
{
4241
this.eventsCollector = requireNonNull(eventsCollector, "eventsCollector is null");
4342
this.queryRunner = requireNonNull(queryRunner, "queryRunner is null");
4443
}
4544

46-
MaterializedResultWithEvents runQueryAndWaitForEvents(@Language("SQL") String sql, Session session)
45+
public MaterializedResultWithEvents runQueryAndWaitForEvents(@Language("SQL") String sql, Session session)
4746
throws Exception
4847
{
4948
return runQueryAndWaitForEvents(sql, session, Optional.empty());
5049
}
5150

52-
MaterializedResultWithEvents runQueryAndWaitForEvents(@Language("SQL") String sql, Session session, boolean requireAnonymizedPlan)
51+
public MaterializedResultWithEvents runQueryAndWaitForEvents(@Language("SQL") String sql, Session session, boolean requireAnonymizedPlan)
5352
throws Exception
5453
{
5554
eventsCollector.setRequiresAnonymizedPlan(requireAnonymizedPlan);
5655
return runQueryAndWaitForEvents(sql, session, Optional.empty());
5756
}
5857

59-
MaterializedResultWithEvents runQueryAndWaitForEvents(@Language("SQL") String sql, Session session, Optional<String> expectedExceptionRegEx)
58+
public MaterializedResultWithEvents runQueryAndWaitForEvents(@Language("SQL") String sql, Session session, Optional<String> expectedExceptionRegEx)
6059
throws Exception
6160
{
6261
QueryId queryId = null;

0 commit comments

Comments
 (0)