Skip to content

Commit 602b59d

Browse files
hantangwangdhantangwangd
committed
Enhance test framework to handle start/rollback/commit in a nature way
1 parent 9ac8e50 commit 602b59d

File tree

7 files changed

+167
-63
lines changed

7 files changed

+167
-63
lines changed

‎presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,12 @@ public void testUpdate()
121121
// Updates are not supported by the connector
122122
}
123123

124+
@Override
125+
public void testNonAutoCommitTransactionWithRollback()
126+
{
127+
// This connector do not support rollback for insert actions
128+
}
129+
124130
@Override
125131
public void testInsert()
126132
{

‎presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,13 @@ public void testDescribeOutputNamedAndUnnamed()
116116
}
117117

118118
@Override
119-
public void testClearTransactionId()
119+
public void testNonAutoCommitTransactionWithRollback()
120+
{
121+
// Catalog iceberg only supports writes using autocommit
122+
}
123+
124+
@Override
125+
public void testNonAutoCommitTransactionWithCommit()
120126
{
121127
// Catalog iceberg only supports writes using autocommit
122128
}

‎presto-main-base/src/main/java/com/facebook/presto/Session.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
import com.facebook.presto.sql.planner.optimizations.OptimizerInformationCollector;
4343
import com.facebook.presto.sql.planner.optimizations.OptimizerResultCollector;
4444
import com.facebook.presto.transaction.TransactionManager;
45+
import com.google.common.annotations.VisibleForTesting;
4546
import com.google.common.collect.ImmutableMap;
4647
import com.google.common.collect.ImmutableSet;
4748
import com.google.common.collect.Maps;
@@ -470,6 +471,50 @@ public Session beginTransactionId(TransactionId transactionId, boolean enableRol
470471
queryType);
471472
}
472473

474+
@VisibleForTesting
475+
public Session clearTransaction(TransactionManager transactionManager, AccessControl accessControl)
476+
{
477+
checkArgument(this.transactionId.isPresent(), "Session does not have an active transaction");
478+
requireNonNull(transactionManager, "transactionManager is null");
479+
requireNonNull(accessControl, "accessControl is null");
480+
481+
for (Entry<String, String> property : systemProperties.entrySet()) {
482+
// verify permissions
483+
accessControl.checkCanSetSystemSessionProperty(identity, context, property.getKey());
484+
485+
// validate session property value
486+
sessionPropertyManager.validateSystemSessionProperty(property.getKey(), property.getValue());
487+
}
488+
489+
return new Session(
490+
queryId,
491+
Optional.empty(),
492+
clientTransactionSupport,
493+
identity,
494+
source,
495+
catalog,
496+
schema,
497+
traceToken,
498+
timeZoneKey,
499+
locale,
500+
remoteUserAddress,
501+
userAgent,
502+
clientInfo,
503+
clientTags,
504+
resourceEstimates,
505+
startTime,
506+
systemProperties,
507+
connectorProperties,
508+
unprocessedCatalogProperties,
509+
sessionPropertyManager,
510+
preparedStatements,
511+
sessionFunctions,
512+
tracer,
513+
warningCollector,
514+
runtimeStats,
515+
queryType);
516+
}
517+
473518
public ConnectorSession toConnectorSession()
474519
{
475520
return new FullConnectorSession(this, identity.toConnectorIdentity());

‎presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,13 @@ public void testDescribeOutputNamedAndUnnamed()
184184
}
185185

186186
@Override
187-
public void testClearTransactionId()
187+
public void testNonAutoCommitTransactionWithRollback()
188+
{
189+
// Catalog singlestore only supports writes using autocommit
190+
}
191+
192+
@Override
193+
public void testNonAutoCommitTransactionWithCommit()
188194
{
189195
// Catalog singlestore only supports writes using autocommit
190196
}

‎presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java‎

Lines changed: 58 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,6 @@
7777
import static com.facebook.presto.testing.TestingSession.TESTING_CATALOG;
7878
import static com.facebook.presto.testing.assertions.Assert.assertEquals;
7979
import static com.facebook.presto.tests.QueryAssertions.assertContains;
80-
import static com.facebook.presto.transaction.TransactionBuilder.transaction;
8180
import static com.google.common.collect.ImmutableList.toImmutableList;
8281
import static com.google.common.collect.Iterables.getOnlyElement;
8382
import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly;
@@ -250,73 +249,71 @@ public void testCreateTable()
250249
}
251250

