Fix events collector

core/trino-main/src/main/java/io/trino/testing/MaterializedResultWithQueryId.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.testing;
+
+import io.trino.spi.QueryId;
+
+import static java.util.Objects.requireNonNull;
+
+public class MaterializedResultWithQueryId
+{
+    private final QueryId queryId;
+    private final MaterializedResult result;
+
+    public MaterializedResultWithQueryId(QueryId queryId, MaterializedResult result)
+    {
+        this.queryId = requireNonNull(queryId, "queryId is null");
+        this.result = requireNonNull(result, "result is null");
+    }
+
+    public QueryId getQueryId()
+    {
+        return queryId;
+    }
+
+    public MaterializedResult getResult()
+    {
+        return result;
+    }
+}

core/trino-main/src/main/java/io/trino/testing/QueryFailedException.java

@@ -0,0 +1,39 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.testing;
+
+import io.trino.spi.QueryId;
+
+public class QueryFailedException
+        extends RuntimeException
+{
+    private final QueryId queryId;
+
+    public QueryFailedException(QueryId queryId, String message)
+    {
+        super(message);
+        this.queryId = queryId;
+    }
+
+    public QueryFailedException(QueryId queryId, String message, Throwable cause)
+    {
+        super(message, cause);
+        this.queryId = queryId;
+    }
+
+    public QueryId getQueryId()
+    {
+        return queryId;
+    }
+}

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java

@@ -30,8 +30,8 @@
 import io.trino.testing.BaseConnectorSmokeTest;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.MaterializedResult;
+import io.trino.testing.MaterializedResultWithQueryId;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import io.trino.testing.TestingConnectorBehavior;
 import io.trino.tpch.TpchTable;
 import org.intellij.lang.annotations.Language;
@@ -327,9 +327,9 @@ public void testInputDataSize()
                         hiveTableName,
                         getLocationForTable(bucketName, "foo")));
 
-        ResultWithQueryId<MaterializedResult> deltaResult = queryRunner.executeWithQueryId(broadcastJoinDistribution(true), "SELECT * FROM foo");
+        MaterializedResultWithQueryId deltaResult = queryRunner.executeWithQueryId(broadcastJoinDistribution(true), "SELECT * FROM foo");
         assertEquals(deltaResult.getResult().getRowCount(), 2);
-        ResultWithQueryId<MaterializedResult> hiveResult = queryRunner.executeWithQueryId(broadcastJoinDistribution(true), format("SELECT * FROM %s.%s.%s", "hive", SCHEMA, hiveTableName));
+        MaterializedResultWithQueryId hiveResult = queryRunner.executeWithQueryId(broadcastJoinDistribution(true), format("SELECT * FROM %s.%s.%s", "hive", SCHEMA, hiveTableName));
         assertEquals(hiveResult.getResult().getRowCount(), 2);
 
         QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
@@ -1544,7 +1544,7 @@ private void invalidateMetadataCache(String tableName)
 
     private void testCountQuery(@Language("SQL") String sql, long expectedRowCount, long expectedSplitCount)
     {
-        ResultWithQueryId<MaterializedResult> result = getDistributedQueryRunner().executeWithQueryId(getSession(), sql);
+        MaterializedResultWithQueryId result = getDistributedQueryRunner().executeWithQueryId(getSession(), sql);
         assertEquals(result.getResult().getOnlyColumnAsSet(), ImmutableSet.of(expectedRowCount));
         verifySplitCount(result.getQueryId(), expectedSplitCount);
     }

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeMinioConnectorTest.java

@@ -23,9 +23,9 @@
 import io.trino.testing.BaseConnectorTest;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.MaterializedResult;
+import io.trino.testing.MaterializedResultWithQueryId;
 import io.trino.testing.MaterializedRow;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import io.trino.testing.TestingConnectorBehavior;
 import io.trino.testing.sql.TestTable;
 import io.trino.tpch.TpchTable;
@@ -336,7 +336,7 @@ public void testTimestampPredicatePushdown(String value)
         assertUpdate("INSERT INTO " + tableName + " VALUES (TIMESTAMP '" + value + "')", 1);
 
         DistributedQueryRunner queryRunner = (DistributedQueryRunner) getQueryRunner();
