Refactor test framework to return transactionId control flags to client

hantangwangd · hantangwangd · commit b016faf6cd8c · 2025-08-27T08:09:06.000+08:00
The test framework client now receives statement executing results with
`clearTransactionId` and `startTransactionId` flags embedded.
diff --git a/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java b/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java
@@ -119,7 +119,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
         }
 
         @Override
-        public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties)
+        public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties, String startTransactionId, boolean clearTransactionId)
         {
             return null;
         }
diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java
@@ -115,6 +115,12 @@ public void testDescribeOutputNamedAndUnnamed()
         assertEqualsIgnoreOrder(actual, expected);
     }
 
+    @Override
+    public void testClearTransactionId()
+    {
+        // Catalog iceberg only supports writes using autocommit
+    }
+
     /**
      * Increased the optimizer timeout from 15000ms to 25000ms
      */
diff --git a/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java b/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java
@@ -120,7 +120,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
         }
 
         @Override
-        public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties)
+        public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties, String startTransactionId, boolean clearTransactionId)
         {
             return null;
         }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java b/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.common.PageBuilder;
 import com.facebook.presto.common.block.Block;
 import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.common.transaction.TransactionId;
 import com.facebook.presto.common.type.ArrayType;
 import com.facebook.presto.common.type.CharType;
 import com.facebook.presto.common.type.MapType;
@@ -99,11 +100,13 @@ public class MaterializedResult
     private final Set<String> resetSessionProperties;
     private final Optional<String> updateType;
     private final OptionalLong updateCount;
+    private final Optional<TransactionId> startedTransactionId;
+    private final boolean clearTransactionId;
     private final List<PrestoWarning> warnings;
 
     public MaterializedResult(List<MaterializedRow> rows, List<? extends Type> types)
     {
-        this(rows, types, ImmutableMap.of(), ImmutableSet.of(), Optional.empty(), OptionalLong.empty(), ImmutableList.of());
+        this(rows, types, ImmutableMap.of(), ImmutableSet.of(), Optional.empty(), OptionalLong.empty(), Optional.empty(), false, ImmutableList.of());
     }
 
     public MaterializedResult(
@@ -113,6 +116,8 @@ public MaterializedResult(
             Set<String> resetSessionProperties,
             Optional<String> updateType,
             OptionalLong updateCount,
+            Optional<TransactionId> startedTransactionId,
+            boolean clearTransactionId,
             List<PrestoWarning> warnings)
     {
         this.rows = ImmutableList.copyOf(requireNonNull(rows, "rows is null"));
@@ -121,6 +126,8 @@ public MaterializedResult(
         this.resetSessionProperties = ImmutableSet.copyOf(requireNonNull(resetSessionProperties, "resetSessionProperties is null"));
         this.updateType = requireNonNull(updateType, "updateType is null");
         this.updateCount = requireNonNull(updateCount, "updateCount is null");
+        this.startedTransactionId = requireNonNull(startedTransactionId, "startedTransactionId is null");
+        this.clearTransactionId = clearTransactionId;
         this.warnings = requireNonNull(warnings, "warnings is null");
     }
 
@@ -165,6 +172,16 @@ public OptionalLong getUpdateCount()
         return updateCount;
     }
 
+    public Optional<TransactionId> getStartedTransactionId()
+    {
+        return startedTransactionId;
+    }
+
+    public boolean isClearTransactionId()
+    {
+        return clearTransactionId;
+    }
+
     public List<PrestoWarning> getWarnings()
     {
         return warnings;
@@ -185,13 +202,15 @@ public boolean equals(Object obj)
                 Objects.equals(setSessionProperties, o.setSessionProperties) &&
                 Objects.equals(resetSessionProperties, o.resetSessionProperties) &&
                 Objects.equals(updateType, o.updateType) &&
-                Objects.equals(updateCount, o.updateCount);
+                Objects.equals(updateCount, o.updateCount) &&
+                Objects.equals(startedTransactionId, o.startedTransactionId) &&
+                Objects.equals(clearTransactionId, o.clearTransactionId);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(rows, types, setSessionProperties, resetSessionProperties, updateType, updateCount);
+        return Objects.hash(rows, types, setSessionProperties, resetSessionProperties, updateType, updateCount, startedTransactionId, clearTransactionId);
     }
 
     @Override
@@ -204,6 +223,8 @@ public String toString()
                 .add("resetSessionProperties", resetSessionProperties)
                 .add("updateType", updateType.orElse(null))
                 .add("updateCount", updateCount.isPresent() ? updateCount.getAsLong() : null)
+                .add("startedTransactionId", startedTransactionId.orElse(null))
+                .add("clearTransactionId", clearTransactionId)
                 .omitNullValues()
                 .toString();
     }