252251
@Test
253-
public void testClearTransactionId()
252+
public void testNonAutoCommitTransactionWithRollback()
254253
{
255-
assertUpdate("create table test_clear_transaction_id_table(a int, b varchar)");
256-
assertUpdate("insert into test_clear_transaction_id_table values(1, '1001'), (2, '1002')", 2);
257-
Session session = getQueryRunner().getDefaultSession();
258-
String defaultCatalog = session.getCatalog().get();
259-
transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl())
260-
.execute(Session.builder(session)
261-
.setIdentity(new Identity("admin",
262-
Optional.empty(),
263-
ImmutableMap.of(defaultCatalog, new SelectedRole(ROLE, Optional.of("admin"))),
264-
ImmutableMap.of(),
265-
ImmutableMap.of(),
266-
Optional.empty(),
267-
Optional.empty()))
268-
.build(),
269-
txnSession -> {
270-
MaterializedResult result = computeActual(txnSession, "select * from test_clear_transaction_id_table");
271-
assertEquals(result.getRowCount(), 2);
272-
assertFalse(result.isClearTransactionId());
273-
274-
result = computeActual(txnSession, "insert into test_clear_transaction_id_table values(1, '1001'), (2, '1002')");
275-
assertEquals(result.getOnlyValue(), 2L);
276-
assertFalse(result.isClearTransactionId());
277-
278-
// `Rollback` executes successfully, and the client gets a flag `clearTransactionId = true`
279-
result = computeActual(txnSession, "rollback");
280-
assertTrue(result.isClearTransactionId());
281-
});
282-
283-
assertQuery("select * from test_clear_transaction_id_table", "values(1, '1001'), (2, '1002')");
284-
assertUpdate("drop table if exists test_clear_transaction_id_table");
254+
assertUpdate("create table multi_statements_transaction_rollback(a int, b varchar)");
255+
assertUpdate("insert into multi_statements_transaction_rollback values(1, '1001'), (2, '1002')", 2);
256+
257+
Session session = assertStartTransaction(getSession(), "start transaction");
258+
259+
assertQuery(session, "select * from multi_statements_transaction_rollback", "values(1, '1001'), (2, '1002')");
260+
assertUpdate(session, "insert into multi_statements_transaction_rollback values(3, '1003'), (4, '1004')", 2);
261+
262+
// `Rollback` executes successfully, and the client gets a flag `clearTransactionId = true`
263+
session = assertEndTransaction(session, "rollback");
264+
265+
assertQuery(session, "select * from multi_statements_transaction_rollback", "values(1, '1001'), (2, '1002')");
266+
assertUpdate(session, "drop table if exists multi_statements_transaction_rollback");
267+
}
268+
269+
@Test
270+
public void testNonAutoCommitTransactionWithCommit()
271+
{
272+
assertUpdate("create table multi_statements_transaction_commit(a int, b varchar)");
273+
assertUpdate("insert into multi_statements_transaction_commit values(1, '1001'), (2, '1002')", 2);
274+
Session session = assertStartTransaction(getSession(), "start transaction");
275+
276+
assertQuery(session, "select * from multi_statements_transaction_commit", "values(1, '1001'), (2, '1002')");
277+
assertUpdate(session, "insert into multi_statements_transaction_commit values(3, '1003'), (4, '1004')", 2);
278+
279+
// `Commit` executes successfully, and the client gets a flag `clearTransactionId = true`
280+
session = assertEndTransaction(session, "commit");
281+
282+
assertQuery(session, "select * from multi_statements_transaction_commit",
283+
"values(1, '1001'), (2, '1002'), (3, '1003'), (4, '1004')");
284+
assertUpdate(session, "drop table if exists multi_statements_transaction_commit");
285285
}
286286

287287
@Test
288288
public void testNonAutoCommitTransactionWithFailAndRollback()
289289
{
290290
assertUpdate("create table test_non_autocommit_table(a int, b varchar)");
291-
Session session = getQueryRunner().getDefaultSession();
292-
String defaultCatalog = session.getCatalog().get();
293-
transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl())
294-
.execute(Session.builder(session)
295-
.setIdentity(new Identity("admin",
296-
Optional.empty(),
297-
ImmutableMap.of(defaultCatalog, new SelectedRole(ROLE, Optional.of("admin"))),
298-
ImmutableMap.of(),
299-
ImmutableMap.of(),
300-
Optional.empty(),
301-
Optional.empty()))
302-
.build(),
303-
txnSession -> {
304-
// simulate failure of SQL statement execution
305-
assertQueryFails(txnSession, "SELECT fail('forced failure')", "forced failure");
306-
307-
// cannot execute any SQLs except `rollback` in current session
308-
assertQueryFails(txnSession, "select count(*) from test_non_autocommit_table", "Current transaction is aborted, commands ignored until end of transaction block");
309-
assertQueryFails(txnSession, "show tables", "Current transaction is aborted, commands ignored until end of transaction block");
310-
assertQueryFails(txnSession, "insert into test_non_autocommit_table values(1, '1001')", "Current transaction is aborted, commands ignored until end of transaction block");
311-
assertQueryFails(txnSession, "create table test_table(a int, b varchar)", "Current transaction is aborted, commands ignored until end of transaction block");
312-
313-
// execute `rollback` successfully
314-
MaterializedResult result = computeActual(txnSession, "rollback");
315-
assertTrue(result.isClearTransactionId());
316-
});
317-
318-
assertQuery("select count(*) from test_non_autocommit_table", "values(0)");
319-
assertUpdate("drop table if exists test_non_autocommit_table");
291+
Session session = Session.builder(getSession())
292+
.setIdentity(new Identity("admin",
293+
Optional.empty(),
294+
ImmutableMap.of(getSession().getCatalog().get(), new SelectedRole(ROLE, Optional.of("admin"))),
295+
ImmutableMap.of(),
296+
ImmutableMap.of(),
297+
Optional.empty(),
298+
Optional.empty()))
299+
.build();
300+
301+
session = assertStartTransaction(session, "start transaction");
302+
303+
// simulate failure of SQL statement execution
304+
assertQueryFails(session, "SELECT fail('forced failure')", "forced failure");
305+
306+
// cannot execute any SQLs except `rollback` in current session
307+
assertQueryFails(session, "select count(*) from test_non_autocommit_table", "Current transaction is aborted, commands ignored until end of transaction block");
308+
assertQueryFails(session, "show tables", "Current transaction is aborted, commands ignored until end of transaction block");
309+
assertQueryFails(session, "insert into test_non_autocommit_table values(1, '1001')", "Current transaction is aborted, commands ignored until end of transaction block");
310+
assertQueryFails(session, "create table test_table(a int, b varchar)", "Current transaction is aborted, commands ignored until end of transaction block");
311+
312+
// execute `rollback` successfully
313+
session = assertEndTransaction(session, "rollback");
314+
315+
assertQuery(session, "select count(*) from test_non_autocommit_table", "values(0)");
316+
assertUpdate(session, "drop table if exists test_non_autocommit_table");
320317
}
321318

