rdsfix.java

import org.flywaydb.core.Flyway;
import org.flywaydb.core.api.MigrationInfo;
import org.flywaydb.core.api.callback.BaseFlywayCallback;
import org.flywaydb.core.api.callback.FlywayCallback;
import org.flywaydb.core.api.configuration.FlywayConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.flyway.FlywayMigrationInitializer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;

@Configuration
@ConditionalOnClass(BaseFlywayCallback.class)
public class RdsFixFlywayCallback extends BaseFlywayCallback {
    private static Logger LOGGER = LoggerFactory.getLogger(RdsFixFlywayCallback.class);

    private FlywayConfiguration flywayConfiguration;

    @Bean
    public FlywayMigrationInitializer flywayInitializer(Flyway flyway) {
        flyway.setCallbacks(flywayCallback());
        return new FlywayMigrationInitializer(flyway);
    }

    @Bean
    public FlywayCallback flywayCallback() {
        return this;
    }

    @Override
    public void beforeValidate(Connection connection) {
        LOGGER.info("Trying to create schemas with owner binding_group");

        Statement statement = null;
        for (String schema : flywayConfiguration.getSchemas()) {
            try {
                statement = connection.createStatement();
                statement.execute(String.format("create schema if not exists %s authorization binding_group;", schema));
            } catch (Exception e) {
                LOGGER.error(String.format("Error when creating schema %s. IF YOU ARE RUNNING LOCALLY YOU CAN IGNORE THIS ERROR", schema), e);
                try {
                    connection.rollback();
                } catch (SQLException e1) {
                    LOGGER.error("Error rolling back", e1);
                }
            } finally {
                if (statement != null) {
                    try {
                        statement.close();
                    } catch (Exception e) {
                        LOGGER.error("Error closing statement", e);
                    }
                }
            }
        }
    }

    @Override
    public void beforeEachMigrate(Connection connection, MigrationInfo info) {
        LOGGER.info("Trying to set role to 'binding_group'");

        Statement statement = null;
        try {
            statement = connection.createStatement();
            statement.execute(
                    "DO\n" +
                            "$body$\n" +
                            "BEGIN\n" +
                            "   IF EXISTS (\n" +
                            "      SELECT *\n" +
                            "      FROM   pg_catalog.pg_group\n" +
                            "      WHERE  groname = 'binding_group') THEN\n" +
                            "      set role binding_group;\n" +
                            "   END IF;\n" +
                            "END\n" +
                            "$body$;");
        } catch (Exception e) {
            LOGGER.error("Error when setting role to binding_group. IF YOU ARE RUNNING LOCALLY YOU CAN IGNORE THIS ERROR", e);
            try {
                connection.rollback();
            } catch (SQLException e1) {
                LOGGER.error("Error rolling back", e1);
            }
        } finally {
            if (statement != null) {
                try {
                    statement.close();
                } catch (Exception e) {
                    LOGGER.error("Error closing statement", e);
                }
            }
        }
    }

    @Override
    public void setFlywayConfiguration(FlywayConfiguration flywayConfiguration) {
        this.flywayConfiguration = flywayConfiguration;
    }

}