@@ -360,6 +381,8 @@ public MaterializedResult toTestTypes()
                 resetSessionProperties,
                 updateType,
                 updateCount,
+                startedTransactionId,
+                clearTransactionId,
                 warnings);
     }
 
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java
@@ -1211,7 +1211,7 @@ public void testSetSessionJavaWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={distinct_aggregation_spill_enabled=false}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 
     @Test
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
@@ -155,7 +155,7 @@ public void testSetNativeWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={driver_cpu_time_slice_limit_ms=500}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 
     @Test
diff --git a/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java b/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java
@@ -143,7 +143,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
         }
 
         @Override
-        public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties)
+        public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties, String startTransactionId, boolean clearTransactionId)
         {
             return null;
         }
diff --git a/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java b/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java
@@ -183,6 +183,12 @@ public void testDescribeOutputNamedAndUnnamed()
         // this connector uses a non-canonical type for varchar columns in tpch
     }
 
+    @Override
+    public void testClearTransactionId()
+    {
+        // Catalog singlestore only supports writes using autocommit
+    }
+
     @Override
     public void testInsert()
     {
diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java
@@ -594,6 +594,8 @@ private MaterializedResult executeWithStrategies(
                         ImmutableSet.of(),
                         p.getUpdateType(),
                         getOnlyElement(getOnlyElement(rows).getFields()) == null ? OptionalLong.empty() : OptionalLong.of((Long) getOnlyElement(getOnlyElement(rows).getFields())),
+                        Optional.empty(),
+                        false,
                         ImmutableList.of());
             }
         }
@@ -606,6 +608,8 @@ else if (execution instanceof PrestoSparkAccessControlCheckerExecution) {
                     ImmutableSet.of(),
                     Optional.empty(),
                     OptionalLong.empty(),
+                    Optional.empty(),
+                    false,
                     ImmutableList.of());
         }
         else {
@@ -616,6 +620,8 @@ else if (execution instanceof PrestoSparkAccessControlCheckerExecution) {
                     ImmutableSet.of(),
                     Optional.empty(),
                     OptionalLong.empty(),
+                    Optional.empty(),
+                    false,
                     ImmutableList.of());
         }
     }
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java
@@ -249,6 +249,41 @@ public void testCreateTable()
         assertFalse(getQueryRunner().tableExists(getSession(), "test_create_like"));
     }
 
