org.apache.shardingsphere.transaction.core.TransactionType Java Examples

@Test
public void assertShardingTransactionAutoCommit() throws SQLException {
    TransactionTypeHolder.set(TransactionType.XA);
    try (ShardingSphereConnection actual = getShardingSphereDataSource().getConnection()) {
        actual.createStatement().executeQuery(sql);
        actual.setAutoCommit(false);
        actual.createStatement().executeQuery(sql);
        assertTrue(actual.getShardingTransactionManager().isInTransaction());
        Multimap<String, Connection> cachedConnections = getCachedConnections(actual);
        assertThat(cachedConnections.size(), is(1));
        for (Connection each : cachedConnections.values()) {
            assertTrue(each.getAutoCommit());
        }
    } finally {
        TransactionTypeHolder.clear();
    }
}
 

@Override
public BackendResponse execute() {
    Optional<ShardingCTLSetStatement> shardingTCLStatement = new ShardingCTLSetParser(sql).doParse();
    if (!shardingTCLStatement.isPresent()) {
        return new ErrorResponse(new InvalidShardingCTLFormatException(sql));
    }
    switch (shardingTCLStatement.get().getKey()) {
        case "TRANSACTION_TYPE":
            try {
                backendConnection.setTransactionType(TransactionType.valueOf(shardingTCLStatement.get().getValue()));
            } catch (final IllegalArgumentException ex) {
                return new ErrorResponse(new UnsupportedShardingCTLTypeException(sql));
            }
            break;
        default:
            return new ErrorResponse(new UnsupportedShardingCTLTypeException(sql));
    }
    return new UpdateResponse();
}
 

@Override
public void processFailure() throws SQLException {
    System.out.println("-------------------- Process Start ---------------------");
    TransactionTypeHolder.set(TransactionType.XA);
    try (Connection connection = dataSource.getConnection()) {
        connection.setAutoCommit(false);
        PreparedStatement preparedStatement = connection.prepareStatement("INSERT INTO t_order (user_id, status) VALUES (?, ?)");
        doInsert(preparedStatement);
        connection.rollback();
        System.out.println("INSERT 10 orders failed");
    } finally {
        TransactionTypeHolder.clear();
    }
    int quantity = selectAll();
    System.out.printf("Rollback, expect:0, actual:%d \n", quantity);
    printData();
    System.out.println("-------------------- Process End -----------------------");
}
 

/**
 * Close cached connection.
 *
 * @param forceClose force close flag
 * @throws SQLException SQL exception
 */
public synchronized void close(final boolean forceClose) throws SQLException {
    Collection<SQLException> exceptions = new LinkedList<>();
    MasterVisitedManager.clear();
    exceptions.addAll(closeResultSets());
    exceptions.addAll(closeStatements());
    if (!stateHandler.isInTransaction() || forceClose || TransactionType.BASE == transactionType) {
        exceptions.addAll(releaseConnections(forceClose));
    }
    stateHandler.doNotifyIfNecessary();
    throwSQLExceptionIfNecessary(exceptions);
}
 

@Before
public void setUp() {
    when(backendConnection.getTransactionType()).thenReturn(TransactionType.LOCAL);
    RuntimeContext runtimeContext = mock(RuntimeContext.class);
    SchemaContext schema = mock(SchemaContext.class);
    when(runtimeContext.getTransactionManagerEngine()).thenReturn(mock(ShardingTransactionManagerEngine.class));
    when(schema.getRuntimeContext()).thenReturn(runtimeContext);
    when(backendConnection.getSchema()).thenReturn(schema);
}
 

@Override
public void commit() throws SQLException {
    if (TransactionType.LOCAL == transactionType) {
        super.commit();
    } else {
        shardingTransactionManager.commit();
    }
}
 

@Test
public void assertShardingTransactionForceCommit() throws SQLException {
    TransactionTypeHolder.set(TransactionType.XA);
    try (ShardingSphereConnection actual = getShardingSphereDataSource().getConnection()) {
        actual.setAutoCommit(false);
        actual.setAutoCommit(true);
        assertTrue(XAShardingTransactionManagerFixture.getInvocations().contains(TransactionOperationType.COMMIT));
    } finally {
        TransactionTypeHolder.clear();
    }
}
 

