Tuesday, November 27, 2012

Execute scripts in multiple databases

Have you ever needed to execute the same script on several databases? I have met such need and developed a simple algorithm for it. This algorithm I'd like to share.

At first I tried to develop it using only JDBC but unfortunately it has the really big limitation: it cannot execute at once several calls, which is pretty common for SQL scripts. So I decided to run scripts merely from a command line.

Here is the code:
package uay.loader;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileFilter;
import java.io.FileReader;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.apache.log4j.Logger;

public class BulkScriptLoader {

    private static final String   DB_PREFIX     = "ua";
    private static final String[] DB_SKIP = { "ua_1", "ua_2" };
    private static final String   FOLDER_NAME   = "scripts";
    private static final String   FILE_EXT      = ".sql";
    private static final String   URL           = HOST + ":" + PORT;
    private static final String   USER_NAME     = "xxx";
    private static final String   PASSWORD      = "xxx";

    private final static Logger   LOGGER        = Logger.getLogger(BulkScriptLoader.class);

    public static void main(String[] args) {
        Connection connection;
        try {
            connection = DriverManager.getConnection(URL, USER_NAME, PASSWORD);
        } catch (SQLException e1) {
            LOGGER.error(e1);
            throw new RuntimeException(e1);
        }

        List<String> dbNames = getDbNames(connection);
        File[] scriptFiles = getScriptFiles();
        for (File file : scriptFiles) {
            LOGGER.info("Work with " + file.getName());
            BufferedReader bf;
            StringBuilder statement = new StringBuilder();
            try {
                bf = new BufferedReader(new FileReader(file));
                String line = "";
                while ((line = bf.readLine()) != null) {
                    statement.append(line).append("\n");
                }
            } catch (IOException e) {
                LOGGER.error(e, e);
            }
            if (statement.toString().length() > 0) {
                execute(connection, dbNames, statement.toString(), URL, USER_NAME, PASSWORD);
            }
        }
    }

    public static File[] getScriptFiles() {
        File folder = new File(FOLDER_NAME);
        if (folder.exists() && folder.isDirectory()) {
            return folder.listFiles(new FileFilter() {
                @Override
                public boolean accept(File pathname) {
                    return pathname.getName().toLowerCase().endsWith(FILE_EXT.toLowerCase());
                }
            });
        }
        return null;
    }

    public static void execute(Connection connection, List<String> dbNames,
            String generalStatement, String url, String userName, String password) {
        for (String dbName : dbNames) {
            String cmdCommand = String.format("mysql -u %s -p%s -h %s -P %s -D %s",userName, password, 
              parseHost(url), parsePort(url), dbName);
            LOGGER.info(cmdCommand);
            try {
                Process process = Runtime.getRuntime().exec(cmdCommand);
                Writer w = new OutputStreamWriter(process.getOutputStream());
                // System.out.println(statement);
                w.write(generalStatement);
                w.flush();
                w.close();
            } catch (IOException e) {
                LOGGER.error(e, e);
            }
        }
    }

    public static String parseHost(String url) {
        return url.substring(0, (url.indexOf(":") > 0 ? url.indexOf(":"): url.length()))

    }

    public static String parsePort(String url) {
        Pattern p = Pattern.compile("(\\d\\d\\d\\d)");
        Matcher m = p.matcher(url);
        if (m.find()) {
            return m.group(0);
        }
        return "3306";
    }

    public static List<String> getDbNames(Connection connection) {
        List<String> dbNames = new ArrayList<String>();
        try {
            PreparedStatement preparedStatement = connection.
               prepareStatement("select schema_name from information_schema.schemata");
            if (preparedStatement.execute()) {
                ResultSet rs = preparedStatement.getResultSet();
                while (rs.next()) {
                    String dbName = rs.getString(1);
                    if (dbName.startsWith(DB_PREFIX) && !Arrays.asList(DB_SKIP).contains(dbName)) {
                        dbNames.add(dbName);
                    }
                }
            }
        } catch (SQLException e) {
            LOGGER.error(e, e);
        }
        return dbNames;
    }
}
As you can see it's really simple. All scripts that I'd like to execute I place in 'scripts' folder with '.sql' extension. I configure database connection, prefix of databases that I want to update and specify the list of databases that I want to skip with this update.