/amps-maven-plugin/src/main/java/com/atlassian/maven/plugins/amps/database/Postgres.java

package com.atlassian.maven.plugins.amps.database;

import com.atlassian.maven.plugins.amps.DataSource;
import com.atlassian.maven.plugins.amps.product.ImportMethod;
import com.google.common.annotations.VisibleForTesting;
import org.apache.commons.lang3.StringUtils;
import org.apache.maven.plugin.MojoExecutionException;
import org.apache.maven.plugin.logging.Log;
import org.codehaus.plexus.util.xml.Xpp3Dom;

import javax.annotation.Nonnull;
import javax.annotation.ParametersAreNonnullByDefault;
import javax.annotation.Nullable;
import java.sql.Driver;
import java.sql.DriverManager;
import java.sql.DriverPropertyInfo;
import java.sql.SQLException;
import java.util.Optional;
import java.util.Properties;

import static java.lang.String.format;
import static org.twdata.maven.mojoexecutor.MojoExecutor.configuration;
import static org.twdata.maven.mojoexecutor.MojoExecutor.element;
import static org.twdata.maven.mojoexecutor.MojoExecutor.name;

@ParametersAreNonnullByDefault
public class Postgres extends AbstractDatabase {

    private static final String ARGUMENT = "argument";
    private static final String DROP_DATABASE = "DROP DATABASE IF EXISTS \"%s\";";
    private static final String DROP_USER = "DROP USER IF EXISTS \"%s\";";
    private static final String CREATE_DATABASE = "CREATE DATABASE \"%s\";";
    private static final String CREATE_USER = "CREATE USER \"%s\" WITH PASSWORD '%s' ;";
    private static final String GRANT_PERMISSION = "ALTER ROLE \"%s\" superuser; ALTER DATABASE \"%s\" OWNER TO \"%s\";";

    @VisibleForTesting
    static final String PGDBNAME = "PGDBNAME";
    @VisibleForTesting
    static final String PGHOST = "PGHOST";
    @VisibleForTesting
    static final String PGPORT = "PGPORT";

    public Postgres(final Log log) {
        super(log, true, "postgres72", "org.postgresql.Driver", "jdbc:postgresql");
    }

    @Override
    protected String dropDatabase(final DataSource dataSource) throws MojoExecutionException {
        return String.format(DROP_DATABASE, getDatabaseName(dataSource));
    }

    @Override
    protected String dropUser(final DataSource dataSource) {
        return String.format(DROP_USER, dataSource.getUsername());
    }

    @Override
    protected String createDatabase(final DataSource dataSource) throws MojoExecutionException {
        return String.format(CREATE_DATABASE, getDatabaseName(dataSource));
    }

    @Override
    protected String createUser(final DataSource dataSource) {
        return String.format(CREATE_USER, dataSource.getUsername(), dataSource.getPassword());
    }

    @Override
    protected String grantPermissionForUser(final DataSource dataSource) throws MojoExecutionException {
        return String.format(GRANT_PERMISSION,
                dataSource.getUsername(), getDatabaseName(dataSource), dataSource.getUsername());
    }

    @Override
    public Xpp3Dom getExecMavenToolImportConfiguration(final DataSource dataSource) throws MojoExecutionException {
        Xpp3Dom configDatabaseTool = null;
        if (ImportMethod.PSQL.equals(ImportMethod.getValueOf(dataSource.getImportMethod()))) {
            configDatabaseTool = configuration(
                    element(name("executable"), "psql"),
                    element(name("arguments"),
                            element(name(ARGUMENT), "-f" + dataSource.getDumpFilePath()),
                            element(name(ARGUMENT), "-U" + dataSource.getUsername()),
                            element(name(ARGUMENT), "-h" + getHostName(dataSource.getUrl())),
                            element(name(ARGUMENT), getDatabaseName(dataSource))
                    )
            );
        }
        return configDatabaseTool;
    }

    private static String getHostName(final String jdbcUrl) {
        return Optional.ofNullable(parseURL(jdbcUrl))
                .map(urlProperties -> urlProperties.getProperty(PGHOST))
                .filter(StringUtils::isNotEmpty)
                .orElseThrow(() -> new IllegalArgumentException(
                        format("Could not parse host name from Postgres URL '%s'", jdbcUrl)));
    }