+    @Test
+    public void testClearTransactionId()
+    {
+        assertUpdate("create table test_clear_transaction_id_table(a int, b varchar)");
+        assertUpdate("insert into test_clear_transaction_id_table values(1, '1001'), (2, '1002')", 2);
+        Session session = getQueryRunner().getDefaultSession();
+        String defaultCatalog = session.getCatalog().get();
+        transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl())
+                .execute(Session.builder(session)
+                                .setIdentity(new Identity("admin",
+                                        Optional.empty(),
+                                        ImmutableMap.of(defaultCatalog, new SelectedRole(ROLE, Optional.of("admin"))),
+                                        ImmutableMap.of(),
+                                        ImmutableMap.of(),
+                                        Optional.empty(),
+                                        Optional.empty()))
+                                .build(),
+                        txnSession -> {
+                            MaterializedResult result = computeActual(txnSession, "select * from test_clear_transaction_id_table");
+                            assertEquals(result.getRowCount(), 2);
+                            assertFalse(result.isClearTransactionId());
+
+                            result = computeActual(txnSession, "insert into test_clear_transaction_id_table values(1, '1001'), (2, '1002')");
+                            assertEquals(result.getOnlyValue(), 2L);
+                            assertFalse(result.isClearTransactionId());
+
+                            // `Rollback` executes successfully, and the client gets a flag `clearTransactionId = true`
+                            result = computeActual(txnSession, "rollback");
+                            assertTrue(result.isClearTransactionId());
+                        });
+
+        assertQuery("select * from test_clear_transaction_id_table", "values(1, '1001'), (2, '1002')");
+        assertUpdate("drop table if exists test_clear_transaction_id_table");
+    }
+
     @Test
     public void testNonAutoCommitTransactionWithFailAndRollback()
     {
@@ -276,7 +311,8 @@ public void testNonAutoCommitTransactionWithFailAndRollback()
                             assertQueryFails(txnSession, "create table test_table(a int, b varchar)", "Current transaction is aborted, commands ignored until end of transaction block");
 
                             // execute `rollback` successfully
-                            assertUpdate(txnSession, "rollback");
+                            MaterializedResult result = computeActual(txnSession, "rollback");
+                            assertTrue(result.isClearTransactionId());
                         });
 
         assertQuery("select count(*) from test_non_autocommit_table", "values(0)");
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
@@ -3218,7 +3218,8 @@ public void testSetSessionNativeWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={native_expression_max_array_size_in_reduce=50000}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, " +
+                        "clearTransactionId=false}");
     }
 
     @Test
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java
@@ -117,7 +117,7 @@ public ResultWithQueryId<T> execute(Session session, @Language("SQL") String sql
 
                 resultsSession.setWarnings(results.getWarnings());
 
-                T result = resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties());
+                T result = resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties(), client.getStartedTransactionId(), client.isClearTransactionId());
                 return new ResultWithQueryId<>(new QueryId(results.getId()), result);
             }
 
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java b/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java
@@ -16,6 +16,7 @@
 import com.facebook.presto.client.QueryData;
 import com.facebook.presto.client.QueryStatusInfo;
 import com.facebook.presto.spi.PrestoWarning;
+import jakarta.annotation.Nullable;
 
 import java.util.List;
 import java.util.Map;
@@ -40,5 +41,5 @@ default void setWarnings(List<PrestoWarning> warnings)
 
     void addResults(QueryStatusInfo statusInfo, QueryData data);
 
-    T build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties);
+    T build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties, @Nullable String startTransactionId, boolean clearTransactionId);
 }
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java b/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.client.IntervalYearMonth;
 import com.facebook.presto.client.QueryData;
 import com.facebook.presto.client.QueryStatusInfo;
+import com.facebook.presto.common.transaction.TransactionId;
 import com.facebook.presto.common.type.ArrayType;
 import com.facebook.presto.common.type.BigintEnumType;
 import com.facebook.presto.common.type.DecimalType;
@@ -41,6 +42,7 @@
 import com.facebook.presto.type.SqlIntervalYearMonth;
 import com.google.common.base.Function;
 import com.google.common.collect.ImmutableList;
+import jakarta.annotation.Nullable;
 
 import java.math.BigDecimal;
 import java.time.LocalDate;
@@ -147,7 +149,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
         }
 
         @Override
-        public MaterializedResult build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties)
+        public MaterializedResult build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties, @Nullable String startTransactionId, boolean clearTransactionId)
         {
             checkState(types.get() != null, "never received types for the query");
             return new MaterializedResult(
@@ -157,6 +159,8 @@ public MaterializedResult build(Map<String, String> setSessionProperties, Set<St
                     resetSessionProperties,
                     updateType.get(),
                     updateCount.get(),
+                    Optional.ofNullable(startTransactionId).map(TransactionId::new),
+                    clearTransactionId,
                     warnings.get());
         }
     }
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesExcludingInvalidProperties.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesExcludingInvalidProperties.java
@@ -51,6 +51,6 @@ public void testSetSessionValidJavaWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={distinct_aggregation_spill_enabled=false}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 }
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesIncludingInvalidProperties.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesIncludingInvalidProperties.java
@@ -42,7 +42,7 @@ public void testSetSessionValidNativeWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={native_expression_max_array_size_in_reduce=50000}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 
     @Test
