Improve db migration so that multiple SELECT can be concurrently executed

2 files changed, 34 insertions, 22 deletions

diff --git a/server/sonar-server/src/main/java/org/sonar/server/db/migrations/BaseDataChange.java b/server/sonar-server/src/main/java/org/sonar/server/db/migrations/BaseDataChange.java
index 32ebbaf443e..58304aa3d19 100644
--- a/server/sonar-server/src/main/java/org/sonar/server/db/migrations/BaseDataChange.java
+++ b/server/sonar-server/src/main/java/org/sonar/server/db/migrations/BaseDataChange.java
@@ -19,10 +19,8 @@
  */
 package org.sonar.server.db.migrations;
 
-import org.apache.commons.dbutils.DbUtils;
 import org.sonar.core.persistence.Database;
 
-import java.sql.Connection;
 import java.sql.SQLException;
 
 public abstract class BaseDataChange implements DataChange, DatabaseMigration {
@@ -35,22 +33,11 @@ public abstract class BaseDataChange implements DataChange, DatabaseMigration {
 
   @Override
   public final void execute() throws SQLException {
-    Connection readConnection = null, writeConnection = null;
+    Context context = new Context(db);
     try {
-      readConnection = db.getDataSource().getConnection();
-      readConnection.setAutoCommit(false);
-      if (readConnection.getMetaData().supportsTransactionIsolationLevel(Connection.TRANSACTION_READ_UNCOMMITTED)) {
-        readConnection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED);
-      }
-
-      writeConnection = db.getDataSource().getConnection();
-      writeConnection.setAutoCommit(false);
-      Context context = new Context(db, readConnection, writeConnection);
       execute(context);
-
     } finally {
-      DbUtils.closeQuietly(readConnection);
-      DbUtils.closeQuietly(writeConnection);
+      context.close();
     }
   }
 }
diff --git a/server/sonar-server/src/main/java/org/sonar/server/db/migrations/DataChange.java b/server/sonar-server/src/main/java/org/sonar/server/db/migrations/DataChange.java
index c901a782594..dff05618251 100644
--- a/server/sonar-server/src/main/java/org/sonar/server/db/migrations/DataChange.java
+++ b/server/sonar-server/src/main/java/org/sonar/server/db/migrations/DataChange.java
@@ -19,33 +19,58 @@
  */
 package org.sonar.server.db.migrations;
 
+import com.google.common.collect.Lists;
+import org.apache.commons.dbutils.DbUtils;
 import org.sonar.core.persistence.Database;
 
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.util.List;
 
 public interface DataChange {
 
   class Context {
     private final Database db;
-    private final Connection readConnection, writeConnection;
+    private final List<Connection> connections = Lists.newArrayList();
 
-    public Context(Database db, Connection readConnection, Connection writeConnection) {
+    Context(Database db) {
       this.db = db;
-      this.readConnection = readConnection;
-      this.writeConnection = writeConnection;
     }
 
     public Select prepareSelect(String sql) throws SQLException {
-      return SelectImpl.create(db, readConnection, sql);
+      return SelectImpl.create(db, openReadConnection(), sql);
     }
 
     public Upsert prepareUpsert(String sql) throws SQLException {
-      return UpsertImpl.create(writeConnection, sql);
+      return UpsertImpl.create(openWriteConnection(), sql);
     }
 
     public MassUpdate prepareMassUpdate() throws SQLException {
-      return new MassUpdate(db, readConnection, writeConnection);
+      return new MassUpdate(db, openReadConnection(), openWriteConnection());
+    }
+
+    private Connection openReadConnection() throws SQLException {
+      Connection readConnection = db.getDataSource().getConnection();
+      readConnection.setAutoCommit(false);
+      if (readConnection.getMetaData().supportsTransactionIsolationLevel(Connection.TRANSACTION_READ_UNCOMMITTED)) {
+        readConnection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED);
+      }
+      connections.add(readConnection);
+      return readConnection;
+    }
+
+    private Connection openWriteConnection() throws SQLException {
+      Connection writeConnection = db.getDataSource().getConnection();
+      writeConnection.setAutoCommit(false);
+      connections.add(writeConnection);
+      return writeConnection;
+    }
+
+    void close() {
+      for (Connection connection : connections) {
+        DbUtils.closeQuietly(connection);
+      }
+      connections.clear();
     }
   }