@Test
public void assertSwitchTransactionTypeXA() {
    backendConnection.setCurrentSchema("schema_0");
    ShardingCTLSetBackendHandler shardingCTLBackendHandler = new ShardingCTLSetBackendHandler("sctl:set transaction_type=XA", backendConnection);
    BackendResponse actual = shardingCTLBackendHandler.execute();
    assertThat(actual, instanceOf(UpdateResponse.class));
    assertThat(backendConnection.getTransactionType(), is(TransactionType.XA));
}
 

@Override
public Object invoke(final MethodInvocation methodInvocation) throws Throwable {
    ShardingTransactionType shardingTransactionType = getAnnotation(methodInvocation);
    Objects.requireNonNull(shardingTransactionType, "could not found sharding transaction type annotation");
    TransactionType preTransactionType = TransactionTypeHolder.get();
    TransactionTypeHolder.set(shardingTransactionType.value());
    try {
        return methodInvocation.proceed();
    } finally {
        TransactionTypeHolder.clear();
        if (null != preTransactionType) {
            TransactionTypeHolder.set(preTransactionType);
        }
    }
}
 

/**
 * Execute XA with exception.
 *
 * @throws SQLException SQL exception
 */
void insertFailed() throws SQLException {
    TransactionTypeHolder.set(TransactionType.BASE);
    try (Connection connection = dataSource.getConnection()) {
        connection.setAutoCommit(false);
        PreparedStatement preparedStatement = connection.prepareStatement("INSERT INTO t_order (user_id, status) VALUES (?, ?)");
        doInsert(preparedStatement);
        connection.rollback();
    } finally {
        TransactionTypeHolder.clear();
    }
}
 

/**
 * Execute XA with exception.
 *
 * @param count insert record count
 */
@Transactional
@ShardingTransactionType(TransactionType.BASE)
public void insertFailed(final int count) {
    jdbcTemplate.execute("INSERT INTO t_order (user_id, status) VALUES (?, ?)", (PreparedStatementCallback<TransactionType>) preparedStatement -> {
        doInsert(count, preparedStatement);
        throw new SQLException("mock transaction failed");
    });
}
 

/**
 * Execute XA.
 *
 * @param count insert record count
 * @return transaction type
 */
@Transactional
@ShardingTransactionType(TransactionType.XA)
public TransactionType insert(final int count) {
    return jdbcTemplate.execute("INSERT INTO t_order (user_id, status) VALUES (?, ?)", (PreparedStatementCallback<TransactionType>) preparedStatement -> {
        doInsert(count, preparedStatement);
        return TransactionTypeHolder.get();
    });
}
 

/**
 * Execute XA with exception.
 *
 * @param count insert record count
 */
@Transactional
@ShardingTransactionType(TransactionType.XA)
public void insertFailed(final int count) {
    jdbcTemplate.execute("INSERT INTO t_order (user_id, status) VALUES (?, ?)", (PreparedStatementCallback<TransactionType>) preparedStatement -> {
        doInsert(count, preparedStatement);
        throw new SQLException("mock transaction failed");
    });
}
 

/**
 * Execute XA.
 *
 * @param count insert record count
 * @return transaction type
 */
@Transactional
@ShardingTransactionType(TransactionType.BASE)
public TransactionType insert(final int count) {
    return jdbcTemplate.execute("INSERT INTO t_order (user_id, status) VALUES (?, ?)", (PreparedStatementCallback<TransactionType>) preparedStatement -> {
        doInsert(count, preparedStatement);
        return TransactionTypeHolder.get();
    });
}
 

@Test
public void assertBeginForLocalTransaction() {
    newBackendTransactionManager(TransactionType.LOCAL, false);
    backendTransactionManager.begin();
    verify(stateHandler).setStatus(ConnectionStatus.TRANSACTION);
    verify(backendConnection).releaseConnections(false);
    verify(localTransactionManager).begin();
}
 

@Test
public void assertRollbackWithoutTransaction() throws SQLException {
    newBackendTransactionManager(TransactionType.LOCAL, false);
    backendTransactionManager.rollback();
    verify(stateHandler, times(0)).setStatus(ConnectionStatus.TERMINATED);
    verify(localTransactionManager, times(0)).rollback();
    verify(shardingTransactionManager, times(0)).rollback();
}
 