-        ResultWithQueryId<MaterializedResult> queryResult = queryRunner.executeWithQueryId(
+        MaterializedResultWithQueryId queryResult = queryRunner.executeWithQueryId(
                 getSession(),
                 "SELECT * FROM " + tableName + " WHERE t < TIMESTAMP '" + value + "'");
         assertEquals(getQueryInfo(queryRunner, queryResult).getQueryStats().getProcessedInputDataSize().toBytes(), 0);
@@ -394,7 +394,7 @@ public void testAddColumnToPartitionedTable()
         }
     }
 
-    private QueryInfo getQueryInfo(DistributedQueryRunner queryRunner, ResultWithQueryId<MaterializedResult> queryResult)
+    private QueryInfo getQueryInfo(DistributedQueryRunner queryRunner, MaterializedResultWithQueryId queryResult)
     {
         return queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryResult.getQueryId());
     }

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeAnalyze.java

@@ -19,9 +19,8 @@
 import io.trino.plugin.hive.containers.HiveMinioDataLake;
 import io.trino.spi.QueryId;
 import io.trino.testing.AbstractTestQueryFramework;
-import io.trino.testing.MaterializedResult;
+import io.trino.testing.MaterializedResultWithQueryId;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import io.trino.testing.sql.TestTable;
 import org.testng.annotations.Test;
 
@@ -498,7 +497,7 @@ public void testDropStatsAccessControl()
 
     private void runAnalyzeVerifySplitCount(String tableName, long expectedSplitCount)
     {
-        ResultWithQueryId<MaterializedResult> analyzeResult = getDistributedQueryRunner().executeWithQueryId(getSession(), "ANALYZE " + tableName);
+        MaterializedResultWithQueryId analyzeResult = getDistributedQueryRunner().executeWithQueryId(getSession(), "ANALYZE " + tableName);
         verifySplitCount(analyzeResult.getQueryId(), expectedSplitCount);
     }
 

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeDynamicFiltering.java

@@ -32,9 +32,8 @@
 import io.trino.split.SplitSource;
 import io.trino.sql.planner.OptimizerConfig.JoinDistributionType;
 import io.trino.testing.AbstractTestQueryFramework;
-import io.trino.testing.MaterializedResult;
+import io.trino.testing.MaterializedResultWithQueryId;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import io.trino.transaction.TransactionId;
 import io.trino.transaction.TransactionManager;
 import org.testng.annotations.DataProvider;