322319
@Test

‎presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,16 @@ protected void assertUpdate(@Language("SQL") String sql)
268268
assertUpdate(getSession(), sql);
269269
}
270270

271+
protected Session assertStartTransaction(Session session, @Language("SQL") String sql)
272+
{
273+
return QueryAssertions.assertStartTransaction(queryRunner, session, sql);
274+
}
275+
276+
protected Session assertEndTransaction(Session session, @Language("SQL") String sql)
277+
{
278+
return QueryAssertions.assertEndTransaction(queryRunner, session, sql);
279+
}
280+
271281
protected void assertUpdate(Session session, @Language("SQL") String sql)
272282
{
273283
QueryAssertions.assertUpdate(queryRunner, session, sql, OptionalLong.empty(), Optional.empty());

‎presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import com.facebook.airlift.units.Duration;
1818
import com.facebook.presto.Session;
1919
import com.facebook.presto.common.QualifiedObjectName;
20+
import com.facebook.presto.common.transaction.TransactionId;
2021
import com.facebook.presto.spi.WarningCollector;
2122
import com.facebook.presto.sql.planner.Plan;
2223
import com.facebook.presto.testing.ExpectedQueryRunner;
@@ -46,7 +47,9 @@
4647
import static java.util.concurrent.TimeUnit.MILLISECONDS;
4748
import static java.util.concurrent.TimeUnit.SECONDS;
4849
import static org.testng.Assert.assertEquals;
50+
import static org.testng.Assert.assertFalse;
4951
import static org.testng.Assert.assertNotNull;
52+
import static org.testng.Assert.assertTrue;
5053
import static org.testng.Assert.fail;
5154

5255
public final class QueryAssertions
@@ -57,6 +60,37 @@ private QueryAssertions()
5760
{
5861
}
5962

63+
public static Session assertStartTransaction(QueryRunner queryRunner, Session session, @Language("SQL") String sql)
64+
{
65+
MaterializedResult results = queryRunner.execute(session, sql);
66+
if (!results.getUpdateType().isPresent()) {
67+
fail("update type is not set");
68+
}
69+
if (!results.getUpdateType().get().equals("START TRANSACTION")) {
70+
fail("not a start transaction statement");
71+
}
72+
assertTrue(results.getStartedTransactionId().isPresent());
73+
assertFalse(results.isClearTransactionId());
74+
75+
TransactionId transactionId = results.getStartedTransactionId().get();
76+
return session.beginTransactionId(transactionId, queryRunner.getTransactionManager(), queryRunner.getAccessControl());
77+
}
78+
79+
public static Session assertEndTransaction(QueryRunner queryRunner, Session session, @Language("SQL") String sql)
80+
{
81+
MaterializedResult results = queryRunner.execute(session, sql);
82+
if (!results.getUpdateType().isPresent()) {
83+
fail("update type is not set");
84+
}
85+
if (!results.getUpdateType().get().equals("ROLLBACK") && !results.getUpdateType().get().equals("COMMIT")) {
86+
fail("not a end transaction statement");
87+
}
88+
assertTrue(results.isClearTransactionId());
89+
assertFalse(results.getStartedTransactionId().isPresent());
90+
91+
return session.clearTransaction(queryRunner.getTransactionManager(), queryRunner.getAccessControl());
92+
}
93+
6094
public static void assertUpdate(QueryRunner queryRunner, Session session, @Language("SQL") String sql, OptionalLong count, Optional<Consumer<Plan>> planAssertion)
6195
{
6296
long start = System.nanoTime();

0 commit comments

Comments
 (0)