/**
 * Get connections.
 *
 * @param schemaName scheme name
 * @param dataSourceName data source name
 * @param connectionSize size of connections to be get
 * @param connectionMode connection mode
 * @param transactionType transaction type
 * @return connections
 * @throws SQLException SQL exception
 */
@SuppressWarnings("SynchronizationOnLocalVariableOrMethodParameter")
public List<Connection> getConnections(final String schemaName, final String dataSourceName, 
                                       final int connectionSize, final ConnectionMode connectionMode, final TransactionType transactionType) throws SQLException {
    DataSource dataSource = schemaContexts.getSchemaContexts().get(schemaName).getSchema().getDataSources().get(dataSourceName);
    if (1 == connectionSize) {
        return Collections.singletonList(createConnection(schemaName, dataSourceName, dataSource, transactionType));
    }
    if (ConnectionMode.CONNECTION_STRICTLY == connectionMode) {
        return createConnections(schemaName, dataSourceName, dataSource, connectionSize, transactionType);
    }
    synchronized (dataSource) {
        return createConnections(schemaName, dataSourceName, dataSource, connectionSize, transactionType);
    }
}
 

/**
 * Execute XA.
 *
 * @throws SQLException SQL exception
 */
void insert() throws SQLException {
    TransactionTypeHolder.set(TransactionType.BASE);
    try (Connection connection = dataSource.getConnection()) {
        connection.setAutoCommit(false);
        PreparedStatement preparedStatement = connection.prepareStatement("INSERT INTO t_order (user_id, status) VALUES (?, ?)");
        doInsert(preparedStatement);
        connection.commit();
    } finally {
        TransactionTypeHolder.clear();
    }
}
 

@Test
public void assertInit() {
    Map<String, DataSource> actual = getDataSourceMap();
    assertThat(actual.size(), is(1));
    assertThat(actual.get("demo_ds"), instanceOf(DataSourceProxy.class));
    assertThat(seataATShardingTransactionManager.getTransactionType(), is(TransactionType.BASE));
}
 

private List<Connection> createConnections(final String schemaName, final String dataSourceName, 
                                           final DataSource dataSource, final int connectionSize, final TransactionType transactionType) throws SQLException {
    List<Connection> result = new ArrayList<>(connectionSize);
    for (int i = 0; i < connectionSize; i++) {
        try {
            result.add(createConnection(schemaName, dataSourceName, dataSource, transactionType));
        } catch (final SQLException ex) {
            for (Connection each : result) {
                each.close();
            }
            throw new SQLException(String.format("Could't get %d connections one time, partition succeed connection(%d) have released!", connectionSize, result.size()), ex);
        }
    }
    return result;
}
 

private Map<String, SchemaContext> getSchemaContextMap() {
    SchemaContext schemaContext = mock(SchemaContext.class);
    ShardingSphereSchema shardingSphereSchema = mock(ShardingSphereSchema.class);
    RuntimeContext runtimeContext = mock(RuntimeContext.class);
    ShardingTransactionManagerEngine transactionManagerEngine = mock(ShardingTransactionManagerEngine.class);
    when(transactionManagerEngine.getTransactionManager(TransactionType.LOCAL)).thenReturn(null);
    when(runtimeContext.getTransactionManagerEngine()).thenReturn(transactionManagerEngine);
    when(shardingSphereSchema.getDataSources()).thenReturn(mockDataSources(2));
    when(schemaContext.getName()).thenReturn("schema");
    when(schemaContext.getSchema()).thenReturn(shardingSphereSchema);
    when(schemaContext.getRuntimeContext()).thenReturn(runtimeContext);
    return Collections.singletonMap("schema", schemaContext);
}
 

@Test
public void assertBASETransactionOperation() throws SQLException {
    connection = new ShardingSphereConnection(dataSourceMap, schemaContexts, TransactionType.BASE);
    connection.setAutoCommit(false);
    assertTrue(BASEShardingTransactionManagerFixture.getInvocations().contains(TransactionOperationType.BEGIN));
    connection.commit();
    assertTrue(BASEShardingTransactionManagerFixture.getInvocations().contains(TransactionOperationType.COMMIT));
    connection.rollback();
    assertTrue(BASEShardingTransactionManagerFixture.getInvocations().contains(TransactionOperationType.ROLLBACK));
}
 