@@ -107,8 +106,8 @@ public Object[][] joinDistributionTypes()
     public void testDynamicFiltering(JoinDistributionType joinDistributionType)
     {
         String query = "SELECT * FROM lineitem JOIN orders ON lineitem.orderkey = orders.orderkey AND orders.totalprice > 59995 AND orders.totalprice < 60000";
-        ResultWithQueryId<MaterializedResult> filteredResult = getDistributedQueryRunner().executeWithQueryId(sessionWithDynamicFiltering(true, joinDistributionType), query);
-        ResultWithQueryId<MaterializedResult> unfilteredResult = getDistributedQueryRunner().executeWithQueryId(sessionWithDynamicFiltering(false, joinDistributionType), query);
+        MaterializedResultWithQueryId filteredResult = getDistributedQueryRunner().executeWithQueryId(sessionWithDynamicFiltering(true, joinDistributionType), query);
+        MaterializedResultWithQueryId unfilteredResult = getDistributedQueryRunner().executeWithQueryId(sessionWithDynamicFiltering(false, joinDistributionType), query);
         assertEqualsIgnoreOrder(filteredResult.getResult().getMaterializedRows(), unfilteredResult.getResult().getMaterializedRows());
 
         QueryInputStats filteredStats = getQueryInputStats(filteredResult.getQueryId());

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestPredicatePushdown.java

@@ -19,9 +19,9 @@
 import io.trino.spi.QueryId;
 import io.trino.testing.AbstractTestQueryFramework;
 import io.trino.testing.MaterializedResult;
+import io.trino.testing.MaterializedResultWithQueryId;
 import io.trino.testing.MaterializedRow;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import org.testng.annotations.Test;
 import org.testng.asserts.SoftAssert;
 
@@ -143,7 +143,7 @@ public void testUpdatePushdown()
      */
     private void assertPushdown(String actual, String expected, long countProcessed)
     {
-        ResultWithQueryId<MaterializedResult> result = executeWithQueryId(actual);
+        MaterializedResultWithQueryId result = executeWithQueryId(actual);
         Set<MaterializedRow> actualRows = Set.copyOf(result.getResult().getMaterializedRows());
         Set<MaterializedRow> expectedRows = Set.copyOf(
                 computeExpected(expected, result.getResult().getTypes()).getMaterializedRows());
@@ -176,7 +176,7 @@ private void assertPushdown(String actual, String expected, long countProcessed)
      */
     private void assertPushdownUpdate(String sql, long count, long countProcessed)
     {
-        ResultWithQueryId<MaterializedResult> result = executeWithQueryId(sql);
+        MaterializedResultWithQueryId result = executeWithQueryId(sql);
         OptionalLong actualCount = result.getResult().getUpdateCount();
 
         SoftAssert softly = new SoftAssert();
@@ -189,7 +189,7 @@ private void assertPushdownUpdate(String sql, long count, long countProcessed)
         softly.assertAll();
     }
 
-    private ResultWithQueryId<MaterializedResult> executeWithQueryId(String sql)
+    private MaterializedResultWithQueryId executeWithQueryId(String sql)
     {
         return getDistributedQueryRunner().executeWithQueryId(getSession(), sql);
     }

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestSplitPruning.java

@@ -20,7 +20,6 @@
 import io.trino.testing.AbstractTestQueryFramework;
 import io.trino.testing.MaterializedResult;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import org.intellij.lang.annotations.Language;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.DataProvider;
@@ -135,10 +134,10 @@ public void testStatsPruningNaN(String type)
                 Set.of(),
                 0);
 
-        ResultWithQueryId<MaterializedResult> result = getDistributedQueryRunner().executeWithQueryId(
+        MaterializedResult result = getDistributedQueryRunner().execute(
                 getSession(),
                 format("SELECT name FROM %s WHERE val IS NOT NULL", tableName));
-        assertEquals(result.getResult().getOnlyColumnAsSet(), Set.of("a5", "b5", "a6", "b6"));
+        assertEquals(result.getOnlyColumnAsSet(), Set.of("a5", "b5", "a6", "b6"));
     }
 
     @Test

plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java

@@ -53,9 +53,9 @@
 import io.trino.testing.BaseConnectorTest;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.MaterializedResult;
+import io.trino.testing.MaterializedResultWithQueryId;
 import io.trino.testing.MaterializedRow;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import io.trino.testing.TestingConnectorBehavior;
 import io.trino.testing.sql.TestTable;
 import io.trino.testing.sql.TrinoSqlExecutor;
@@ -4725,7 +4725,7 @@ private void doTestParquetTimestampPredicatePushdown(Session baseSession, HiveTi
         assertQuery(session, "SELECT * FROM " + tableName, format("VALUES (%s)", formatTimestamp(value)));
 
         DistributedQueryRunner queryRunner = (DistributedQueryRunner) getQueryRunner();
-        ResultWithQueryId<MaterializedResult> queryResult = queryRunner.executeWithQueryId(
+        MaterializedResultWithQueryId queryResult = queryRunner.executeWithQueryId(
                 session,
                 format("SELECT * FROM %s WHERE t < %s", tableName, formatTimestamp(value)));
         assertEquals(getQueryInfo(queryRunner, queryResult).getQueryStats().getProcessedInputDataSize().toBytes(), 0);
@@ -4756,7 +4756,7 @@ public void testOrcTimestampPredicatePushdown(HiveTimestampPrecision timestampPr
         // to account for the fact that ORC stats are stored at millisecond precision and Trino rounds timestamps,
         // we filter by timestamps that differ from the actual value by at least 1ms, to observe pruning
         DistributedQueryRunner queryRunner = getDistributedQueryRunner();
-        ResultWithQueryId<MaterializedResult> queryResult = queryRunner.executeWithQueryId(
+        MaterializedResultWithQueryId queryResult = queryRunner.executeWithQueryId(
                 session,
                 format("SELECT * FROM test_orc_timestamp_predicate_pushdown WHERE t < %s", formatTimestamp(value.minusNanos(MILLISECONDS.toNanos(1)))));
         assertEquals(getQueryInfo(queryRunner, queryResult).getQueryStats().getProcessedInputDataSize().toBytes(), 0);
@@ -4845,7 +4845,7 @@ private void assertNoDataRead(@Language("SQL") String sql)
                 results -> assertThat(results.getRowCount()).isEqualTo(0));
     }
 
-    private QueryInfo getQueryInfo(DistributedQueryRunner queryRunner, ResultWithQueryId<MaterializedResult> queryResult)
+    private QueryInfo getQueryInfo(DistributedQueryRunner queryRunner, MaterializedResultWithQueryId queryResult)
     {
         return queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryResult.getQueryId());
     }

plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveDistributedJoinQueries.java

@@ -18,9 +18,8 @@
 import io.trino.metadata.QualifiedObjectName;
 import io.trino.operator.OperatorStats;
 import io.trino.testing.AbstractTestJoinQueries;
-import io.trino.testing.MaterializedResult;
+import io.trino.testing.MaterializedResultWithQueryId;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import org.testng.annotations.Test;
 
 import static com.google.common.base.Verify.verify;
@@ -59,7 +58,7 @@ public void testJoinWithEmptyBuildSide()
         Session session = Session.builder(getSession())
                 .setSystemProperty(JOIN_DISTRIBUTION_TYPE, BROADCAST.name())
                 .build();
-        ResultWithQueryId<MaterializedResult> result = getDistributedQueryRunner().executeWithQueryId(
+        MaterializedResultWithQueryId result = getDistributedQueryRunner().executeWithQueryId(
                 session,
                 "SELECT * FROM lineitem JOIN orders ON lineitem.orderkey = orders.orderkey AND orders.totalprice = 123.4567");
         assertEquals(result.getResult().getRowCount(), 0);

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java

@@ -36,9 +36,9 @@
 import io.trino.testing.BaseConnectorTest;
 import io.trino.testing.DataProviders;
 import io.trino.testing.MaterializedResult;
+import io.trino.testing.MaterializedResultWithQueryId;
 import io.trino.testing.MaterializedRow;
 import io.trino.testing.QueryRunner;
-import io.trino.testing.ResultWithQueryId;
 import io.trino.testing.TestingConnectorBehavior;
 import io.trino.testing.sql.TestTable;
 import io.trino.tpch.TpchTable;
@@ -3277,15 +3277,15 @@ public void testSplitPruningFromDataFileStatistics(DataMappingTestSetup testSetu
 
     private void verifySplitCount(String query, int expectedSplitCount)
     {
-        ResultWithQueryId<MaterializedResult> selectAllPartitionsResult = getDistributedQueryRunner().executeWithQueryId(getSession(), query);
+        MaterializedResultWithQueryId selectAllPartitionsResult = getDistributedQueryRunner().executeWithQueryId(getSession(), query);
         assertEqualsIgnoreOrder(selectAllPartitionsResult.getResult().getMaterializedRows(), computeActual(withoutPredicatePushdown(getSession()), query).getMaterializedRows());
         verifySplitCount(selectAllPartitionsResult.getQueryId(), expectedSplitCount);
     }
 
     private void verifyPredicatePushdownDataRead(@Language("SQL") String query, boolean supportsPushdown)
     {
-        ResultWithQueryId<MaterializedResult> resultWithPredicatePushdown = getDistributedQueryRunner().executeWithQueryId(getSession(), query);
-        ResultWithQueryId<MaterializedResult> resultWithoutPredicatePushdown = getDistributedQueryRunner().executeWithQueryId(
+        MaterializedResultWithQueryId resultWithPredicatePushdown = getDistributedQueryRunner().executeWithQueryId(getSession(), query);
+        MaterializedResultWithQueryId resultWithoutPredicatePushdown = getDistributedQueryRunner().executeWithQueryId(
                 withoutPredicatePushdown(getSession()),
                 query);