@@ -56,6 +56,6 @@ public void testSetSessionValidJavaWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={distinct_aggregation_spill_enabled=false}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -119,7 +119,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)`
`119`	`119`	`}`
`120`	`120`
`121`	`121`	`@Override`
`122`		`- public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties)`
	`122`	`+ public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties, String startTransactionId, boolean clearTransactionId)`
`123`	`123`	`{`
`124`	`124`	`return null;`
`125`	`125`	`}`
Original file line number	Diff line number	Diff line change
`@@ -120,7 +120,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)`
`120`	`120`	`}`
`121`	`121`
`122`	`122`	`@Override`
`123`		`- public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties)`
	`123`	`+ public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties, String startTransactionId, boolean clearTransactionId)`
`124`	`124`	`{`
`125`	`125`	`return null;`
`126`	`126`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1211,7 +1211,7 @@ public void testSetSessionJavaWorkerSessionProperty()`
`1211`	`1211`	`"MaterializedResult{rows=[[true]], " +`
`1212`	`1212`	`"types=[boolean], " +`
`1213`	`1213`	`"setSessionProperties={distinct_aggregation_spill_enabled=false}, " +`
`1214`		`- "resetSessionProperties=[], updateType=SET SESSION}");`
	`1214`	`+ "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");`
`1215`	`1215`	`}`
`1216`	`1216`
`1217`	`1217`	`@Test`
Original file line number	Diff line number	Diff line change
`@@ -155,7 +155,7 @@ public void testSetNativeWorkerSessionProperty()`
`155`	`155`	`"MaterializedResult{rows=[[true]], " +`
`156`	`156`	`"types=[boolean], " +`
`157`	`157`	`"setSessionProperties={driver_cpu_time_slice_limit_ms=500}, " +`
`158`		`- "resetSessionProperties=[], updateType=SET SESSION}");`
	`158`	`+ "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");`
`159`	`159`	`}`
`160`	`160`
`161`	`161`	`@Test`
Original file line number	Diff line number	Diff line change
`@@ -143,7 +143,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)`
`143`	`143`	`}`
`144`	`144`
`145`	`145`	`@Override`
`146`		`- public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties)`
	`146`	`+ public Void build(Map<String, String> setSessionProperties, Set<String> resetSessionProperties, String startTransactionId, boolean clearTransactionId)`
`147`	`147`	`{`
`148`	`148`	`return null;`
`149`	`149`	`}`
Original file line number	Diff line number	Diff line change
`@@ -183,6 +183,12 @@ public void testDescribeOutputNamedAndUnnamed()`
`183`	`183`	`// this connector uses a non-canonical type for varchar columns in tpch`
`184`	`184`	`}`
`185`	`185`
	`186`	`+ @Override`
	`187`	`+ public void testClearTransactionId()`
	`188`	`+ {`
	`189`	`+ // Catalog singlestore only supports writes using autocommit`
	`190`	`+ }`
	`191`	`+`
`186`	`192`	`@Override`
`187`	`193`	`public void testInsert()`
`188`	`194`	`{`
Original file line number	Diff line number	Diff line change
`@@ -594,6 +594,8 @@ private MaterializedResult executeWithStrategies(`
`594`	`594`	`ImmutableSet.of(),`
`595`	`595`	`p.getUpdateType(),`
`596`	`596`	`getOnlyElement(getOnlyElement(rows).getFields()) == null ? OptionalLong.empty() : OptionalLong.of((Long) getOnlyElement(getOnlyElement(rows).getFields())),`
	`597`	`+ Optional.empty(),`
	`598`	`+ false,`
`597`	`599`	`ImmutableList.of());`
`598`	`600`	`}`
`599`	`601`	`}`
`@@ -606,6 +608,8 @@ else if (execution instanceof PrestoSparkAccessControlCheckerExecution) {`
`606`	`608`	`ImmutableSet.of(),`
`607`	`609`	`Optional.empty(),`
`608`	`610`	`OptionalLong.empty(),`
	`611`	`+ Optional.empty(),`
	`612`	`+ false,`
`609`	`613`	`ImmutableList.of());`
`610`	`614`	`}`
`611`	`615`	`else {`
`@@ -616,6 +620,8 @@ else if (execution instanceof PrestoSparkAccessControlCheckerExecution) {`
`616`	`620`	`ImmutableSet.of(),`
`617`	`621`	`Optional.empty(),`
`618`	`622`	`OptionalLong.empty(),`
	`623`	`+ Optional.empty(),`
	`624`	`+ false,`
`619`	`625`	`ImmutableList.of());`
`620`	`626`	`}`
`621`	`627`	`}`