package com.coreoz.plume.db.transaction;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.function.Consumer;
import java.util.function.Function;

import javax.inject.Inject;
import javax.inject.Singleton;
import javax.sql.DataSource;

import com.google.common.base.Throwables;
import com.typesafe.config.Config;

/**
 * Handle transactions over a classic JDBC {@link Connection}.
 */
@Singleton
public class TransactionManager {

	private final DataSource dataSource;

	@Inject
	public TransactionManager(Config config) {
		this(config, "db");
	}

	public TransactionManager(Config config, String prefix) {
		this(HikariDataSources.fromConfig(config, prefix + ".hikari"));
	}

	public TransactionManager(DataSource dataSource) {
		this.dataSource = dataSource;
	}

	// API

	public DataSource dataSource() {
		return dataSource;
	}

	public <T> T executeAndReturn(Function<Connection, T> toExecuteOnDb) {
		Connection connection = null;
		Boolean initialAutoCommit = null;
		try {
			connection = dataSource.getConnection();
			initialAutoCommit = connection.getAutoCommit();
			connection.setAutoCommit(false);
			T result = toExecuteOnDb.apply(connection);
			connection.commit();
			return result;
		} catch(Throwable e) {
			try {
				if(connection != null) {
					connection.rollback();
				}
			} catch (Throwable e2) {
				// never mind if the connection cannot be rolled back
			}
			Throwables.throwIfUnchecked(e);
			throw new RuntimeException(e);
		} finally {
			if(connection != null) {
				try {
					if(initialAutoCommit != null) {
						connection.setAutoCommit(initialAutoCommit);
					}
					connection.close();
				} catch (SQLException e) {
					// never mind if the connection cannot be closed
				}
			}
		}
	}

	public void execute(Consumer<Connection> toExecuteOnDb) {
		executeAndReturn(connection -> {
			toExecuteOnDb.accept(connection);
			return null;
		});
	}

}

@@ -27,32 +27,60 @@
 b.transaction;%0A%0A
+import java.io.IOException;%0A
 import java.sql.
@@ -1055,16 +1055,18 @@
 taSource
+()
 .getConn
@@ -1387,58 +1387,260 @@
 %09// 
-never mind if the connection cannot be rolled back
+if the rollback failed, raise an exception about the rollback failure%0A%09%09%09%09// and the original error%0A%09%09%09%09RuntimeException combinedException = new RuntimeException(e2);%0A%09%09%09%09combinedException.addSuppressed(new IOException(e));%0A%09%09%09%09throw combinedException;
 %0A%09%09%09

Handle correctly rollback exception + simplify executeAndReturn override