    /**
     * reference postgres 9.3 documentation: Connecting to the Database http://jdbc.postgresql.org/documentation/93/connect.html
     * With JDBC, a database is represented by a URL (Uniform Resource Locator) takes one of the following forms:
     * jdbc:postgresql:database jdbc:postgresql://host/database jdbc:postgresql://host:port/database
     *
     * @return database name
     */
    @Override
    protected String getDatabaseName(final DataSource dataSource) throws MojoExecutionException {
        String databaseName = StringUtils.EMPTY;
        try {
            Class.forName(dataSource.getDriver());
        } catch (ClassNotFoundException e) {
            throw new MojoExecutionException("Could not load Postgresql database library from AMPS classpath");
        }
        try {
            final String url = dataSource.getUrl();
            final Driver driver = DriverManager.getDriver(url);
            final DriverPropertyInfo[] driverPropertyInfos = driver.getPropertyInfo(url, null);
            if (null != driverPropertyInfos) {
                for (DriverPropertyInfo driverPropertyInfo : driverPropertyInfos) {
                    if (PGDBNAME.equals(driverPropertyInfo.name)) {
                        databaseName = driverPropertyInfo.value;
                        break;
                    }
                }
            }
            // bug from postgresql JDBC <= 9.3
            if (databaseName == null) {
                // apply local monkey-patch
                final Properties driverProps = parseURL(url);
                if (driverProps != null) {
                    databaseName = driverProps.getProperty(PGDBNAME);
                }
            }
        } catch (SQLException e) {
            throw new MojoExecutionException("No suitable driver");
        }
        return databaseName;
    }

    @Override
    @Nonnull
    public Xpp3Dom getSqlMavenCreateConfiguration(final DataSource dataSource) throws MojoExecutionException {
        final String sql = dropDatabase(dataSource) + dropUser(dataSource) + createDatabase(dataSource) +
                createUser(dataSource) + grantPermissionForUser(dataSource);
        log.info("Postgres initialization database sql: " + sql);
        Xpp3Dom pluginConfiguration = systemDatabaseConfiguration(dataSource);
        pluginConfiguration.addChild(
                element(name("sqlCommand"), sql).toDom()
        );
        return pluginConfiguration;
    }

    /**
     * Parses the given Postgres JDBC URL into its component parts.
     *
     * Copied from org.postgresql:postgresql:9.3-1102-jdbc41.jar!/org/postgresql/Driver.java
     *
     * @param url the JDBC URL to parse
     * @return Properties with elements added from the url, or {@code null} if the given URL can't be parsed
     */
    @Nullable
    @VisibleForTesting
    static Properties parseURL(final String url) {
        final Properties urlProps = new Properties();
        String l_urlServer = url;
        String l_urlArgs = "";
        final int l_qPos = url.indexOf('?');
        if (l_qPos != -1) {
            l_urlServer = url.substring(0, l_qPos);
            l_urlArgs = url.substring(l_qPos + 1);
        }
        if (!l_urlServer.startsWith("jdbc:postgresql:")) {
            return null;
        }
        l_urlServer = l_urlServer.substring("jdbc:postgresql:".length());
        if (l_urlServer.startsWith("//")) {
            l_urlServer = l_urlServer.substring(2);
            final int slash = l_urlServer.indexOf('/');
            if (slash == -1) {
                return null;
            }
            urlProps.setProperty(PGDBNAME, l_urlServer.substring(slash + 1));
            final String[] addresses = l_urlServer.substring(0, slash).split(",");
            final StringBuilder hosts = new StringBuilder();
            final StringBuilder ports = new StringBuilder();
            for (final String address : addresses) {
                final int portIdx = address.lastIndexOf(':');
                if (portIdx != -1 && address.lastIndexOf(']') < portIdx) {
                    final String portStr = address.substring(portIdx + 1);
                    try {
                        Integer.parseInt(portStr);
                    } catch (NumberFormatException ex) {
                        return null;
                    }
                    ports.append(portStr);
                    hosts.append(address.subSequence(0, portIdx));
                } else {
                    ports.append("5432");
                    hosts.append(address);
                }
                ports.append(',');
                hosts.append(',');
            }
            ports.setLength(ports.length() - 1);
            hosts.setLength(hosts.length() - 1);
            urlProps.setProperty(PGPORT, ports.toString());
            urlProps.setProperty(PGHOST, hosts.toString());
        } else {
            urlProps.setProperty(PGPORT, "5432");
            urlProps.setProperty(PGHOST, "localhost");
            urlProps.setProperty(PGDBNAME, l_urlServer);
        }
        //parse the args part of the url
        final String[] args = l_urlArgs.split("&");
        for (final String token : args) {
            if (token.length() == 0) {
                continue;
            }
            int l_pos = token.indexOf('=');
            if (l_pos == -1) {
                urlProps.setProperty(token, "");
            } else {
                urlProps.setProperty(token.substring(0, l_pos), token.substring(l_pos + 1));
            }
        }
        return urlProps;
    }
}