@Test
public void assertCloseShouldNotClearTransactionType() throws SQLException {
    TransactionTypeHolder.set(TransactionType.XA);
    TransactionType currentTransactionType = TransactionTypeHolder.get();
    try (ShardingSphereConnection actual = getShardingSphereDataSource().getConnection()) {
        actual.createStatement().executeQuery(sql);
    }
    assertThat(TransactionTypeHolder.get(), is(currentTransactionType));
}
 

@Test
public void assertXATransactionOperation() throws SQLException {
    connection = new ShardingSphereConnection(dataSourceMap, schemaContexts, TransactionType.XA);
    connection.setAutoCommit(false);
    assertTrue(XAShardingTransactionManagerFixture.getInvocations().contains(TransactionOperationType.BEGIN));
    connection.commit();
    assertTrue(XAShardingTransactionManagerFixture.getInvocations().contains(TransactionOperationType.COMMIT));
    connection.rollback();
    assertTrue(XAShardingTransactionManagerFixture.getInvocations().contains(TransactionOperationType.ROLLBACK));
}
 

@Test
public void assertSwitchTransactionTypeLOCAL() {
    backendConnection.setCurrentSchema("schema_0");
    ShardingCTLSetBackendHandler shardingCTLBackendHandler = new ShardingCTLSetBackendHandler("sctl:set transaction_type=LOCAL", backendConnection);
    BackendResponse actual = shardingCTLBackendHandler.execute();
    assertThat(actual, instanceOf(UpdateResponse.class));
    assertThat(backendConnection.getTransactionType(), is(TransactionType.LOCAL));
}
 

@Test(expected = ShardingException.class)
public void assertSwitchTransactionTypeFailGetConnection() throws SQLException {
    when(dataSource.getConnection()).thenThrow(new SQLException("Mock get connection failed"));
    try {
        dataSourceTransactionManagerHandler.switchTransactionType(TransactionType.XA);
    } finally {
        TransactionSynchronizationManager.unbindResourceIfPossible(dataSource);
    }
}
 

@Test
public void assertSwitchTransactionTypeSuccess() throws SQLException {
    Connection connection = mock(Connection.class);
    Statement statement = mock(Statement.class);
    when(dataSource.getConnection()).thenReturn(connection);
    when(connection.createStatement()).thenReturn(statement);
    dataSourceTransactionManagerHandler.switchTransactionType(TransactionType.XA);
    verify(statement).execute(anyString());
    TransactionSynchronizationManager.unbindResourceIfPossible(dataSource);
}
 

@Test(expected = ShardingException.class)
public void assertSwitchTransactionTypeFailExecute() throws SQLException {
    when(statement.execute(anyString())).thenThrow(new SQLException("Mock send switch transaction type SQL failed"));
    try {
        jpaTransactionManagerHandler.switchTransactionType(TransactionType.XA);
    } finally {
        TransactionSynchronizationManager.unbindResourceIfPossible(entityManagerFactory);
    }
}
 

@Override
public void begin() {
    if (!connection.getStateHandler().isInTransaction()) {
        connection.getStateHandler().setStatus(ConnectionStatus.TRANSACTION);
        connection.releaseConnections(false);
    }
    if (TransactionType.LOCAL == transactionType || null == shardingTransactionManager) {
        localTransactionManager.begin();
    } else {
        shardingTransactionManager.begin();
    }
    MetricsUtils.buriedTransactionMetric("begin");
}
 

@Override
public void rollback() throws SQLException {
    if (connection.getStateHandler().isInTransaction()) {
        try {
            if (TransactionType.LOCAL == transactionType || null == shardingTransactionManager) {
                localTransactionManager.rollback();
            } else {
                shardingTransactionManager.rollback();
            }
            MetricsUtils.buriedTransactionMetric("rollback");
        } finally {
            connection.getStateHandler().setStatus(ConnectionStatus.TERMINATED);
        }
    }
}